platform/agent_api
5
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Compare commits

base: platform:master
Branches Tags
platform:master
platform:#9-clientside-tool-call
..
compare: platform:#9-clientside-tool-call
Branches Tags
platform:master
platform:#9-clientside-tool-call

No commits in common. "master" and "#9-clientside-tool-call" have entirely different histories.
master ... #9-clients

5 changed files with 112 additions and 216 deletions
Download patch file Download diff file Expand all files Collapse all files

63
README.md
View file

@ -1,16 +1,9 @@
from lambda_agent_api import AgentApi
# Lambda Agent API # Lambda Agent API
WebSocket API SDK для взаимодействия с AI-агентом. WebSocket API SDK для взаимодействия с AI-агентом.
## Release Notes
# v1.1
- **CRITICAL**: `AgentAPI` вместо `url` принимает `base_url` и сам дописывает нужный эндпоинт.
Раньше: `AgentAPI(url="ws://localhost:8000/agent_ws/")`. Сейчас: `AgentAPI(base_url="ws://localhost:8000/")`
- Добавлен параметр `chat_id` в конструктор `AgentAPI`. Нужен для разделения истории сообщений по чатам/веткам.
- `AgentAPI.connect()` вызывает `AgentBusyException`, если выбранный чат уже занят другим API клиентом.
- Добавлен новый тип события `MsgEventSendFile` для отправки файлов пользователю. Поле `path` — путь к файлу относительно `/workspace`.
- `send_message(text, attachments)` теперь принимает `attachments` — список путей к файлам относительно `/workspace`.
## Установка ## Установка
В `master` всегда будет актуальная рабочая версия. В `master` всегда будет актуальная рабочая версия.
```bash ```bash
@ -21,10 +14,42 @@ pip install git+https://git.lambda.coredump.ru/platform/agent_api.git
## Быстрый старт (с использованием AgentApi) ## Быстрый старт (с использованием AgentApi)
**Рабочий REPL пример: [tests/manual.py](tests/manual.py).** ```python
import asyncio
from lambda_agent_api.agent_api import AgentApi
from lambda_agent_api.server import MsgEventTextChunk
## Предполагаемое управление подключениями async def main():
api = AgentApi("agent-1", "ws://localhost:8000/ws")
await api.connect()
try:
response = await api.send_message("Привет, агент!")
async for chunk in response:
if isinstance(chunk, MsgEventTextChunk):
print(chunk.text, end="", flush=True)
elif isinstance(chunk, MsgEventToolCallChunk):
print(f"Tool call started: {chunk.tool_name}")
elif isinstance(chunk, MsgEventToolResult):
print(f"Tool result: {chunk.result}")
elif isinstance(chunk, MsgEventCustomUpdate):
print(f"Progress update: {chunk.payload}")
elif isinstance(chunk, MsgEventEnd):
print(f"Generation ended, tokens used: {chunk.tokens_used}")
finally:
await api.close()
asyncio.run(main())
```
> `AgentApi.send_message()` возвращает стриминг-итерируемый объект, который может выдавать не только текстовые чанки, но и события инструментов (`MsgEventToolCallChunk`, `MsgEventToolResult`, `MsgEventCustomUpdate`) и финальный `MsgEventEnd`.
## Предполагаемое использование
```python ```python
from lambda_agent_api.agent_api import AgentApi from lambda_agent_api.agent_api import AgentApi
@ -166,22 +191,6 @@ async def on_telegram_message(from_user: int, text: str):
| payload | object | Любые данные о прогрессе | | payload | object | Любые данные о прогрессе |
| source | string | Источник события (по умолчанию "main") | | source | string | Источник события (по умолчанию "main") |
#### AGENT_EVENT_SEND_FILE
Агент отправляет файл пользователю. Путь к файлу относительно `/workspace`.
```json
{
"type": "AGENT_EVENT_SEND_FILE",
"path": "reports/2024/report.pdf"
}
```
| Поле | Тип | Описание |
|--------|--------|-----------------------------------------------|
| type | string | Всегда `AGENT_EVENT_SEND_FILE` |
| path | string | Путь к файлу относительно `/workspace` |
#### AGENT_EVENT_END #### AGENT_EVENT_END
Агент закончил генерацию ответа. Агент закончил генерацию ответа.

117
lambda_agent_api/agent_api.py
View file

@ -2,9 +2,6 @@ import logging
from typing import Callable, Optional, AsyncIterator from typing import Callable, Optional, AsyncIterator
import aiohttp import aiohttp
import asyncio import asyncio
from urllib.parse import urljoin
from aiohttp import WSServerHandshakeError
from lambda_agent_api.server import * from lambda_agent_api.server import *
from lambda_agent_api.client import * from lambda_agent_api.client import *
@ -28,14 +25,12 @@ class AgentApi:
def __init__( def __init__(
self, self,
agent_id: str, agent_id: str,
base_url: str, url: str,
callback: Optional[Callable[[ServerMessage], None]] = None, callback: Optional[Callable[[ServerMessage], None]] = None,
on_disconnect: Optional[Callable[["AgentApi"], None]] = None, on_disconnect: Optional[Callable[['AgentApi'], None]] = None
chat_id: int = 0, # значение по умолчанию для обратной совместимости # значение по умолчанию для обратной совместимости
): ):
self.id = agent_id # ID агента для словаря self.id = agent_id # ID агента для словаря
self.chat_id = chat_id self.url = url
self.url = urljoin(base_url, f"v1/agent_ws/{chat_id}/")
self.callback = callback self.callback = callback
self.on_disconnect = on_disconnect self.on_disconnect = on_disconnect
@ -48,11 +43,7 @@ class AgentApi:
self._listen_task: asyncio.Task | None = None self._listen_task: asyncio.Task | None = None
async def connect(self): async def connect(self):
"""Явное подключение к агенту. """Явное подключение к агенту."""
:raise AgentBusyException: Чат занят другим клиентом.
:raise AgentException: Непредвиденная ошибка протокола, см. code и details
"""
self._session = aiohttp.ClientSession() self._session = aiohttp.ClientSession()
try: try:
self._ws = await self._session.ws_connect(self.url, heartbeat=30) self._ws = await self._session.ws_connect(self.url, heartbeat=30)
@ -67,12 +58,10 @@ class AgentApi:
logger.info(f"Agent {self.id} is ready") logger.info(f"Agent {self.id} is ready")
else: else:
raise AgentException( raise AgentException(
"INVALID_STATUS", f"Expected SM.Status, got {status_msg.type}" "INVALID_STATUS", f"Expected SM.Status, got {status_msg.type}")
)
else: else:
raise AgentException( raise AgentException("UNEXPECTED_MSG_TYPE",
"UNEXPECTED_MSG_TYPE", f"Unexpected message type: {msg.type}" f"Unexpected message type: {msg.type}")
)
self._connected = True self._connected = True
self._listen_task = asyncio.create_task(self._listen()) self._listen_task = asyncio.create_task(self._listen())
@ -85,8 +74,7 @@ class AgentApi:
await self._session.close() await self._session.close()
raise AgentException( raise AgentException(
"TIMEOUT", "Agent did not send initial Status message within 5 seconds" "TIMEOUT", "Agent did not send initial Status message within 5 seconds") from e
) from e
except AgentException: except AgentException:
# Перехватываем наши собственные ошибки (INVALID_STATUS, UNEXPECTED_MSG_TYPE), # Перехватываем наши собственные ошибки (INVALID_STATUS, UNEXPECTED_MSG_TYPE),
@ -97,25 +85,6 @@ class AgentApi:
await self._session.close() await self._session.close()
raise raise
except WSServerHandshakeError as e: # если при открытии подключения сервер вернул какую-то ошибку
if self._session and not self._session.closed:
await self._session.close()
# во-первых, aiohttp зачем-то приводит WS коды ошибок к HTTP.
# во-вторых, делает он это некорректно. Любой неизвестный код WS ошибки становится 403 HTTP
# в-третьих, он не передает оригинальное сообщение об ошибке.
# Т. е. какое бы сообщение сервер не отправлял, тут всегда будет "Invalid response status"
# см. site-packages\aiohttp\client.py, line 1104, in _ws_connect
# итого понять реальную причину ошибки почти невозможно. Нужно менять библиотеку
# сейчас сервер специально кидает только ошибку с WS кодом 1008 Policy Violation, когда внутри ловит ChatBusyError
# поэтому скорее всего, если мы получили WSServerHandshakeError с 403, то это внутренний ChatBusyError
if e.status != 403:
# обрабатываем как обычную ошибку WS по примеру except блока ниже
raise AgentException(code="CONNECTION_ERROR",
details=f"Failed to connect agent {self.id}: {e}") from e
raise AgentBusyException(f"Chat {self.chat_id} is already in use by other client")
except Exception as e: except Exception as e:
# Обработка всех остальных ошибок (например, aiohttp.ClientConnectionError) # Обработка всех остальных ошибок (например, aiohttp.ClientConnectionError)
if self._ws and not self._ws.closed: if self._ws and not self._ws.closed:
@ -124,10 +93,8 @@ class AgentApi:
await self._session.close() await self._session.close()
# Можно оставить RuntimeError, а можно тоже завернуть в AgentException # Можно оставить RuntimeError, а можно тоже завернуть в AgentException
raise AgentException( raise AgentException(code="CONNECTION_ERROR",
code="CONNECTION_ERROR", details=f"Failed to connect agent {self.id}: {e}") from e
details=f"Failed to connect agent {self.id}: {e}",
) from e
async def close(self): async def close(self):
"""Явное ручное закрытие соединения.""" """Явное ручное закрытие соединения."""
@ -164,25 +131,19 @@ class AgentApi:
except Exception as e: except Exception as e:
logger.error(f"Error in on_disconnect: {e}") logger.error(f"Error in on_disconnect: {e}")
async def send_message( async def send_message(self, text: str) -> AsyncIterator[AgentEventUnion]:
self, text: str, attachments: list[str] | None = None
) -> AsyncIterator[AgentEventUnion]:
""" """
Нативный асинхронный генератор. Нативный асинхронный генератор.
Не требует отдельного класса ResponseIterator. Не требует отдельного класса ResponseIterator.
Гарантированно освобождает блокировку. Гарантированно освобождает блокировку.
Args:
text: Текст сообщения.
attachments: Список путей к файлам относительно /workspace.
""" """
if not self._connected or not self._ws: if not self._connected or not self._ws:
raise AgentException( raise AgentException(code="NOT_CONNECTED",
code="NOT_CONNECTED", details="Not connected. Call connect() first." details="Not connected. Call connect() first.")
)
if self._request_lock.locked(): if self._request_lock.locked():
raise AgentBusyException("Agent is currently processing another request") raise AgentBusyException(
"Agent is currently processing another request")
# Блокируем параллельные запросы # Блокируем параллельные запросы
# если идет стриминг ответа, то при попытки отправить новое сообщение будет ошибка - ее рейзим(делаем AgentBusyError) # если идет стриминг ответа, то при попытки отправить новое сообщение будет ошибка - ее рейзим(делаем AgentBusyError)
@ -193,8 +154,7 @@ class AgentApi:
message = MsgUserMessage( message = MsgUserMessage(
type=EClientMessage.USER_MESSAGE, type=EClientMessage.USER_MESSAGE,
text=text, text=text
attachments=attachments or [],
) )
await self._ws.send_str(message.model_dump_json()) await self._ws.send_str(message.model_dump_json())
@ -238,8 +198,7 @@ class AgentApi:
# Если это исключение, просто логируем, в коллбек его кидать не стоит # Если это исключение, просто логируем, в коллбек его кидать не стоит
if isinstance(orphan_msg, Exception): if isinstance(orphan_msg, Exception):
logger.debug( logger.debug(
f"[{self.id}] Dropped exception from queue during cleanup: {orphan_msg}" f"[{self.id}] Dropped exception from queue during cleanup: {orphan_msg}")
)
continue continue
# 3. Отправляем "мусорные/осиротевшие" куски в callback # 3. Отправляем "мусорные/осиротевшие" куски в callback
@ -247,8 +206,7 @@ class AgentApi:
self.callback(orphan_msg) self.callback(orphan_msg)
else: else:
logger.debug( logger.debug(
f"[{self.id}] Dropped orphaned message during cleanup" f"[{self.id}] Dropped orphaned message during cleanup")
)
except asyncio.QueueEmpty: except asyncio.QueueEmpty:
break break
@ -258,41 +216,34 @@ class AgentApi:
self._request_lock.release() self._request_lock.release()
async def _listen(self): async def _listen(self):
""" """"
Прослушивание вебсокета. Прослушивание вебсокета.
""" """
try: try:
async for msg in self._ws: async for msg in self._ws:
if msg.type == aiohttp.WSMsgType.TEXT: if msg.type == aiohttp.WSMsgType.TEXT:
try: try:
outgoing_msg = ServerMessage.validate_json(msg.data) outgoing_msg = ServerMessage.validate_json(
msg.data)
if isinstance( if isinstance(outgoing_msg, (MsgEventTextChunk,
outgoing_msg, MsgEventToolCallChunk,
( MsgEventToolResult,
MsgEventTextChunk, MsgEventCustomUpdate,
MsgEventToolCallChunk, MsgEventEnd)):
MsgEventToolResult,
MsgEventCustomUpdate,
MsgEventSendFile,
MsgEventEnd,
),
):
if self._current_queue: if self._current_queue:
await self._current_queue.put(outgoing_msg) await self._current_queue.put(outgoing_msg)
elif self.callback: elif self.callback:
self.callback(outgoing_msg) self.callback(outgoing_msg)
else: else:
logger.warning( logger.warning(
f"[{self.id}] AgentEvent without active request" f"[{self.id}] AgentEvent without active request")
)
elif isinstance(outgoing_msg, MsgError): elif isinstance(outgoing_msg, MsgError):
if self.callback: if self.callback:
self.callback(outgoing_msg) self.callback(outgoing_msg)
error = AgentException( error = AgentException(
outgoing_msg.code, outgoing_msg.details outgoing_msg.code, outgoing_msg.details)
)
logger.error(f"[{self.id}] Agent error: {error}") logger.error(f"[{self.id}] Agent error: {error}")
if self._current_queue: if self._current_queue:
await self._current_queue.put(error) await self._current_queue.put(error)
@ -300,7 +251,8 @@ class AgentApi:
elif isinstance(outgoing_msg, MsgGracefulDisconnect): elif isinstance(outgoing_msg, MsgGracefulDisconnect):
if self.callback: if self.callback:
self.callback(outgoing_msg) self.callback(outgoing_msg)
logger.info(f"[{self.id}] Gracefully disconnecting") logger.info(
f"[{self.id}] Gracefully disconnecting")
break # Выход из цикла приведет к finally -> _cleanup break # Выход из цикла приведет к finally -> _cleanup
else: else:
@ -312,14 +264,17 @@ class AgentApi:
self.callback(outgoing_msg) self.callback(outgoing_msg)
except Exception as e: except Exception as e:
logger.error(f"[{self.id}] Failed to deserialize message: {e}") logger.error(
f"[{self.id}] Failed to deserialize message: {e}")
if self._current_queue: if self._current_queue:
await self._current_queue.put( await self._current_queue.put(
AgentException("PARSE_ERROR", f"Validation failed: {e}") AgentException(
"PARSE_ERROR", f"Validation failed: {e}")
) )
elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED): elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
logger.error(f"[{self.id}] WebSocket closed/error: {msg.type}") logger.error(
f"[{self.id}] WebSocket closed/error: {msg.type}")
break break
except asyncio.CancelledError: except asyncio.CancelledError:

6
lambda_agent_api/client.py
View file

@ -19,12 +19,6 @@ class MsgUserMessage(BaseModel):
""" """
Текст сообщения. Текст сообщения.
""" """
attachments: list[str] = Field(default_factory=list)
"""
Список вложений (файлов) к сообщению.
Передается путь до файла относительно /workspace .
Файлы уже должны быть загружены в директорию.
"""
ClientMessage = TypeAdapter(Annotated[ ClientMessage = TypeAdapter(Annotated[

99
lambda_agent_api/server.py
View file

@ -4,18 +4,10 @@ from typing import Literal, Annotated, Union, Any, Dict, Optional
__all__ = [ __all__ = [
"EServerMessage", 'EServerMessage', 'MsgStatus', 'MsgError', 'MsgGracefulDisconnect',
"MsgStatus", 'MsgEventTextChunk', 'MsgEventToolCallChunk', 'MsgEventToolResult',
"MsgError", 'MsgEventCustomUpdate', 'MsgEventEnd',
"MsgGracefulDisconnect", 'AgentEventUnion', 'ServerMessage'
"MsgEventTextChunk",
"MsgEventToolCallChunk",
"MsgEventToolResult",
"MsgEventCustomUpdate",
"MsgEventSendFile",
"MsgEventEnd",
"AgentEventUnion",
"ServerMessage",
] ]
@ -27,21 +19,18 @@ class EServerMessage(str, Enum):
# Ивенты агента # Ивенты агента
AGENT_EVENT_TEXT_CHUNK = "AGENT_EVENT_TEXT_CHUNK" AGENT_EVENT_TEXT_CHUNK = "AGENT_EVENT_TEXT_CHUNK"
AGENT_EVENT_TOOL_CALL_CHUNK = "AGENT_EVENT_TOOL_CALL_CHUNK" # Новое AGENT_EVENT_TOOL_CALL_CHUNK = "AGENT_EVENT_TOOL_CALL_CHUNK" # Новое
AGENT_EVENT_TOOL_RESULT = "AGENT_EVENT_TOOL_RESULT" # Новоеы AGENT_EVENT_TOOL_RESULT = "AGENT_EVENT_TOOL_RESULT" # Новоеы
AGENT_EVENT_CUSTOM_UPDATE = "AGENT_EVENT_CUSTOM_UPDATE" # Новое AGENT_EVENT_CUSTOM_UPDATE = "AGENT_EVENT_CUSTOM_UPDATE" # Новое
AGENT_EVENT_SEND_FILE = "AGENT_EVENT_SEND_FILE" # Новое
AGENT_EVENT_END = "AGENT_EVENT_END" AGENT_EVENT_END = "AGENT_EVENT_END"
class MsgStatus(BaseModel): class MsgStatus(BaseModel):
"""Отправляется сервером при открытии соединения с клиентом.""" """Отправляется сервером при открытии соединения с клиентом."""
type: Literal[EServerMessage.STATUS] = EServerMessage.STATUS type: Literal[EServerMessage.STATUS] = EServerMessage.STATUS
class MsgError(BaseModel): class MsgError(BaseModel):
"""Неопределенная ошибка в работе агента.""" """Неопределенная ошибка в работе агента."""
type: Literal[EServerMessage.ERROR] = EServerMessage.ERROR type: Literal[EServerMessage.ERROR] = EServerMessage.ERROR
code: str code: str
details: str details: str
@ -49,23 +38,16 @@ class MsgError(BaseModel):
class MsgGracefulDisconnect(BaseModel): class MsgGracefulDisconnect(BaseModel):
"""Отправляется перед завершением работы контейнера с агентом.""" """Отправляется перед завершением работы контейнера с агентом."""
type: Literal[EServerMessage.GRACEFUL_DISCONNECT] = EServerMessage.GRACEFUL_DISCONNECT
type: Literal[EServerMessage.GRACEFUL_DISCONNECT] = (
EServerMessage.GRACEFUL_DISCONNECT
)
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# AGENT EVENTS (События генерации) # AGENT EVENTS (События генерации)
# ------------------------------------------------------------------ # ------------------------------------------------------------------
class MsgEventTextChunk(BaseModel): class MsgEventTextChunk(BaseModel):
"""Чанк текста ответа агента.""" """Чанк текста ответа агента."""
type: Literal[EServerMessage.AGENT_EVENT_TEXT_CHUNK] = EServerMessage.AGENT_EVENT_TEXT_CHUNK
type: Literal[EServerMessage.AGENT_EVENT_TEXT_CHUNK] = (
EServerMessage.AGENT_EVENT_TEXT_CHUNK
)
text: str text: str
# Новое: "main" (главный агент) или "tools:..." (субагент, если будем использовать) # Новое: "main" (главный агент) или "tools:..." (субагент, если будем использовать)
source: str = "main" # пока везде будет main source: str = "main" # пока везде будет main
@ -73,23 +55,17 @@ class MsgEventTextChunk(BaseModel):
class MsgEventToolCallChunk(BaseModel): class MsgEventToolCallChunk(BaseModel):
"""Агент решил использовать инструмент и генерирует аргументы.""" """Агент решил использовать инструмент и генерирует аргументы."""
type: Literal[EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK] = EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK
type: Literal[EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK] = (
EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK
)
tool_name: Optional[str] = Field( tool_name: Optional[str] = Field(
None, description="Имя инструмента (приходит обычно в первом чанке)" None, description="Имя инструмента (приходит обычно в первом чанке)")
) args_chunk: Optional[str] = Field(
args_chunk: Optional[str] = Field(None, description="Кусок JSON-аргументов") None, description="Кусок JSON-аргументов")
source: str = "main" source: str = "main"
class MsgEventToolResult(BaseModel): class MsgEventToolResult(BaseModel):
"""Инструмент отработал и вернул результат.""" """Инструмент отработал и вернул результат."""
type: Literal[EServerMessage.AGENT_EVENT_TOOL_RESULT] = EServerMessage.AGENT_EVENT_TOOL_RESULT
type: Literal[EServerMessage.AGENT_EVENT_TOOL_RESULT] = (
EServerMessage.AGENT_EVENT_TOOL_RESULT
)
tool_name: str tool_name: str
result: Any # Может быть строкой, словарем или списком result: Any # Может быть строкой, словарем или списком
source: str = "main" source: str = "main"
@ -97,28 +73,14 @@ class MsgEventToolResult(BaseModel):
class MsgEventCustomUpdate(BaseModel): class MsgEventCustomUpdate(BaseModel):
"""Кастомный прогресс (например, скачивание файла) изнутри инструмента.""" """Кастомный прогресс (например, скачивание файла) изнутри инструмента."""
type: Literal[EServerMessage.AGENT_EVENT_CUSTOM_UPDATE] = EServerMessage.AGENT_EVENT_CUSTOM_UPDATE
type: Literal[EServerMessage.AGENT_EVENT_CUSTOM_UPDATE] = (
EServerMessage.AGENT_EVENT_CUSTOM_UPDATE
)
payload: Dict[str, Any] = Field( payload: Dict[str, Any] = Field(
..., description="Любые данные о прогрессе (status, progress и т.д.)" ..., description="Любые данные о прогрессе (status, progress и т.д.)")
)
source: str = "main" source: str = "main"
class MsgEventSendFile(BaseModel):
"""Агент отправляет файл пользователю."""
type: Literal[EServerMessage.AGENT_EVENT_SEND_FILE] = (
EServerMessage.AGENT_EVENT_SEND_FILE
)
path: str = Field(..., description="Путь к файлу относительно /workspace")
class MsgEventEnd(BaseModel): class MsgEventEnd(BaseModel):
"""Агент закончил генерацию ответа.""" """Агент закончил генерацию ответа."""
type: Literal[EServerMessage.AGENT_EVENT_END] = EServerMessage.AGENT_EVENT_END type: Literal[EServerMessage.AGENT_EVENT_END] = EServerMessage.AGENT_EVENT_END
tokens_used: int tokens_used: int
@ -133,24 +95,25 @@ AgentEventUnion = Union[
MsgEventToolCallChunk, MsgEventToolCallChunk,
MsgEventToolResult, MsgEventToolResult,
MsgEventCustomUpdate, MsgEventCustomUpdate,
MsgEventSendFile, MsgEventEnd
MsgEventEnd,
] ]
# ServerMessage использует AgentEventUnion + остальные типы # Обновлено: добавили новые модели в Union адаптера
ServerMessage = TypeAdapter( ServerMessage = TypeAdapter(Annotated[
Annotated[ Union[
Union[ MsgStatus,
MsgStatus, MsgError,
MsgError, MsgGracefulDisconnect,
MsgGracefulDisconnect, MsgEventTextChunk,
AgentEventUnion, MsgEventToolCallChunk,
], MsgEventToolResult,
Field(discriminator="type"), MsgEventCustomUpdate,
] MsgEventEnd
) ],
Field(discriminator="type")
])
""" """
TypeAdapter для десериализации всех входящих сообщений (от сервера). Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.
Pydantic сам определит нужный тип в зависимости от поля `type`. Pydantic сам определит нужный тип в зависимости от поля `type`.
Использование: Использование:
msg = ServerMessage.validate_json(json_str) msg = ServerMessage.validate_json(json_str)

39
tests/manual.py
View file

@ -1,58 +1,33 @@
import asyncio import asyncio
import traceback import traceback
from lambda_agent_api.agent_api import AgentApi, AgentBusyException from lambda_agent_api.agent_api import AgentApi
from lambda_agent_api.server import ( from lambda_agent_api.server import MsgEventTextChunk, MsgEventToolCallChunk, MsgEventToolResult
MsgEventTextChunk,
MsgEventToolCallChunk,
MsgEventToolResult,
MsgEventSendFile,
)
async def main(): async def main():
chat_id = input("Chat id: ") or 0 api = AgentApi("agent-1", "ws://localhost:8000/agent_ws/")
api = AgentApi("agent-1", "ws://localhost:8000/", chat_id=chat_id)
try:
await api.connect()
except AgentBusyException:
print(f"Чат {chat_id} занят другим клиентом")
return
await api.connect()
while True: while True:
try: try:
prompt = await asyncio.get_event_loop().run_in_executor(None, input, ">>> ") prompt = await asyncio.get_event_loop().run_in_executor(None, input, ">>> ")
attachments_input = await asyncio.get_event_loop().run_in_executor(
None, input, "Attachments (comma-separated, empty for none): "
)
attachments = (
[a.strip() for a in attachments_input.split(",") if a.strip()]
if attachments_input.strip()
else None
)
print("Agent: ", end="") print("Agent: ", end="")
is_tool = False is_tool = False
async for chunk in api.send_message(prompt, attachments): async for chunk in api.send_message(prompt):
match chunk: match chunk:
case MsgEventTextChunk(): case MsgEventTextChunk():
is_tool = False is_tool = False
print(chunk.text, end="", flush=True) print(chunk.text, end="", flush=True)
case MsgEventToolCallChunk(): case MsgEventToolCallChunk():
if not is_tool: if not is_tool:
print( print(f"\n\n### TOOL CALL: ({chunk.tool_name}) ", end="", flush=True)
f"\n\n### TOOL CALL: ({chunk.tool_name}) ",
end="",
flush=True,
)
is_tool = True is_tool = True
print(chunk.args_chunk, end="", flush=True) print(chunk.args_chunk, end="", flush=True)
case MsgEventToolResult(): case MsgEventToolResult():
is_tool = False is_tool = False
print(f"\nResult: {chunk.result}\n\n", end="", flush=True) print(f"\nResult: {chunk.result}\n\n", end="", flush=True)
case MsgEventSendFile():
print(f"\n[SEND FILE] {chunk.path}\n", end="", flush=True)
print("\n") print("\n")
except KeyboardInterrupt: except KeyboardInterrupt: