swarmkit 0.1.34__py3-none-any.whl

swarmkit/schema.py

@@ -0,0 +1,107 @@
+"""Schema utilities for Pydantic models, dataclasses, and JSON Schema.
+
+Provides unified detection, conversion, and validation for schema types.
+Uses Pydantic's TypeAdapter for dataclass support.
+"""
+
+import dataclasses
+from typing import Any, Dict, Optional
+
+
+# =============================================================================
+# DETECTION
+# =============================================================================
+
+
+def is_pydantic_model(obj: Any) -> bool:
+    """Check if object is a Pydantic model class."""
+    try:
+        from pydantic import BaseModel
+        return isinstance(obj, type) and issubclass(obj, BaseModel)
+    except ImportError:
+        return False
+
+
+def is_dataclass(obj: Any) -> bool:
+    """Check if object is a dataclass."""
+    return dataclasses.is_dataclass(obj) and isinstance(obj, type)
+
+
+def is_json_schema(obj: Any) -> bool:
+    """Check if object is a JSON Schema dict."""
+    return isinstance(obj, dict)
+
+
+# =============================================================================
+# CONVERSION
+# =============================================================================
+
+
+def to_json_schema(schema: Any) -> Optional[Dict[str, Any]]:
+    """Convert a schema to JSON Schema format.
+
+    Supports:
+    - Dict (JSON Schema) - passed through
+    - Pydantic models - uses model_json_schema()
+    - Dataclasses - uses Pydantic TypeAdapter
+
+    Args:
+        schema: Pydantic model, dataclass, or JSON Schema dict
+
+    Returns:
+        JSON Schema dict, or None if schema is None
+    """
+    if schema is None:
+        return None
+
+    if is_json_schema(schema):
+        return schema
+
+    if is_pydantic_model(schema):
+        return schema.model_json_schema()
+
+    if is_dataclass(schema):
+        from pydantic import TypeAdapter
+        return TypeAdapter(schema).json_schema()
+
+    raise TypeError(
+        f"Schema must be a Pydantic model, dataclass, or dict. "
+        f"Got {type(schema).__name__}"
+    )
+
+
+# =============================================================================
+# VALIDATION
+# =============================================================================
+
+
+def validate_and_parse(
+    raw_json: str,
+    schema: Any,
+    strict: bool = False,
+) -> Any:
+    """Validate JSON string and parse into schema type.
+
+    Args:
+        raw_json: Raw JSON string to validate
+        schema: Pydantic model or dataclass (returns instance)
+        strict: Use strict validation mode
+
+    Returns:
+        Model/dataclass instance, or None for dict schemas
+
+    Raises:
+        ValidationError: If validation fails
+    """
+    if schema is None or is_json_schema(schema):
+        # JSON Schema validation is handled by TS SDK
+        return None
+
+    if is_pydantic_model(schema):
+        return schema.model_validate_json(raw_json, strict=strict)
+
+    if is_dataclass(schema):
+        from pydantic import TypeAdapter
+        return TypeAdapter(schema).validate_json(raw_json, strict=strict)
+
+    return None

swarmkit/swarm/__init__.py

@@ -0,0 +1,75 @@
+"""Swarm abstractions for parallel AI agent execution."""
+
+from .swarm import Swarm
+from .types import (
+    FileMap,
+    SwarmConfig,
+    BestOfConfig,
+    VerifyConfig,
+    IndexedMeta,
+    ReduceMeta,
+    JudgeMeta,
+    VerifyMeta,
+    BaseMeta,
+    OperationType,
+    Prompt,
+    PromptFn,
+    ItemInput,
+    SchemaType,
+    JudgeDecision,
+    VerifyDecision,
+    # Callback types
+    OnCandidateCompleteCallback,
+    OnJudgeCompleteCallback,
+    OnWorkerCompleteCallback,
+    OnVerifierCompleteCallback,
+)
+from .results import (
+    SwarmResult,
+    SwarmResultList,
+    ReduceResult,
+    BestOfResult,
+    BestOfInfo,
+    VerifyInfo,
+    is_swarm_result,
+    SWARM_RESULT_BRAND,
+)
+
+__all__ = [
+    # Main class
+    'Swarm',
+    # Config types
+    'SwarmConfig',
+    'BestOfConfig',
+    'VerifyConfig',
+    # Result types
+    'SwarmResult',
+    'SwarmResultList',
+    'ReduceResult',
+    'BestOfResult',
+    'BestOfInfo',
+    'VerifyInfo',
+    # Meta types
+    'IndexedMeta',
+    'ReduceMeta',
+    'JudgeMeta',
+    'VerifyMeta',
+    'BaseMeta',
+    'OperationType',
+    # Other types
+    'FileMap',
+    'Prompt',
+    'PromptFn',
+    'ItemInput',
+    'SchemaType',
+    'JudgeDecision',
+    'VerifyDecision',
+    # Callback types
+    'OnCandidateCompleteCallback',
+    'OnJudgeCompleteCallback',
+    'OnWorkerCompleteCallback',
+    'OnVerifierCompleteCallback',
+    # Helpers
+    'is_swarm_result',
+    'SWARM_RESULT_BRAND',
+]

swarmkit/swarm/results.py

@@ -0,0 +1,140 @@
+"""Result types for Swarm abstractions."""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Generic, List, Literal, Optional, TypeVar, Union
+
+from .types import FileMap, IndexedMeta, ReduceMeta, JudgeMeta, VerifyMeta
+
+
+T = TypeVar('T')
+
+# Brand for runtime detection of SwarmResult (for chaining)
+SWARM_RESULT_BRAND = "__swarm_result__"
+
+
+# =============================================================================
+# SWARM RESULT
+# =============================================================================
+
+@dataclass
+class SwarmResult(Generic[T]):
+    """Result from a single worker (map, filter, bestof candidate).
+
+    Status meanings:
+    - "success": Positive outcome (agent succeeded / condition passed)
+    - "filtered": Neutral outcome (evaluated but didn't pass condition) - filter only
+    - "error": Negative outcome (agent error)
+    """
+    status: Literal["success", "filtered", "error"]
+    data: Optional[T]
+    files: FileMap
+    meta: IndexedMeta
+    error: Optional[str] = None
+    raw_data: Optional[str] = None
+    best_of: Optional['BestOfInfo'] = None
+    verify: Optional['VerifyInfo'] = None
+
+    # Brand for runtime detection (used by is_swarm_result() function)
+    __swarm_result__: bool = field(default=True, repr=False)
+
+
+@dataclass
+class BestOfInfo(Generic[T]):
+    """BestOf information attached to SwarmResult when map used bestOf option."""
+    winner_index: int
+    judge_reasoning: str
+    judge_meta: JudgeMeta
+    candidates: List[SwarmResult[T]]
+
+
+@dataclass
+class VerifyInfo:
+    """Verification info attached to results when verify option used."""
+    passed: bool
+    reasoning: str
+    verify_meta: VerifyMeta
+    attempts: int
+
+
+# =============================================================================
+# SWARM RESULT LIST
+# =============================================================================
+
+class SwarmResultList(List[SwarmResult[T]], Generic[T]):
+    """List of SwarmResults with helper properties.
+
+    Extends list so all normal list operations work.
+
+    Getters:
+    - `.success` - items with positive outcome
+    - `.filtered` - items that didn't pass condition (filter only)
+    - `.error` - items that encountered errors
+
+    Chaining examples:
+    - `swarm.reduce(results.success, ...)` - forward only successful
+    - `swarm.reduce([*results.success, *results.filtered], ...)` - forward all evaluated
+    """
+
+    @property
+    def success(self) -> List[SwarmResult[T]]:
+        """Returns items with status 'success'."""
+        return [r for r in self if r.status == "success"]
+
+    @property
+    def filtered(self) -> List[SwarmResult[T]]:
+        """Returns items with status 'filtered' (didn't pass condition)."""
+        return [r for r in self if r.status == "filtered"]
+
+    @property
+    def error(self) -> List[SwarmResult[T]]:
+        """Returns items with status 'error'."""
+        return [r for r in self if r.status == "error"]
+
+    @classmethod
+    def from_results(cls, results: List[SwarmResult[T]]) -> 'SwarmResultList[T]':
+        """Create SwarmResultList from a list of SwarmResults."""
+        result_list = cls()
+        result_list.extend(results)
+        return result_list
+
+
+# =============================================================================
+# REDUCE RESULT
+# =============================================================================
+
+@dataclass
+class ReduceResult(Generic[T]):
+    """Result from reduce operation."""
+    status: Literal["success", "error"]
+    data: Optional[T]
+    files: FileMap
+    meta: ReduceMeta
+    error: Optional[str] = None
+    raw_data: Optional[str] = None
+    verify: Optional['VerifyInfo'] = None
+
+
+# =============================================================================
+# BESTOF RESULT
+# =============================================================================
+
+@dataclass
+class BestOfResult(Generic[T]):
+    """Result from bestOf operation."""
+    winner: SwarmResult[T]
+    winner_index: int
+    judge_reasoning: str
+    judge_meta: JudgeMeta
+    candidates: List[SwarmResult[T]]
+
+
+# =============================================================================
+# HELPER FUNCTIONS
+# =============================================================================
+
+def is_swarm_result(obj: Any) -> bool:
+    """Check if an object is a SwarmResult (for chaining detection)."""
+    return (
+        isinstance(obj, SwarmResult) or
+        (hasattr(obj, SWARM_RESULT_BRAND) and getattr(obj, SWARM_RESULT_BRAND) is True)
+    )