CallMeDaniel commited on
Commit
b20b4ef
·
1 Parent(s): cdccbe3

refactor: convert AgentDef from dataclass to Pydantic BaseModel

Browse files
Files changed (2) hide show
  1. agents/definitions.py +3 -3
  2. tests/test_types.py +14 -0
agents/definitions.py CHANGED
@@ -4,12 +4,12 @@ Agent data is loaded from config.yaml. The AGENTS dict, AGENT_COLORS,
4
  AGENT_NAMES, and AGENT_AVATARS are kept for backward compatibility.
5
  """
6
 
7
- from dataclasses import dataclass
 
8
  from config.settings import settings
9
 
10
 
11
- @dataclass
12
- class AgentDef:
13
  """Definition of a chat agent."""
14
  id: str
15
  name: str
 
4
  AGENT_NAMES, and AGENT_AVATARS are kept for backward compatibility.
5
  """
6
 
7
+ from pydantic import BaseModel
8
+
9
  from config.settings import settings
10
 
11
 
12
+ class AgentDef(BaseModel):
 
13
  """Definition of a chat agent."""
14
  id: str
15
  name: str
tests/test_types.py CHANGED
@@ -1,5 +1,6 @@
1
  """Tests for core/types.py — enums, dataclasses, ABC."""
2
  import pytest
 
3
  from core.types import BackendName, AgentId, AgentResponse, ChatResult, LLMBackend
4
 
5
 
@@ -86,3 +87,16 @@ class TestLLMBackendABC:
86
  system, rest = LLMBackend.split_system_message(msgs)
87
  assert system == ""
88
  assert len(rest) == 1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  """Tests for core/types.py — enums, dataclasses, ABC."""
2
  import pytest
3
+ from agents.definitions import AgentDef
4
  from core.types import BackendName, AgentId, AgentResponse, ChatResult, LLMBackend
5
 
6
 
 
87
  system, rest = LLMBackend.split_system_message(msgs)
88
  assert system == ""
89
  assert len(rest) == 1
90
+
91
+
92
+ class TestAgentDefModel:
93
+ def test_create(self):
94
+ ad = AgentDef(id="design", name="Design", role="Designer", color="#fff", avatar="D", goal="g", backstory="b")
95
+ assert ad.id == "design"
96
+ assert ad.name == "Design"
97
+
98
+ def test_model_dump(self):
99
+ ad = AgentDef(id="cad", name="CAD", role="Coder", color="#000", avatar="C", goal="g", backstory="b")
100
+ d = ad.model_dump()
101
+ assert d["id"] == "cad"
102
+ assert "role" in d