Compare commits
18 commits
#9-clients
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
| 8a4f4db6d3 | |||
| aa480bbec5 | |||
| 972fc23bfc | |||
| bd2923ffca | |||
| 483a63b999 | |||
| 700f7fa0e7 | |||
| 0cf3117fde | |||
| 4c3e7253c7 | |||
| 8256453bbf | |||
| 1b8efdb4a4 | |||
| 9ac4ef1ba1 | |||
| 146dcdf21d | |||
| aac37db672 | |||
| 0c9906ecb4 | |||
| 234050df9f | |||
| ee98eb09d9 | |||
| 649834beae | |||
| bb20a84741 |
5 changed files with 216 additions and 112 deletions
63
README.md
63
README.md
|
|
@ -1,9 +1,16 @@
|
||||||
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
|
||||||
|
|
@ -14,42 +21,10 @@ pip install git+https://git.lambda.coredump.ru/platform/agent_api.git
|
||||||
|
|
||||||
## Быстрый старт (с использованием AgentApi)
|
## Быстрый старт (с использованием AgentApi)
|
||||||
|
|
||||||
```python
|
**Рабочий REPL пример: [tests/manual.py](tests/manual.py).**
|
||||||
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
|
||||||
|
|
@ -191,6 +166,22 @@ 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
|
||||||
|
|
||||||
Агент закончил генерацию ответа.
|
Агент закончил генерацию ответа.
|
||||||
|
|
|
||||||
|
|
@ -2,6 +2,9 @@ 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 *
|
||||||
|
|
@ -25,12 +28,14 @@ class AgentApi:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
agent_id: str,
|
agent_id: str,
|
||||||
url: str,
|
base_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.url = url
|
self.chat_id = chat_id
|
||||||
|
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
|
||||||
|
|
||||||
|
|
@ -43,7 +48,11 @@ 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)
|
||||||
|
|
@ -58,10 +67,12 @@ 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("UNEXPECTED_MSG_TYPE",
|
raise AgentException(
|
||||||
f"Unexpected message type: {msg.type}")
|
"UNEXPECTED_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())
|
||||||
|
|
@ -74,7 +85,8 @@ 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") from e
|
"TIMEOUT", "Agent did not send initial Status message within 5 seconds"
|
||||||
|
) from e
|
||||||
|
|
||||||
except AgentException:
|
except AgentException:
|
||||||
# Перехватываем наши собственные ошибки (INVALID_STATUS, UNEXPECTED_MSG_TYPE),
|
# Перехватываем наши собственные ошибки (INVALID_STATUS, UNEXPECTED_MSG_TYPE),
|
||||||
|
|
@ -85,6 +97,25 @@ 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:
|
||||||
|
|
@ -93,8 +124,10 @@ class AgentApi:
|
||||||
await self._session.close()
|
await self._session.close()
|
||||||
|
|
||||||
# Можно оставить RuntimeError, а можно тоже завернуть в AgentException
|
# Можно оставить RuntimeError, а можно тоже завернуть в AgentException
|
||||||
raise AgentException(code="CONNECTION_ERROR",
|
raise AgentException(
|
||||||
details=f"Failed to connect agent {self.id}: {e}") from e
|
code="CONNECTION_ERROR",
|
||||||
|
details=f"Failed to connect agent {self.id}: {e}",
|
||||||
|
) from e
|
||||||
|
|
||||||
async def close(self):
|
async def close(self):
|
||||||
"""Явное ручное закрытие соединения."""
|
"""Явное ручное закрытие соединения."""
|
||||||
|
|
@ -131,19 +164,25 @@ 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(self, text: str) -> AsyncIterator[AgentEventUnion]:
|
async def send_message(
|
||||||
|
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(code="NOT_CONNECTED",
|
raise AgentException(
|
||||||
details="Not connected. Call connect() first.")
|
code="NOT_CONNECTED", details="Not connected. Call connect() first."
|
||||||
|
)
|
||||||
|
|
||||||
if self._request_lock.locked():
|
if self._request_lock.locked():
|
||||||
raise AgentBusyException(
|
raise AgentBusyException("Agent is currently processing another request")
|
||||||
"Agent is currently processing another request")
|
|
||||||
|
|
||||||
# Блокируем параллельные запросы
|
# Блокируем параллельные запросы
|
||||||
# если идет стриминг ответа, то при попытки отправить новое сообщение будет ошибка - ее рейзим(делаем AgentBusyError)
|
# если идет стриминг ответа, то при попытки отправить новое сообщение будет ошибка - ее рейзим(делаем AgentBusyError)
|
||||||
|
|
@ -154,7 +193,8 @@ 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())
|
||||||
|
|
@ -198,7 +238,8 @@ 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
|
||||||
|
|
@ -206,7 +247,8 @@ 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
|
||||||
|
|
@ -216,34 +258,41 @@ 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(
|
outgoing_msg = ServerMessage.validate_json(msg.data)
|
||||||
msg.data)
|
|
||||||
|
|
||||||
if isinstance(outgoing_msg, (MsgEventTextChunk,
|
if isinstance(
|
||||||
MsgEventToolCallChunk,
|
outgoing_msg,
|
||||||
MsgEventToolResult,
|
(
|
||||||
MsgEventCustomUpdate,
|
MsgEventTextChunk,
|
||||||
MsgEventEnd)):
|
MsgEventToolCallChunk,
|
||||||
|
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)
|
||||||
|
|
@ -251,8 +300,7 @@ 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(
|
logger.info(f"[{self.id}] Gracefully disconnecting")
|
||||||
f"[{self.id}] Gracefully disconnecting")
|
|
||||||
break # Выход из цикла приведет к finally -> _cleanup
|
break # Выход из цикла приведет к finally -> _cleanup
|
||||||
|
|
||||||
else:
|
else:
|
||||||
|
|
@ -264,17 +312,14 @@ class AgentApi:
|
||||||
self.callback(outgoing_msg)
|
self.callback(outgoing_msg)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(
|
logger.error(f"[{self.id}] Failed to deserialize message: {e}")
|
||||||
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(
|
AgentException("PARSE_ERROR", f"Validation failed: {e}")
|
||||||
"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(
|
logger.error(f"[{self.id}] WebSocket closed/error: {msg.type}")
|
||||||
f"[{self.id}] WebSocket closed/error: {msg.type}")
|
|
||||||
break
|
break
|
||||||
|
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
|
|
|
||||||
|
|
@ -19,6 +19,12 @@ class MsgUserMessage(BaseModel):
|
||||||
"""
|
"""
|
||||||
Текст сообщения.
|
Текст сообщения.
|
||||||
"""
|
"""
|
||||||
|
attachments: list[str] = Field(default_factory=list)
|
||||||
|
"""
|
||||||
|
Список вложений (файлов) к сообщению.
|
||||||
|
Передается путь до файла относительно /workspace .
|
||||||
|
Файлы уже должны быть загружены в директорию.
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
ClientMessage = TypeAdapter(Annotated[
|
ClientMessage = TypeAdapter(Annotated[
|
||||||
|
|
|
||||||
|
|
@ -4,10 +4,18 @@ from typing import Literal, Annotated, Union, Any, Dict, Optional
|
||||||
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
'EServerMessage', 'MsgStatus', 'MsgError', 'MsgGracefulDisconnect',
|
"EServerMessage",
|
||||||
'MsgEventTextChunk', 'MsgEventToolCallChunk', 'MsgEventToolResult',
|
"MsgStatus",
|
||||||
'MsgEventCustomUpdate', 'MsgEventEnd',
|
"MsgError",
|
||||||
'AgentEventUnion', 'ServerMessage'
|
"MsgGracefulDisconnect",
|
||||||
|
"MsgEventTextChunk",
|
||||||
|
"MsgEventToolCallChunk",
|
||||||
|
"MsgEventToolResult",
|
||||||
|
"MsgEventCustomUpdate",
|
||||||
|
"MsgEventSendFile",
|
||||||
|
"MsgEventEnd",
|
||||||
|
"AgentEventUnion",
|
||||||
|
"ServerMessage",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -19,18 +27,21 @@ 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
|
||||||
|
|
@ -38,16 +49,23 @@ 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
|
||||||
|
|
@ -55,17 +73,23 @@ 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(
|
)
|
||||||
None, description="Кусок JSON-аргументов")
|
args_chunk: Optional[str] = Field(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"
|
||||||
|
|
@ -73,14 +97,28 @@ 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
|
||||||
|
|
||||||
|
|
@ -95,25 +133,24 @@ AgentEventUnion = Union[
|
||||||
MsgEventToolCallChunk,
|
MsgEventToolCallChunk,
|
||||||
MsgEventToolResult,
|
MsgEventToolResult,
|
||||||
MsgEventCustomUpdate,
|
MsgEventCustomUpdate,
|
||||||
MsgEventEnd
|
MsgEventSendFile,
|
||||||
|
MsgEventEnd,
|
||||||
]
|
]
|
||||||
|
|
||||||
# Обновлено: добавили новые модели в Union адаптера
|
# ServerMessage использует AgentEventUnion + остальные типы
|
||||||
ServerMessage = TypeAdapter(Annotated[
|
ServerMessage = TypeAdapter(
|
||||||
Union[
|
Annotated[
|
||||||
MsgStatus,
|
Union[
|
||||||
MsgError,
|
MsgStatus,
|
||||||
MsgGracefulDisconnect,
|
MsgError,
|
||||||
MsgEventTextChunk,
|
MsgGracefulDisconnect,
|
||||||
MsgEventToolCallChunk,
|
AgentEventUnion,
|
||||||
MsgEventToolResult,
|
],
|
||||||
MsgEventCustomUpdate,
|
Field(discriminator="type"),
|
||||||
MsgEventEnd
|
]
|
||||||
],
|
)
|
||||||
Field(discriminator="type")
|
|
||||||
])
|
|
||||||
"""
|
"""
|
||||||
Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.
|
TypeAdapter для десериализации всех входящих сообщений (от сервера).
|
||||||
Pydantic сам определит нужный тип в зависимости от поля `type`.
|
Pydantic сам определит нужный тип в зависимости от поля `type`.
|
||||||
Использование:
|
Использование:
|
||||||
msg = ServerMessage.validate_json(json_str)
|
msg = ServerMessage.validate_json(json_str)
|
||||||
|
|
|
||||||
|
|
@ -1,33 +1,58 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
import traceback
|
import traceback
|
||||||
|
|
||||||
from lambda_agent_api.agent_api import AgentApi
|
from lambda_agent_api.agent_api import AgentApi, AgentBusyException
|
||||||
from lambda_agent_api.server import MsgEventTextChunk, MsgEventToolCallChunk, MsgEventToolResult
|
from lambda_agent_api.server import (
|
||||||
|
MsgEventTextChunk,
|
||||||
|
MsgEventToolCallChunk,
|
||||||
|
MsgEventToolResult,
|
||||||
|
MsgEventSendFile,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
async def main():
|
async def main():
|
||||||
api = AgentApi("agent-1", "ws://localhost:8000/agent_ws/")
|
chat_id = input("Chat id: ") or 0
|
||||||
|
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):
|
async for chunk in api.send_message(prompt, attachments):
|
||||||
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(f"\n\n### TOOL CALL: ({chunk.tool_name}) ", end="", flush=True)
|
print(
|
||||||
|
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:
|
||||||
|
|
@ -38,4 +63,4 @@ async def main():
|
||||||
await api.close()
|
await api.close()
|
||||||
|
|
||||||
|
|
||||||
asyncio.run(main())
|
asyncio.run(main())
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue