lumera 0.4.22__py3-none-any.whl → 0.5.0__py3-none-any.whl

lumera/__init__.py

@@ -5,6 +5,10 @@ This SDK provides helpers for agents running within the Lumera Notebook environm
 to interact with the Lumera API and define dynamic user interfaces.
 """
 
+__version__ = "0.5.0"
+
+# Import new modules (as modules, not individual functions)
+from . import exceptions, llm, locks, pb, storage
 from ._utils import (
     LumeraAPIError,
     RecordNotUniqueError,
@@ -13,6 +17,15 @@ from ._utils import (
     log_timed,
 )
 
+# Import specific exceptions for convenience
+from .exceptions import (
+    LockHeldError,
+    LumeraError,
+    RecordNotFoundError,
+    UniqueConstraintError,
+    ValidationError,
+)
+
 # Import key SDK helpers to expose them at the package root.
 from .sdk import (
     CollectionField,
@@ -39,29 +52,46 @@ from .sdk import (
 
 # Define what `from lumera import *` imports.
 __all__ = [
+    # Authentication & utilities
     "get_access_token",
-    "save_to_lumera",
     "get_google_access_token",  # Kept for backwards compatibility
     "log_timed",
+    # Collections (low-level API)
     "list_collections",
     "get_collection",
     "create_collection",
     "update_collection",
     "delete_collection",
+    # Records (low-level API)
     "list_records",
     "get_record",
     "get_record_by_external_id",
+    "create_record",
+    "update_record",
+    "upsert_record",
+    "delete_record",
+    # Other operations
     "replay_hook",
     "query_sql",
     "run_agent",
     "get_agent_run",
     "upload_lumera_file",
-    "create_record",
-    "update_record",
-    "upsert_record",
-    "delete_record",
+    "save_to_lumera",
+    # Type definitions
     "CollectionField",
     "HookReplayResult",
+    # Exceptions
     "LumeraAPIError",
     "RecordNotUniqueError",
+    "LumeraError",
+    "ValidationError",
+    "UniqueConstraintError",
+    "RecordNotFoundError",
+    "LockHeldError",
+    # New modules (use as lumera.pb, lumera.storage, etc.)
+    "pb",
+    "storage",
+    "llm",
+    "locks",
+    "exceptions",
 ]

lumera/exceptions.py

@@ -0,0 +1,72 @@
+"""
+Custom exceptions for the Lumera SDK.
+
+Exception hierarchy:
+    LumeraError (base)
+    ├── RecordNotFoundError  - Record doesn't exist (404)
+    ├── ValidationError      - Data doesn't match schema
+    ├── UniqueConstraintError - Unique field violation
+    └── LockHeldError        - Lock already held by another process
+
+Example:
+    >>> from lumera import pb
+    >>> from lumera.exceptions import RecordNotFoundError
+    >>> try:
+    ...     deposit = pb.get("deposits", "invalid_id")
+    ... except RecordNotFoundError as e:
+    ...     print(f"Not found: {e.record_id}")
+"""
+
+__all__ = [
+    "LumeraError",
+    "RecordNotFoundError",
+    "ValidationError",
+    "UniqueConstraintError",
+    "LockHeldError",
+]
+
+
+class LumeraError(Exception):
+    """Base exception for all Lumera SDK errors."""
+
+    pass
+
+
+class RecordNotFoundError(LumeraError):
+    """Record doesn't exist in the collection."""
+
+    def __init__(self, collection: str, record_id: str) -> None:
+        super().__init__(f"Record '{record_id}' not found in collection '{collection}'")
+        self.collection = collection
+        self.record_id = record_id
+
+
+class ValidationError(LumeraError):
+    """Data doesn't match collection schema."""
+
+    def __init__(self, collection: str, errors: dict[str, str]) -> None:
+        super().__init__(f"Validation failed for '{collection}': {errors}")
+        self.collection = collection
+        self.errors = errors
+
+
+class UniqueConstraintError(LumeraError):
+    """Unique field constraint violation."""
+
+    def __init__(self, collection: str, field: str, value: object) -> None:
+        super().__init__(f"Record with {field}='{value}' already exists in '{collection}'")
+        self.collection = collection
+        self.field = field
+        self.value = value
+
+
+class LockHeldError(LumeraError):
+    """Lock is already held by another process."""
+
+    def __init__(self, lock_name: str, held_by: str | None = None) -> None:
+        msg = f"Lock '{lock_name}' is already held"
+        if held_by:
+            msg += f" by {held_by}"
+        super().__init__(msg)
+        self.lock_name = lock_name
+        self.held_by = held_by

lumera/llm.py

@@ -0,0 +1,481 @@
+"""
+LLM operations for AI completions and embeddings.
+
+This module provides a unified interface for LLM operations with pluggable
+provider support. Currently implements OpenAI, with extensibility for other
+providers (Anthropic, Google, etc.) in the future.
+
+Available functions:
+    complete()  - Single-turn LLM completion with prompt
+    chat()      - Multi-turn chat completion with message history
+    embed()     - Generate embeddings for text (single or batch)
+
+Configuration:
+    OPENAI_API_KEY     - Required for OpenAI provider
+    LUMERA_LLM_PROVIDER - Provider to use (default: "openai")
+
+Example:
+    >>> from lumera import llm
+    >>> response = llm.complete("What is 2+2?", model="gpt-5.2-mini")
+    >>> print(response["content"])
+"""
+
+from __future__ import annotations
+
+import os
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Literal, NotRequired, TypedDict, Unpack
+
+if TYPE_CHECKING:
+    import openai
+
+__all__ = [
+    "complete",
+    "chat",
+    "embed",
+    "Message",
+    "LLMResponse",
+    "ProviderConfig",
+    "LLMProvider",
+    "get_provider",
+    "set_provider",
+]
+
+
+# ---------------------------------------------------------------------------
+# Type definitions
+# ---------------------------------------------------------------------------
+
+
+class Message(TypedDict):
+    """Chat message format compatible with OpenAI and other providers."""
+
+    role: Literal["system", "user", "assistant"]
+    content: str
+
+
+class LLMResponse(TypedDict, total=False):
+    """LLM completion response."""
+
+    content: str  # Response text (always present)
+    model: str  # Model used
+    usage: dict[str, int]  # Token usage: prompt_tokens, completion_tokens, total_tokens
+    finish_reason: str  # "stop", "length", "content_filter", etc.
+    provider: str  # Provider name (e.g., "openai", "anthropic")
+
+
+class ProviderConfig(TypedDict, total=False):
+    """Configuration options for LLM providers."""
+
+    api_key: NotRequired[str]  # API key (overrides Lumera/env lookup)
+    provider_name: NotRequired[str]  # Provider name for get_access_token (default: "openai")
+
+
+# ---------------------------------------------------------------------------
+# Provider interface (for future extensibility)
+# ---------------------------------------------------------------------------
+
+
+class LLMProvider(ABC):
+    """Abstract base class for LLM providers.
+
+    Subclass this to add support for new providers (Anthropic, Google, etc.).
+    """
+
+    name: str = "base"
+
+    @abstractmethod
+    def complete(
+        self,
+        prompt: str,
+        *,
+        model: str,
+        temperature: float,
+        max_tokens: int | None,
+        system_prompt: str | None,
+        json_mode: bool,
+    ) -> LLMResponse:
+        """Generate a completion for a single prompt."""
+        ...
+
+    @abstractmethod
+    def chat(
+        self,
+        messages: list[Message],
+        *,
+        model: str,
+        temperature: float,
+        max_tokens: int | None,
+        json_mode: bool,
+    ) -> LLMResponse:
+        """Generate a chat completion from message history."""
+        ...
+
+    @abstractmethod
+    def embed(
+        self,
+        text: str | list[str],
+        *,
+        model: str,
+    ) -> list[float] | list[list[float]]:
+        """Generate embeddings for text."""
+        ...
+
+
+# ---------------------------------------------------------------------------
+# OpenAI provider implementation
+# ---------------------------------------------------------------------------
+
+
+class OpenAIProvider(LLMProvider):
+    """OpenAI provider implementation using the openai Python SDK."""
+
+    name = "openai"
+
+    # Model aliases for convenience
+    MODEL_ALIASES: dict[str, str] = {
+        "gpt-5.2": "gpt-5.2",
+        "gpt-5.2-mini": "gpt-5.2-mini",
+        "gpt-5.2-nano": "gpt-5.2-nano",
+        # Embedding models
+        "text-embedding-3-small": "text-embedding-3-small",
+        "text-embedding-3-large": "text-embedding-3-large",
+    }
+
+    DEFAULT_CHAT_MODEL = "gpt-5.2-mini"
+    DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small"
+    DEFAULT_PROVIDER_NAME = "openai"
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        provider_name: str | None = None,
+    ) -> None:
+        """Initialize OpenAI provider.
+
+        Args:
+            api_key: OpenAI API key. If not provided, fetches from Lumera
+                     using get_access_token(provider_name), or falls back
+                     to OPENAI_API_KEY env var.
+            provider_name: Provider name for get_access_token lookup.
+                          Defaults to "openai".
+        """
+        self._explicit_api_key = api_key
+        self._provider_name = provider_name or self.DEFAULT_PROVIDER_NAME
+        self._client: openai.OpenAI | None = None  # noqa: F821
+
+    def _get_api_key(self) -> str:
+        """Get API key from explicit config, Lumera, or environment."""
+        # 1. Use explicitly provided key
+        if self._explicit_api_key:
+            return self._explicit_api_key
+
+        # 2. Try to fetch from Lumera platform
+        try:
+            from ._utils import get_access_token
+
+            return get_access_token(self._provider_name)
+        except Exception:
+            pass  # Fall through to env var
+
+        # 3. Fall back to environment variable
+        env_key = os.environ.get("OPENAI_API_KEY")
+        if env_key:
+            return env_key
+
+        raise ValueError(
+            "OpenAI API key not configured. Either:\n"
+            f"  1. Configure '{self._provider_name}' provider in Lumera platform\n"
+            "  2. Set OPENAI_API_KEY environment variable\n"
+            "  3. Pass api_key to set_provider()"
+        )
+
+    @property
+    def client(self) -> openai.OpenAI:  # noqa: F821
+        """Lazy-initialize OpenAI client."""
+        if self._client is None:
+            try:
+                import openai
+            except ImportError as e:
+                raise ImportError(
+                    "OpenAI package not installed. Install with: pip install 'lumera[full]'"
+                ) from e
+
+            api_key = self._get_api_key()
+            self._client = openai.OpenAI(api_key=api_key)
+        return self._client
+
+    def _resolve_model(self, model: str) -> str:
+        """Resolve model alias to actual model name."""
+        return self.MODEL_ALIASES.get(model, model)
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        model: str,
+        temperature: float,
+        max_tokens: int | None,
+        system_prompt: str | None,
+        json_mode: bool,
+    ) -> LLMResponse:
+        """Generate a completion using OpenAI."""
+        messages: list[Message] = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+
+        return self.chat(
+            messages,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            json_mode=json_mode,
+        )
+
+    def chat(
+        self,
+        messages: list[Message],
+        *,
+        model: str,
+        temperature: float,
+        max_tokens: int | None,
+        json_mode: bool,
+    ) -> LLMResponse:
+        """Generate a chat completion using OpenAI."""
+        resolved_model = self._resolve_model(model)
+
+        # Build request kwargs
+        kwargs: dict = {
+            "model": resolved_model,
+            "messages": messages,  # type: ignore[arg-type]
+            "temperature": temperature,
+        }
+
+        if max_tokens is not None:
+            kwargs["max_tokens"] = max_tokens
+
+        if json_mode:
+            kwargs["response_format"] = {"type": "json_object"}
+
+        # Make API call
+        response = self.client.chat.completions.create(**kwargs)
+
+        # Extract response
+        choice = response.choices[0]
+        content = choice.message.content or ""
+
+        result: LLMResponse = {
+            "content": content,
+            "model": response.model,
+            "provider": self.name,
+        }
+
+        if choice.finish_reason:
+            result["finish_reason"] = choice.finish_reason
+
+        if response.usage:
+            result["usage"] = {
+                "prompt_tokens": response.usage.prompt_tokens,
+                "completion_tokens": response.usage.completion_tokens,
+                "total_tokens": response.usage.total_tokens,
+            }
+
+        return result
+
+    def embed(
+        self,
+        text: str | list[str],
+        *,
+        model: str,
+    ) -> list[float] | list[list[float]]:
+        """Generate embeddings using OpenAI."""
+        resolved_model = self._resolve_model(model)
+
+        # Normalize input to list
+        input_texts = [text] if isinstance(text, str) else text
+
+        response = self.client.embeddings.create(
+            model=resolved_model,
+            input=input_texts,
+        )
+
+        # Extract embeddings
+        embeddings = [item.embedding for item in response.data]
+
+        # Return single embedding if single input
+        if isinstance(text, str):
+            return embeddings[0]
+        return embeddings
+
+
+# ---------------------------------------------------------------------------
+# Provider registry and module-level state
+# ---------------------------------------------------------------------------
+
+# Registry of available providers
+_PROVIDERS: dict[str, type[LLMProvider]] = {
+    "openai": OpenAIProvider,
+}
+
+# Current active provider instance
+_current_provider: LLMProvider | None = None
+
+
+def get_provider() -> LLMProvider:
+    """Get the current LLM provider instance.
+
+    Returns the configured provider, initializing it if necessary.
+    Provider is determined by LUMERA_LLM_PROVIDER env var (default: "openai").
+    """
+    global _current_provider
+
+    if _current_provider is None:
+        provider_name = os.environ.get("LUMERA_LLM_PROVIDER", "openai").lower()
+
+        if provider_name not in _PROVIDERS:
+            available = ", ".join(_PROVIDERS.keys())
+            raise ValueError(f"Unknown LLM provider: {provider_name}. Available: {available}")
+
+        provider_class = _PROVIDERS[provider_name]
+        _current_provider = provider_class()
+
+    return _current_provider
+
+
+def set_provider(provider: LLMProvider | str, **kwargs: Unpack[ProviderConfig]) -> None:
+    """Set the active LLM provider.
+
+    Args:
+        provider: Either a provider instance or provider name string.
+        **kwargs: If provider is a string, kwargs are passed to provider constructor.
+
+    Example:
+        >>> llm.set_provider("openai", api_key="sk-...")
+        >>> # Or with a custom provider instance
+        >>> llm.set_provider(MyCustomProvider())
+    """
+    global _current_provider
+
+    if isinstance(provider, str):
+        if provider not in _PROVIDERS:
+            available = ", ".join(_PROVIDERS.keys())
+            raise ValueError(f"Unknown provider: {provider}. Available: {available}")
+        _current_provider = _PROVIDERS[provider](**kwargs)
+    else:
+        _current_provider = provider
+
+
+# ---------------------------------------------------------------------------
+# Public API functions
+# ---------------------------------------------------------------------------
+
+
+def complete(
+    prompt: str,
+    *,
+    model: str = "gpt-5.2-mini",
+    temperature: float = 0.7,
+    max_tokens: int | None = None,
+    system_prompt: str | None = None,
+    json_mode: bool = False,
+) -> LLMResponse:
+    """Get LLM completion for a prompt.
+
+    Args:
+        prompt: User prompt/question
+        model: Model to use (default: gpt-5.2-mini)
+        temperature: Sampling temperature 0.0 to 2.0 (default: 0.7)
+        max_tokens: Max tokens in response (None = model default)
+        system_prompt: Optional system message to set behavior
+        json_mode: Force JSON output (default: False)
+
+    Returns:
+        LLM response with content and metadata
+
+    Example:
+        >>> response = llm.complete(
+        ...     prompt="Classify this deposit: ...",
+        ...     system_prompt="You are an expert accountant.",
+        ...     model="gpt-5.2-mini",
+        ...     json_mode=True
+        ... )
+        >>> data = json.loads(response["content"])
+    """
+    provider = get_provider()
+    return provider.complete(
+        prompt,
+        model=model,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        system_prompt=system_prompt,
+        json_mode=json_mode,
+    )
+
+
+def chat(
+    messages: list[Message],
+    *,
+    model: str = "gpt-5.2-mini",
+    temperature: float = 0.7,
+    max_tokens: int | None = None,
+    json_mode: bool = False,
+) -> LLMResponse:
+    """Multi-turn chat completion.
+
+    Args:
+        messages: Conversation history with role and content
+        model: Model to use (default: gpt-5.2-mini)
+        temperature: Sampling temperature 0.0 to 2.0 (default: 0.7)
+        max_tokens: Max tokens in response (None = model default)
+        json_mode: Force JSON output (default: False)
+
+    Returns:
+        LLM response with assistant's message
+
+    Example:
+        >>> response = llm.chat([
+        ...     {"role": "system", "content": "You are a helpful assistant."},
+        ...     {"role": "user", "content": "What is 2+2?"},
+        ...     {"role": "assistant", "content": "4"},
+        ...     {"role": "user", "content": "What about 3+3?"}
+        ... ])
+        >>> print(response["content"])
+    """
+    provider = get_provider()
+    return provider.chat(
+        messages,
+        model=model,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        json_mode=json_mode,
+    )
+
+
+def embed(
+    text: str | list[str],
+    *,
+    model: str = "text-embedding-3-small",
+) -> list[float] | list[list[float]]:
+    """Generate embeddings for text.
+
+    Args:
+        text: Single string or list of strings to embed
+        model: Embedding model (default: text-embedding-3-small)
+
+    Returns:
+        Embedding vector (for single string) or list of vectors (for list)
+
+    Example:
+        >>> embedding = llm.embed("deposit payment notice")
+        >>> # Use for similarity search, semantic matching, etc.
+        >>>
+        >>> # Batch embeddings
+        >>> embeddings = llm.embed([
+        ...     "payment notice",
+        ...     "direct deposit",
+        ...     "apportionment"
+        ... ])
+    """
+    provider = get_provider()
+    return provider.embed(text, model=model)