amiai-sdk 0.1.0
AMIAI SDK - Build AI agents with local tool execution
Requires Python: >=3.9
Keywords: agents, ai, amiai, llm, sdk, tools

# 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

```bash
pip install amiai-sdk
```

## Quick Start

### 1. Create a tool

```python
# 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

```bash
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

```bash
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.

```python
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

```python
@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.

```python
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.

```python
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.

```bash
# 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

```python
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

```python
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

```python
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

```python
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:

```python
# 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