feat(phase3): Task1 embedding链路验证 - embo-01(1536维)+pgvector检索打通
This commit is contained in:
@@ -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")
|
_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"))
|
_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:
|
if not _DATABASE_URL:
|
||||||
raise RuntimeError(f"[config] DATABASE_URL 未设置。请检查 {_env_path} 是否存在且内容正确。")
|
raise RuntimeError(f"[config] DATABASE_URL 未设置。请检查 {_env_path} 是否存在且内容正确。")
|
||||||
@@ -25,6 +29,8 @@ class Settings:
|
|||||||
DATABASE_URL: str = _DATABASE_URL
|
DATABASE_URL: str = _DATABASE_URL
|
||||||
SECRET_KEY: str = _SECRET_KEY
|
SECRET_KEY: str = _SECRET_KEY
|
||||||
SESSION_MAX_AGE: int = _SESSION_MAX_AGE
|
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()
|
settings = Settings()
|
||||||
|
|||||||
@@ -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()),
|
||||||
|
)
|
||||||
@@ -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)
|
||||||
@@ -0,0 +1,70 @@
|
|||||||
|
"""
|
||||||
|
探路脚本 — 调 MiniMax embo-01,打印原始返回 JSON
|
||||||
|
确认向量字段位置和维度后再写正式 service。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# 加载 .env
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
_env_path = Path(__file__).parent.parent / ".env"
|
||||||
|
load_dotenv(str(_env_path))
|
||||||
|
|
||||||
|
api_key = os.environ.get("MINIMAX_EMBED_API_KEY", "")
|
||||||
|
group_id = os.environ.get("MINIMAX_GROUP_ID", "")
|
||||||
|
|
||||||
|
if not api_key or api_key == "your_api_key_here":
|
||||||
|
print("[ERROR] MINIMAX_EMBED_API_KEY not configured, please edit backend/.env")
|
||||||
|
exit(1)
|
||||||
|
if not group_id or group_id == "your_group_id_here":
|
||||||
|
print("[ERROR] MINIMAX_GROUP_ID not configured, please edit backend/.env")
|
||||||
|
exit(1)
|
||||||
|
|
||||||
|
print(f"API Key (first 4 chars): {api_key[:4]}...")
|
||||||
|
print(f"GroupId: {group_id}")
|
||||||
|
print()
|
||||||
|
|
||||||
|
# 最小调用
|
||||||
|
test_text = "这是一段测试文本,用于验证 embo-01 接口返回结构。"
|
||||||
|
|
||||||
|
print(f"Sending request, test text: {test_text}")
|
||||||
|
print("-" * 60)
|
||||||
|
|
||||||
|
try:
|
||||||
|
resp = httpx.post(
|
||||||
|
"https://api.minimax.chat/v1/embeddings",
|
||||||
|
headers={
|
||||||
|
"Authorization": f"Bearer {api_key}",
|
||||||
|
"GroupId": group_id,
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
},
|
||||||
|
json={"model": "embo-01", "texts": [test_text], "type": "db"},
|
||||||
|
timeout=30.0,
|
||||||
|
)
|
||||||
|
print(f"HTTP status: {resp.status_code}")
|
||||||
|
print()
|
||||||
|
data = resp.json()
|
||||||
|
print(json.dumps(data, indent=2, ensure_ascii=False))
|
||||||
|
|
||||||
|
# 提取向量,验证维度
|
||||||
|
print()
|
||||||
|
print("-" * 60)
|
||||||
|
vectors = data.get("vectors", [])
|
||||||
|
if vectors and len(vectors) > 0:
|
||||||
|
embedding = vectors[0]
|
||||||
|
dim = len(embedding)
|
||||||
|
print(f"[OK] Embedding field: vectors[0]")
|
||||||
|
print(f"[OK] Embedding dimension: {dim}")
|
||||||
|
if dim != 1536:
|
||||||
|
print(f"[STOP] Dimension is NOT 1536! Got {dim} - stopping here")
|
||||||
|
else:
|
||||||
|
print(f"[OK] Dimension correct: 1536")
|
||||||
|
print(f"[OK] API call successful, structure confirmed.")
|
||||||
|
else:
|
||||||
|
print("[WARNING] vectors not found in response")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"[ERROR] Request failed: {e}")
|
||||||
@@ -0,0 +1,79 @@
|
|||||||
|
"""
|
||||||
|
全链路验证脚本 — TPS 知识库 embedding 最小链路
|
||||||
|
|
||||||
|
验证步骤:
|
||||||
|
1. 读取 backend/sample_md/ 下的 5 篇 .md 文件
|
||||||
|
2. 调用 embo-01 转成向量(打印维度)
|
||||||
|
3. 存入 knowledge_items + knowledge_embeddings(打印行数)
|
||||||
|
4. 执行语义检索(打印查询句 + 最相似笔记)
|
||||||
|
5. 查 episodes 表行数(打印,只读不动)
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
from sqlmodel import text
|
||||||
|
|
||||||
|
# 加载 .env
|
||||||
|
_env_path = Path(__file__).parent.parent / ".env"
|
||||||
|
load_dotenv(str(_env_path))
|
||||||
|
|
||||||
|
from app.services.knowledge_service import KnowledgeService
|
||||||
|
from app.db.session import engine
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
print("=" * 60)
|
||||||
|
print("TPS Knowledge Base — Embedding Full链路验证")
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
sample_dir = Path(__file__).parent.parent / "sample_md"
|
||||||
|
md_files = sorted(sample_dir.glob("*.md"))
|
||||||
|
print(f"\n[FIND] Found {len(md_files)} .md files in sample_md/")
|
||||||
|
|
||||||
|
ks = KnowledgeService()
|
||||||
|
|
||||||
|
# 1. 写入知识库
|
||||||
|
print("\n[STEP 1] Storing MD files into knowledge base...")
|
||||||
|
items_stored = []
|
||||||
|
for mf in md_files:
|
||||||
|
title = mf.stem # 文件名(不含扩展名)作为标题
|
||||||
|
content = mf.read_text(encoding="utf-8")
|
||||||
|
item = ks.store_md_file(
|
||||||
|
title=title,
|
||||||
|
content_md=content,
|
||||||
|
source_file_name=mf.name,
|
||||||
|
source_type="manual",
|
||||||
|
)
|
||||||
|
items_stored.append(item)
|
||||||
|
print(f" - Stored: {item.title} (id={item.id})")
|
||||||
|
|
||||||
|
ki_count = ks.get_item_count()
|
||||||
|
ke_count = ks.get_embedding_count()
|
||||||
|
print(f"\n[OK] knowledge_items rows: {ki_count}")
|
||||||
|
print(f"[OK] knowledge_embeddings rows: {ke_count}")
|
||||||
|
|
||||||
|
# 2. 语义检索
|
||||||
|
print("\n[STEP 2] Semantic search test...")
|
||||||
|
query = "五代战斗机的隐身技术有哪些关键要素?"
|
||||||
|
print(f"Query: {query}")
|
||||||
|
results = ks.search_similar(query, top_k=3)
|
||||||
|
print(f"\n[OK] Top 3 similar notes:")
|
||||||
|
for i, r in enumerate(results, 1):
|
||||||
|
print(f" {i}. [{r['similarity']}] {r['title']}")
|
||||||
|
|
||||||
|
# 3. 查 episodes 表行数(只读不动)
|
||||||
|
print("\n[STEP 3] Episodes table (read-only)...")
|
||||||
|
with engine.connect() as conn:
|
||||||
|
result = conn.execute(text("SELECT COUNT(*) FROM episodes"))
|
||||||
|
episode_count = result.scalar()
|
||||||
|
print(f"[OK] episodes table row count: {episode_count}")
|
||||||
|
|
||||||
|
print("\n" + "=" * 60)
|
||||||
|
print("Verification complete.")
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user