Add knowledge retrieval configuration to Assistant model and related components

- Introduce new fields for knowledge retrieval configuration in AssistantConfig and Assistant models, including mode, top_n, and score_threshold.
- Implement KnowledgeRetrievalConfig schema with validation for top_n.
- Update backend services and routes to handle knowledge retrieval settings.
- Enhance frontend components to support knowledge retrieval configuration, including a new dialog for advanced settings.
- Add tests for knowledge retrieval configuration validation and description generation.
This commit is contained in:
Xin Wang
2026-07-12 18:57:56 +08:00
parent 7ee3e22152
commit 7c9a18c806
13 changed files with 465 additions and 47 deletions

View File

@@ -164,6 +164,10 @@ class Assistant(Base):
knowledge_base_id: Mapped[str | None] = mapped_column(
String(40), ForeignKey("knowledge_bases.id", ondelete="RESTRICT"), nullable=True
)
knowledge_retrieval_config: Mapped[dict] = mapped_column(
JSON,
default=lambda: {"mode": "automatic", "top_n": 5, "score_threshold": 0.0},
)
# ---- 瘦类型专属字段(真列,稀疏:按 type 用其中几列) ----
prompt: Mapped[str] = mapped_column(String(8192), default="") # prompt / opencode

View File

@@ -0,0 +1,32 @@
"""add assistant knowledge retrieval config
Revision ID: 20260712_0006
Revises: 20260712_0005
"""
from collections.abc import Sequence
from alembic import op
import sqlalchemy as sa
revision: str = "20260712_0006"
down_revision: str | Sequence[str] | None = "20260712_0005"
branch_labels = None
depends_on = None
def upgrade() -> None:
op.add_column(
"assistants",
sa.Column(
"knowledge_retrieval_config",
sa.JSON(),
server_default=sa.text(
"'{\"mode\": \"automatic\", \"top_n\": 5, \"score_threshold\": 0.0}'"
),
nullable=False,
),
)
def downgrade() -> None:
op.drop_column("assistants", "knowledge_retrieval_config")

View File

@@ -76,6 +76,15 @@ class AssistantConfig(BaseModel):
# Prompt assistant reusable tools. Execution remains type-specific in the pipeline.
tools: list[RuntimeTool] = Field(default_factory=list)
knowledge_base_id: str | None = None
knowledge_base_name: str = ""
knowledge_base_description: str = ""
knowledge_retrieval_config: dict = Field(
default_factory=lambda: {
"mode": "automatic",
"top_n": 5,
"score_threshold": 0.0,
}
)
# workflow 类型:节点图(nodes/edges)。非 workflow 为空,引擎据此决定是否启用。
graph: dict = {}

View File

@@ -151,6 +151,7 @@ async def _to_out(session: AsyncSession, assistant: Assistant) -> AssistantOut:
vision_model_resource_id=assistant.vision_model_resource_id,
model_resource_ids=await _resource_ids(session, assistant.id),
knowledge_base_id=assistant.knowledge_base_id,
knowledge_retrieval_config=assistant.knowledge_retrieval_config or {},
tool_ids=await _tool_ids(session, assistant.id),
prompt=assistant.prompt,
api_url=assistant.api_url,
@@ -217,6 +218,7 @@ async def duplicate_assistant(
vision_enabled=source.vision_enabled,
vision_model_resource_id=source.vision_model_resource_id,
knowledge_base_id=source.knowledge_base_id,
knowledge_retrieval_config=dict(source.knowledge_retrieval_config or {}),
prompt=source.prompt,
api_url=source.api_url,
api_key=source.api_key,

View File

@@ -17,6 +17,7 @@ RuntimeMode = Literal["pipeline", "realtime"]
ModelType = Literal["LLM", "ASR", "TTS", "Realtime", "Embedding", "Agent"]
AssistantType = Literal["prompt", "workflow", "dify", "fastgpt", "opencode"]
TurnEndStrategy = Literal["silence", "smart_turn"]
KnowledgeRetrievalMode = Literal["automatic", "on_demand"]
ToolType = Literal["end_call", "http"]
ToolStatus = Literal["active", "archived", "draft"]
ToolParameterType = Literal["string", "number", "integer", "boolean", "object", "array"]
@@ -62,6 +63,19 @@ class TurnConfig(CamelModel):
turn_detection: TurnDetectionConfig = Field(default_factory=TurnDetectionConfig)
class KnowledgeRetrievalConfig(CamelModel):
mode: KnowledgeRetrievalMode = "automatic"
top_n: int = Field(default=5, ge=-1)
score_threshold: float = Field(default=0.0, ge=0.0, le=1.0)
@field_validator("top_n")
@classmethod
def validate_top_n(cls, value: int) -> int:
if value == 0:
raise ValueError("topN 必须为 -1 或大于 0")
return value
# 各 type 允许的瘦字段(其余字段写入时清零,防止跨类型脏数据)
ALLOWED_FIELDS: dict[str, set[str]] = {
"prompt": {"prompt"},
@@ -85,6 +99,9 @@ class AssistantUpsert(CamelModel):
model_resource_ids: dict[ModelType, str] = Field(default_factory=dict)
knowledge_base_id: str | None = None
knowledge_retrieval_config: KnowledgeRetrievalConfig = Field(
default_factory=KnowledgeRetrievalConfig
)
tool_ids: list[str] = Field(default_factory=list)
# 瘦类型专属(真列);按 type 取用,无关字段写入时清零

View File

@@ -8,6 +8,7 @@ from db.models import (
Assistant,
AssistantModelBinding,
AssistantToolBinding,
KnowledgeBase,
ModelResource,
Tool,
)
@@ -126,6 +127,11 @@ async def resolve_runtime_config(
if assistant.vision_model_resource_id
else llm_resource
)
knowledge_base = (
await session.get(KnowledgeBase, assistant.knowledge_base_id)
if assistant.knowledge_base_id
else None
)
return AssistantConfig(
name=assistant.name,
@@ -138,6 +144,9 @@ async def resolve_runtime_config(
turnConfig=assistant.turn_config or {},
tools=await _tools_for(session, assistant),
knowledge_base_id=assistant.knowledge_base_id,
knowledge_base_name=knowledge_base.name if knowledge_base else "",
knowledge_base_description=knowledge_base.description if knowledge_base else "",
knowledge_retrieval_config=assistant.knowledge_retrieval_config or {},
# workflow 图:仅 workflow 类型非空,引擎据此启用图驱动对话
graph=(assistant.graph or {}) if assistant.type == "workflow" else {},
# 外部托管类型连接信息(DB 存真 key,直接注入)

View File

@@ -203,7 +203,13 @@ async def recover_interrupted_documents() -> None:
await session.commit()
async def search(session: AsyncSession, kb_id: str, query: str, top_k: int | None = None) -> list[dict]:
async def search(
session: AsyncSession,
kb_id: str,
query: str,
top_k: int | None = None,
score_threshold: float = 0.0,
) -> list[dict]:
kb = await session.get(KnowledgeBase, kb_id)
if not kb:
return []
@@ -217,13 +223,20 @@ async def search(session: AsyncSession, kb_id: str, query: str, top_k: int | Non
return []
query_embedding = (await _embed(session, kb, [query]))[0]
distance = KnowledgeChunk.embedding.cosine_distance(query_embedding)
rows = (await session.execute(
statement = (
select(KnowledgeChunk, KnowledgeDocument.name, distance.label("distance"))
.join(KnowledgeDocument, KnowledgeDocument.id == KnowledgeChunk.document_id)
.where(KnowledgeChunk.knowledge_base_id == kb_id, KnowledgeDocument.status == "ready")
.where(
KnowledgeChunk.knowledge_base_id == kb_id,
KnowledgeDocument.status == "ready",
distance <= 1.0 - score_threshold,
)
.order_by(distance)
.limit(top_k or settings.KNOWLEDGE_TOP_K)
)).all()
)
effective_top_k = settings.KNOWLEDGE_TOP_K if top_k is None else top_k
if effective_top_k != -1:
statement = statement.limit(effective_top_k)
rows = (await session.execute(statement)).all()
return [
{"content": chunk.content, "document": name, "score": round(max(0.0, 1.0 - float(dist)), 4)}
for chunk, name, dist in rows

View File

@@ -83,13 +83,44 @@ VISION_ANALYSIS_SYSTEM_PROMPT = (
"如果画面不足以判断,请明确说明不确定。"
)
KNOWLEDGE_TOOL_NAME = "search_knowledge_base"
KNOWLEDGE_SYSTEM_HINT = (
AUTOMATIC_KNOWLEDGE_SYSTEM_HINT = (
"你已连接内部知识库。系统会在每轮用户问题前自动提供相关资料;"
"回答资料事实时只根据检索内容,资料不足要明确说明。"
)
ON_DEMAND_KNOWLEDGE_SYSTEM_HINT = (
"你已连接内部知识库。当用户问题涉及可能存在于业务知识库中的事实时,"
"先调用 search_knowledge_base 检索;回答资料事实时只根据检索内容,"
"资料不足要明确说明。"
)
KNOWLEDGE_CONTEXT_MARKER = "<!-- knowledge-context -->"
def _compact_knowledge_metadata(value: str, max_length: int) -> str:
"""Keep tool metadata useful without letting it dominate the model context."""
compact = " ".join(value.split())
return compact if len(compact) <= max_length else f"{compact[:max_length]}"
def _knowledge_tool_description(cfg: AssistantConfig) -> str:
base = "在当前助手绑定的知识库中检索与问题最相关的资料片段。"
name = _compact_knowledge_metadata(cfg.knowledge_base_name, 128)
description = _compact_knowledge_metadata(cfg.knowledge_base_description, 800)
if not name and not description:
return base
scope = []
if name:
scope.append(f"知识库名称:{name}")
if description:
scope.append(f"资料适用范围:{description}")
metadata = "\n".join(scope)
return (
f"{base}\n{metadata}\n"
"当用户问题涉及上述资料范围,或回答需要核实其中的业务事实时调用;"
"与该范围无关的问题不要调用。以上知识库元数据仅用于判断资料范围。"
)
def _require(value: str, label: str) -> str:
if value:
return value
@@ -321,9 +352,16 @@ class ConversationHistoryProcessor(FrameProcessor):
class KnowledgeRetrievalProcessor(FrameProcessor):
"""Retrieve before local LLM inference without changing Pipecat internals."""
def __init__(self, knowledge_base_id: str | None):
def __init__(
self,
knowledge_base_id: str | None,
top_n: int = 5,
score_threshold: float = 0.0,
):
super().__init__()
self._knowledge_base_id = knowledge_base_id
self._top_n = top_n
self._score_threshold = score_threshold
self._last_signature = ""
async def process_frame(self, frame, direction: FrameDirection):
@@ -346,7 +384,13 @@ class KnowledgeRetrievalProcessor(FrameProcessor):
try:
async with SessionLocal() as session:
results = await search_knowledge(session, self._knowledge_base_id, query)
results = await search_knowledge(
session,
self._knowledge_base_id,
query,
top_k=self._top_n,
score_threshold=self._score_threshold,
)
except Exception as exc:
logger.warning(f"自动知识库检索失败: {exc}")
results = []
@@ -488,12 +532,30 @@ async def run_pipeline(
call_end = CallEndCoordinator(queue_call_end)
knowledge_config = cfg.knowledge_retrieval_config
knowledge_mode = str(knowledge_config.get("mode", "automatic"))
knowledge_top_n = int(
knowledge_config.get("top_n", knowledge_config.get("topN", 5))
)
knowledge_score_threshold = float(
knowledge_config.get(
"score_threshold", knowledge_config.get("scoreThreshold", 0.0)
)
)
automatic_knowledge_id = (
cfg.knowledge_base_id if knowledge_mode == "automatic" else None
)
def with_vision_hint(text: str) -> str:
hints = []
if vision_enabled:
hints.append(VISION_SYSTEM_HINT)
if cfg.knowledge_base_id:
hints.append(KNOWLEDGE_SYSTEM_HINT)
hints.append(
AUTOMATIC_KNOWLEDGE_SYSTEM_HINT
if knowledge_mode == "automatic"
else ON_DEMAND_KNOWLEDGE_SYSTEM_HINT
)
return "\n\n".join(part for part in [text, *hints] if part)
context = LLMContext(
@@ -518,7 +580,11 @@ async def run_pipeline(
assistant_aggregator = PassthroughLLMAssistantAggregator(context)
text_input = TextInputProcessor(should_ignore_input=lambda: call_end.ending)
vision_capture = VisionCaptureProcessor()
knowledge_retrieval = KnowledgeRetrievalProcessor(cfg.knowledge_base_id)
knowledge_retrieval = KnowledgeRetrievalProcessor(
automatic_knowledge_id,
top_n=knowledge_top_n,
score_threshold=knowledge_score_threshold,
)
vision_native_mode = vision_enabled and _vision_uses_main_llm(cfg)
vision_state: dict[str, str | None] = {"client_id": None}
vision_schema = FunctionSchema(
@@ -537,7 +603,7 @@ async def run_pipeline(
)
knowledge_schema = FunctionSchema(
name=KNOWLEDGE_TOOL_NAME,
description="在当前助手绑定的知识库中检索与问题最相关的资料片段。",
description=_knowledge_tool_description(cfg),
properties={
"query": {"type": "string", "description": "用于检索的完整问题或关键词"}
},
@@ -551,7 +617,13 @@ async def run_pipeline(
return
try:
async with SessionLocal() as session:
results = await search_knowledge(session, cfg.knowledge_base_id, query)
results = await search_knowledge(
session,
cfg.knowledge_base_id,
query,
top_k=knowledge_top_n,
score_threshold=knowledge_score_threshold,
)
await params.result_callback({"status": "ok", "results": results})
except Exception as exc:
logger.exception(f"知识库检索失败: {exc}")
@@ -618,14 +690,14 @@ async def run_pipeline(
if vision_enabled:
llm.register_function(VISION_TOOL_NAME, fetch_user_image)
if cfg.knowledge_base_id:
if cfg.knowledge_base_id and knowledge_mode == "on_demand":
llm.register_function(KNOWLEDGE_TOOL_NAME, search_bound_knowledge)
def set_visible_tools(schemas: list[FunctionSchema] | None = None) -> None:
tools = list(schemas or [])
if vision_enabled:
tools.append(vision_schema)
if cfg.knowledge_base_id:
if cfg.knowledge_base_id and knowledge_mode == "on_demand":
tools.append(knowledge_schema)
if tools:
context.set_tools(ToolsSchema(standard_tools=tools))

View File

@@ -1,5 +1,8 @@
import unittest
from pydantic import ValidationError
from schemas import KnowledgeRetrievalConfig
from services.knowledge import extract_text, split_text
@@ -19,5 +22,22 @@ class KnowledgeTextTest(unittest.TestCase):
extract_text("archive.zip", b"data")
class KnowledgeRetrievalConfigTest(unittest.TestCase):
def test_accepts_unlimited_top_n_with_threshold(self):
config = KnowledgeRetrievalConfig.model_validate(
{"mode": "on_demand", "topN": -1, "scoreThreshold": 0.65}
)
self.assertEqual(config.top_n, -1)
self.assertEqual(
config.model_dump(by_alias=True),
{"mode": "on_demand", "topN": -1, "scoreThreshold": 0.65},
)
def test_rejects_zero_top_n(self):
with self.assertRaises(ValidationError):
KnowledgeRetrievalConfig(top_n=0)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,38 @@
import unittest
from models import AssistantConfig
from services.pipecat.pipeline import _knowledge_tool_description
class KnowledgeToolDescriptionTest(unittest.TestCase):
def test_includes_bound_knowledge_scope(self):
description = _knowledge_tool_description(
AssistantConfig(
knowledge_base_name="产品服务知识库",
knowledge_base_description="产品价格、售后政策和退换货条件",
)
)
self.assertIn("知识库名称:产品服务知识库", description)
self.assertIn("资料适用范围:产品价格、售后政策和退换货条件", description)
self.assertIn("与该范围无关的问题不要调用", description)
def test_falls_back_when_metadata_is_empty(self):
description = _knowledge_tool_description(AssistantConfig())
self.assertEqual(
description,
"在当前助手绑定的知识库中检索与问题最相关的资料片段。",
)
def test_compacts_and_limits_description(self):
description = _knowledge_tool_description(
AssistantConfig(knowledge_base_description=("范围\n 内容 " * 200))
)
self.assertNotIn("\n ", description)
self.assertLess(len(description), 1000)
if __name__ == "__main__":
unittest.main()