# -*- coding: utf-8 -*-
"""
智能体模块
"""

# 列举导入模块
from pydantic_ai import Agent
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.models.openai import OpenAIChatModel

from typing import Any, Optional, List


class BaseAgent:
    """
    通用智能体基类
    """

    def __init__(self, system_prompt: str, output_type: Any, tools: List = []):
        """
        初始化智能体
        :param system_prompt: 系统提示词
        :param output_type: 输出类型
        :param tools: 工具列表
        :return: 智能体实例
        """
        self.agent = self._instantiate_agent(
            system_prompt=system_prompt, output_type=output_type, tools=tools
        )

    def _instantiate_agent(
        self, system_prompt: str, output_type: Any, tools: list = []
    ) -> Agent:
        """
        实例化智能体
        """
        agent = Agent(
            model=OpenAIChatModel(
                model_name="deepseek-v4-flash",
                provider=OpenAIProvider(
                    base_url="https://tokenhub.tencentmaas.com/v1",
                    api_key="sk-D9Y1mCe8VlvNqLuSC4mAjqEwxJ2nW4C0h8a7EPn8kg9RLsHq",
                ),
            ),
            system_prompt=system_prompt,
            output_type=output_type,
        )
        for tool in tools:
            agent.tool(tool)  # 注册工具
        return agent

    async def run(self, user_prompt: str, **kwargs):
        """
        运行智能体
        :param user_prompt: 用户提示词
        :param kwargs: 其它参数
        :return: 智能体回复
        """
        return await self.agent.run(user_prompt=user_prompt, **kwargs)