diff --git a/src/agent/service.py b/src/agent/service.py index e0d2cc3..05f50c4 100644 --- a/src/agent/service.py +++ b/src/agent/service.py @@ -1,7 +1,10 @@ from typing import AsyncIterator from src.agent.base import create_agent - +from lambda_agent_api.server import ( + AgentEventUnion, MsgEventTextChunk, MsgEventToolCallChunk, + MsgEventToolResult, MsgEventEnd +) class AgentService: _instance = None @@ -13,21 +16,44 @@ class AgentService: cls._instance._thread_id = "default" return cls._instance - async def astream(self, text: str) -> AsyncIterator[str]: + async def astream(self, text: str) -> AsyncIterator[AgentEventUnion]: config = {"configurable": {"thread_id": self._thread_id}} - async for event in self._agent.astream( + # Используем astream_events для перехвата детальных событий (инструменты, чанки и т.д.) + async for event in self._agent.astream_events( {"messages": [{"role": "user", "content": text}]}, config=config, + version="v2" # Обязательно v2 для современных версий LangChain ): - messages = event.get("messages") or event.get("model", {}).get( - "messages", [] - ) - if messages: - last_msg = messages[-1] - content = getattr(last_msg, "content", None) - if isinstance(content, str) and content.strip(): - yield content + kind = event["event"] + + # 1. Агент генерирует токены (текст или аргументы для инструмента) + if kind == "on_chat_model_stream": + chunk = event["data"]["chunk"] + + # Если генерируется обычный текст + if chunk.content: + yield MsgEventTextChunk(text=chunk.content) + + # Если агент решил использовать инструмент (Langchain выдает tool_call_chunks) + if hasattr(chunk, "tool_call_chunks") and chunk.tool_call_chunks: + for tool_chunk in chunk.tool_call_chunks: + yield MsgEventToolCallChunk( + tool_name=tool_chunk.get("name"), + args_chunk=tool_chunk.get("args") + ) + + # 2. Инструмент завершил работу и вернул результат + elif kind == "on_tool_end": + yield MsgEventToolResult( + tool_name=event["name"], + result=event["data"].get("output") + ) + + # 3. В конце генерации отправляем событие завершения + yield MsgEventEnd(tokens_used=0) # потом заменить на метадату + + def get_agent_service() -> AgentService: