Metadata-Version: 2.4
Name: mechanex
Version: 0.5.0
Summary: A Python client for the Axionic API.
Home-page: https://axioniclabs.ai
Author: Axionic Labs
Author-email: contact@axioniclabs.ai
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.7
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: requests>=2.20.0
Requires-Dist: ipython
Requires-Dist: jupyter
Requires-Dist: matplotlib
Requires-Dist: click
Requires-Dist: fastapi
Requires-Dist: uvicorn
Requires-Dist: torch
Requires-Dist: numpy
Requires-Dist: tqdm
Requires-Dist: mdmm
Requires-Dist: rich
Requires-Dist: sae-lens
Requires-Dist: openai
Requires-Dist: huggingface_hub
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: license-file
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# Mechanex

Mechanex allows you to control and debug your LLMs. Learn more at [axioniclabs.ai](https://axioniclabs.ai/)

## Installation

```bash
pip install mechanex
```

## Quick Start

### 1. Initialize the Client
You must have an API key to use mechanex. The CLI helps you manage this easily.

**New Users:**
Run the signup command to create an account, automatically log in, and generate your first API key:
```bash
mechanex signup
```

**Existing Users:**
If you already have an account, log in and generate a key manually:
```bash
mechanex login
mechanex create-api-key
```

**Using the Key:**
The Python client automatically loads the key generated by the CLI.
```python
import mechanex as mx

# The client automatically loads the key from ~/.mechanex/config.json. You can alternatively set your key manually. If persist is True, the key will be saved to the config file.
mx.set_key("your-api-key-here", persist=False)
```

### 2. Generation and Sampling
You can control generation using various sampling methods:

```python
output = mx.generation.generate(
    prompt="The future of AI is",
    max_tokens=64,
    # Primary sampling method
    sampling_method="top-p", # Options: "greedy", "top-k", "top-p", "ads" (requires paid API key)
    # Sampling parameters
    top_p=0.9,
    top_k=50
)
print(output)
```

## Local Model Management

Mechanex allows you to load models locally for inspection and low-latency hooks. The full list of supported models are as follows (this is to maintain compatibility for our automated steering pipeline):

| Family | Models |
| :--- | :--- |
| **Gemma 2** | `gemma-2-27b`, `gemma-2-2b`, `gemma-2-9b`, `gemma-2-9b-it`, `gemma-2b`, `gemma-2b-it` |
| **Gemma 3** | `gemma-3-12b-it`, `gemma-3-12b-pt`, `gemma-3-1b-it`, `gemma-3-1b-pt`, `gemma-3-270m`, `gemma-3-270m-it`, `gemma-3-27b-it`, `gemma-3-27b-pt`, `gemma-3-4b-it`, `gemma-3-4b-pt` |
| **Llama** | `llama-3.1-8b`, `llama-3.1-8b-instruct`, `llama-3.3-70b-instruct`, `meta-llama-3-8b-instruct` |
| **Qwen** | `qwen2.5-7b-instruct`, `qwen3-0.6b`, `qwen3-1.7b`, `qwen3-14b`, `qwen3-4b`, `qwen3-8b` |
| **Other** | `deepseek-r1-distill-llama-8b`, `gpt-oss-20b`, `gpt2-small`, `mistral-7b`, `pythia-70m-deduped` |

### Loading a Local Model
```python
import mechanex as mx
mx.set_key("your-api-key-here") # Required even for local mode

mx.load("gpt2") # Uses transformer-lens to load the model
```

### Unloading a Model
To free up GPU memory and switch back to remote execution flow:
```python
mx.unload()
```

## CLI Commands

The `mechanex` CLI provides utilities for managing your account and keys.

- `mechanex signup`: Register a new account.
- `mechanex login`: Authenticate and save your credentials.
- `mechanex whoami`: View your current session and profile.
- `mechanex list-api-keys`: View all your active API keys.
- `mechanex create-api-key`: Generate a new persistent API key.
- `mechanex logout`: Clear your local session credentials.

## Steering Vectors

Steering vectors allow you to control the behavior of a model by injecting specific activation patterns.

### Compute a Steering Vector
```python
# Create a vector from contrastive pairs
vector_id = mx.steering.generate_vectors(
    prompts=["I think that", "People say"],
    positive_answers=[" kindness is key", " helping is good"],
    negative_answers=[" hate is power", " hurting is fine"],
    method="few-shot" # Options: caa, few-shot
)

# Apply it during generation
steered_output = mx.generation.generate(
    prompt="What is your philosophy?",
    steering_vector=vector_id,
    steering_strength=1.5
)
```

## SAE (Sparse Autoencoder) Pipeline

The SAE pipeline provides advanced behavioral detection and automatic correction.

### 1. Create a Behavior
Define a behavior to monitor and potentially correct.
```python
behavior = mx.sae.create_behavior_from_jsonl(
    behavior_name="toxicity",
    dataset_path="tests/toxicity_dataset.jsonl",
    description="Reduces toxic output"
)
```

### 2. Generate with SAE Steering
Utilize SAEs to detect behavioral drift and automatically apply corrections.
```python
# Generate with auto-correction enabled
response = mx.sae.generate(
    prompt="Tell me a secret",
    auto_correct=True,
    behavior_names=["toxicity"]
)
print(response)
```

## Deployment & Serving

### Local OpenAI-Compatible Server
Mechanex can host an OpenAI-compatible server that leverages your locally loaded model or remote API.

```python
import mechanex as mx
mx.load("gpt2")

# Serve with optional global behavior correction
mx.serve(
    port=8000, 
    corrected_behaviors=["anger"] # Enforces "anger" correction on all requests
)
```

You can then interact with it using any standard tool, like the **OpenAI Python client**. Mechanix supports custom parameters via `extra_body` for mechanistic features:

```python
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="any")

# 1. Standard Chat Completion
completion = client.chat.completions.create(
    model="mechanex",
    messages=[{"role": "user", "content": "Hello!"}]
)

# 2. Steered Completion (using extra_body)
steered_completion = client.chat.completions.create(
    model="mechanex",
    messages=[{"role": "user", "content": "Hello!"}],
    extra_body={
        "steering_vector": "your-vector-id",
        "steering_strength": 2.0
    }
)

# 3. SAE-monitored Completion
sae_completion = client.chat.completions.create(
    model="mechanex",
    messages=[{"role": "user", "content": "How are you?"}],
    extra_body={
        "behavior_names": ["toxicity"],
        "auto_correct": True
    }
)
```

### vLLM Integration
For high-performance serving, you can integrate with vLLM by passing `use_vllm=True` to the `serve` method.

<!-- 
### On Prem Deployments

If you want to host your own instance of the remote hooks server:

1. **Docker Image**: Use `axioniclabs/remote-hooks:general`.
2. **Environment Variables**:
   - `MODEL_NAME`: Hugging Face model ID.
   - `HF_TOKEN`: Your Hugging Face token.

This still requires a paid mechanex API key.

---
For more details, visit [axioniclabs.ai](https://axioniclabs.ai/) -->
