simonkoson
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) |