From 13258072573b35f6a08624597ff8b692b622bac3 Mon Sep 17 00:00:00 2001 From: simonkoson <28867558@qq.com> Date: Tue, 26 May 2026 10:33:25 +0800 Subject: [PATCH] =?UTF-8?q?feat(phase3):=20Task1=20embedding=E9=93=BE?= =?UTF-8?q?=E8=B7=AF=E9=AA=8C=E8=AF=81=20-=20embo-01(1536=E7=BB=B4)+pgvect?= =?UTF-8?q?or=E6=A3=80=E7=B4=A2=E6=89=93=E9=80=9A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/core/config.py | 6 ++ backend/app/models/knowledge.py | 55 +++++++++++ backend/app/services/embedding_service.py | 69 ++++++++++++++ backend/app/services/knowledge_service.py | 111 ++++++++++++++++++++++ backend/scripts/test_embo01_api.py | 70 ++++++++++++++ backend/scripts/verify_embedding.py | 79 +++++++++++++++ 6 files changed, 390 insertions(+) create mode 100644 backend/app/models/knowledge.py create mode 100644 backend/app/services/embedding_service.py create mode 100644 backend/app/services/knowledge_service.py create mode 100644 backend/scripts/test_embo01_api.py create mode 100644 backend/scripts/verify_embedding.py diff --git a/backend/app/core/config.py b/backend/app/core/config.py index 65258cb..5c48b13 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -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() diff --git a/backend/app/models/knowledge.py b/backend/app/models/knowledge.py new file mode 100644 index 0000000..5135edf --- /dev/null +++ b/backend/app/models/knowledge.py @@ -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()), + ) \ No newline at end of file diff --git a/backend/app/services/embedding_service.py b/backend/app/services/embedding_service.py new file mode 100644 index 0000000..9092b43 --- /dev/null +++ b/backend/app/services/embedding_service.py @@ -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] \ No newline at end of file diff --git a/backend/app/services/knowledge_service.py b/backend/app/services/knowledge_service.py new file mode 100644 index 0000000..3e50aad --- /dev/null +++ b/backend/app/services/knowledge_service.py @@ -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) \ No newline at end of file diff --git a/backend/scripts/test_embo01_api.py b/backend/scripts/test_embo01_api.py new file mode 100644 index 0000000..dc06109 --- /dev/null +++ b/backend/scripts/test_embo01_api.py @@ -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}") diff --git a/backend/scripts/verify_embedding.py b/backend/scripts/verify_embedding.py new file mode 100644 index 0000000..c42b3cc --- /dev/null +++ b/backend/scripts/verify_embedding.py @@ -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() \ No newline at end of file