113 lines
5.8 KiB
Python
113 lines
5.8 KiB
Python
from functools import lru_cache
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import cast, delete, func, or_, select
|
|
|
|
from .config import get_settings
|
|
from .knowledge_pages import page_text, relative_image_path
|
|
from .models import KnowledgeChunk, KnowledgeDocument, KnowledgePage
|
|
|
|
|
|
INDEX_VERSION = 2
|
|
|
|
|
|
def split_pages(pages: list[dict], *, size=650, overlap=100):
|
|
for page in pages:
|
|
text = page.get("text", "").strip()
|
|
blocks = [text] if text else []
|
|
for table in page.get("tables", []):
|
|
blocks.append("表格: " + str(table.get("rows", table)))
|
|
for block in blocks:
|
|
start = 0
|
|
while start < len(block):
|
|
yield {"page": page["page_number"], "section": page.get("section", ""), "text": block[start:start + size]}
|
|
if start + size >= len(block):
|
|
break
|
|
start += size - overlap
|
|
|
|
|
|
@lru_cache(maxsize=2)
|
|
def local_encoder(model_path: str):
|
|
path = Path(model_path)
|
|
if not path.is_dir() or not any(path.iterdir()):
|
|
raise RuntimeError(f"离线嵌入模型不存在: {path}；请挂载本地模型包")
|
|
from sentence_transformers import SentenceTransformer
|
|
return SentenceTransformer(str(path), local_files_only=True, trust_remote_code=False)
|
|
|
|
|
|
def embed(texts: list[str], *, model_path: str | None = None):
|
|
if not texts:
|
|
return []
|
|
model = local_encoder(model_path or str(get_settings().embedding_model_path))
|
|
vectors = model.encode(texts, normalize_embeddings=True)
|
|
output = []
|
|
for vector in vectors:
|
|
values = vector.tolist()
|
|
if len(values) > 1024:
|
|
raise RuntimeError("嵌入维度超过pgvector列容量1024")
|
|
output.append(values + [0.0] * (1024 - len(values)))
|
|
return output
|
|
|
|
|
|
def ingest(db, *, course_version_id: int, project_id: int | None, submission_file_id: int | None,
|
|
title: str, category: str, source_hash: str, pages: list[dict], embedder=embed):
|
|
existing = db.scalar(select(KnowledgeDocument).where(KnowledgeDocument.course_version_id == course_version_id,
|
|
KnowledgeDocument.project_id == project_id, KnowledgeDocument.submission_file_id == submission_file_id,
|
|
KnowledgeDocument.title == title))
|
|
if existing and existing.source_hash == source_hash and existing.index_version == INDEX_VERSION and existing.status == "ready":
|
|
return existing
|
|
chunks = list(split_pages(pages))
|
|
vectors = embedder([chunk["text"] for chunk in chunks])
|
|
if len(vectors) != len(chunks):
|
|
raise RuntimeError("嵌入向量数量与知识片段数量不一致")
|
|
if existing:
|
|
db.execute(delete(KnowledgePage).where(KnowledgePage.document_id == existing.id))
|
|
db.execute(delete(KnowledgeChunk).where(KnowledgeChunk.document_id == existing.id))
|
|
doc = existing
|
|
doc.document_type, doc.source_hash = category, source_hash
|
|
else:
|
|
doc = KnowledgeDocument(course_version_id=course_version_id, project_id=project_id, submission_file_id=submission_file_id,
|
|
title=title, document_type=category, source_hash=source_hash)
|
|
db.add(doc)
|
|
doc.index_version, doc.status = INDEX_VERSION, "indexing"
|
|
db.flush()
|
|
for source in pages:
|
|
db.add(KnowledgePage(document_id=doc.id, page=source["page_number"], text=page_text(source),
|
|
image_path=relative_image_path(source.get("rendered_image_path"))))
|
|
for chunk, vector in zip(chunks, vectors):
|
|
db.add(KnowledgeChunk(document_id=doc.id, course_version_id=course_version_id, project_id=project_id,
|
|
document_type=category, page=chunk["page"], section=chunk["section"], text=chunk["text"], embedding=vector))
|
|
doc.status = "ready"
|
|
db.commit()
|
|
return doc
|
|
|
|
|
|
def retrieve(db, *, course_version_id: int, project_id: int, query: str, category: str | None = None,
|
|
limit=10, embedder=embed, include_public=True, public_only=False):
|
|
"""Hybrid retrieval with project isolation and an optional public course namespace."""
|
|
scope = (KnowledgeChunk.project_id.is_(None) if public_only else
|
|
or_(KnowledgeChunk.project_id == project_id, KnowledgeChunk.project_id.is_(None)) if include_public else
|
|
KnowledgeChunk.project_id == project_id)
|
|
filters = [KnowledgeChunk.course_version_id == course_version_id, scope]
|
|
if category:
|
|
filters.append(KnowledgeChunk.document_type == category)
|
|
vector = embedder([query])[0]
|
|
if db.bind.dialect.name == "postgresql":
|
|
from pgvector.sqlalchemy import Vector
|
|
semantic = (1 - cast(KnowledgeChunk.embedding, Vector(1024)).cosine_distance(vector)).label("semantic")
|
|
lexical = func.ts_rank(func.to_tsvector("simple", KnowledgeChunk.text), func.plainto_tsquery("simple", query)).label("lexical")
|
|
rows = db.execute(select(KnowledgeChunk, semantic, lexical).where(*filters).order_by(semantic.desc(), lexical.desc()).limit(limit * 3)).all()
|
|
ranked = [(row[0], float(row.semantic or 0) * .7 + float(row.lexical or 0) * .3) for row in rows]
|
|
else:
|
|
chunks = db.scalars(select(KnowledgeChunk).where(*filters)).all()
|
|
terms = [query[i:i+2] for i in range(max(0, len(query)-1)) if query[i:i+2].strip()]
|
|
ranked = []
|
|
for chunk in chunks:
|
|
semantic = sum(a*b for a, b in zip(chunk.embedding or [], vector))
|
|
lexical = sum(term in chunk.text for term in terms) / max(len(terms), 1)
|
|
ranked.append((chunk, semantic * .7 + lexical * .3))
|
|
ranked.sort(key=lambda pair: pair[1], reverse=True)
|
|
return [{"chunk_id": chunk.id, "document_id": chunk.document_id, "project_id": chunk.project_id,
|
|
"document_type": chunk.document_type, "page": chunk.page, "section": chunk.section,
|
|
"text": chunk.text, "rank": float(rank)} for chunk, rank in ranked[:limit]]
|