64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
任务模块
|
||
"""
|
||
from typing import AsyncGenerator, List
|
||
|
||
from pydantic_ai import Agent, ModelMessage
|
||
from pydantic_ai.messages import AgentStreamEvent
|
||
from pydantic_ai.run import AgentRunResultEvent
|
||
|
||
from application.states.models import TaskNodeResultEvent, TaskType, Task, Message
|
||
from application.tasks.models import DEEPSEEK_V4_FLASH_MODEL
|
||
|
||
|
||
instruction = """
|
||
# 角色
|
||
专业友好AI助手，结构化解答各类问题。
|
||
|
||
# 输出硬性规则
|
||
1. 全文强制标准Markdown，禁止纯文本；不要额外说明排版格式，直接输出内容；
|
||
2. 层级使用 `#/##/###`，列表用 `-` 无序列表或数字有序列表；
|
||
3. 代码块用 ```语言名``` 包裹；
|
||
4. 重点内容标注 **粗体**/*斜体*；
|
||
5. 思考、工具日志仅输出文本，适配前端折叠面板，禁止输出HTML标签；
|
||
6. 内容分点拆分，排版整洁适配前端Markdown渲染。
|
||
|
||
# 行文要求
|
||
语言通俗，逻辑完整简洁，无多余废话。
|
||
"""
|
||
|
||
|
||
async def run_stream_events(
|
||
work: Work | None,
|
||
user_prompt: Message,
|
||
message_history: List[ModelMessage],
|
||
) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent | TaskNodeResultEvent, None]:
|
||
"""
|
||
以流式事件模式运行
|
||
"""
|
||
if task:
|
||
match task.type:
|
||
case TaskType.BOOK_FLIGHT:
|
||
from application.tasks.book_flight import run_stream_events
|
||
|
||
async for event in run_stream_events(
|
||
task=task,
|
||
user_prompt=user_prompt,
|
||
message_history=message_history,
|
||
):
|
||
print(event)
|
||
yield event
|
||
|
||
else:
|
||
agent = Agent(
|
||
model=DEEPSEEK_V4_FLASH_MODEL,
|
||
instructions=instruction,
|
||
)
|
||
async with agent.run_stream_events(
|
||
user_prompt=user_prompt,
|
||
message_history=message_history,
|
||
) as events:
|
||
async for event in events:
|
||
yield event
|