tracera-axolotl 0.1.0
A modular, extensible Python toolkit for machine learning operations with unified interfaces for LLMs, embeddings, and vector stores across multiple providers and frameworks (LangChain, LlamaIndex). Provides AWS integrations and follows the ports and adapters pattern for clean, extensible architecture.
Requires Python: >=3.12.11

# Axolotl ML Toolkit

A modular, extensible Python toolkit for machine learning operations, with initial support for embeddings and vector stores.

## Overview

Axolotl provides a unified interface for common ML operations across different providers and frameworks. The toolkit follows the ports and adapters pattern to ensure clean separation of concerns and easy extensibility.

### Key Features

- **Unified Interfaces**: Common APIs across different ML providers and frameworks
- **Extensible Architecture**: Easy to add new providers and frameworks
- **Configuration-Driven**: Simple configuration-based adapter creation
- **Type Safety**: Full type hints and Pydantic models
- **Error Handling**: Comprehensive error handling with detailed context
- **Framework Agnostic**: Supports multiple underlying frameworks (LlamaIndex, etc.)

## Modules

### 🧠 Embeddings (`axolotl.embeddings`)

Text-to-vector conversion with support for multiple providers:

- **AWS Bedrock** (via LlamaIndex)
- **Future**: OpenAI, Cohere, Hugging Face, etc.

```python
from axolotl.embeddings import EmbeddingAdapterFactory, BedrockConfig

# Configure Bedrock
config = BedrockConfig(
    model_name="amazon.titan-embed-text-v1",
    region="us-east-1"
)

# Create adapter
embeddings = EmbeddingAdapterFactory.create_adapter(config)

# Generate embeddings
text = "Hello, world!"
vector = embeddings.embed_text(text)
```

### 🗄️ Vector Store (`axolotl.vector_store`)

Distributed vector search and retrieval:

- **OpenSearch** (via LlamaIndex)
- **Future**: Pinecone, Chroma, Weaviate, Qdrant, etc.

```python
from axolotl.vector_store import VectorStoreAdapterFactory, OpenSearchConfig, DocumentChunk, TextQuery

# Configure OpenSearch
config = OpenSearchConfig(
    endpoint="http://localhost:9200",
    dimension=384,
    embedding_config={"provider": "bedrock", "framework": "llamaindex"}
)

# Create adapter
adapter = VectorStoreAdapterFactory.create_adapter(config)

# Add documents
documents = [DocumentChunk(id="doc1", content="Example content")]
adapter.add_documents(documents, "my_index")

# Search
query = TextQuery(text="example query")
results = adapter.search(query, "my_index")
```

## Quick Start

### Installation

```bash
# Install core package
pip install tracera-axolotl

# Install with optional dependencies
pip install tracera-axolotl[llamaindex]
```

### Basic Usage

```python
from axolotl.embeddings import EmbeddingAdapterFactory, BedrockConfig
from axolotl.vector_store import VectorStoreAdapterFactory, OpenSearchConfig, DocumentChunk, TextQuery

# 1. Create embedding service
embedding_config = BedrockConfig(
    model_name="amazon.titan-embed-text-v1",
    region="us-east-1"
)
embeddings = EmbeddingAdapterFactory.create_adapter(embedding_config)

# 2. Create vector store
vector_config = OpenSearchConfig(
    endpoint="http://localhost:9200",
    dimension=embeddings.dimension,
    embedding_config={"provider": "bedrock", "framework": "llamaindex"}
)
vector_store = VectorStoreAdapterFactory.create_adapter(vector_config)

# 3. Add documents
documents = [
    DocumentChunk(
        id="doc1",
        content="Machine learning is fascinating",
        metadata={"topic": "AI"}
    )
]
vector_store.add_documents(documents, "knowledge_base")

# 4. Search
query = TextQuery(text="artificial intelligence", limit=5)
results = vector_store.search(query, "knowledge_base")
```

## Architecture

The toolkit follows a consistent architecture across all modules:

### Domain Layer
- Pure data models and business logic
- Configuration classes for each provider
- Type definitions and validation

### Ports Layer
- Abstract interfaces defining contracts
- Framework-agnostic operation definitions
- Consistent error handling

### Adapters Layer
- Concrete implementations for specific providers
- Framework-specific integrations
- Error translation and handling

### Factory Pattern
- Configuration-driven adapter creation
- Automatic discovery via decorators
- Type-safe configuration validation

## Adapter System

The toolkit uses a decorator-based adapter registration pattern that allows for easy extension and automatic discovery of new providers and frameworks.

### How Adapters Work

Adapters are concrete implementations of the port interfaces. Each adapter is responsible for:

- **Provider-specific logic**: Implementing operations for a specific backend
- **Configuration handling**: Converting generic configs to provider-specific settings
- **Error translation**: Converting provider errors to standardized exceptions
- **Framework integration**: Working with underlying frameworks like LlamaIndex

### Adapter Registration

Adapters are automatically registered using decorators:

```python
# For embeddings
@register_embedding_adapter("bedrock", "llamaindex")
class LlamaindexBedrockEmbeddingAdapter(BaseEmbeddingAdapter):
    # Implementation
    pass

# For vector stores
@register_vector_store_adapter("aws", "llamaindex")
class LlamaIndexOpenSearchAdapter(BaseVectorStoreAdapter):
    # Implementation
    pass
```

### Registration Process

1. **Decorator Application**: The `@register_*_adapter("provider", "framework")` decorator registers the adapter
2. **Automatic Discovery**: When the module is imported, adapters are automatically discovered and registered
3. **Factory Integration**: The respective factory can then create instances based on configuration

**Note**: The automatic discovery happens when the `adapters.factory` module is imported, which triggers the `_discover_adapters()` function that imports all adapter modules. This means adapters are registered as soon as the package is imported.

### Current Adapters

#### Embeddings
- **LlamaindexBedrockEmbeddingAdapter**: Bedrock implementation using LlamaIndex framework
  - Provider: `"bedrock"`
  - Framework: `"llamaindex"`
  - Configuration: `BedrockConfig`

#### Vector Store
- **LlamaIndexOpenSearchAdapter**: OpenSearch implementation using LlamaIndex framework
  - Provider: `"aws"`
  - Framework: `"llamaindex"`
  - Configuration: `OpenSearchConfig`

### Adding New Adapters

To add a new provider:

1. **Create Configuration Class**: Define a Pydantic model with `provider` and `framework` fields
2. **Implement Adapter**: Extend the appropriate base adapter class and implement all abstract methods
3. **Register with Decorator**: Use the appropriate `@register_*_adapter("provider", "framework")` decorator
4. **Import Module**: Ensure the adapter module is imported to trigger registration

### Factory Integration

The factories automatically discover and use registered adapters:

```python
# Embeddings
from axolotl.embeddings import EmbeddingAdapterFactory, BedrockConfig
config = BedrockConfig(model_name="amazon.titan-embed-text-v1")
embeddings = EmbeddingAdapterFactory.create_adapter(config)

# Vector Store
from axolotl.vector_store import VectorStoreAdapterFactory, OpenSearchConfig
config = OpenSearchConfig(endpoint="http://localhost:9200", dimension=384)
vector_store = VectorStoreAdapterFactory.create_adapter(config)
```

### Factory Pattern Benefits

- **Automatic Discovery**: Adapters register themselves via decorators
- **Configuration Validation**: Type-safe configuration objects
- **Consistent Interface**: All adapters implement the same ports
- **Error Handling**: Centralized error handling and reporting

## Configuration

### Environment Variables

```bash
# AWS Configuration
export AWS_REGION="us-east-1"
export AWS_ACCESS_KEY_ID="your-key"
export AWS_SECRET_ACCESS_KEY="your-secret" # pragma: allowlist secret

# OpenSearch Configuration
export OPENSEARCH_ENDPOINT="http://localhost:9200"
export OPENSEARCH_USERNAME="admin"
export OPENSEARCH_PASSWORD="admin" # pragma: allowlist secret
```

### Configuration Classes

All configuration classes support:
- Environment variable fallbacks
- Type validation with Pydantic
- Default values and optional fields
- Provider-specific settings

## Development

### Prerequisites

- Python 3.12+
- Required dependencies (see pyproject.toml)

### Running Examples

```bash
# Embeddings examples
uv run python -m axolotl.embeddings.examples.bedrock_usage
uv run python -m axolotl.embeddings.examples.list_models

# Vector store examples
uv run python -m axolotl.vector_store.examples.opensearch_cli
```

### Testing

```bash
# Run all tests
uv run pytest

# Run with coverage
uv run pytest --cov=axolotl --cov-report=term-missing
```

## Dependencies

### Core Dependencies
- **pydantic**: Data validation and configuration
- **opensearch-py**: OpenSearch Python client

### Optional Dependencies
- **llama-index**: LlamaIndex framework (for llamaindex adapters)
- **llama-index-embeddings-bedrock**: Bedrock integration via LlamaIndex
- **llama-index-vector-stores-opensearch**: OpenSearch integration via LlamaIndex