de58f30014
- 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.
236 lines
8.8 KiB
Python
236 lines
8.8 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
|
|
|
|
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)
|