215 lines
8.0 KiB
Python
215 lines
8.0 KiB
Python
"""Minimal knowledge ingestion and retrieval, isolated from the voice pipeline."""
|
|
|
|
import asyncio
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
import re
|
|
import uuid
|
|
|
|
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
|
|
from services.object_storage import delete_object, get_object, put_object
|
|
|
|
|
|
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(put_object, 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,
|
|
score_threshold: float = 0.0,
|
|
) -> 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)
|
|
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",
|
|
distance <= 1.0 - score_threshold,
|
|
)
|
|
.order_by(distance)
|
|
)
|
|
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
|
|
]
|
|
|
|
|
|
async def delete_storage_object(document: KnowledgeDocument) -> None:
|
|
if document.storage_key:
|
|
await asyncio.to_thread(delete_object, document.storage_key)
|