20260510

产品需求文档AI生成/utils/agnet.py

@@ -0,0 +1,62 @@
+# -*- 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:
+    """
+    通用智能体基类
+    """
+
+    model = OpenAIChatModel(
+        model_name="deepseek-v4-flash",
+        provider=OpenAIProvider(
+            base_url="https://tokenhub.tencentmaas.com/v1",
+            api_key="sk-D9Y1mCe8VlvNqLuSC4mAjqEwxJ2nW4C0h8a7EPn8kg9RLsHq",
+        ),
+    )
+
+    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=self.model,
+            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)
+