swarmkit 0.1.34__py3-none-any.whl

swarmkit/pipeline/types.py

@@ -0,0 +1,272 @@
+"""Pipeline Types - Fluent API for chaining Swarm operations."""
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Generic, List, Literal, Optional, TypeVar, Union
+
+from ..swarm.types import (
+    BestOfConfig,
+    VerifyConfig,
+    SchemaType,
+    Prompt,
+)
+from ..swarm.results import SwarmResult, ReduceResult
+from ..retry import RetryConfig
+
+
+T = TypeVar('T')
+
+
+# =============================================================================
+# EMIT OPTION (filter only)
+# =============================================================================
+
+EmitOption = Literal["success", "filtered", "all"]
+"""What filter emits to the next step.
+
+- "success": Items that passed condition (default)
+- "filtered": Items that failed condition
+- "all": Both success and filtered
+"""
+
+
+# =============================================================================
+# STEP CONFIGURATIONS
+# =============================================================================
+
+@dataclass
+class MapConfig(Generic[T]):
+    """Map step configuration."""
+    prompt: Prompt
+    """Task prompt."""
+    name: Optional[str] = None
+    """Step name for observability (appears in events)."""
+    system_prompt: Optional[str] = None
+    """System prompt override."""
+    schema: Optional[SchemaType] = None
+    """Schema for structured output."""
+    schema_options: Optional[Dict[str, Any]] = None
+    """Validation options for JSON Schema."""
+    agent: Optional[Any] = None
+    """Agent override."""
+    mcp_servers: Optional[Dict[str, Any]] = None
+    """MCP servers override (replaces swarm default for this step)."""
+    best_of: Optional[BestOfConfig] = None
+    """BestOf configuration (mutually exclusive with verify)."""
+    verify: Optional[VerifyConfig] = None
+    """Verify configuration (mutually exclusive with bestOf)."""
+    retry: Optional[RetryConfig] = None
+    """Retry configuration."""
+    timeout_ms: Optional[int] = None
+    """Timeout in ms."""
+
+
+@dataclass
+class FilterConfig(Generic[T]):
+    """Filter step configuration."""
+    prompt: str
+    """Evaluation prompt."""
+    schema: SchemaType
+    """Schema for structured output (required)."""
+    condition: Callable[[Any], bool]
+    """Condition function to determine pass/fail."""
+    name: Optional[str] = None
+    """Step name for observability (appears in events)."""
+    system_prompt: Optional[str] = None
+    """System prompt override."""
+    schema_options: Optional[Dict[str, Any]] = None
+    """Validation options for JSON Schema."""
+    agent: Optional[Any] = None
+    """Agent override."""
+    mcp_servers: Optional[Dict[str, Any]] = None
+    """MCP servers override (replaces swarm default for this step)."""
+    emit: EmitOption = "success"
+    """What to emit to next step (default: "success")."""
+    verify: Optional[VerifyConfig] = None
+    """Verify configuration."""
+    retry: Optional[RetryConfig] = None
+    """Retry configuration."""
+    timeout_ms: Optional[int] = None
+    """Timeout in ms."""
+
+
+@dataclass
+class ReduceConfig(Generic[T]):
+    """Reduce step configuration."""
+    prompt: str
+    """Synthesis prompt."""
+    name: Optional[str] = None
+    """Step name for observability (appears in events)."""
+    system_prompt: Optional[str] = None
+    """System prompt override."""
+    schema: Optional[SchemaType] = None
+    """Schema for structured output."""
+    schema_options: Optional[Dict[str, Any]] = None
+    """Validation options for JSON Schema."""
+    agent: Optional[Any] = None
+    """Agent override."""
+    mcp_servers: Optional[Dict[str, Any]] = None
+    """MCP servers override (replaces swarm default for this step)."""
+    verify: Optional[VerifyConfig] = None
+    """Verify configuration."""
+    retry: Optional[RetryConfig] = None
+    """Retry configuration."""
+    timeout_ms: Optional[int] = None
+    """Timeout in ms."""
+
+
+# =============================================================================
+# INTERNAL
+# =============================================================================
+
+StepType = Literal["map", "filter", "reduce"]
+
+
+@dataclass
+class Step:
+    """Internal step representation."""
+    type: StepType
+    config: Union[MapConfig, FilterConfig, ReduceConfig]
+
+
+# =============================================================================
+# RESULTS
+# =============================================================================
+
+@dataclass
+class StepResult(Generic[T]):
+    """Result of a single pipeline step."""
+    type: StepType
+    index: int
+    duration_ms: int
+    results: Union[List[SwarmResult[T]], ReduceResult[T]]
+
+
+@dataclass
+class PipelineResult(Generic[T]):
+    """Final result from pipeline execution."""
+    run_id: str
+    steps: List[StepResult]
+    output: Union[List[SwarmResult[T]], ReduceResult[T]]
+    total_duration_ms: int
+
+
+# =============================================================================
+# EVENTS
+# =============================================================================
+
+@dataclass
+class StepEvent:
+    """Step lifecycle event (base class)."""
+    type: StepType
+    index: int
+    name: Optional[str]
+
+
+@dataclass
+class StepStartEvent(StepEvent):
+    """Emitted when step starts."""
+    item_count: int
+
+
+@dataclass
+class StepCompleteEvent(StepEvent):
+    """Emitted when step completes."""
+    duration_ms: int
+    success_count: int
+    error_count: int
+    filtered_count: int
+
+
+@dataclass
+class StepErrorEvent(StepEvent):
+    """Emitted when step errors."""
+    error: Exception
+
+
+@dataclass
+class ItemRetryEvent:
+    """Emitted on item retry."""
+    step_index: int
+    step_name: Optional[str]
+    item_index: int
+    attempt: int
+    error: str
+
+
+@dataclass
+class WorkerCompleteEvent:
+    """Emitted when verify worker completes."""
+    step_index: int
+    step_name: Optional[str]
+    item_index: int
+    attempt: int
+    status: Literal["success", "error"]
+
+
+@dataclass
+class VerifierCompleteEvent:
+    """Emitted when verifier completes."""
+    step_index: int
+    step_name: Optional[str]
+    item_index: int
+    attempt: int
+    passed: bool
+    feedback: Optional[str]
+
+
+@dataclass
+class CandidateCompleteEvent:
+    """Emitted when bestOf candidate completes."""
+    step_index: int
+    step_name: Optional[str]
+    item_index: int
+    candidate_index: int
+    status: Literal["success", "error"]
+
+
+@dataclass
+class JudgeCompleteEvent:
+    """Emitted when bestOf judge completes."""
+    step_index: int
+    step_name: Optional[str]
+    item_index: int
+    winner_index: int
+    reasoning: str
+
+
+@dataclass
+class PipelineEvents:
+    """Event handlers."""
+    on_step_start: Optional[Callable[[StepStartEvent], None]] = None
+    on_step_complete: Optional[Callable[[StepCompleteEvent], None]] = None
+    on_step_error: Optional[Callable[[StepErrorEvent], None]] = None
+    on_item_retry: Optional[Callable[[ItemRetryEvent], None]] = None
+    on_worker_complete: Optional[Callable[[WorkerCompleteEvent], None]] = None
+    on_verifier_complete: Optional[Callable[[VerifierCompleteEvent], None]] = None
+    on_candidate_complete: Optional[Callable[[CandidateCompleteEvent], None]] = None
+    on_judge_complete: Optional[Callable[[JudgeCompleteEvent], None]] = None
+
+
+# Event name mapping for chainable .on() style
+PipelineEventMap = {
+    "step_start": "on_step_start",
+    "step_complete": "on_step_complete",
+    "step_error": "on_step_error",
+    "item_retry": "on_item_retry",
+    "worker_complete": "on_worker_complete",
+    "verifier_complete": "on_verifier_complete",
+    "candidate_complete": "on_candidate_complete",
+    "judge_complete": "on_judge_complete",
+}
+
+# Event names (for type hints)
+EventName = Literal[
+    "step_start",
+    "step_complete",
+    "step_error",
+    "item_retry",
+    "worker_complete",
+    "verifier_complete",
+    "candidate_complete",
+    "judge_complete",
+]

swarmkit/prompts/__init__.py

@@ -0,0 +1,126 @@
+"""Prompt templates for Swarm abstractions.
+
+Prompts are stored as markdown files for easy editing.
+They are loaded at import time using importlib.resources.
+"""
+
+from importlib import resources
+from typing import Dict, Union
+
+
+def _load_prompt(subdir: str, filename: str) -> str:
+    """Load a prompt template from a .md file in a subdirectory."""
+    return resources.files(__package__).joinpath(subdir, filename).read_text(encoding='utf-8')
+
+
+# Load agent prompts (system prompts - goes into CLAUDE.md-like context)
+JUDGE_PROMPT: str = _load_prompt('agent_md', 'judge.md')
+VERIFY_PROMPT: str = _load_prompt('agent_md', 'verify.md')
+REDUCE_PROMPT: str = _load_prompt('agent_md', 'reduce.md')
+
+# Load user prompts (task prompts - passed to .run())
+JUDGE_USER_PROMPT: str = _load_prompt('user', 'judge.md')
+VERIFY_USER_PROMPT: str = _load_prompt('user', 'verify.md')
+RETRY_FEEDBACK_PROMPT: str = _load_prompt('user', 'retry_feedback.md')
+
+
+def apply_template(template: str, variables: Dict[str, str]) -> str:
+    """Apply template variables to a template string.
+
+    Replaces {{variable}} with the corresponding value from variables dict.
+
+    Args:
+        template: Template string with {{variable}} placeholders
+        variables: Dict mapping variable names to values
+
+    Returns:
+        Template with placeholders replaced
+    """
+    result = template
+    for key, value in variables.items():
+        result = result.replace(f"{{{{{key}}}}}", value)
+    return result
+
+
+def build_file_tree(files: Dict[str, Union[str, bytes]]) -> str:
+    """Build a formatted file tree string for judge context.
+
+    Generates an ASCII tree representation of the files structure,
+    with comments explaining the purpose of each section.
+
+    Args:
+        files: Dict mapping file paths to content
+
+    Returns:
+        Formatted file tree string
+    """
+    if not files:
+        return "context/\n  (empty)"
+
+    # Get unique top-level folders, sorted with worker_task first
+    folders = sorted(
+        set(path.split("/")[0] for path in files.keys()),
+        key=lambda f: ("" if f == "worker_task" else f)
+    )
+
+    if not folders:
+        return "context/\n  (empty)"
+
+    # Check what exists in worker_task/
+    has_system_prompt = "worker_task/system_prompt.txt" in files
+    has_schema = "worker_task/schema.json" in files
+    has_input = any(p.startswith("worker_task/input/") for p in files.keys())
+
+    # Build entries with comments
+    entries: list[tuple[str, str]] = []  # (line, comment)
+
+    for i, folder in enumerate(folders):
+        is_last_folder = i == len(folders) - 1
+        folder_prefix = "└── " if is_last_folder else "├── "
+        child_indent = "    " if is_last_folder else "│   "
+
+        if folder == "worker_task":
+            entries.append((f"{folder_prefix}{folder}/", "task given to workers"))
+
+            # Build worker_task children (only show what exists)
+            children: list[tuple[str, str]] = []
+            if has_system_prompt:
+                children.append(("system_prompt.txt", "worker system prompt"))
+            children.append(("user_prompt.txt", "worker task prompt"))
+            if has_schema:
+                children.append(("schema.json", "expected output schema"))
+            if has_input:
+                children.append(("input/", "worker input files"))
+
+            for j, (name, comment) in enumerate(children):
+                is_last_child = j == len(children) - 1
+                child_prefix = "└── " if is_last_child else "├── "
+                entries.append((f"{child_indent}{child_prefix}{name}", comment))
+
+        elif folder.startswith("candidate_"):
+            idx = folder.replace("candidate_", "")
+            entries.append((f"{folder_prefix}{folder}/", f"worker {idx} solution"))
+        elif folder == "worker_output":
+            entries.append((f"{folder_prefix}{folder}/", "output to verify"))
+        elif folder.startswith("item_"):
+            idx = folder.replace("item_", "")
+            entries.append((f"{folder_prefix}{folder}/", f"input {idx}"))
+        else:
+            entries.append((f"{folder_prefix}{folder}/", ""))
+
+    # Calculate max line width for alignment
+    max_width = max(len(line) for line, _ in entries)
+
+    # Build final output with aligned comments
+    lines = ["context/"]
+    for line, comment in entries:
+        if comment:
+            padding = " " * (max_width - len(line) + 3)
+            lines.append(f"{line}{padding}# {comment}")
+        else:
+            lines.append(line)
+
+    return "\n".join(lines)
+
+
+__all__ = ['JUDGE_PROMPT', 'JUDGE_USER_PROMPT', 'VERIFY_PROMPT', 'VERIFY_USER_PROMPT', 'REDUCE_PROMPT', 'RETRY_FEEDBACK_PROMPT', 'apply_template', 'build_file_tree']

swarmkit/prompts/agent_md/judge.md

@@ -0,0 +1,30 @@
+### 1. YOUR ROLE: BEST OF N JUDGE
+
+You are a judge. {{candidateCount}} AI workers attempted the same task independently. Your job is to analyze their solution attempts and pick the best one based on the evaluation criteria below.
+
+### 2. CONTEXT STRUCTURE
+
+```
+{{fileTree}}
+```
+
+### 3. YOUR EVALUATION CRITERIA
+
+You must judge their work based on:
+
+```
+{{criteria}}
+```
+
+### 4. YOUR PROCESS
+
+1. Read `worker_task/` to understand the task:
+   - Review the worker system prompt and task prompt
+   - Check the expected output schema (if present)
+   - Examine the worker input files in `input/`
+2. Carefully review EACH solution attempt in `candidate_i/`
+3. Compare outputs against the evaluation criteria
+4. Reason through your findings — perform all necessary evidence-based analyses and verifications before deciding
+5. Pick the best candidate (0-indexed)
+
+**IMPORTANT:** Be thorough. Do not skip steps. Your judgment must be evidence-based — cite specific files, outputs, or discrepancies to justify your decision.

swarmkit/prompts/agent_md/reduce.md

@@ -0,0 +1,7 @@
+### CONTEXT STRUCTURE
+
+```
+{{fileTree}}
+```
+
+Examine all files within each item folder before producing your output.

swarmkit/prompts/agent_md/verify.md

@@ -0,0 +1,33 @@
+### 1. YOUR ROLE: OUTPUT VERIFIER
+
+You are a quality verifier. An AI worker produced output for a task. Your job is to verify whether the output meets the specified quality criteria.
+
+### 2. CONTEXT STRUCTURE
+
+```
+{{fileTree}}
+```
+
+### 3. VERIFICATION CRITERIA
+
+The output must satisfy:
+
+```
+{{criteria}}
+```
+
+### 4. YOUR PROCESS
+
+1. Read `worker_task/` to understand what was asked:
+   - Review the worker system prompt (if present)
+   - Review the task prompt
+   - Check the expected output schema (if present)
+   - Examine any input files in `input/`
+2. Carefully review the worker's output in `worker_output/`
+3. Evaluate against the verification criteria
+4. Reason through your findings
+5. Make your decision
+
+**IMPORTANT:** Be thorough and fair. Cite specific evidence. If the output generally achieves the goal with minor issues, consider passing. Only fail if there are significant problems that violate the criteria.
+
+If failing, provide specific, actionable feedback explaining what needs to be fixed.

swarmkit/prompts/user/judge.md

@@ -0,0 +1 @@
+Evaluate the candidates and select the best one. You must save your decision to the file `output/result.json`.

swarmkit/prompts/user/retry_feedback.md

@@ -0,0 +1,9 @@
+{{originalPrompt}}
+
+## IMPORTANT: Previous Attempt Failed Verification
+
+Your previous attempt did not pass quality verification. Please address this feedback:
+
+{{feedback}}
+
+Make sure your output addresses all the feedback points above.

swarmkit/prompts/user/verify.md

@@ -0,0 +1 @@
+Verify the worker output against the criteria. You must save your decision to the file `output/result.json`.

swarmkit/results.py

@@ -0,0 +1,45 @@
+"""Result types for SwarmKit SDK."""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Union
+
+
+@dataclass
+class AgentResponse:
+    """Response from agent execution.
+
+    Matches TypeScript SDK's AgentResponse for exact parity.
+
+    Attributes:
+        sandbox_id: Sandbox ID
+        exit_code: Command exit code
+        stdout: Standard output
+        stderr: Standard error
+    """
+    sandbox_id: str
+    exit_code: int
+    stdout: str
+    stderr: str
+
+
+# Backward compatibility alias
+ExecuteResult = AgentResponse
+
+
+@dataclass
+class OutputResult:
+    """Result from get_output_files() with optional schema validation.
+
+    Matches TypeScript SDK's OutputResult<T> for exact parity.
+    Evidence: sdk-ts/src/types.ts lines 258-268
+
+    Attributes:
+        files: Output files from output/ folder
+        data: Parsed and validated result.json data (None if no schema or validation failed)
+        error: Validation or parse error message, if any
+        raw_data: Raw result.json string when parse or validation failed (for debugging)
+    """
+    files: Dict[str, Union[str, bytes]] = field(default_factory=dict)
+    data: Optional[Any] = None
+    error: Optional[str] = None
+    raw_data: Optional[str] = None

swarmkit/retry.py

@@ -0,0 +1,133 @@
+"""Retry Utility for Swarm operations.
+
+Generic retry with exponential backoff.
+Works with any result type that has a status field.
+"""
+
+import asyncio
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable, Optional, TypeVar
+
+# =============================================================================
+# CONSTANTS
+# =============================================================================
+
+DEFAULT_MAX_ATTEMPTS = 3
+DEFAULT_BACKOFF_MS = 1000
+DEFAULT_BACKOFF_MULTIPLIER = 2.0
+
+# =============================================================================
+# TYPES
+# =============================================================================
+
+# TypeVar for result types (SwarmResult, ReduceResult, etc.)
+# Results must have a `status` field for default retry behavior.
+TResult = TypeVar('TResult')
+
+# Callback type for item retry events (must be defined before RetryConfig)
+OnItemRetryCallback = Callable[[int, int, str], None]  # (item_index, attempt, error)
+
+
+def _get_field(obj: Any, field: str, default: Any = None) -> Any:
+    """Get field from dict or object (duck typing helper)."""
+    if isinstance(obj, dict):
+        return obj.get(field, default)
+    return getattr(obj, field, default)
+
+
+@dataclass
+class RetryConfig:
+    """Per-item retry configuration.
+
+    Example:
+        # Basic retry on error
+        RetryConfig(max_attempts=3)
+
+        # With exponential backoff
+        RetryConfig(max_attempts=3, backoff_ms=1000, backoff_multiplier=2)
+
+        # Custom retry condition
+        RetryConfig(max_attempts=3, retry_on=lambda r: r.status == "error" or "timeout" in (r.error or ""))
+
+        # With callback
+        RetryConfig(max_attempts=3, on_item_retry=lambda i, a, e: print(f"Item {i} retry {a}: {e}"))
+
+    Args:
+        max_attempts: Maximum retry attempts (default: 3)
+        backoff_ms: Initial backoff in ms (default: 1000)
+        backoff_multiplier: Exponential backoff multiplier (default: 2)
+        retry_on: Custom retry condition (default: status == "error")
+        on_item_retry: Callback when retry occurs (item_index, attempt, error)
+    """
+    max_attempts: int = DEFAULT_MAX_ATTEMPTS
+    backoff_ms: int = DEFAULT_BACKOFF_MS
+    backoff_multiplier: float = DEFAULT_BACKOFF_MULTIPLIER
+    retry_on: Optional[Callable[[Any], bool]] = None
+    on_item_retry: Optional[OnItemRetryCallback] = None
+
+    def should_retry(self, result: Any) -> bool:
+        """Check if result should be retried."""
+        if self.retry_on is not None:
+            return self.retry_on(result)
+        # Default: retry on error status
+        return _get_field(result, 'status') == "error"
+
+
+# =============================================================================
+# RETRY LOGIC
+# =============================================================================
+
+async def execute_with_retry(
+    fn: Callable[[int], Awaitable[TResult]],
+    config: Optional[RetryConfig] = None,
+    item_index: int = 0,
+) -> TResult:
+    """Execute a function with retry and exponential backoff.
+
+    Works with any result type that has a `status` field (SwarmResult, ReduceResult, etc.).
+
+    Args:
+        fn: Async function that receives attempt number (1-based) and returns a result
+        config: Retry configuration (optional, uses defaults if not provided)
+        item_index: Item index for callback (default: 0)
+
+    Returns:
+        Result from the function
+
+    Example:
+        result = await execute_with_retry(
+            lambda attempt: self._execute_map_item(item, prompt, index, run_id, params, timeout, attempt),
+            RetryConfig(max_attempts=3, backoff_ms=1000),
+            item_index=index,
+        )
+    """
+    resolved = config or RetryConfig()
+
+    last_result: Optional[TResult] = None
+    attempts = 0
+    backoff = resolved.backoff_ms
+
+    while attempts < resolved.max_attempts:
+        attempts += 1
+        last_result = await fn(attempts)
+
+        # Check if we should retry
+        if not resolved.should_retry(last_result):
+            return last_result
+
+        # Don't retry if we've exhausted attempts
+        if attempts >= resolved.max_attempts:
+            break
+
+        # Notify of retry via callback in config
+        if resolved.on_item_retry is not None:
+            error = _get_field(last_result, 'error') or "Unknown error"
+            resolved.on_item_retry(item_index, attempts, error)
+
+        # Wait before retrying (convert ms to seconds)
+        await asyncio.sleep(backoff / 1000)
+        backoff = backoff * resolved.backoff_multiplier
+
+    # Return last result
+    assert last_result is not None
+    return last_result