#5 Переименованы IM->CM(IncomigMessage->ClientMessage),OM->SM(OutgoingMessage->ServerMessage)
This commit is contained in:
parent
8fd5c462ed
commit
42c6571d3a
3 changed files with 31 additions and 31 deletions
|
|
@ -8,13 +8,13 @@
|
|||
"""
|
||||
|
||||
from .agent_api import AgentApi, AgentException
|
||||
from .models import IM, OM, ClientMessage, ServerMessage
|
||||
from .models import CM, SM, ClientMessage, ServerMessage
|
||||
|
||||
__all__ = [
|
||||
"AgentApi",
|
||||
"AgentException",
|
||||
"IM",
|
||||
"OM",
|
||||
"CM",
|
||||
"SM",
|
||||
"ClientMessage",
|
||||
"ServerMessage",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -3,10 +3,10 @@ from enum import Enum
|
|||
from typing import Literal, Annotated, Union
|
||||
|
||||
|
||||
class IM:
|
||||
class CM:
|
||||
"""
|
||||
Namespace для моделей входящих сообщений (от клиента к серверу).\n
|
||||
IM = Incoming Message
|
||||
CM = Client Message
|
||||
"""
|
||||
|
||||
class Type(str, Enum):
|
||||
|
|
@ -18,7 +18,7 @@ class IM:
|
|||
"""
|
||||
Полное сообщение от пользователя.
|
||||
"""
|
||||
type: Literal[IM.Type.USER_MESSAGE]
|
||||
type: Literal[CM.Type.USER_MESSAGE]
|
||||
text: str
|
||||
"""
|
||||
Текст сообщения.
|
||||
|
|
@ -26,7 +26,7 @@ class IM:
|
|||
|
||||
|
||||
ClientMessage = Annotated[
|
||||
Union[IM.UserMessage,],
|
||||
Union[CM.UserMessage,],
|
||||
Field(discriminator="type")
|
||||
]
|
||||
"""
|
||||
|
|
@ -37,10 +37,10 @@ msg = ClientMessage.model_validate_json(json)
|
|||
"""
|
||||
|
||||
|
||||
class OM:
|
||||
class SM:
|
||||
"""
|
||||
Namespace для моделей исходящих сообщений (от сервера к клиенту).\n
|
||||
OM = Outgoing Message
|
||||
SM = Server Message
|
||||
"""
|
||||
|
||||
class Type(str, Enum):
|
||||
|
|
@ -54,7 +54,7 @@ class OM:
|
|||
Отправляется сервером при открытии соединения с клиентом.
|
||||
Будет дополнен информацией о готовности агента принимать сообщения.
|
||||
"""
|
||||
type: Literal[OM.Type.STATUS]
|
||||
type: Literal[SM.Type.STATUS]
|
||||
|
||||
class AgentEventType(str, Enum):
|
||||
TEXT_CHUNK = "TEXT_CHUNK"
|
||||
|
|
@ -65,28 +65,28 @@ class OM:
|
|||
Базовый класс для ивентов, которые стримит агент во время генерации ответа.
|
||||
Конкретный класс для ивента определяется по ``subtype``.
|
||||
"""
|
||||
type: Literal[OM.Type.AGENT_EVENT]
|
||||
subtype: OM.AgentEventType
|
||||
type: Literal[SM.Type.AGENT_EVENT]
|
||||
subtype: SM.AgentEventType
|
||||
|
||||
class EventTextChunk(AgentEvent):
|
||||
"""
|
||||
Чанк текста ответа агента.
|
||||
"""
|
||||
subtype: Literal[OM.AgentEventType.TEXT_CHUNK]
|
||||
subtype: Literal[SM.AgentEventType.TEXT_CHUNK]
|
||||
text: str
|
||||
|
||||
class EventEnd(AgentEvent):
|
||||
"""
|
||||
Агент закончил генерацию ответа.
|
||||
"""
|
||||
subtype: Literal[OM.AgentEventType.END]
|
||||
subtype: Literal[SM.AgentEventType.END]
|
||||
tokens_used: int
|
||||
|
||||
class Error(BaseModel):
|
||||
"""
|
||||
Неопределенная ошибка в работе агента.
|
||||
"""
|
||||
type: Literal[OM.Type.ERROR]
|
||||
type: Literal[SM.Type.ERROR]
|
||||
code: str
|
||||
details: str
|
||||
|
||||
|
|
@ -97,11 +97,11 @@ class OM:
|
|||
Приход этого сообщения означает, что агент осознанно завершает работу с клиентом по какой-то причине.
|
||||
Для дальнейшего взаимодействия нужно снова обратиться к мастеру.
|
||||
"""
|
||||
type: Literal[OM.Type.GRACEFUL_DISCONNECT]
|
||||
type: Literal[SM.Type.GRACEFUL_DISCONNECT]
|
||||
|
||||
|
||||
ServerMessage = Annotated[
|
||||
Union[OM.Status, OM.AgentEvent, OM.Error, OM.GracefulDisconnect],
|
||||
Union[SM.Status, SM.AgentEvent, SM.Error, SM.GracefulDisconnect],
|
||||
Field(discriminator="type")
|
||||
]
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue