Upload mythos/pantheon.py

mythos/pantheon.py

@@ -0,0 +1,96 @@
+"""Pantheon — a collection of agents with defined relationships and workflows."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from .agent import Agent
+from .message import Message, MessageType
+
+
+class Relationship(BaseModel):
+    """A relationship between two agents in a Pantheon."""
+
+    source: str      # Agent name
+    target: str      # Agent name
+    relation: str    # e.g., "reports_to", "collaborates_with", "verifies"
+    weight: float = 1.0
+    bidirectional: bool = False
+
+
+class Pantheon(BaseModel):
+    """A collection of agents with relationships and shared context."""
+
+    name: str = "default"
+    agents: dict[str, Agent] = Field(default_factory=dict)
+    relationships: list[Relationship] = Field(default_factory=list)
+    shared_memory: dict[str, str] = Field(default_factory=dict)
+
+    def add_agent(self, agent: Agent) -> None:
+        """Add an agent to the pantheon."""
+        self.agents[agent.name] = agent
+
+    def remove_agent(self, name: str) -> None:
+        """Remove an agent from the pantheon."""
+        self.agents.pop(name, None)
+
+    def get_agent(self, name: str) -> Optional[Agent]:
+        """Get an agent by name."""
+        return self.agents.get(name)
+
+    def add_relationship(self, relationship: Relationship) -> None:
+        """Add a relationship between agents."""
+        self.relationships.append(relationship)
+
+    def get_related(self, agent_name: str, relation: Optional[str] = None) -> list[str]:
+        """Get agents related to the given agent."""
+        related = []
+        for rel in self.relationships:
+            if rel.source == agent_name:
+                if relation is None or rel.relation == relation:
+                    related.append(rel.target)
+            if rel.bidirectional and rel.target == agent_name:
+                if relation is None or rel.relation == relation:
+                    related.append(rel.source)
+        return related
+
+    async def broadcast(self, content: str, sender: str = "system") -> list[Message]:
+        """Broadcast a message to all agents."""
+        msg = Message(
+            type=MessageType.BROADCAST,
+            sender=sender,
+            recipient="broadcast",
+            content=content,
+        )
+        responses = []
+        for agent in self.agents.values():
+            response = await agent.handle(msg)
+            responses.append(response)
+        return responses
+
+    async def send(self, sender: str, recipient: str, content: str) -> Message:
+        """Send a direct message from one agent to another."""
+        sender_agent = self.agents.get(sender)
+        recipient_agent = self.agents.get(recipient)
+        if not sender_agent or not recipient_agent:
+            raise ValueError(f"Unknown agent: {sender} or {recipient}")
+
+        msg = Message(
+            type=MessageType.TASK,
+            sender=sender,
+            recipient=recipient,
+            content=content,
+        )
+        return await recipient_agent.handle(msg)
+
+    def list_agents(self) -> list[str]:
+        """List all agent names."""
+        return list(self.agents.keys())
+
+    def __len__(self) -> int:
+        return len(self.agents)
+
+    def __str__(self) -> str:
+        return f"Pantheon({self.name}, agents={list(self.agents.keys())})"