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

aix/__init__.py

@@ -1,62 +1,159 @@
-"""
-Facade to key AI tools.
-
-Get a list of available chat functions (this will depend on the AI packages
-you have installed locally):
-
->>> 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']
-
-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'
+"""AIX: Artificial Intelligence eXtensions
+
+A clean, pythonic facade for common AI operations that abstracts away
+provider-specific details and complexities.
+
+Quick Start:
+    >>> from aix import chat, embeddings, prompt_func, models
+
+    # 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'
+
+    # 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', ...]
+
+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.gen_ai import chat, chat_models, chat_funcs
-
-# TODO: Change this so that there's a load_pkg function that loads the packages dynamically
-#   if and when use wants.
-
-# from aix.pd import *
-# from aix.np import *
-# from aix.sk import *
-
-# from aix import pd
-# from aix import np
-# from aix import sk
-
-
-#
-# 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)
+# 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,
+)
+
+# Legacy interfaces (for backward compatibility)
+from aix.gen_ai import chat_models, chat_funcs
+
+# Version info
+__version__ = "0.1.0"
+
+# 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/base.py

@@ -4,8 +4,9 @@ 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
-from typing import Any, Mapping, MutableMapping, Iterator, Callable, Iterable
+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
@@ -14,7 +15,7 @@ from pathlib import Path
 @dataclass
 class Model:
     """Represents an AI model with its metadata.
-    
+
     >>> model = Model(
     ...     id="gpt-4",
     ...     provider="openai",
@@ -24,6 +25,7 @@ class Model:
     >>> model.id
     'gpt-4'
     """
+
     id: str
     provider: str
     context_size: int | None = None
@@ -33,14 +35,14 @@ class Model:
     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
@@ -54,10 +56,25 @@ class Model:
                 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
@@ -65,23 +82,25 @@ class ModelRegistry(MutableMapping[str, Model]):
     >>> 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]:
+
+    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"]]
@@ -95,21 +114,21 @@ class ModelRegistry(MutableMapping[str, Model]):
             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,
         *,
@@ -119,10 +138,10 @@ class ModelRegistry(MutableMapping[str, Model]):
         max_context_size: int | None = None,
         has_capabilities: Iterable[str] | None = None,
         tags: Iterable[str] | None = None,
-        custom_filter: Callable[[Model], bool] | 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)
@@ -130,14 +149,19 @@ class ModelRegistry(MutableMapping[str, Model]):
         >>> 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):
+            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):
+            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:
@@ -148,46 +172,46 @@ class ModelRegistry(MutableMapping[str, Model]):
             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, 'r') as f:
+
+        with open(self._storage_path) as f:
             data = json.load(f)
-            for model_data in data['models']:
+            for model_data in data["models"]:
                 # Reconstruct set for tags
-                model_data['tags'] = set(model_data.get('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'])
+            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)
+        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
@@ -195,15 +219,15 @@ class ModelSource(ABC):
 
 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:
@@ -213,7 +237,7 @@ class Connector(ABC):
 
 class ConnectorRegistry(MutableMapping[str, Connector]):
     """Registry for managing model connectors.
-    
+
     >>> registry = ConnectorRegistry()
     >>> class MyConnector(Connector):
     ...     @property
@@ -226,21 +250,21 @@ class ConnectorRegistry(MutableMapping[str, 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)

aix/ai_models/examples.py

@@ -4,7 +4,8 @@ This module demonstrates common patterns and use cases.
 """
 
 from pathlib import Path
-from typing import Any, Iterable, Callable
+from typing import Any
+from collections.abc import Iterable, Callable
 from aix.ai_models.base import ModelSource, Model
 from aix.ai_models.manager import ModelManager
 
@@ -103,7 +104,7 @@ def example_advanced_filtering():
         """Check if model is suitable for a specific task."""
         has_function_calling = model.capabilities.get("function_calling", False)
         has_sufficient_context = (model.context_size or 0) >= 32000
-        is_affordable = model.cost_per_token.get("input", float('inf')) < 0.01
+        is_affordable = model.cost_per_token.get("input", float("inf")) < 0.01
         return has_function_calling and has_sufficient_context and is_affordable
 
     suitable_models = manager.list_models(custom_filter=_is_suitable_for_task)

aix/ai_models/manager.py

@@ -4,7 +4,8 @@ Provides a high-level API for managing AI models across providers.
 """
 
 from pathlib import Path
-from typing import Any, Iterable, Callable
+from typing import Any
+from collections.abc import Iterable, Callable
 
 from aix.ai_models.base import Model, ModelRegistry, ConnectorRegistry, ModelSource
 from aix.ai_models.sources import (

aix/ai_models/sources.py

@@ -5,7 +5,8 @@ and connectors for formatting metadata for different clients.
 """
 
 import requests
-from typing import Any, Iterable
+from typing import Any
+from collections.abc import Iterable
 from dataclasses import dataclass
 
 from aix.ai_models.base import Model, ModelSource, Connector