feat(phase3): Task1 embedding链路验证 - embo-01(1536维)+pgvector检索打通

This commit is contained in:
simonkoson
2026-05-26 10:33:25 +08:00
parent d40d46a434
commit 1325807257
6 changed files with 390 additions and 0 deletions
+6
View File
@@ -16,6 +16,10 @@ _DATABASE_URL = os.environ.get("DATABASE_URL")
_SECRET_KEY = os.environ.get("SECRET_KEY", "change-me-to-a-random-string-in-production")
_SESSION_MAX_AGE = int(os.environ.get("SESSION_MAX_AGE", "86400"))
# MiniMax Embedding API 凭证
_MINIMAX_EMBED_API_KEY = os.environ.get("MINIMAX_EMBED_API_KEY", "")
_MINIMAX_GROUP_ID = os.environ.get("MINIMAX_GROUP_ID", "")
# 验证必需配置
if not _DATABASE_URL:
raise RuntimeError(f"[config] DATABASE_URL 未设置。请检查 {_env_path} 是否存在且内容正确。")
@@ -25,6 +29,8 @@ class Settings:
DATABASE_URL: str = _DATABASE_URL
SECRET_KEY: str = _SECRET_KEY
SESSION_MAX_AGE: int = _SESSION_MAX_AGE
MINIMAX_EMBED_API_KEY: str = _MINIMAX_EMBED_API_KEY
MINIMAX_GROUP_ID: str = _MINIMAX_GROUP_ID
settings = Settings()
+55
View File
@@ -0,0 +1,55 @@
"""
知识库模型 — SQLModel
对应 knowledge_items 和 knowledge_embeddings 两张表
embedding 字段使用 pgvector.Vector(对应 PG vector(1536)
"""
from datetime import datetime, date
from typing import Optional, Any
from sqlalchemy import Column, DateTime as SADateTime, Text, Integer
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.sql import func as sa_func
from sqlmodel import Field, SQLModel
from pgvector.sqlalchemy import Vector
class KnowledgeItem(SQLModel, table=True):
"""知识库条目(knowledge_items"""
__tablename__ = "knowledge_items"
id: Optional[int] = Field(default=None, primary_key=True)
title: str = Field(max_length=300)
content_md: Optional[str] = Field(default=None)
source_type: str = Field(default="manual", max_length=30)
source_file_name: Optional[str] = Field(default=None, max_length=300)
source_url: Optional[str] = Field(default=None, max_length=1000)
author: Optional[str] = Field(default=None, max_length=100)
publish_date: Optional[date] = Field(default=None)
tags: Any = Field(default=None, sa_column=Column(JSONB, default=[]))
related_entities: Any = Field(default=None, sa_column=Column(JSONB, default=[]))
related_concepts: Any = Field(default=None, sa_column=Column(JSONB, default=[]))
created_at: datetime | None = Field(
default=None,
sa_column=Column(SADateTime(timezone=True), nullable=False, server_default=sa_func.now()),
)
updated_at: datetime | None = Field(
default=None,
sa_column=Column(SADateTime(timezone=True), nullable=False, server_default=sa_func.now()),
)
class KnowledgeEmbedding(SQLModel, table=True):
"""知识库向量(knowledge_embeddings"""
__tablename__ = "knowledge_embeddings"
id: Optional[int] = Field(default=None, primary_key=True)
knowledge_id: int = Field(foreign_key="knowledge_items.id", index=True)
chunk_index: int = Field(default=0)
chunk_text: str = Field(sa_column=Column(Text, nullable=False))
embedding: Any = Field(sa_column=Column(Vector(1536), nullable=False))
created_at: datetime | None = Field(
default=None,
sa_column=Column(SADateTime(timezone=True), nullable=False, server_default=sa_func.now()),
)
+69
View File
@@ -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]
+111
View File
@@ -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
"""
# 调用 embeddingtype="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(单 chunkchunk_index=0
# 直接传 listpgvector.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)