Metadata-Version: 2.4
Name: litmus-sdk
Version: 0.1.3
Summary: LLM regression testing via embedding drift detection
License: MIT
Keywords: llm,testing,regression,embedding,drift
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Software Development :: Testing
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: chromadb>=1.4.1
Requires-Dist: click>=8.1
Requires-Dist: litellm>=1.81
Requires-Dist: numpy>=1.24
Requires-Dist: openai>=2.0
Requires-Dist: pydantic>=2.0
Requires-Dist: python-dotenv>=1.0
Requires-Dist: PyYAML>=6.0
Requires-Dist: requests>=2.31
Requires-Dist: scikit-learn>=1.6
Requires-Dist: scipy>=1.11
Requires-Dist: tomlkit>=0.12
Provides-Extra: dev
Requires-Dist: pytest>=9.0; extra == "dev"

# litmus-sdk

LLM regression testing via embedding drift detection.

Litmus instruments your LiteLLM calls passively, stores traces, and uses embedding-space distance to detect whether your LLM outputs have semantically drifted between versions — not just whether they match character-for-character.

---

## Installation

```bash
pip install litmus-sdk
```

With optional providers:

```bash
pip install litmus-sdk[openai]           # OpenAI embedding provider
pip install litmus-sdk[chromadb]         # ChromaDB vector storage
pip install litmus-sdk[openai,chromadb]  # both
```

---

## Quickstart

### Passive instrumentation

```python
import litellm
from litmus_sdk import LitmusSDK

sdk = LitmusSDK(db_path="./litmus.db", project_id="my-project")
sdk.init("v1.0.0")

# All litellm.completion() calls are now automatically captured
response = litellm.completion(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": "Summarise this article: ..."}]
)

sdk.init("v2.0.0")
# ... more calls with a different prompt or model ...

report = sdk.compare("v1.0.0", "v2.0.0")
report.display()
sdk.close()
```

### With GoldenTestRunner

```python
import litellm
from litmus_sdk import LitmusSDK, GoldenTestRunner

def my_agent(question: str) -> str:
    resp = litellm.completion(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": question}]
    )
    return resp.choices[0].message.content

with LitmusSDK(project_id="qa-bot") as sdk:
    sdk.init("v1.0.0")
    runner = GoldenTestRunner(sdk)

    runner.add_test_case(
        name="capital_cities",
        inputs={"question": "What is the capital of France?"},
        run_function=my_agent,
        expected_pattern=r"Paris",
    )

    report = runner.run_tests(num_runs=5)
    report.display()

    sdk.init("v2.0.0")
    runner.run_tests(num_runs=5)

    drift = sdk.compare("v1.0.0", "v2.0.0")
    drift.display()
```

---

## Environment variables

```env
OPENAI_API_KEY=sk-...              # required for embedding computation
LITMUS_DB_PATH=./litmus.db         # SQLite file path (created automatically)
LITMUS_CHROMA_PATH=./chroma        # ChromaDB directory (defaults to ./chroma next to DB)
```

---

## API Reference

### `LitmusSDK`

The main entry point.

```python
sdk = LitmusSDK(
    db_path="./litmus.db",     # path to SQLite file
    chroma_path="./chroma",    # path to ChromaDB directory (optional)
    project_id="my-project",   # tags all traces (optional)
)
```

| Method | Description |
|--------|-------------|
| `init(version)` | Sets the active version and registers the LiteLLM callback. Returns `self` for chaining. |
| `get_traces(version)` | Returns all `LLMTrace` objects for the given version. |
| `get_versions()` | Returns all version strings in insertion order. |
| `compare(version_a, version_b)` | Runs drift detection and returns a `DriftReport`. Triggers embedding backfill on first call. |
| `close()` | Deregisters the callback and closes the DB. |

Supports context manager: `with LitmusSDK() as sdk:` calls `close()` automatically.

---

### `GoldenTestRunner`

Executes test cases multiple times to measure stability and populate the drift pipeline.

```python
runner = GoldenTestRunner(sdk)
```

| Method | Description |
|--------|-------------|
| `add_test_case(name, inputs, run_function, expected_pattern, expected_behavior)` | Registers a test case. `run_function` is called with `**inputs` on each run. |
| `run_tests(num_runs)` | Runs all registered tests `num_runs` times. Returns a `TestReport`. |

---

### `DriftDetector`

Can be used directly for more control over drift comparison.

```python
from litmus_sdk.drift import DriftDetector

detector = DriftDetector(sdk, drift_threshold=0.10)
report = detector.compare("v1.0.0", "v2.0.0")
```

The `drift_threshold` sets the sensitivity for flagging individual tests (cosine distance). For full control over classification thresholds, see `DriftConfig` below.

---

### `DriftConfig` — classification thresholds

| Attribute | Default | Meaning |
|-----------|---------|---------|
| `threshold_safe` | `0.05` | Output drift below this → `SAFE` |
| `threshold_review` | `0.15` | Below this → `REVIEW` |
| `threshold_warning` | `0.30` | Below this → `WARNING`, at or above → `REGRESSION` |
| `stability_min_safe` | `0.90` | Output stability must be ≥ this for `SAFE` |
| `stability_min_review` | `0.80` | Output stability must be ≥ this for `REVIEW` |

---

### `DriftReport`

Returned by `sdk.compare()` and `detector.compare()`.

| Field | Type | Description |
|-------|------|-------------|
| `version_a` / `version_b` | `str` | The two versions being compared |
| `avg_input_drift` | `float` | Mean cosine distance of prompt embeddings [0, 2] |
| `avg_output_drift` | `float` | Mean cosine distance of response embeddings [0, 2] |
| `stability_a` / `stability_b` | `float` | Output consistency per version [0, 1] |
| `flagged_tests` | `list[str]` | Test names that exceeded the drift threshold |
| `recommendation` | `str` | One of `SAFE`, `REVIEW`, `WARNING`, `REGRESSION` |
| `details` | `str` | Human-readable explanation |

Call `report.display()` to pretty-print with status icons.

---

### `LitmusCallback`

The LiteLLM callback that captures traces. Registered automatically by `sdk.init()` — you typically don't need to use it directly. If you do need custom integration:

```python
from litmus_sdk.callback import LitmusCallback

callback = LitmusCallback(sdk)
litellm.callbacks = [callback]
```

---

### Data models

| Model | Description |
|-------|-------------|
| `LLMTrace` | One captured LLM call: prompt, response, model, tokens, latency, cost, embeddings |
| `PromptComponents` | Decomposed prompt: `system_prompt`, `user_input`, `context`, `final_prompt` |
| `TestCase` | A golden test definition with `name`, `inputs`, `run_function`, `expected_pattern` |
| `TestResult` | One test execution: `output`, `passed`, `latency_ms`, `error` |
| `TestReport` | Full suite result with per-test summaries and pass rates |

---

## How it works

```
litellm.completion()
        │
        ▼
LitmusCallback (LiteLLM hook)
  - Extracts prompt + response
  - Stores metadata → SQLite
  - Computes embeddings eagerly (OpenAI API call)
  - Stores vectors → ChromaDB
        │
        ▼
sdk.compare("v1", "v2")
        │
        ▼
DriftDetector
  1. Load traces for both versions
  2. Compute centroids of embeddings + cosine distances
  3. Classify: SAFE / REVIEW / WARNING / REGRESSION
  4. Return DriftReport
```

**Embeddings are computed eagerly** during callback capture (via OpenAI API, ~50-200 ms per call) so that drift detection works immediately without additional setup. Embedding computation happens asynchronously in the LiteLLM callback hook, so it doesn't block your agent's latency.

**Hybrid storage** — SQLite stores structured metadata (fast queries, no infrastructure). ChromaDB stores vectors (optimized cosine distance). Both are linked by `trace_id`.

---

## Storage backends

Default: SQLite + ChromaDB, both local files.

| Env var | Default | Purpose |
|---------|---------|---------|
| `LITMUS_DB_PATH` | `./litmus.db` | SQLite file |
| `LITMUS_CHROMA_PATH` | `./chroma` (next to DB) | ChromaDB directory |

SQLite uses WAL mode for thread safety.

---

## Embedding providers

Default: `text-embedding-3-small` via OpenAI.

To swap the provider, subclass `EmbeddingProvider` and pass it to `LitmusConfig`:

```python
from litmus_sdk.embeddings.base import EmbeddingProvider

class MyProvider(EmbeddingProvider):
    def embed(self, text: str) -> list[float]: ...
    def embed_batch(self, texts: list[str]) -> list[list[float]]: ...
    @property
    def dimensions(self) -> int: return 1536
```

---

## License

MIT
