79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
|
"""
|
||
|
|
全链路验证脚本 — 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()
|