celeste-ai 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

celeste/http.py

@@ -0,0 +1,201 @@
+"""HTTP client with persistent connection pooling for AI provider APIs."""
+
+import json
+import logging
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+from httpx_sse import aconnect_sse
+
+from celeste.core import Capability, Provider
+
+logger = logging.getLogger(__name__)
+
+MAX_CONNECTIONS = 20
+MAX_KEEPALIVE_CONNECTIONS = 10
+DEFAULT_TIMEOUT = 60.0
+
+
+class HTTPClient:
+    """Async HTTP client with persistent connection pooling."""
+
+    def __init__(
+        self,
+        max_connections: int = MAX_CONNECTIONS,
+        max_keepalive_connections: int = MAX_KEEPALIVE_CONNECTIONS,
+    ) -> None:
+        """Initialize HTTP client with connection pool limits.
+
+        Args:
+            max_connections: Maximum total connections in pool.
+            max_keepalive_connections: Maximum idle keepalive connections.
+        """
+        self._client: httpx.AsyncClient | None = None
+        self._max_connections = max_connections
+        self._max_keepalive_connections = max_keepalive_connections
+
+    async def _get_client(self) -> httpx.AsyncClient:
+        """Get or create httpx.AsyncClient with connection pooling."""
+        if self._client is None:
+            limits = httpx.Limits(
+                max_connections=self._max_connections,
+                max_keepalive_connections=self._max_keepalive_connections,
+            )
+            self._client = httpx.AsyncClient(limits=limits)  # nosec B113
+        return self._client
+
+    async def post(
+        self,
+        url: str,
+        headers: dict[str, str],
+        json_body: dict[str, Any],
+        timeout: float = DEFAULT_TIMEOUT,
+    ) -> httpx.Response:
+        """Make POST request with connection pooling.
+
+        Args:
+            url: Full URL to POST to.
+            headers: HTTP headers including authentication.
+            json_body: JSON request body.
+            timeout: Request timeout in seconds.
+
+        Returns:
+            HTTP response from the server.
+
+        Raises:
+            httpx.HTTPError: On network or timeout errors.
+        """
+        client = await self._get_client()
+        return await client.post(
+            url,
+            headers=headers,
+            json=json_body,
+            timeout=timeout,
+        )
+
+    async def get(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+        timeout: float = DEFAULT_TIMEOUT,
+        follow_redirects: bool = True,
+    ) -> httpx.Response:
+        """Make GET request with connection pooling.
+
+        Args:
+            url: Full URL to GET.
+            headers: HTTP headers including authentication (optional).
+            timeout: Request timeout in seconds.
+            follow_redirects: Whether to follow HTTP redirects (default: True).
+
+        Returns:
+            HTTP response from the server.
+
+        Raises:
+            httpx.HTTPError: On network or timeout errors.
+        """
+        client = await self._get_client()
+        return await client.get(
+            url,
+            headers=headers or {},
+            timeout=timeout,
+            follow_redirects=follow_redirects,
+        )
+
+    async def stream_post(
+        self,
+        url: str,
+        headers: dict[str, str],
+        json_body: dict[str, Any],
+        timeout: float = DEFAULT_TIMEOUT,
+    ) -> AsyncIterator[dict[str, Any]]:
+        """Stream POST request using Server-Sent Events.
+
+        Args:
+            url: API endpoint URL.
+            headers: HTTP headers (including authentication).
+            json_body: JSON request body.
+            timeout: Timeout in seconds (default: DEFAULT_TIMEOUT).
+
+        Yields:
+            Parsed JSON events from SSE stream.
+        """
+        client = await self._get_client()
+
+        async with aconnect_sse(
+            client,
+            "POST",
+            url,
+            json=json_body,
+            headers=headers,
+            timeout=timeout,
+        ) as event_source:
+            async for sse in event_source.aiter_sse():
+                try:
+                    yield json.loads(sse.data)
+                except json.JSONDecodeError:
+                    continue  # Skip non-JSON control messages (provider-agnostic)
+
+    async def aclose(self) -> None:
+        """Close HTTP client and cleanup all connections."""
+        if self._client is not None:
+            await self._client.aclose()
+            self._client = None
+
+    async def __aenter__(self) -> "HTTPClient":
+        """Enter async context manager."""
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:  # noqa: ANN401
+        """Exit async context manager and cleanup connections."""
+        await self.aclose()
+
+
+# Module-level registry of shared HTTPClient instances
+_http_clients: dict[tuple[Provider, Capability], HTTPClient] = {}
+
+
+def get_http_client(provider: Provider, capability: Capability) -> HTTPClient:
+    """Get or create shared HTTP client for provider and capability combination.
+
+    Args:
+        provider: The AI provider.
+        capability: The capability being used.
+
+    Returns:
+        Shared HTTPClient instance for this provider and capability.
+    """
+    key = (provider, capability)
+
+    if key not in _http_clients:
+        _http_clients[key] = HTTPClient()
+
+    return _http_clients[key]
+
+
+async def close_all_http_clients() -> None:
+    """Close all HTTP clients gracefully and clear registry."""
+    for key, client in list(_http_clients.items()):
+        try:
+            await client.aclose()
+        except Exception as e:
+            logger.warning(f"Failed to close HTTP client for {key}: {e}")
+
+    _http_clients.clear()
+
+
+def clear_http_clients() -> None:
+    """Clear HTTP client registry without closing connections."""
+    _http_clients.clear()
+
+
+__all__ = [
+    "DEFAULT_TIMEOUT",
+    "MAX_CONNECTIONS",
+    "MAX_KEEPALIVE_CONNECTIONS",
+    "HTTPClient",
+    "clear_http_clients",
+    "close_all_http_clients",
+    "get_http_client",
+]

celeste/io.py

@@ -0,0 +1,43 @@
+"""Input and output types for generation operations."""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class Input(BaseModel):
+    """Base class for capability-specific input types."""
+
+    pass
+
+
+class FinishReason(BaseModel):
+    """Base class for capability-specific finish reasons (used in streaming chunks)."""
+
+    pass
+
+
+class Usage(BaseModel):
+    """Base class for capability-specific usage metrics."""
+
+    pass
+
+
+class Output[Content](BaseModel):
+    """Base output class with generic content type."""
+
+    content: Content
+    usage: Usage = Field(default_factory=Usage)
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class Chunk[Content](BaseModel):
+    """Incremental chunk from streaming response with generic content type."""
+
+    content: Content
+    finish_reason: FinishReason | None = None
+    usage: Usage | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+__all__ = ["Chunk", "FinishReason", "Input", "Output", "Usage"]

celeste/mime_types.py

@@ -0,0 +1,47 @@
+"""MIME type enumerations for Celeste."""
+
+from enum import StrEnum
+
+
+class MimeType(StrEnum):
+    """Base class for all MIME types."""
+
+    pass
+
+
+class ImageMimeType(MimeType):
+    """Standard MIME types for images."""
+
+    PNG = "image/png"
+    JPEG = "image/jpeg"
+    WEBP = "image/webp"
+
+
+class VideoMimeType(MimeType):
+    """Standard MIME types for videos."""
+
+    MP4 = "video/mp4"
+    AVI = "video/x-msvideo"
+    MOV = "video/quicktime"
+
+
+class AudioMimeType(MimeType):
+    """Standard MIME types for audio."""
+
+    MP3 = "audio/mpeg"
+    WAV = "audio/wav"
+    OGG = "audio/ogg"
+    WEBM = "audio/webm"
+    AAC = "audio/aac"
+    FLAC = "audio/flac"
+    AIFF = "audio/aiff"
+    M4A = "audio/mp4"
+    WMA = "audio/x-ms-wma"
+
+
+__all__ = [
+    "AudioMimeType",
+    "ImageMimeType",
+    "MimeType",
+    "VideoMimeType",
+]

celeste/models.py

@@ -0,0 +1,91 @@
+"""Models and model registry for Celeste."""
+
+from pydantic import BaseModel, Field
+
+from celeste.constraints import Constraint
+from celeste.core import Capability, Provider
+
+
+class Model(BaseModel):
+    """Represents an AI model with its capabilities and metadata."""
+
+    id: str
+    provider: Provider
+    capabilities: set[Capability] = Field(default_factory=set)
+    display_name: str
+    parameter_constraints: dict[str, Constraint] = Field(default_factory=dict)
+
+    @property
+    def supported_parameters(self) -> set[str]:
+        """Compute supported parameter names from parameter_constraints."""
+        return set(self.parameter_constraints.keys())
+
+
+# Module-level registry mapping (model_id, provider) to model
+_models: dict[tuple[str, Provider], Model] = {}
+
+
+def register_models(models: Model | list[Model]) -> None:
+    """Register one or more models in the global registry.
+
+    Args:
+        models: Single Model instance or list of Models to register.
+               Each model is indexed by (model_id, provider) tuple.
+
+    Raises:
+        ValueError: If a model with the same (id, provider) is already registered.
+    """
+    if isinstance(models, Model):
+        models = [models]
+
+    for model in models:
+        key = (model.id, model.provider)
+        if key in _models:
+            msg = f"Model '{model.id}' for provider {model.provider.value} is already registered"
+            raise ValueError(msg)
+        _models[key] = model
+
+
+def get_model(model_id: str, provider: Provider) -> Model | None:
+    """Get a registered model by ID and provider.
+
+    Args:
+        model_id: The model identifier.
+        provider: The provider that owns the model.
+
+    Returns:
+        Model instance if found, None otherwise.
+    """
+    return _models.get((model_id, provider))
+
+
+def list_models(
+    provider: Provider | None = None,
+    capability: Capability | None = None,
+) -> list[Model]:
+    """List all registered models, optionally filtered by provider and/or capability.
+
+    Args:
+        provider: Optional provider filter. If provided, only models from this provider are returned.
+        capability: Optional capability filter. If provided, only models supporting this capability are returned.
+
+    Returns:
+        List of Model instances matching the filters.
+    """
+    filtered = list(_models.values())
+
+    if provider is not None:
+        filtered = [m for m in filtered if m.provider == provider]
+
+    if capability is not None:
+        filtered = [m for m in filtered if capability in m.capabilities]
+
+    return filtered
+
+
+def clear() -> None:
+    """Clear all registered models from the registry."""
+    _models.clear()
+
+
+__all__ = ["Model", "clear", "get_model", "list_models", "register_models"]

celeste/parameters.py

@@ -0,0 +1,51 @@
+"""Parameter system for Celeste."""
+
+from abc import ABC, abstractmethod
+from enum import StrEnum
+from typing import Any, TypedDict
+
+from celeste.models import Model
+
+
+class Parameters(TypedDict, total=False):
+    """Base parameters for all capabilities."""
+
+
+class ParameterMapper(ABC):
+    """Base class for provider-specific parameter transformation."""
+
+    name: StrEnum
+    """Parameter name matching capability TypedDict key. Must be StrEnum for type safety."""
+
+    @abstractmethod
+    def map(self, request: dict[str, Any], value: Any, model: Model) -> dict[str, Any]:  # noqa: ANN401
+        """Transform parameter value into provider's request structure.
+
+        Args:
+            request: Provider request dict.
+            value: Parameter value.
+            model: Model instance containing parameter_constraints for validation.
+
+        Returns:
+            Updated request dict.
+        """
+        ...
+
+    def parse_output(self, content: object, value: object | None) -> object:
+        """Optionally transform parsed content based on parameter value (default: return unchanged)."""
+        return content
+
+    def _validate_value(self, value: Any, model: Model) -> Any:  # noqa: ANN401
+        """Validate parameter value using model constraint, raising ValueError if no constraint exists."""
+        if value is None:
+            return None
+
+        constraint = model.parameter_constraints.get(self.name)
+        if constraint is None:
+            msg = f"Parameter {self.name.value} is not supported by model {model.id}"
+            raise ValueError(msg)
+
+        return constraint(value)
+
+
+__all__ = ["ParameterMapper", "Parameters"]

celeste/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561 - this package supports type checking

celeste/streaming.py

@@ -0,0 +1,114 @@
+"""Streaming support for Celeste."""
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from types import TracebackType
+from typing import Any, Self
+
+from celeste.io import Chunk, Output
+
+
+class Stream[Out: Output](ABC):
+    """Async iterator wrapper providing final Output access after stream exhaustion."""
+
+    def __init__(
+        self,
+        sse_iterator: AsyncIterator[dict[str, Any]],
+    ) -> None:
+        """Initialize stream with SSE iterator."""
+        self._sse_iterator = sse_iterator
+        self._chunks: list[Chunk] = []
+        self._closed = False
+        self._output: Out | None = None
+
+    @abstractmethod
+    def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
+        """Parse SSE event into Chunk (returns None to filter lifecycle events)."""
+        ...
+
+    @abstractmethod
+    def _parse_output(self, chunks: list[Chunk]) -> Out:
+        """Parse final Output from accumulated chunks."""
+        ...
+
+    def __repr__(self) -> str:
+        """Developer-friendly representation showing stream state."""
+        if self._output:
+            state = "done"
+        elif self._closed:
+            state = "closed"
+        elif self._chunks:
+            state = "streaming"
+        else:
+            state = "idle"
+
+        chunks = f", {len(self._chunks)} chunks" if self._chunks else ""
+        return f"<{self.__class__.__name__}: {state}{chunks}>"
+
+    # AsyncIterator protocol
+    def __aiter__(self) -> Self:
+        """Return self as async iterator."""
+        return self
+
+    async def __anext__(self) -> Chunk:
+        """Yield next chunk from stream."""
+        if self._closed:
+            raise StopAsyncIteration
+
+        try:
+            async for event in self._sse_iterator:
+                chunk = self._parse_chunk(event)
+                if chunk is not None:
+                    self._chunks.append(chunk)
+                    return chunk
+
+            # Stream exhausted - validate and parse final output
+            if not self._chunks:
+                msg = "Stream completed but no chunks were produced"
+                raise RuntimeError(msg)
+
+            self._output = self._parse_output(self._chunks)
+        except Exception:
+            await self.aclose()
+            raise
+
+        # Only reached on successful exhaustion
+        await self.aclose()
+        raise StopAsyncIteration
+
+    # AsyncContextManager protocol
+    async def __aenter__(self) -> Self:
+        """Enter context - return self for iteration."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        """Exit context - ensure cleanup even on exception."""
+        await self.aclose()
+        return False  # Propagate exceptions
+
+    @property
+    def output(self) -> Out:
+        """Access final Output after stream exhaustion (raises RuntimeError if not ready)."""
+        if self._output is None:
+            msg = "Stream not exhausted. Consume all chunks before accessing .output"
+            raise RuntimeError(msg)
+        return self._output
+
+    async def aclose(self) -> None:
+        """Explicitly close stream and cleanup resources."""
+        if self._closed:
+            return
+
+        self._closed = True
+
+        # Close SSE iterator (httpx-sse connection)
+        if hasattr(self._sse_iterator, "aclose"):
+            await self._sse_iterator.aclose()
+
+
+__all__ = ["Stream"]