diff --git a/pyproject.toml b/pyproject.toml index b91e16c..73e1dfe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,8 +46,6 @@ known-first-party = ["smolswarms"] [tool.mypy] python_version = "3.10" -warn_return_type = true -warn_unused_configs = true disallow_untyped_defs = true check_untyped_defs = true strict_optional = true diff --git a/smolswarms/factory.py b/smolswarms/factory.py new file mode 100644 index 0000000..9a4c81e --- /dev/null +++ b/smolswarms/factory.py @@ -0,0 +1,43 @@ +"""Factory for creating and configuring swarms.""" + +from typing import Optional + +from .specs import AgentSpec +from .specs import BusinessUnit +from .swarm import Swarm + + +class SwarmFactory: + """Factory class for creating and configuring swarms.""" + + @staticmethod + def create_swarm(name: str, description: Optional[str] = None) -> Swarm: + """Create a new swarm with the given name and description.""" + return Swarm(name=name, description=description) + + @staticmethod + def create_business_unit( + name: str, description: str, objectives: list[str] + ) -> BusinessUnit: + """Create a new business unit with the given specifications.""" + return BusinessUnit( + name=name, + description=description, + agents=[], + objectives=objectives, + ) + + @staticmethod + def create_agent( + name: str, + role: str, + capabilities: list[str], + memory_size: Optional[int] = None, + ) -> AgentSpec: + """Create a new agent with the given specifications.""" + return AgentSpec( + name=name, + role=role, + capabilities=capabilities, + memory_size=memory_size if memory_size is not None else 1024, + ) diff --git a/smolswarms/specs.py b/smolswarms/specs.py new file mode 100644 index 0000000..71a8f49 --- /dev/null +++ b/smolswarms/specs.py @@ -0,0 +1,21 @@ +"""Specifications for agents and business units in the swarm.""" + +from pydantic import BaseModel + + +class AgentSpec(BaseModel): + """Specification for an individual agent in the swarm.""" + + name: str + role: str + capabilities: list[str] + memory_size: int = 1024 + + +class BusinessUnit(BaseModel): + """A business unit that groups agents together for a specific purpose.""" + + name: str + description: str + agents: list[AgentSpec] + objectives: list[str] diff --git a/smolswarms/swarm.py b/smolswarms/swarm.py new file mode 100644 index 0000000..d80806d --- /dev/null +++ b/smolswarms/swarm.py @@ -0,0 +1,29 @@ +"""Core swarm implementation for orchestrating agents.""" + +from typing import Optional + +from .specs import AgentSpec +from .specs import BusinessUnit + + +class Swarm: + """A swarm of agents working together to achieve objectives.""" + + def __init__(self, name: str, description: Optional[str] = None) -> None: + """Initialize a new swarm with a name and optional description.""" + self.name = name + self.description = description + self.business_units: list[BusinessUnit] = [] + self.agents: list[AgentSpec] = [] + + def add_business_unit(self, unit: BusinessUnit) -> None: + """Add a business unit to the swarm.""" + self.business_units.append(unit) + self.agents.extend(unit.agents) + + def get_agent(self, name: str) -> Optional[AgentSpec]: + """Get an agent by name.""" + for agent in self.agents: + if agent.name == name: + return agent + return None