"""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)