diff --git a/src/agent/base.py b/src/agent/base.py index d85775f..17707ee 100644 --- a/src/agent/base.py +++ b/src/agent/base.py @@ -3,12 +3,20 @@ import os from deepagents import create_deep_agent from langchain_openai import ChatOpenAI from langgraph.checkpoint.memory import MemorySaver +from langgraph.graph.state import CompiledStateGraph from src.agent.backends import IsolatedShellBackend from src.agent.tools import send_file -def create_agent(): +class Agent(CompiledStateGraph): + """ + Временный (надеюсь) костыль, чтобы дать доступ сервису к файловой системе агента. + """ + backend: IsolatedShellBackend + + +def create_agent() -> Agent: model = ChatOpenAI( model=os.environ["PROVIDER_MODEL"], base_url=os.environ["PROVIDER_URL"], @@ -24,10 +32,15 @@ def create_agent(): virtual_mode=True, ) - return create_deep_agent( + # noinspection PyTypeChecker + # create_deep_agent возвращает CompiledStateGraph, но ниже мы его дополняем так, чтобы он соответствовал сигнатуре Agent + agent: Agent = create_deep_agent( model=model, tools=[send_file], system_prompt="You are a helpful assistant.", checkpointer=MemorySaver(), backend=backend, ) + agent.backend = backend + + return agent diff --git a/src/agent/service.py b/src/agent/service.py index 90c4663..11c50be 100644 --- a/src/agent/service.py +++ b/src/agent/service.py @@ -16,7 +16,13 @@ class ChatBusyError(Exception): """ Чат занят в другом блоке ``with`` """ + pass + +class AttachmentError(Exception): + """ + Ошибка при работе с вложением. + """ pass @@ -33,7 +39,7 @@ class AgentChat(AsyncContextManager[Self]): chat_id: int @abstractmethod - def astream(self, text: str) -> AsyncIterator[AgentEventUnion]: ... + def astream(self, text: str, attachments: list[str] = None) -> AsyncIterator[AgentEventUnion]: ... class AgentService: @@ -71,11 +77,11 @@ class AgentService: async def __aexit__(self, exc_type, exc_val, exc_tb): self.__locks.remove(self.__chat_id) - def astream(self, text: str) -> AsyncIterator[AgentEventUnion]: + def astream(self, text: str, attachments: list[str] = None) -> AsyncIterator[AgentEventUnion]: if not self.__chat_id in self.__locks: raise RuntimeError("Chat must be used in `with` statement") - return self.__service._AgentService__astream(self.__chat_id, text) + return self.__service._AgentService__astream(self.__chat_id, text, attachments) def chat(self, chat_id: int) -> AgentChat: """ @@ -84,13 +90,18 @@ class AgentService: return self.__AgentChat(self, chat_id) async def __astream( - self, chat_id: int, text: str + self, chat_id: int, text: str, attachments: list[str] = None ) -> AsyncIterator[AgentEventUnion]: config = {"configurable": {"thread_id": chat_id}} + new_message = text + if attachments: + attachments_description = await self.__describe_attachments(attachments) + new_message += "\n" + attachments_description + # Используем astream_events для перехвата детальных событий (инструменты, чанки и т.д.) async for event in self._agent.astream_events( - {"messages": [{"role": "user", "content": text}]}, + {"messages": [{"role": "user", "content": new_message}]}, config=config, version="v2", # Обязательно v2 для современных версий LangChain ): @@ -125,3 +136,20 @@ class AgentService: # 3. В конце генерации отправляем событие завершения yield MsgEventEnd(tokens_used=0) # потом заменить на метадату + + async def __describe_attachments(self, raw_paths: list[str]) -> str: + lines = [] + for raw_path in raw_paths: + try: + p = self._agent.backend._resolve_path(raw_path) + rel = p.relative_to(self._agent.backend.cwd) + if not p.exists(): + raise FileNotFoundError(f"File {rel.as_posix()} not found") + lines.append(f"- {rel.as_posix()}") + except Exception as e: + raise AttachmentError(f"Failed to validate attachment {raw_path}: {str(e)}") from e + + return (f"К сообщению приложены {len(lines)} файлов. " + f"Ниже даны пути до этих файлов относительно рабочей директории:\n") + "\n".join(lines) + + diff --git a/src/agent/tools.py b/src/agent/tools.py index 4afd6e0..195fdf6 100644 --- a/src/agent/tools.py +++ b/src/agent/tools.py @@ -25,7 +25,7 @@ async def send_file(path: str) -> str: Используй этот tool без явного запроса от пользователя. Args: - path: Путь к файлу относительно /workspace (например: 'report.pdf', 'docs/readme.txt', 'output/data.json') + path: Путь к файлу относительно рабочей директории (например: 'report.pdf', 'docs/readme.txt', 'output/data.json') Returns: Подтверждение отправки или сообщение об ошибке @@ -39,7 +39,7 @@ async def send_file(path: str) -> str: full_path = Path(workspace) / input_path if not full_path.exists(): - return f"Ошибка: файл '{path}' не найден в /workspace" + return f"Ошибка: файл '{path}' не найден" if not full_path.is_file(): return f"Ошибка: '{path}' не является файлом" diff --git a/src/api/external.py b/src/api/external.py index d0d9445..ecd8bf1 100644 --- a/src/api/external.py +++ b/src/api/external.py @@ -43,6 +43,6 @@ async def websocket_endpoint( async def process_message(ws: WebSocket, chat: AgentChat, msg): match msg: case MsgUserMessage(): - async for chunk in chat.astream(msg.text): + async for chunk in chat.astream(msg.text, msg.attachments): await ws.send_text(chunk.model_dump_json()) await ws.send_text(MsgEventEnd(tokens_used=0).model_dump_json())