diff --git a/.gitignore b/.gitignore index 274911b..974ad1b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ .idea/ +workspace/ # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/docker-compose.yml b/docker-compose.yml index 32ec3ae..95960f4 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -13,6 +13,8 @@ services: build: context: . target: development + tags: + - mrkan0/lambda-agent:dev additional_contexts: agent_api: ${AGENT_API_PATH} volumes: 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: diff --git a/src/api/external.py b/src/api/external.py index 219d5f7..c93f2cb 100644 --- a/src/api/external.py +++ b/src/api/external.py @@ -40,5 +40,5 @@ async def process_message(ws: WebSocket, msg, agent_service: AgentService): match msg: case MsgUserMessage(): async for chunk in agent_service.astream(msg.text): - await ws.send_text(MsgEventTextChunk(text=chunk).model_dump_json()) + await ws.send_text(chunk.model_dump_json()) await ws.send_text(MsgEventEnd(tokens_used=0).model_dump_json())