feat(phase3): Task1 embedding链路验证 - embo-01(1536维)+pgvector检索打通
This commit is contained in:
simonkoson
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
Embedding 调用服务 — 封装 MiniMax embo-01
|
||||
|
||||
请求格式（确认自探路脚本）：
|
||||
POST /v1/embeddings
|
||||
Body: {"model": "embo-01", "texts": [...], "type": "db"|"query"}
|
||||
响应格式：
|
||||
{"vectors": [[...1536 floats...]], "total_tokens": N, "base_resp": {"status_code": 0, "status_msg": "success"}}
|
||||
"""
|
||||
|
||||
import httpx
|
||||
from typing import List
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
|
||||
class EmbeddingService:
|
||||
"""MiniMax embo-01 embedding 调用封装"""
|
||||
|
||||
def __init__(self):
|
||||
self.api_key = settings.MINIMAX_EMBED_API_KEY
|
||||
self.group_id = settings.MINIMAX_GROUP_ID
|
||||
self.endpoint = "https://api.minimax.chat/v1/embeddings"
|
||||
|
||||
def embed(self, texts: List[str], embed_type: str = "db") -> List[List[float]]:
|
||||
"""
|
||||
调用 embo-01 将文本列表转为向量
|
||||
|
||||
Args:
|
||||
texts: 文本列表（支持批量）
|
||||
embed_type: "db" = 存入库，"query" = 查询
|
||||
Returns:
|
||||
List[List[float]]，每个元素是一组 1536 维向量
|
||||
"""
|
||||
if not self.api_key or self.api_key == "your_api_key_here":
|
||||
raise RuntimeError("MINIMAX_EMBED_API_KEY not configured in .env")
|
||||
if not self.group_id or self.group_id == "your_group_id_here":
|
||||
raise RuntimeError("MINIMAX_GROUP_ID not configured in .env")
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"GroupId": self.group_id,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"model": "embo-01",
|
||||
"texts": texts,
|
||||
"type": embed_type,
|
||||
}
|
||||
|
||||
resp = httpx.post(self.endpoint, headers=headers, json=payload, timeout=60.0)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
# 检查业务错误
|
||||
base_resp = data.get("base_resp", {})
|
||||
if base_resp.get("status_code", 0) != 0:
|
||||
raise RuntimeError(f"Embedding API error: {base_resp.get('status_msg', 'unknown')}")
|
||||
|
||||
vectors = data.get("vectors", [])
|
||||
if not vectors:
|
||||
raise RuntimeError("No vectors returned from embedding API")
|
||||
|
||||
return vectors
|
||||
|
||||
def embed_single(self, text: str, embed_type: str = "db") -> List[float]:
|
||||
"""单文本 embedding，返回 1536 维向量列表（Python list）"""
|
||||
vectors = self.embed([text], embed_type=embed_type)
|
||||
return vectors[0]
|
||||
@@ -0,0 +1,111 @@
|
||||
"""
|
||||
知识库服务 — 写入向量 + 语义检索
|
||||
使用 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)
|
||||
Reference in New Issue
Block a user