amiai 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- amiai-0.1.0.dist-info/METADATA +305 -0
- amiai-0.1.0.dist-info/RECORD +9 -0
- amiai-0.1.0.dist-info/WHEEL +4 -0
- amiai-0.1.0.dist-info/entry_points.txt +2 -0
- amiai_sdk/__init__.py +25 -0
- amiai_sdk/cli.py +177 -0
- amiai_sdk/harness.py +321 -0
- amiai_sdk/py.typed +0 -0
- amiai_sdk/tool.py +147 -0
|
@@ -0,0 +1,305 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: amiai
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: AMIAI SDK - Build AI agents with local tool execution
|
|
5
|
+
Project-URL: Homepage, https://github.com/amiai/sdk-python
|
|
6
|
+
Project-URL: Documentation, https://github.com/amiai/sdk-python#readme
|
|
7
|
+
Project-URL: Repository, https://github.com/amiai/sdk-python
|
|
8
|
+
Project-URL: Issues, https://github.com/amiai/sdk-python/issues
|
|
9
|
+
Author-email: AMIAI <hello@amiai.com>
|
|
10
|
+
License-Expression: MIT
|
|
11
|
+
Keywords: agents,ai,amiai,llm,sdk,tools
|
|
12
|
+
Classifier: Development Status :: 4 - Beta
|
|
13
|
+
Classifier: Intended Audience :: Developers
|
|
14
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
20
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
21
|
+
Requires-Python: >=3.9
|
|
22
|
+
Requires-Dist: pydantic>=2.0
|
|
23
|
+
Requires-Dist: websockets>=12.0
|
|
24
|
+
Provides-Extra: dev
|
|
25
|
+
Requires-Dist: black>=23.0; extra == 'dev'
|
|
26
|
+
Requires-Dist: mypy>=1.0; extra == 'dev'
|
|
27
|
+
Requires-Dist: pytest-asyncio>=0.21; extra == 'dev'
|
|
28
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
29
|
+
Requires-Dist: ruff>=0.1; extra == 'dev'
|
|
30
|
+
Description-Content-Type: text/markdown
|
|
31
|
+
|
|
32
|
+
# amiai
|
|
33
|
+
|
|
34
|
+
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.
|
|
35
|
+
|
|
36
|
+
## Installation
|
|
37
|
+
|
|
38
|
+
```bash
|
|
39
|
+
pip install amiai
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
## Quick Start
|
|
43
|
+
|
|
44
|
+
### 1. Create a tool
|
|
45
|
+
|
|
46
|
+
```python
|
|
47
|
+
# tools/search.tool.py
|
|
48
|
+
from amiai_sdk import tool
|
|
49
|
+
|
|
50
|
+
@tool(name="web_search", description="Search the web for information")
|
|
51
|
+
async def search(query: str) -> dict:
|
|
52
|
+
# Your API key stays local - never sent to AMIAI
|
|
53
|
+
import httpx
|
|
54
|
+
async with httpx.AsyncClient() as client:
|
|
55
|
+
response = await client.get(
|
|
56
|
+
"https://api.example.com/search",
|
|
57
|
+
params={"q": query},
|
|
58
|
+
headers={"x-api-key": os.environ["SEARCH_API_KEY"]}
|
|
59
|
+
)
|
|
60
|
+
return response.json()
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
### 2. Run the harness
|
|
64
|
+
|
|
65
|
+
```bash
|
|
66
|
+
export AMIAI_API_KEY=amiai_xxx
|
|
67
|
+
amiai dev
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
Output:
|
|
71
|
+
```
|
|
72
|
+
✓ Connected to AMIAI
|
|
73
|
+
✓ Registered 1 tool: web_search
|
|
74
|
+
✓ Waiting for invocations...
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
### 3. Create an agent that uses your tool
|
|
78
|
+
|
|
79
|
+
```bash
|
|
80
|
+
curl -X POST https://backend.amiai.com/api/agents \
|
|
81
|
+
-H "Authorization: Bearer $AMIAI_API_KEY" \
|
|
82
|
+
-d '{
|
|
83
|
+
"name": "Search Agent",
|
|
84
|
+
"systemPrompt": "Help users search the web.",
|
|
85
|
+
"tools": [{"name": "web_search", "source": "live"}]
|
|
86
|
+
}'
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
When the agent calls `web_search`, it executes **on your machine** with your local API keys.
|
|
90
|
+
|
|
91
|
+
## How It Works
|
|
92
|
+
|
|
93
|
+
```
|
|
94
|
+
Your Machine AMIAI Backend
|
|
95
|
+
───────────── ─────────────
|
|
96
|
+
search.tool.py
|
|
97
|
+
↓
|
|
98
|
+
amiai dev ──────────WebSocket────────▶ ToolSessionDO
|
|
99
|
+
│
|
|
100
|
+
◀───────── "invoke web_search" ──────────┤
|
|
101
|
+
│ │
|
|
102
|
+
├── execute locally with your API keys │
|
|
103
|
+
│ │
|
|
104
|
+
└─────────── "result" ──────────────────▶│──▶ Agent gets result
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
**Benefits:**
|
|
108
|
+
- 🔐 API keys never leave your machine
|
|
109
|
+
- 🚀 No webhooks or public servers needed
|
|
110
|
+
- 🔄 Hot reload - just save and reconnect
|
|
111
|
+
- 🛠️ Full access to local resources (files, databases, etc.)
|
|
112
|
+
|
|
113
|
+
## API Reference
|
|
114
|
+
|
|
115
|
+
### `@tool` decorator
|
|
116
|
+
|
|
117
|
+
Define a tool for AMIAI agents to call.
|
|
118
|
+
|
|
119
|
+
```python
|
|
120
|
+
from amiai_sdk import tool
|
|
121
|
+
|
|
122
|
+
@tool(
|
|
123
|
+
name="my_tool",
|
|
124
|
+
description="What this tool does"
|
|
125
|
+
)
|
|
126
|
+
async def my_tool(
|
|
127
|
+
param1: str, # Required parameter
|
|
128
|
+
param2: int = 10 # Optional with default
|
|
129
|
+
) -> dict:
|
|
130
|
+
return {"result": "data"}
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
Input schema is automatically derived from function signature and type hints.
|
|
134
|
+
|
|
135
|
+
### Explicit schema
|
|
136
|
+
|
|
137
|
+
```python
|
|
138
|
+
@tool(
|
|
139
|
+
name="calculator",
|
|
140
|
+
description="Evaluate math expressions",
|
|
141
|
+
input_schema={
|
|
142
|
+
"type": "object",
|
|
143
|
+
"properties": {
|
|
144
|
+
"expression": {"type": "string", "description": "e.g., 2+2"}
|
|
145
|
+
},
|
|
146
|
+
"required": ["expression"]
|
|
147
|
+
}
|
|
148
|
+
)
|
|
149
|
+
def calculate(expression: str) -> dict:
|
|
150
|
+
return {"result": eval(expression)}
|
|
151
|
+
```
|
|
152
|
+
|
|
153
|
+
### `Harness` class
|
|
154
|
+
|
|
155
|
+
Programmatic control over the WebSocket connection.
|
|
156
|
+
|
|
157
|
+
```python
|
|
158
|
+
from amiai_sdk import Harness
|
|
159
|
+
|
|
160
|
+
harness = Harness(
|
|
161
|
+
api_key="amiai_xxx",
|
|
162
|
+
tools=[my_tool],
|
|
163
|
+
url="wss://backend.amiai.com/api/tools/live", # optional
|
|
164
|
+
on_connect=lambda sid: print(f"Connected: {sid}"),
|
|
165
|
+
on_disconnect=lambda: print("Disconnected"),
|
|
166
|
+
on_invoke=lambda tool, args: print(f"Invoking {tool}"),
|
|
167
|
+
on_result=lambda tool, result: print(f"Result from {tool}"),
|
|
168
|
+
on_error=lambda e: print(f"Error: {e}"),
|
|
169
|
+
)
|
|
170
|
+
|
|
171
|
+
# Start and run forever
|
|
172
|
+
await harness.start()
|
|
173
|
+
await harness.run_forever()
|
|
174
|
+
|
|
175
|
+
# Or manage manually
|
|
176
|
+
await harness.start()
|
|
177
|
+
# ... do other things
|
|
178
|
+
await harness.stop()
|
|
179
|
+
```
|
|
180
|
+
|
|
181
|
+
### `start_harness` function
|
|
182
|
+
|
|
183
|
+
Convenience function that creates and starts a harness.
|
|
184
|
+
|
|
185
|
+
```python
|
|
186
|
+
from amiai_sdk import start_harness
|
|
187
|
+
|
|
188
|
+
harness = await start_harness(
|
|
189
|
+
api_key="amiai_xxx",
|
|
190
|
+
tools=[my_tool],
|
|
191
|
+
)
|
|
192
|
+
```
|
|
193
|
+
|
|
194
|
+
## CLI Usage
|
|
195
|
+
|
|
196
|
+
The `amiai` CLI automatically discovers tools in your project.
|
|
197
|
+
|
|
198
|
+
```bash
|
|
199
|
+
# Discover and register all *.tool.py and tools/*.py files
|
|
200
|
+
amiai dev
|
|
201
|
+
|
|
202
|
+
# Specify custom patterns
|
|
203
|
+
amiai dev --pattern "src/**/*.py"
|
|
204
|
+
|
|
205
|
+
# Use a different API endpoint (for local development)
|
|
206
|
+
AMIAI_URL=ws://localhost:8787/api/tools/live amiai dev
|
|
207
|
+
```
|
|
208
|
+
|
|
209
|
+
## Examples
|
|
210
|
+
|
|
211
|
+
### Calculator Tool
|
|
212
|
+
|
|
213
|
+
```python
|
|
214
|
+
from amiai_sdk import tool
|
|
215
|
+
|
|
216
|
+
@tool(name="calculator", description="Evaluate mathematical expressions")
|
|
217
|
+
def calculate(expression: str) -> dict:
|
|
218
|
+
"""
|
|
219
|
+
Calculate the result of a math expression.
|
|
220
|
+
|
|
221
|
+
Args:
|
|
222
|
+
expression: Math expression like "2 + 2 * 3"
|
|
223
|
+
"""
|
|
224
|
+
result = eval(expression) # In production, use a safe evaluator
|
|
225
|
+
return {"expression": expression, "result": result}
|
|
226
|
+
```
|
|
227
|
+
|
|
228
|
+
### Web Search with Parallel AI
|
|
229
|
+
|
|
230
|
+
```python
|
|
231
|
+
import os
|
|
232
|
+
import httpx
|
|
233
|
+
from amiai_sdk import tool
|
|
234
|
+
|
|
235
|
+
@tool(name="web_search", description="Search the web using Parallel AI")
|
|
236
|
+
async def search(query: str) -> dict:
|
|
237
|
+
async with httpx.AsyncClient() as client:
|
|
238
|
+
response = await client.post(
|
|
239
|
+
"https://api.parallel.ai/v1beta/search",
|
|
240
|
+
headers={
|
|
241
|
+
"Content-Type": "application/json",
|
|
242
|
+
"x-api-key": os.environ["PARALLEL_API_KEY"],
|
|
243
|
+
"parallel-beta": "search-extract-2025-10-10",
|
|
244
|
+
},
|
|
245
|
+
json={"search_queries": [query]},
|
|
246
|
+
)
|
|
247
|
+
data = response.json()
|
|
248
|
+
return {
|
|
249
|
+
"query": query,
|
|
250
|
+
"results": [
|
|
251
|
+
{"title": r["title"], "url": r["url"]}
|
|
252
|
+
for r in data.get("results", [])[:5]
|
|
253
|
+
]
|
|
254
|
+
}
|
|
255
|
+
```
|
|
256
|
+
|
|
257
|
+
### File Reader
|
|
258
|
+
|
|
259
|
+
```python
|
|
260
|
+
from pathlib import Path
|
|
261
|
+
from amiai_sdk import tool
|
|
262
|
+
|
|
263
|
+
@tool(name="read_file", description="Read a file from the local filesystem")
|
|
264
|
+
def read_file(path: str) -> dict:
|
|
265
|
+
content = Path(path).read_text()
|
|
266
|
+
return {"path": path, "content": content}
|
|
267
|
+
```
|
|
268
|
+
|
|
269
|
+
### Database Query
|
|
270
|
+
|
|
271
|
+
```python
|
|
272
|
+
import sqlite3
|
|
273
|
+
from amiai_sdk import tool
|
|
274
|
+
|
|
275
|
+
@tool(name="query_db", description="Query the local SQLite database")
|
|
276
|
+
def query_database(sql: str) -> dict:
|
|
277
|
+
conn = sqlite3.connect("local.db")
|
|
278
|
+
cursor = conn.execute(sql)
|
|
279
|
+
columns = [d[0] for d in cursor.description]
|
|
280
|
+
rows = [dict(zip(columns, row)) for row in cursor.fetchall()]
|
|
281
|
+
conn.close()
|
|
282
|
+
return {"columns": columns, "rows": rows}
|
|
283
|
+
```
|
|
284
|
+
|
|
285
|
+
## Sync vs Async Tools
|
|
286
|
+
|
|
287
|
+
Both sync and async functions work:
|
|
288
|
+
|
|
289
|
+
```python
|
|
290
|
+
# Async tool (recommended for I/O operations)
|
|
291
|
+
@tool(name="async_search", description="Async web search")
|
|
292
|
+
async def async_search(query: str) -> dict:
|
|
293
|
+
async with httpx.AsyncClient() as client:
|
|
294
|
+
response = await client.get(f"https://api.example.com/search?q={query}")
|
|
295
|
+
return response.json()
|
|
296
|
+
|
|
297
|
+
# Sync tool (fine for CPU-bound operations)
|
|
298
|
+
@tool(name="sync_calculate", description="Calculate expression")
|
|
299
|
+
def sync_calculate(expression: str) -> dict:
|
|
300
|
+
return {"result": eval(expression)}
|
|
301
|
+
```
|
|
302
|
+
|
|
303
|
+
## License
|
|
304
|
+
|
|
305
|
+
MIT
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
amiai_sdk/__init__.py,sha256=mN_LdpdQPAiubDnoN9-Q3xgV6K9lF-QFaAywdHmoklo,656
|
|
2
|
+
amiai_sdk/cli.py,sha256=xJHElzrAHz7TAvIXQRTMAZDIy18Wchf9fSSi4ELq7zU,5097
|
|
3
|
+
amiai_sdk/harness.py,sha256=ZOsvSshindvzHnjWpJc2L2FVnUO30zBVYkSogj_K0e4,10164
|
|
4
|
+
amiai_sdk/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
5
|
+
amiai_sdk/tool.py,sha256=cl6XG4EKgjAWqfTJRrB0KJkdxvM0HEmaGi5wY4vow0M,4124
|
|
6
|
+
amiai-0.1.0.dist-info/METADATA,sha256=YT_Ge4mxoLSoXPn7U5ax1g0J6ioA0_3KQ0QEolXS1Vk,8228
|
|
7
|
+
amiai-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
|
|
8
|
+
amiai-0.1.0.dist-info/entry_points.txt,sha256=uFFJcsvdAxpkTrDa_rXR6TgH_gXg-dwBhsV65NJtzu0,45
|
|
9
|
+
amiai-0.1.0.dist-info/RECORD,,
|
amiai_sdk/__init__.py
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
"""
|
|
2
|
+
AMIAI SDK - Build AI agents with local tool execution
|
|
3
|
+
|
|
4
|
+
Example:
|
|
5
|
+
from amiai_sdk import tool, Harness
|
|
6
|
+
|
|
7
|
+
@tool(
|
|
8
|
+
name="web_search",
|
|
9
|
+
description="Search the web for information"
|
|
10
|
+
)
|
|
11
|
+
async def search(query: str) -> dict:
|
|
12
|
+
# Your API key stays local
|
|
13
|
+
response = await fetch(f"https://api.example.com/search?q={query}")
|
|
14
|
+
return response.json()
|
|
15
|
+
|
|
16
|
+
# Start the harness
|
|
17
|
+
harness = Harness(api_key="amiai_xxx", tools=[search])
|
|
18
|
+
await harness.start()
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
from .tool import tool, Tool
|
|
22
|
+
from .harness import Harness, start_harness
|
|
23
|
+
|
|
24
|
+
__version__ = "0.1.0"
|
|
25
|
+
__all__ = ["tool", "Tool", "Harness", "start_harness"]
|
amiai_sdk/cli.py
ADDED
|
@@ -0,0 +1,177 @@
|
|
|
1
|
+
"""
|
|
2
|
+
AMIAI CLI - Run the tool harness
|
|
3
|
+
|
|
4
|
+
Usage:
|
|
5
|
+
amiai dev # Discover and run tools
|
|
6
|
+
amiai dev --pattern "*.py" # Custom pattern
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
import argparse
|
|
10
|
+
import asyncio
|
|
11
|
+
import importlib.util
|
|
12
|
+
import os
|
|
13
|
+
import sys
|
|
14
|
+
from glob import glob
|
|
15
|
+
from pathlib import Path
|
|
16
|
+
from typing import List
|
|
17
|
+
|
|
18
|
+
from .tool import Tool
|
|
19
|
+
from .harness import Harness
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def discover_tools(patterns: List[str]) -> List[Tool]:
|
|
23
|
+
"""Discover tools from Python files matching patterns."""
|
|
24
|
+
tools: List[Tool] = []
|
|
25
|
+
seen_files: set = set()
|
|
26
|
+
|
|
27
|
+
for pattern in patterns:
|
|
28
|
+
for file_path in glob(pattern, recursive=True):
|
|
29
|
+
abs_path = os.path.abspath(file_path)
|
|
30
|
+
if abs_path in seen_files:
|
|
31
|
+
continue
|
|
32
|
+
seen_files.add(abs_path)
|
|
33
|
+
|
|
34
|
+
# Skip __pycache__ and hidden files
|
|
35
|
+
if "__pycache__" in file_path or file_path.startswith("."):
|
|
36
|
+
continue
|
|
37
|
+
|
|
38
|
+
try:
|
|
39
|
+
tools.extend(load_tools_from_file(file_path))
|
|
40
|
+
except Exception as e:
|
|
41
|
+
print(f"Warning: Failed to load {file_path}: {e}")
|
|
42
|
+
|
|
43
|
+
return tools
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def load_tools_from_file(file_path: str) -> List[Tool]:
|
|
47
|
+
"""Load Tool instances from a Python file."""
|
|
48
|
+
tools: List[Tool] = []
|
|
49
|
+
|
|
50
|
+
spec = importlib.util.spec_from_file_location("tool_module", file_path)
|
|
51
|
+
if spec is None or spec.loader is None:
|
|
52
|
+
return tools
|
|
53
|
+
|
|
54
|
+
module = importlib.util.module_from_spec(spec)
|
|
55
|
+
|
|
56
|
+
try:
|
|
57
|
+
spec.loader.exec_module(module)
|
|
58
|
+
except Exception as e:
|
|
59
|
+
print(f"Warning: Error loading {file_path}: {e}")
|
|
60
|
+
return tools
|
|
61
|
+
|
|
62
|
+
# Find all Tool instances in the module
|
|
63
|
+
for name in dir(module):
|
|
64
|
+
obj = getattr(module, name)
|
|
65
|
+
if isinstance(obj, Tool):
|
|
66
|
+
tools.append(obj)
|
|
67
|
+
|
|
68
|
+
return tools
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
async def run_harness(api_key: str, tools: List[Tool], url: str) -> None:
|
|
72
|
+
"""Run the harness with discovered tools."""
|
|
73
|
+
print(f"\n{'='*50}")
|
|
74
|
+
print("AMIAI Tool Harness")
|
|
75
|
+
print(f"{'='*50}\n")
|
|
76
|
+
|
|
77
|
+
if not tools:
|
|
78
|
+
print("No tools found!")
|
|
79
|
+
print("\nCreate a tool file like this:")
|
|
80
|
+
print(" # tools/search.py")
|
|
81
|
+
print(" from amiai_sdk import tool")
|
|
82
|
+
print("")
|
|
83
|
+
print(" @tool(name='web_search', description='Search the web')")
|
|
84
|
+
print(" async def search(query: str) -> dict:")
|
|
85
|
+
print(" return {'results': [...]}")
|
|
86
|
+
return
|
|
87
|
+
|
|
88
|
+
print(f"Found {len(tools)} tool(s):")
|
|
89
|
+
for t in tools:
|
|
90
|
+
print(f" - {t.name}: {t.description[:50]}...")
|
|
91
|
+
print()
|
|
92
|
+
|
|
93
|
+
def on_connect(session_id: str) -> None:
|
|
94
|
+
print(f"✓ Connected to AMIAI")
|
|
95
|
+
print(f" Session ID: {session_id}")
|
|
96
|
+
print(f"\n✓ Waiting for invocations...\n")
|
|
97
|
+
|
|
98
|
+
def on_disconnect() -> None:
|
|
99
|
+
print("⚠ Disconnected from AMIAI")
|
|
100
|
+
|
|
101
|
+
def on_invoke(tool_name: str, args: dict) -> None:
|
|
102
|
+
print(f"→ Invoking: {tool_name}")
|
|
103
|
+
print(f" Args: {args}")
|
|
104
|
+
|
|
105
|
+
def on_result(tool_name: str, result: any) -> None:
|
|
106
|
+
result_str = str(result)
|
|
107
|
+
if len(result_str) > 100:
|
|
108
|
+
result_str = result_str[:100] + "..."
|
|
109
|
+
print(f"← Result: {result_str}")
|
|
110
|
+
|
|
111
|
+
def on_error(error: Exception) -> None:
|
|
112
|
+
print(f"✗ Error: {error}")
|
|
113
|
+
|
|
114
|
+
harness = Harness(
|
|
115
|
+
api_key=api_key,
|
|
116
|
+
tools=tools,
|
|
117
|
+
url=url,
|
|
118
|
+
on_connect=on_connect,
|
|
119
|
+
on_disconnect=on_disconnect,
|
|
120
|
+
on_invoke=on_invoke,
|
|
121
|
+
on_result=on_result,
|
|
122
|
+
on_error=on_error,
|
|
123
|
+
)
|
|
124
|
+
|
|
125
|
+
try:
|
|
126
|
+
await harness.run_forever()
|
|
127
|
+
except KeyboardInterrupt:
|
|
128
|
+
print("\n\nStopping...")
|
|
129
|
+
await harness.stop()
|
|
130
|
+
print("Done!")
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def main() -> None:
|
|
134
|
+
"""CLI entry point."""
|
|
135
|
+
parser = argparse.ArgumentParser(
|
|
136
|
+
description="AMIAI Tool Harness - Run local tools for AI agents"
|
|
137
|
+
)
|
|
138
|
+
subparsers = parser.add_subparsers(dest="command", help="Commands")
|
|
139
|
+
|
|
140
|
+
# dev command
|
|
141
|
+
dev_parser = subparsers.add_parser("dev", help="Run the tool harness")
|
|
142
|
+
dev_parser.add_argument(
|
|
143
|
+
"--pattern",
|
|
144
|
+
"-p",
|
|
145
|
+
action="append",
|
|
146
|
+
default=None,
|
|
147
|
+
help="Glob pattern for tool files (default: **/*.tool.py, **/tools/*.py)",
|
|
148
|
+
)
|
|
149
|
+
dev_parser.add_argument(
|
|
150
|
+
"--url",
|
|
151
|
+
default=os.environ.get("AMIAI_URL", "wss://backend.amiai.com/api/tools/live"),
|
|
152
|
+
help="AMIAI WebSocket URL",
|
|
153
|
+
)
|
|
154
|
+
|
|
155
|
+
args = parser.parse_args()
|
|
156
|
+
|
|
157
|
+
if args.command is None:
|
|
158
|
+
parser.print_help()
|
|
159
|
+
sys.exit(1)
|
|
160
|
+
|
|
161
|
+
if args.command == "dev":
|
|
162
|
+
api_key = os.environ.get("AMIAI_API_KEY")
|
|
163
|
+
if not api_key:
|
|
164
|
+
print("Error: AMIAI_API_KEY environment variable is required")
|
|
165
|
+
print("\nUsage:")
|
|
166
|
+
print(" export AMIAI_API_KEY=amiai_xxx")
|
|
167
|
+
print(" amiai dev")
|
|
168
|
+
sys.exit(1)
|
|
169
|
+
|
|
170
|
+
patterns = args.pattern or ["**/*.tool.py", "**/tools/*.py"]
|
|
171
|
+
tools = discover_tools(patterns)
|
|
172
|
+
|
|
173
|
+
asyncio.run(run_harness(api_key, tools, args.url))
|
|
174
|
+
|
|
175
|
+
|
|
176
|
+
if __name__ == "__main__":
|
|
177
|
+
main()
|
amiai_sdk/harness.py
ADDED
|
@@ -0,0 +1,321 @@
|
|
|
1
|
+
"""
|
|
2
|
+
AMIAI Harness - WebSocket client for local tool execution
|
|
3
|
+
|
|
4
|
+
Connects to AMIAI backend and registers tools. When the agent needs to
|
|
5
|
+
call a tool, the invocation is routed here and executed locally.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import asyncio
|
|
9
|
+
import json
|
|
10
|
+
import logging
|
|
11
|
+
from typing import Any, Callable, Dict, List, Optional
|
|
12
|
+
from dataclasses import dataclass, field
|
|
13
|
+
|
|
14
|
+
import websockets
|
|
15
|
+
from websockets.client import WebSocketClientProtocol
|
|
16
|
+
|
|
17
|
+
from .tool import Tool
|
|
18
|
+
|
|
19
|
+
logger = logging.getLogger("amiai_sdk")
|
|
20
|
+
|
|
21
|
+
DEFAULT_URL = "wss://backend.amiai.com/api/tools/live"
|
|
22
|
+
RECONNECT_DELAY_SEC = 5
|
|
23
|
+
HEARTBEAT_INTERVAL_SEC = 30
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
@dataclass
|
|
27
|
+
class HarnessOptions:
|
|
28
|
+
"""Configuration for the Harness."""
|
|
29
|
+
api_key: str
|
|
30
|
+
tools: List[Tool]
|
|
31
|
+
url: str = DEFAULT_URL
|
|
32
|
+
on_connect: Optional[Callable[[str], None]] = None
|
|
33
|
+
on_disconnect: Optional[Callable[[], None]] = None
|
|
34
|
+
on_invoke: Optional[Callable[[str, Dict[str, Any]], None]] = None
|
|
35
|
+
on_result: Optional[Callable[[str, Any], None]] = None
|
|
36
|
+
on_error: Optional[Callable[[Exception], None]] = None
|
|
37
|
+
auto_reconnect: bool = True
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class Harness:
|
|
41
|
+
"""
|
|
42
|
+
WebSocket client for local tool execution.
|
|
43
|
+
|
|
44
|
+
Example:
|
|
45
|
+
harness = Harness(
|
|
46
|
+
api_key="amiai_xxx",
|
|
47
|
+
tools=[my_tool],
|
|
48
|
+
on_connect=lambda sid: print(f"Connected: {sid}"),
|
|
49
|
+
)
|
|
50
|
+
await harness.start()
|
|
51
|
+
"""
|
|
52
|
+
|
|
53
|
+
def __init__(
|
|
54
|
+
self,
|
|
55
|
+
api_key: str,
|
|
56
|
+
tools: List[Tool],
|
|
57
|
+
url: str = DEFAULT_URL,
|
|
58
|
+
on_connect: Optional[Callable[[str], None]] = None,
|
|
59
|
+
on_disconnect: Optional[Callable[[], None]] = None,
|
|
60
|
+
on_invoke: Optional[Callable[[str, Dict[str, Any]], None]] = None,
|
|
61
|
+
on_result: Optional[Callable[[str, Any], None]] = None,
|
|
62
|
+
on_error: Optional[Callable[[Exception], None]] = None,
|
|
63
|
+
auto_reconnect: bool = True,
|
|
64
|
+
):
|
|
65
|
+
self.api_key = api_key
|
|
66
|
+
self.tools = {t.name: t for t in tools}
|
|
67
|
+
self.tools_list = tools
|
|
68
|
+
self.url = url
|
|
69
|
+
self.on_connect = on_connect
|
|
70
|
+
self.on_disconnect = on_disconnect
|
|
71
|
+
self.on_invoke = on_invoke
|
|
72
|
+
self.on_result = on_result
|
|
73
|
+
self.on_error = on_error
|
|
74
|
+
self.auto_reconnect = auto_reconnect
|
|
75
|
+
|
|
76
|
+
self._ws: Optional[WebSocketClientProtocol] = None
|
|
77
|
+
self._session_id: Optional[str] = None
|
|
78
|
+
self._running = False
|
|
79
|
+
self._heartbeat_task: Optional[asyncio.Task] = None
|
|
80
|
+
self._receive_task: Optional[asyncio.Task] = None
|
|
81
|
+
|
|
82
|
+
@property
|
|
83
|
+
def session_id(self) -> Optional[str]:
|
|
84
|
+
"""Get the current session ID."""
|
|
85
|
+
return self._session_id
|
|
86
|
+
|
|
87
|
+
@property
|
|
88
|
+
def is_connected(self) -> bool:
|
|
89
|
+
"""Check if connected to AMIAI."""
|
|
90
|
+
return self._ws is not None and self._session_id is not None
|
|
91
|
+
|
|
92
|
+
async def start(self) -> str:
|
|
93
|
+
"""
|
|
94
|
+
Start the harness and connect to AMIAI.
|
|
95
|
+
|
|
96
|
+
Returns:
|
|
97
|
+
The session ID
|
|
98
|
+
"""
|
|
99
|
+
self._running = True
|
|
100
|
+
return await self._connect()
|
|
101
|
+
|
|
102
|
+
async def stop(self) -> None:
|
|
103
|
+
"""Stop the harness and disconnect."""
|
|
104
|
+
self._running = False
|
|
105
|
+
self.auto_reconnect = False
|
|
106
|
+
|
|
107
|
+
if self._heartbeat_task:
|
|
108
|
+
self._heartbeat_task.cancel()
|
|
109
|
+
try:
|
|
110
|
+
await self._heartbeat_task
|
|
111
|
+
except asyncio.CancelledError:
|
|
112
|
+
pass
|
|
113
|
+
|
|
114
|
+
if self._receive_task:
|
|
115
|
+
self._receive_task.cancel()
|
|
116
|
+
try:
|
|
117
|
+
await self._receive_task
|
|
118
|
+
except asyncio.CancelledError:
|
|
119
|
+
pass
|
|
120
|
+
|
|
121
|
+
if self._ws:
|
|
122
|
+
await self._ws.close()
|
|
123
|
+
self._ws = None
|
|
124
|
+
|
|
125
|
+
self._session_id = None
|
|
126
|
+
|
|
127
|
+
async def run_forever(self) -> None:
|
|
128
|
+
"""Run the harness until stopped."""
|
|
129
|
+
await self.start()
|
|
130
|
+
|
|
131
|
+
while self._running:
|
|
132
|
+
await asyncio.sleep(1)
|
|
133
|
+
|
|
134
|
+
async def _connect(self) -> str:
|
|
135
|
+
"""Connect to AMIAI and register tools."""
|
|
136
|
+
try:
|
|
137
|
+
self._ws = await websockets.connect(self.url)
|
|
138
|
+
|
|
139
|
+
# Send hello
|
|
140
|
+
await self._send({
|
|
141
|
+
"type": "hello",
|
|
142
|
+
"apiKey": self.api_key,
|
|
143
|
+
"tools": [t.to_manifest() for t in self.tools_list],
|
|
144
|
+
})
|
|
145
|
+
|
|
146
|
+
# Wait for ack
|
|
147
|
+
response = await self._ws.recv()
|
|
148
|
+
message = json.loads(response)
|
|
149
|
+
|
|
150
|
+
if message.get("type") == "ack":
|
|
151
|
+
self._session_id = message.get("sessionId")
|
|
152
|
+
if self._session_id and self.on_connect:
|
|
153
|
+
self.on_connect(self._session_id)
|
|
154
|
+
|
|
155
|
+
# Start background tasks
|
|
156
|
+
self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
|
|
157
|
+
self._receive_task = asyncio.create_task(self._receive_loop())
|
|
158
|
+
|
|
159
|
+
return self._session_id or ""
|
|
160
|
+
elif message.get("type") == "error":
|
|
161
|
+
raise Exception(f"Connection error: {message.get('message')}")
|
|
162
|
+
else:
|
|
163
|
+
raise Exception(f"Unexpected response: {message}")
|
|
164
|
+
|
|
165
|
+
except Exception as e:
|
|
166
|
+
if self.on_error:
|
|
167
|
+
self.on_error(e)
|
|
168
|
+
raise
|
|
169
|
+
|
|
170
|
+
async def _send(self, message: Dict[str, Any]) -> None:
|
|
171
|
+
"""Send a message to the server."""
|
|
172
|
+
if self._ws:
|
|
173
|
+
await self._ws.send(json.dumps(message))
|
|
174
|
+
|
|
175
|
+
async def _heartbeat_loop(self) -> None:
|
|
176
|
+
"""Send periodic heartbeats."""
|
|
177
|
+
while self._running and self._ws:
|
|
178
|
+
try:
|
|
179
|
+
await asyncio.sleep(HEARTBEAT_INTERVAL_SEC)
|
|
180
|
+
await self._send({"type": "heartbeat"})
|
|
181
|
+
except asyncio.CancelledError:
|
|
182
|
+
break
|
|
183
|
+
except Exception as e:
|
|
184
|
+
logger.debug(f"Heartbeat error: {e}")
|
|
185
|
+
|
|
186
|
+
async def _receive_loop(self) -> None:
|
|
187
|
+
"""Receive and handle messages from server."""
|
|
188
|
+
while self._running and self._ws:
|
|
189
|
+
try:
|
|
190
|
+
message_str = await self._ws.recv()
|
|
191
|
+
message = json.loads(message_str)
|
|
192
|
+
await self._handle_message(message)
|
|
193
|
+
except websockets.ConnectionClosed:
|
|
194
|
+
logger.info("Connection closed")
|
|
195
|
+
self._session_id = None
|
|
196
|
+
if self.on_disconnect:
|
|
197
|
+
self.on_disconnect()
|
|
198
|
+
|
|
199
|
+
if self.auto_reconnect and self._running:
|
|
200
|
+
await self._reconnect()
|
|
201
|
+
break
|
|
202
|
+
except asyncio.CancelledError:
|
|
203
|
+
break
|
|
204
|
+
except Exception as e:
|
|
205
|
+
logger.error(f"Receive error: {e}")
|
|
206
|
+
if self.on_error:
|
|
207
|
+
self.on_error(e)
|
|
208
|
+
|
|
209
|
+
async def _reconnect(self) -> None:
|
|
210
|
+
"""Attempt to reconnect after disconnection."""
|
|
211
|
+
logger.info(f"Reconnecting in {RECONNECT_DELAY_SEC}s...")
|
|
212
|
+
await asyncio.sleep(RECONNECT_DELAY_SEC)
|
|
213
|
+
|
|
214
|
+
if self._running:
|
|
215
|
+
try:
|
|
216
|
+
await self._connect()
|
|
217
|
+
except Exception as e:
|
|
218
|
+
logger.error(f"Reconnection failed: {e}")
|
|
219
|
+
if self._running:
|
|
220
|
+
await self._reconnect()
|
|
221
|
+
|
|
222
|
+
async def _handle_message(self, message: Dict[str, Any]) -> None:
|
|
223
|
+
"""Handle an incoming message."""
|
|
224
|
+
msg_type = message.get("type")
|
|
225
|
+
|
|
226
|
+
if msg_type == "invoke":
|
|
227
|
+
call_id = message.get("callId")
|
|
228
|
+
tool_name = message.get("tool")
|
|
229
|
+
args = message.get("args", {})
|
|
230
|
+
|
|
231
|
+
if call_id and tool_name:
|
|
232
|
+
await self._handle_invoke(call_id, tool_name, args)
|
|
233
|
+
|
|
234
|
+
elif msg_type == "cancel":
|
|
235
|
+
call_id = message.get("callId")
|
|
236
|
+
logger.info(f"Cancellation requested for {call_id}")
|
|
237
|
+
|
|
238
|
+
elif msg_type == "ping":
|
|
239
|
+
await self._send({"type": "heartbeat"})
|
|
240
|
+
|
|
241
|
+
elif msg_type == "error":
|
|
242
|
+
error_msg = message.get("message", "Unknown error")
|
|
243
|
+
logger.error(f"Server error: {error_msg}")
|
|
244
|
+
if self.on_error:
|
|
245
|
+
self.on_error(Exception(error_msg))
|
|
246
|
+
|
|
247
|
+
async def _handle_invoke(
|
|
248
|
+
self,
|
|
249
|
+
call_id: str,
|
|
250
|
+
tool_name: str,
|
|
251
|
+
args: Dict[str, Any],
|
|
252
|
+
) -> None:
|
|
253
|
+
"""Handle a tool invocation."""
|
|
254
|
+
tool = self.tools.get(tool_name)
|
|
255
|
+
|
|
256
|
+
if not tool:
|
|
257
|
+
await self._send({
|
|
258
|
+
"type": "error",
|
|
259
|
+
"callId": call_id,
|
|
260
|
+
"message": f"Tool not found: {tool_name}",
|
|
261
|
+
})
|
|
262
|
+
return
|
|
263
|
+
|
|
264
|
+
if self.on_invoke:
|
|
265
|
+
self.on_invoke(tool_name, args)
|
|
266
|
+
|
|
267
|
+
try:
|
|
268
|
+
# Send started event
|
|
269
|
+
await self._send({
|
|
270
|
+
"type": "event",
|
|
271
|
+
"callId": call_id,
|
|
272
|
+
"event": "started",
|
|
273
|
+
"data": {"tool": tool_name},
|
|
274
|
+
})
|
|
275
|
+
|
|
276
|
+
# Execute the tool
|
|
277
|
+
if asyncio.iscoroutinefunction(tool.execute):
|
|
278
|
+
result = await tool.execute(**args)
|
|
279
|
+
else:
|
|
280
|
+
result = tool.execute(**args)
|
|
281
|
+
|
|
282
|
+
# Send result
|
|
283
|
+
await self._send({
|
|
284
|
+
"type": "result",
|
|
285
|
+
"callId": call_id,
|
|
286
|
+
"output": result,
|
|
287
|
+
})
|
|
288
|
+
|
|
289
|
+
if self.on_result:
|
|
290
|
+
self.on_result(tool_name, result)
|
|
291
|
+
|
|
292
|
+
except Exception as e:
|
|
293
|
+
error_msg = str(e)
|
|
294
|
+
await self._send({
|
|
295
|
+
"type": "error",
|
|
296
|
+
"callId": call_id,
|
|
297
|
+
"message": error_msg,
|
|
298
|
+
})
|
|
299
|
+
|
|
300
|
+
if self.on_error:
|
|
301
|
+
self.on_error(e)
|
|
302
|
+
|
|
303
|
+
|
|
304
|
+
async def start_harness(
|
|
305
|
+
api_key: str,
|
|
306
|
+
tools: List[Tool],
|
|
307
|
+
url: str = DEFAULT_URL,
|
|
308
|
+
**kwargs,
|
|
309
|
+
) -> Harness:
|
|
310
|
+
"""
|
|
311
|
+
Create and start a harness.
|
|
312
|
+
|
|
313
|
+
Example:
|
|
314
|
+
harness = await start_harness(
|
|
315
|
+
api_key="amiai_xxx",
|
|
316
|
+
tools=[my_tool],
|
|
317
|
+
)
|
|
318
|
+
"""
|
|
319
|
+
harness = Harness(api_key=api_key, tools=tools, url=url, **kwargs)
|
|
320
|
+
await harness.start()
|
|
321
|
+
return harness
|
amiai_sdk/py.typed
ADDED
|
File without changes
|
amiai_sdk/tool.py
ADDED
|
@@ -0,0 +1,147 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Tool definition and decorator for AMIAI SDK
|
|
3
|
+
|
|
4
|
+
Example:
|
|
5
|
+
from amiai_sdk import tool
|
|
6
|
+
|
|
7
|
+
@tool(name="calculator", description="Evaluate math expressions")
|
|
8
|
+
async def calculate(expression: str) -> dict:
|
|
9
|
+
result = eval(expression)
|
|
10
|
+
return {"expression": expression, "result": result}
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from typing import Any, Callable, Dict, List, Optional, Type, get_type_hints
|
|
14
|
+
from dataclasses import dataclass
|
|
15
|
+
import inspect
|
|
16
|
+
import json
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@dataclass
|
|
20
|
+
class InputField:
|
|
21
|
+
"""Schema for a tool input field."""
|
|
22
|
+
type: str
|
|
23
|
+
description: Optional[str] = None
|
|
24
|
+
required: bool = True
|
|
25
|
+
default: Any = None
|
|
26
|
+
enum: Optional[List[str]] = None
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
@dataclass
|
|
30
|
+
class Tool:
|
|
31
|
+
"""A tool that can be registered with AMIAI."""
|
|
32
|
+
name: str
|
|
33
|
+
description: str
|
|
34
|
+
input_schema: Dict[str, Any]
|
|
35
|
+
execute: Callable[..., Any]
|
|
36
|
+
|
|
37
|
+
def to_manifest(self) -> Dict[str, Any]:
|
|
38
|
+
"""Convert to manifest format for registration."""
|
|
39
|
+
return {
|
|
40
|
+
"name": self.name,
|
|
41
|
+
"description": self.description,
|
|
42
|
+
"inputSchema": self.input_schema,
|
|
43
|
+
}
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _python_type_to_json_type(python_type: Type) -> str:
|
|
47
|
+
"""Convert Python type hints to JSON Schema types."""
|
|
48
|
+
type_map = {
|
|
49
|
+
str: "string",
|
|
50
|
+
int: "number",
|
|
51
|
+
float: "number",
|
|
52
|
+
bool: "boolean",
|
|
53
|
+
list: "array",
|
|
54
|
+
dict: "object",
|
|
55
|
+
List: "array",
|
|
56
|
+
Dict: "object",
|
|
57
|
+
}
|
|
58
|
+
|
|
59
|
+
# Handle generic types like List[str], Dict[str, Any]
|
|
60
|
+
origin = getattr(python_type, "__origin__", None)
|
|
61
|
+
if origin is not None:
|
|
62
|
+
if origin is list:
|
|
63
|
+
return "array"
|
|
64
|
+
if origin is dict:
|
|
65
|
+
return "object"
|
|
66
|
+
|
|
67
|
+
return type_map.get(python_type, "string")
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def _derive_schema_from_function(func: Callable) -> Dict[str, Any]:
|
|
71
|
+
"""Derive JSON Schema from function signature."""
|
|
72
|
+
sig = inspect.signature(func)
|
|
73
|
+
hints = get_type_hints(func) if hasattr(func, "__annotations__") else {}
|
|
74
|
+
|
|
75
|
+
properties: Dict[str, Any] = {}
|
|
76
|
+
required: List[str] = []
|
|
77
|
+
|
|
78
|
+
for name, param in sig.parameters.items():
|
|
79
|
+
if name in ("self", "cls"):
|
|
80
|
+
continue
|
|
81
|
+
|
|
82
|
+
# Get type from hints or default to string
|
|
83
|
+
param_type = hints.get(name, str)
|
|
84
|
+
json_type = _python_type_to_json_type(param_type)
|
|
85
|
+
|
|
86
|
+
properties[name] = {
|
|
87
|
+
"type": json_type,
|
|
88
|
+
}
|
|
89
|
+
|
|
90
|
+
# Check if parameter has a default
|
|
91
|
+
if param.default is inspect.Parameter.empty:
|
|
92
|
+
required.append(name)
|
|
93
|
+
else:
|
|
94
|
+
properties[name]["default"] = param.default
|
|
95
|
+
|
|
96
|
+
return {
|
|
97
|
+
"type": "object",
|
|
98
|
+
"properties": properties,
|
|
99
|
+
"required": required,
|
|
100
|
+
}
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
def tool(
|
|
104
|
+
name: str,
|
|
105
|
+
description: str,
|
|
106
|
+
input_schema: Optional[Dict[str, Any]] = None,
|
|
107
|
+
) -> Callable[[Callable], Tool]:
|
|
108
|
+
"""
|
|
109
|
+
Decorator to define a tool for AMIAI agents.
|
|
110
|
+
|
|
111
|
+
Args:
|
|
112
|
+
name: Unique name for the tool
|
|
113
|
+
description: Description shown to the AI
|
|
114
|
+
input_schema: Optional JSON Schema for inputs (auto-derived if not provided)
|
|
115
|
+
|
|
116
|
+
Example:
|
|
117
|
+
@tool(name="web_search", description="Search the web")
|
|
118
|
+
async def search(query: str) -> dict:
|
|
119
|
+
response = await fetch(f"https://api.example.com/search?q={query}")
|
|
120
|
+
return response.json()
|
|
121
|
+
|
|
122
|
+
Example with explicit schema:
|
|
123
|
+
@tool(
|
|
124
|
+
name="calculator",
|
|
125
|
+
description="Evaluate math expressions",
|
|
126
|
+
input_schema={
|
|
127
|
+
"type": "object",
|
|
128
|
+
"properties": {
|
|
129
|
+
"expression": {"type": "string", "description": "e.g., 2+2"}
|
|
130
|
+
},
|
|
131
|
+
"required": ["expression"]
|
|
132
|
+
}
|
|
133
|
+
)
|
|
134
|
+
async def calculate(expression: str) -> dict:
|
|
135
|
+
return {"result": eval(expression)}
|
|
136
|
+
"""
|
|
137
|
+
def decorator(func: Callable) -> Tool:
|
|
138
|
+
schema = input_schema or _derive_schema_from_function(func)
|
|
139
|
+
|
|
140
|
+
return Tool(
|
|
141
|
+
name=name,
|
|
142
|
+
description=description,
|
|
143
|
+
input_schema=schema,
|
|
144
|
+
execute=func,
|
|
145
|
+
)
|
|
146
|
+
|
|
147
|
+
return decorator
|