import asyncio
import json
import requests
import codecs
from typing import Dict, Any

from mars_toolkit.core.llm_tools import llm_tool
from mars_toolkit.core.config import config

@llm_tool(
    name="retrieval_from_knowledge_base",
    description="Retrieve information from local materials science literature knowledge base"
)
async def retrieval_from_knowledge_base(query: str, topk: int = 3) -> str:
    """
    检索本地材料科学文献知识库中的相关信息

    Args:
        query: 查询字符串，如材料名称"CsPbBr3"
        topk: 返回结果数量，默认3条

    Returns:
        包含文档ID、标题和相关性分数的字典
    """
    # 设置Dify API的URL端点
    url = f'{config.DIFY_ROOT_URL}/v1/chat-messages'

    # 配置请求头，包含API密钥和内容类型
    headers = {
        'Authorization': f'Bearer {config.DIFY_API_KEY}',
        'Content-Type': 'application/json'
    }

    # 准备请求数据
    data = {
        "inputs": {"topK": topk},  # 设置返回的最大结果数量
        "query": query,  # 设置查询字符串
        "response_mode": "blocking",  # 使用阻塞模式，等待并获取完整响应
        "conversation_id": "",  # 不使用会话ID，每次都是独立查询
        "user": "abc-123"  # 用户标识符
    }

    try:
        # 发送POST请求到Dify API并获取响应
        # 设置较长的超时时间(1111秒)以处理可能的长时间响应
        response = requests.post(url, headers=headers, json=data, timeout=1111)

        # 获取响应文本
        response_text = response.text

        # 解码响应文本中的Unicode转义序列
        response_text = codecs.decode(response_text, 'unicode_escape')

        # 将响应文本解析为JSON对象
        result_json = json.loads(response_text)

        # 从响应中提取元数据
        metadata = result_json.get("metadata", {})

        # 构建包含关键信息的结果字典
        useful_info = {
            "id": metadata.get("document_id"),  # 文档ID
            "title": result_json.get("title"),  # 文档标题
            "content": result_json.get("answer", ""),  # 内容字段，使用'answer'字段存储内容
            "score": metadata.get("score")  # 相关性分数
        }

        # 返回提取的有用信息
        return json.dumps(useful_info, ensure_ascii=False, indent=2)

    except Exception as e:
        # 捕获并处理所有可能的异常，返回错误信息
        return f"错误: {str(e)}"