import os
from tools.core.common_functions import read_json_file
from tools.memory_mag.rag.rag import set_embeddings
from langchain.vectorstores import Chroma

# 1. 读取 rag_config.json，确保和 groundcontrol 用的是同一套配置
ai_agent_root = "/home/huangfukk/AI_Agent/scripts/ai_agent"
rag_config_path = os.path.join(ai_agent_root, "config", "rag_config.json")
rag_config = read_json_file(rag_config_path)

rag_mag = rag_config["rag_mag"]
persist_directory = rag_mag["vectorstore_persist_directory"]
collection_name = rag_mag["collection_name"]           # 应该是 "drone_docs"
embedding_model_path = rag_mag["embedding_model_path"]
embedding_type = rag_mag["embedding_framework_type"]   # "llamacpp_embedding" -> 我们要转成 "llamacpp_embeddings"

if embedding_type == "llamacpp_embedding":
    embedding_type = "llamacpp_embeddings"

# 2. 初始化 embeddings（和正式流程一致）
model_config = rag_mag.get("model_config_llamacpp", {})
embeddings = set_embeddings(
    embedding_model_path=embedding_model_path,
    embedding_type=embedding_type,
    model_config=model_config
)

# 3. 加载向量数据库（同样的目录和 collection_name）
vector_db = Chroma(
    persist_directory=persist_directory,
    embedding_function=embeddings,
    collection_name=collection_name
)

# 4. 测试一个你关心的 query，看原始分数
query = "基础实验楼"  # 你可以换成各种中文地标测试
docs_and_scores = vector_db.similarity_search_with_score(query, k=5)

print(f"共返回 {len(docs_and_scores)} 条（不加阈值，只看原始分数）：\n")
for i, (doc, score) in enumerate(docs_and_scores, 1):
    print(f"=== 结果 {i} ===")
    print(f"score: {score}")
    print(f"source: {doc.metadata.get('source', '未知')}")
    text_preview = doc.page_content[:200].replace("\n", " ")
    print(f"text : {text_preview}")
    print("-" * 80)