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.
This commit is contained in:
Xin Wang
235
backend/services/knowledge.py
Normal file
235
backend/services/knowledge.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user