ai-infra

Build AI applications in minutes, not months.

Overview

One unified SDK for LLMs, agents, RAG, voice, images, and MCP—across 10+ providers.

Key Features

• LLM Chat - Chat, streaming, structured output, retries across providers
• Agents - Tool calling, human-in-the-loop, deep research mode
• RAG - Embeddings, vector stores, retrieval pipelines
• MCP - Client/server, OpenAPI->MCP conversion, tool discovery
• Voice - Text-to-speech, speech-to-text, realtime conversations
• Tracing - OpenTelemetry distributed tracing built-in

Why ai-infra?

Building AI apps means juggling OpenAI, Anthropic, Google, embeddings, vector stores, tool calling, MCP servers... each with different APIs and gotchas.

ai-infra gives you one clean interface that works everywhere:

from ai_infra import Agent

def search_web(query: str) -> str:
    """Search the web."""
    return f"Results for: {query}"

agent = Agent(tools=[search_web])
result = agent.run("Find the latest news about AI")
# Works with OpenAI, Anthropic, Google—same code.

Quick Install

pip install ai-infra

What's Included

30-Second Examples

Chat with any LLM

from ai_infra import LLM

llm = LLM()  # Uses OPENAI_API_KEY by default
response = llm.chat("Explain quantum computing in one sentence")
print(response)

# Switch providers instantly
llm = LLM(provider="anthropic", model="claude-sonnet-4-20250514")
response = llm.chat("Same question, different model")

Build an Agent with Tools

from ai_infra import Agent

def get_weather(city: str) -> str:
    """Get current weather for a city."""
    return f"72F and sunny in {city}"

def search_web(query: str) -> str:
    """Search the web for information."""
    return f"Top results for: {query}"

agent = Agent(tools=[get_weather, search_web])
result = agent.run("What's the weather in Tokyo and find me restaurants there")
# Agent automatically calls both tools and synthesizes the answer

RAG in 5 Lines

from ai_infra import Retriever

retriever = Retriever()
retriever.add_file("company_docs.pdf")
retriever.add_file("product_manual.md")

results = retriever.search("How do I reset my password?")
print(results[0].content)

Connect to MCP Servers

from ai_infra import MCPClient

async with MCPClient("http://localhost:8080") as client:
    tools = await client.list_tools()
    result = await client.call_tool("search", {"query": "AI news"})

Create an MCP Server

from ai_infra import mcp_from_functions

def search_docs(query: str) -> str:
    """Search documentation."""
    return f"Found: {query}"

mcp = mcp_from_functions(name="my-mcp", functions=[search_docs])
mcp.run(transport="stdio")

Supported Providers

Setup

# Set your API keys (use whichever providers you need)
export OPENAI_API_KEY=sk-...
export ANTHROPIC_API_KEY=sk-ant-...
export GOOGLE_API_KEY=...

# That's it. ai-infra auto-detects available providers.

Feature Highlights

Deep Agent (Autonomous Mode)

For complex, multi-step tasks:

from ai_infra import DeepAgent

agent = DeepAgent(
    goal="Analyze this codebase and generate documentation",
    tools=[read_file, write_file, search],
    max_iterations=50,
)

result = await agent.run()
print(result.output)

Includes: Planning, self-correction, progress tracking, human approval gates.

MCP Client with Interceptors

Advanced MCP features:

from ai_infra import MCPClient
from ai_infra.mcp import RetryInterceptor, CachingInterceptor, LoggingInterceptor

async with MCPClient(
    "http://localhost:8080",
    interceptors=[
        RetryInterceptor(max_retries=3),
        CachingInterceptor(ttl=300),
        LoggingInterceptor(),
    ]
) as client:
    # Automatic retries, caching, and logging for all tool calls
    result = await client.call_tool("expensive_operation", {...})

Includes: Callbacks, interceptors, prompts, resources, progress tracking.

RAG with Multiple Backends

from ai_infra import Retriever

# In-memory (development)
retriever = Retriever(backend="memory")

# SQLite (local persistence)
retriever = Retriever(backend="sqlite", path="./vectors.db")

# PostgreSQL with pgvector (production)
retriever = Retriever(backend="postgres", connection_string="...")

# Pinecone (managed cloud)
retriever = Retriever(backend="pinecone", index_name="my-index")

Voice & Multimodal

from ai_infra import TTS, STT

# Text to speech
tts = TTS(provider="elevenlabs")
audio = tts.speak("Hello, world!")

# Speech to text
stt = STT(provider="deepgram")
text = stt.transcribe("audio.mp3")

Image Generation

from ai_infra import ImageGen

gen = ImageGen(provider="openai")  # or "stability", "replicate"
image = gen.generate("A futuristic city at sunset")
image.save("city.png")

CLI Tools

# Test MCP connections
ai-infra mcp test --url http://localhost:8080

# List MCP tools
ai-infra mcp tools --url http://localhost:8080

# Call an MCP tool
ai-infra mcp call --url http://localhost:8080 --tool search --args '{"query": "test"}'

# Server info
ai-infra mcp info --url http://localhost:8080

Documentation

Running Examples

git clone https://github.com/nfraxlab/ai-infra.git
cd ai-infra
poetry install

# Chat
poetry run python -c "from ai_infra import LLM; print(LLM().chat('Hello!'))"

# Agent
poetry run python examples/agents/01_basic_tools.py

# See more examples
ls examples/

Related Packages

ai-infra is part of the nfrax infrastructure suite:

License

MIT License - use it for anything.

Built by nfraxlab

Star us on GitHub | View on PyPI