From 69f67a7f42bf05fe87b4360f6f501b26a881a756 Mon Sep 17 00:00:00 2001
From: collhoun <2904yr@mail.ru>
Date: Mon, 13 Apr 2026 21:11:55 +0300
Subject: [PATCH] =?UTF-8?q?#6=20=D0=9E=D0=B1=D1=80=D0=B0=D0=B1=D0=BE=D1=82?=
 =?UTF-8?q?=D0=B0=D1=82=D1=8C=20=D1=81=D1=8B=D1=80=D1=8B=D0=B5=20=D0=B2?=
 =?UTF-8?q?=D1=8B=D1=85=D0=BE=D0=B4=D1=8B=20=D0=B8=D0=B7=20=D0=B0=D0=B3?=
 =?UTF-8?q?=D0=B5=D0=BD=D1=82=D0=B0.=20=D0=94=D0=BE=D0=B1=D0=B0=D0=B2?=
 =?UTF-8?q?=D0=BB=D0=B5=D0=BD=D1=8B=20=D0=BC=D0=BE=D0=B4=D0=B5=D0=BB=D0=B8?=
 =?UTF-8?q?=20=D0=B8=D0=B7=20lambda=5Fagent=5Fapi.server=20=D0=B4=D0=BB?=
 =?UTF-8?q?=D1=8F=20=D0=BF=D1=80=D0=B0=D0=B2=D0=B8=D0=BB=D1=8C=D0=BD=D0=BE?=
 =?UTF-8?q?=D0=B9=20=D0=BE=D0=B1=D1=80=D0=B0=D0=B1=D0=BE=D1=82=D0=BA=D0=B8?=
 =?UTF-8?q?=20=D0=B2=D1=8B=D1=85=D0=BE=D0=B4=D0=BE=D0=B2=20=D0=B0=D0=B3?=
 =?UTF-8?q?=D0=B5=D0=BD=D1=82=D0=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/agent/service.py | 48 ++++++++++++++++++++++++++++++++++----------
 1 file changed, 37 insertions(+), 11 deletions(-)

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: