|
10 | 10 | Dict, |
11 | 11 | List, |
12 | 12 | TypedDict, |
13 | | - TypeVar, |
14 | 13 | Union, |
15 | 14 | ) |
16 | 15 |
|
17 | 16 | import yaml |
| 17 | +from pydantic import BaseModel |
18 | 18 | from tqdm import tqdm |
19 | 19 |
|
20 | 20 | from swarms.agents.create_agents_from_yaml import ( |
21 | 21 | create_agents_from_yaml, |
22 | 22 | ) |
23 | | -from swarms.schemas.swarms_api_schemas import AgentSpec |
24 | 23 | from swarms.utils.types import ReturnTypes |
25 | 24 | from swarms.structs.agent import Agent |
26 | 25 | from swarms.utils.agent_loader_markdown import ( |
|
29 | 28 | MarkdownAgentLoader, |
30 | 29 | ) |
31 | 30 |
|
32 | | -# Type variable for agent configuration |
33 | | -AgentConfigType = TypeVar( |
34 | | - "AgentConfigType", bound=Union[AgentSpec, Dict[str, Any]] |
35 | | -) |
36 | | - |
37 | 31 |
|
38 | 32 | class ModelName(str, Enum): |
39 | 33 | """Valid model names for swarms agents""" |
@@ -101,12 +95,11 @@ class AgentValidator: |
101 | 95 |
|
102 | 96 | @staticmethod |
103 | 97 | def validate_config( |
104 | | - config: Union[AgentSpec, Dict[str, Any]], |
| 98 | + config: Union[BaseModel, Dict[str, Any]], |
105 | 99 | ) -> AgentConfigDict: |
106 | | - """Validate and convert agent configuration from either AgentSpec or Dict""" |
| 100 | + """Validate a config supplied as a pydantic model or a plain dict.""" |
107 | 101 | try: |
108 | | - # Convert AgentSpec to dict if needed |
109 | | - if isinstance(config, AgentSpec): |
| 102 | + if isinstance(config, BaseModel): |
110 | 103 | config = config.model_dump() |
111 | 104 |
|
112 | 105 | # Deferred: litellm must not load at `import swarms` time (#1754). |
@@ -222,7 +215,7 @@ def file_type(self) -> FileType: |
222 | 215 | raise ValueError(f"Unsupported file type: {ext}") |
223 | 216 |
|
224 | 217 | def create_agent_file( |
225 | | - self, agents: List[Union[AgentSpec, Dict[str, Any]]] |
| 218 | + self, agents: List[Union[BaseModel, Dict[str, Any]]] |
226 | 219 | ) -> None: |
227 | 220 | """Create a file with validated agent configurations""" |
228 | 221 | validated_agents = [] |
@@ -291,7 +284,7 @@ def load_agents(self) -> List[Agent]: |
291 | 284 | return agents |
292 | 285 |
|
293 | 286 | def _process_agent( |
294 | | - self, agent_data: Union[AgentSpec, Dict[str, Any]] |
| 287 | + self, agent_data: Union[BaseModel, Dict[str, Any]] |
295 | 288 | ) -> Union[Agent, None]: |
296 | 289 | """Process a single agent configuration""" |
297 | 290 | try: |
|
0 commit comments