Enhance knowledge base functionality and integrate S3 storage support

- Add new models for `KnowledgeDocument` and `KnowledgeChunk` to manage document ingestion and chunking. - Implement S3-compatible storage integration for knowledge documents, allowing for file uploads and retrieval. - Introduce API endpoints for managing knowledge bases and documents, including creation, deletion, and searching. - Update frontend components to support knowledge base configuration and document management, improving user interaction. - Enhance backend services for knowledge processing and retrieval, ensuring robust handling of document statuses and errors.
2026-07-12 13:58:47 +08:00
parent 01c563a3e7
commit de58f30014
21 changed files with 995 additions and 34 deletions
--- a/backend/services/config_resolver.py
+++ b/backend/services/config_resolver.py
@@ -137,6 +137,7 @@ async def resolve_runtime_config(
        enableInterrupt=assistant.enable_interrupt,
        turnConfig=assistant.turn_config or {},
        tools=await _tools_for(session, assistant),
+        knowledge_base_id=assistant.knowledge_base_id,
        # workflow 图:仅 workflow 类型非空,引擎据此启用图驱动对话
        graph=(assistant.graph or {}) if assistant.type == "workflow" else {},
        # 外部托管类型连接信息(DB 存真 key,直接注入)
--- a/backend/services/knowledge.py
+++ b/backend/services/knowledge.py
@@ -0,0 +1,235 @@
+"""Minimal knowledge ingestion and retrieval, isolated from the voice pipeline."""
+
+import asyncio
+from io import BytesIO
+from pathlib import Path
+import re
+import uuid
+
+import boto3
+from botocore.config import Config
+from botocore.exceptions import ClientError
+from docx import Document as DocxDocument
+from openai import AsyncOpenAI
+from pypdf import PdfReader
+from sqlalchemy import delete, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+import settings
+from db.models import KnowledgeBase, KnowledgeChunk, KnowledgeDocument, ModelResource
+
+
+def _s3_client():
+    return boto3.client(
+        "s3",
+        endpoint_url=settings.S3_ENDPOINT_URL,
+        aws_access_key_id=settings.S3_ACCESS_KEY,
+        aws_secret_access_key=settings.S3_SECRET_KEY,
+        region_name=settings.S3_REGION,
+        config=Config(signature_version="s3v4", s3={"addressing_style": "path"}),
+    )
+
+
+def _ensure_bucket_and_put(key: str, data: bytes, mime_type: str) -> None:
+    client = _s3_client()
+    try:
+        client.head_bucket(Bucket=settings.S3_BUCKET)
+    except ClientError as exc:
+        status = exc.response.get("ResponseMetadata", {}).get("HTTPStatusCode")
+        if status != 404:
+            raise
+        client.create_bucket(Bucket=settings.S3_BUCKET)
+    client.put_object(Bucket=settings.S3_BUCKET, Key=key, Body=data, ContentType=mime_type)
+
+
+def _delete_object(key: str) -> None:
+    _s3_client().delete_object(Bucket=settings.S3_BUCKET, Key=key)
+
+
+def _get_object(key: str) -> bytes:
+    response = _s3_client().get_object(Bucket=settings.S3_BUCKET, Key=key)
+    return response["Body"].read()
+
+
+def extract_text(filename: str, data: bytes) -> str:
+    suffix = Path(filename).suffix.lower()
+    if suffix == ".pdf":
+        return "\n".join(page.extract_text() or "" for page in PdfReader(BytesIO(data)).pages)
+    if suffix == ".docx":
+        document = DocxDocument(BytesIO(data))
+        return "\n".join(paragraph.text for paragraph in document.paragraphs)
+    if suffix in {".txt", ".md", ".csv", ".json", ".html", ".htm"}:
+        return data.decode("utf-8", errors="replace")
+    raise ValueError("暂仅支持 PDF、DOCX、TXT、Markdown、CSV、JSON 和 HTML 文件")
+
+
+def split_text(text: str, chunk_size: int = 800, overlap: int = 120) -> list[str]:
+    normalized = re.sub(r"[ \t]+", " ", text).strip()
+    if not normalized:
+        return []
+    chunks: list[str] = []
+    start = 0
+    while start < len(normalized):
+        end = min(start + chunk_size, len(normalized))
+        if end < len(normalized):
+            boundary = max(normalized.rfind("\n", start, end), normalized.rfind("。", start, end))
+            if boundary > start + chunk_size // 2:
+                end = boundary + 1
+        chunks.append(normalized[start:end].strip())
+        if end >= len(normalized):
+            break
+        start = max(end - overlap, start + 1)
+    return [chunk for chunk in chunks if chunk]
+
+
+async def _embedding_resource(session: AsyncSession, kb: KnowledgeBase) -> ModelResource:
+    resource = (
+        await session.get(ModelResource, kb.embedding_model_resource_id)
+        if kb.embedding_model_resource_id
+        else None
+    )
+    if not resource or resource.capability != "Embedding" or not resource.enabled:
+        raise ValueError("请先为知识库选择已启用的 Embedding 模型")
+    return resource
+
+
+async def _embed(session: AsyncSession, kb: KnowledgeBase, texts: list[str]) -> list[list[float]]:
+    resource = await _embedding_resource(session, kb)
+    values, secrets = resource.values or {}, resource.secrets or {}
+    client = AsyncOpenAI(
+        api_key=str(secrets.get("apiKey") or ""),
+        base_url=str(values.get("apiUrl") or "") or None,
+    )
+    try:
+        response = await client.embeddings.create(
+            model=str(values.get("modelId") or "text-embedding-3-small"), input=texts
+        )
+        return [item.embedding for item in response.data]
+    finally:
+        await client.close()
+
+
+async def create_document(
+    session: AsyncSession,
+    kb: KnowledgeBase,
+    *,
+    name: str,
+    source_type: str,
+    raw_data: bytes,
+    mime_type: str = "text/plain",
+) -> KnowledgeDocument:
+    document_id = f"doc_{uuid.uuid4().hex[:12]}"
+    safe_name = Path(name).name
+    extension = ".txt" if source_type == "text" else Path(safe_name).suffix
+    storage_key = f"knowledge/{kb.id}/{document_id}/source{extension}"
+    await asyncio.to_thread(_ensure_bucket_and_put, storage_key, raw_data, mime_type)
+
+    document = KnowledgeDocument(
+        id=document_id,
+        knowledge_base_id=kb.id,
+        name=name,
+        source_type=source_type,
+        mime_type=mime_type,
+        storage_key=storage_key,
+        size_bytes=len(raw_data),
+        status="pending",
+    )
+    session.add(document)
+    await session.commit()
+    await session.refresh(document)
+    return document
+
+
+async def process_document(document_id: str) -> None:
+    """Process one persisted source. Safe to call again after a failure."""
+    from db.session import SessionLocal
+
+    async with SessionLocal() as session:
+        document = await session.get(KnowledgeDocument, document_id)
+        if not document:
+            return
+        kb = await session.get(KnowledgeBase, document.knowledge_base_id)
+        if not kb or not document.storage_key:
+            return
+        document.status = "processing"
+        document.error_message = ""
+        document.chunk_count = 0
+        await session.execute(delete(KnowledgeChunk).where(KnowledgeChunk.document_id == document.id))
+        await session.commit()
+
+        try:
+            data = await asyncio.to_thread(_get_object, document.storage_key)
+            text = (
+                data.decode("utf-8", errors="replace")
+                if document.source_type == "text"
+                else await asyncio.to_thread(extract_text, document.name, data)
+            )
+            chunks = split_text(text)
+            if not chunks:
+                raise ValueError("文档中没有可入库的文字")
+            embeddings: list[list[float]] = []
+            for start in range(0, len(chunks), 64):
+                embeddings.extend(await _embed(session, kb, chunks[start : start + 64]))
+            if len(embeddings) != len(chunks):
+                raise ValueError("Embedding 服务返回的向量数量与分块数量不一致")
+            for index, (content, embedding) in enumerate(zip(chunks, embeddings)):
+                session.add(KnowledgeChunk(
+                    id=f"chunk_{uuid.uuid4().hex[:12]}", knowledge_base_id=kb.id,
+                    document_id=document.id, chunk_index=index, content=content, embedding=embedding,
+                ))
+            document.chunk_count = len(chunks)
+            document.status = "ready"
+            await session.commit()
+        except Exception as exc:
+            await session.rollback()
+            document = await session.get(KnowledgeDocument, document_id)
+            if document:
+                document.status = "failed"
+                document.error_message = str(exc)[:2048]
+                document.chunk_count = 0
+                await session.commit()
+
+
+async def recover_interrupted_documents() -> None:
+    """Make restart-interrupted work visible and retryable."""
+    from db.session import SessionLocal
+
+    async with SessionLocal() as session:
+        await session.execute(
+            update(KnowledgeDocument)
+            .where(KnowledgeDocument.status.in_(["pending", "processing"]))
+            .values(status="failed", error_message="服务重启导致处理被中断，请点击重试")
+        )
+        await session.commit()
+
+
+async def search(session: AsyncSession, kb_id: str, query: str, top_k: int | None = None) -> list[dict]:
+    kb = await session.get(KnowledgeBase, kb_id)
+    if not kb:
+        return []
+    has_chunks = (await session.execute(
+        select(KnowledgeChunk.id)
+        .join(KnowledgeDocument, KnowledgeDocument.id == KnowledgeChunk.document_id)
+        .where(KnowledgeChunk.knowledge_base_id == kb_id, KnowledgeDocument.status == "ready")
+        .limit(1)
+    )).scalar_one_or_none()
+    if not has_chunks:
+        return []
+    query_embedding = (await _embed(session, kb, [query]))[0]
+    distance = KnowledgeChunk.embedding.cosine_distance(query_embedding)
+    rows = (await session.execute(
+        select(KnowledgeChunk, KnowledgeDocument.name, distance.label("distance"))
+        .join(KnowledgeDocument, KnowledgeDocument.id == KnowledgeChunk.document_id)
+        .where(KnowledgeChunk.knowledge_base_id == kb_id, KnowledgeDocument.status == "ready")
+        .order_by(distance)
+        .limit(top_k or settings.KNOWLEDGE_TOP_K)
+    )).all()
+    return [
+        {"content": chunk.content, "document": name, "score": round(max(0.0, 1.0 - float(dist)), 4)}
+        for chunk, name, dist in rows
+    ]
+
+
+async def delete_storage_object(document: KnowledgeDocument) -> None:
+    if document.storage_key:
+        await asyncio.to_thread(_delete_object, document.storage_key)
--- a/backend/services/pipecat/pipeline.py
+++ b/backend/services/pipecat/pipeline.py
@@ -27,6 +27,8 @@ from services.pipecat.service_factory import (
    create_stt,
    create_tts,
 )
+from db.session import SessionLocal
+from services.knowledge import search as search_knowledge

 from pipecat.adapters.schemas.function_schema import FunctionSchema
 from pipecat.adapters.schemas.tools_schema import ToolsSchema
@@ -36,6 +38,7 @@ from pipecat.frames.frames import (
    InterruptionFrame,
    LLMFullResponseEndFrame,
    LLMFullResponseStartFrame,
+    LLMContextFrame,
    LLMTextFrame,
    LLMMessagesAppendFrame,
    OutputTransportMessageUrgentFrame,
@@ -79,6 +82,12 @@ VISION_ANALYSIS_SYSTEM_PROMPT = (
    "你是一个视觉理解模型。请只根据图片内容和用户问题给出准确、简洁的中文观察结果。"
    "如果画面不足以判断,请明确说明不确定。"
 )
+KNOWLEDGE_TOOL_NAME = "search_knowledge_base"
+KNOWLEDGE_SYSTEM_HINT = (
+    "你已连接内部知识库。系统会在每轮用户问题前自动提供相关资料；"
+    "回答资料事实时只根据检索内容，资料不足要明确说明。"
+)
+KNOWLEDGE_CONTEXT_MARKER = "<!-- knowledge-context -->"


 def _require(value: str, label: str) -> str:
@@ -309,6 +318,54 @@ class ConversationHistoryProcessor(FrameProcessor):
            await self._recorder.record_transport_message(frame.message)


+class KnowledgeRetrievalProcessor(FrameProcessor):
+    """Retrieve before local LLM inference without changing Pipecat internals."""
+
+    def __init__(self, knowledge_base_id: str | None):
+        super().__init__()
+        self._knowledge_base_id = knowledge_base_id
+        self._last_signature = ""
+
+    async def process_frame(self, frame, direction: FrameDirection):
+        await super().process_frame(frame, direction)
+        if not self._knowledge_base_id or not isinstance(frame, LLMContextFrame):
+            await self.push_frame(frame, direction)
+            return
+
+        messages = frame.context.get_messages()
+        user_messages = [message for message in messages if message.get("role") == "user"]
+        if not user_messages:
+            await self.push_frame(frame, direction)
+            return
+        query = str(user_messages[-1].get("content") or "").strip()
+        signature = f"{len(user_messages)}:{query}"
+        if not query or signature == self._last_signature:
+            await self.push_frame(frame, direction)
+            return
+        self._last_signature = signature
+
+        try:
+            async with SessionLocal() as session:
+                results = await search_knowledge(session, self._knowledge_base_id, query)
+        except Exception as exc:
+            logger.warning(f"自动知识库检索失败: {exc}")
+            results = []
+
+        sources = "\n\n".join(
+            f"[{index + 1}] 来源：{item['document']}（相关度 {item['score']}）\n{item['content']}"
+            for index, item in enumerate(results)
+        ) or "未检索到相关资料。"
+        block = f"{KNOWLEDGE_CONTEXT_MARKER}\n当前问题的知识库检索结果：\n{sources}"
+        system_message = next((message for message in messages if message.get("role") == "system"), None)
+        if system_message is None:
+            messages.insert(0, {"role": "system", "content": block})
+        else:
+            content = str(system_message.get("content") or "")
+            base = content.split(KNOWLEDGE_CONTEXT_MARKER, 1)[0].rstrip()
+            system_message["content"] = f"{base}\n\n{block}" if base else block
+        await self.push_frame(frame, direction)
+
+
 class PassthroughLLMAssistantAggregator(LLMAssistantAggregator):
    """聚合 LLM 回复进上下文,同时继续把回复帧交给下游 TTS。"""

@@ -432,11 +489,12 @@ async def run_pipeline(
    call_end = CallEndCoordinator(queue_call_end)

    def with_vision_hint(text: str) -> str:
-        if not vision_enabled:
-            return text
-        if not text:
-            return VISION_SYSTEM_HINT
-        return f"{text}\n\n{VISION_SYSTEM_HINT}"
+        hints = []
+        if vision_enabled:
+            hints.append(VISION_SYSTEM_HINT)
+        if cfg.knowledge_base_id:
+            hints.append(KNOWLEDGE_SYSTEM_HINT)
+        return "\n\n".join(part for part in [text, *hints] if part)

    context = LLMContext(
        messages=[{"role": "system", "content": with_vision_hint(system_content)}]
@@ -460,6 +518,7 @@ 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)
    vision_native_mode = vision_enabled and _vision_uses_main_llm(cfg)
    vision_state: dict[str, str | None] = {"client_id": None}
    vision_schema = FunctionSchema(
@@ -476,6 +535,27 @@ async def run_pipeline(
        },
        required=["question"],
    )
+    knowledge_schema = FunctionSchema(
+        name=KNOWLEDGE_TOOL_NAME,
+        description="在当前助手绑定的知识库中检索与问题最相关的资料片段。",
+        properties={
+            "query": {"type": "string", "description": "用于检索的完整问题或关键词"}
+        },
+        required=["query"],
+    )
+
+    async def search_bound_knowledge(params: FunctionCallParams):
+        query = str(params.arguments.get("query") or "").strip()
+        if not query or not cfg.knowledge_base_id:
+            await params.result_callback({"status": "error", "message": "检索问题为空或未绑定知识库"})
+            return
+        try:
+            async with SessionLocal() as session:
+                results = await search_knowledge(session, cfg.knowledge_base_id, query)
+            await params.result_callback({"status": "ok", "results": results})
+        except Exception as exc:
+            logger.exception(f"知识库检索失败: {exc}")
+            await params.result_callback({"status": "error", "message": "知识库检索暂时不可用"})

    async def fetch_user_image(params: FunctionCallParams):
        question = str(params.arguments.get("question") or "请描述当前画面。")
@@ -538,11 +618,15 @@ async def run_pipeline(

    if vision_enabled:
        llm.register_function(VISION_TOOL_NAME, fetch_user_image)
+    if cfg.knowledge_base_id:
+        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:
+            tools.append(knowledge_schema)
        if tools:
            context.set_tools(ToolsSchema(standard_tools=tools))
        else:
@@ -561,6 +645,7 @@ async def run_pipeline(
            text_input,
            stt,
            user_aggregator,
+            knowledge_retrieval,
            llm,
            # Aggregate the streamed LLM text before TTS. On interruption,
            # Pipecat commits the generated prefix immediately instead of