feat: add durable document indexing pipeline
This commit is contained in:
gqt
@@ -0,0 +1 @@
|
||||
"""Retrieval-augmented generation infrastructure."""
|
||||
@@ -0,0 +1,63 @@
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
|
||||
from kbqa.api.errors import AppError
|
||||
from kbqa.config import Settings, validate_live_settings
|
||||
|
||||
|
||||
class EmbeddingProvider:
|
||||
def __init__(self, settings: Settings) -> None:
|
||||
self.settings = settings
|
||||
self._client: OpenAIEmbeddings | None = None
|
||||
|
||||
def _get_client(self) -> OpenAIEmbeddings:
|
||||
if self._client is None:
|
||||
validate_live_settings(self.settings)
|
||||
self._client = OpenAIEmbeddings(
|
||||
api_key=self.settings.dashscope_api_key,
|
||||
base_url=self.settings.dashscope_base_url,
|
||||
model=self.settings.embedding_model,
|
||||
dimensions=self.settings.embedding_dim,
|
||||
chunk_size=20,
|
||||
max_retries=2,
|
||||
timeout=60.0,
|
||||
check_embedding_ctx_length=False,
|
||||
)
|
||||
return self._client
|
||||
|
||||
def _validate_vectors(self, vectors: list[list[float]]) -> list[list[float]]:
|
||||
if any(len(vector) != self.settings.embedding_dim for vector in vectors):
|
||||
raise AppError(
|
||||
"EMBEDDING_DIMENSION_MISMATCH",
|
||||
"Embedding 返回的向量维度不正确",
|
||||
502,
|
||||
retryable=True,
|
||||
)
|
||||
return vectors
|
||||
|
||||
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
try:
|
||||
vectors = await self._get_client().aembed_documents(texts)
|
||||
except AppError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise AppError(
|
||||
"EMBEDDING_UPSTREAM_ERROR",
|
||||
"文档向量化失败，请稍后重试",
|
||||
502,
|
||||
retryable=True,
|
||||
) from exc
|
||||
return self._validate_vectors(vectors)
|
||||
|
||||
async def embed_query(self, text: str) -> list[float]:
|
||||
try:
|
||||
vector = await self._get_client().aembed_query(text)
|
||||
except AppError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise AppError(
|
||||
"EMBEDDING_UPSTREAM_ERROR",
|
||||
"查询向量化失败，请稍后重试",
|
||||
502,
|
||||
retryable=True,
|
||||
) from exc
|
||||
return self._validate_vectors([vector])[0]
|
||||
@@ -0,0 +1,56 @@
|
||||
import asyncio
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker
|
||||
|
||||
from kbqa.documents.models import Chunk, Document
|
||||
from kbqa.rag.embeddings import EmbeddingProvider
|
||||
from kbqa.rag.types import SourceCandidate
|
||||
from kbqa.rag.vector_store import MilvusVectorStore
|
||||
|
||||
|
||||
class KnowledgeRetriever:
|
||||
def __init__(
|
||||
self,
|
||||
embeddings: EmbeddingProvider,
|
||||
vector_store: MilvusVectorStore,
|
||||
session_factory: async_sessionmaker,
|
||||
) -> None:
|
||||
self.embeddings = embeddings
|
||||
self.vector_store = vector_store
|
||||
self.session_factory = session_factory
|
||||
|
||||
async def retrieve(self, query: str, top_k: int) -> list[SourceCandidate]:
|
||||
vector = await self.embeddings.embed_query(query)
|
||||
candidate_limit = max(top_k * 3, top_k)
|
||||
hits = await asyncio.to_thread(self.vector_store.search, vector, candidate_limit)
|
||||
if not hits:
|
||||
return []
|
||||
ids = [hit.chunk_id for hit in hits]
|
||||
async with self.session_factory() as session:
|
||||
rows = await session.execute(
|
||||
select(Chunk, Document)
|
||||
.join(Document, Document.id == Chunk.document_id)
|
||||
.where(Chunk.id.in_(ids), Document.status == "indexed")
|
||||
)
|
||||
lookup = {chunk.id: (chunk, document) for chunk, document in rows}
|
||||
sources: list[SourceCandidate] = []
|
||||
for hit in hits:
|
||||
record = lookup.get(hit.chunk_id)
|
||||
if record is None:
|
||||
continue
|
||||
chunk, document = record
|
||||
sources.append(
|
||||
SourceCandidate(
|
||||
chunk_id=chunk.id,
|
||||
document_id=document.id,
|
||||
filename=document.filename,
|
||||
order_index=chunk.order_index,
|
||||
page_number=chunk.page_number,
|
||||
content=chunk.content,
|
||||
score=hit.score,
|
||||
)
|
||||
)
|
||||
if len(sources) == top_k:
|
||||
break
|
||||
return sources
|
||||
@@ -0,0 +1,29 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SourceCandidate:
|
||||
chunk_id: str
|
||||
document_id: str
|
||||
filename: str
|
||||
order_index: int
|
||||
page_number: int | None
|
||||
content: str
|
||||
score: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PreparedChunk:
|
||||
id: str
|
||||
document_id: str
|
||||
order_index: int
|
||||
page_number: int | None
|
||||
content: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VectorHit:
|
||||
chunk_id: str
|
||||
document_id: str
|
||||
order_index: int
|
||||
score: float
|
||||
@@ -0,0 +1,94 @@
|
||||
import json
|
||||
|
||||
from pymilvus import DataType, MilvusClient
|
||||
|
||||
from kbqa.rag.types import PreparedChunk, VectorHit
|
||||
|
||||
|
||||
class MilvusVectorStore:
|
||||
collection_name = "kbqa_chunks"
|
||||
|
||||
def __init__(self, client: MilvusClient, embedding_dim: int) -> None:
|
||||
self.client = client
|
||||
self.embedding_dim = embedding_dim
|
||||
|
||||
def ensure_collection(self) -> None:
|
||||
if self.client.has_collection(self.collection_name):
|
||||
return
|
||||
schema = self.client.create_schema(auto_id=False, enable_dynamic_field=False)
|
||||
schema.add_field(
|
||||
field_name="chunk_id",
|
||||
datatype=DataType.VARCHAR,
|
||||
is_primary=True,
|
||||
max_length=36,
|
||||
)
|
||||
schema.add_field(
|
||||
field_name="document_id",
|
||||
datatype=DataType.VARCHAR,
|
||||
max_length=36,
|
||||
)
|
||||
schema.add_field(field_name="order_index", datatype=DataType.INT64)
|
||||
schema.add_field(
|
||||
field_name="vector",
|
||||
datatype=DataType.FLOAT_VECTOR,
|
||||
dim=self.embedding_dim,
|
||||
)
|
||||
index_params = self.client.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="vector",
|
||||
index_type="AUTOINDEX",
|
||||
metric_type="COSINE",
|
||||
)
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
schema=schema,
|
||||
index_params=index_params,
|
||||
)
|
||||
|
||||
def upsert(self, chunks: list[PreparedChunk], vectors: list[list[float]]) -> None:
|
||||
if len(chunks) != len(vectors):
|
||||
raise ValueError("chunks and vectors must have equal lengths")
|
||||
if not chunks:
|
||||
return
|
||||
data = [
|
||||
{
|
||||
"chunk_id": chunk.id,
|
||||
"document_id": chunk.document_id,
|
||||
"order_index": chunk.order_index,
|
||||
"vector": vector,
|
||||
}
|
||||
for chunk, vector in zip(chunks, vectors, strict=True)
|
||||
]
|
||||
self.client.upsert(collection_name=self.collection_name, data=data)
|
||||
|
||||
def search(self, vector: list[float], limit: int) -> list[VectorHit]:
|
||||
results = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
data=[vector],
|
||||
anns_field="vector",
|
||||
limit=limit,
|
||||
output_fields=["document_id", "order_index"],
|
||||
search_params={"metric_type": "COSINE"},
|
||||
)
|
||||
hits: list[VectorHit] = []
|
||||
for item in results[0] if results else []:
|
||||
entity = item.get("entity", {})
|
||||
hits.append(
|
||||
VectorHit(
|
||||
chunk_id=str(item["id"]),
|
||||
document_id=str(entity["document_id"]),
|
||||
order_index=int(entity["order_index"]),
|
||||
score=float(item["distance"]),
|
||||
)
|
||||
)
|
||||
return hits
|
||||
|
||||
def delete_document(self, document_id: str) -> None:
|
||||
encoded = json.dumps(document_id)
|
||||
self.client.delete(
|
||||
collection_name=self.collection_name,
|
||||
filter=f"document_id == {encoded}",
|
||||
)
|
||||
|
||||
def ping(self) -> None:
|
||||
self.client.list_collections()
|
||||
Reference in New Issue
Block a user