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

77 lines
2.0 KiB
Python

# -*- 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())