aigateway-sdk

Ai_Gateway

A unified AI Gateway SDK for LLM routing, observability, and evaluation.

Features

• Multi-Provider Gateway: Unified API for OpenAI, Anthropic, Google Gemini with automatic fallbacks
• Full Observability: Distributed tracing, logging, cost tracking, and analytics
• Evaluation Metrics: Built-in metrics for relevance, coherence, toxicity detection
• Smart Caching: Response caching with semantic similarity matching
• Self-Hosted or Cloud: Store telemetry locally (SQLite) or in the cloud

Installation

pip install Ai_Gateway

# With all providers
pip install Ai_Gateway[all-providers]

# With semantic caching
pip install Ai_Gateway[semantic-cache]

# Everything
pip install Ai_Gateway[full]

Quick Start

Basic Usage

import aigateway

# Initialize with your API keys
aigateway.init(
    openai_api_key="sk-...",
    anthropic_api_key="sk-ant-...",
    google_api_key="...",
)

# Use the unified gateway
response = aigateway.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "Hello!"}]
)

print(response.choices[0].message.content)

Monitor Existing Clients

from openai import OpenAI
import aigateway

# Initialize AIGateway
aigateway.init()

# Wrap your existing client
client = OpenAI()
aigateway.monitor(client)

# All calls are now automatically traced
response = client.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "Hello!"}]
)

Automatic Fallbacks

import aigateway

aigateway.init(
    openai_api_key="sk-...",
    anthropic_api_key="sk-ant-...",
    fallback_providers=["anthropic"],  # Fallback to Anthropic if OpenAI fails
)

# Will automatically retry with Anthropic if OpenAI fails
response = aigateway.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "Hello!"}]
)

Tracing

import aigateway
from aigateway import trace, span

@trace(name="my_workflow")
async def process_request(query: str):
    # Create spans for sub-operations
    with span("retrieve_context"):
        context = await retrieve_documents(query)

    with span("generate_response"):
        response = await generate_answer(query, context)

    return response

Evaluation

import aigateway

# Evaluate LLM outputs
result = await aigateway.evaluate(
    input_text="What is Python?",
    output_text="Python is a programming language...",
    metrics=["relevance", "coherence", "toxicity"]
)

print(f"Overall Score: {result.overall_score}")
print(f"Passed: {result.passed}")

for metric in result.metrics:
    print(f"  {metric.name}: {metric.score:.2f}")

Track Feedback

import aigateway

# Track user feedback on responses
aigateway.track_feedback(
    trace_id="...",
    value="thumbs_up",
    comment="Very helpful response!"
)

Configuration

Environment Variables

AIGATEWAY_API_KEY=your-cloud-api-key
OPENAI_API_KEY=sk-...
ANTHROPIC_API_KEY=sk-ant-...
GOOGLE_API_KEY=...

AIGATEWAY_DEFAULT_PROVIDER=openai
AIGATEWAY_STORAGE_BACKEND=sqlite
AIGATEWAY_TRACING_ENABLED=true
AIGATEWAY_CACHE_ENABLED=true

Programmatic Configuration

import aigateway

aigateway.init(
    # Provider keys
    openai_api_key="sk-...",
    anthropic_api_key="sk-ant-...",

    # Gateway settings
    default_provider="openai",
    fallback_providers=["anthropic", "google"],

    # Tracing
    tracing_enabled=True,
    sample_rate=1.0,

    # Storage (sqlite, postgresql, cloud, memory)
    storage_backend="sqlite",
    storage_path="./aigateway.db",

    # Caching
    cache_enabled=True,
    cache_backend="memory",

    # Debug
    debug=False,
)

API Reference

Gateway Client

from aigateway import AIGateway

gateway = AIGateway()

# Async completion
response = await gateway.complete(
    model="gpt-4",
    messages=[{"role": "user", "content": "Hello"}],
    temperature=0.7,
    max_tokens=1000,
)

# Streaming
async for chunk in gateway.complete_stream(
    model="gpt-4",
    messages=[{"role": "user", "content": "Tell me a story"}],
):
    print(chunk.choices[0].delta.content, end="")

# Embeddings
embeddings = await gateway.embed(
    model="text-embedding-3-small",
    input="Hello, world!"
)

Tracing

from aigateway import Trace, Span, SpanKind

# Manual tracing
with Trace(name="my_request", user_id="user123") as trace:
    with trace.span("step1", kind=SpanKind.LLM) as span:
        span.set_input({"query": "Hello"})
        result = do_llm_call()
        span.set_output({"response": result})
        span.set_llm_info(model="gpt-4", provider="openai")

Evaluation Metrics

from aigateway.evaluation import (
    Evaluator,
    RelevanceMetric,
    CoherenceMetric,
    ToxicityMetric,
    CustomMetric,
)

# Custom evaluator
evaluator = Evaluator(metrics=[
    RelevanceMetric(threshold=0.8),
    CoherenceMetric(threshold=0.6),
    ToxicityMetric(threshold=0.2),
])

result = await evaluator.evaluate(
    input_text="...",
    output_text="...",
    context="...",  # Optional
)

Storage Backends

from aigateway.storage import SQLiteStorage, CloudStorage, MemoryStorage

# SQLite (local)
storage = SQLiteStorage(db_path="./data.db")
await storage.initialize()

# Cloud (remote)
storage = CloudStorage(
    endpoint="https://api.aigateway.io",
    api_key="your-api-key"
)
await storage.initialize()

Development

# Install development dependencies
pip install -e ".[dev]"

# Run tests
pytest

# Format code
black src tests
ruff check src tests

# Type checking
mypy src

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

Links

• GitHub: https://github.com/UditMishra-Webby/Ai_gateway
• PyPI: https://pypi.org/project/Ai_Gateway/

License

MIT License - see LICENSE for details.