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

# 列举导入模块
import asyncio
from typing import List

from pydantic_ai import Agent, AgentRunResult
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.output import OutputSpec
from pydantic_ai.providers.openai import OpenAIProvider


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

    def __init__(self, instructions: str, output_type: OutputSpec = str):
        """
        初始化智能体
        :param instructions: 指令
        :param output_type: 输出类型
        :return: 智能体实例
        """
        # 实例智能体
        self.agent = self._instantiate(
            instructions=instructions,
            output_type=output_type,
        )

    def _instantiate(self, instructions: str, output_type: OutputSpec) -> Agent:
        """
        实例智能体
        :param instructions: 指令
        :param output_type: 输出类型
        :return: 智能体实例
        """
        agent = Agent(
            model=OpenAIChatModel(
                model_name="deepseek-v4-flash",
                provider=OpenAIProvider(
                    base_url="https://tokenhub.tencentmaas.com/v1",
                    api_key="sk-D9Y1mCe8VlvNqLuSC4mAjqEwxJ2nW4C0h8a7EPn8kg9RLsHq",
                ),
            ),
            instructions=instructions,
            output_type=output_type,
        )
        return agent

    async def run(self, user_prompt: str | List[str]) -> AgentRunResult:
        """
        异步运行智能体
        :param user_prompt: 用户提示词
        :return: 智能体回复
        """
        return await self.agent.run(user_prompt=user_prompt)


async def test():
    # 1. 创建智能体（给系统提示词）
    agent = BaseAgent(instructions="请用一句话简洁回复。")

    # 2. 运行智能体（给用户问题）
    result = await agent.run(user_prompt="Hello World 最早出现在哪里？")

    # 3. 输出结果
    print("答案：", result.output)


# 运行
if __name__ == "__main__":
    asyncio.run(test())