from llama_index.core import PromptTemplate
from llama_index.llms.ollama import Ollama 

# 定义路由选项
choices = [
    "Useful for questions related to apples",
    "Useful for questions related to oranges",
]

# 格式化选项为字符串（带编号）
def get_choice_str(choices):
    choices_str = "\n\n".join(
        [f"{idx+1}. {c}" for idx, c in enumerate(choices)]
    )
    return choices_str

choices_str = get_choice_str(choices)

# 构建严格的路由提示词（仅返回数字）
router_prompt = PromptTemplate(
    """
    请根据以下问题，选择最合适的选项编号（仅返回数字）：  
    问题：{query_str} 
    选项：
    {choices_str}
    """
)

# 初始化 Ollama LLM
llm = Ollama(model="qwen2.5:7b", request_timeout=60.0)  
# 测试查询
query = "What color are oranges?" 

# 填充提示词并调用 LLM
filled_prompt = router_prompt.format(
    query_str=query,
    choices_str=choices_str
)
response = llm.complete(filled_prompt)
print("LLM 选择的选项编号：", response.text.strip())

from llama_index.core import PromptTemplate
from llama_index.llms.ollama import Ollama 
from llama_index.core.types import BaseOutputParser 
import json  

choices = [
    "Useful for questions related to apples",
    "Useful for questions related to oranges",
]

def get_choice_str(choices):
    choices_str = "\n\n".join(
        [f"{idx+1}. {c}" for idx, c in enumerate(choices)]
    )
    return choices_str

choices_str = get_choice_str(choices)

# 升级提示词：要求仅返回 JSON 字符串
router_prompt = PromptTemplate(
    """
    请根据以下问题，选择最合适的选项，并严格按照指定格式返回结果（仅返回JSON字符串，无其他内容）：  
    问题：{query_str} 
    选项：
    {choices_str}
    
    输出格式（JSON）：
    {"choice": 选中的数字（如1/2）}
    """
)

# 自定义输出解析器
class SimpleRouterParser(BaseOutputParser):
    def __init__(self, choices: list):
        self.choices = choices

    def parse(self, output: str) -> int:
        try:
            # 尝试解析 JSON 格式
            output_dict = json.loads(output.strip())
            print(f"LLM 输出 JSON：{output_dict}")
            choice_num = int(output_dict["choice"]) 
            
            # 校验选项范围
            if choice_num < 1 or choice_num > len(self.choices):
                raise ValueError(f"选项编号{choice_num}超出范围（1-{len(self.choices)}）")
            
            return choice_num
        
        # 兼容纯数字输出的降级处理
        except (json.JSONDecodeError, KeyError, ValueError) as e:
            try:
                choice_num = int(output.strip())
                if choice_num < 1 or choice_num > len(self.choices):
                    raise ValueError(f"选项编号{choice_num}超出范围（1-{len(self.choices)}）")
                return choice_num
            except:
                raise RuntimeError(f"解析失败：{str(e)}，原始输出：{output}")

# 初始化 LLM 并执行路由
llm = Ollama(model="qwen2.5:7b", request_timeout=60.0)  
query = "What color are oranges?" 

filled_prompt = router_prompt.format(
    query_str=query,
    choices_str=choices_str
)
response = llm.complete(filled_prompt)

# 解析输出
parser = SimpleRouterParser(choices=choices)
try:
    choice_num = parser.parse(response.text)
    print("LLM 选择的选项编号：", choice_num)
except RuntimeError as e:
    print("解析失败：", e)

from dataclasses import dataclass
from llama_index.core import PromptTemplate
from llama_index.llms.ollama import Ollama

# 定义 Answer 数据类：封装选择结果和理由
@dataclass
class Answer:
    choice: int   # 选中的选项编号
    reason: str   # 选择理由

# 路由选项
route_choices = [
    "Useful for questions related to apples",
    "Useful for questions related to oranges",
]

# 格式化选项
def format_choices(choices):
    choices_str = "\n\n".join([f"{idx+1}. {c}" for idx, c in enumerate(choices)])
    return choices_str

# 构建带格式要求的提示词
router_prompt = PromptTemplate(
    """
    你的任务是根据用户问题，选择最匹配的选项，并按指定格式返回结果。
    
    【用户问题】：{query_str}
    【可选选项】：
    {choices_str}
    
    【输出要求】：
    1. 第一行仅返回选中的选项编号（数字，如1、2）；
    2. 第二行开始返回选择该选项的理由（需清晰说明匹配逻辑）；
    3. 严格遵守格式，不要添加额外内容。
    """
)

# 初始化 LLM
llm = Ollama(model="qwen2.5:7b", request_timeout=60.0)  

# 核心路由函数
def run_llm_router(query: str, choices: list) -> Answer:
    """
    执行LLM路由，返回包含选择结果和理由的Answer对象
    :param query: 用户问题字符串
    :param choices: 路由选项列表
    :return: 封装好的Answer对象
    """
    formatted_choices = format_choices(choices)
    filled_prompt = router_prompt.format(
        query_str=query,
        choices_str=formatted_choices
    )
    response = llm.complete(filled_prompt)
    response_text = response.text.strip()
    
    try:
        # 分割编号和理由
        lines = response_text.split("\n", 1) 
        choice_num = int(lines[0].strip()) 
        reason = lines[1].strip() if len(lines) > 1 else "未提供理由"
        
        # 校验选项范围
        if choice_num < 1 or choice_num > len(choices):
            raise ValueError(f"选择的编号{choice_num}超出选项范围（1-{len(choices)}）")
        
        return Answer(choice=choice_num, reason=reason)
    
    except (ValueError, IndexError) as e:
        raise RuntimeError(f"解析LLM路由结果失败：{str(e)}，原始响应：{response_text}")

# 测试执行
if __name__ == "__main__":
    test_query = "What color are oranges?"  
    try:
        answer = run_llm_router(test_query, route_choices)
        print("=== LLM路由结果 ===")
        print(f"选中选项编号：{answer.choice}")
        print(f"选择理由：{answer.reason}")
    except RuntimeError as e:
        print("路由执行失败：", e)

from dataclasses import dataclass
from llama_index.core import PromptTemplate, Settings 
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding 
from llama_index.core.types import BaseOutputParser
import json 
from pathlib import Path
from llama_index.readers.file import PyMuPDFReader
from llama_index.core import VectorStoreIndex, SummaryIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.query_engine import CustomQueryEngine 
from llama_index.core.response_synthesizers import TreeSummarize
from llama_index.core.base.response.schema import Response 
from pydantic import Field

# 全局配置：初始化嵌入模型和LLM
Settings.embed_model = OllamaEmbedding(
    model_name="nomic-embed-text:latest",
    base_url="http://localhost:11434",
    request_timeout=360.0
)

Settings.llm = Ollama(
    model="qwen2.5:7b", 
    base_url="http://localhost:11434",
    request_timeout=360.0,
    additional_kwargs={"temperature": 0.0, "num_ctx": 2048},
)

# 加载本地文档
loader = PyMuPDFReader()
pdf_path = Path("D:\\LlamaIndex\\LLmaRouter\\data\\flag.txt")
if not pdf_path.exists():
    raise FileNotFoundError(f"文件不存在：{pdf_path.absolute()}")

documents = loader.load(file_path=str(pdf_path))
splitter = SentenceSplitter(chunk_size=1024)  # 文本分割器

# 构建两种索引
vector_index = VectorStoreIndex.from_documents(
    documents, transformations=[splitter], embed_model=Settings.embed_model
)  # 向量索引：用于事实性查询
summary_index = SummaryIndex.from_documents(
    documents, transformations=[splitter]
)  # 摘要索引：用于总结性查询

# 定义 Answer 数据类
@dataclass
class Answer:
    choice: int  # 选中的索引编号
    reason: str  # 选择理由

# 路由选项（索引映射）
route_choices = [
    'Useful for specific/factual questions (e.g., "What is Llama2\'s context window?") → use VectorStoreIndex',
    'Useful for summary/general questions (e.g., "Summarize Llama2\'s key features") → use SummaryIndex',
]

# 格式化选项
def format_choices(choices):
    choices_str = "\n\n".join([f"{idx+1}. {c}" for idx, c in enumerate(choices)])
    print(f"路由选项：{choices_str}")
    return choices_str

# JSON 输出格式模板
JSON_FORMAT_STR = """
The output should be formatted as a JSON instance that conforms to the JSON schema below.
Do NOT add any extra text, explanation, or comments—only output the JSON string.

JSON Schema:
{
  "type": "object",
  "properties": {
    "choice": {"type": "integer"},
    "reason": {"type": "string"}
  },
  "required": ["choice", "reason"],
  "additionalProperties": false
}
"""

# 自定义 JSON 解析器
class RouterOutputParser(BaseOutputParser):
    def parse(self, output: str) -> Answer:
        try:
            output_dict = json.loads(output.strip())
            print(f"LLM路由原始输出：{output_dict}")
            
            # 类型校验
            if not isinstance(output_dict.get("choice"), int):
                raise ValueError(f"choice必须是整数，实际是{type(output_dict.get('choice'))}")
            if not isinstance(output_dict.get("reason"), str):
                raise ValueError(f"reason必须是字符串，实际是{type(output_dict.get('reason'))}")
            
            # 范围校验
            choice_num = output_dict["choice"]
            if choice_num < 1 or choice_num > len(route_choices):
                raise ValueError(f"选择的编号{choice_num}超出选项范围（1-{len(route_choices)}）")
            
            return Answer(choice=choice_num, reason=output_dict["reason"])
        
        except (json.JSONDecodeError, KeyError, ValueError) as e:
            raise RuntimeError(f"解析LLM输出失败：{str(e)}，原始响应：{output}")

# 自定义路由查询引擎
class CustomRouterQueryEngine(CustomQueryEngine):
    vector_index: VectorStoreIndex = Field(description="向量索引，用于事实性查询")
    summary_index: SummaryIndex = Field(description="摘要索引，用于总结性查询")
    llm: Ollama = Field(description="Ollama LLM实例")
    parser: RouterOutputParser = Field(description="路由输出解析器")
    route_choices: list = Field(description="路由选项列表")
    
    def custom_query(self, query_str: str) -> Response:
        try:
            # 1. 格式化路由选项
            formatted_choices = format_choices(self.route_choices)
            
            # 2. 构建并填充提示词
            router_prompt = PromptTemplate(
                f"""
                你的任务是根据用户问题类型，选择最匹配的索引，并严格按照指定JSON格式返回结果。
                
                【用户问题】：{{query_str}}
                【可选索引】：
                {formatted_choices}
                
                {JSON_FORMAT_STR}
                """
            )   
            filled_prompt = router_prompt.format(query_str=query_str)
            
            # 3. 调用 LLM 做路由决策
            llm_response = self.llm.complete(filled_prompt)
            route_answer = self.parser.parse(llm_response.text.strip())
            
            print(f"\n=== 路由决策 ===")
            print(f"选中索引编号：{route_answer.choice}")
            print(f"路由理由：{route_answer.reason}")
            
            # 4. 根据路由结果选择对应索引执行查询
            if route_answer.choice == 1:
                query_engine = self.vector_index.as_query_engine(llm=Settings.llm)
                used_index = "VectorStoreIndex（事实查询）"
            elif route_answer.choice == 2:
                summarizer = TreeSummarize(llm=Settings.llm)
                query_engine = self.summary_index.as_query_engine(
                    response_synthesizer=summarizer,
                    llm=Settings.llm
                )
                used_index = "SummaryIndex（总结查询）"
            else:
                raise ValueError(f"无效的索引编号：{route_answer.choice}")
            
            # 5. 执行查询并返回结果
            final_response = query_engine.query(query_str)
            print(f"\n=== 查询执行 ===")
            print(f"使用索引：{used_index}")
            
            return Response(
                response=str(final_response),
                metadata={
                    "route_choice": route_answer.choice,
                    "route_reason": route_answer.reason,
                    "used_index": used_index,
                    "error": False
                }
            )
        
        except Exception as e:
            error_msg = f"路由查询失败：{str(e)}"
            print(error_msg)
            return Response(
                response=error_msg,
                metadata={"error": True, "error_detail": str(e)}
            )

# 主函数：测试路由查询引擎
if __name__ == "__main__":
    router_parser = RouterOutputParser()
    router_query_engine = CustomRouterQueryEngine(
        vector_index=vector_index,
        summary_index=summary_index,
        llm=Settings.llm,  
        parser=router_parser,
        route_choices=route_choices
    )
    
    # 测试查询
    test_queries = [
        "28.7Blog其网址域名是多少？他的大学是什么专业？目前他学习那个方向？",  
    ]
    
    for idx, test_query in enumerate(test_queries, 1):
        print(f"\n===================== 测试查询 {idx} =====================")
        print(f"查询问题：{test_query}")
        response = router_query_engine.query(test_query)
        
        print(f"\n=== 最终回答 ===")
        print(response.response)
        
        print(f"\n=== 路由元信息 ===")
        print(f"选中索引：{response.metadata.get('used_index')}")
        print(f"路由理由：{response.metadata.get('route_reason')}")
        print(f"是否出错：{response.metadata.get('error')}")