# 加载配置
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core.settings import Settings
import os
from dotenv import load_dotenv
# 加载环境变量
load_dotenv()

# 设置为全局默认Embedding模型
Settings.embed_model = OpenAIEmbedding(
    model="text-embedding-3-small",
    api_key=os.getenv("OPENAI_API_KEY"),
    api_base=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
)

# 设置为全局默认 LLM
Settings.llm = OpenAI(
    model="gpt-4-turbo",
    api_key=os.getenv("OPENAI_API_KEY"),
    api_base=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
)

text = """证 券 研 究 报 告 行 业 研 究

中小市值

行业研究/行业点评

海外科技巨头持续发力 AI，龙头公司中报业绩亮眼

——AI 行业点评报告

行业： 日期： 中小市值 yxzqdatemark 2024年08月20日 
分析师： 彭毅 E-mail： pengyi@yongxingsec.c om SAC编号： S1760523090003 分析师： 张恬 
E-mail： zhangtian@yongxings ec.com SAC编号： S1760524070002

◼ 核心观点

4k: Ps AB
海外 AI 视角：
（1）英伟达推出 B200A，2025 年 Blackwell GPU 有望 上量。《科创板日报》8 月 7 日讯，据 TrendForce 集邦咨询，
英伟达仍 计划在 2024 年下半年推出 B100 及 B200，供应 CSPs（云端服务业者） 客户，并规划于 2024 年第三季后陆续供货。
在 CoWoS-L 良率和量产 尚待整备的情况下，英伟达同步规划降规版 B200A 给其他企业客户， 并转为采用 CoWoS-S 封装技术。
B200A 的存储器规格将采用 4 颗 HBM3e（第五代高带宽内存）12hi（12 层堆叠），总容量为 144GB。 
预期 OEMs（原始设备制造商）应会于 2025 年上半年正式拿到 B200A 芯片。
到 2025 年 Blackwell 平台将占英伟达高端 GPU 逾八成，并促 使英伟达高端 GPU 系列的出货年增率上升至 55%。
（2）马斯克旗下 xAI 公司发布 Grok-2 测试版。《科创板日报》8 月 14 日讯，马斯克旗 下 xAI 正式发布语言模型 Grok 2 早期预览版。
据介绍，该系列模型具 有聊天、编码和推理等功能，包括 Grok 2 和 Grok 2 mini 两个版本， 目前正在 X 平台上进行测试，并将于本月晚些时候通过企业 API 提供 上述模型。
（3）苹果计划推出 AI 桌面机器人。据环球网援引彭博社报 道，苹果公司正在加速其桌面机器人项目的研发工作，并计划最早于 2026 年推出这款创新设备。
据报道，苹果的桌面机器人将配备一块类 似 iPad 的大尺寸显示屏，由一个纤薄的机械臂支撑，能够实现上下倾 斜和 360 度旋转。
这款设备将集成智能家居控制中心、视频会议终端 和家庭安全监控等多项功能。桌面机器人将搭载 Siri 和 Apple Intelligence技术，具备响应多种语音指令、识别不同声音的能力，
并 能自动调整显示屏方向以面向房间内的用户，提供更加智能化和个性 化的交互体验。
分析师： 彭毅

近一年行业与沪深 300 比较

10% 中小市值 沪深300 2% -6% -14% -22% -30% 08/23 10/23 01/24 03/24 06/24 08/24 资料来源：Wind，甬兴证券研究所

行 业 点 评

AI 重点公司中报业绩：
（1）工业富联：2024 年上半年，公司营业收入 2600.9 亿元，同比增长 28.7%，归母净利润 87.4 亿元，同比增长 22%， 扣非净利润若扣除汇兑损益影响后，二季度利润总额同比增长 25%。 
收入增速加快原因系品牌客户不断增加 AI 服务器订单。公司凭借先 进的 AI 技术及产品，科学的生产布局和强大的供应链管理能力，不断 提升科技创新能力，发展新质生产力。
（2）海光信息：2024 年上半年， 公司营业收入 376,291.26 万元，同比增长 44.08%；归母净利润 85,339.59 万元，同比增长 25.97%；扣非归母净利润 81,763.04 万元， 同比增长 32.09%。
公司整体经营指标持续增长，高端处理器产品的产 业生态持续扩展，涉及的行业应用以及新兴人工智能大模型产业逐步 增加。
公司立足通用计算市场，依靠自身的高端处理器设计能力、产 品迭代能力、行业引领能力，持续提高公司的知名度及竞争优势，公 司营收规模及利润等经营指标实现较好增长。
（3）腾讯控股：财联社 8 月 14 日电，2024 年上半年，腾讯控股实现营收 3206.2 亿元，同比 增长 7%；实现调整后净利润 1075.8 亿元，同比增长 53%。
第二季度 营收 1,611.2 亿元人民币，预估 1,613.5 亿元人民币，同比增长 8%；第 二季度调整后净利润 573.1 亿元人民币，同比增长 53%。

相关报告：
《工信部提议加快发展新质生产力， 关注商业航天、人工智能产业》 
——2024 年 08 月 02 日 《低空经济产业方兴未艾，特斯拉二 代 Optimus 亮相 WAIC》 
——2024 年 07 月 09 日 《政策利好低空经济发展，AI PC 助 力消费电子成长》 ——2024 年 05 月 27 日
◼ 投资建议

我们认为，英伟达、苹果等海外科技巨头在 AI 领域加大投入，有望促 成 AI 应用实现商业化落地；
随着 AI 大模型及应用日趋成熟，AI 基础 设施投资有望持续增长，GPU、服务器、AI 大模型及应用环节有望充 分受益。
建议关注：英伟达、苹果、工业富联、海光信息。

◼ 风险提示

硬件快速迭代的风险；地缘政治冲突与出口管制；AI 应用商业化进度 不及预期。

Oats
分析师声明

本报告署名分析师具有中国证券业协会授予的证券投资咨询执业资格并注册为证券分析师，
以勤勉尽责的职业态度， 专业审慎的研究方法，独立、客观地出具本报告，保证报告采用的信息均来自合规渠道，并对本报告的内容和观点负 责。
负责准备以及撰写本报告的所有研究人员在此保证，本报告所发表的任何观点均清晰、准确、如实地反映了研究 人员的观点和结论，并不受任何第三方的授意或影响。
此外，所有研究人员薪酬的任何部分不曾、不与、也将不会与 本报告中的具体推荐意见或观点直接或间接相关。
公司业务资格说明

甬兴证券有限公司经中国证券监督管理委员会核准，取得证券投资咨询业务许可，具备证券投资咨询业务资格。

投资评级体系与评级定义

股票投资评级： 分析师给出下列评级中的其中一项代表其根据公司基本面及（或）估值预期以报告日起 6 个月 内公司股价相对于同期市场基准指数表现的看法。

买入

股价表现将强于基准指数 20%以上

增持

股价表现将强于基准指数 5-20%

中性

股价表现将介于基准指数±5%之间

BAF

减持

股价表现将弱于基准指数 5%以上

行业投资评级： 分析师给出下列评级中的其中一项代表其根据行业历史基本面及（或）估值对所研究行业以报告 日起 12 个月内的基本面和行业指数相对于同期市场基准指数表现的看法。

增持

行业基本面看好，相对表现优于同期基准指数

中性

行业基本面稳定，相对表现与同期基准指数持平

减持

行业基本面看淡，相对表现弱于同期基准指数

相关证券市场基准指数说明：A 股市场以沪深 300 指数为基准；港股市场以恒生指数为基准；新三板市场以三板成 指（针对协议转让标的）或三板做市指数（针对做市转让标的）为基准指数。

投资评级说明：

不同证券研究机构采用不同的评级术语及评级标准，投资者应区分不同机构在相同评级名称下的定义差异。本评级体 系采用的是相对评级体系。
投资者买卖证券的决定取决于个人的实际情况。投资者应阅读整篇报告，以获取比较完整 的观点与信息，投资者不应以分析师的投资评级取代个人的分析与判断。
特别声明

在法律许可的情况下，甬兴证券有限公司(以下简称“本公司”）或其关联机构可能会持有报告中涉及的公司所发行的 证券或期权并进行交易，也可能为这些公司提供或争取提供投资银行、财务顾问以及金融产品等各种服务。
因此，投 资者应当考虑到本公司或其相关人员可能存在影响本报告观点客观性的潜在利益冲突，投资者请勿将本报告视为投资 或其他决定的唯一参考依据。也不应当认为本报告可以取代自己的判断。
版权声明

本报告版权归属于本公司所有，属于非公开资料。本公司对本报告保留一切权利。未经本公司事先书面许可，任何机 构或个人不得以任何形式翻版、复制、转载、刊登和引用本报告中的任何内容。
否则由此造成的一切不良后果及法律 责任由私自翻版、复制、转载、刊登和引用者承担。

请务必阅读报告正文后各项声明

行业点评

2

行业点评
重要声明

本报告由本公司发布，仅供本公司的客户使用，且对于接收人而言具有保密义务。本公司并不因相关人员通过其他途 径收到或阅读本报告而视其为本公司的客户。
客户应当认识到有关本报告的短信提示、电话推荐及其他交流方式等只 是研究观点的简要沟通，需以本公司发布的完整报告为准，本公司接受客户的后续问询。
本报告首页列示的联系人， 除非另有说明，仅作为本公司就本报告与客户的联络人，承担联络工作，不从事任何证券投资咨询服务业务。

本报告中的信息均来源于我们认为可靠的已公开资料，本公司对这些信息的真实性、准确性及完整性不作任何保证。 
本报告中的信息、意见等均仅供客户参考，该等信息、意见并未考虑到获取本报告人员的具体投资目的、财务状况以 及特定需求，在任何时候均不构成对任何人的个人推荐。
客户应当对本报告中的信息和意见进行独立评估，并应同时 思量各自的投资目的、财务状况以及特定需求，必要时就法律、商业、财务、税收等方面咨询专家的意见。
客户应自 主作出投资决策并自行承担投资风险。本公司特别提示，本公司不会与任何客户以任何形式分享证券投资收益或分担 证券投资损失，任何形式的分享证券投资收益或者分担证券投资损失的书面或口头承诺均为无效。
市场有风险，投资 须谨慎。对依据或者使用本报告所造成的一切后果，本公司和关联人员均不承担任何法律责任。
本报告所载的意见、评估及预测仅反映本公司于发布本报告当日的判断。该等意见、评估及预测无需通知即可随时更 改。
过往的表现亦不应作为日后表现的预示和担保。在不同时期，本公司可发出与本报告所载意见、评估及预测不一 致的研究报告。
本公司不保证本报告所含信息保持在最新状态。
同时，本公司的销售人员、交易人员以及其他专业人 士可能会依据不同假设和标准、采用不同的分析方法而口头或书面发表与本报告意见及建议不一致的市场评论或交易 观点。
本公司没有将此意见及建议向报告所有接收者进行更新的义务。投资者应当自行关注相应的更新或修改。

请务必阅读报告正文后各项声明"""

import json
from llama_index.core.schema import Document

# 讲文档转化为documents信息
documents = []

# 创建Document对象
doc = Document(
    text=text,
    metadata={
        'element_type':'Text',
        'file_name':'甬兴证券-AI行业点评报告：海外科技巨头持续发力AI，龙头公司中报业绩亮眼.pdf'
    },
)

documents.append(doc)
documents

# 加载句子切分器
from llama_index.core.node_parser import SentenceSplitter

# 初始化 SentenceSplitter
sentence_splitter = SentenceSplitter(
    chunk_size=512,      # 这里 chunk_size 表示 token 近似或字符近似，视版本调整
    chunk_overlap=64     # 重叠部分选择大小为64
)

nodes = sentence_splitter.get_nodes_from_documents(documents)
len(nodes)

from llama_index.core import VectorStoreIndex,StorageContext

# 直接基于nodes构建索引
index = VectorStoreIndex(nodes)

# 2. 索引持久化
index.storage_context.vector_store.persist("vector_store.json")

# 学术研究场景，可以多检索几篇文章
retriever = index.as_retriever(
    similarity_top_k=5,  # 学术研究需要更多参考资料
)

# 返回检索到的节点
response_nodes = retriever.retrieve("B200A 芯片的存储器规格和预计何时被原始设备制造商拿到是什么？")

# 打印检索到的信息
for node in response_nodes:
    print(f"相似度: {node.score:.4f}")
    print(f"内容: {node.text[:200]}...")
    print(f"元数据: {node.metadata}")
    print("=" * 60)

from llama_index.core.retrievers import VectorIndexRetriever

# 4. 显式创建VectorIndexRetriever以获得更多控制
vector_retriever = VectorIndexRetriever(
    index=index,
    similarity_top_k=5,
    # filters=MetadataFilters(...) # 可选的元数据过滤
)

query = "B200A 芯片的存储器规格和预计何时被原始设备制造商拿到是什么？"
retriever_nodes = vector_retriever.retrieve(query)

# 打印结果
for node in retriever_nodes:
    print(f"Score: {node.score:.4f}")
    print(f"Content: {node.get_content()}")
    print("=" * 60)

#!pip install rank-bm25 llama-index-retrievers-bm25

# 关键词策略检索
from llama_index.retrievers.bm25 import BM25Retriever

# 1.创建向量检索器
vector_retriever = index.as_retriever(similarity_top_k=5)

# 2.BM25检索器
bm25_retriever = BM25Retriever.from_defaults(
        nodes=nodes,
        similarity_top_k=5
)

# 3. 执行BM25检索
query = "B200A 的存储器规格"
results = bm25_retriever.retrieve(query)

# 打印结果
for node in results:
    print(f"Score: {node.score:.4f}")
    print(f"Content: {node.get_content()}")
    print("-" * 20)

from llama_index.core.retrievers import QueryFusionRetriever

# 1. 创建QueryFusionRetriever
retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],  # 集成不同的检索器
    similarity_top_k=5,                 # 最终返回的文档数量
    num_queries=1,                      # 为原始查询生成的变体数量，（默认为1，不生成变体）
    mode="reciprocal_rerank",           # 结果融合模式，'reciprocal_rerank'是常用且效果好的模式
    use_async=True,                     # 是否异步执行
    verbose=True,                       # 是否打印调试信息
    retriever_weights=[0.6,0.4]         # 检索器权重，用于加权融合
)

# 2. 执行检索
query = "B200A 芯片的存储器规格和预计何时被原始设备制造商拿到是什么？"
query_nodes = retriever.retrieve(query)

# 打印结果
for node in query_nodes:
    print(f"Score: {node.score:.4f}")
    print(f"Content: {node.get_content()}")
    print("=" * 60)

#!pip install huggingface_hub modelscope
# modelscope download --model BAAI/bge-reranker-base  --local_dir ./

from llama_index.core.postprocessor import SentenceTransformerRerank
from llama_index.core.query_engine import RetrieverQueryEngine

# 模型本地路径
local_model_path = "/root/autodl-tmp/bge-reranker-base"

# 1.定义精排序器 (Reranker)
reranker = SentenceTransformerRerank(
    model=local_model_path,
    top_n=3, # 精选出最相关的4个文档送给LLM,
)

# 2.构建QueryEngine
query_engine = RetrieverQueryEngine.from_args(
    retriever=retriever,
    node_postprocessors=[reranker]
)

# 3. 执行查询
query = "B200A 芯片的存储器规格和预计何时被原始设备制造商拿到是什么？"
response = query_engine.query(query)

response.get_formatted_sources

response_synthesizer = get_response_synthesizer(
    response_mode="compact",  # 选择合成模式
    llm=llm,                  # 指定使用的LLM
    text_qa_template=my_qa_prompt,  # 自定义提示模板
    streaming=True            # 是否启用流式输出
)

# 基础查询
query_engine = index.as_query_engine(
    similarity_top_k=3,
    response_mode="refine",  # 对多个检索结果进行精炼
    streaming=True  # 流式输出
)

# 正常响应
response = query_engine.query("B200A 芯片的存储器规格和预计何时被原始设备制造商拿到是什么？")
print(response)

# 流式响应处理
# streaming_response = query_engine.query("B200A 芯片的存储器规格和预计何时被原始设备制造商拿到是什么？")
# for token in streaming_response.response_gen:
#     print(token, end=" ", flush=True)

from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever

# 1. 从已有的 index 创建 retriever
retriever = VectorIndexRetriever(
    index=index,
    similarity_top_k=3
)

# 2. 构造 QueryEngine
query_engine = RetrieverQueryEngine(
        
    retriever=retriever,          # 传入检索器
    response_synthesizer=None,    # 可以自定义响应合成器
    node_postprocessors=None     # 可传入 reranker、过滤器
)

# 3. 直接查询
response = query_engine.query("B200A 芯片的存储器规格和预计何时被原始设备制造商拿到是什么？")
print(response)

chat_engine = index.as_chat_engine(
    chat_mode="context",  # "context", "condense_question", "react", "openai"
    system_prompt="你是一个关于LlamaIndex的问答助手。请基于提供的上下文回答问题。",  # 系统提示词
    memory=None          # 对话记忆管理
)

# 多轮对话
response1 = chat_engine.chat("文章主要讲述了哪几家海外AI巨头？")
print(f"AI: {response1}")
print("=" * 60)
     
response2 = chat_engine.chat("这几家AI巨头都研发了哪些产品？")
print(f"AI: {response2}")  # 能理解这是基于前一个问题的后续

# 重置对话历史
chat_engine.reset()

from llama_index.core.memory import ChatMemoryBuffer

# 创建对话记忆
memory = ChatMemoryBuffer.from_defaults(token_limit=3000)

chat_engine = index.as_chat_engine(
    chat_mode="context",  # "context", "condense_question", "react", "openai"
    system_prompt="你是一个关于LlamaIndex的问答助手。请基于提供的上下文回答问题。",  # 系统提示词
    memory=memory          # 对话记忆管理
)

# 多轮对话
response1 = chat_engine.chat("文章主要讲述了哪几家海外AI巨头？")
print(f"AI: {response1}")
print("=" * 60)
     
response2 = chat_engine.chat("这几家AI巨头都研发了哪些产品？")
print(f"AI: {response2}")  # 能理解这是基于前一个问题的后续

# 重置对话历史
chat_engine.reset()

from llama_index.core import VectorStoreIndex
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.postprocessor import SentenceTransformerRerank
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.response_synthesizers import get_response_synthesizer

# 1. 加载向量索引
index = VectorStoreIndex(nodes=nodes)

# 2. 向量检索器 (粗召回)
vector_retriever = index.as_retriever(
    similarity_top_k=5 # 扩大召回范围，为Reranker提供充足的候选
)

# 3. BM25检索器 (关键词补充)
bm25_retriever = BM25Retriever.from_defaults(
    nodes=nodes,
    similarity_top_k=5
)

# 4. 混合检索器 (融合结果)
hybrid_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    similarity_top_k=5,         # 融合后保留的候选数量
    num_queries=1,              # 保持简单，不生成子查询
    mode="reciprocal_rerank",   # 使用RRF融合策略
    use_async=True              # 必须开启异步
)

from llama_index.core.postprocessor import SentenceTransformerRerank

# 模型本地路径
local_model_path = "/root/autodl-tmp/bge-reranker-base"

# 5.定义精排序器 (Reranker)
reranker = SentenceTransformerRerank(
    model=local_model_path,
    top_n=3, # 精选出最相关的4个文档送给LLM,
)

# 6. 创建响应合成，用大模型（LLM）根据 top-K 文档 + 用户问题，生成一个最终的回答。
response_synthesizer = get_response_synthesizer(
    response_mode="compact", streaming=True
)

# 7.构建QueryEngine
query_engine = RetrieverQueryEngine.from_args(
    retriever=hybrid_retriever,
    node_postprocessors=[reranker],
    response_synthesizer=response_synthesizer
)

# 8. 执行查询
query = "B200A 芯片的存储器规格和预计何时被原始设备制造商拿到是什么？"
response = query_engine.query(query)

print("大模型回复内容：")
print(response)
print("=" * 60)

# 打印最终被送入LLM的Source Nodes
print("\n--- Final Source Nodes (after Reranking) ---")
for node in response.source_nodes:
    print(f"Score: {node.score:.4f}, File: {node.text}")
    print("=" * 60)

from llama_index.core.evaluation import generate_question_context_pairs

custom_prompt = """
请基于以下文本内容，生成 {num_questions_per_chunk} 个有代表性的问题。
问题应当：
1. 覆盖文本关键信息；
2. 不要过于宽泛；
3. 使用自然中文。

---------------------
文本内容如下：
{context_str}
---------------------
"""
# 3. 自动生成评估数据集
# LlamaIndex可以从文档中自动生成“问题-上下文”对，非常适合冷启动
# num_questions_per_chunk: 每个文档块生成多少问题
qa_dataset = generate_question_context_pairs(
    # llamaindex_documents,
    nodes=nodes,
    llm=Settings.llm,
    num_questions_per_chunk=1,
    qa_generate_prompt_tmpl=custom_prompt,
)

for value in qa_dataset.queries.values():
    print(value)
    print("=" * 60)

for value in qa_dataset.corpus.keys():
    print(value)
    print("=" * 60)

# 建立问答数据字典
qa_dict = dict()

for i in range(0,len(qa_dataset.corpus.keys())):
    qa_dict[list(qa_dataset.queries.values())[i]] = [list(qa_dataset.corpus.keys())[i]]

qa_dict

from llama_index.core import VectorStoreIndex


# 1. 加载向量索引
index = VectorStoreIndex(nodes=nodes)

# 2. 准备检索器：仅使用向量检索（避免 'hybrid' 报错）
retriever = index.as_retriever(
    similarity_top_k=10  # top_k 要足够大，避免命中率为 0
)

# 3. 探针检索，观察系统返回的 ID 形式（用于制作 ground truth）
probe_nodes = retriever.retrieve("B200A 芯片的存储器规格和预计何时被原始设备制造商拿到是什么？")

probe_ids = [n.node.node_id for n in probe_nodes]
print("检索返回的 ID 例子：", probe_ids)

from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.evaluation import RetrieverEvaluator
from tqdm import tqdm

# 加载评估器
from llama_index.core.evaluation.retrieval.metrics import (
    Precision,   # 精确率
    Recall,      # 召回率
    MRR,         # 平均倒数排名
    HitRate,     # 命中率
    NDCG,        # 归一化折损累积增益
    AveragePrecision, # 平均精确率
)

# 1. 建立评估矩阵
metrics = [
    Precision(),
    Recall(),
    MRR(),
    HitRate(),
    NDCG(),
    AveragePrecision(),
]

# 2. 实例化RetrieverEvaluator评估器
retriever_evaluator = RetrieverEvaluator(metrics=metrics, retriever=retriever)


# 3. 运行评估
eval_results = []
for i, question in enumerate(list(qa_dataset.queries.values())):
    # 获取当前问题的真实相关文档ID
    expected_ids = qa_dict[question]

    # 评估单个查询的检索结果
    eval_result = retriever_evaluator.evaluate(
        query=question,             # 提问问题数据
        expected_ids=expected_ids,  # 标准答案id，list格式
        retriever=retriever         # 检索器
    )
    eval_results.append(eval_result)

eval_results

import pandas as pd
import matplotlib.pyplot as plt

# 1. 数据整理
results_data = []
for result in eval_results:
    row = {"query": result.query}
    for metric_name, metric_score in result.metric_dict.items():
        row[metric_name] = metric_score.score
    results_data.append(row)

# 2. 转换成DataFrame格式
df_results = pd.DataFrame(results_data)

# 3. 计算每个指标的平均值
mean_scores = df_results.mean(numeric_only=True)
print("=== 平均评估指标 ===")
print(mean_scores)

# 4. 将平均结果也加入DataFrame用于绘图
df_mean = mean_scores.reset_index()
df_mean.columns = ['Metric', 'Score']

# 5. 可视化评估结果
# 绘制平均指标条形图
plt.figure(figsize=(10, 6))
bars = plt.bar(df_mean['Metric'], df_mean['Score'], color=['skyblue', 'lightgreen', 'coral', 'gold', 'violet', 'lightgray'])
plt.title('Retriever Evaluation Metrics (Average Scores)')
plt.ylabel('Score')
plt.ylim(0, 1.5)

# 在柱子上添加数值标签
for bar in bars:
    height = bar.get_height()
    plt.text(bar.get_x() + bar.get_width()/2., height + 0.01,
             f'{height:.3f}', ha='center', va='bottom')
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

# 打印详细的评估结果表格
print("\n=== 详细评估结果 ===")
df_results

from llama_index.core.evaluation import (
    SemanticSimilarityEvaluator,
    FaithfulnessEvaluator,
    RelevancyEvaluator
)
from llama_index.core.node_parser import SentenceSplitter


def quick_evaluation_demo():
    """快速评估演示"""

    print("=== LlamaIndex文档切片质量评估器演示 ===\n")

    # 1. 配置API密钥（请替换为您的实际密钥）
    print("1. 配置API密钥...")
    if not os.getenv("OPENAI_API_KEY"):
        print("请设置环境变量 OPENAI_API_KEY")
        print("例如: export OPENAI_API_KEY='your-api-key-here'")
        return

    # 2. 初始化模型和评估器
    print("2. 初始化模型和评估器...")
    # embed_model = OpenAIEmbedding(model="text-embedding-3-large")
    # llm = OpenAI(model="gpt-4")

    # 创建语义相似度评估器，通过比较生成的答案与参考答案，两者在语义上的接近程度。（答案 vs 参考答案）
    semantic_evaluator = SemanticSimilarityEvaluator(
        embed_model=Settings.embed_model,
        similarity_threshold=0.8
    )
    # 创建忠实度评估器，检测幻觉，确保答案源于给定上下文（答案 vs 上下文）
    faithfulness_evaluator = FaithfulnessEvaluator(llm=Settings.llm)

    # 创建相关性评估器，评估答案与问题的匹配度及上下文的辅助作用（问题 vs 答案 vs 上下文）
    relevancy_evaluator = RelevancyEvaluator(llm=Settings.llm)

    # 3. 准备测试数据
    print("3. 准备测试数据...")
    sample_document = """
    人工智能（Artificial Intelligence，AI）是计算机科学的一个分支，
    它企图了解智能的实质，并生产出一种新的能以人类智能相似的方式做出反应的智能机器。

    机器学习是人工智能的一个核心研究领域，它使计算机有能力在不被明确编程的情况下进行学习。
    机器学习专注于计算机程序的开发，这些程序可以访问数据并使用它学习为自己。

    深度学习是机器学习的子集，它基于人工神经网络的表征学习方法。
    深度学习在图像、语音、文本等感知任务上取得了突破性进展。
    """

    # 切分文档
    splitter = SentenceSplitter(chunk_size=200, chunk_overlap=30)
    chunks = splitter.split_text(sample_document)

    print(f"文档已切分为 {len(chunks)} 个片段:")
    for i, chunk in enumerate(chunks):
        print(f"  片段 {i+1}: {chunk[:50]}...")

    # 4. 执行各项评估
    print("\n4. 执行质量评估...")

    # 测试查询和响应
    test_query = "什么是人工智能？"
    test_response = "人工智能是计算机科学的一个分支，研究如何让机器表现出智能行为。"
    test_reference = "人工智能是计算机科学的一个分支，企图了解智能的实质"

    # 语义相似度评估
    print("\n--- 语义相似度评估 ---")
    semantic_result = semantic_evaluator.evaluate(
        query=test_query,
        response=test_response,
        reference=test_reference,
        contexts=[chunks[0]]
    )
    print(f"相似度分数: {semantic_result.score:.3f}")
    print(f"是否通过: {semantic_result.passing}")
    print(f"反馈: {semantic_result.feedback}")

    # 忠实度评估
    print("\n--- 忠实度评估 ---")
    faithfulness_result = faithfulness_evaluator.evaluate(
        query=test_query,
        response=test_response,
        contexts=[chunks[0]]
    )
    print(f"忠实度分数: {faithfulness_result.score:.3f}")
    print(f"是否通过: {faithfulness_result.passing}")
    print(f"反馈: {faithfulness_result.feedback}")

    # 相关性评估
    print("\n--- 相关性评估 ---")
    relevancy_result = relevancy_evaluator.evaluate(
        query=test_query,
        response=test_response,
        contexts=[chunks[0]]
    )
    print(f"相关性分数: {relevancy_result.score:.3f}")
    print(f"是否通过: {relevancy_result.passing}")
    print(f"反馈: {relevancy_result.feedback}")

    # 5. 综合评估结果
    print("\n=== 评估总结 ===")
    metrics = {
        '语义相似度': semantic_result.score,
        '忠实度': faithfulness_result.score,
        '相关性': relevancy_result.score
    }

    for metric, score in metrics.items():
        status = "✓ 通过" if score >= 0.7 else "✗ 未通过"
        print(f"{metric}: {score:.3f} {status}")

    avg_score = sum(metrics.values()) / len(metrics)
    print(f"\n平均分数: {avg_score:.3f}")
    overall_status = "✓ 整体通过" if avg_score >= 0.7 else "✗ 整体未通过"
    print(f"整体评估: {overall_status}")

    return "评估完成！"

quick_evaluation_demo()

#!pip install deepeval

#构建黄金数据集格式
qa_res_dict = dict()

for i in range(0,len(qa_dataset.corpus.keys())):
    qa_res_dict[list(qa_dataset.queries.values())[i]] = list(qa_dataset.corpus.values())[i]

# 展示数据集
qa_res_dict

from deepeval import evaluate
from deepeval.metrics import (
    AnswerRelevancyMetric,  # 答案相关性
    FaithfulnessMetric,      # 忠诚度（防幻觉）
    ContextualPrecisionMetric, # 上下文精确度
    ContextualRecallMetric,  # 上下文召回率
    # MultimodalAnswerRelevancyMetric # 多模态答案相关性
)
from deepeval.test_case import LLMTestCase
from deepeval.dataset import EvaluationDataset, Golden
from llama_index.core import VectorStoreIndex
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.postprocessor import SentenceTransformerRerank
from llama_index.core.retrievers import QueryFusionRetriever


# 1. 加载向量索引
index = VectorStoreIndex(nodes=nodes)

# 2. 向量检索器 (粗召回)
vector_retriever = index.as_retriever(
    similarity_top_k=6 # 扩大召回范围，为Reranker提供充足的候选
)

# 3. BM25检索器 (关键词补充)
bm25_retriever = BM25Retriever.from_defaults(
    nodes=nodes,
    similarity_top_k=6
)

# 4. 混合检索器 (融合结果)
retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    similarity_top_k=6, # 融合后保留的候选数量
    num_queries=1, # 保持简单，不生成子查询
    mode="reciprocal_rerank",
    use_async=True # 必须开启异步
)

# 5. 定义精排序器 (Reranker)
reranker = SentenceTransformerRerank(
    model="/root/autodl-tmp/bge-reranker-base",
    top_n=3, # 精选出最相关的4个文档送给LLM,
)

# 6. 定义评估指标
metrics = [
    AnswerRelevancyMetric(threshold=0.7),  # 答案相关性阈值0.7
    FaithfulnessMetric(threshold=0.8),     # 忠诚度阈值0.8
    ContextualPrecisionMetric(threshold=0.6),  # 上下文精确度
    ContextualRecallMetric(threshold=0.7),    # 上下文召回率
]

# 7. 准备测试用例（黄金数据集）
dataset = EvaluationDataset(
    goldens=[
        Golden(
            input=list(qa_res_dict.keys())[0],
            expected_output=list(qa_res_dict.values())[0]
        ),
        Golden(
            input=list(qa_res_dict.keys())[1],
            expected_output=list(qa_res_dict.values())[1]
        )
    ]
)

from llama_index.core.query_engine import RetrieverQueryEngine

# 8.构建查询引擎
# query_engine = RetrieverQueryEngine(
#     retriever=retriever,
#     node_postprocessors=[reranker],
# )
query_engine = index.as_query_engine(similarity_top_k=5)

# 9. 执行端到端评估
def evaluate_rag():
    test_cases = []
    for golden in dataset.evals_iterator():
        # 运行查询
        response = query_engine.query(golden.input)

        # 提取检索上下文
        retrieval_context = [
            node.node.text for node in response.source_nodes
        ]

        # 创建测试用例
        test_case = LLMTestCase(
            input=golden.input,
            actual_output=response.response,
            expected_output=golden.expected_output,
            retrieval_context=retrieval_context
        )
        test_cases.append(test_case)

    # 一次性评估整批用例（关键：传入列表）
    run_result = evaluate(test_cases=test_cases, metrics=metrics)

    # 简要打印每个用例的得分（不同版本的 deepeval 输出结构可能稍有差异）
    try:
        for i, case_result in enumerate(run_result.results):
            print(f"Case {i+1}:")
            for m in case_result.metrics:
                print(f"- {getattr(m, 'name', 'metric')}: {getattr(m, 'score', 'N/A')}")
    except Exception:
        print(run_result)

    return run_result


# 10. 运行
evaluate_rag()

#!pip install trulens trulens-apps-llamaindex trulens-providers-openai trulens-eval streamlit

from trulens.providers.openai import OpenAI as TruOpenAI
from trulens.core import Feedback, TruSession, Select
import numpy as np

def build_feedbacks(provider: TruOpenAI | None = None):
    """
    构建 RAG 三元组反馈函数：
    - Groundedness：答案是否基于提供的上下文，有推理链（CoT）
    - Answer Relevance：答案与用户查询的相关性
    - Context Relevance：上下文与用户查询的相关性

    使用新版 TruLens API（OTEL 模式）：
    - 使用 .on_input(), .on_output(), .on_context() 方法
    - 这些方法会自动从 span attributes 中提取数据
    """

    # 若未显式传入 Provider，则默认使用 OpenAI Provider
    oai = provider or TruOpenAI()

    # Answer Relevance：答案与查询的相关性
    # 从 RECORD_ROOT.INPUT 和 RECORD_ROOT.OUTPUT 获取数据
    f_answer_relevance = (
        Feedback(oai.relevance_with_cot_reasons, name="Answer Relevance")
        .on_input()   # 自动从 SpanAttributes.RECORD_ROOT.INPUT 获取
        .on_output()  # 自动从 SpanAttributes.RECORD_ROOT.OUTPUT 获取
    )

    # Groundedness：答案与上下文的一致性
    # 从 RETRIEVAL.RETRIEVED_CONTEXTS 和 RECORD_ROOT.OUTPUT 获取数据
    # 注意：groundedness_measure_with_cot_reasons 的参数顺序是 (source, statement)
    # source 是上下文列表，statement 是答案
    f_groundedness = (
        Feedback(oai.groundedness_measure_with_cot_reasons, name="Groundedness")
        .on_context(collect_list=True)  # source: 从 SpanAttributes.RETRIEVAL.RETRIEVED_CONTEXTS 获取
        .on_output()  # statement: 从 SpanAttributes.RECORD_ROOT.OUTPUT 获取
    )

    # Context Relevance：上下文与查询的相关性
    # 从 RECORD_ROOT.INPUT 和 RETRIEVAL.RETRIEVED_CONTEXTS 获取数据
    # 注意：context_relevance_with_cot_reasons 的参数顺序是 (question, context)
    f_context_relevance = (
        Feedback(oai.context_relevance_with_cot_reasons, name="Context Relevance")
        .on_input()  # question: 从 SpanAttributes.RECORD_ROOT.INPUT 获取
        .on_context(collect_list=False)  # context: 从 SpanAttributes.RETRIEVAL.RETRIEVED_CONTEXTS 获取，每个上下文片段单独评估
        .aggregate(np.mean)  # 对所有上下文片段的分数取平均
    )

    return [f_groundedness, f_answer_relevance, f_context_relevance]


f_list = build_feedbacks()

from llama_index.core import VectorStoreIndex

# 建立index索引
index = VectorStoreIndex(nodes=nodes)

# 创建查询引擎
query_engine = index.as_query_engine(similarity_top_k=5)

from trulens.core.otel.instrument import instrument
from trulens.otel.semconv.trace import SpanAttributes
from typing import List

class LlamaIndexRAG:
    """
    LlamaIndex RAG 的最小封装：包含检索和生成两个步骤
    - 使用 @instrument 为方法打点，以便 TruApp 记录调用的入参与返回值
    - 明确不依赖 TruLlama 或 LangChain 插件，减少版本兼容问题
    """

    def __init__(self, query_engine):
        # 保存由 LlamaIndex 构建的 QueryEngine（例如 index.as_query_engine(...))
        self.query_engine = query_engine

    @instrument(
        span_type=SpanAttributes.SpanType.RETRIEVAL,
        attributes={
            SpanAttributes.RETRIEVAL.QUERY_TEXT: "query",
            SpanAttributes.RETRIEVAL.RETRIEVED_CONTEXTS: "return",
        }
    )
    def retrieve(self, query: str) -> List[str]:
        """检索相关上下文"""
        # 发起 LlamaIndex 查询
        resp = self.query_engine.query(query)

        # 兼容地提取上下文文本：常见在 resp.source_nodes
        contexts: List[str] = []
        source_nodes = getattr(resp, "source_nodes", [])

        # 遍历每个 SourceNode，尽可能获取原文内容
        for sn in source_nodes:
            text = None

            # 常见字段：sn.node 或 sn.source_node，内部可能有 text 或 get_content()
            node_obj = getattr(sn, "node", None) or getattr(sn, "source_node", None)
            if node_obj is not None:
                text = getattr(node_obj, "text", None)
                if text is None:
                    try:
                        text = node_obj.get_content()
                    except Exception:
                        pass

            # 进一步的回退尝试：直接取 sn.text 或调用 node.get_content()
            if text is None:
                text = getattr(sn, "text", None)
            if text is None:
                try:
                    node = getattr(sn, "node", None)
                    if node is not None and hasattr(node, "get_content"):
                        text = node.get_content()
                except Exception:
                    pass

            if text:
                contexts.append(text)

        return contexts

    @instrument(span_type=SpanAttributes.SpanType.GENERATION)
    def generate_answer(self, query: str, contexts: List[str]) -> str:
        """基于上下文生成答案"""
        # 这里我们直接使用 query_engine 的结果
        # 在实际应用中，这个方法会调用 LLM 生成答案
        resp = self.query_engine.query(query)

        # 兼容不同版本的响应对象：优先取 response，其次 message，最后转字符串
        ans = getattr(resp, "response", None)
        return ans

    # 把这个方法标记为“可观测”的应用调用
    @instrument(
        span_type=SpanAttributes.SpanType.RECORD_ROOT,
        attributes={
            SpanAttributes.RECORD_ROOT.INPUT: "query",
            SpanAttributes.RECORD_ROOT.OUTPUT: "return",
        }
    )
    def query(self, query: str) -> str:
        """执行完整的 RAG 流程：检索 + 生成"""
        contexts = self.retrieve(query)
        answer = self.generate_answer(query, contexts)
        return answer

from trulens.core import Feedback, TruSession
from trulens.apps.app import TruApp

def setup_rag_evaluator(index: VectorStoreIndex, top_k: int = 3):
    """
    设置 RAG 评估器（一次性初始化）：
    1) 创建 QueryEngine
    2) 构建反馈函数
    3) 封装为 TruApp（将反馈函数传入）
    4) 初始化 TruSession

    返回：
    (tru_rag, session) - 可复用的评估组件
    """
    # 从索引创建 QueryEngine，指定 top_k 控制检索的上下文条数
    qe = index.as_query_engine(similarity_top_k=top_k)

    # 封装为自定义 RAG
    rag = LlamaIndexRAG(qe)

    # 初始化会话（连接默认 sqlite 或指定数据库）
    session = TruSession()

    # 构建反馈函数
    feedbacks = build_feedbacks()

    # 创建 TruApp，将反馈函数传入（OTEL 模式下反馈函数会自动运行）
    tru_rag = TruApp(
        rag,
        app_name="LlamaIndex RAG",
        app_version="v1",
        feedbacks=feedbacks  # 在 OTEL 模式下，反馈函数在这里传入
    )

    return tru_rag, session

def evaluate_single_query(tru_rag, session, query: str):
    """
    评估单个查询（可多次调用）：
    1) 通过 TruApp 记录一次调用
    2) 反馈函数会自动运行（已在 TruApp 中配置）
    3) 返回结果

    参数：
    - tru_rag: TruApp 实例（由 setup_rag_evaluator 创建）
    - session: TruSession 实例
    - query: 用户查询

    返回：
    (answer, record)
    """
    # 记录一次查询调用：with 语句块结束后可通过 recording.get() 获取 Record
    # 在 OTEL 模式下，反馈函数会自动运行
    with tru_rag as recording:
        answer = tru_rag.app.query(query)

    record = recording.get()

    # 注意：在 OTEL 模式下，反馈函数已经自动运行
    # 不需要手动调用 session.run_feedback_functions()
    # 反馈结果会自动保存到数据库，可以通过 Dashboard 查看

    return answer, record

from trulens.dashboard import run_dashboard

def evaluate_query(
    index: VectorStoreIndex,
    query: str,
    top_k: int = 4,
    run_ui: bool = True,
):
    """
    对单个查询做端到端评估（兼容旧版 API）：
    1) 由索引构建 QueryEngine；
    2) 用 LlamaIndexRAG 封装并通过 TruApp 记录一次调用；
    3) 反馈函数自动运行（在 OTEL 模式下）；
    4) 可选启动 Dashboard 展示评估结果。

    返回：
    (record, session, answer)
    """
    tru_rag, session = setup_rag_evaluator(index, top_k)
    answer, record = evaluate_single_query(
        tru_rag, session, query
    )

    # 可选启动 Dashboard，查看记录、反馈分数与详细轨迹
    if run_ui:
        run_dashboard(session)

    return record, session, answer

query = "B200A 芯片的存储器规格是多少，预计何时由 OEMs 正式拿到？"
record, session, answer = evaluate_query(index, query)

# 打印模型回复内容
record.main_output

# 打印单次对话的评估指标
record.retrieve_feedback_results()

# 打印同一个版本的评估指标
session.get_leaderboard()