lumera 0.4.6__py3-none-any.whl → 0.9.6__py3-none-any.whl

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)

lumera/locks.py

@@ -0,0 +1,216 @@
+"""
+Lock management for preventing concurrent operations.
+
+Provides two types of locks:
+1. Record-level locks: Lock specific records (uses platform lm_locks table)
+2. Operation-level locks: Lock entire operations globally (requires custom collection)
+
+Available functions:
+    claim_record_locks()    - Lock specific records for processing
+    release_record_locks()  - Release previously claimed record locks
+    acquire_operation_lock() - Lock an entire operation (NOT YET IMPLEMENTED)
+    release_operation_lock() - Release operation lock (NOT YET IMPLEMENTED)
+    operation_lock()        - Context manager for operation locks (NOT YET IMPLEMENTED)
+
+Example:
+    >>> from lumera import locks
+    >>> result = locks.claim_record_locks("export", "deposits", ["dep_1", "dep_2"])
+    >>> for id in result["claimed"]:
+    ...     process(id)
+    >>> locks.release_record_locks("export", record_ids=result["claimed"])
+"""
+
+__all__ = [
+    "claim_record_locks",
+    "release_record_locks",
+    "acquire_operation_lock",
+    "release_operation_lock",
+    "operation_lock",
+]
+
+from contextlib import contextmanager
+from typing import Any, Iterator
+
+# Import platform lock primitives from the main SDK module
+from .sdk import claim_locks as _claim_locks
+from .sdk import release_locks as _release_locks
+
+
+def claim_record_locks(
+    job_type: str,
+    collection: str,
+    record_ids: list[str],
+    *,
+    ttl_seconds: int = 900,
+    job_id: str | None = None,
+) -> dict[str, Any]:
+    """Claim record-level locks (using platform lm_locks).
+
+    Prevents multiple workers from processing the same records concurrently.
+    Uses the platform's built-in lm_locks table.
+
+    Args:
+        job_type: Workflow name (e.g., "deposit_processing")
+        collection: Collection name
+        record_ids: List of record IDs to lock
+        ttl_seconds: Lock duration in seconds (default 900 = 15 minutes)
+        job_id: Optional job identifier for grouping locks
+
+    Returns:
+        {
+            "claimed": ["id1", "id2"],  # Successfully locked
+            "skipped": ["id3"],         # Already locked by another process
+            "ttl_seconds": 900
+        }
+
+    Example:
+        >>> result = claim_record_locks(
+        ...     job_type="export",
+        ...     collection="deposits",
+        ...     record_ids=["dep_1", "dep_2", "dep_3"]
+        ... )
+        >>> for dep_id in result["claimed"]:
+        ...     process(dep_id)
+        >>> # Release when done
+        >>> release_record_locks("export", record_ids=result["claimed"])
+    """
+    return _claim_locks(
+        job_type=job_type,
+        collection=collection,
+        record_ids=record_ids,
+        ttl_seconds=ttl_seconds,
+        job_id=job_id,
+    )
+
+
+def release_record_locks(
+    job_type: str,
+    *,
+    collection: str | None = None,
+    record_ids: list[str] | None = None,
+    job_id: str | None = None,
+) -> int:
+    """Release record-level locks.
+
+    Args:
+        job_type: Workflow name (required)
+        collection: Optional collection filter
+        record_ids: Optional specific records to release
+        job_id: Optional job identifier filter
+
+    Returns:
+        Number of locks released
+
+    Example:
+        >>> released = release_record_locks(
+        ...     job_type="export",
+        ...     record_ids=["dep_1", "dep_2"]
+        ... )
+        >>> print(f"Released {released} locks")
+    """
+    return _release_locks(
+        job_type=job_type, collection=collection, record_ids=record_ids, job_id=job_id
+    )
+
+
+# Operation-level locks (simple key-value locks)
+# Note: These would need a custom collection like "export_locks" to be implemented
+# For now, providing the interface that should be implemented
+
+
+def acquire_operation_lock(
+    lock_name: str, *, ttl_seconds: int = 600, wait: bool = False, wait_timeout: int = 30
+) -> bool:
+    """Acquire an operation-level lock.
+
+    For preventing concurrent execution of entire operations (like exports).
+    Uses a simple key-value lock, not tied to specific records.
+
+    Note: This requires a custom locks collection to be created.
+    See the Charter Impact export_locks collection as an example.
+
+    Args:
+        lock_name: Unique lock identifier (e.g., "csv_export")
+        ttl_seconds: Lock duration in seconds (default 600 = 10 minutes)
+        wait: If True, wait for lock to become available
+        wait_timeout: Max seconds to wait (if wait=True)
+
+    Returns:
+        True if lock acquired, False if already held (when wait=False)
+
+    Raises:
+        TimeoutError: If wait=True and timeout exceeded
+        NotImplementedError: If operation locks collection doesn't exist
+
+    Example:
+        >>> if acquire_operation_lock("csv_export"):
+        ...     try:
+        ...         perform_export()
+        ...     finally:
+        ...         release_operation_lock("csv_export")
+        ... else:
+        ...     print("Export already in progress")
+    """
+    raise NotImplementedError(
+        "Operation-level locks require a custom locks collection. "
+        "Create a collection like 'operation_locks' with fields: "
+        "lock_name (text, unique), held_by (text), acquired_at (date), expires_at (date). "
+        "Then implement acquire/release using pb.search/create/delete."
+    )
+
+
+def release_operation_lock(lock_name: str) -> bool:
+    """Release an operation-level lock.
+
+    Args:
+        lock_name: Lock identifier
+
+    Returns:
+        True if lock was released, False if wasn't held
+
+    Raises:
+        NotImplementedError: If operation locks collection doesn't exist
+
+    Example:
+        >>> release_operation_lock("csv_export")
+    """
+    raise NotImplementedError(
+        "Operation-level locks require a custom locks collection. "
+        "See acquire_operation_lock() for details."
+    )
+
+
+@contextmanager
+def operation_lock(
+    lock_name: str, *, ttl_seconds: int = 600, wait: bool = False, wait_timeout: int = 30
+) -> Iterator[None]:
+    """Context manager for operation locks.
+
+    Args:
+        lock_name: Lock identifier
+        ttl_seconds: Lock duration in seconds
+        wait: Wait for lock if held
+        wait_timeout: Max wait time in seconds
+
+    Raises:
+        NotImplementedError: If operation locks collection doesn't exist
+        TimeoutError: If wait=True and timeout exceeded
+
+    Example:
+        >>> with operation_lock("csv_export"):
+        ...     perform_export()
+        ...     # Lock automatically released on exit
+    """
+    # This would acquire the lock
+    acquired = acquire_operation_lock(
+        lock_name, ttl_seconds=ttl_seconds, wait=wait, wait_timeout=wait_timeout
+    )
+
+    if not acquired:
+        raise RuntimeError(f"Failed to acquire lock: {lock_name}")
+
+    try:
+        yield
+    finally:
+        # Always release the lock
+        release_operation_lock(lock_name)