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

2026-05-26 10:33:25 +08:00
parent d40d46a434
commit 1325807257
6 changed files with 390 additions and 0 deletions
@@ -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()