#6 Обработать сырые выходы из агента. Добавлены модели из lambda_agent_api.server для правильной обработки выходов агента

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: