import os 
target = '/home/worker/tengniu/rag/ragas'
os.chdir(target)          # 切换工作目录
print(os.getcwd())        # 验证
# 这个txt需要提前准备好
with open('恐龙.txt', 'r', encoding='utf-8') as f:
    text_content = f.read()
    print(f"文件读取成功，内容长度: {len(text_content)} 字符")

# 修复导入：使用新的导入路径
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import InMemoryStore
from langchain.schema import Document
 
# 创建embedding
bge_embeddings = HuggingFaceBgeEmbeddings(
    model_name="/home/worker/models/bge-large-zh-v1.5",
    model_kwargs={'device': 'cuda'}
)

# 创建文档分割器
parent_splitter = RecursiveCharacterTextSplitter(chunk_size=1000)
child_splitter = RecursiveCharacterTextSplitter(chunk_size=400)

# 创建向量数据库
vectorstore = Chroma(
    collection_name="split_parents", 
    embedding_function=bge_embeddings
)

# 创建内存存储
store = InMemoryStore()

# 创建父文档检索器
retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=store,
    child_splitter=child_splitter,
    parent_splitter=parent_splitter,
    search_kwargs={"k": 2}
)

# 将文本转换为Document对象
documents = [Document(page_content=text_content, metadata={"source": "恐龙.txt"})]

# 添加文档集
print("正在添加文档到检索器...")
retriever.add_documents(documents)
print("文档添加完成！")

# 测试检索功能
query = "恐龙灭绝的原因"
print(f"\n测试查询: {query}")
results = retriever.invoke(query)
print(f"检索到 {len(results)} 个相关文档")
for i, doc in enumerate(results, 1):
    print(f"结果 {i}: {doc.page_content[:200]}...")


from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnableMap
from langchain.schema.output_parser import StrOutputParser
from langchain.chat_models import ChatOpenAI
# from langchain_google_genai import ChatGoogleGenerativeAI
 
#创建gemini model
# model = ChatGoogleGenerativeAI(model="gemini-pro")
 
#创建openai model

from langchain.chat_models import ChatOpenAI

model = ChatOpenAI(
    model="qwen3-14b-local",
    openai_api_base="http://172.25.67.61:9998/v1",  # 你本地API地址
    openai_api_key="not-needed",  # 可以随便填
    temperature=0.7,
    max_tokens=2048
)
#创建prompt模板
template = """You are an assistant for question-answering tasks. 
Use the following pieces of retrieved context to answer the question. 
If you don't know the answer, just say that you don't know. 
Use two sentences maximum and keep the answer concise.
Question: {question} 
Context: {context} 
Answer:
"""
 
#由模板生成prompt
prompt = ChatPromptTemplate.from_template(template)

 
 
#创建chain
chain = RunnableMap({
    "context": lambda x: retriever.get_relevant_documents(x["question"]),
    "question": lambda x: x["question"]
}) | prompt | model | StrOutputParser()

from datasets import Dataset
 
questions = ["恐龙是怎么被命名的？", 
             "恐龙怎么分类的？",
             "体型最大的是哪种恐龙?",
             "体型最长的是哪种恐龙？它在哪里被发现？",
             "恐龙采样什么样的方式繁殖？",
             "恐龙是冷血动物吗？",
             "陨石撞击是导致恐龙灭绝的原因吗？",
             "恐龙是在什么时候灭绝的？",
             "鳄鱼是恐龙的近亲吗？",
             "恐龙在英语中叫什么？"
            ]
ground_truths = [["1841年，英国科学家理查德·欧文在研究几块样子像蜥蜴骨头化石时，认为它们是某种史前动物留下来的，并命名为恐龙，意思是“恐怖的蜥蜴”。"],
                ["恐龙可分为鸟类和非鸟恐龙。"],
                 ["恐龙整体而言的体型很大。以恐龙作为标准来看，蜥脚下目是其中的巨无霸。"],
                 ["最长的恐龙是27米长的梁龙，是在1907年发现于美国怀俄明州。"],
                 ["恐龙采样产卵、孵蛋的方式繁殖。"],
                 ["恐龙是介于冷血和温血之间的动物"],
                 ["科学家最新研究显示，0.65亿年前小行星碰撞地球时间或早或晚都可能不会导致恐龙灭绝，真实灭绝原因是当时恐龙处于较脆弱的生态系统中，环境剧变易导致灭绝。"],
                 ["恐龙灭绝的时间是在距今约6500万年前，地质年代为中生代白垩纪末或新生代第三纪初。"],
                 ["鳄鱼是另一群恐龙的现代近亲，但两者关系较非鸟恐龙与鸟类远。"],
                ["1842年，英国古生物学家理查德·欧文创建了“dinosaur”这一名词。英文的dinosaur来自希腊文deinos（恐怖的）Saurosc（蜥蜴或爬行动物）。对当时的欧文来说，这“恐怖的蜥蜴”或“恐怖的爬行动物”是指大的灭绝的爬行动物（实则不是）"]]
answers = []
contexts = []
 
# Inference
for query in questions:
    answers.append(chain.invoke({"question": query}))
    contexts.append([docs.page_content for docs in retriever.get_relevant_documents(query)])
reference = []
for ref in ground_truths:
    reference.append(ref[0])
# To dict
data = {
    "question": questions,
    "answer": answers,
    "contexts": contexts,
    "reference": reference,
    "ground_truths": ground_truths,
}

# Convert dict to dataset
dataset = Dataset.from_dict(data)
result = evaluate(
    dataset=dataset,
    metrics=[context_precision, context_recall, faithfulness, answer_relevancy],
    llm=model,            # 之前创建的 ChatOpenAI，本地 Qwen3
    embeddings=bge_embeddings  # 使用本地 embeddings，避免 OpenAIEmbeddings
)
df = result.to_pandas()
print(df)