ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/tests/runners/stream_run.py

@@ -0,0 +1,122 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import asyncio
+import os
+from dataclasses import dataclass, field
+from dotenv import load_dotenv
+from rich.table import Table
+from rich.status import Status
+from rich.console import Console
+
+from aworld.agents.llm_agent import Agent
+from aworld.config.conf import AgentConfig, TaskConfig
+from aworld.core.task import Task
+from aworld.output import MessageOutput, WorkSpace
+from aworld.output.base import StepOutput, ToolResultOutput
+from aworld.output.ui.base import AworldUI
+from aworld.output.utils import consume_content
+from aworld.runner import Runners
+
+
+@dataclass
+class RichAworldUI(AworldUI):
+    console: Console = field(default_factory=Console)
+    status: Status = None
+    workspace: WorkSpace = None
+
+    async def message_output(self, __output__: MessageOutput):
+        result = []
+
+        async def __log_item(item):
+            result.append(item)
+            self.console.print(item, end="")
+
+        if __output__.reason_generator or __output__.response_generator:
+            if __output__.reason_generator:
+                await consume_content(__output__.reason_generator, __log_item)
+            if __output__.reason_generator:
+                await consume_content(__output__.response_generator, __log_item)
+        else:
+            await consume_content(__output__.reasoning, __log_item)
+            await consume_content(__output__.response, __log_item)
+        # if __output__.tool_calls:
+        #     await consume_content(__output__.tool_calls, __log_item)
+        self.console.print("")
+
+    async def tool_result(self, output: ToolResultOutput):
+        """
+            tool_result
+        """
+        table = Table(show_header=False, header_style="bold magenta",
+                      title=f"Call Tools#ID_{output.origin_tool_call.id}")
+        table.add_column("name", style="dim", width=12)
+        table.add_column("content")
+        table.add_row("function_name", output.origin_tool_call.function.name)
+        table.add_row("arguments", output.origin_tool_call.function.arguments)
+        table.add_row("result", output.data)
+        self.console.print(table)
+
+    async def step(self, output: StepOutput):
+        if output.status == "START":
+            self.console.print(f"[bold green]{output.name} ✈️START ...")
+            self.status = self.console.status(f"[bold green]{output.name} RUNNING ...")
+            self.status.start()
+        elif output.status == "FINISHED":
+            self.status.stop()
+            self.console.print(f"[bold green]{output.name} 🛬FINISHED ...")
+        elif output.status == "FAILED":
+            self.status.stop()
+            self.console.print(f"[bold red]{output.name} 💥FAILED ...")
+        else:
+            self.status.stop()
+            self.console.print(f"============={output.name} ❓❓❓UNKNOWN#{output.status} ======================")
+
+
+def run():
+    load_dotenv()
+    agent_config = AgentConfig(
+        llm_provider="openai",
+        llm_model_name=os.environ["LLM_MODEL_NAME"],
+        llm_api_key=os.environ["LLM_API_KEY"],
+        llm_base_url=os.environ["LLM_BASE_URL"]
+    )
+
+    AMAP_API_KEY = os.environ['AMAP_API_KEY']
+    amap_sys_prompt = "You are a helpful agent."
+    amap_agent = Agent(
+        conf=agent_config,
+        name="amap_agent",
+        system_prompt=amap_sys_prompt,
+        mcp_servers=["amap-amap-sse"],  # MCP server name for agent to use
+        history_messages=100,
+        mcp_config={
+            "mcpServers": {
+                "amap-amap-sse": {
+                    "url": f"https://mcp.amap.com/sse?key={AMAP_API_KEY}",
+                    "timeout": 5.0,
+                    "sse_read_timeout": 300.0
+                }
+            }
+        }
+    )
+
+    user_input = (
+        "How long does it take to drive from Hangzhou of Zhejiang to  Weihai of Shandong (generate a table with columns for starting point, destination, duration, distance), "
+        "which cities are passed along the way, what interesting places are there along the route, "
+        "and finally generate the content as markdown and save it")
+
+
+    async def _run(agent, input):
+        task = Task(
+            input=input,
+            agent=agent,
+            conf=TaskConfig()
+        )
+
+        rich_ui = RichAworldUI()
+
+        async for output in Runners.streamed_run_task(task).stream_events():
+            await AworldUI.parse_output(output, rich_ui)
+
+
+    asyncio.run(_run(amap_agent, user_input))