spaik-sdk 0.6.2__py3-none-any.whl

spaik_sdk/orchestration/base_orchestrator.py

@@ -0,0 +1,238 @@
+import asyncio
+from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Generic,
+    Optional,
+    TypeVar,
+    Union,
+)
+
+from spaik_sdk.orchestration.checkpoint import CheckpointProvider
+from spaik_sdk.orchestration.models import OrchestratorEvent
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+T_State = TypeVar("T_State")
+T_Result = TypeVar("T_Result")
+
+
+class BaseOrchestrator(ABC, Generic[T_State, T_Result]):
+    """
+    Code-first orchestration without graph DSLs.
+
+    Subclass this and implement `run()` to define your orchestration logic.
+    Use `step()` to execute steps with automatic status emission and optional checkpointing.
+
+    Example:
+        class MyOrchestrator(BaseOrchestrator[MyState, MyResult]):
+            async def run(self) -> AsyncIterator[OrchestratorEvent[MyResult]]:
+                state = MyState(items=[])
+
+                # Run a step - yields status events automatically
+                async for event in self.step("fetch", "Fetching data", self.fetch_data, state):
+                    yield event
+                    if event.result:
+                        state = event.result
+
+                # Emit progress during processing
+                for i, item in enumerate(state.items):
+                    yield self.progress("process", i + 1, len(state.items))
+                    await self.process_item(item)
+
+                yield self.ok(MyResult(processed=len(state.items)))
+
+            async def fetch_data(self, state: MyState) -> MyState:
+                # Your logic here
+                return state.copy(items=fetched_items)
+    """
+
+    def __init__(
+        self,
+        checkpoint_provider: Optional[CheckpointProvider[T_State]] = None,
+        resume_from: Optional[str] = None,
+    ) -> None:
+        """
+        Args:
+            checkpoint_provider: Optional provider for state persistence.
+                                 If None, no checkpointing is performed.
+            resume_from: Step ID to resume from. Steps up to and including
+                        this ID will be skipped, using checkpointed state.
+        """
+        self.checkpoint_provider = checkpoint_provider
+        self.resume_from = resume_from
+        self._completed_steps: set[str] = set()
+        self._passed_resume_point = False
+
+        if resume_from is not None:
+            logger.info(f"Will resume after step '{resume_from}'")
+
+    @abstractmethod
+    def run(self) -> AsyncIterator[OrchestratorEvent[T_Result]]:
+        """
+        Implement your orchestration logic here.
+
+        Yield OrchestratorEvent instances to emit status updates, progress,
+        messages, and the final result.
+        """
+        ...
+
+    def run_sync(self) -> OrchestratorEvent[T_Result]:
+        """
+        Run the orchestration synchronously and return the final event.
+
+        Returns the last event emitted (typically result or error).
+        """
+
+        async def _runner() -> OrchestratorEvent[T_Result]:
+            last_event: Optional[OrchestratorEvent[T_Result]] = None
+            async for event in self.run():
+                last_event = event
+                if event.is_terminal():
+                    return event
+            if last_event is None:
+                return self.fail("No events emitted during orchestration")
+            return last_event
+
+        return asyncio.run(_runner())
+
+    async def step(
+        self,
+        step_id: str,
+        name: str,
+        fn: Union[
+            Callable[[T_State], T_State],
+            Callable[[T_State], Awaitable[T_State]],
+        ],
+        state: T_State,
+    ) -> AsyncIterator[OrchestratorEvent[T_State]]:
+        """
+        Execute a step with automatic status emission and checkpointing.
+
+        Yields:
+            - step_started event
+            - step_completed/step_failed event
+            - result event with new state (or error event on failure)
+
+        If resuming and this step was already completed, yields step_skipped
+        and the checkpointed state instead.
+        """
+        if self._should_skip(step_id):
+            yield OrchestratorEvent.step_skipped(step_id, name, "Resumed from checkpoint")
+            loaded_state = self._load_checkpoint(step_id)
+            if loaded_state is not None:
+                yield OrchestratorEvent.state_update(loaded_state)
+            else:
+                yield OrchestratorEvent.fail(f"Checkpoint not found for step '{step_id}'")
+            return
+
+        yield OrchestratorEvent.step_started(step_id, name)
+
+        try:
+            result = fn(state)
+            if asyncio.iscoroutine(result):
+                new_state = await result
+            else:
+                new_state = result
+
+            self._save_checkpoint(step_id, new_state)
+            self._completed_steps.add(step_id)
+
+            yield OrchestratorEvent.step_completed(step_id, name)
+            yield OrchestratorEvent.state_update(new_state)
+
+        except Exception as e:
+            logger.exception(f"Step '{step_id}' failed")
+            yield OrchestratorEvent.step_failed(step_id, name, str(e))
+            yield OrchestratorEvent.fail(str(e))
+
+    # --- Convenience factory methods ---
+
+    def ok(self, result: T_Result) -> OrchestratorEvent[T_Result]:
+        """Create a success result event"""
+        return OrchestratorEvent.ok(result)
+
+    def fail(self, error: str) -> OrchestratorEvent[T_Result]:
+        """Create an error event"""
+        return OrchestratorEvent.fail(error)
+
+    def msg(self, message: str) -> OrchestratorEvent[T_Result]:
+        """Create a message event"""
+        return OrchestratorEvent.msg(message)
+
+    def progress(self, step_id: str, current: int, total: int, detail: Optional[str] = None) -> OrchestratorEvent[T_Result]:
+        """Create a progress update event"""
+        return OrchestratorEvent.progress_update(step_id, current, total, detail)
+
+    def step_started(self, step_id: str, name: str, detail: Optional[str] = None) -> OrchestratorEvent[T_Result]:
+        """Create a step started event (for manual step management)"""
+        return OrchestratorEvent.step_started(step_id, name, detail)
+
+    def step_completed(self, step_id: str, name: str, detail: Optional[str] = None) -> OrchestratorEvent[T_Result]:
+        """Create a step completed event (for manual step management)"""
+        return OrchestratorEvent.step_completed(step_id, name, detail)
+
+    def step_failed(self, step_id: str, name: str, error: str) -> OrchestratorEvent[T_Result]:
+        """Create a step failed event (for manual step management)"""
+        return OrchestratorEvent.step_failed(step_id, name, error)
+
+    # --- Internal helpers ---
+
+    def _should_skip(self, step_id: str) -> bool:
+        """Check if a step should be skipped due to checkpoint resume"""
+        if self.resume_from is None:
+            return False
+        if self._passed_resume_point:
+            return False
+        # We skip until we hit the resume_from step, then skip that one too
+        if step_id == self.resume_from:
+            self._passed_resume_point = True
+        return True
+
+    def _load_checkpoint(self, step_id: str) -> Optional[T_State]:
+        """Load checkpointed state for a step"""
+        if self.checkpoint_provider is None:
+            return None
+        return self.checkpoint_provider.load(step_id)
+
+    def _save_checkpoint(self, step_id: str, state: T_State) -> None:
+        """Save state to checkpoint after step completion"""
+        if self.checkpoint_provider is None:
+            return
+        self.checkpoint_provider.save(step_id, state)
+        logger.debug(f"Saved checkpoint for step '{step_id}'")
+
+
+class SimpleOrchestrator(BaseOrchestrator[None, T_Result]):
+    """
+    Orchestrator without state management.
+
+    Use this when you don't need to pass state between steps,
+    or when steps manage their own state internally.
+    """
+
+    async def run_step(
+        self,
+        step_id: str,
+        name: str,
+        fn: Union[Callable[[], Any], Callable[[], Awaitable[Any]]],
+    ) -> AsyncIterator[OrchestratorEvent[Any]]:
+        """Execute a step without state management"""
+        yield OrchestratorEvent.step_started(step_id, name)
+
+        try:
+            result = fn()
+            if asyncio.iscoroutine(result):
+                await result
+
+            self._completed_steps.add(step_id)
+            yield OrchestratorEvent.step_completed(step_id, name)
+
+        except Exception as e:
+            logger.exception(f"Step '{step_id}' failed")
+            yield OrchestratorEvent.step_failed(step_id, name, str(e))
+            yield OrchestratorEvent.fail(str(e))

spaik_sdk/orchestration/checkpoint.py

@@ -0,0 +1,80 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Generic, Optional, TypeVar
+
+T_State = TypeVar("T_State")
+
+
+class CheckpointProvider(ABC, Generic[T_State]):
+    """
+    Protocol for checkpoint storage/retrieval.
+
+    Implement this to add persistence for orchestration state.
+    The orchestrator will call these methods automatically during step execution.
+    """
+
+    @abstractmethod
+    def save(self, step_id: str, state: T_State) -> None:
+        """Save state after a step completes"""
+        pass
+
+    @abstractmethod
+    def load(self, step_id: str) -> Optional[T_State]:
+        """Load state for a specific step. Returns None if not found."""
+        pass
+
+    @abstractmethod
+    def get_completed_steps(self) -> set[str]:
+        """Get all step IDs that have been completed"""
+        pass
+
+    @abstractmethod
+    def clear(self) -> None:
+        """Clear all checkpoints"""
+        pass
+
+
+class InMemoryCheckpointProvider(CheckpointProvider[T_State]):
+    """Simple in-memory checkpoint provider for testing/development"""
+
+    def __init__(self) -> None:
+        self._checkpoints: Dict[str, T_State] = {}
+
+    def save(self, step_id: str, state: T_State) -> None:
+        self._checkpoints[step_id] = state
+
+    def load(self, step_id: str) -> Optional[T_State]:
+        return self._checkpoints.get(step_id)
+
+    def get_completed_steps(self) -> set[str]:
+        return set(self._checkpoints.keys())
+
+    def clear(self) -> None:
+        self._checkpoints.clear()
+
+
+class DictCheckpointProvider(CheckpointProvider[Dict[str, Any]]):
+    """
+    Checkpoint provider that serializes state as dicts.
+
+    Useful when you want to persist to JSON/database but your state
+    objects have to_dict()/from_dict() methods.
+    """
+
+    def __init__(self, storage: Optional[Dict[str, Dict[str, Any]]] = None) -> None:
+        self._storage = storage if storage is not None else {}
+
+    def save(self, step_id: str, state: Dict[str, Any]) -> None:
+        self._storage[step_id] = state
+
+    def load(self, step_id: str) -> Optional[Dict[str, Any]]:
+        return self._storage.get(step_id)
+
+    def get_completed_steps(self) -> set[str]:
+        return set(self._storage.keys())
+
+    def clear(self) -> None:
+        self._storage.clear()
+
+    def get_all(self) -> Dict[str, Dict[str, Any]]:
+        """Get all stored checkpoints (for serialization)"""
+        return dict(self._storage)

spaik_sdk/orchestration/models.py

@@ -0,0 +1,103 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, Generic, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+class StepStatus(Enum):
+    """Status of an orchestration step"""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    SKIPPED = "skipped"
+
+
+@dataclass
+class StepInfo:
+    """Information about a step's current state"""
+
+    step_id: str
+    name: str
+    status: StepStatus
+    detail: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class OrchestratorEvent(Generic[T]):
+    """
+    Event emitted during orchestration.
+
+    This is a union-style dataclass - only one field should be set at a time.
+    Use the factory methods for cleaner construction.
+
+    Note: `state` is for intermediate step outputs, `result` is for final orchestration result.
+    """
+
+    step: Optional[StepInfo] = None
+    message: Optional[str] = None
+    progress: Optional["ProgressUpdate"] = None
+    state: Optional[T] = None  # Intermediate state from steps
+    result: Optional[T] = None  # Final result
+    error: Optional[str] = None
+
+    @staticmethod
+    def step_started(step_id: str, name: str, detail: Optional[str] = None) -> "OrchestratorEvent[Any]":
+        return OrchestratorEvent(step=StepInfo(step_id, name, StepStatus.RUNNING, detail))
+
+    @staticmethod
+    def step_completed(step_id: str, name: str, detail: Optional[str] = None) -> "OrchestratorEvent[Any]":
+        return OrchestratorEvent(step=StepInfo(step_id, name, StepStatus.COMPLETED, detail))
+
+    @staticmethod
+    def step_failed(step_id: str, name: str, error: str) -> "OrchestratorEvent[Any]":
+        return OrchestratorEvent(step=StepInfo(step_id, name, StepStatus.FAILED, error))
+
+    @staticmethod
+    def step_skipped(step_id: str, name: str, reason: Optional[str] = None) -> "OrchestratorEvent[Any]":
+        return OrchestratorEvent(step=StepInfo(step_id, name, StepStatus.SKIPPED, reason))
+
+    @staticmethod
+    def msg(message: str) -> "OrchestratorEvent[Any]":
+        return OrchestratorEvent(message=message)
+
+    @staticmethod
+    def state_update(state: T) -> "OrchestratorEvent[Any]":
+        """Emit intermediate state from a step"""
+        return OrchestratorEvent(state=state)
+
+    @staticmethod
+    def ok(result: T) -> "OrchestratorEvent[Any]":
+        """Emit final result"""
+        return OrchestratorEvent(result=result)
+
+    @staticmethod
+    def fail(error: str) -> "OrchestratorEvent[Any]":
+        return OrchestratorEvent(error=error)
+
+    @staticmethod
+    def progress_update(step_id: str, current: int, total: int, detail: Optional[str] = None) -> "OrchestratorEvent[Any]":
+        return OrchestratorEvent(progress=ProgressUpdate(step_id, current, total, detail))
+
+    def is_terminal(self) -> bool:
+        """Returns True if this event represents a final state (result or error)"""
+        return self.result is not None or self.error is not None
+
+
+@dataclass
+class ProgressUpdate:
+    """Progress within a step (e.g., processing item 3/10)"""
+
+    step_id: str
+    current: int
+    total: int
+    detail: Optional[str] = None
+
+    @property
+    def percent(self) -> float:
+        if self.total == 0:
+            return 0.0
+        return (self.current / self.total) * 100

spaik_sdk/prompt/get_prompt_loader.py

@@ -0,0 +1,13 @@
+from typing import Optional
+
+from spaik_sdk.config.env import env_config
+from spaik_sdk.prompt.local_prompt_loader import LocalPromptLoader
+from spaik_sdk.prompt.prompt_loader import PromptLoader
+from spaik_sdk.prompt.prompt_loader_mode import PromptLoaderMode
+
+
+def get_prompt_loader(mode: Optional[PromptLoaderMode] = None) -> PromptLoader:
+    mode = mode or env_config.get_prompt_loader_mode()
+    if mode == PromptLoaderMode.LOCAL:
+        return LocalPromptLoader()
+    raise ValueError(f"Unknown PromptLoaderMode: {mode}")

spaik_sdk/prompt/local_prompt_loader.py

@@ -0,0 +1,21 @@
+import os
+from typing import Optional
+
+from spaik_sdk.config.env import env_config
+from spaik_sdk.prompt.prompt_loader import PromptLoader
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class LocalPromptLoader(PromptLoader):
+    def __init__(self, prompts_dir: Optional[str] = None):
+        super().__init__()
+        self.prompts_dir = prompts_dir if prompts_dir else env_config.get_prompts_dir()
+
+    def _load_raw_prompt(self, prompt_path: str) -> str:
+        full_path = os.path.join(self.prompts_dir, f"{prompt_path}.md")
+        if not os.path.exists(full_path):
+            raise FileNotFoundError(f"Prompt file {full_path} not found")
+        with open(full_path, "r") as f:
+            return f.read().strip()

spaik_sdk/prompt/prompt_loader.py

@@ -0,0 +1,48 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type
+
+if TYPE_CHECKING:
+    from spaik_sdk.agent.base_agent import BaseAgent
+
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class PromptLoader(ABC):
+    def __init__(self):
+        self._prompts_cache: Dict[str, str] = {}
+
+    @abstractmethod
+    def _load_raw_prompt(self, prompt_path: str) -> str:
+        pass
+
+    def _get_raw_prompt(self, prompt_path: str) -> str:
+        if prompt_path in self._prompts_cache:
+            return self._prompts_cache[prompt_path]
+        else:
+            raw_prompt = self._load_raw_prompt(prompt_path)
+            self._prompts_cache[prompt_path] = raw_prompt
+            return raw_prompt
+
+    def _format_prompt(self, prompt: str, params: Dict[str, Any]) -> str:
+        try:
+            return prompt.format(**params)
+        except KeyError as e:
+            missing_key = str(e).strip("'")
+            raise KeyError(f"Missing required variable '{missing_key}' in params: {params} for prompt: '{prompt}'")
+
+    def get_prompt(self, prompt_path: str, params: Dict[str, Any]) -> str:
+        return self._format_prompt(self._get_raw_prompt(prompt_path), params)
+
+    def _get_raw_agent_prompt(self, agent_class: Type["BaseAgent"], prompt_name: str, version: Optional[str] = None) -> str:
+        return self._get_raw_prompt(f"agent/{agent_class.__name__}/{prompt_name}{f'-{version}' if version else ''}")
+
+    def get_system_prompt(self, agent_class: Type["BaseAgent"], params: Dict[str, Any], version: Optional[str] = None) -> str:
+        return self.get_agent_prompt(agent_class, "system", params, version)
+
+    def get_agent_prompt(
+        self, agent_class: Type["BaseAgent"], prompt_name: str, params: Dict[str, Any], version: Optional[str] = None
+    ) -> str:
+        raw_prompt = self._get_raw_agent_prompt(agent_class, prompt_name, version)
+        return self._format_prompt(raw_prompt, params)

spaik_sdk/prompt/prompt_loader_mode.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class PromptLoaderMode(Enum):
+    LOCAL = "local"
+
+    @classmethod
+    def from_name(cls, name: str) -> "PromptLoaderMode":
+        for mode in cls:
+            if mode.value == name:
+                return mode
+
+        available_modes = [mode.value for mode in cls]
+        raise ValueError(f"Unknown PromptLoaderMode '{name}'. Available: {', '.join(available_modes)}")

spaik_sdk/py.typed

@@ -0,0 +1 @@
+ 

spaik_sdk/recording/__init__.py

@@ -0,0 +1 @@
+

spaik_sdk/recording/base_playback.py

@@ -0,0 +1,90 @@
+import time
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Iterator, Optional
+
+from spaik_sdk.recording.langchain_serializer import deserialize_token_data
+
+
+class BasePlayback(ABC):
+    """Abstract base class for playing back recorded LLM interactions."""
+
+    def __init__(self, recording_name: str = "default", delay: float = 0.001):
+        self.recording_name = recording_name
+        self.current_session = 1
+        self._iterator: Optional[Iterator[Dict[str, Any]]] = None
+        self.delay = delay
+
+    @abstractmethod
+    def _load_session_data_impl(self, session_num: int) -> Iterator[Dict[str, Any]]:
+        """Load raw data for a specific session number. Returns plain dicts."""
+        pass
+
+    def _load_session_data(self, session_num: int) -> Iterator[Dict[str, Any]]:
+        """Load data for a specific session number and deserialize LangChain objects."""
+        for raw_data in self._load_session_data_impl(session_num):
+            # Deserialize LangChain objects after loading from implementation
+            yield deserialize_token_data(raw_data)
+
+    @abstractmethod
+    def _session_exists(self, session_num: int) -> bool:
+        """Check if a session file exists."""
+        pass
+
+    @abstractmethod
+    def is_available(self) -> bool:
+        """Check if playback data is available."""
+        pass
+
+    def __iter__(self) -> Iterator[Dict[str, Any]]:
+        """Make the playback object iterable."""
+        return self
+
+    def __aiter__(self):
+        """Make the playback object async iterable."""
+        return self
+
+    async def __anext__(self) -> Dict[str, Any]:
+        """Async version of __next__."""
+        try:
+            return self.__next__()
+        except StopIteration:
+            raise StopAsyncIteration
+
+    def __next__(self) -> Dict[str, Any]:
+        """Get the next token/response from current session only."""
+        # Initialize iterator if not set
+        if self._iterator is None:
+            if not self._session_exists(self.current_session):
+                raise StopIteration("No recorded sessions found")
+            self._iterator = self._load_session_data(self.current_session)
+
+        try:
+            time.sleep(self.delay)
+            return next(self._iterator)
+        except StopIteration:
+            # Current session exhausted, auto-bump to next session
+            self.current_session += 1
+            self._iterator = None
+            raise StopIteration("Current session exhausted")
+
+    def reset(self) -> None:
+        """Reset playback to first session."""
+        self.current_session = 1
+        self._iterator = None
+
+    def get_current_session(self) -> int:
+        """Get the current session number."""
+        return self.current_session
+
+    def get_recording_name(self) -> str:
+        """Get the recording name."""
+        return self.recording_name
+
+    def next_session(self) -> None:
+        """Manually advance to next session."""
+        self.current_session += 1
+        self._iterator = None
+
+    def has_next_session(self) -> bool:
+        """Check if there's a next session available."""
+        return self._session_exists(self.current_session + 1)

spaik_sdk/recording/base_recorder.py

@@ -0,0 +1,50 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from spaik_sdk.recording.langchain_serializer import ensure_json_serializable, serialize_token_data
+
+
+class BaseRecorder(ABC):
+    """Abstract base class for recording LLM interactions."""
+
+    def __init__(self, recording_name: str = "default"):
+        self.recording_name = recording_name
+        self.current_session = 1
+
+    def record_token(self, token_data: Dict[str, Any]) -> None:
+        """Record a streaming token from LLM response."""
+        # Serialize LangChain objects before passing to implementation
+        serialized_data = serialize_token_data(token_data)
+        safe_data = ensure_json_serializable(serialized_data)
+        self._record_token_impl(safe_data)
+
+    def record_structured(self, data: Dict[str, Any]) -> None:
+        """Record structured response and immediately bump session counter."""
+        # Serialize LangChain objects before passing to implementation
+        serialized_data = serialize_token_data(data)
+        safe_data = ensure_json_serializable(serialized_data)
+        self._record_structured_impl(safe_data)
+
+    @abstractmethod
+    def _record_token_impl(self, token_data: Dict[str, Any]) -> None:
+        """Implementation-specific token recording. Data is already serialized."""
+        pass
+
+    @abstractmethod
+    def _record_structured_impl(self, data: Dict[str, Any]) -> None:
+        """Implementation-specific structured recording. Data is already serialized."""
+        pass
+
+    @abstractmethod
+    def request_completed(self) -> None:
+        """Mark current request as completed and bump to next session."""
+        pass
+
+    @abstractmethod
+    def get_current_session(self) -> int:
+        """Get the current session number."""
+        pass
+
+    def get_recording_name(self) -> str:
+        """Get the recording name."""
+        return self.recording_name