153 lines
4.4 KiB
Python
153 lines
4.4 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
智能体模块
|
||
"""
|
||
|
||
# 列举导入模块
|
||
from pathlib import Path
|
||
|
||
from typing import List, Optional, Union, cast
|
||
|
||
from starlette.applications import Starlette
|
||
from pydantic_ai import Agent as BaseAgent, AgentRunResult
|
||
from pydantic_ai.models import Model
|
||
from pydantic_ai.models.openai import OpenAIChatModel
|
||
from pydantic_ai.capabilities import AgentCapability
|
||
from pydantic_ai.output import OutputSpec
|
||
from pydantic_ai.providers.openai import OpenAIProvider
|
||
from pydantic_ai.ui._web.api import ConfigureFrontend, ModelInfo, BuiltinToolInfo
|
||
from pydantic_ai_skills import SkillsCapability
|
||
|
||
from starlette.routing import Route
|
||
from starlette.requests import Request
|
||
from starlette.responses import Response
|
||
|
||
|
||
|
||
|
||
class Agent:
|
||
"""
|
||
智能体，支持：
|
||
1）实例智能体
|
||
2）异步运行
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
instructions: str,
|
||
output_type: OutputSpec = str,
|
||
capabilities: Optional[List[AgentCapability]] = None,
|
||
):
|
||
"""
|
||
初始化智能体
|
||
:param instructions: 指令
|
||
:param skills: 智能体技能列表，默认为不使用技能
|
||
:param output_type: 输出类型
|
||
:return: 智能体实例
|
||
"""
|
||
# 引入相应模块
|
||
from uuid import uuid4
|
||
from .memory import Memory
|
||
|
||
# 创建智能体
|
||
self.agent = self._create_agent(
|
||
instructions=instructions,
|
||
capabilities=capabilities,
|
||
output_type=output_type,
|
||
)
|
||
|
||
# 生成会话唯一标识
|
||
self.session_id = uuid4().hex.lower()
|
||
|
||
# 实例记忆体
|
||
self.memory = Memory()
|
||
|
||
def _create_agent(
|
||
self,
|
||
instructions: str,
|
||
capabilities: Optional[List[AgentCapability]],
|
||
output_type: OutputSpec,
|
||
) -> BaseAgent:
|
||
"""
|
||
创建智能体
|
||
:param instructions: 指令
|
||
:param capabilities: 智能体能力列表
|
||
:param output_type: 输出类型
|
||
:return: 智能体实例
|
||
"""
|
||
agent = BaseAgent(
|
||
model=OpenAIChatModel(
|
||
model_name="deepseek-v4-flash",
|
||
provider=OpenAIProvider(
|
||
base_url="https://tokenhub.tencentmaas.com/v1",
|
||
api_key="sk-D9Y1mCe8VlvNqLuSC4mAjqEwxJ2nW4C0h8a7EPn8kg9RLsHq",
|
||
),
|
||
),
|
||
instructions=instructions,
|
||
capabilities=capabilities,
|
||
output_type=output_type,
|
||
retries=1,
|
||
)
|
||
return agent
|
||
|
||
async def run(self, user_prompt: str | List[str]) -> AgentRunResult:
|
||
"""
|
||
异步运行
|
||
:param user_prompt: 用户提示词
|
||
:return: 智能体回复
|
||
"""
|
||
# 查询会话历史消息
|
||
message_history = self.memory.read(session_id=self.session_id)
|
||
result = await self.agent.run(
|
||
user_prompt=user_prompt, message_history=message_history
|
||
)
|
||
# 记录会话历史消息
|
||
self.memory.create(
|
||
session_id=self.session_id,
|
||
dialogue_message=result.new_messages(),
|
||
)
|
||
return result
|
||
|
||
def start_starlette_app(self) -> None:
|
||
"""
|
||
启动 Starlette 应用，为智能体提供网页交互式对话界面
|
||
:return: None
|
||
"""
|
||
self.agent.to_web()
|
||
|
||
model = cast(Model, self.agent.model)
|
||
|
||
model_infos = [ModelInfo(id=model.model_id, name=model.label, builtin_tools=[])]
|
||
|
||
async def options_chat(request: Request) -> Response:
|
||
"""处理 OPTIONS 请求"""
|
||
return Response()
|
||
|
||
async def configure_frontend(request: Request) -> Response:
|
||
"""向前端提供模型和技能"""
|
||
config = ConfigureFrontend(
|
||
models=[
|
||
ModelInfo(
|
||
id=model.model_id,
|
||
name=model.label,
|
||
builtin_tools=model.profile.supported_builtin_tools,
|
||
)
|
||
],
|
||
builtin_tools=[],
|
||
)
|
||
return JSONResponse(config.model_dump(by_alias=True))
|
||
|
||
Starlette(
|
||
routes=[
|
||
Route("/chat", options_chat, methods=["OPTIONS"]),
|
||
Route("/chat", post_chat, methods=["POST"]),
|
||
Route("/configure", configure_frontend, methods=["GET"]),
|
||
Route("/health", health, methods=["GET"]),
|
||
]
|
||
)
|
||
|
||
|
||
a = Agent(instructions="你是一个专业的翻译")
|
||
|
||
print(a)
|