корректные pydantic модели для автоматического определения класса по полю type

tests/manual.py

@@ -0,0 +1,24 @@
+import asyncio
+
+from lambda_agent_api.agent_api import AgentApi
+from lambda_agent_api.server import MsgEventTextChunk
+
+
+def my_callback(message):
+    print(f"Callback: {message}")
+
+
+async def main():
+    api = AgentApi("agent-1", "ws://localhost:8000/agent_ws/", callback=my_callback)
+
+    await api.connect()
+    try:
+        async for chunk in api.send_message("Привет, агент!"):
+            if isinstance(chunk, MsgEventTextChunk):
+                print(chunk.text, end="", flush=True)
+
+    finally:
+        await api.close()
+
+
+asyncio.run(main())

tests/test_models.py

@@ -0,0 +1,74 @@
+import pytest
+from pydantic import ValidationError
+
+from lambda_agent_api.server import *
+from lambda_agent_api.client import *
+
+
+@pytest.mark.parametrize(
+    "data, expected_type",
+    [
+        ({"type": "USER_MESSAGE", "text": "hello"}, MsgUserMessage),
+    ],
+)
+def test_client_message_valid(data, expected_type):
+    msg = ClientMessage.validate_python(data)
+    assert isinstance(msg, expected_type)
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        {"type": "UNKNOWN", "text": "hello"},
+        {"type": "USER_MESSAGE"},  # нет text
+    ],
+)
+def test_client_message_invalid(data):
+    with pytest.raises(ValidationError):
+        ClientMessage.validate_python(data)
+
+
+@pytest.mark.parametrize(
+    "data, expected_type",
+    [
+        ({"type": "STATUS"}, MsgStatus),
+        ({"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hi"}, MsgEventTextChunk),
+        ({"type": "AGENT_EVENT_END", "tokens_used": 10}, MsgEventEnd),
+        ({"type": "ERROR", "code": "E1", "details": "fail"}, MsgError),
+        ({"type": "GRACEFUL_DISCONNECT"}, MsgGracefulDisconnect),
+    ],
+)
+def test_server_message_valid(data, expected_type):
+    msg = ServerMessage.validate_python(data)
+    assert isinstance(msg, expected_type)
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        {"type": "AGENT_EVENT_TEXT_CHUNK"},  # нет text
+        {"type": "AGENT_EVENT_END"},  # нет tokens_used
+        {"type": "ERROR", "code": "E1"},  # нет details
+        {"type": "UNKNOWN"},
+    ],
+)
+def test_server_message_invalid(data):
+    with pytest.raises(ValidationError):
+        ServerMessage.validate_python(data)
+
+
+def test_validate_json():
+    json_data = '{"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hello"}'
+    msg = ServerMessage.validate_json(json_data)
+
+    assert isinstance(msg, MsgEventTextChunk)
+    assert msg.text == "hello"
+
+
+def test_model_dump_roundtrip():
+    data = {"type": "ERROR", "code": "E1", "details": "fail"}
+
+    msg = ServerMessage.validate_python(data)
+    dumped = msg.model_dump()
+
+    assert dumped == data