amiai 0.1.0

AMIAI SDK - Build AI agents with local tool execution

amiai

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

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