from mx_rag.document import LoaderMng
from mx_rag.document.loader import TextLoader, DocxLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from mx_rag.embedding.local import TextEmbedding
from mx_rag.storage.document_store import SQLiteDocstore
from mx_rag.storage.vectorstore import MindFAISS
from mx_rag.knowledge import KnowledgeDB, KnowledgeStore, upload_files

def build_knowledge():
    # 1. 文档解析与切分管理
    loader_mng = LoaderMng()
    loader_mng.register_loader(loader_class=TextLoader, file_types=[".txt"])
    
    # 🔴【挖空预测 1：切分器参数】必考 chunk_size 和 chunk_overlap
    loader_mng.register_splitter(
        splitter_class=RecursiveCharacterTextSplitter,
        file_types=[".txt", ".docx"],
        splitter_params={"chunk_size": 500, "chunk_overlap": 50, "keep_separator": False}
    )

    # 2. 核心组件初始化
    # 🔴【挖空预测 2：设备 ID】本地 Embedding 模型必须指定 NPU 设备 ID
    emb = TextEmbedding(model_path="/path/to/embedding", dev_id=0)
    
    # 🔴【挖空预测 3：向量维度】x_dim 必须填，且要与你的 Embedding 模型输出维度一致（如 768 或 1024）
    vector_store = MindFAISS(x_dim=1024, devs=[0], load_local_index="./faiss.index", auto_save=True)
    chunk_store = SQLiteDocstore(db_path="./sql.db")
    
    knowledge_store = KnowledgeStore(db_path="./sql.db")
    knowledge_store.add_knowledge("default_kb", "user1", "admin")
    
    knowledge_db = KnowledgeDB(
        knowledge_store=knowledge_store, chunk_store=chunk_store, 
        vector_store=vector_store, knowledge_name="default_kb", 
        white_paths=["/home/"], user_id="user1"
    )

    # 3. 触发入库动作
    # 🔴【挖空预测 4：一键上传方法】注意传给 embed_func 的是 emb.embed_documents
    upload_files(
        knowledge=knowledge_db, 
        files=["/home/data/test.txt"], 
        loader_mng=loader_mng,
        embed_func=emb.embed_documents, 
        force=True
    )

from mx_rag.retrievers import Retriever
from mx_rag.llm import Text2TextLLM
from mx_rag.chain import SingleText2TextChain
from mx_rag.utils import ClientParam
# (此处省略与建库时完全一样的 emb, vector_store, chunk_store 初始化代码)

def query_knowledge():
    # ... (初始化 emb, vector_store, chunk_store)
    
    # 1. 组装检索器
    # 🔴【挖空预测 5：Retriever 必填参数】
    text_retriever = Retriever(
        vector_store=vector_store, 
        document_store=chunk_store,
        embed_func=emb.embed_documents, 
        k=3,                  # 召回 Top-K 的数量
        score_threshold=0.5   # 相似度阈值
    )

    # 2. 组装大模型客户端
    # 🔴【挖空预测 6：大模型网络参数】
    client_param = ClientParam(use_http=True, timeout=120)
    llm = Text2TextLLM(
        base_url="http://127.0.0.1:1025/v1/chat/completions", 
        model_name="Qwen-14B-Chat", 
        client_param=client_param
    )

    # 3. 串联为 RAG 工作流并触发提问
    # 🔴【挖空预测 7：串联 API】全场最经典的封装链！
    rag_chain = SingleText2TextChain(retriever=text_retriever, llm=llm)
    
    # 🔴【挖空预测 8：触发推理】
    response = rag_chain.query("请总结文章的核心观点")
    print(response)

import numpy as np
import faiss

def native_faiss_demo():
    d = 1024  # 向量维度
    # 生成一些随机假数据作为文档向量 (50个文档，1024维)
    xb = np.random.random((50, d)).astype('float32')
    # 生成 1 个用户提问的向量
    xq = np.random.random((1, d)).astype('float32')

    # 🔴【挖空预测 1：创建索引】最基础的精确 L2 距离索引
    index = faiss.IndexFlatL2(d) 
    
    # 🔴【挖空预测 2：NPU 硬件加速机制】如何把 CPU 上的索引挪到 NPU 上？
    # 昇腾环境下可能会引入 ascendfaiss，调用 index_cpu_to_npu (参考 CUDA 的 gpu 逻辑)
    # 此处若挖空，记住寻找带 npu 字样的 API
    # index = ascendfaiss.index_cpu_to_npu(0, index) 

    # 🔴【挖空预测 3：添加数据与检索】
    index.add(xb)             # 将知识库向量添加到索引中
    k = 4                     # 我们需要返回 4 个最相似的文档
    D, I = index.search(xq, k) # D 是距离(Distance)，I 是索引号(Index)
    
    print("相似文档的距离:", D)
    print("相似文档的ID:", I)

import requests
import json

def call_mindie():
    # 🔴【挖空预测 1：请求端点】标准的 completions 接口路径
    url = "http://127.0.0.1:1025/v1/chat/completions"
    
    # 🔴【挖空预测 2：HTTP Header】必须要声明 Content-Type
    headers = {
        "Content-Type": "application/json"
    }
    
    # 🔴【挖空预测 3：消息体结构】必须是 role 和 content 的字典列表
    payload = {
        "model": "Llama3-8B-Chinese-Chat",
        "messages": [
            {"role": "system", "content": "你是一个有用的AI助手。"},
            {"role": "user", "content": "贵阳的美食有哪些？"}
        ],
        "max_tokens": 512,
        "temperature": 0.7,
        "stream": False  # 🔴【填空点】如果不写流式回调，这里必须是 False
    }
    
    # 🔴【挖空预测 4：发送 POST 请求】
    response = requests.post(url, headers=headers, json=payload)
    
    # 解析返回值
    result = response.json()
    # 🔴【挖空预测 5：结果路径解析】标准 OpenAI 格式的返回值获取路径
    answer = result["choices"][0]["message"]["content"]
    print(answer)

from mx_rag.document import LoaderMng
from mx_rag.document.loader import ImageLoader # 引入图片加载器
from mx_rag.llm import MultiModalLLM # 引入多模态大模型封装

def multi_modal_rag_demo():
    # 1. 注册图片解析器
    loader_mng = LoaderMng()
    # 🔴【挖空预测 1：图片加载器注册】
    loader_mng.register_loader(loader_class=ImageLoader, file_types=[".jpg", ".png"])
    
    # 2. 多模态大模型连接
    # 🔴【挖空预测 2：使用 MultiModalLLM】不再是 Text2TextLLM
    client_param = ClientParam(use_http=True)
    mllm = MultiModalLLM(
        base_url="http://127.0.0.1:1025/v1/chat/completions", 
        model_name="Qwen-VL", 
        client_param=client_param
    )
    
    # 3. 构造图文混排的提问
    # 🔴【挖空预测 3：多模态内容构造】在标准的 messages 结构里，content 变成一个列表，包含 text 和 image_url
    query_messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "请描述这张图片里的设备型号是什么？"},
                # 传入本地图片或 base64
                {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,....."}} 
            ]
        }
    ]
    
    # 直接调用大模型
    response = mllm.chat(query_messages)
    print(response)