111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
|
"""
|
|||
|
|
知识库服务 — 写入向量 + 语义检索
|
|||
|
|
使用 pgvector 原生 SQL 向量检索（<=> 余弦距离算子），不在 Python 侧计算
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from typing import Optional
|
|||
|
|
|
|||
|
|
from sqlalchemy import text
|
|||
|
|
from sqlmodel import Session, select
|
|||
|
|
from pgvector.sqlalchemy import Vector
|
|||
|
|
|
|||
|
|
from app.models.knowledge import KnowledgeItem, KnowledgeEmbedding
|
|||
|
|
from app.services.embedding_service import EmbeddingService
|
|||
|
|
from app.db.session import engine
|
|||
|
|
|
|||
|
|
|
|||
|
|
class KnowledgeService:
|
|||
|
|
"""知识库 CRUD + 语义检索"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self.embedder = EmbeddingService()
|
|||
|
|
|
|||
|
|
def store_md_file(
|
|||
|
|
self,
|
|||
|
|
title: str,
|
|||
|
|
content_md: str,
|
|||
|
|
source_file_name: Optional[str] = None,
|
|||
|
|
source_type: str = "manual",
|
|||
|
|
author: Optional[str] = None,
|
|||
|
|
) -> KnowledgeItem:
|
|||
|
|
"""
|
|||
|
|
读取一篇 md 内容，调用 embo-01 拿到向量，写入 knowledge_items + knowledge_embeddings
|
|||
|
|
"""
|
|||
|
|
# 调用 embedding（type="db" 表示存入知识库）
|
|||
|
|
embedding_list = self.embedder.embed_single(content_md, embed_type="db")
|
|||
|
|
|
|||
|
|
with Session(engine) as session:
|
|||
|
|
# 写入 knowledge_items
|
|||
|
|
item = KnowledgeItem(
|
|||
|
|
title=title,
|
|||
|
|
content_md=content_md,
|
|||
|
|
source_type=source_type,
|
|||
|
|
source_file_name=source_file_name,
|
|||
|
|
author=author,
|
|||
|
|
)
|
|||
|
|
session.add(item)
|
|||
|
|
session.flush() # 拿到 id
|
|||
|
|
|
|||
|
|
# 写入 knowledge_embeddings（单 chunk，chunk_index=0）
|
|||
|
|
# 直接传 list，pgvector.sqlalchemy.Vector 会自动处理转换
|
|||
|
|
emb = KnowledgeEmbedding(
|
|||
|
|
knowledge_id=item.id,
|
|||
|
|
chunk_index=0,
|
|||
|
|
chunk_text=content_md,
|
|||
|
|
embedding=embedding_list,
|
|||
|
|
)
|
|||
|
|
session.add(emb)
|
|||
|
|
session.commit()
|
|||
|
|
session.refresh(item)
|
|||
|
|
return item
|
|||
|
|
|
|||
|
|
def search_similar(self, query_text: str, top_k: int = 5) -> list[dict]:
|
|||
|
|
"""
|
|||
|
|
语义检索：查询句转为向量，用 SQL 余弦距离（<=>）在数据库层检索
|
|||
|
|
返回 top_k 条相似笔记，含相似度分数
|
|||
|
|
"""
|
|||
|
|
# 查询向量（type="query"）
|
|||
|
|
query_vector = self.embedder.embed_single(query_text, embed_type="query")
|
|||
|
|
|
|||
|
|
# 将向量列表转为 pgvector SQL 字符串格式
|
|||
|
|
vec_str = "[" + ",".join(str(v) for v in query_vector) + "]"
|
|||
|
|
|
|||
|
|
with Session(engine) as session:
|
|||
|
|
# pgvector 原生 SQL：<=> 是余弦距离，1 - 距离 = 相似度
|
|||
|
|
# 用字符串插注向量，避免 psycopg2 参数化问题
|
|||
|
|
sql = f"""
|
|||
|
|
SELECT
|
|||
|
|
ki.id,
|
|||
|
|
ki.title,
|
|||
|
|
ki.source_type,
|
|||
|
|
1 - (ke.embedding <=> '{vec_str}'::vector) AS similarity
|
|||
|
|
FROM knowledge_embeddings ke
|
|||
|
|
JOIN knowledge_items ki ON ke.knowledge_id = ki.id
|
|||
|
|
WHERE ke.chunk_index = 0
|
|||
|
|
ORDER BY ke.embedding <=> '{vec_str}'::vector
|
|||
|
|
LIMIT {top_k}
|
|||
|
|
"""
|
|||
|
|
stmt = text(sql)
|
|||
|
|
rows = session.execute(stmt).all()
|
|||
|
|
|
|||
|
|
results = []
|
|||
|
|
for row in rows:
|
|||
|
|
results.append({
|
|||
|
|
"id": row.id,
|
|||
|
|
"title": row.title,
|
|||
|
|
"source_type": row.source_type,
|
|||
|
|
"similarity": round(row.similarity, 4),
|
|||
|
|
})
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
def get_item_count(self) -> int:
|
|||
|
|
"""返回 knowledge_items 表行数"""
|
|||
|
|
with Session(engine) as session:
|
|||
|
|
count = session.exec(select(KnowledgeItem)).all()
|
|||
|
|
return len(count)
|
|||
|
|
|
|||
|
|
def get_embedding_count(self) -> int:
|
|||
|
|
"""返回 knowledge_embeddings 表行数"""
|
|||
|
|
with Session(engine) as session:
|
|||
|
|
count = session.exec(select(KnowledgeEmbedding)).all()
|
|||
|
|
return len(count)
|