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

#5 Переименованы IM->CM(IncomigMessage->ClientMessage),OM->SM(OutgoingMessage->ServerMessage)

Browse source
This commit is contained in:
Ярослав Малинин 2026-03-31 11:02:45 +03:00
parent 8fd5c462ed
commit 42c6571d3a
3 changed files with 31 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

28
api/agent_api.py
View file

@ -2,7 +2,7 @@ import logging
from typing import Callable, Optional
import aiohttp
import asyncio
from models import IM, OM, ClientMessage, ServerMessage
from models import CM, SM, ClientMessage, ServerMessage
logger = logging.getLogger(__name__)
@ -13,7 +13,7 @@ class AgentException(Exception):
Кастомное исключение для ошибок, полученных от агента.
Атрибуты:
code: Код ошибки из сообщения OM.Error
code: Код ошибки из сообщения SM.Error
details: Детали ошибки
"""
@ -28,14 +28,14 @@ class AgentApi:
Асинхронный клиент для взаимодействия с AI-агентом через WebSocket.
Класс инкапсулирует обмен сообщениями согласно контракту Pydantic
(IM для входящих, OM для исходящих сообщений).
(CM для входящих, SM для исходящих сообщений).
Пример использования:
async with AgentApi("ws://localhost:8000", callback=my_callback) as agent:
response = await agent.send_message("Hello, agent!")
async for chunk in response:
match chunk:
case OM.EventTextChunk():
case SM.EventTextChunk():
print(chunk.text, end="")
print(f" [{response.tokens} токенов]")
@ -81,17 +81,17 @@ class AgentApi:
self._connected = True
logger.info(f"Connected to agent at {self.url}")
# Ожидаем OM.Status при открытии соединения
# Ожидаем SM.Status при открытии соединения
msg = await self._ws.receive()
if msg.type == aiohttp.WSMsgType.TEXT:
status_msg = ServerMessage.model_validate_json(msg.data)
if isinstance(status_msg, OM.Status):
if isinstance(status_msg, SM.Status):
if self.callback:
self.callback(status_msg)
logger.info("Agent is ready to accept messages")
else:
raise RuntimeError(
f"Expected OM.Status on connection, got {status_msg.type}")
f"Expected SM.Status on connection, got {status_msg.type}")
else:
raise RuntimeError(
f"Unexpected message type on connection: {msg.type}")
@ -154,8 +154,8 @@ class AgentApi:
raise RuntimeError(
"Not connected to agent. Use 'async with' context manager.")
message = IM.UserMessage(
type=IM.Type.USER_MESSAGE,
message = CM.UserMessage(
type=CM.Type.USER_MESSAGE,
text=text
)
@ -181,20 +181,20 @@ class AgentApi:
logger.debug(
f"Received message of type: {outgoing_msg.type}")
if isinstance(outgoing_msg, OM.AgentEvent):
if isinstance(outgoing_msg, SM.AgentEvent):
await self._queue.put(outgoing_msg)
elif isinstance(outgoing_msg, OM.Status):
elif isinstance(outgoing_msg, SM.Status):
if self.callback:
self.callback(outgoing_msg)
logger.info("Agent status update")
elif isinstance(outgoing_msg, OM.Error):
elif isinstance(outgoing_msg, SM.Error):
if self.callback:
self.callback(outgoing_msg)
error = AgentException(
outgoing_msg.code, outgoing_msg.details)
logger.error(f"Agent error: {error}")
# Не бросаем исключение, а вызываем callback
elif isinstance(outgoing_msg, OM.GracefulDisconnect):
elif isinstance(outgoing_msg, SM.GracefulDisconnect):
if self.callback:
self.callback(outgoing_msg)
logger.info("Agent gracefully disconnecting")
@ -242,7 +242,7 @@ class ResponseIterator:
async def __anext__(self):
try:
chunk = await self._queue.get()
if isinstance(chunk, OM.EventEnd):
if isinstance(chunk, SM.EventEnd):
self.tokens = chunk.tokens_used
raise StopAsyncIteration
return chunk