prefactor-core 0.1.0__py3-none-any.whl

prefactor_core/__init__.py

@@ -0,0 +1,57 @@
+"""Public API for prefactor-core.
+
+This module exports the main classes and functions for the prefactor-core SDK.
+"""
+
+from .client import PrefactorCoreClient
+from .config import PrefactorCoreConfig, QueueConfig
+from .context_stack import SpanContextStack
+from .exceptions import (
+    ClientAlreadyInitializedError,
+    ClientNotInitializedError,
+    InstanceNotFoundError,
+    OperationError,
+    PrefactorCoreError,
+    SpanNotFoundError,
+)
+from .managers.agent_instance import AgentInstanceHandle
+from .models import AgentInstance, Span
+from .operations import Operation, OperationType
+from .queue import InMemoryQueue, Queue, QueueClosedError, TaskExecutor
+from .schema_registry import SchemaRegistry
+from .span_context import SpanContext
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Client
+    "PrefactorCoreClient",
+    # Config
+    "PrefactorCoreConfig",
+    "QueueConfig",
+    # Context
+    "SpanContext",
+    "SpanContextStack",
+    # Exceptions
+    "PrefactorCoreError",
+    "ClientNotInitializedError",
+    "ClientAlreadyInitializedError",
+    "OperationError",
+    "InstanceNotFoundError",
+    "SpanNotFoundError",
+    # Models
+    "AgentInstance",
+    "Span",
+    # Operations
+    "Operation",
+    "OperationType",
+    # Queue
+    "Queue",
+    "QueueClosedError",
+    "InMemoryQueue",
+    "TaskExecutor",
+    # Handle
+    "AgentInstanceHandle",
+    # Schema Registry
+    "SchemaRegistry",
+]

prefactor_core/client.py

@@ -0,0 +1,405 @@
+"""Main client for prefactor-core.
+
+This module provides the PrefactorCoreClient, which is the main entry point
+for the SDK. It manages the complete lifecycle of agent instances and spans
+through an async queue-based architecture.
+"""
+
+from __future__ import annotations
+
+import time
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+from prefactor_http.client import PrefactorHttpClient
+
+from .config import PrefactorCoreConfig
+from .context_stack import SpanContextStack
+from .exceptions import (
+    ClientAlreadyInitializedError,
+    ClientNotInitializedError,
+)
+from .managers.agent_instance import AgentInstanceManager
+from .managers.span import SpanManager
+from .operations import Operation, OperationType
+from .queue.base import Queue
+from .queue.executor import TaskExecutor
+from .queue.memory import InMemoryQueue
+
+if TYPE_CHECKING:
+    from .managers.agent_instance import AgentInstanceHandle
+
+
+class PrefactorCoreClient:
+    """Main entry point for the prefactor-core SDK.
+
+    This client provides a high-level interface for managing agent instances
+    and spans. All operations are queued and processed asynchronously, ensuring
+    minimal impact on agent execution flow.
+
+    The client must be initialized before use, either by calling initialize()
+    or using it as an async context manager.
+
+    Example:
+        config = PrefactorCoreConfig(http_config=...)
+
+        async with PrefactorCoreClient(config) as client:
+            instance = await client.create_agent_instance(...)
+            await instance.start()
+
+            async with instance.span("agent:llm") as span:
+                span.set_payload({"model": "gpt-4"})
+                # Your agent logic here
+
+            await instance.finish()
+    """
+
+    def __init__(
+        self,
+        config: PrefactorCoreConfig,
+        queue: Queue[Operation] | None = None,
+    ) -> None:
+        """Initialize the client.
+
+        Args:
+            config: Configuration for the client.
+            queue: Optional custom queue implementation. If not provided,
+                an InMemoryQueue is used.
+        """
+        self._config = config
+        self._queue = queue or InMemoryQueue()
+        self._http: PrefactorHttpClient | None = None
+        self._executor: TaskExecutor | None = None
+        self._instance_manager: AgentInstanceManager | None = None
+        self._span_manager: SpanManager | None = None
+        self._initialized = False
+
+    async def __aenter__(self) -> "PrefactorCoreClient":
+        """Enter async context manager."""
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit async context manager."""
+        await self.close()
+
+    async def initialize(self) -> None:
+        """Initialize the client and start processing.
+
+        This method:
+        1. Initializes the HTTP client
+        2. Starts the task executor
+        3. Initializes managers
+
+        Raises:
+            ClientAlreadyInitializedError: If already initialized.
+        """
+        if self._initialized:
+            raise ClientAlreadyInitializedError("Client is already initialized")
+
+        # Initialize HTTP client
+        self._http = PrefactorHttpClient(self._config.http_config)
+        await self._http.__aenter__()
+
+        # Initialize executor
+        self._executor = TaskExecutor(
+            queue=self._queue,
+            handler=self._process_operation,
+            num_workers=self._config.queue_config.num_workers,
+            max_retries=self._config.queue_config.max_retries,
+        )
+        self._executor.start()
+
+        # Initialize managers
+        self._instance_manager = AgentInstanceManager(
+            http_client=self._http,
+            enqueue=self._enqueue,
+        )
+        self._span_manager = SpanManager(
+            http_client=self._http,
+            enqueue=self._enqueue,
+        )
+
+        self._initialized = True
+
+    async def close(self) -> None:
+        """Close the client and cleanup resources.
+
+        This method gracefully shuts down the executor and closes the
+        HTTP client. It should be called when the client is no longer needed.
+        """
+        if not self._initialized:
+            return
+
+        # Stop executor
+        if self._executor:
+            await self._executor.stop()
+
+        # Close HTTP client
+        if self._http:
+            await self._http.__aexit__(None, None, None)
+
+        self._initialized = False
+
+    def _ensure_initialized(self) -> None:
+        """Ensure the client is initialized.
+
+        Raises:
+            ClientNotInitializedError: If not initialized.
+        """
+        if not self._initialized:
+            raise ClientNotInitializedError(
+                "Client is not initialized. Call initialize() or "
+                "use as context manager."
+            )
+
+    async def _enqueue(self, operation: Operation) -> None:
+        """Add an operation to the queue.
+
+        Args:
+            operation: The operation to queue.
+        """
+        await self._queue.put(operation)
+
+    async def _process_operation(self, operation: Operation) -> None:
+        """Process a single operation from the queue.
+
+        This method routes operations to the appropriate handler based on type.
+
+        Args:
+            operation: The operation to process.
+        """
+        if not self._http:
+            return
+
+        try:
+            if operation.type == OperationType.REGISTER_AGENT_INSTANCE:
+                await self._http.agent_instances.register(
+                    agent_id=operation.payload["agent_id"],
+                    agent_version=operation.payload["agent_version"],
+                    agent_schema_version=operation.payload["agent_schema_version"],
+                    id=operation.payload.get("id"),
+                )
+
+            elif operation.type == OperationType.START_AGENT_INSTANCE:
+                await self._http.agent_instances.start(
+                    agent_instance_id=operation.payload["instance_id"],
+                    timestamp=operation.timestamp,
+                )
+
+            elif operation.type == OperationType.FINISH_AGENT_INSTANCE:
+                await self._http.agent_instances.finish(
+                    agent_instance_id=operation.payload["instance_id"],
+                    timestamp=operation.timestamp,
+                )
+            elif operation.type == OperationType.CREATE_SPAN:
+                await self._http.agent_spans.create(
+                    agent_instance_id=operation.payload["instance_id"],
+                    schema_name=operation.payload["schema_name"],
+                    status=operation.payload.get("status", "pending"),
+                    id=operation.payload.get("span_id"),
+                    parent_span_id=operation.payload.get("parent_span_id"),
+                    payload=operation.payload.get("payload"),
+                )
+
+            elif operation.type == OperationType.FINISH_SPAN:
+                await self._http.agent_spans.finish(
+                    agent_span_id=operation.payload["span_id"],
+                    status=operation.payload.get("status", "complete"),
+                    result_payload=operation.payload.get("result_payload"),
+                )
+
+            # Note: UPDATE_SPAN_PAYLOAD requires API support or batching
+
+        except Exception as e:
+            # Log error but don't re-raise - we don't want to crash the worker
+            import logging
+
+            logger = logging.getLogger(__name__)
+            logger.error(
+                f"Failed to process operation {operation.type}: {e}",
+                exc_info=True,
+            )
+            raise
+
+    async def create_agent_instance(
+        self,
+        agent_id: str,
+        agent_version: dict[str, Any],
+        agent_schema_version: dict[str, Any] | None = None,
+        instance_id: str | None = None,
+        external_schema_version_id: str | None = None,
+    ) -> "AgentInstanceHandle":
+        """Create a new agent instance.
+
+        Returns immediately with a handle. The actual registration happens
+        asynchronously via the queue.
+
+        If agent_schema_version is not provided but the client has a schema_registry,
+        the registry's schemas will be used automatically.
+
+        Args:
+            agent_id: ID of the agent to create an instance for.
+            agent_version: Version information (name, etc.).
+            agent_schema_version: Schema version. Uses registry if not provided
+                and registry is configured.
+            instance_id: Optional custom ID for the instance.
+            external_schema_version_id: Optional external identifier for the
+                schema version. Defaults to "auto-generated" when using registry.
+
+        Returns:
+            AgentInstanceHandle for the created instance.
+
+        Raises:
+            ClientNotInitializedError: If the client is not initialized.
+            ValueError: If no schema version provided and registry not configured.
+        """
+        self._ensure_initialized()
+        assert self._instance_manager is not None
+
+        # Determine the agent_schema_version to use
+        final_schema_version: dict[str, Any]
+        if agent_schema_version is not None:
+            final_schema_version = agent_schema_version
+        elif self._config.schema_registry is not None:
+            # Use registry to generate schema version
+            ext_id = external_schema_version_id
+            if ext_id is None:
+                ext_id = f"auto-{agent_id}-{time.time()}"
+            final_schema_version = self._config.schema_registry.to_agent_schema_version(
+                ext_id
+            )
+        else:
+            msg1 = "agent_schema_version required when no schema_registry configured"
+            msg2 = "Either provide agent_schema_version or configure a SchemaRegistry"
+            raise ValueError(f"{msg1}. {msg2}.")
+
+        # Import here to avoid circular import
+        from .managers.agent_instance import AgentInstanceHandle
+
+        instance_id = await self._instance_manager.register(
+            agent_id=agent_id,
+            agent_version=agent_version,
+            agent_schema_version=final_schema_version,
+            instance_id=instance_id,
+        )
+
+        return AgentInstanceHandle(
+            instance_id=instance_id,
+            client=self,
+        )
+
+    async def create_span(
+        self,
+        instance_id: str,
+        schema_name: str,
+        parent_span_id: str | None = None,
+        payload: dict[str, Any] | None = None,
+    ) -> str:
+        """Create a span and return its ID without finishing it.
+
+        Use this for spans that need to stay open across multiple operations.
+        Call finish_span() when done.
+
+        Args:
+            instance_id: ID of the agent instance this span belongs to.
+            schema_name: Name of the schema for this span.
+            parent_span_id: Optional explicit parent span ID.
+            payload: Optional initial payload (params/inputs) stored on creation.
+
+        Returns:
+            The span ID.
+        """
+        self._ensure_initialized()
+        assert self._span_manager is not None
+
+        if parent_span_id is None:
+            parent_span_id = SpanContextStack.peek()
+
+        return await self._span_manager.create(
+            instance_id=instance_id,
+            schema_name=schema_name,
+            parent_span_id=parent_span_id,
+            payload=payload,
+        )
+
+    async def finish_span(
+        self,
+        span_id: str,
+        result_payload: dict[str, Any] | None = None,
+    ) -> None:
+        """Finish a previously created span.
+
+        Args:
+            span_id: The ID of the span to finish.
+            result_payload: Optional result data to store on the span.
+        """
+        self._ensure_initialized()
+        assert self._span_manager is not None
+
+        await self._span_manager.finish(span_id, result_payload=result_payload)
+
+    @asynccontextmanager
+    async def span(
+        self,
+        instance_id: str,
+        schema_name: str,
+        parent_span_id: str | None = None,
+        span_id: str | None = None,
+        payload: dict[str, Any] | None = None,
+    ):
+        """Context manager for creating and finishing a span.
+
+        If parent_span_id is not provided, the current span from the
+        SpanContextStack is used as the parent.
+
+        The returned :class:`SpanContext` supports an explicit lifecycle:
+
+        1. ``await span.start(payload)`` — POST the span to the API.
+        2. Do work.
+        3. ``await span.complete(result)`` / ``span.fail(result)`` /
+           ``span.cancel()`` — finish with a specific status.
+
+        If ``start()`` or a finish method is not called explicitly, the
+        context manager handles them automatically on exit.
+
+        Args:
+            instance_id: ID of the agent instance this span belongs to.
+            schema_name: Name of the schema for this span.
+            parent_span_id: Optional explicit parent span ID.
+            span_id: Ignored (API generates IDs).
+            payload: Optional initial payload sent via auto-start on exit
+                if ``start()`` is never called explicitly.
+
+        Yields:
+            SpanContext for the created span.
+        """
+        self._ensure_initialized()
+        assert self._span_manager is not None
+
+        # Import here to avoid circular import
+        from .span_context import SpanContext
+
+        # Auto-detect parent from stack if not explicit
+        if parent_span_id is None:
+            parent_span_id = SpanContextStack.peek()
+
+        temp_id = self._span_manager.prepare(
+            instance_id=instance_id,
+            schema_name=schema_name,
+            parent_span_id=parent_span_id,
+        )
+
+        context = SpanContext(
+            temp_id=temp_id,
+            span_manager=self._span_manager,
+            default_payload=payload,
+        )
+
+        try:
+            yield context
+        finally:
+            await context.finish()
+
+
+__all__ = ["PrefactorCoreClient"]

prefactor_core/config.py

@@ -0,0 +1,60 @@
+"""Configuration for prefactor-core.
+
+This module contains configuration classes for the prefactor-core SDK.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from prefactor_http.config import HttpClientConfig
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    pass
+
+
+class QueueConfig(BaseModel):
+    """Configuration for queue processing.
+
+    Attributes:
+        num_workers: Number of concurrent worker tasks.
+        max_retries: Maximum retry attempts per failed operation.
+        retry_delay_base: Base delay for exponential backoff (seconds).
+    """
+
+    num_workers: int = Field(default=3, ge=1, le=20)
+    max_retries: int = Field(default=3, ge=0)
+    retry_delay_base: float = Field(default=1.0, gt=0)
+
+
+class PrefactorCoreConfig(BaseModel):
+    """Complete configuration for PrefactorCoreClient.
+
+    Attributes:
+        http_config: Configuration for the HTTP client.
+        queue_config: Configuration for queue processing.
+        schema_registry: Optional schema registry for aggregating span type definitions.
+
+    Example:
+        from prefactor_core.schema_registry import SchemaRegistry
+        from prefactor_core import PrefactorCoreConfig
+
+        registry = SchemaRegistry()
+        registry.register("langchain:llm", {"type": "object"})
+
+        config = PrefactorCoreConfig(
+            http_config=HttpClientConfig(...),
+            schema_registry=registry
+        )
+    """
+
+    http_config: HttpClientConfig
+    queue_config: QueueConfig = Field(default_factory=QueueConfig)
+    schema_registry: Any = Field(
+        default=None,
+        description="Optional SchemaRegistry for aggregating span type definitions",
+    )
+
+
+__all__ = ["QueueConfig", "PrefactorCoreConfig"]

prefactor_core/context_stack.py

@@ -0,0 +1,118 @@
+"""Span context stack for managing nested span relationships.
+
+The SpanContextStack provides a stack-based context tracking system for
+nested spans. Each async context maintains its own stack of active span IDs,
+allowing automatic parent detection when creating new spans.
+
+This module uses contextvars to ensure proper isolation between concurrent
+async operations.
+"""
+
+from contextvars import ContextVar
+from typing import Optional
+
+# Context variable to hold the span stack for the current async context
+# Each async task gets its own copy of this variable
+_current_span_stack: ContextVar[list[str]] = ContextVar("span_stack", default=[])
+
+
+class SpanContextStack:
+    """Manages a stack of active span IDs within an async context.
+
+    The stack tracks the hierarchy of nested spans. The top of the stack
+    is always the current (innermost) active span, which serves as the
+    default parent for any new spans created within the same context.
+
+    This class uses contextvars to ensure that each async task maintains
+    its own independent stack, preventing interference between concurrent
+    operations.
+
+    Example:
+        # Root span
+        SpanContextStack.push("span-1")
+        assert SpanContextStack.peek() == "span-1"
+
+        # Nested span (child of span-1)
+        SpanContextStack.push("span-2")
+        assert SpanContextStack.peek() == "span-2"
+
+        # Exit nested span
+        SpanContextStack.pop()
+        assert SpanContextStack.peek() == "span-1"
+
+        # Exit root span
+        SpanContextStack.pop()
+        assert SpanContextStack.peek() is None
+    """
+
+    @classmethod
+    def get_stack(cls) -> list[str]:
+        """Get the current span stack for this async context.
+
+        Returns:
+            A list of span IDs, from outermost to innermost.
+            Returns an empty list if no spans are active.
+        """
+        return _current_span_stack.get()
+
+    @classmethod
+    def push(cls, span_id: str) -> None:
+        """Push a span ID onto the stack.
+
+        This marks the span as the current (innermost) active span.
+
+        Args:
+            span_id: The ID of the span to push onto the stack.
+        """
+        stack = cls.get_stack()
+        new_stack = stack + [span_id]
+        _current_span_stack.set(new_stack)
+
+    @classmethod
+    def pop(cls) -> Optional[str]:
+        """Pop and return the current span ID from the stack.
+
+        Returns:
+            The span ID that was removed from the stack, or None
+            if the stack was empty.
+        """
+        stack = cls.get_stack()
+        if not stack:
+            return None
+
+        span_id = stack[-1]
+        new_stack = stack[:-1]
+        _current_span_stack.set(new_stack)
+        return span_id
+
+    @classmethod
+    def peek(cls) -> Optional[str]:
+        """Get the current span ID without removing it from the stack.
+
+        Returns:
+            The ID of the current (innermost) span, or None if no
+            spans are active in this context.
+        """
+        stack = cls.get_stack()
+        return stack[-1] if stack else None
+
+    @classmethod
+    def depth(cls) -> int:
+        """Get the current nesting depth.
+
+        Returns:
+            The number of active spans in the stack (0 if empty).
+        """
+        return len(cls.get_stack())
+
+    @classmethod
+    def is_empty(cls) -> bool:
+        """Check if the stack is empty.
+
+        Returns:
+            True if no spans are currently active.
+        """
+        return len(cls.get_stack()) == 0
+
+
+__all__ = ["SpanContextStack"]

prefactor_core/exceptions.py

@@ -0,0 +1,49 @@
+"""Custom exceptions for prefactor-core."""
+
+
+class PrefactorCoreError(Exception):
+    """Base exception for all prefactor-core errors."""
+
+    pass
+
+
+class ClientNotInitializedError(PrefactorCoreError):
+    """Raised when attempting to use a client that hasn't been initialized."""
+
+    pass
+
+
+class ClientAlreadyInitializedError(PrefactorCoreError):
+    """Raised when attempting to initialize a client that's already initialized."""
+
+    pass
+
+
+class OperationError(PrefactorCoreError):
+    """Raised when an operation fails to process."""
+
+    def __init__(self, message: str, operation_type: str | None = None) -> None:
+        super().__init__(message)
+        self.operation_type = operation_type
+
+
+class InstanceNotFoundError(PrefactorCoreError):
+    """Raised when an agent instance is not found."""
+
+    pass
+
+
+class SpanNotFoundError(PrefactorCoreError):
+    """Raised when a span is not found."""
+
+    pass
+
+
+__all__ = [
+    "PrefactorCoreError",
+    "ClientNotInitializedError",
+    "ClientAlreadyInitializedError",
+    "OperationError",
+    "InstanceNotFoundError",
+    "SpanNotFoundError",
+]

prefactor_core/managers/__init__.py

@@ -0,0 +1,6 @@
+"""Manager exports for prefactor-core."""
+
+from .agent_instance import AgentInstanceHandle, AgentInstanceManager
+from .span import SpanManager
+
+__all__ = ["AgentInstanceManager", "AgentInstanceHandle", "SpanManager"]