from helper_utils import word_wrap

# 导入PDF阅读器
from pypdf import PdfReader

# 使用《北京2024年政府工作报告》作为示例文件
reader = PdfReader("./data/2024年北京市政府工作报告.pdf")
# 从该文件中提取文本，并跳过空格
pdf_texts=[p.extract_text().strip() for p in reader.pages]

# 过滤空行，因为检索系统不能接受空行
pdf_texts=[text for text in pdf_texts if text]
# 打印文档第一页作为输出示例
print(word_wrap(pdf_texts[0]))

# 在LangChain工具集中，使用递归字符文本拆分器和句子转换器令牌文本拆分器。
# 字符拆分器可以根据特定的分隔符递归地划分文本，使得它可以在文本中查找指定的字符并在这些字符处将文本分割成更小的片段。
# 如果报错：SSLError:(MaxRetryEror("SOCKSHTTPSomectionpool(host='huggingface.co', port-443)，注意挂上梯子并添加以下代码：（将1080改为自己的vpn端口号）
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter, SentenceTransformersTokenTextSplitter
# import os
# os.environ['HTTPS_PROXY']='http://127.0.0.1:1080'
# os.environ["HTTP_PROXY"]='http://127.0.0.1:1080'

character_splitter = RecursiveCharacterTextSplitter(
    #separators=["\n\n", "\n", ". ", " ", ""],英文文档的分割符示例
    separators=["\n\n", "\n", "。", "，", "、", "；", ""],
    chunk_size=1000,
    chunk_overlap=0
)
character_split_texts = character_splitter.split_text('\n\n'.join(pdf_texts))

print(word_wrap(character_split_texts[10]))
print(f"\nTotal chunks: {len(character_split_texts)}")

# 创建分词器实例，假设可以通过某种方式传入model实例
# 这里的代码需要根据实际的API和参数进行调整
# 为了适应中文文档的要求，这里指定多语言模型 all-mpnet-base-v2 其初始化格式如下：
# __init__(chunk_overlap: int = 50, model_name: str = 'sentence-transformers/all-mpnet-base-v2', tokens_per_chunk: Optional[int] = None, **kwargs: Any) → None
token_splitter = SentenceTransformersTokenTextSplitter(chunk_overlap=0,
                                                       model_name='sentence-transformers/paraphrase-xlm-r-multilingual-v1', 
                                                       tokens_per_chunk=128)

token_split_texts = []
# for text in preprocessed_texts:
for text in character_split_texts:
    token_split_texts += token_splitter.split_text(text)

print(word_wrap(token_split_texts[10]))
print(f"\nTotal chunks: {len(token_split_texts)}")

# 下面的工作是为了创建一个句子转换器嵌入函数，使其能够和 Chorma 一起使用。
import chromadb
import chromadb.utils.embedding_functions as embedding_functions
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from dotenv import load_dotenv, find_dotenv
import os

loaded = load_dotenv(find_dotenv(), override=True)
# 从环境变量中获取 OpenAI API Key 或者直接赋值
API_KEY = os.getenv("API_KEY")

# 如果您使用的是官方 API，就直接用 https://api.siliconflow.cn/v1 就行。
BASE_URL = "https://api.siliconflow.cn/v1"

# chromadb支持的嵌入函数有许多种，这里介绍常用的几种：
# 参考资料：https://docs.trychroma.com/embeddings

# 方式1：默认嵌入函数，需要下载模型，本地计算。英文文本表现不错，中文文本表现一般
# embedding_function = SentenceTransformerEmbeddingFunction()

# 方式2：OpenAI的嵌入函数，直接调用OpenAI的接口，无需下载模型（推荐）
embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=API_KEY,
        api_base=BASE_URL,
        model_name="BAAI/bge-m3",
        dimensions=1024
    )

# 方式3：HuggingFace的嵌入函数，需要下载模型，本地计算，对网络要求高
# embedding_function = embedding_functions.HuggingFaceEmbeddingFunction(
#     api_key="hf_",  # 填入你的 huggingface Access Token
#     model_name="jinaai/jina-embeddings-v2-base-zh"  # 指定模型
# )
## 中文备选模型
# jinaai/jina-embeddings-v2-base-zh
# GanymedeNil/text2vec-large-chinese
# BAAI/bge-large-zh-v1.5
# BAAI/bge-small-zh-v1.5

# print(embedding_function([token_split_texts[10]]))

chroma_client = chromadb.Client()
# 注意这里的名称（beijing_annual_report_2024）可以自己定义，但要求格式是str
chroma_collection = chroma_client.get_or_create_collection("beijing_annual_report_2024",
                                                           embedding_function=embedding_function)

batch_size = 64
for i in range(0, len(token_split_texts), batch_size):
    batch_texts = token_split_texts[i:i + batch_size]
    batch_ids = [str(j) for j in range(i, i + len(batch_texts))]
    chroma_collection.add(ids=batch_ids, documents=batch_texts)

chroma_collection.count()

query = "财政总收入是多少?"

# 查询Chorma来获取结果，请求5个结果
results = chroma_collection.query(query_texts=[query], n_results=5)
retrieved_documents = results['documents'][0]

for document in retrieved_documents:
    print(word_wrap(document))
    print('\n')

import os
import openai
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv

loaded = load_dotenv(find_dotenv(), override=True)
# 从环境变量中获取 OpenAI API Key 或者直接赋值
API_KEY = os.getenv("API_KEY")

# 如果您使用的是官方 API，就直接用 https://api.siliconflow.cn/v1 就行。
BASE_URL = "https://api.siliconflow.cn/v1"

openai_client = OpenAI(api_key=API_KEY, base_url=BASE_URL, max_retries=3)

# 使用GPT3.5 Turbo完成后续操作
def rag(query, retrieved_documents, model="Qwen/Qwen3-8B"):
    information = "\n\n".join(retrieved_documents)
    
    extra_body = {}
    if "Qwen3" in model:
        extra_body={
            "enable_thinking": False
        }
    
# 英文prompt
#   messages = [
#        {
#           "role": "system",
#           "content": "You are a helpful expert financial research assistant. Your users are asking questions about information contained in an annual report."
#           "You will be shown the user's question, and the relevant information from the annual report. Answer the user's question using only this information."
#       },
#        {"role": "user", "content": f"Question: {query}. \n Information: {information}"}
#   ]

# 中文prompt
    messages = [
        {
            "role": "system",
            "content": "你是一个有用的专家级财务研究助手。你的用户正在询问年度报告中包含的信息相关的问题。"
                       "你将看到用户的问题和年度报告中的相关信息。请仅使用这些信息回答用户的问题。"
        },
        {"role": "user", "content": f"问题：{query}。\n信息：{information}"}
    ]
    response = openai_client.chat.completions.create(
        model=model,
        messages=messages,
        extra_body=extra_body,
        temperature=0, seed=42,
        presence_penalty=0, frequency_penalty=0,
    )
    content = response.choices[0].message.content
    return content

# 得到最终的回答
output = rag(query=query, retrieved_documents=retrieved_documents)

print(word_wrap(output))