amiai-sdk

Build AI agents with local tool execution. Write tools that run on your machine, connect to AMIAI, and let agents call them securely over WebSocket.

Installation

pip install amiai-sdk

Quick Start

1. Create a tool

# tools/search.tool.py
from amiai_sdk import tool

@tool(name="web_search", description="Search the web for information")
async def search(query: str) -> dict:
    # Your API key stays local - never sent to AMIAI
    import httpx
    async with httpx.AsyncClient() as client:
        response = await client.get(
            "https://api.example.com/search",
            params={"q": query},
            headers={"x-api-key": os.environ["SEARCH_API_KEY"]}
        )
        return response.json()

2. Run the harness

export AMIAI_API_KEY=amiai_xxx
amiai dev

Output:

✓ Connected to AMIAI
✓ Registered 1 tool: web_search
✓ Waiting for invocations...

3. Create an agent that uses your tool

curl -X POST https://backend.amiai.com/api/agents \
  -H "Authorization: Bearer $AMIAI_API_KEY" \
  -d '{
    "name": "Search Agent",
    "systemPrompt": "Help users search the web.",
    "tools": [{"name": "web_search", "source": "live"}]
  }'

When the agent calls web_search, it executes on your machine with your local API keys.

How It Works

Your Machine                              AMIAI Backend
─────────────                             ─────────────
search.tool.py
     ↓
amiai dev ──────────WebSocket────────▶ ToolSessionDO
                                              │
     ◀───────── "invoke web_search" ──────────┤
     │                                        │
     ├── execute locally with your API keys   │
     │                                        │
     └─────────── "result" ──────────────────▶│──▶ Agent gets result

Benefits:

• 🔐 API keys never leave your machine
• 🚀 No webhooks or public servers needed
• 🔄 Hot reload - just save and reconnect
• 🛠️ Full access to local resources (files, databases, etc.)

API Reference

@tool decorator

Define a tool for AMIAI agents to call.

from amiai_sdk import tool

@tool(
    name="my_tool",
    description="What this tool does"
)
async def my_tool(
    param1: str,      # Required parameter
    param2: int = 10  # Optional with default
) -> dict:
    return {"result": "data"}

Input schema is automatically derived from function signature and type hints.

Explicit schema

@tool(
    name="calculator",
    description="Evaluate math expressions",
    input_schema={
        "type": "object",
        "properties": {
            "expression": {"type": "string", "description": "e.g., 2+2"}
        },
        "required": ["expression"]
    }
)
def calculate(expression: str) -> dict:
    return {"result": eval(expression)}

Harness class

Programmatic control over the WebSocket connection.

from amiai_sdk import Harness

harness = Harness(
    api_key="amiai_xxx",
    tools=[my_tool],
    url="wss://backend.amiai.com/api/tools/live",  # optional
    on_connect=lambda sid: print(f"Connected: {sid}"),
    on_disconnect=lambda: print("Disconnected"),
    on_invoke=lambda tool, args: print(f"Invoking {tool}"),
    on_result=lambda tool, result: print(f"Result from {tool}"),
    on_error=lambda e: print(f"Error: {e}"),
)

# Start and run forever
await harness.start()
await harness.run_forever()

# Or manage manually
await harness.start()
# ... do other things
await harness.stop()

start_harness function

Convenience function that creates and starts a harness.

from amiai_sdk import start_harness

harness = await start_harness(
    api_key="amiai_xxx",
    tools=[my_tool],
)

CLI Usage

The amiai CLI automatically discovers tools in your project.

# Discover and register all *.tool.py and tools/*.py files
amiai dev

# Specify custom patterns
amiai dev --pattern "src/**/*.py"

# Use a different API endpoint (for local development)
AMIAI_URL=ws://localhost:8787/api/tools/live amiai dev

Examples

Calculator Tool

from amiai_sdk import tool

@tool(name="calculator", description="Evaluate mathematical expressions")
def calculate(expression: str) -> dict:
    """
    Calculate the result of a math expression.
    
    Args:
        expression: Math expression like "2 + 2 * 3"
    """
    result = eval(expression)  # In production, use a safe evaluator
    return {"expression": expression, "result": result}

Web Search with Parallel AI

import os
import httpx
from amiai_sdk import tool

@tool(name="web_search", description="Search the web using Parallel AI")
async def search(query: str) -> dict:
    async with httpx.AsyncClient() as client:
        response = await client.post(
            "https://api.parallel.ai/v1beta/search",
            headers={
                "Content-Type": "application/json",
                "x-api-key": os.environ["PARALLEL_API_KEY"],
                "parallel-beta": "search-extract-2025-10-10",
            },
            json={"search_queries": [query]},
        )
        data = response.json()
        return {
            "query": query,
            "results": [
                {"title": r["title"], "url": r["url"]}
                for r in data.get("results", [])[:5]
            ]
        }

File Reader

from pathlib import Path
from amiai_sdk import tool

@tool(name="read_file", description="Read a file from the local filesystem")
def read_file(path: str) -> dict:
    content = Path(path).read_text()
    return {"path": path, "content": content}

Database Query

import sqlite3
from amiai_sdk import tool

@tool(name="query_db", description="Query the local SQLite database")
def query_database(sql: str) -> dict:
    conn = sqlite3.connect("local.db")
    cursor = conn.execute(sql)
    columns = [d[0] for d in cursor.description]
    rows = [dict(zip(columns, row)) for row in cursor.fetchall()]
    conn.close()
    return {"columns": columns, "rows": rows}

Sync vs Async Tools

Both sync and async functions work:

# Async tool (recommended for I/O operations)
@tool(name="async_search", description="Async web search")
async def async_search(query: str) -> dict:
    async with httpx.AsyncClient() as client:
        response = await client.get(f"https://api.example.com/search?q={query}")
        return response.json()

# Sync tool (fine for CPU-bound operations)
@tool(name="sync_calculate", description="Calculate expression")
def sync_calculate(expression: str) -> dict:
    return {"result": eval(expression)}

License

MIT