aix 0.0.23__py3-none-any.whl → 0.0.25__py3-none-any.whl

aix/__init__.py

@@ -1,68 +1,159 @@
-"""
-Facade to key AI tools.
+"""AIX: Artificial Intelligence eXtensions
 
-Get a list of available chat functions (this will depend on the AI packages
-you have installed locally):
+A clean, pythonic facade for common AI operations that abstracts away
+provider-specific details and complexities.
 
->>> from aix import chat_funcs
->>> list(chat_funcs)  # doctest: +SKIP
-['gemini-1.5-flash',
- 'gpt-4',
- 'gpt-4-32k',
- 'gpt-4-turbo',
- 'gpt-3.5-turbo',
- 'o1-preview',
- 'o1-mini',
- 'gpt-4o',
- 'gpt-4o-mini']
+Quick Start:
+    >>> from aix import chat, embeddings, prompt_func, models
 
-Choose a chat function and chat with it:
->>> google_ai_chat = chat_funcs['gemini-1.5-flash']  # doctest: +SKIP
->>> google_ai_chat("What is the meaning of life? Respond with a number.")  # doctest: +SKIP
-'42'
->>> openai_chat = chat_funcs['gpt-3.5-turbo']  # doctest: +SKIP
->>> openai_chat("What is the meaning of life? Respond with a number.")  # doctest: +SKIP
-'42'
+    # Simple chat
+    >>> response = chat("What is 2+2?")  # doctest: +SKIP
+    'The answer is 4.'
 
-"""
+    # Create prompt-based functions
+    >>> translate = prompt_func("Translate to French: {text}")
+    >>> translate(text="Hello world")  # doctest: +SKIP
+    'Bonjour le monde'
 
-from aix.gen_ai import chat, chat_models, chat_funcs
-from aix import contexts
-from aix.contexts import (
-    bytes_to_markdown,  # Convert bytes to markdown (with plugin support for different types of files)
-    bytes_store_to_markdown_store,  # Convert a bytes store to a markdown store based on extensions
-    aggregate_store,  # Aggregate a store into a single value (say, a string)
-)
+    # Get embeddings
+    >>> vecs = list(embeddings(["hello", "world"]))  # doctest: +SKIP
+    >>> len(vecs)  # doctest: +SKIP
+    2
+
+    # Discover models
+    >>> models.discover()  # doctest: +SKIP
+    >>> list(models)[:5]  # doctest: +SKIP
+    ['openai/gpt-4o', 'openai/gpt-4o-mini', ...]
 
-# TODO: Change this so that there's a load_pkg function that loads the packages dynamically
-#   if and when use wants.
+Main Features:
+    - chat(): Simple chat interface across providers
+    - embeddings(): Vector embeddings for text
+    - prompt_func(): Create functions from prompt templates
+    - models: Model discovery and selection
+    - generate_image(): Text-to-image generation
+    - text_to_speech(), transcribe(): Audio operations
+    - generate_video(): Text-to-video generation (provider-dependent)
+    - Batch operations for efficiency
+    - Clean, i2mint-style Mapping interfaces
+
+Backends:
+    - Uses LiteLLM for provider interactions
+    - Supports OpenAI, Anthropic, Google, and 100+ models
+    - OpenRouter integration for multi-provider access
+
+For detailed documentation, see: https://github.com/thorwhalen/aix
+"""
 
-# from aix.pd import *
-# from aix.np import *
-# from aix.sk import *
+# Core interfaces (new clean API)
+from aix.chat import chat, ask, chat_with_history, ChatSession
+from aix.embeddings import (
+    embeddings,
+    embed,
+    cosine_similarity,
+    find_most_similar,
+    EmbeddingCache,
+)
+from aix.prompts import (
+    prompt_func,
+    prompt_to_text,
+    prompt_to_json,
+    PromptFuncs,
+    common_funcs,
+    constrained_answer,
+)
+from aix.models import (
+    models,
+    ModelStore,
+    discover_available_models,
+    get_model_info,
+    find_models,
+)
+from aix.batches import (
+    batch_chat,
+    batch_embeddings,
+    batch_process,
+    BatchProcessor,
+)
+from aix.image import (
+    generate_image,
+    generate_images,
+    edit_image,
+    create_variation,
+    GeneratedImage,
+)
+from aix.audio import (
+    text_to_speech,
+    transcribe,
+    transcribe_with_timestamps,
+    translate_audio,
+    GeneratedAudio,
+    TranscriptionResult,
+)
+from aix.video import (
+    generate_video,
+    animate_image as animate_image_to_video,
+    extend_video,
+    GeneratedVideo,
+    get_available_providers as get_video_providers,
+)
 
-# from aix import pd
-# from aix import np
-# from aix import sk
+# Legacy interfaces (for backward compatibility)
+from aix.gen_ai import chat_models, chat_funcs
 
+# Version info
+__version__ = "0.1.0"
 
-#
-# from contextlib import suppress
-#
-# preferred_order = ['sk', 'np', 'pd']
-#
-# with suppress(ModuleNotFoundError):
-#     from aix import sk
-#
-# with suppress(ModuleNotFoundError):
-#     from aix import np
-#
-# with suppress(ModuleNotFoundError):
-#     from aix import pd
-#
-# for _module_name in preferred_order[::-1]:
-#     print(f"------ {_module_name}")
-#     _module = __import__(f'aix.{_module_name}')
-#     for _name in filter(lambda x: not x.startswith('__'), dir(_module)):
-#         print(_name, _module)
-#         locals()[_name] = getattr(_module, _name)
+# Public API
+__all__ = [
+    # Core chat
+    "chat",
+    "ask",
+    "chat_with_history",
+    "ChatSession",
+    # Embeddings
+    "embeddings",
+    "embed",
+    "cosine_similarity",
+    "find_most_similar",
+    "EmbeddingCache",
+    # Prompts
+    "prompt_func",
+    "prompt_to_text",
+    "prompt_to_json",
+    "PromptFuncs",
+    "common_funcs",
+    "constrained_answer",
+    # Models
+    "models",
+    "ModelStore",
+    "discover_available_models",
+    "get_model_info",
+    "find_models",
+    # Batches
+    "batch_chat",
+    "batch_embeddings",
+    "batch_process",
+    "BatchProcessor",
+    # Image
+    "generate_image",
+    "generate_images",
+    "edit_image",
+    "create_variation",
+    "GeneratedImage",
+    # Audio
+    "text_to_speech",
+    "transcribe",
+    "transcribe_with_timestamps",
+    "translate_audio",
+    "GeneratedAudio",
+    "TranscriptionResult",
+    # Video
+    "generate_video",
+    "animate_image_to_video",
+    "extend_video",
+    "GeneratedVideo",
+    "get_video_providers",
+    # Legacy
+    "chat_models",
+    "chat_funcs",
+]

aix/ai_models/__init__.py

@@ -0,0 +1,136 @@
+"""AI Model Management Module.
+
+A unified interface for managing AI models across multiple providers.
+
+Basic usage:
+    >>> from aix.ai_models import get_manager
+    >>> manager = get_manager()
+    >>> _ = manager.discover_from_source("openrouter", auto_register=True, verbose=False)
+    >>> models = manager.list_models(provider="openai")
+
+Custom filtering:
+    >>> cheap_models = manager.list_models(
+    ...     custom_filter=lambda m: m.cost_per_token.get("input", 0) < 0.001
+    ... )
+
+Get connector-specific metadata:
+    >>> openai_meta = manager.get_connector_metadata("openai/gpt-4", "openai")
+    >>> # Use with: openai.ChatCompletion.create(**openai_meta, messages=[...])
+
+"""
+
+# Core types
+from aix.ai_models.base import (
+    Model,
+    ModelRegistry,
+    ModelSource,
+    Connector,
+    ConnectorRegistry,
+)
+
+# Concrete sources and connectors
+from aix.ai_models.sources import (
+    Connector,
+    OpenRouterSource,
+    OllamaSource,
+    ProviderAPISource,
+    OpenAIConnector,
+    OpenRouterConnector,
+    LangChainConnector,
+    OllamaConnector,
+    DSPyConnector,
+)
+
+# Main facade
+from aix.ai_models.manager import (
+    ModelManager,
+    get_manager,
+)
+
+# Version info
+__version__ = "0.1.0"
+
+# Public API
+__all__ = [
+    # Core types
+    "Model",
+    "ModelRegistry",
+    "ModelSource",
+    "Connector",
+    "ConnectorRegistry",
+    # Sources
+    "OpenRouterSource",
+    "OllamaSource",
+    "ProviderAPISource",
+    # Connectors
+    "OpenAIConnector",
+    "OpenRouterConnector",
+    "LangChainConnector",
+    "OllamaConnector",
+    "DSPyConnector",
+    # Main API
+    "ModelManager",
+    "get_manager",
+]
+
+
+# Convenience functions for common operations
+
+
+def list_available_models(
+    *,
+    provider: str | None = None,
+    is_local: bool | None = None,
+    storage_path: str | None = None,
+) -> list[Model]:
+    """Quick function to list available models without explicit manager.
+
+    >>> models = list_available_models(provider="openai")
+    >>> len(models) >= 0
+    True
+    """
+    manager = get_manager(storage_path=storage_path)
+    return manager.list_models(provider=provider, is_local=is_local)
+
+
+def discover_models(
+    source_name: str = "openrouter",
+    *,
+    storage_path: str | None = None,
+    auto_register: bool = True,
+    verbose: bool = True,
+) -> list[Model]:
+    """Quick function to discover models from a source.
+
+    >>> models = discover_models("openrouter", auto_register=False, verbose=False)
+    >>> len(models) > 0
+    True
+    """
+    manager = get_manager(storage_path=storage_path)
+    return manager.discover_from_source(
+        source_name, auto_register=auto_register, verbose=verbose
+    )
+
+
+def get_model_metadata(
+    model_id: str, connector_name: str, *, storage_path: str | None = None
+) -> dict:
+    """Quick function to get formatted metadata for a model.
+
+    >>> import tempfile, os, json
+    >>> if 'OPENROUTER_API_KEY' in os.environ:
+    ...     from aix.ai_models.manager import get_manager
+    ...     with tempfile.NamedTemporaryFile(mode='w+', suffix=".json", delete=False) as temp:
+    ...         storage_path = temp.name
+    ...         json.dump({'models': []}, temp)
+    ...     try:
+    ...         manager = get_manager(storage_path=storage_path)
+    ...         _ = manager.discover_from_source("openrouter", auto_register=True, verbose=False)
+    ...         assert 'openai/gpt-3.5-turbo' in manager.models
+    ...         metadata = get_model_metadata("openai/gpt-3.5-turbo", "openrouter", storage_path=storage_path)
+    ...         assert 'model' in metadata
+    ...     finally:
+    ...         os.remove(storage_path)
+    """
+    manager = get_manager(storage_path=storage_path)
+    return manager.get_connector_metadata(model_id, connector_name)

aix/ai_models/base.py

@@ -0,0 +1,270 @@
+"""Core types for AI model management.
+
+This module provides a unified interface for managing, discovering, and
+connecting to AI models across multiple providers and deployment methods.
+"""
+
+from dataclasses import dataclass, field, asdict, fields
+from typing import Any
+from collections.abc import Mapping, MutableMapping, Iterator, Callable, Iterable
+from abc import ABC, abstractmethod
+import json
+from pathlib import Path
+
+
+@dataclass
+class Model:
+    """Represents an AI model with its metadata.
+
+    >>> model = Model(
+    ...     id="gpt-4",
+    ...     provider="openai",
+    ...     context_size=8192,
+    ...     is_local=False
+    ... )
+    >>> model.id
+    'gpt-4'
+    """
+
+    id: str
+    provider: str
+    context_size: int | None = None
+    is_local: bool = False
+    capabilities: dict[str, Any] = field(default_factory=dict)
+    cost_per_token: dict[str, float] = field(default_factory=dict)
+    tags: set[str] = field(default_factory=set)
+    connector_metadata: dict[str, dict[str, Any]] = field(default_factory=dict)
+    custom_metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert model to dictionary representation."""
+        return asdict(self)
+
+    def matches_filter(self, **criteria) -> bool:
+        """Check if model matches given criteria.
+
+        >>> model = Model(id="gpt-4", provider="openai", is_local=False)
+        >>> model.matches_filter(provider="openai")
+        True
+        >>> model.matches_filter(is_local=True)
+        False
+        """
+        for key, value in criteria.items():
+            if not hasattr(self, key):
+                return False
+            if getattr(self, key) != value:
+                return False
+        return True
+
+    def __getitem__(self, key: str) -> Any:
+        """Get field value by name."""
+        try:
+            return getattr(self, key)
+        except AttributeError:
+            raise KeyError(key)
+
+    def __iter__(self) -> Iterator[str]:
+        """Iterate over field names."""
+        return iter(field.name for field in fields(self))
+
+    def __len__(self) -> int:
+        """Return number of fields."""
+        return len(fields(self))
+
+
+class ModelRegistry(MutableMapping[str, Model]):
+    """Registry for managing AI models using Mapping interface.
+
+    >>> registry = ModelRegistry()
+    >>> registry["gpt-4"] = Model(id="gpt-4", provider="openai")
+    >>> "gpt-4" in registry
+    True
+    >>> len(registry)
+    1
+    """
+
+    def __init__(self, *, storage_path: Path | None = None):
+        """Initialize registry with optional persistent storage."""
+        self._models: dict[str, Model] = {}
+        self._storage_path = storage_path
+        if storage_path and storage_path.exists():
+            self._load()
+
+    def __setitem__(self, model_id: str, model: Model) -> None:
+        """Add or update a model in the registry."""
+        self._models[model_id] = model
+        if self._storage_path:
+            self._save()
+
+    def __getitem__(
+        self, key: str | list[str] | Callable[[Model], bool]
+    ) -> Model | list[Model]:
+        """Get model(s) by ID, list of IDs, or filter function.
+
+        Supports:
+        - Single ID: registry["gpt-4"]
+        - Multiple IDs: registry[["gpt-4", "claude-3"]]
+        - Filter function: registry[lambda m: m.is_local]
+        """
+        if isinstance(key, str):
+            return self._models[key]
+        elif isinstance(key, list):
+            return [self._models[k] for k in key if k in self._models]
+        elif callable(key):
+            return [m for m in self._models.values() if key(m)]
+        else:
+            raise TypeError(f"Unsupported key type: {type(key)}")
+
+    def __delitem__(self, model_id: str) -> None:
+        """Remove a model from the registry."""
+        del self._models[model_id]
+        if self._storage_path:
+            self._save()
+
+    def __iter__(self) -> Iterator[str]:
+        """Iterate over model IDs."""
+        yield from self._models.keys()
+
+    def __len__(self) -> int:
+        """Return number of models in registry."""
+        return len(self._models)
+
+    def filter(
+        self,
+        *,
+        provider: str | None = None,
+        is_local: bool | None = None,
+        min_context_size: int | None = None,
+        max_context_size: int | None = None,
+        has_capabilities: Iterable[str] | None = None,
+        tags: Iterable[str] | None = None,
+        custom_filter: Callable[[Model], bool] | None = None,
+    ) -> list[Model]:
+        """Filter models by multiple criteria.
+
+        >>> registry = ModelRegistry()
+        >>> registry["gpt-4"] = Model(id="gpt-4", provider="openai", context_size=8192)
+        >>> registry["llama2"] = Model(id="llama2", provider="ollama", is_local=True)
+        >>> local_models = registry.filter(is_local=True)
+        >>> len(local_models)
+        1
+        """
+
+        def _matches(model: Model) -> bool:
+            if provider and model.provider != provider:
+                return False
+            if is_local is not None and model.is_local != is_local:
+                return False
+            if min_context_size and (
+                not model.context_size or model.context_size < min_context_size
+            ):
+                return False
+            if max_context_size and (
+                not model.context_size or model.context_size > max_context_size
+            ):
+                return False
+            if has_capabilities:
+                for cap in has_capabilities:
+                    if not model.capabilities.get(cap):
+                        return False
+            if tags and not model.tags.issuperset(tags):
+                return False
+            if custom_filter and not custom_filter(model):
+                return False
+            return True
+
+        return [m for m in self._models.values() if _matches(m)]
+
+    def _load(self) -> None:
+        """Load models from persistent storage."""
+        if not self._storage_path:
+            return
+
+        with open(self._storage_path) as f:
+            data = json.load(f)
+            for model_data in data["models"]:
+                # Reconstruct set for tags
+                model_data["tags"] = set(model_data.get("tags", []))
+                model = Model(**model_data)
+                self._models[model.id] = model
+
+    def _save(self) -> None:
+        """Save models to persistent storage."""
+        if not self._storage_path:
+            return
+
+        models_data = []
+        for model in self._models.values():
+            model_dict = model.to_dict()
+            # Convert set to list for JSON serialization
+            model_dict["tags"] = list(model_dict["tags"])
+            models_data.append(model_dict)
+
+        self._storage_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(self._storage_path, "w") as f:
+            json.dump({"models": models_data}, f, indent=2)
+
+
+class ModelSource(ABC):
+    """Abstract base for model discovery sources."""
+
+    @abstractmethod
+    def discover_models(self) -> Iterable[Model]:
+        """Discover available models from this source.
+
+        Yields Model instances for each discovered model.
+        """
+        pass
+
+
+class Connector(ABC):
+    """Abstract base for model connectors/clients."""
+
+    @abstractmethod
+    def format_metadata(self, model: Model) -> dict[str, Any]:
+        """Format model metadata for this connector.
+
+        Returns a dict that can be used to instantiate/connect via this connector.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Unique identifier for this connector."""
+        pass
+
+
+class ConnectorRegistry(MutableMapping[str, Connector]):
+    """Registry for managing model connectors.
+
+    >>> registry = ConnectorRegistry()
+    >>> class MyConnector(Connector):
+    ...     @property
+    ...     def name(self) -> str:
+    ...         return "my_connector"
+    ...     def format_metadata(self, model: Model) -> dict[str, Any]:
+    ...         return {"model": model.id}
+    >>> connector = MyConnector()
+    >>> registry[connector.name] = connector
+    >>> "my_connector" in registry
+    True
+    """
+
+    def __init__(self):
+        self._connectors: dict[str, Connector] = {}
+
+    def __setitem__(self, name: str, connector: Connector) -> None:
+        self._connectors[name] = connector
+
+    def __getitem__(self, name: str) -> Connector:
+        return self._connectors[name]
+
+    def __delitem__(self, name: str) -> None:
+        del self._connectors[name]
+
+    def __iter__(self) -> Iterator[str]:
+        yield from self._connectors.keys()
+
+    def __len__(self) -> int:
+        return len(self._connectors)