aragora-client 2.1.10__py3-none-any.whl

aragora_client/types.py

@@ -0,0 +1,323 @@
+"""Type definitions for the Aragora SDK."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class DebateStatus(str, Enum):
+    """Status of a debate."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class VerificationStatus(str, Enum):
+    """Status of a verification result."""
+
+    VALID = "valid"
+    INVALID = "invalid"
+    UNKNOWN = "unknown"
+    ERROR = "error"
+
+
+class ConsensusResult(BaseModel):
+    """Result of consensus detection."""
+
+    reached: bool
+    conclusion: str | None = None
+    confidence: float = 0.0
+    supporting_agents: list[str] = Field(default_factory=list)
+    dissenting_agents: list[str] = Field(default_factory=list)
+    reasoning: str | None = None
+
+
+class AgentMessage(BaseModel):
+    """A message from an agent during debate."""
+
+    agent_id: str
+    content: str
+    round_number: int
+    timestamp: datetime
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class Debate(BaseModel):
+    """A debate instance."""
+
+    id: str
+    task: str
+    status: DebateStatus
+    agents: list[str]
+    rounds: list[list[AgentMessage]] = Field(default_factory=list)
+    consensus: ConsensusResult | None = None
+    created_at: datetime
+    updated_at: datetime
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class GraphBranch(BaseModel):
+    """A branch in a graph debate."""
+
+    id: str
+    parent_id: str | None = None
+    approach: str
+    agents: list[str]
+    rounds: list[list[AgentMessage]] = Field(default_factory=list)
+    consensus: ConsensusResult | None = None
+    divergence_score: float = 0.0
+
+
+class GraphDebate(BaseModel):
+    """A graph debate with branching."""
+
+    id: str
+    task: str
+    status: DebateStatus
+    branches: list[GraphBranch] = Field(default_factory=list)
+    created_at: datetime
+    updated_at: datetime
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class MatrixConclusion(BaseModel):
+    """Conclusions from a matrix debate."""
+
+    universal: list[str] = Field(default_factory=list)
+    conditional: dict[str, list[str]] = Field(default_factory=dict)
+    contradictions: list[str] = Field(default_factory=list)
+
+
+class MatrixScenario(BaseModel):
+    """A scenario in a matrix debate."""
+
+    name: str
+    parameters: dict[str, Any]
+    is_baseline: bool = False
+    consensus: ConsensusResult | None = None
+
+
+class MatrixDebate(BaseModel):
+    """A matrix debate across scenarios."""
+
+    id: str
+    task: str
+    status: DebateStatus
+    scenarios: list[MatrixScenario] = Field(default_factory=list)
+    conclusions: MatrixConclusion | None = None
+    created_at: datetime
+    updated_at: datetime
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class AgentProfile(BaseModel):
+    """Profile of an agent."""
+
+    id: str
+    name: str
+    provider: str
+    elo_rating: float = 1500.0
+    matches_played: int = 0
+    win_rate: float = 0.5
+    specialties: list[str] = Field(default_factory=list)
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class VerificationResult(BaseModel):
+    """Result of formal verification."""
+
+    status: VerificationStatus
+    claim: str
+    formal_translation: str | None = None
+    proof: str | None = None
+    counterexample: str | None = None
+    backend: str
+    duration_ms: int = 0
+
+
+class GauntletFinding(BaseModel):
+    """A finding from gauntlet validation."""
+
+    severity: str
+    category: str
+    description: str
+    location: str | None = None
+    suggestion: str | None = None
+
+
+class GauntletReceipt(BaseModel):
+    """Receipt from gauntlet validation."""
+
+    id: str
+    score: float
+    findings: list[GauntletFinding] = Field(default_factory=list)
+    persona: str
+    created_at: datetime
+    hash: str | None = None
+
+
+class MemoryTierStats(BaseModel):
+    """Statistics for a memory tier."""
+
+    tier: str
+    entries: int
+    size_bytes: int
+    hit_rate: float
+    avg_age_seconds: float
+
+
+class MemoryAnalytics(BaseModel):
+    """Analytics for the memory system."""
+
+    total_entries: int
+    total_size_bytes: int
+    learning_velocity: float
+    tiers: list[MemoryTierStats] = Field(default_factory=list)
+    period_days: int
+
+
+class HealthStatus(BaseModel):
+    """Server health status."""
+
+    status: str
+    version: str
+    uptime_seconds: float
+    components: dict[str, str] = Field(default_factory=dict)
+
+
+class DebateEvent(BaseModel):
+    """A WebSocket event from a debate."""
+
+    type: str
+    data: dict[str, Any] = Field(default_factory=dict)
+    timestamp: datetime = Field(default_factory=datetime.now)
+    loop_id: str | None = None
+
+
+class ScorerInfo(BaseModel):
+    """Information about a scorer plugin."""
+
+    name: str
+    description: str
+
+
+class TeamSelectorInfo(BaseModel):
+    """Information about a team selector plugin."""
+
+    name: str
+    description: str
+
+
+class RoleAssignerInfo(BaseModel):
+    """Information about a role assigner plugin."""
+
+    name: str
+    description: str
+
+
+class SelectionPlugins(BaseModel):
+    """Available selection plugins."""
+
+    scorers: list[ScorerInfo] = Field(default_factory=list)
+    team_selectors: list[TeamSelectorInfo] = Field(default_factory=list)
+    role_assigners: list[RoleAssignerInfo] = Field(default_factory=list)
+
+
+class AgentScore(BaseModel):
+    """Score for an agent."""
+
+    name: str
+    score: float
+    elo_rating: float
+    breakdown: dict[str, float] = Field(default_factory=dict)
+
+
+class TeamMember(BaseModel):
+    """A member of a selected team."""
+
+    name: str
+    role: str
+    score: float
+
+
+class TeamSelection(BaseModel):
+    """Result of team selection."""
+
+    agents: list[TeamMember] = Field(default_factory=list)
+    expected_quality: float
+    diversity_score: float
+    rationale: str
+
+
+# Request/Response types for API calls
+class CreateDebateRequest(BaseModel):
+    """Request to create a debate."""
+
+    task: str
+    agents: list[str] | None = None
+    max_rounds: int = 5
+    consensus_threshold: float = 0.8
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class CreateGraphDebateRequest(BaseModel):
+    """Request to create a graph debate."""
+
+    task: str
+    agents: list[str] | None = None
+    max_rounds: int = 5
+    branch_threshold: float = 0.5
+    max_branches: int = 10
+
+
+class CreateMatrixDebateRequest(BaseModel):
+    """Request to create a matrix debate."""
+
+    task: str
+    scenarios: list[dict[str, Any]]
+    agents: list[str] | None = None
+    max_rounds: int = 3
+
+
+class VerifyClaimRequest(BaseModel):
+    """Request to verify a claim."""
+
+    claim: str
+    backend: str = "z3"
+    timeout: int = 30
+
+
+class RunGauntletRequest(BaseModel):
+    """Request to run gauntlet validation."""
+
+    input_content: str
+    input_type: str = "spec"
+    persona: str = "security"
+
+
+class ScoreAgentsRequest(BaseModel):
+    """Request to score agents."""
+
+    task_description: str
+    primary_domain: str | None = None
+    scorer: str | None = None
+
+
+class SelectTeamRequest(BaseModel):
+    """Request to select a team."""
+
+    task_description: str
+    min_agents: int = 2
+    max_agents: int = 5
+    diversity_preference: float = 0.5
+    quality_priority: float = 0.5
+    scorer: str | None = None
+    team_selector: str | None = None
+    role_assigner: str | None = None

aragora_client/websocket.py

@@ -0,0 +1,262 @@
+"""WebSocket streaming for the Aragora SDK."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import AsyncIterator, Callable
+
+import websockets
+from websockets.client import WebSocketClientProtocol
+
+from aragora_client.types import DebateEvent
+
+
+class DebateStream:
+    """
+    WebSocket stream for debate events.
+
+    Example:
+        >>> stream = DebateStream("ws://localhost:8765", "debate-123")
+        >>> stream.on("agent_message", lambda e: print(e.data))
+        >>> stream.on("consensus", lambda e: print("Consensus reached!"))
+        >>> await stream.connect()
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        debate_id: str,
+        *,
+        reconnect: bool = True,
+        reconnect_interval: float = 1.0,
+        max_reconnect_attempts: int = 5,
+        heartbeat_interval: float = 30.0,
+    ) -> None:
+        """
+        Initialize the debate stream.
+
+        Args:
+            base_url: WebSocket server URL.
+            debate_id: ID of the debate to stream.
+            reconnect: Whether to auto-reconnect on disconnect.
+            reconnect_interval: Base reconnect delay in seconds.
+            max_reconnect_attempts: Maximum reconnect attempts.
+            heartbeat_interval: Heartbeat ping interval in seconds.
+        """
+        # Normalize URL
+        if base_url.startswith("http://"):
+            base_url = base_url.replace("http://", "ws://")
+        elif base_url.startswith("https://"):
+            base_url = base_url.replace("https://", "wss://")
+        elif not base_url.startswith(("ws://", "wss://")):
+            base_url = f"ws://{base_url}"
+
+        self.base_url = base_url.rstrip("/")
+        self.debate_id = debate_id
+        self.reconnect = reconnect
+        self.reconnect_interval = reconnect_interval
+        self.max_reconnect_attempts = max_reconnect_attempts
+        self.heartbeat_interval = heartbeat_interval
+
+        self._ws: WebSocketClientProtocol | None = None
+        self._handlers: dict[str, list[Callable[[DebateEvent], None]]] = {}
+        self._error_handlers: list[Callable[[Exception], None]] = []
+        self._connected = False
+        self._should_stop = False
+        self._reconnect_attempts = 0
+
+    def on(
+        self, event_type: str, handler: Callable[[DebateEvent], None]
+    ) -> DebateStream:
+        """
+        Register an event handler.
+
+        Args:
+            event_type: Event type to handle (e.g., "agent_message", "consensus").
+            handler: Callback function to handle the event.
+
+        Returns:
+            Self for chaining.
+        """
+        if event_type not in self._handlers:
+            self._handlers[event_type] = []
+        self._handlers[event_type].append(handler)
+        return self
+
+    def on_error(self, handler: Callable[[Exception], None]) -> DebateStream:
+        """
+        Register an error handler.
+
+        Args:
+            handler: Callback function to handle errors.
+
+        Returns:
+            Self for chaining.
+        """
+        self._error_handlers.append(handler)
+        return self
+
+    async def connect(self) -> None:
+        """Connect to the WebSocket and start receiving events."""
+        self._should_stop = False
+        ws_url = f"{self.base_url}/ws/debates/{self.debate_id}"
+
+        while not self._should_stop:
+            try:
+                async with websockets.connect(ws_url) as ws:
+                    self._ws = ws
+                    self._connected = True
+                    self._reconnect_attempts = 0
+
+                    # Start heartbeat
+                    heartbeat_task = asyncio.create_task(self._heartbeat())
+
+                    try:
+                        async for message in ws:
+                            if self._should_stop:
+                                break
+                            await self._handle_message(message)
+                    finally:
+                        heartbeat_task.cancel()
+                        try:
+                            await heartbeat_task
+                        except asyncio.CancelledError:
+                            pass
+
+            except websockets.ConnectionClosed:
+                self._connected = False
+                if not self.reconnect or self._should_stop:
+                    break
+                if self._reconnect_attempts >= self.max_reconnect_attempts:
+                    max_attempts = self.max_reconnect_attempts
+                    self._emit_error(
+                        Exception(f"Max reconnect attempts ({max_attempts}) exceeded")
+                    )
+                    break
+                self._reconnect_attempts += 1
+                delay = self.reconnect_interval * (2 ** (self._reconnect_attempts - 1))
+                await asyncio.sleep(delay)
+
+            except Exception as e:
+                self._connected = False
+                self._emit_error(e)
+                if not self.reconnect or self._should_stop:
+                    break
+                self._reconnect_attempts += 1
+                if self._reconnect_attempts >= self.max_reconnect_attempts:
+                    break
+                delay = self.reconnect_interval * (2 ** (self._reconnect_attempts - 1))
+                await asyncio.sleep(delay)
+
+    async def disconnect(self) -> None:
+        """Disconnect from the WebSocket."""
+        self._should_stop = True
+        self._connected = False
+        if self._ws:
+            await self._ws.close()
+            self._ws = None
+
+    async def _heartbeat(self) -> None:
+        """Send periodic heartbeat pings."""
+        while self._connected and self._ws:
+            try:
+                await asyncio.sleep(self.heartbeat_interval)
+                if self._ws and self._connected:
+                    await self._ws.ping()
+            except Exception:
+                break
+
+    async def _handle_message(self, message: str | bytes) -> None:
+        """Handle an incoming WebSocket message."""
+        try:
+            if isinstance(message, bytes):
+                message = message.decode("utf-8")
+
+            data = json.loads(message)
+            event = DebateEvent(
+                type=data.get("type", "unknown"),
+                data=data.get("data", {}),
+                loop_id=data.get("loop_id") or data.get("data", {}).get("debate_id"),
+            )
+
+            # Call handlers for this event type
+            handlers = self._handlers.get(event.type, [])
+            for handler in handlers:
+                try:
+                    handler(event)
+                except Exception as e:
+                    self._emit_error(e)
+
+            # Call handlers for "*" (all events)
+            for handler in self._handlers.get("*", []):
+                try:
+                    handler(event)
+                except Exception as e:
+                    self._emit_error(e)
+
+        except json.JSONDecodeError as e:
+            self._emit_error(e)
+
+    def _emit_error(self, error: Exception) -> None:
+        """Emit an error to all error handlers."""
+        for handler in self._error_handlers:
+            try:
+                handler(error)
+            except Exception:
+                pass  # Avoid infinite loops
+
+    @property
+    def connected(self) -> bool:
+        """Whether the stream is currently connected."""
+        return self._connected
+
+
+async def stream_debate(
+    base_url: str,
+    debate_id: str,
+    *,
+    reconnect: bool = True,
+) -> AsyncIterator[DebateEvent]:
+    """
+    Stream debate events as an async iterator.
+
+    Example:
+        >>> async for event in stream_debate("ws://localhost:8765", "debate-123"):
+        ...     print(event.type, event.data)
+        ...     if event.type == "debate_end":
+        ...         break
+
+    Args:
+        base_url: WebSocket server URL.
+        debate_id: ID of the debate to stream.
+        reconnect: Whether to auto-reconnect on disconnect.
+
+    Yields:
+        DebateEvent objects.
+    """
+    # Normalize URL
+    if base_url.startswith("http://"):
+        base_url = base_url.replace("http://", "ws://")
+    elif base_url.startswith("https://"):
+        base_url = base_url.replace("https://", "wss://")
+    elif not base_url.startswith(("ws://", "wss://")):
+        base_url = f"ws://{base_url}"
+
+    ws_url = f"{base_url.rstrip('/')}/ws/debates/{debate_id}"
+
+    async with websockets.connect(ws_url) as ws:
+        async for message in ws:
+            if isinstance(message, bytes):
+                message = message.decode("utf-8")
+
+            try:
+                data = json.loads(message)
+                yield DebateEvent(
+                    type=data.get("type", "unknown"),
+                    data=data.get("data", {}),
+                    loop_id=data.get("loop_id")
+                    or data.get("data", {}).get("debate_id"),
+                )
+            except json.JSONDecodeError:
+                continue