import os
from typing import Dict, List, Optional
from qwen_agent.agents import Assistant
from qwen_agent.tools import BaseTool

class Qwen3Agent:
    """優(yōu)化后的Qwen3智能體封裝類"""
    
    def __init__(
        self,
        model_name: str = 'Qwen3-30B-A3B',
        endpoint_type: str = 'local',  # 'dashscope' or 'local'
        custom_endpoint: Optional[str] = None,
        enable_tools: bool = True
    ):
        """
        初始化Qwen3智能體
        :param model_name: 模型名稱
        :param endpoint_type: 服務(wù)端點(diǎn)類型 ('dashscope' 或 'local')
        :param custom_endpoint: 自定義OpenAI兼容端點(diǎn)URL
        :param enable_tools: 是否啟用工具功能
        """
        self.llm_cfg = self._init_llm_config(model_name, endpoint_type, custom_endpoint)
        self.tools = self._init_tools() if enable_tools else []
        self.agent = Assistant(llm=self.llm_cfg, function_list=self.tools)
        
    def _init_llm_config(
        self,
        model_name: str,
        endpoint_type: str,
        custom_endpoint: Optional[str]
    ) -> Dict:
        """初始化LLM配置"""
        config = {
            'model': model_name,
            'generate_cfg': {
                'thought_in_content': False,  # 推薦分離思維鏈和響應(yīng)內(nèi)容
                'max_tokens': 2048,
                'temperature': 0.7
            }
        }
        
        if endpoint_type == 'dashscope':
            config.update({
                'model_type': 'qwen_dashscope',
                'api_key': os.getenv('DASHSCOPE_API_KEY')
            })
        else:
            config.update({
                'model_server': custom_endpoint or'http://localhost:8000/v1',
                'api_key': 'EMPTY'
            })
            
        return config
    
    def _init_tools(self) -> List[BaseTool]:
        """初始化工具配置"""
        return [
            {
                'mcpServers': {
                    'time': {
                        'command': 'uvx',
                        'args': ['mcp-server-time', '--local-timeznotallow=Asia/Shanghai']
                    },
                    "fetch": {
                        "command": "uvx",
                        "args": ["mcp-server-fetch"]
                    }
                }
            },
            'code_interpreter',  # 內(nèi)置代碼解釋器
            'web_search',       # 新增網(wǎng)絡(luò)搜索工具
        ]
    
    def stream_response(
        self,
        messages: List[Dict],
        timeout: int = 30
    ) -> Dict:
        """
        流式響應(yīng)生成
        :param messages: 對(duì)話消息列表
        :param timeout: 超時(shí)時(shí)間(秒)
        :return: 最終響應(yīng)字典
        """
        try:
            for response in self.agent.run(messages=messages, timeout=timeout):
                # 可在此處添加實(shí)時(shí)處理邏輯
                pass
            return response
        except Exception as e:
            self._handle_error(e)
            return {'error': str(e)}
    
    def _handle_error(self, error: Exception):
        """統(tǒng)一錯(cuò)誤處理"""
        # 可擴(kuò)展為特定錯(cuò)誤類型的處理
        print(f"Agent Error: {str(error)}")
        # 可添加日志記錄或警報(bào)邏輯

# 使用示例
if __name__ == "__main__":
    # 初始化智能體 (生產(chǎn)環(huán)境建議使用環(huán)境變量配置API密鑰)
    agent = Qwen3Agent(
        model_name='Qwen3-30B-A3B',
        endpoint_type='local',
        custom_endpoint='http://api.example.com/v1'
    )
    
    # 構(gòu)造對(duì)話消息
    messages = [{
        'role': 'user',
        'content': 'https://qwenlm.github.io/blog/ 介紹Qwen的最新發(fā)展動(dòng)態(tài)'
    }]
    
    # 獲取響應(yīng)
    response = agent.stream_response(messages)
    
    # 處理響應(yīng)
    if'error'notin response:
        print("完整響應(yīng):")
        print(response)
    else:
        print("請(qǐng)求失敗:", response['error'])