Compare commits
No commits in common. "483a63b99962ec6668d9ceb125eaf35deaf66cf4" and "aac37db672ed60d614f90e746514e3e42903c214" have entirely different histories.
483a63b999
...
aac37db672
4 changed files with 73 additions and 144 deletions
19
README.md
19
README.md
|
|
@ -5,10 +5,9 @@ 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/")`
|
||||
Раньше: `AgentAPI(url="ws://localhost:8000/agent_ws/")`. Сейчас: `AgentAPI(base_url="ws://localhost:8000/")`
|
||||
- Добавлен параметр `chat_id` в конструктор `AgentAPI`. Нужен для разделения истории сообщений по чатам/веткам.
|
||||
- `AgentAPI.connect()` вызывает `AgentBusyException`, если выбранный чат уже занят другим API клиентом.
|
||||
- Добавлен новый тип события `MsgEventSendFile` для отправки файлов пользователю. Поле `path` — путь к файлу относительно `/workspace`.
|
||||
|
||||
## Установка
|
||||
В `master` всегда будет актуальная рабочая версия.
|
||||
|
|
@ -165,22 +164,6 @@ async def on_telegram_message(from_user: int, text: str):
|
|||
| payload | object | Любые данные о прогрессе |
|
||||
| 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
|
||||
|
||||
Агент закончил генерацию ответа.
|
||||
|
|
|
|||
|
|
@ -30,8 +30,8 @@ class AgentApi:
|
|||
agent_id: str,
|
||||
base_url: str,
|
||||
callback: Optional[Callable[[ServerMessage], None]] = None,
|
||||
on_disconnect: Optional[Callable[["AgentApi"], None]] = None,
|
||||
chat_id: int = 0, # значение по умолчанию для обратной совместимости # значение по умолчанию для обратной совместимости
|
||||
on_disconnect: Optional[Callable[['AgentApi'], None]] = None,
|
||||
chat_id: int = 0 # значение по умолчанию для обратной совместимости
|
||||
):
|
||||
self.id = agent_id # ID агента для словаря
|
||||
self.chat_id = chat_id
|
||||
|
|
@ -67,12 +67,10 @@ class AgentApi:
|
|||
logger.info(f"Agent {self.id} is ready")
|
||||
else:
|
||||
raise AgentException(
|
||||
"INVALID_STATUS", f"Expected SM.Status, got {status_msg.type}"
|
||||
)
|
||||
"INVALID_STATUS", f"Expected SM.Status, got {status_msg.type}")
|
||||
else:
|
||||
raise AgentException(
|
||||
"UNEXPECTED_MSG_TYPE", f"Unexpected message type: {msg.type}"
|
||||
)
|
||||
raise AgentException("UNEXPECTED_MSG_TYPE",
|
||||
f"Unexpected message type: {msg.type}")
|
||||
|
||||
self._connected = True
|
||||
self._listen_task = asyncio.create_task(self._listen())
|
||||
|
|
@ -85,8 +83,7 @@ class AgentApi:
|
|||
await self._session.close()
|
||||
|
||||
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:
|
||||
# Перехватываем наши собственные ошибки (INVALID_STATUS, UNEXPECTED_MSG_TYPE),
|
||||
|
|
@ -124,10 +121,8 @@ class AgentApi:
|
|||
await self._session.close()
|
||||
|
||||
# Можно оставить RuntimeError, а можно тоже завернуть в AgentException
|
||||
raise AgentException(
|
||||
code="CONNECTION_ERROR",
|
||||
details=f"Failed to connect agent {self.id}: {e}",
|
||||
) from e
|
||||
raise AgentException(code="CONNECTION_ERROR",
|
||||
details=f"Failed to connect agent {self.id}: {e}") from e
|
||||
|
||||
async def close(self):
|
||||
"""Явное ручное закрытие соединения."""
|
||||
|
|
@ -171,12 +166,12 @@ class AgentApi:
|
|||
Гарантированно освобождает блокировку.
|
||||
"""
|
||||
if not self._connected or not self._ws:
|
||||
raise AgentException(
|
||||
code="NOT_CONNECTED", details="Not connected. Call connect() first."
|
||||
)
|
||||
raise AgentException(code="NOT_CONNECTED",
|
||||
details="Not connected. Call connect() first.")
|
||||
|
||||
if self._request_lock.locked():
|
||||
raise AgentBusyException("Agent is currently processing another request")
|
||||
raise AgentBusyException(
|
||||
"Agent is currently processing another request")
|
||||
|
||||
# Блокируем параллельные запросы
|
||||
# если идет стриминг ответа, то при попытки отправить новое сообщение будет ошибка - ее рейзим(делаем AgentBusyError)
|
||||
|
|
@ -185,7 +180,10 @@ class AgentApi:
|
|||
try:
|
||||
self._current_queue = asyncio.Queue()
|
||||
|
||||
message = MsgUserMessage(type=EClientMessage.USER_MESSAGE, text=text)
|
||||
message = MsgUserMessage(
|
||||
type=EClientMessage.USER_MESSAGE,
|
||||
text=text
|
||||
)
|
||||
|
||||
await self._ws.send_str(message.model_dump_json())
|
||||
logger.debug(f"[{self.id}] Sent message: {text[:50]}...")
|
||||
|
|
@ -228,8 +226,7 @@ class AgentApi:
|
|||
# Если это исключение, просто логируем, в коллбек его кидать не стоит
|
||||
if isinstance(orphan_msg, Exception):
|
||||
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
|
||||
|
||||
# 3. Отправляем "мусорные/осиротевшие" куски в callback
|
||||
|
|
@ -237,8 +234,7 @@ class AgentApi:
|
|||
self.callback(orphan_msg)
|
||||
else:
|
||||
logger.debug(
|
||||
f"[{self.id}] Dropped orphaned message during cleanup"
|
||||
)
|
||||
f"[{self.id}] Dropped orphaned message during cleanup")
|
||||
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
|
@ -255,34 +251,27 @@ class AgentApi:
|
|||
async for msg in self._ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
outgoing_msg = ServerMessage.validate_json(msg.data)
|
||||
outgoing_msg = ServerMessage.validate_json(
|
||||
msg.data)
|
||||
|
||||
if isinstance(
|
||||
outgoing_msg,
|
||||
(
|
||||
MsgEventTextChunk,
|
||||
MsgEventToolCallChunk,
|
||||
MsgEventToolResult,
|
||||
MsgEventCustomUpdate,
|
||||
MsgEventSendFile,
|
||||
MsgEventEnd,
|
||||
),
|
||||
):
|
||||
if isinstance(outgoing_msg, (MsgEventTextChunk,
|
||||
MsgEventToolCallChunk,
|
||||
MsgEventToolResult,
|
||||
MsgEventCustomUpdate,
|
||||
MsgEventEnd)):
|
||||
if self._current_queue:
|
||||
await self._current_queue.put(outgoing_msg)
|
||||
elif self.callback:
|
||||
self.callback(outgoing_msg)
|
||||
else:
|
||||
logger.warning(
|
||||
f"[{self.id}] AgentEvent without active request"
|
||||
)
|
||||
f"[{self.id}] AgentEvent without active request")
|
||||
|
||||
elif isinstance(outgoing_msg, MsgError):
|
||||
if self.callback:
|
||||
self.callback(outgoing_msg)
|
||||
error = AgentException(
|
||||
outgoing_msg.code, outgoing_msg.details
|
||||
)
|
||||
outgoing_msg.code, outgoing_msg.details)
|
||||
logger.error(f"[{self.id}] Agent error: {error}")
|
||||
if self._current_queue:
|
||||
await self._current_queue.put(error)
|
||||
|
|
@ -290,7 +279,8 @@ class AgentApi:
|
|||
elif isinstance(outgoing_msg, MsgGracefulDisconnect):
|
||||
if self.callback:
|
||||
self.callback(outgoing_msg)
|
||||
logger.info(f"[{self.id}] Gracefully disconnecting")
|
||||
logger.info(
|
||||
f"[{self.id}] Gracefully disconnecting")
|
||||
break # Выход из цикла приведет к finally -> _cleanup
|
||||
|
||||
else:
|
||||
|
|
@ -302,14 +292,17 @@ class AgentApi:
|
|||
self.callback(outgoing_msg)
|
||||
|
||||
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:
|
||||
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):
|
||||
logger.error(f"[{self.id}] WebSocket closed/error: {msg.type}")
|
||||
logger.error(
|
||||
f"[{self.id}] WebSocket closed/error: {msg.type}")
|
||||
break
|
||||
|
||||
except asyncio.CancelledError:
|
||||
|
|
|
|||
|
|
@ -4,18 +4,10 @@ from typing import Literal, Annotated, Union, Any, Dict, Optional
|
|||
|
||||
|
||||
__all__ = [
|
||||
"EServerMessage",
|
||||
"MsgStatus",
|
||||
"MsgError",
|
||||
"MsgGracefulDisconnect",
|
||||
"MsgEventTextChunk",
|
||||
"MsgEventToolCallChunk",
|
||||
"MsgEventToolResult",
|
||||
"MsgEventCustomUpdate",
|
||||
"MsgEventSendFile",
|
||||
"MsgEventEnd",
|
||||
"AgentEventUnion",
|
||||
"ServerMessage",
|
||||
'EServerMessage', 'MsgStatus', 'MsgError', 'MsgGracefulDisconnect',
|
||||
'MsgEventTextChunk', 'MsgEventToolCallChunk', 'MsgEventToolResult',
|
||||
'MsgEventCustomUpdate', 'MsgEventEnd',
|
||||
'AgentEventUnion', 'ServerMessage'
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -27,21 +19,18 @@ class EServerMessage(str, Enum):
|
|||
# Ивенты агента
|
||||
AGENT_EVENT_TEXT_CHUNK = "AGENT_EVENT_TEXT_CHUNK"
|
||||
AGENT_EVENT_TOOL_CALL_CHUNK = "AGENT_EVENT_TOOL_CALL_CHUNK" # Новое
|
||||
AGENT_EVENT_TOOL_RESULT = "AGENT_EVENT_TOOL_RESULT" # Новоеы
|
||||
AGENT_EVENT_CUSTOM_UPDATE = "AGENT_EVENT_CUSTOM_UPDATE" # Новое
|
||||
AGENT_EVENT_SEND_FILE = "AGENT_EVENT_SEND_FILE" # Новое
|
||||
AGENT_EVENT_TOOL_RESULT = "AGENT_EVENT_TOOL_RESULT" # Новоеы
|
||||
AGENT_EVENT_CUSTOM_UPDATE = "AGENT_EVENT_CUSTOM_UPDATE" # Новое
|
||||
AGENT_EVENT_END = "AGENT_EVENT_END"
|
||||
|
||||
|
||||
class MsgStatus(BaseModel):
|
||||
"""Отправляется сервером при открытии соединения с клиентом."""
|
||||
|
||||
type: Literal[EServerMessage.STATUS] = EServerMessage.STATUS
|
||||
|
||||
|
||||
class MsgError(BaseModel):
|
||||
"""Неопределенная ошибка в работе агента."""
|
||||
|
||||
type: Literal[EServerMessage.ERROR] = EServerMessage.ERROR
|
||||
code: str
|
||||
details: str
|
||||
|
|
@ -49,23 +38,16 @@ class MsgError(BaseModel):
|
|||
|
||||
class MsgGracefulDisconnect(BaseModel):
|
||||
"""Отправляется перед завершением работы контейнера с агентом."""
|
||||
|
||||
type: Literal[EServerMessage.GRACEFUL_DISCONNECT] = (
|
||||
EServerMessage.GRACEFUL_DISCONNECT
|
||||
)
|
||||
type: Literal[EServerMessage.GRACEFUL_DISCONNECT] = EServerMessage.GRACEFUL_DISCONNECT
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# AGENT EVENTS (События генерации)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
|
||||
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
|
||||
# Новое: "main" (главный агент) или "tools:..." (субагент, если будем использовать)
|
||||
source: str = "main" # пока везде будет main
|
||||
|
|
@ -73,23 +55,17 @@ class MsgEventTextChunk(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(
|
||||
None, description="Имя инструмента (приходит обычно в первом чанке)"
|
||||
)
|
||||
args_chunk: Optional[str] = Field(None, description="Кусок JSON-аргументов")
|
||||
None, description="Имя инструмента (приходит обычно в первом чанке)")
|
||||
args_chunk: Optional[str] = Field(
|
||||
None, description="Кусок JSON-аргументов")
|
||||
source: str = "main"
|
||||
|
||||
|
||||
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
|
||||
result: Any # Может быть строкой, словарем или списком
|
||||
source: str = "main"
|
||||
|
|
@ -97,28 +73,14 @@ class MsgEventToolResult(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(
|
||||
..., description="Любые данные о прогрессе (status, progress и т.д.)"
|
||||
)
|
||||
..., description="Любые данные о прогрессе (status, progress и т.д.)")
|
||||
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):
|
||||
"""Агент закончил генерацию ответа."""
|
||||
|
||||
type: Literal[EServerMessage.AGENT_EVENT_END] = EServerMessage.AGENT_EVENT_END
|
||||
tokens_used: int
|
||||
|
||||
|
|
@ -133,24 +95,25 @@ AgentEventUnion = Union[
|
|||
MsgEventToolCallChunk,
|
||||
MsgEventToolResult,
|
||||
MsgEventCustomUpdate,
|
||||
MsgEventSendFile,
|
||||
MsgEventEnd,
|
||||
MsgEventEnd
|
||||
]
|
||||
|
||||
# ServerMessage использует AgentEventUnion + остальные типы
|
||||
ServerMessage = TypeAdapter(
|
||||
Annotated[
|
||||
Union[
|
||||
MsgStatus,
|
||||
MsgError,
|
||||
MsgGracefulDisconnect,
|
||||
AgentEventUnion,
|
||||
],
|
||||
Field(discriminator="type"),
|
||||
]
|
||||
)
|
||||
# Обновлено: добавили новые модели в Union адаптера
|
||||
ServerMessage = TypeAdapter(Annotated[
|
||||
Union[
|
||||
MsgStatus,
|
||||
MsgError,
|
||||
MsgGracefulDisconnect,
|
||||
MsgEventTextChunk,
|
||||
MsgEventToolCallChunk,
|
||||
MsgEventToolResult,
|
||||
MsgEventCustomUpdate,
|
||||
MsgEventEnd
|
||||
],
|
||||
Field(discriminator="type")
|
||||
])
|
||||
"""
|
||||
TypeAdapter для десериализации всех входящих сообщений (от сервера).
|
||||
Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.
|
||||
Pydantic сам определит нужный тип в зависимости от поля `type`.
|
||||
Использование:
|
||||
msg = ServerMessage.validate_json(json_str)
|
||||
|
|
|
|||
|
|
@ -2,12 +2,8 @@ import asyncio
|
|||
import traceback
|
||||
|
||||
from lambda_agent_api.agent_api import AgentApi, AgentBusyException
|
||||
from lambda_agent_api.server import (
|
||||
MsgEventTextChunk,
|
||||
MsgEventToolCallChunk,
|
||||
MsgEventToolResult,
|
||||
MsgEventSendFile,
|
||||
)
|
||||
from lambda_agent_api.server import MsgEventTextChunk, MsgEventToolCallChunk, MsgEventToolResult
|
||||
|
||||
|
||||
|
||||
async def main():
|
||||
|
|
@ -32,18 +28,12 @@ async def main():
|
|||
print(chunk.text, end="", flush=True)
|
||||
case MsgEventToolCallChunk():
|
||||
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
|
||||
print(chunk.args_chunk, end="", flush=True)
|
||||
case MsgEventToolResult():
|
||||
is_tool = False
|
||||
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")
|
||||
except KeyboardInterrupt:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue