60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
# -*- 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)
|