1、使用异步API

src/mini_agent/agent/async_llm_client.py

@@ -0,0 +1,39 @@
+import os
+
+from openai import AsyncOpenAI
+from openai.types.chat import (
+    ChatCompletion,
+    ChatCompletionMessageParam,
+    ChatCompletionToolChoiceOptionParam,
+    ChatCompletionToolParam,
+)
+
+
+class AsyncLLMClient:
+    def __init__(
+        self,
+        model: str | None = None,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        timeout: int | None = None,
+    ) -> None:
+        self.model = model or os.getenv("LLM_MODEL_ID")
+        self.base_url = base_url or os.getenv("LLM_BASE_URL")
+        self.api_key = api_key or os.getenv("LLM_API_KEY")
+        self.timeout = timeout if timeout is not None else int(os.getenv("LLM_TIMEOUT", "60"))
+
+        if not all([self.model, self.base_url, self.api_key]):
+            raise ValueError("模型ID、API密钥和服务地址必须被提供或在.env文件中定义。")
+        self.client = AsyncOpenAI(base_url=self.base_url, api_key=self.api_key, timeout=self.timeout)
+
+    async def think(
+        self,
+        messages: list[ChatCompletionMessageParam],
+        tools: list[ChatCompletionToolParam],
+        tool_choice: ChatCompletionToolChoiceOptionParam = "auto",
+        **kwargs,
+    ) -> ChatCompletion:
+        assert self.model is not None
+        return await self.client.chat.completions.create(
+            model=self.model, messages=messages, tools=tools, tool_choice=tool_choice, **kwargs
+        )

src/mini_agent/agent/base_agent.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from collections.abc import Iterator
+from collections.abc import AsyncIterator
 import json
 from typing import Any
 
@@ -65,5 +65,5 @@ class BaseAgent:
         return message_param
 
     @abstractmethod
-    def run(self, question: str, context: ConversationContext) -> Iterator[Message]:
+    def run(self, question: str, context: ConversationContext) -> AsyncIterator[Message]:
         pass

src/mini_agent/agent/context.py

@@ -1,3 +1,5 @@
+from uuid import uuid7
+
 from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel, Field, PrivateAttr
 
@@ -5,6 +7,12 @@ from mini_agent.tools.tool import Tool
 
 
 class ConversationContext(BaseModel):
+
+    @staticmethod
+    def create(run_id:str = "")->ConversationContext:
+        context = ConversationContext(conversation_id=str(uuid7()),run_id=run_id)
+        return context
+
     conversation_id: str = Field(description="会话ID")
     run_id: str = Field(description="当前会话轮id")
     messages: list[ChatCompletionMessageParam] = Field(description="消息历史", default_factory=list)

src/mini_agent/agent/react_agent.py

@@ -1,11 +1,11 @@
-from collections.abc import Iterator
+from collections.abc import AsyncIterator
 import os
 from typing import override
 from uuid import uuid7
 
+from mini_agent.agent.async_llm_client import AsyncLLMClient
 from mini_agent.agent.base_tool_agent import BaseToolAgent
 from mini_agent.agent.context import ConversationContext
-from mini_agent.agent.llm_client import LLMClient
 from mini_agent.agent.message import (
     ErrorContent,
     Message,
@@ -24,7 +24,7 @@ class ReActAgent(BaseToolAgent):
     def __init__(
         self,
         tool_repository: ToolRepository,
-        llm_client: LLMClient,
+        llm_client: AsyncLLMClient,
         max_steps: int = 10,
     ) -> None:
 
@@ -80,14 +80,14 @@ class ReActAgent(BaseToolAgent):
             """
 
     @override
-    def run(self, question: str, context: ConversationContext) -> Iterator[Message]:
+    async def run(self, question: str, context: ConversationContext) -> AsyncIterator[Message]:
         self._initialize_context(context, question)
         run_id = str(uuid7())
 
         for _ in range(self.max_steps):
             tools = [tool.to_openai() for tool in context.discovered_tools]
 
-            response = self.client.think(
+            response = await self.client.think(
                 messages=context.messages,
                 tools=tools,
                 tool_choice="auto",

src/mini_agent/agent/route_agent.py

@@ -1,12 +1,12 @@
-from collections.abc import Iterator
+from collections.abc import AsyncIterator
 from typing import Literal, override
 from uuid import uuid7
 
 from pydantic import BaseModel, Field
 
+from mini_agent.agent.async_llm_client import AsyncLLMClient
 from mini_agent.agent.base_agent import BaseAgent
 from mini_agent.agent.context import ConversationContext
-from mini_agent.agent.llm_client import LLMClient
 from mini_agent.agent.message import (
     ErrorContent,
     Message,
@@ -26,7 +26,7 @@ class RouteDecision(BaseModel):
 
 
 class RouteAgent(BaseAgent):
-    def __init__(self, llm_client: LLMClient, max_steps: int = 10) -> None:
+    def __init__(self, llm_client: AsyncLLMClient, max_steps: int = 10) -> None:
         super().__init__()
         self.sub_agents: dict[str, BaseAgent] = {}
         self.llm_client = llm_client
@@ -143,13 +143,13 @@ class RouteAgent(BaseAgent):
         return template.format("\n".join(agent_descriptions))
 
     @override
-    def run(self, question: str, context: ConversationContext) -> Iterator[Message]:
+    async def run(self, question: str, context: ConversationContext) -> AsyncIterator[Message]:
         self._initialize_context(context=context, question=question)
         run_id = context.run_id
         for _ in range(self.max_steps):
-            response = self.llm_client.think(
+            response = await self.llm_client.think(
                 messages=context.messages,
-                tools = [],
+                tools=[],
                 tool_choice="auto",
                 response_format={
                     "type": "json_schema",
@@ -190,11 +190,13 @@ class RouteAgent(BaseAgent):
                             conversation_id=context.conversation_id,
                             run_id=context.run_id,
                             run_status=RunStatus.RUNNING,
-                            role = MessageRole.ASSISTANT,
-                            content= [ThinkContent(text=decision.reason  or "正在路由请求")]
+                            role=MessageRole.ASSISTANT,
+                            content=[ThinkContent(text=decision.reason or "正在路由请求")],
                         )
-                        agent_context = ConversationContext(conversation_id=context.conversation_id,run_id=str(uuid7()))
-                        for agent_message in agent.run(
+                        agent_context = ConversationContext(
+                            conversation_id=context.conversation_id, run_id=str(uuid7())
+                        )
+                        async for agent_message in agent.run(
                             context=agent_context, question=decision.task or question
                         ):
                             yield Message(

src/mini_agent/api/websocket.py

@@ -12,7 +12,7 @@ router = APIRouter()
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
     agent = websocket.app.state.agent
-    context = ConversationContext(conversation_id=str(uuid7()),run_id="")
+    context = ConversationContext.create()
 
     try:
         while True:
@@ -40,8 +40,17 @@ async def websocket_endpoint(websocket: WebSocket):
                 )
                 continue
 
-            for message in agent.run(question, context):
-                await websocket.send_json(message.model_dump(mode="json"))
+            try:
+                async for message in agent.run(question, context):
+                    await websocket.send_json(message.model_dump(mode="json"))
+            except Exception as e:
+                await websocket.send_json(Message(
+                    conversation_id=context.conversation_id,
+                    role=MessageRole.ASSISTANT,
+                    run_id=context.run_id,
+                    run_status=RunStatus.FAILED,
+                    content=[TextContent(text="执行出错")],
+                ).model_dump(mode="json"))
 
     except WebSocketDisconnect:
         print("客户端断开连接")

src/mini_agent/main.py

@@ -5,8 +5,8 @@ from uuid import uuid7
 from dotenv import load_dotenv
 from fastapi import FastAPI
 
+from mini_agent.agent.async_llm_client import AsyncLLMClient
 from mini_agent.agent.context import ConversationContext
-from mini_agent.agent.llm_client import LLMClient
 from mini_agent.agent.message import TextContent, ToolCallContent, ToolResultContent
 from mini_agent.agent.react_agent import ReActAgent
 from mini_agent.agent.route_agent import RouteAgent
@@ -14,14 +14,15 @@ from mini_agent.api.websocket import router
 from mini_agent.tools.tool_repository import ToolRepository
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
-print(PROJECT_ROOT)
 load_dotenv(PROJECT_ROOT / ".env")
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    react_agent =  ReActAgent(tool_repository=ToolRepository(), llm_client=LLMClient(),max_steps=50)
-    route_agent = RouteAgent(llm_client=LLMClient())
+    react_agent = ReActAgent(
+        tool_repository=ToolRepository(), llm_client=AsyncLLMClient(), max_steps=50
+    )
+    route_agent = RouteAgent(llm_client=AsyncLLMClient())
     route_agent.register_agent(react_agent)
     app.state.agent = route_agent
 
@@ -41,14 +42,14 @@ async def health():
     }
 
 
-if __name__ == "__main__":
-    agent = ReActAgent(tool_repository=ToolRepository(), llm_client=LLMClient())
-    context = ConversationContext(conversation_id="1",run_id="")
+async def main():
+    agent = ReActAgent(tool_repository=ToolRepository(), llm_client=AsyncLLMClient())
+    context = ConversationContext(conversation_id="1", run_id="")
     while True:
         question = input("提问：")
         context.run_id = str(uuid7())
         try:
-            for message in agent.run(question=question, context=context):
+            async for message in agent.run(question=question, context=context):
                 for content in message.content:
                     match content:
                         case TextContent(text=text):
@@ -64,3 +65,8 @@ if __name__ == "__main__":
                             pass
         except Exception as e:
             print(f"❌{e}")
+
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())