mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
feat: add hanzo-tools-api package and hanzo.infra SDK
Add generic REST API tool for calling any API via OpenAPI specs: - 30+ built-in provider configs with env var auto-detection - 1100+ auto-generated providers from APIs.guru + oapis.org - Secure credential management with pluggable storage - OpenAPI spec parsing with ETag caching - Agent-friendly interface with search, overview, call actions Add hanzo.infra unified infrastructure SDK: - Vector (Qdrant), KV (Redis/Valkey), DocumentDB (MongoDB) - Storage (S3/MinIO), Search (Meilisearch), PubSub (NATS) - Tasks (Temporal), Queues, Cron, Functions (Nuclio) - Async clients with lazy initialization and health checks - Optional deps via [infra-*] extras Bump hanzo-mcp to v0.11.4 with api tool integration.
This commit is contained in:
@@ -0,0 +1,91 @@
|
||||
name: Generate API Providers
|
||||
|
||||
on:
|
||||
schedule:
|
||||
# Run weekly on Sunday at 00:00 UTC
|
||||
- cron: '0 0 * * 0'
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
create_pr:
|
||||
description: 'Create a PR with the changes'
|
||||
type: boolean
|
||||
default: true
|
||||
preload_specs:
|
||||
description: 'Pre-download popular OpenAPI specs'
|
||||
type: boolean
|
||||
default: false
|
||||
|
||||
jobs:
|
||||
generate:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
pip install httpx aiofiles pyyaml
|
||||
|
||||
- name: Generate providers from APIs.guru + oapis.org
|
||||
run: |
|
||||
cd pkg/hanzo-tools-api
|
||||
python scripts/generate_providers.py > hanzo_tools/api/apis_guru_providers.py
|
||||
|
||||
- name: Pre-download popular specs
|
||||
if: github.event.inputs.preload_specs == 'true'
|
||||
run: |
|
||||
cd pkg/hanzo-tools-api
|
||||
pip install -e .
|
||||
python scripts/preload_specs.py
|
||||
echo "Specs cached at ~/.hanzo/api/specs/"
|
||||
ls -la ~/.hanzo/api/specs/ | head -20
|
||||
|
||||
- name: Check for changes
|
||||
id: changes
|
||||
run: |
|
||||
if git diff --quiet pkg/hanzo-tools-api/hanzo_tools/api/apis_guru_providers.py; then
|
||||
echo "changed=false" >> $GITHUB_OUTPUT
|
||||
else
|
||||
echo "changed=true" >> $GITHUB_OUTPUT
|
||||
echo "Provider configs changed"
|
||||
git diff --stat pkg/hanzo-tools-api/hanzo_tools/api/apis_guru_providers.py
|
||||
fi
|
||||
|
||||
- name: Create Pull Request
|
||||
if: steps.changes.outputs.changed == 'true' && (github.event.inputs.create_pr == 'true' || github.event_name == 'schedule')
|
||||
uses: peter-evans/create-pull-request@v6
|
||||
with:
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
commit-message: 'chore(api): regenerate provider configs from APIs.guru + oapis.org'
|
||||
title: 'chore(api): Update API provider configurations'
|
||||
body: |
|
||||
## Summary
|
||||
Auto-generated provider configurations from:
|
||||
- [APIs.guru](https://apis.guru/) - OpenAPI specification directory (2500+ APIs)
|
||||
- [oapis.org](https://oapis.org/) - LLM-optimized API descriptions
|
||||
|
||||
This PR updates the `apis_guru_providers.py` file with the latest API configurations.
|
||||
|
||||
## Stats
|
||||
- Total providers: ~1100+
|
||||
- With OpenAPI specs: ~1100+
|
||||
- LLM-optimized descriptions for popular APIs
|
||||
|
||||
## Changes
|
||||
- Updated provider configurations
|
||||
- Refreshed spec URLs
|
||||
- Updated LLM-friendly descriptions for popular APIs
|
||||
branch: chore/update-api-providers
|
||||
delete-branch: true
|
||||
labels: |
|
||||
automated
|
||||
dependencies
|
||||
@@ -49,6 +49,7 @@ PACKAGE_TOOL_PREFIXES: dict[str, list[str]] = {
|
||||
"config": ["config", "mode"],
|
||||
"mcp_tools": ["mcp"], # Consolidated MCP tool
|
||||
"computer": ["computer", "screen"],
|
||||
"api": ["api"], # Generic API tool for OpenAPI specs
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hanzo-mcp"
|
||||
version = "0.11.0"
|
||||
version = "0.11.4"
|
||||
description = "The Zen of Hanzo MCP: One server to rule them all. The ultimate MCP that orchestrates all others."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
@@ -39,6 +39,7 @@ dependencies = [
|
||||
"hanzo-tools-refactor>=0.1.0",
|
||||
"hanzo-tools-llm>=0.1.0",
|
||||
"hanzo-tools-agent>=0.3.1",
|
||||
"hanzo-tools-api>=0.3.1",
|
||||
"hanzo-persona>=1.0.0",
|
||||
"uvloop>=0.22.1",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
# hanzo-tools-api
|
||||
|
||||
Generic API tool for calling any REST API via OpenAPI specs.
|
||||
|
||||
## Features
|
||||
|
||||
- **Credential Management**: Securely store and manage API keys with env var fallback
|
||||
- **OpenAPI Support**: Parse OpenAPI 3.x specs to discover and call operations
|
||||
- **Auto-Detection**: Automatically detects credentials from environment variables
|
||||
- **Multi-Provider**: Built-in support for 30+ cloud providers
|
||||
- **Raw Requests**: Make raw HTTP requests when you need more control
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install hanzo-tools-api
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Using Environment Variables (Recommended)
|
||||
|
||||
Set your API keys as environment variables:
|
||||
|
||||
```bash
|
||||
export CLOUDFLARE_API_TOKEN=your-token
|
||||
export GITHUB_TOKEN=ghp_xxx
|
||||
export STRIPE_API_KEY=sk_xxx
|
||||
```
|
||||
|
||||
The tool automatically detects these when you make calls.
|
||||
|
||||
### Using the Tool
|
||||
|
||||
```python
|
||||
from hanzo_tools.api import APITool
|
||||
|
||||
tool = APITool()
|
||||
|
||||
# List available providers and their status
|
||||
await tool.call(ctx, action="list")
|
||||
|
||||
# Configure a provider manually
|
||||
await tool.call(ctx,
|
||||
action="config",
|
||||
provider="cloudflare",
|
||||
api_key="your-key"
|
||||
)
|
||||
|
||||
# Load OpenAPI spec for a provider
|
||||
await tool.call(ctx, action="spec", provider="cloudflare")
|
||||
|
||||
# List available operations
|
||||
await tool.call(ctx, action="ops", provider="cloudflare", search="zones")
|
||||
|
||||
# Call an operation
|
||||
await tool.call(ctx,
|
||||
action="call",
|
||||
provider="cloudflare",
|
||||
operation="listZones"
|
||||
)
|
||||
|
||||
# Make a raw request
|
||||
await tool.call(ctx,
|
||||
action="raw",
|
||||
provider="github",
|
||||
method="GET",
|
||||
path="/user/repos"
|
||||
)
|
||||
```
|
||||
|
||||
### MCP Tool Usage
|
||||
|
||||
When registered with hanzo-mcp, use like:
|
||||
|
||||
```
|
||||
api # List all providers
|
||||
api --action config --provider github --api_key ghp_xxx
|
||||
api --action ops --provider cloudflare --search zones
|
||||
api --action call --provider cloudflare --operation listZones
|
||||
api --action raw --provider github --method GET --path /user/repos
|
||||
```
|
||||
|
||||
## Supported Providers
|
||||
|
||||
Built-in configurations for:
|
||||
|
||||
| Provider | Env Variables |
|
||||
|----------|---------------|
|
||||
| Cloudflare | `CLOUDFLARE_API_TOKEN`, `CF_API_TOKEN` |
|
||||
| GitHub | `GITHUB_TOKEN`, `GH_TOKEN` |
|
||||
| Stripe | `STRIPE_API_KEY` |
|
||||
| OpenAI | `OPENAI_API_KEY` |
|
||||
| Anthropic | `ANTHROPIC_API_KEY` |
|
||||
| Vercel | `VERCEL_TOKEN` |
|
||||
| DigitalOcean | `DIGITALOCEAN_TOKEN`, `DO_TOKEN` |
|
||||
| Fly.io | `FLY_API_TOKEN` |
|
||||
| Supabase | `SUPABASE_API_KEY` |
|
||||
| AWS | `AWS_ACCESS_KEY_ID` |
|
||||
| GCP | `GOOGLE_API_KEY` |
|
||||
| Azure | `AZURE_API_KEY` |
|
||||
| Slack | `SLACK_TOKEN` |
|
||||
| Discord | `DISCORD_TOKEN` |
|
||||
| ... and more |
|
||||
|
||||
## API Reference
|
||||
|
||||
### Actions
|
||||
|
||||
- **list**: Show all providers and their configuration status
|
||||
- **config**: Set credentials for a provider
|
||||
- **delete**: Remove credentials for a provider
|
||||
- **spec**: Load/refresh OpenAPI spec for a provider
|
||||
- **ops**: List available operations for a provider
|
||||
- **call**: Call an API operation by operation ID
|
||||
- **raw**: Make a raw HTTP request to any endpoint
|
||||
|
||||
### Parameters
|
||||
|
||||
| Parameter | Type | Description |
|
||||
|-----------|------|-------------|
|
||||
| `action` | str | Action to perform (default: "list") |
|
||||
| `provider` | str | Provider name (e.g., "cloudflare") |
|
||||
| `api_key` | str | API key for config action |
|
||||
| `api_secret` | str | API secret (if needed) |
|
||||
| `account_id` | str | Account/org ID |
|
||||
| `base_url` | str | Override base URL |
|
||||
| `spec_url` | str | URL to OpenAPI spec |
|
||||
| `operation` | str | Operation ID for call action |
|
||||
| `params` | str | JSON parameters |
|
||||
| `body` | str | JSON request body |
|
||||
| `method` | str | HTTP method for raw action |
|
||||
| `path` | str | URL path for raw action |
|
||||
| `search` | str | Search filter for ops action |
|
||||
| `tag` | str | Tag filter for ops action |
|
||||
|
||||
## Credential Storage
|
||||
|
||||
Credentials are stored in `~/.hanzo/api/credentials.json` with basic obfuscation.
|
||||
For production use, consider using system keyring or a secrets manager.
|
||||
|
||||
OpenAPI specs are cached in `~/.hanzo/api/specs/`.
|
||||
|
||||
## Adding Custom Providers
|
||||
|
||||
You can add any provider by providing a base URL and API key:
|
||||
|
||||
```python
|
||||
from hanzo_tools.api import get_credential_manager
|
||||
|
||||
cred_manager = get_credential_manager()
|
||||
cred_manager.set_credential(
|
||||
provider="custom-api",
|
||||
api_key="your-key",
|
||||
base_url="https://api.custom.com/v1"
|
||||
)
|
||||
```
|
||||
|
||||
Then load a spec:
|
||||
|
||||
```python
|
||||
from hanzo_tools.api import get_client
|
||||
|
||||
client = await get_client(
|
||||
"custom-api",
|
||||
spec_url="https://api.custom.com/openapi.json"
|
||||
)
|
||||
```
|
||||
|
||||
## Security Notes
|
||||
|
||||
- API keys are stored with basic base64 obfuscation (not encryption)
|
||||
- Credentials file has restrictive permissions (0600)
|
||||
- Environment variables are preferred for sensitive credentials
|
||||
- Never commit credentials to version control
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,4 @@
|
||||
# Namespace package
|
||||
from pkgutil import extend_path
|
||||
|
||||
__path__ = extend_path(__path__, __name__)
|
||||
@@ -0,0 +1,163 @@
|
||||
"""Hanzo API Tools - Generic API tool for calling any REST API.
|
||||
|
||||
This package provides a unified interface for calling any REST API
|
||||
via OpenAPI specs with automatic credential management.
|
||||
|
||||
Key Features:
|
||||
- Auto-detection of 30+ cloud provider credentials from environment
|
||||
- OpenAPI spec parsing with smart caching (ETag support)
|
||||
- Structured errors with recovery hints
|
||||
- Pluggable credential storage (memory, file, environment)
|
||||
- Agent-friendly typed interface
|
||||
|
||||
Quick Start:
|
||||
from hanzo_tools.api import APIClient
|
||||
|
||||
async with APIClient() as client:
|
||||
# List providers (auto-detects CLOUDFLARE_API_TOKEN, etc.)
|
||||
providers = await client.list_providers()
|
||||
|
||||
# Configure if needed
|
||||
await client.config("cloudflare", api_key="...")
|
||||
|
||||
# Load spec and call operations
|
||||
await client.spec("cloudflare")
|
||||
result = await client.call("cloudflare", "listZones")
|
||||
|
||||
MCP Tool Usage:
|
||||
from hanzo_tools.api import APITool, TOOLS
|
||||
|
||||
# The tool handles everything via actions
|
||||
tool = APITool()
|
||||
result = await tool.call(ctx, action="list")
|
||||
result = await tool.call(ctx, action="call", provider="cloudflare", operation="listZones")
|
||||
"""
|
||||
|
||||
# Models (structured data types)
|
||||
from .models import (
|
||||
AuthType,
|
||||
Credential,
|
||||
CredentialSource,
|
||||
EffectiveCredential,
|
||||
ProviderConfig,
|
||||
ProviderStatus,
|
||||
ProviderListResult,
|
||||
Operation,
|
||||
Parameter,
|
||||
OperationListResult,
|
||||
APICallResult,
|
||||
ToolSchema,
|
||||
ToolParameter,
|
||||
)
|
||||
|
||||
# Errors (structured errors with hints)
|
||||
from .errors import (
|
||||
APIError,
|
||||
CredentialError,
|
||||
OperationNotFoundError,
|
||||
ProviderNotFoundError,
|
||||
SpecNotLoadedError,
|
||||
SpecLoadError,
|
||||
AuthenticationError,
|
||||
RateLimitError,
|
||||
ValidationError,
|
||||
)
|
||||
|
||||
# Storage (pluggable credential storage)
|
||||
from .storage import (
|
||||
CredentialStorage,
|
||||
MemoryCredentialStorage,
|
||||
FileCredentialStorage,
|
||||
EnvironmentCredentialStorage,
|
||||
ChainedCredentialStorage,
|
||||
)
|
||||
|
||||
# Providers (provider configs and env var mappings)
|
||||
from .providers import (
|
||||
PROVIDER_CONFIGS,
|
||||
ENV_VAR_MAPPINGS,
|
||||
get_provider_config,
|
||||
list_providers as list_provider_names,
|
||||
)
|
||||
|
||||
# Credentials (credential management)
|
||||
from .credentials import (
|
||||
CredentialManager,
|
||||
get_credential_manager,
|
||||
reset_credential_manager,
|
||||
)
|
||||
|
||||
# OpenAPI (spec parsing and caching)
|
||||
from .openapi_client import (
|
||||
OpenAPIClient,
|
||||
SpecCache,
|
||||
get_client,
|
||||
clear_clients,
|
||||
)
|
||||
|
||||
# Main client (typed interface)
|
||||
from .client import (
|
||||
APIClient,
|
||||
get_api_client,
|
||||
reset_api_client,
|
||||
)
|
||||
|
||||
# MCP Tool
|
||||
from .api_tool import APITool
|
||||
|
||||
# Tools list for entry point discovery - must be classes, not instances
|
||||
TOOLS = [APITool]
|
||||
|
||||
__all__ = [
|
||||
# Models
|
||||
"AuthType",
|
||||
"Credential",
|
||||
"CredentialSource",
|
||||
"EffectiveCredential",
|
||||
"ProviderConfig",
|
||||
"ProviderStatus",
|
||||
"ProviderListResult",
|
||||
"Operation",
|
||||
"Parameter",
|
||||
"OperationListResult",
|
||||
"APICallResult",
|
||||
"ToolSchema",
|
||||
"ToolParameter",
|
||||
# Errors
|
||||
"APIError",
|
||||
"CredentialError",
|
||||
"OperationNotFoundError",
|
||||
"ProviderNotFoundError",
|
||||
"SpecNotLoadedError",
|
||||
"SpecLoadError",
|
||||
"AuthenticationError",
|
||||
"RateLimitError",
|
||||
"ValidationError",
|
||||
# Storage
|
||||
"CredentialStorage",
|
||||
"MemoryCredentialStorage",
|
||||
"FileCredentialStorage",
|
||||
"EnvironmentCredentialStorage",
|
||||
"ChainedCredentialStorage",
|
||||
# Providers
|
||||
"PROVIDER_CONFIGS",
|
||||
"ENV_VAR_MAPPINGS",
|
||||
"get_provider_config",
|
||||
"list_provider_names",
|
||||
# Credentials
|
||||
"CredentialManager",
|
||||
"get_credential_manager",
|
||||
"reset_credential_manager",
|
||||
# OpenAPI
|
||||
"OpenAPIClient",
|
||||
"SpecCache",
|
||||
"get_client",
|
||||
"clear_clients",
|
||||
# Client
|
||||
"APIClient",
|
||||
"get_api_client",
|
||||
"reset_api_client",
|
||||
# Tool
|
||||
"APITool",
|
||||
"TOOLS",
|
||||
]
|
||||
@@ -0,0 +1,658 @@
|
||||
"""Unified API tool for calling any REST API.
|
||||
|
||||
This is a thin MCP wrapper over APIClient. The real logic lives in client.py.
|
||||
|
||||
Provides a single tool interface for:
|
||||
- Managing API credentials
|
||||
- Loading OpenAPI specs
|
||||
- Calling API operations
|
||||
- Making raw HTTP requests
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Annotated, Any, Optional, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, auto_timeout, create_tool_context
|
||||
|
||||
from .client import APIClient, get_api_client
|
||||
from .models import APICallResult, OperationListResult, ProviderListResult
|
||||
from .errors import APIError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _format_provider_list(result: ProviderListResult) -> str:
|
||||
"""Format provider list for display."""
|
||||
lines = ["=== API Providers ===", ""]
|
||||
|
||||
configured = [p for p in result.providers if p.configured]
|
||||
unconfigured = [p for p in result.providers if not p.configured]
|
||||
|
||||
if configured:
|
||||
lines.append("Configured:")
|
||||
for p in sorted(configured, key=lambda x: x.name):
|
||||
source = f"[{p.source.value}]" if p.source else ""
|
||||
lines.append(f" {p.name}: {p.display_name} {source}")
|
||||
|
||||
if unconfigured:
|
||||
lines.append("\nAvailable (not configured):")
|
||||
for p in sorted(unconfigured, key=lambda x: x.name)[:20]: # Limit display
|
||||
lines.append(f" {p.name}: {p.display_name}")
|
||||
if len(unconfigured) > 20:
|
||||
lines.append(f" ... and {len(unconfigured) - 20} more")
|
||||
|
||||
lines.append("")
|
||||
lines.append(f"Total: {result.total_count} providers, {result.configured_count} configured")
|
||||
lines.append("")
|
||||
lines.append("Use 'api --provider <name>' for details")
|
||||
lines.append("Use 'api --action config --provider <name> --api_key <key>' to configure")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _format_operation_list(result: OperationListResult, provider: str) -> str:
|
||||
"""Format operation list for display."""
|
||||
if not result.operations:
|
||||
return f"No operations found for {provider}. Try loading the spec first."
|
||||
|
||||
lines = [f"=== {provider} Operations ==="]
|
||||
|
||||
# Group by tag
|
||||
tags: dict[str, list] = {}
|
||||
for op in result.operations:
|
||||
for t in op.tags or ["untagged"]:
|
||||
if t not in tags:
|
||||
tags[t] = []
|
||||
tags[t].append(op)
|
||||
|
||||
for tag_name, tag_ops in sorted(tags.items()):
|
||||
lines.append(f"\n[{tag_name}]")
|
||||
for op in tag_ops[:10]: # Limit per tag
|
||||
lines.append(f" {op.operation_id}: {op.method} {op.path}")
|
||||
if op.summary:
|
||||
lines.append(f" {op.summary[:80]}")
|
||||
if len(tag_ops) > 10:
|
||||
lines.append(f" ... and {len(tag_ops) - 10} more")
|
||||
|
||||
lines.append(f"\nTotal: {result.total_count} operations")
|
||||
if tags:
|
||||
lines.append(f"Tags: {', '.join(sorted(tags.keys()))}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _format_api_result(result: APICallResult) -> str:
|
||||
"""Format API call result for display."""
|
||||
lines = [f"Status: {result.status_code} {'✓' if result.success else '✗'}"]
|
||||
|
||||
if result.body is not None:
|
||||
if isinstance(result.body, dict):
|
||||
formatted = json.dumps(result.body, indent=2)
|
||||
else:
|
||||
formatted = str(result.body)
|
||||
|
||||
# Truncate if too long
|
||||
if len(formatted) > 5000:
|
||||
formatted = formatted[:5000] + "\n... (truncated)"
|
||||
|
||||
lines.append("")
|
||||
lines.append(formatted)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@final
|
||||
class APITool(BaseTool):
|
||||
"""Unified tool for calling any REST API.
|
||||
|
||||
This is a thin wrapper over APIClient. See client.py for the real logic.
|
||||
|
||||
Actions:
|
||||
- list: List available/configured providers
|
||||
- config: Configure credentials for a provider
|
||||
- delete: Remove stored credentials
|
||||
- spec: Load/refresh OpenAPI spec
|
||||
- ops: List available operations for a provider
|
||||
- call: Call an API operation
|
||||
- raw: Make a raw HTTP request
|
||||
"""
|
||||
|
||||
name = "api"
|
||||
|
||||
def __init__(self, client: Optional[APIClient] = None):
|
||||
"""Initialize with optional client."""
|
||||
self._client = client
|
||||
|
||||
@property
|
||||
def client(self) -> APIClient:
|
||||
"""Get the API client (lazy initialization)."""
|
||||
if self._client is None:
|
||||
self._client = get_api_client()
|
||||
return self._client
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
return """Generic API tool for calling any REST API via OpenAPI specs.
|
||||
|
||||
Supports searching, exploring, and dynamically using ANY public API.
|
||||
|
||||
Actions:
|
||||
- list: Show all providers and their configuration status
|
||||
- config: Set credentials for a provider
|
||||
- delete: Remove credentials for a provider
|
||||
- spec: Load/refresh OpenAPI spec for a provider
|
||||
- ops: List available operations for a provider
|
||||
- call: Call an API operation by ID
|
||||
- raw: Make a raw HTTP request to any endpoint
|
||||
- search: Search for APIs in public registries (openapisearch.com)
|
||||
- register: Register a custom API from any OpenAPI spec URL
|
||||
- overview: Get agent-friendly overview of an API
|
||||
- preload: Download and cache common OpenAPI specs
|
||||
|
||||
DISCOVER ANY API:
|
||||
api --action search --search notion
|
||||
api --action register --provider notion --spec_url <url>
|
||||
api --action overview --provider notion
|
||||
api --action call --provider notion --operation <op>
|
||||
|
||||
BUILT-IN PROVIDERS:
|
||||
Auto-detects credentials from environment variables:
|
||||
- Cloudflare: CLOUDFLARE_API_TOKEN, CF_API_TOKEN, CLOUDFLARE_API_KEY
|
||||
- GitHub: GITHUB_TOKEN, GH_TOKEN
|
||||
- Stripe: STRIPE_API_KEY, STRIPE_SECRET_KEY
|
||||
- OpenAI: OPENAI_API_KEY
|
||||
- Anthropic: ANTHROPIC_API_KEY
|
||||
- etc. (26+ providers)
|
||||
|
||||
Configure manually:
|
||||
api --action config --provider cloudflare --api_key "your-key"
|
||||
|
||||
MAKING CALLS:
|
||||
api --action call --provider cloudflare --operation zones-get
|
||||
api --action raw --provider github --method GET --path /user
|
||||
|
||||
Examples:
|
||||
api # List all providers
|
||||
api --action search --search weather # Search for weather APIs
|
||||
api --action register --provider petstore --spec_url https://petstore.swagger.io/v2/swagger.json
|
||||
api --action overview --provider cloudflare # Agent-friendly summary
|
||||
api --action ops --provider cloudflare --search zones
|
||||
api --action call --provider cloudflare --operation zones-get
|
||||
"""
|
||||
|
||||
@override
|
||||
@auto_timeout("api")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
action: str = "list",
|
||||
provider: Optional[str] = None,
|
||||
# Config params
|
||||
api_key: Optional[str] = None,
|
||||
api_secret: Optional[str] = None,
|
||||
account_id: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
# Spec params
|
||||
spec_url: Optional[str] = None,
|
||||
force_refresh: bool = False,
|
||||
# Operation params
|
||||
operation: Optional[str] = None,
|
||||
params: Optional[str] = None, # JSON string
|
||||
body: Optional[str] = None, # JSON string
|
||||
# Raw call params
|
||||
method: str = "GET",
|
||||
path: Optional[str] = None,
|
||||
# List/search params
|
||||
search: Optional[str] = None,
|
||||
tag: Optional[str] = None,
|
||||
configured_only: bool = False,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""Execute API action."""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
try:
|
||||
if action == "list":
|
||||
return await self._handle_list(provider, configured_only)
|
||||
|
||||
elif action == "config":
|
||||
return await self._handle_config(
|
||||
provider, api_key, api_secret, account_id, base_url
|
||||
)
|
||||
|
||||
elif action == "delete":
|
||||
return await self._handle_delete(provider)
|
||||
|
||||
elif action == "spec":
|
||||
return await self._handle_spec(provider, spec_url, force_refresh)
|
||||
|
||||
elif action == "ops":
|
||||
return await self._handle_ops(provider, search, tag)
|
||||
|
||||
elif action == "call":
|
||||
return await self._handle_call(provider, operation, params, body)
|
||||
|
||||
elif action == "raw":
|
||||
return await self._handle_raw(provider, method, path, params, body)
|
||||
|
||||
elif action == "search":
|
||||
return await self._handle_search(search)
|
||||
|
||||
elif action == "register":
|
||||
return await self._handle_register(
|
||||
provider, spec_url, base_url, api_key, kwargs.get("auth_type", "bearer")
|
||||
)
|
||||
|
||||
elif action == "overview":
|
||||
return await self._handle_overview(provider)
|
||||
|
||||
elif action == "preload":
|
||||
return await self._handle_preload()
|
||||
|
||||
else:
|
||||
return f"Unknown action: {action}. Valid: list, config, delete, spec, ops, call, raw, search, register, overview, preload"
|
||||
|
||||
except APIError as e:
|
||||
# Structured error with hints
|
||||
return str(e)
|
||||
except Exception as e:
|
||||
logger.exception(f"API tool error: {e}")
|
||||
return f"Error: {e}"
|
||||
|
||||
async def _handle_list(self, provider: Optional[str], configured_only: bool) -> str:
|
||||
"""List providers and their status."""
|
||||
if provider:
|
||||
# Show details for specific provider
|
||||
try:
|
||||
status = await self.client.get_provider(provider)
|
||||
lines = [f"=== {provider} ==="]
|
||||
lines.append(f"Display name: {status.display_name}")
|
||||
lines.append(f"Configured: {status.configured}")
|
||||
if status.source:
|
||||
lines.append(f"Source: {status.source.value}")
|
||||
if status.base_url:
|
||||
lines.append(f"Base URL: {status.base_url}")
|
||||
if status.auth_type:
|
||||
lines.append(f"Auth type: {status.auth_type}")
|
||||
if status.spec_url:
|
||||
lines.append(f"Spec URL: {status.spec_url}")
|
||||
if status.env_vars:
|
||||
lines.append(f"Env vars: {', '.join(status.env_vars)}")
|
||||
return "\n".join(lines)
|
||||
except Exception as e:
|
||||
return f"Error getting provider info: {e}"
|
||||
|
||||
# List all providers
|
||||
result = await self.client.list_providers(configured_only=configured_only)
|
||||
return _format_provider_list(result)
|
||||
|
||||
async def _handle_config(
|
||||
self,
|
||||
provider: Optional[str],
|
||||
api_key: Optional[str],
|
||||
api_secret: Optional[str],
|
||||
account_id: Optional[str],
|
||||
base_url: Optional[str],
|
||||
) -> str:
|
||||
"""Configure credentials for a provider."""
|
||||
if not provider:
|
||||
return "Error: --provider required for config action"
|
||||
|
||||
if not api_key:
|
||||
return "Error: --api_key required for config action"
|
||||
|
||||
await self.client.config(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
api_secret=api_secret,
|
||||
account_id=account_id,
|
||||
base_url=base_url,
|
||||
)
|
||||
|
||||
return f"Configured credentials for {provider}"
|
||||
|
||||
async def _handle_delete(self, provider: Optional[str]) -> str:
|
||||
"""Delete credentials for a provider."""
|
||||
if not provider:
|
||||
return "Error: --provider required for delete action"
|
||||
|
||||
if await self.client.delete_config(provider):
|
||||
return f"Deleted credentials for {provider}"
|
||||
return f"No stored credentials for {provider}"
|
||||
|
||||
async def _handle_spec(
|
||||
self,
|
||||
provider: Optional[str],
|
||||
spec_url: Optional[str],
|
||||
force_refresh: bool,
|
||||
) -> str:
|
||||
"""Load or refresh OpenAPI spec."""
|
||||
if not provider:
|
||||
return "Error: --provider required for spec action"
|
||||
|
||||
try:
|
||||
count = await self.client.spec(
|
||||
provider=provider,
|
||||
spec_url=spec_url,
|
||||
force_refresh=force_refresh,
|
||||
)
|
||||
return f"Loaded spec for {provider}: {count} operations"
|
||||
except Exception as e:
|
||||
return f"Error loading spec: {e}"
|
||||
|
||||
async def _handle_ops(
|
||||
self,
|
||||
provider: Optional[str],
|
||||
search: Optional[str],
|
||||
tag: Optional[str],
|
||||
) -> str:
|
||||
"""List operations for a provider."""
|
||||
if not provider:
|
||||
return "Error: --provider required for ops action"
|
||||
|
||||
try:
|
||||
result = await self.client.ops(
|
||||
provider=provider,
|
||||
search=search,
|
||||
tag=tag,
|
||||
)
|
||||
return _format_operation_list(result, provider)
|
||||
except Exception as e:
|
||||
return f"Error listing operations: {e}"
|
||||
|
||||
async def _handle_call(
|
||||
self,
|
||||
provider: Optional[str],
|
||||
operation: Optional[str],
|
||||
params: Optional[str],
|
||||
body: Optional[str],
|
||||
) -> str:
|
||||
"""Call an API operation."""
|
||||
if not provider:
|
||||
return "Error: --provider required for call action"
|
||||
if not operation:
|
||||
return "Error: --operation required for call action"
|
||||
|
||||
try:
|
||||
# Parse params and body
|
||||
parsed_params = json.loads(params) if params else None
|
||||
parsed_body = json.loads(body) if body else None
|
||||
|
||||
result = await self.client.call(
|
||||
provider=provider,
|
||||
operation_id=operation,
|
||||
params=parsed_params,
|
||||
body=parsed_body,
|
||||
)
|
||||
|
||||
return _format_api_result(result)
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
return f"Error parsing JSON: {e}"
|
||||
except Exception as e:
|
||||
return f"Error calling operation: {e}"
|
||||
|
||||
async def _handle_raw(
|
||||
self,
|
||||
provider: Optional[str],
|
||||
method: str,
|
||||
path: Optional[str],
|
||||
params: Optional[str],
|
||||
body: Optional[str],
|
||||
) -> str:
|
||||
"""Make a raw API request."""
|
||||
if not provider:
|
||||
return "Error: --provider required for raw action"
|
||||
if not path:
|
||||
return "Error: --path required for raw action"
|
||||
|
||||
try:
|
||||
# Parse params and body
|
||||
parsed_params = json.loads(params) if params else None
|
||||
parsed_body = json.loads(body) if body else None
|
||||
|
||||
result = await self.client.raw(
|
||||
provider=provider,
|
||||
method=method,
|
||||
path=path,
|
||||
params=parsed_params,
|
||||
body=parsed_body,
|
||||
)
|
||||
|
||||
return _format_api_result(result)
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
return f"Error parsing JSON: {e}"
|
||||
except Exception as e:
|
||||
return f"Error making request: {e}"
|
||||
|
||||
async def _handle_search(self, query: Optional[str]) -> str:
|
||||
"""Search for APIs in public registries."""
|
||||
if not query:
|
||||
return "Error: --search <query> required for search action"
|
||||
|
||||
try:
|
||||
results = await self.client.search(query)
|
||||
|
||||
if not results:
|
||||
return f"No APIs found matching '{query}'"
|
||||
|
||||
lines = [f"=== API Search: {query} ===", ""]
|
||||
for r in results[:20]: # Limit results
|
||||
name = r.get("name") or r.get("id", "Unknown")
|
||||
desc = r.get("description", "")[:60]
|
||||
spec_url = r.get("spec_url", "")
|
||||
lines.append(f" {r.get('id', name)}: {name}")
|
||||
if desc:
|
||||
lines.append(f" {desc}")
|
||||
if spec_url:
|
||||
lines.append(f" Spec: {spec_url}")
|
||||
lines.append("")
|
||||
|
||||
lines.append(f"Found {len(results)} API(s)")
|
||||
lines.append("")
|
||||
lines.append("To use an API:")
|
||||
lines.append(" api --action register --provider <id> --spec_url <url>")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error searching: {e}"
|
||||
|
||||
async def _handle_register(
|
||||
self,
|
||||
name: Optional[str],
|
||||
spec_url: Optional[str],
|
||||
base_url: Optional[str],
|
||||
api_key: Optional[str],
|
||||
auth_type: str,
|
||||
) -> str:
|
||||
"""Register a custom API provider."""
|
||||
if not name:
|
||||
return "Error: --provider <name> required for register action"
|
||||
|
||||
if not spec_url:
|
||||
return "Error: --spec_url required for register action"
|
||||
|
||||
try:
|
||||
count = await self.client.register(
|
||||
name=name,
|
||||
spec_url=spec_url,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
auth_type=auth_type,
|
||||
)
|
||||
|
||||
lines = [
|
||||
f"Registered API: {name}",
|
||||
f"Loaded {count} operations from spec",
|
||||
"",
|
||||
"Next steps:",
|
||||
f" api --action overview --provider {name}",
|
||||
f" api --action ops --provider {name}",
|
||||
]
|
||||
|
||||
if not api_key:
|
||||
lines.append(f" api --action config --provider {name} --api_key <your-key>")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error registering API: {e}"
|
||||
|
||||
async def _handle_overview(self, provider: Optional[str]) -> str:
|
||||
"""Get API overview."""
|
||||
if not provider:
|
||||
return "Error: --provider required for overview action"
|
||||
|
||||
try:
|
||||
return await self.client.overview(provider)
|
||||
except Exception as e:
|
||||
return f"Error getting overview: {e}"
|
||||
|
||||
async def _handle_preload(self) -> str:
|
||||
"""Preload and cache common OpenAPI specs."""
|
||||
try:
|
||||
results = await self.client.preload_specs()
|
||||
|
||||
lines = ["=== Preloading OpenAPI Specs ===", ""]
|
||||
|
||||
success = []
|
||||
failed = []
|
||||
for name, count in sorted(results.items()):
|
||||
if count >= 0:
|
||||
success.append(f" {name}: {count} operations")
|
||||
else:
|
||||
failed.append(f" {name}: failed")
|
||||
|
||||
if success:
|
||||
lines.append("Loaded:")
|
||||
lines.extend(success)
|
||||
|
||||
if failed:
|
||||
lines.append("\nFailed:")
|
||||
lines.extend(failed)
|
||||
|
||||
lines.append(f"\nTotal: {len(success)} loaded, {len(failed)} failed")
|
||||
lines.append(f"Cache: ~/.hanzo/api/specs/")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error preloading specs: {e}"
|
||||
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register with MCP server."""
|
||||
tool_instance = self
|
||||
|
||||
@mcp_server.tool()
|
||||
async def api(
|
||||
action: Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Action: list, config, delete, spec, ops, call, raw",
|
||||
default="list",
|
||||
),
|
||||
] = "list",
|
||||
provider: Annotated[
|
||||
Optional[str],
|
||||
Field(description="Provider name (e.g., cloudflare, github)"),
|
||||
] = None,
|
||||
api_key: Annotated[
|
||||
Optional[str],
|
||||
Field(description="API key or token for config action"),
|
||||
] = None,
|
||||
api_secret: Annotated[
|
||||
Optional[str],
|
||||
Field(description="API secret (if needed) for config action"),
|
||||
] = None,
|
||||
account_id: Annotated[
|
||||
Optional[str],
|
||||
Field(description="Account/org ID for config action"),
|
||||
] = None,
|
||||
base_url: Annotated[
|
||||
Optional[str],
|
||||
Field(description="Override base URL for provider"),
|
||||
] = None,
|
||||
spec_url: Annotated[
|
||||
Optional[str],
|
||||
Field(description="URL to OpenAPI spec for spec action"),
|
||||
] = None,
|
||||
force_refresh: Annotated[
|
||||
bool,
|
||||
Field(description="Force refresh cached spec"),
|
||||
] = False,
|
||||
operation: Annotated[
|
||||
Optional[str],
|
||||
Field(description="Operation ID for call action"),
|
||||
] = None,
|
||||
params: Annotated[
|
||||
Optional[str],
|
||||
Field(description="JSON parameters for call/raw action"),
|
||||
] = None,
|
||||
body: Annotated[
|
||||
Optional[str],
|
||||
Field(description="JSON body for call/raw action"),
|
||||
] = None,
|
||||
method: Annotated[
|
||||
str,
|
||||
Field(description="HTTP method for raw action"),
|
||||
] = "GET",
|
||||
path: Annotated[
|
||||
Optional[str],
|
||||
Field(description="URL path for raw action"),
|
||||
] = None,
|
||||
search: Annotated[
|
||||
Optional[str],
|
||||
Field(description="Search filter for ops action"),
|
||||
] = None,
|
||||
tag: Annotated[
|
||||
Optional[str],
|
||||
Field(description="Tag filter for ops action"),
|
||||
] = None,
|
||||
configured_only: Annotated[
|
||||
bool,
|
||||
Field(description="Only list configured providers"),
|
||||
] = False,
|
||||
ctx: MCPContext = None,
|
||||
) -> str:
|
||||
"""Generic API tool for calling any REST API via OpenAPI specs.
|
||||
|
||||
Manage credentials and call APIs for various cloud providers.
|
||||
Auto-detects credentials from environment variables.
|
||||
|
||||
Actions:
|
||||
- list: Show all providers and their status
|
||||
- config: Set credentials for a provider
|
||||
- delete: Remove credentials for a provider
|
||||
- spec: Load/refresh OpenAPI spec for a provider
|
||||
- ops: List available operations for a provider
|
||||
- call: Call an API operation by operation ID
|
||||
- raw: Make a raw HTTP request
|
||||
"""
|
||||
return await tool_instance.call(
|
||||
ctx,
|
||||
action=action,
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
api_secret=api_secret,
|
||||
account_id=account_id,
|
||||
base_url=base_url,
|
||||
spec_url=spec_url,
|
||||
force_refresh=force_refresh,
|
||||
operation=operation,
|
||||
params=params,
|
||||
body=body,
|
||||
method=method,
|
||||
path=path,
|
||||
search=search,
|
||||
tag=tag,
|
||||
configured_only=configured_only,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,306 @@
|
||||
"""Credential manager with pluggable storage.
|
||||
|
||||
Provides unified credential management with:
|
||||
- Explicit resolution order (override > memory > file > env)
|
||||
- Pluggable storage backends
|
||||
- Clear source introspection
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from .models import Credential, CredentialSource, EffectiveCredential, ProviderConfig, ProviderStatus
|
||||
from .providers import PROVIDER_CONFIGS, ENV_VAR_MAPPINGS, get_env_vars, get_provider_config
|
||||
from .storage import (
|
||||
CredentialStorage,
|
||||
ChainedCredentialStorage,
|
||||
FileCredentialStorage,
|
||||
EnvironmentCredentialStorage,
|
||||
MemoryCredentialStorage,
|
||||
)
|
||||
from .errors import CredentialError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CredentialManager:
|
||||
"""Manages API credentials with explicit resolution order.
|
||||
|
||||
Resolution order:
|
||||
1. Per-call overrides (via override_for_provider)
|
||||
2. In-memory config (set at runtime)
|
||||
3. Stored file (~/.hanzo/api/credentials.json)
|
||||
4. Environment variables
|
||||
|
||||
Example:
|
||||
manager = CredentialManager()
|
||||
|
||||
# Check what credentials are available
|
||||
effective = await manager.get_effective_credentials("cloudflare")
|
||||
print(f"Source: {effective.source}") # e.g., "environment"
|
||||
|
||||
# Configure credentials
|
||||
await manager.set_credential("cloudflare", api_key="my-key")
|
||||
|
||||
# Make a call with per-call override
|
||||
async with manager.override_for_provider("cloudflare", api_key="temp-key"):
|
||||
# Uses temp-key for this call only
|
||||
cred = await manager.get_credential("cloudflare")
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config_dir: Optional[Path] = None,
|
||||
storage: Optional[CredentialStorage] = None,
|
||||
):
|
||||
"""Initialize credential manager.
|
||||
|
||||
Args:
|
||||
config_dir: Directory for file storage. Defaults to ~/.hanzo/api/
|
||||
storage: Custom storage backend. If None, uses default chain.
|
||||
"""
|
||||
self.config_dir = config_dir or Path.home() / ".hanzo" / "api"
|
||||
self.specs_dir = self.config_dir / "specs"
|
||||
|
||||
# Set up storage chain
|
||||
# Note: MemoryCredentialStorage is NOT included here because
|
||||
# per-call overrides are handled separately by self._overrides
|
||||
if storage:
|
||||
self._storage = storage
|
||||
else:
|
||||
file_path = self.config_dir / "credentials.json"
|
||||
self._storage = ChainedCredentialStorage([
|
||||
FileCredentialStorage(file_path), # Persistent storage
|
||||
EnvironmentCredentialStorage(ENV_VAR_MAPPINGS), # Read-only fallback
|
||||
])
|
||||
|
||||
# Separate override storage for context manager
|
||||
self._overrides = MemoryCredentialStorage()
|
||||
|
||||
async def get_credential(self, provider: str) -> Credential:
|
||||
"""Get credential for a provider.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
Credential (may be empty if not found)
|
||||
"""
|
||||
effective = await self.get_effective_credentials(provider)
|
||||
return effective.credential
|
||||
|
||||
async def get_effective_credentials(self, provider: str) -> EffectiveCredential:
|
||||
"""Get credential with full resolution info.
|
||||
|
||||
Shows exactly where the credential came from and what
|
||||
alternatives are available.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
EffectiveCredential with source information
|
||||
"""
|
||||
# Check overrides first
|
||||
override = await self._overrides.get(provider)
|
||||
if override and override.has_credentials:
|
||||
return EffectiveCredential(
|
||||
credential=override,
|
||||
source=CredentialSource.OVERRIDE,
|
||||
)
|
||||
|
||||
# Check chained storage
|
||||
if isinstance(self._storage, ChainedCredentialStorage):
|
||||
cred, store = await self._storage.get_with_source(provider)
|
||||
if cred and cred.has_credentials:
|
||||
source = CredentialSource.NONE
|
||||
env_var = None
|
||||
|
||||
if isinstance(store, MemoryCredentialStorage):
|
||||
source = CredentialSource.MEMORY
|
||||
elif isinstance(store, FileCredentialStorage):
|
||||
source = CredentialSource.STORED
|
||||
elif isinstance(store, EnvironmentCredentialStorage):
|
||||
source = CredentialSource.ENVIRONMENT
|
||||
# Find which env var was used
|
||||
for var in get_env_vars(provider):
|
||||
if os.environ.get(var):
|
||||
env_var = var
|
||||
break
|
||||
|
||||
return EffectiveCredential(
|
||||
credential=cred,
|
||||
source=source,
|
||||
env_var_used=env_var,
|
||||
)
|
||||
else:
|
||||
cred = await self._storage.get(provider)
|
||||
if cred and cred.has_credentials:
|
||||
return EffectiveCredential(
|
||||
credential=cred,
|
||||
source=CredentialSource.STORED,
|
||||
)
|
||||
|
||||
# Check provider config for base URL
|
||||
config = get_provider_config(provider)
|
||||
base_url = config.base_url if config else None
|
||||
|
||||
return EffectiveCredential(
|
||||
credential=Credential(provider=provider, base_url=base_url),
|
||||
source=CredentialSource.NONE,
|
||||
)
|
||||
|
||||
async def set_credential(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: Optional[str] = None,
|
||||
api_secret: Optional[str] = None,
|
||||
account_id: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
**extra,
|
||||
) -> None:
|
||||
"""Store a credential.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
api_key: API key or token
|
||||
api_secret: API secret (for providers that need it)
|
||||
account_id: Account/organization ID
|
||||
base_url: Custom base URL override
|
||||
**extra: Additional provider-specific fields
|
||||
"""
|
||||
# Get existing to merge
|
||||
existing = await self._storage.get(provider)
|
||||
|
||||
credential = Credential(
|
||||
provider=provider,
|
||||
api_key=api_key or (existing.api_key if existing else None),
|
||||
api_secret=api_secret or (existing.api_secret if existing else None),
|
||||
account_id=account_id or (existing.account_id if existing else None),
|
||||
base_url=base_url or (existing.base_url if existing else None),
|
||||
extra={**(existing.extra if existing else {}), **extra},
|
||||
)
|
||||
|
||||
await self._storage.set(credential)
|
||||
|
||||
async def delete_credential(self, provider: str) -> bool:
|
||||
"""Delete a stored credential.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
True if deleted, False if not found
|
||||
"""
|
||||
return await self._storage.delete(provider)
|
||||
|
||||
async def list_credentials(self) -> list[str]:
|
||||
"""List all providers with stored credentials."""
|
||||
return await self._storage.list()
|
||||
|
||||
async def list_providers(self) -> list[ProviderStatus]:
|
||||
"""List all providers and their status.
|
||||
|
||||
Returns:
|
||||
List of provider status objects
|
||||
"""
|
||||
result = []
|
||||
all_providers = set(PROVIDER_CONFIGS.keys()) | set(ENV_VAR_MAPPINGS.keys())
|
||||
|
||||
# Add stored providers too
|
||||
stored = await self._storage.list()
|
||||
all_providers.update(stored)
|
||||
|
||||
for provider in sorted(all_providers):
|
||||
effective = await self.get_effective_credentials(provider)
|
||||
config = get_provider_config(provider)
|
||||
|
||||
# Check spec cache
|
||||
spec_cached = False
|
||||
spec_age = None
|
||||
spec_file = self.specs_dir / f"{provider}.json"
|
||||
if spec_file.exists():
|
||||
spec_cached = True
|
||||
import time
|
||||
|
||||
spec_age = time.time() - spec_file.stat().st_mtime
|
||||
|
||||
result.append(
|
||||
ProviderStatus(
|
||||
name=provider,
|
||||
display_name=config.display_name if config else provider.title(),
|
||||
configured=effective.has_credentials,
|
||||
source=effective.source if effective.has_credentials else None,
|
||||
base_url=config.base_url if config else "",
|
||||
has_spec=bool(config and config.spec_url),
|
||||
spec_cached=spec_cached,
|
||||
spec_age_seconds=spec_age,
|
||||
)
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
async def require_credential(self, provider: str) -> Credential:
|
||||
"""Get credential or raise CredentialError.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
Credential with valid api_key
|
||||
|
||||
Raises:
|
||||
CredentialError: If no credential is configured
|
||||
"""
|
||||
effective = await self.get_effective_credentials(provider)
|
||||
|
||||
if not effective.has_credentials:
|
||||
env_vars = get_env_vars(provider)
|
||||
raise CredentialError(
|
||||
message=f"No credentials configured for {provider}",
|
||||
provider=provider,
|
||||
env_vars=env_vars,
|
||||
)
|
||||
|
||||
return effective.credential
|
||||
|
||||
def get_provider_config(self, provider: str) -> Optional[ProviderConfig]:
|
||||
"""Get provider configuration."""
|
||||
return get_provider_config(provider)
|
||||
|
||||
|
||||
# Singleton instance
|
||||
_credential_manager: Optional[CredentialManager] = None
|
||||
|
||||
|
||||
def get_credential_manager() -> CredentialManager:
|
||||
"""Get the global credential manager instance."""
|
||||
global _credential_manager
|
||||
if _credential_manager is None:
|
||||
_credential_manager = CredentialManager()
|
||||
return _credential_manager
|
||||
|
||||
|
||||
def reset_credential_manager() -> None:
|
||||
"""Reset the global credential manager (for testing)."""
|
||||
global _credential_manager
|
||||
_credential_manager = None
|
||||
|
||||
|
||||
# Re-export for backwards compatibility
|
||||
__all__ = [
|
||||
"CredentialManager",
|
||||
"get_credential_manager",
|
||||
"reset_credential_manager",
|
||||
"Credential",
|
||||
"EffectiveCredential",
|
||||
"CredentialSource",
|
||||
"ProviderConfig",
|
||||
"ProviderStatus",
|
||||
"PROVIDER_CONFIGS",
|
||||
"ENV_VAR_MAPPINGS",
|
||||
]
|
||||
@@ -0,0 +1,355 @@
|
||||
"""Structured errors for hanzo-tools-api.
|
||||
|
||||
Provides detailed error objects with hints for resolution,
|
||||
making it easy for agents to understand and recover from errors.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class APIErrorInfo(BaseModel):
|
||||
"""Structured error information."""
|
||||
|
||||
error_type: str = Field(description="Error category")
|
||||
message: str = Field(description="Human-readable error message")
|
||||
provider: Optional[str] = Field(default=None, description="Provider involved")
|
||||
operation_id: Optional[str] = Field(default=None, description="Operation involved")
|
||||
status_code: Optional[int] = Field(default=None, description="HTTP status code")
|
||||
missing_fields: list[str] = Field(default_factory=list, description="Required fields that are missing")
|
||||
hint: Optional[str] = Field(default=None, description="Suggested fix")
|
||||
env_vars: list[str] = Field(default_factory=list, description="Environment variables to set")
|
||||
config_command: Optional[str] = Field(default=None, description="CLI command to fix")
|
||||
details: Optional[dict[str, Any]] = Field(default=None, description="Additional error details")
|
||||
|
||||
|
||||
class APIError(Exception):
|
||||
"""Base exception for API errors with structured info."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
error_type: str = "api_error",
|
||||
provider: Optional[str] = None,
|
||||
operation_id: Optional[str] = None,
|
||||
status_code: Optional[int] = None,
|
||||
body: Any = None,
|
||||
hint: Optional[str] = None,
|
||||
missing_fields: Optional[list[str]] = None,
|
||||
env_vars: Optional[list[str]] = None,
|
||||
config_command: Optional[str] = None,
|
||||
details: Optional[dict[str, Any]] = None,
|
||||
):
|
||||
self.message = message
|
||||
self.error_type = error_type
|
||||
self.provider = provider
|
||||
self.operation_id = operation_id
|
||||
self.status_code = status_code
|
||||
self.body = body
|
||||
self.hint = hint
|
||||
self.missing_fields = missing_fields or []
|
||||
self.env_vars = env_vars or []
|
||||
self.config_command = config_command
|
||||
self.details = details
|
||||
|
||||
super().__init__(message)
|
||||
|
||||
@property
|
||||
def info(self) -> APIErrorInfo:
|
||||
"""Get structured error info."""
|
||||
return APIErrorInfo(
|
||||
error_type=self.error_type,
|
||||
message=self.message,
|
||||
provider=self.provider,
|
||||
operation_id=self.operation_id,
|
||||
status_code=self.status_code,
|
||||
missing_fields=self.missing_fields,
|
||||
hint=self.hint,
|
||||
env_vars=self.env_vars,
|
||||
config_command=self.config_command,
|
||||
details=self.details,
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
parts = [self.message]
|
||||
if self.hint:
|
||||
parts.append(f"Hint: {self.hint}")
|
||||
if self.env_vars:
|
||||
parts.append(f"Try setting: {', '.join(self.env_vars)}")
|
||||
if self.config_command:
|
||||
parts.append(f"Or run: {self.config_command}")
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
class CredentialError(APIError):
|
||||
"""Error related to credentials."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
provider: str,
|
||||
missing_fields: Optional[list[str]] = None,
|
||||
env_vars: Optional[list[str]] = None,
|
||||
):
|
||||
config_cmd = f'api config --provider {provider} --api_key "YOUR_KEY"'
|
||||
hint = f"Configure credentials for {provider}"
|
||||
|
||||
super().__init__(
|
||||
message=message,
|
||||
error_type="credential_error",
|
||||
provider=provider,
|
||||
missing_fields=missing_fields or ["api_key"],
|
||||
hint=hint,
|
||||
env_vars=env_vars or [],
|
||||
config_command=config_cmd,
|
||||
)
|
||||
|
||||
|
||||
class ProviderNotFoundError(APIError):
|
||||
"""Provider is not recognized."""
|
||||
|
||||
def __init__(self, provider: str, available_providers: Optional[list[str]] = None):
|
||||
hint = None
|
||||
if available_providers:
|
||||
hint = f"Available providers: {', '.join(available_providers[:10])}"
|
||||
|
||||
super().__init__(
|
||||
message=f"Unknown provider: {provider}",
|
||||
error_type="provider_not_found",
|
||||
provider=provider,
|
||||
hint=hint,
|
||||
details={"available_providers": available_providers} if available_providers else None,
|
||||
)
|
||||
|
||||
|
||||
class OperationNotFoundError(APIError):
|
||||
"""Operation is not found in the spec."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operation_id: str,
|
||||
provider: str,
|
||||
similar_operations: Optional[list[str]] = None,
|
||||
):
|
||||
hint = None
|
||||
if similar_operations:
|
||||
hint = f"Similar operations: {', '.join(similar_operations[:5])}"
|
||||
else:
|
||||
hint = f"Run 'api ops --provider {provider}' to list available operations"
|
||||
|
||||
super().__init__(
|
||||
message=f"Operation not found: {operation_id}",
|
||||
error_type="operation_not_found",
|
||||
provider=provider,
|
||||
operation_id=operation_id,
|
||||
hint=hint,
|
||||
details={"similar_operations": similar_operations} if similar_operations else None,
|
||||
)
|
||||
|
||||
|
||||
class SpecNotLoadedError(APIError):
|
||||
"""OpenAPI spec has not been loaded."""
|
||||
|
||||
def __init__(self, provider: str):
|
||||
super().__init__(
|
||||
message=f"OpenAPI spec not loaded for {provider}",
|
||||
error_type="spec_not_loaded",
|
||||
provider=provider,
|
||||
hint=f"Load the spec first with 'api spec --provider {provider}'",
|
||||
config_command=f"api spec --provider {provider}",
|
||||
)
|
||||
|
||||
|
||||
class SpecParseError(APIError):
|
||||
"""Error parsing OpenAPI spec."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
provider: str,
|
||||
path_in_spec: Optional[str] = None,
|
||||
parse_error: Optional[str] = None,
|
||||
):
|
||||
details = {}
|
||||
if path_in_spec:
|
||||
details["path_in_spec"] = path_in_spec
|
||||
if parse_error:
|
||||
details["parse_error"] = parse_error
|
||||
|
||||
super().__init__(
|
||||
message=message,
|
||||
error_type="spec_parse_error",
|
||||
provider=provider,
|
||||
hint="The OpenAPI spec may be invalid or use unsupported features",
|
||||
details=details if details else None,
|
||||
)
|
||||
|
||||
|
||||
class ParameterValidationError(APIError):
|
||||
"""Parameter validation failed."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
provider: str,
|
||||
operation_id: str,
|
||||
missing_params: Optional[list[str]] = None,
|
||||
invalid_params: Optional[dict[str, str]] = None,
|
||||
):
|
||||
details = {}
|
||||
if missing_params:
|
||||
details["missing_params"] = missing_params
|
||||
if invalid_params:
|
||||
details["invalid_params"] = invalid_params
|
||||
|
||||
hint_parts = []
|
||||
if missing_params:
|
||||
hint_parts.append(f"Missing required params: {', '.join(missing_params)}")
|
||||
if invalid_params:
|
||||
for param, reason in invalid_params.items():
|
||||
hint_parts.append(f"{param}: {reason}")
|
||||
|
||||
super().__init__(
|
||||
message=message,
|
||||
error_type="parameter_validation_error",
|
||||
provider=provider,
|
||||
operation_id=operation_id,
|
||||
missing_fields=missing_params or [],
|
||||
hint="; ".join(hint_parts) if hint_parts else None,
|
||||
details=details if details else None,
|
||||
)
|
||||
|
||||
|
||||
class NetworkError(APIError):
|
||||
"""Network-related error."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
provider: str,
|
||||
url: Optional[str] = None,
|
||||
original_error: Optional[Exception] = None,
|
||||
):
|
||||
details = {}
|
||||
if url:
|
||||
details["url"] = url
|
||||
if original_error:
|
||||
details["original_error"] = str(original_error)
|
||||
|
||||
super().__init__(
|
||||
message=message,
|
||||
error_type="network_error",
|
||||
provider=provider,
|
||||
hint="Check your network connection and the provider's API status",
|
||||
details=details if details else None,
|
||||
)
|
||||
|
||||
|
||||
class RateLimitError(APIError):
|
||||
"""Rate limit exceeded."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
retry_after: Optional[int] = None,
|
||||
status_code: int = 429,
|
||||
):
|
||||
hint = "Rate limit exceeded"
|
||||
if retry_after:
|
||||
hint = f"Rate limit exceeded. Retry after {retry_after} seconds"
|
||||
|
||||
super().__init__(
|
||||
message=f"Rate limit exceeded for {provider}",
|
||||
error_type="rate_limit_error",
|
||||
provider=provider,
|
||||
status_code=status_code,
|
||||
hint=hint,
|
||||
details={"retry_after": retry_after} if retry_after else None,
|
||||
)
|
||||
|
||||
|
||||
class AuthenticationError(APIError):
|
||||
"""Authentication failed."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
status_code: int = 401,
|
||||
env_vars: Optional[list[str]] = None,
|
||||
):
|
||||
super().__init__(
|
||||
message=f"Authentication failed for {provider}",
|
||||
error_type="authentication_error",
|
||||
provider=provider,
|
||||
status_code=status_code,
|
||||
hint="Check that your API key is valid and has the required permissions",
|
||||
env_vars=env_vars or [],
|
||||
config_command=f'api config --provider {provider} --api_key "YOUR_KEY"',
|
||||
)
|
||||
|
||||
|
||||
class SpecLoadError(APIError):
|
||||
"""Error loading OpenAPI spec."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
spec_url: Optional[str] = None,
|
||||
original_error: Optional[Exception] = None,
|
||||
):
|
||||
message = f"Failed to load OpenAPI spec for {provider}"
|
||||
if spec_url:
|
||||
message = f"Failed to load OpenAPI spec from {spec_url}"
|
||||
|
||||
details = {}
|
||||
if spec_url:
|
||||
details["spec_url"] = spec_url
|
||||
if original_error:
|
||||
details["original_error"] = str(original_error)
|
||||
|
||||
super().__init__(
|
||||
message=message,
|
||||
error_type="spec_load_error",
|
||||
provider=provider,
|
||||
hint="Check the spec URL is accessible and returns valid OpenAPI JSON/YAML",
|
||||
details=details if details else None,
|
||||
)
|
||||
|
||||
|
||||
class ValidationError(APIError):
|
||||
"""General validation error."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
provider: Optional[str] = None,
|
||||
operation_id: Optional[str] = None,
|
||||
field: Optional[str] = None,
|
||||
value: Any = None,
|
||||
expected: Optional[str] = None,
|
||||
):
|
||||
details = {}
|
||||
if field:
|
||||
details["field"] = field
|
||||
if value is not None:
|
||||
details["value"] = str(value)
|
||||
if expected:
|
||||
details["expected"] = expected
|
||||
|
||||
hint = None
|
||||
if field and expected:
|
||||
hint = f"Field '{field}' should be {expected}"
|
||||
elif field:
|
||||
hint = f"Invalid value for field '{field}'"
|
||||
|
||||
super().__init__(
|
||||
message=message,
|
||||
error_type="validation_error",
|
||||
provider=provider,
|
||||
operation_id=operation_id,
|
||||
hint=hint,
|
||||
details=details if details else None,
|
||||
)
|
||||
@@ -0,0 +1,388 @@
|
||||
"""Pydantic models for hanzo-tools-api.
|
||||
|
||||
Provides strongly-typed, structured data objects for:
|
||||
- API call results
|
||||
- Operations and parameters
|
||||
- Credentials and provider configs
|
||||
- Tool schemas for agent integration
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Any, Literal, Optional
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Enums
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class AuthType(str, Enum):
|
||||
"""Authentication type for API providers."""
|
||||
|
||||
BEARER = "bearer" # Authorization: Bearer <token>
|
||||
BASIC = "basic" # Basic auth (username:password)
|
||||
HEADER = "header" # Custom header (e.g., x-api-key)
|
||||
API_KEY = "api_key" # X-API-Key header
|
||||
QUERY = "query" # Query parameter
|
||||
|
||||
|
||||
class CredentialSource(str, Enum):
|
||||
"""Source of credential resolution."""
|
||||
|
||||
OVERRIDE = "override" # Per-call override
|
||||
MEMORY = "memory" # In-memory config
|
||||
STORED = "stored" # File storage
|
||||
ENVIRONMENT = "environment" # Environment variable
|
||||
NONE = "none" # Not found
|
||||
|
||||
|
||||
class ParameterLocation(str, Enum):
|
||||
"""Location of API parameter."""
|
||||
|
||||
PATH = "path"
|
||||
QUERY = "query"
|
||||
HEADER = "header"
|
||||
COOKIE = "cookie"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Credential Models
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class Credential(BaseModel):
|
||||
"""API credential for a provider."""
|
||||
|
||||
model_config = ConfigDict(frozen=False)
|
||||
|
||||
provider: str = Field(description="Provider name")
|
||||
api_key: Optional[str] = Field(default=None, description="API key or token")
|
||||
api_secret: Optional[str] = Field(default=None, description="API secret (if needed)")
|
||||
account_id: Optional[str] = Field(default=None, description="Account/organization ID")
|
||||
base_url: Optional[str] = Field(default=None, description="Custom base URL override")
|
||||
extra: dict[str, Any] = Field(default_factory=dict, description="Provider-specific fields")
|
||||
|
||||
@property
|
||||
def has_credentials(self) -> bool:
|
||||
"""Check if credential has at least an API key."""
|
||||
return bool(self.api_key)
|
||||
|
||||
|
||||
class EffectiveCredential(BaseModel):
|
||||
"""Resolved credential with source information."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
credential: Credential
|
||||
source: CredentialSource = Field(description="Where the credential was resolved from")
|
||||
env_var_used: Optional[str] = Field(default=None, description="Environment variable used (if any)")
|
||||
|
||||
@property
|
||||
def has_credentials(self) -> bool:
|
||||
return self.credential.has_credentials
|
||||
|
||||
|
||||
class ProviderConfig(BaseModel):
|
||||
"""Provider-specific configuration."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
name: str = Field(description="Provider identifier")
|
||||
display_name: str = Field(description="Human-readable name")
|
||||
base_url: str = Field(description="Default base URL")
|
||||
auth_type: AuthType = Field(default=AuthType.BEARER, description="Authentication method")
|
||||
auth_header: str = Field(default="Authorization", description="Header name for auth")
|
||||
auth_prefix: str = Field(default="Bearer", description="Prefix for auth value")
|
||||
auth_query_param: str = Field(default="api_key", description="Query param for auth")
|
||||
spec_url: Optional[str] = Field(default=None, description="URL to OpenAPI spec")
|
||||
env_vars: list[str] = Field(default_factory=list, description="Environment variables to check")
|
||||
extra_headers: dict[str, str] = Field(default_factory=dict, description="Additional headers")
|
||||
|
||||
|
||||
class ProviderStatus(BaseModel):
|
||||
"""Status of a provider configuration."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
name: str
|
||||
display_name: str
|
||||
configured: bool = Field(description="Whether credentials are available")
|
||||
source: Optional[CredentialSource] = Field(default=None, description="Credential source")
|
||||
base_url: str
|
||||
has_spec: bool = Field(description="Whether OpenAPI spec is available")
|
||||
spec_cached: bool = Field(default=False, description="Whether spec is cached locally")
|
||||
spec_age_seconds: Optional[float] = Field(default=None, description="Age of cached spec")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# OpenAPI Operation Models
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class Parameter(BaseModel):
|
||||
"""API operation parameter."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
name: str = Field(description="Parameter name")
|
||||
location: ParameterLocation = Field(description="Where the parameter goes")
|
||||
required: bool = Field(default=False, description="Whether parameter is required")
|
||||
schema_type: str = Field(default="string", description="Parameter data type")
|
||||
description: str = Field(default="", description="Parameter description")
|
||||
default: Any = Field(default=None, description="Default value")
|
||||
enum: list[Any] = Field(default_factory=list, description="Allowed values")
|
||||
json_schema: Optional[dict[str, Any]] = Field(default=None, description="Full JSON schema")
|
||||
|
||||
|
||||
class Operation(BaseModel):
|
||||
"""API operation definition."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
operation_id: str = Field(description="Unique operation identifier")
|
||||
method: str = Field(description="HTTP method (GET, POST, etc.)")
|
||||
path: str = Field(description="URL path pattern")
|
||||
summary: str = Field(default="", description="Short description")
|
||||
description: str = Field(default="", description="Detailed description")
|
||||
parameters: list[Parameter] = Field(default_factory=list, description="Operation parameters")
|
||||
request_body_schema: Optional[dict[str, Any]] = Field(default=None, description="Request body JSON schema")
|
||||
request_body_required: bool = Field(default=False, description="Whether request body is required")
|
||||
response_schema: Optional[dict[str, Any]] = Field(default=None, description="Response JSON schema")
|
||||
tags: list[str] = Field(default_factory=list, description="Operation tags")
|
||||
deprecated: bool = Field(default=False, description="Whether operation is deprecated")
|
||||
|
||||
@property
|
||||
def params_schema(self) -> dict[str, Any]:
|
||||
"""Get JSON schema for parameters."""
|
||||
properties = {}
|
||||
required = []
|
||||
|
||||
for param in self.parameters:
|
||||
prop = {"type": param.schema_type}
|
||||
if param.description:
|
||||
prop["description"] = param.description
|
||||
if param.enum:
|
||||
prop["enum"] = param.enum
|
||||
if param.default is not None:
|
||||
prop["default"] = param.default
|
||||
if param.json_schema:
|
||||
prop.update(param.json_schema)
|
||||
|
||||
properties[param.name] = prop
|
||||
if param.required:
|
||||
required.append(param.name)
|
||||
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": properties,
|
||||
"required": required,
|
||||
}
|
||||
|
||||
def compact_repr(self) -> str:
|
||||
"""Get compact representation for agent prompts."""
|
||||
return f"{self.operation_id}: {self.method} {self.path} - {self.summary}"
|
||||
|
||||
|
||||
class OperationSummary(BaseModel):
|
||||
"""Compact operation summary for agent discovery."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
operation_id: str
|
||||
method: str
|
||||
path: str
|
||||
summary: str
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
deprecated: bool = False
|
||||
|
||||
@classmethod
|
||||
def from_operation(cls, op: Operation) -> "OperationSummary":
|
||||
return cls(
|
||||
operation_id=op.operation_id,
|
||||
method=op.method,
|
||||
path=op.path,
|
||||
summary=op.summary,
|
||||
tags=op.tags,
|
||||
deprecated=op.deprecated,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# API Call Result Models
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class APICallResult(BaseModel):
|
||||
"""Structured result from an API call."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
success: bool = Field(description="Whether the call succeeded (2xx status)")
|
||||
status_code: int = Field(description="HTTP status code")
|
||||
headers: dict[str, str] = Field(default_factory=dict, description="Response headers")
|
||||
body: Any = Field(default=None, description="Response body (parsed JSON or raw text)")
|
||||
raw_body: Optional[str] = Field(default=None, description="Raw response body")
|
||||
elapsed_ms: Optional[float] = Field(default=None, description="Request duration in milliseconds")
|
||||
request_id: Optional[str] = Field(default=None, description="Request ID from headers")
|
||||
|
||||
# Metadata
|
||||
provider: str = Field(description="Provider that was called")
|
||||
operation_id: Optional[str] = Field(default=None, description="Operation ID (if using spec)")
|
||||
method: str = Field(description="HTTP method used")
|
||||
url: str = Field(description="Full URL that was called")
|
||||
|
||||
@property
|
||||
def data(self) -> Any:
|
||||
"""Alias for body for backwards compatibility."""
|
||||
return self.body
|
||||
|
||||
@property
|
||||
def is_error(self) -> bool:
|
||||
"""Check if response indicates an error."""
|
||||
return self.status_code >= 400
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
"""Raise APIError if response indicates an error."""
|
||||
if self.is_error:
|
||||
from .errors import APIError
|
||||
|
||||
raise APIError(
|
||||
message=f"API call failed with status {self.status_code}",
|
||||
status_code=self.status_code,
|
||||
body=self.body,
|
||||
provider=self.provider,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Spec Cache Models
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class SpecCacheEntry(BaseModel):
|
||||
"""Cached OpenAPI spec entry."""
|
||||
|
||||
model_config = ConfigDict(frozen=False)
|
||||
|
||||
provider: str
|
||||
spec: dict[str, Any]
|
||||
fetched_at: datetime
|
||||
etag: Optional[str] = None
|
||||
last_modified: Optional[str] = None
|
||||
source_url: Optional[str] = None
|
||||
|
||||
@property
|
||||
def age_seconds(self) -> float:
|
||||
"""Get age of cache entry in seconds."""
|
||||
return (datetime.now() - self.fetched_at).total_seconds()
|
||||
|
||||
@property
|
||||
def is_stale(self) -> bool:
|
||||
"""Check if cache is stale (>24 hours)."""
|
||||
return self.age_seconds > 86400
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Tool Schema Models (for agent integration)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class ToolParameter(BaseModel):
|
||||
"""Parameter definition for tool schema."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
name: str
|
||||
type: str
|
||||
description: str
|
||||
required: bool = False
|
||||
default: Any = None
|
||||
enum: Optional[list[str]] = None
|
||||
|
||||
|
||||
class ToolSchema(BaseModel):
|
||||
"""Machine-readable tool schema for agent integration."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
name: str = Field(description="Tool/method name")
|
||||
description: str = Field(description="What the tool does")
|
||||
parameters: list[ToolParameter] = Field(default_factory=list)
|
||||
returns: str = Field(description="Return type description")
|
||||
examples: list[str] = Field(default_factory=list, description="Example usage")
|
||||
|
||||
def to_json_schema(self) -> dict[str, Any]:
|
||||
"""Convert to JSON schema format."""
|
||||
properties = {}
|
||||
required = []
|
||||
|
||||
for param in self.parameters:
|
||||
prop = {"type": param.type, "description": param.description}
|
||||
if param.enum:
|
||||
prop["enum"] = param.enum
|
||||
if param.default is not None:
|
||||
prop["default"] = param.default
|
||||
properties[param.name] = prop
|
||||
if param.required:
|
||||
required.append(param.name)
|
||||
|
||||
return {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": properties,
|
||||
"required": required,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# List/Discovery Results
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class ProviderListResult(BaseModel):
|
||||
"""Result of listing providers."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
providers: list[ProviderStatus]
|
||||
configured_count: int
|
||||
total_count: int
|
||||
|
||||
@property
|
||||
def configured(self) -> list[ProviderStatus]:
|
||||
"""Get only configured providers."""
|
||||
return [p for p in self.providers if p.configured]
|
||||
|
||||
@property
|
||||
def unconfigured(self) -> list[ProviderStatus]:
|
||||
"""Get only unconfigured providers."""
|
||||
return [p for p in self.providers if not p.configured]
|
||||
|
||||
|
||||
class OperationListResult(BaseModel):
|
||||
"""Result of listing operations."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
provider: str
|
||||
operations: list[OperationSummary]
|
||||
total_count: int
|
||||
filter_applied: Optional[str] = None
|
||||
|
||||
def by_tag(self) -> dict[str, list[OperationSummary]]:
|
||||
"""Group operations by tag."""
|
||||
result: dict[str, list[OperationSummary]] = {}
|
||||
for op in self.operations:
|
||||
for tag in op.tags or ["untagged"]:
|
||||
if tag not in result:
|
||||
result[tag] = []
|
||||
result[tag].append(op)
|
||||
return result
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,416 @@
|
||||
"""Provider configurations and environment variable mappings.
|
||||
|
||||
Contains built-in configurations for 30+ cloud providers,
|
||||
plus 1100+ auto-generated configs from APIs.guru + oapis.org.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .apis_guru_providers import APIS_GURU_PROVIDERS
|
||||
from .models import AuthType, ProviderConfig
|
||||
|
||||
# =============================================================================
|
||||
# Environment Variable Mappings
|
||||
# =============================================================================
|
||||
|
||||
ENV_VAR_MAPPINGS: dict[str, list[str]] = {
|
||||
# Cloud Providers
|
||||
"cloudflare": ["CLOUDFLARE_API_TOKEN", "CF_API_TOKEN", "CLOUDFLARE_API_KEY", "CF_API_KEY"],
|
||||
"aws": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],
|
||||
"gcp": ["GOOGLE_APPLICATION_CREDENTIALS", "GCP_API_KEY", "GOOGLE_API_KEY"],
|
||||
"azure": ["AZURE_API_KEY", "AZURE_SUBSCRIPTION_KEY"],
|
||||
"digitalocean": ["DIGITALOCEAN_TOKEN", "DO_TOKEN", "DIGITALOCEAN_ACCESS_TOKEN"],
|
||||
"linode": ["LINODE_TOKEN", "LINODE_API_TOKEN"],
|
||||
"vultr": ["VULTR_API_KEY"],
|
||||
"hetzner": ["HETZNER_API_TOKEN", "HCLOUD_TOKEN"],
|
||||
# AI Providers
|
||||
"openai": ["OPENAI_API_KEY"],
|
||||
"anthropic": ["ANTHROPIC_API_KEY"],
|
||||
"together": ["TOGETHER_API_KEY", "TOGETHER_AI_KEY"],
|
||||
"replicate": ["REPLICATE_API_TOKEN", "REPLICATE_API_KEY"],
|
||||
"huggingface": ["HF_TOKEN", "HUGGINGFACE_TOKEN", "HUGGING_FACE_HUB_TOKEN"],
|
||||
"cohere": ["COHERE_API_KEY", "CO_API_KEY"],
|
||||
"perplexity": ["PERPLEXITY_API_KEY", "PPLX_API_KEY"],
|
||||
"groq": ["GROQ_API_KEY"],
|
||||
"mistral": ["MISTRAL_API_KEY"],
|
||||
"fireworks": ["FIREWORKS_API_KEY"],
|
||||
# Developer Platforms
|
||||
"github": ["GITHUB_TOKEN", "GH_TOKEN", "GITHUB_API_TOKEN"],
|
||||
"gitlab": ["GITLAB_TOKEN", "GITLAB_API_TOKEN"],
|
||||
"bitbucket": ["BITBUCKET_TOKEN", "BITBUCKET_API_TOKEN"],
|
||||
# Deployment Platforms
|
||||
"vercel": ["VERCEL_TOKEN", "VERCEL_API_TOKEN"],
|
||||
"netlify": ["NETLIFY_AUTH_TOKEN", "NETLIFY_TOKEN"],
|
||||
"fly": ["FLY_API_TOKEN", "FLY_ACCESS_TOKEN"],
|
||||
"railway": ["RAILWAY_TOKEN", "RAILWAY_API_TOKEN"],
|
||||
"render": ["RENDER_API_KEY", "RENDER_TOKEN"],
|
||||
"heroku": ["HEROKU_API_KEY", "HEROKU_TOKEN"],
|
||||
# Payment/Commerce
|
||||
"stripe": ["STRIPE_API_KEY", "STRIPE_SECRET_KEY"],
|
||||
"shopify": ["SHOPIFY_API_KEY", "SHOPIFY_ACCESS_TOKEN"],
|
||||
"paypal": ["PAYPAL_CLIENT_ID", "PAYPAL_CLIENT_SECRET"],
|
||||
# Communication
|
||||
"twilio": ["TWILIO_ACCOUNT_SID", "TWILIO_AUTH_TOKEN"],
|
||||
"sendgrid": ["SENDGRID_API_KEY"],
|
||||
"resend": ["RESEND_API_KEY"],
|
||||
"postmark": ["POSTMARK_API_TOKEN", "POSTMARK_SERVER_TOKEN"],
|
||||
"mailgun": ["MAILGUN_API_KEY"],
|
||||
"slack": ["SLACK_TOKEN", "SLACK_BOT_TOKEN", "SLACK_API_TOKEN"],
|
||||
"discord": ["DISCORD_TOKEN", "DISCORD_BOT_TOKEN"],
|
||||
# Databases/Backend
|
||||
"supabase": ["SUPABASE_API_KEY", "SUPABASE_SERVICE_KEY"],
|
||||
"planetscale": ["PLANETSCALE_TOKEN", "PSCALE_TOKEN"],
|
||||
"neon": ["NEON_API_KEY"],
|
||||
"upstash": ["UPSTASH_REDIS_REST_TOKEN", "UPSTASH_API_KEY"],
|
||||
"mongodb": ["MONGODB_API_KEY", "ATLAS_API_KEY"],
|
||||
"redis": ["REDIS_PASSWORD", "REDIS_API_KEY"],
|
||||
"fauna": ["FAUNA_SECRET", "FAUNA_KEY"],
|
||||
# Search/Analytics
|
||||
"algolia": ["ALGOLIA_API_KEY", "ALGOLIA_ADMIN_KEY"],
|
||||
"elasticsearch": ["ELASTIC_API_KEY", "ELASTICSEARCH_API_KEY"],
|
||||
"meilisearch": ["MEILI_MASTER_KEY", "MEILISEARCH_API_KEY"],
|
||||
"typesense": ["TYPESENSE_API_KEY"],
|
||||
# Monitoring/Observability
|
||||
"datadog": ["DD_API_KEY", "DATADOG_API_KEY"],
|
||||
"newrelic": ["NEW_RELIC_API_KEY", "NEWRELIC_API_KEY"],
|
||||
"sentry": ["SENTRY_AUTH_TOKEN", "SENTRY_DSN"],
|
||||
"grafana": ["GRAFANA_API_KEY", "GF_SECURITY_ADMIN_TOKEN"],
|
||||
# Storage/CDN
|
||||
"cloudinary": ["CLOUDINARY_API_KEY", "CLOUDINARY_API_SECRET"],
|
||||
"imgix": ["IMGIX_API_KEY"],
|
||||
"bunny": ["BUNNY_API_KEY", "BUNNY_ACCESS_KEY"],
|
||||
"backblaze": ["B2_APPLICATION_KEY_ID", "B2_APPLICATION_KEY"],
|
||||
# Hanzo
|
||||
"hanzo": ["HANZO_API_KEY", "HANZO_TOKEN"],
|
||||
}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Provider Configurations
|
||||
# =============================================================================
|
||||
|
||||
PROVIDER_CONFIGS: dict[str, ProviderConfig] = {
|
||||
# Cloud Providers
|
||||
"cloudflare": ProviderConfig(
|
||||
name="cloudflare",
|
||||
display_name="Cloudflare",
|
||||
base_url="https://api.cloudflare.com/client/v4",
|
||||
auth_type=AuthType.BEARER,
|
||||
spec_url="https://raw.githubusercontent.com/cloudflare/api-schemas/main/openapi.json",
|
||||
env_vars=["CLOUDFLARE_API_TOKEN", "CF_API_TOKEN", "CLOUDFLARE_API_KEY", "CF_API_KEY"],
|
||||
),
|
||||
"digitalocean": ProviderConfig(
|
||||
name="digitalocean",
|
||||
display_name="DigitalOcean",
|
||||
base_url="https://api.digitalocean.com/v2",
|
||||
auth_type=AuthType.BEARER,
|
||||
spec_url="https://api-engineering.nyc3.cdn.digitaloceanspaces.com/spec-ci/DigitalOcean-public.v2.yaml",
|
||||
env_vars=["DIGITALOCEAN_TOKEN", "DO_TOKEN"],
|
||||
),
|
||||
"hetzner": ProviderConfig(
|
||||
name="hetzner",
|
||||
display_name="Hetzner Cloud",
|
||||
base_url="https://api.hetzner.cloud/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["HETZNER_API_TOKEN", "HCLOUD_TOKEN"],
|
||||
),
|
||||
# AI Providers
|
||||
"openai": ProviderConfig(
|
||||
name="openai",
|
||||
display_name="OpenAI",
|
||||
base_url="https://api.openai.com/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
spec_url="https://raw.githubusercontent.com/openai/openai-openapi/refs/heads/manual_spec/openapi.yaml",
|
||||
env_vars=["OPENAI_API_KEY"],
|
||||
),
|
||||
"anthropic": ProviderConfig(
|
||||
name="anthropic",
|
||||
display_name="Anthropic",
|
||||
base_url="https://api.anthropic.com/v1",
|
||||
auth_type=AuthType.HEADER,
|
||||
auth_header="x-api-key",
|
||||
auth_prefix="",
|
||||
env_vars=["ANTHROPIC_API_KEY"],
|
||||
extra_headers={"anthropic-version": "2023-06-01"},
|
||||
),
|
||||
"together": ProviderConfig(
|
||||
name="together",
|
||||
display_name="Together AI",
|
||||
base_url="https://api.together.xyz/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["TOGETHER_API_KEY"],
|
||||
),
|
||||
"groq": ProviderConfig(
|
||||
name="groq",
|
||||
display_name="Groq",
|
||||
base_url="https://api.groq.com/openai/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
spec_url="https://raw.githubusercontent.com/janwilmake/handmade-openapis/main/groq.json",
|
||||
env_vars=["GROQ_API_KEY"],
|
||||
),
|
||||
"mistral": ProviderConfig(
|
||||
name="mistral",
|
||||
display_name="Mistral AI",
|
||||
base_url="https://api.mistral.ai/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["MISTRAL_API_KEY"],
|
||||
),
|
||||
# Developer Platforms
|
||||
"github": ProviderConfig(
|
||||
name="github",
|
||||
display_name="GitHub",
|
||||
base_url="https://api.github.com",
|
||||
auth_type=AuthType.BEARER,
|
||||
spec_url="https://raw.githubusercontent.com/github/rest-api-description/main/descriptions/api.github.com/api.github.com.json",
|
||||
env_vars=["GITHUB_TOKEN", "GH_TOKEN"],
|
||||
extra_headers={"Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28"},
|
||||
),
|
||||
"gitlab": ProviderConfig(
|
||||
name="gitlab",
|
||||
display_name="GitLab",
|
||||
base_url="https://gitlab.com/api/v4",
|
||||
auth_type=AuthType.HEADER,
|
||||
auth_header="PRIVATE-TOKEN",
|
||||
auth_prefix="",
|
||||
env_vars=["GITLAB_TOKEN"],
|
||||
),
|
||||
# Deployment Platforms
|
||||
"vercel": ProviderConfig(
|
||||
name="vercel",
|
||||
display_name="Vercel",
|
||||
base_url="https://api.vercel.com",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["VERCEL_TOKEN"],
|
||||
),
|
||||
"netlify": ProviderConfig(
|
||||
name="netlify",
|
||||
display_name="Netlify",
|
||||
base_url="https://api.netlify.com/api/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["NETLIFY_AUTH_TOKEN"],
|
||||
),
|
||||
"fly": ProviderConfig(
|
||||
name="fly",
|
||||
display_name="Fly.io",
|
||||
base_url="https://api.fly.io/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["FLY_API_TOKEN"],
|
||||
),
|
||||
"railway": ProviderConfig(
|
||||
name="railway",
|
||||
display_name="Railway",
|
||||
base_url="https://backboard.railway.app/graphql/v2",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["RAILWAY_TOKEN"],
|
||||
),
|
||||
"render": ProviderConfig(
|
||||
name="render",
|
||||
display_name="Render",
|
||||
base_url="https://api.render.com/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["RENDER_API_KEY"],
|
||||
),
|
||||
# Payment
|
||||
"stripe": ProviderConfig(
|
||||
name="stripe",
|
||||
display_name="Stripe",
|
||||
base_url="https://api.stripe.com/v1",
|
||||
auth_type=AuthType.BASIC,
|
||||
spec_url="https://raw.githubusercontent.com/stripe/openapi/master/openapi/spec3.json",
|
||||
env_vars=["STRIPE_API_KEY", "STRIPE_SECRET_KEY"],
|
||||
),
|
||||
# Communication
|
||||
"sendgrid": ProviderConfig(
|
||||
name="sendgrid",
|
||||
display_name="SendGrid",
|
||||
base_url="https://api.sendgrid.com/v3",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["SENDGRID_API_KEY"],
|
||||
),
|
||||
"resend": ProviderConfig(
|
||||
name="resend",
|
||||
display_name="Resend",
|
||||
base_url="https://api.resend.com",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["RESEND_API_KEY"],
|
||||
),
|
||||
"slack": ProviderConfig(
|
||||
name="slack",
|
||||
display_name="Slack",
|
||||
base_url="https://slack.com/api",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["SLACK_TOKEN", "SLACK_BOT_TOKEN"],
|
||||
),
|
||||
# Databases
|
||||
"supabase": ProviderConfig(
|
||||
name="supabase",
|
||||
display_name="Supabase",
|
||||
base_url="https://api.supabase.com/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["SUPABASE_API_KEY"],
|
||||
),
|
||||
"neon": ProviderConfig(
|
||||
name="neon",
|
||||
display_name="Neon",
|
||||
base_url="https://console.neon.tech/api/v2",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["NEON_API_KEY"],
|
||||
),
|
||||
"planetscale": ProviderConfig(
|
||||
name="planetscale",
|
||||
display_name="PlanetScale",
|
||||
base_url="https://api.planetscale.com/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["PLANETSCALE_TOKEN"],
|
||||
),
|
||||
# Monitoring
|
||||
"datadog": ProviderConfig(
|
||||
name="datadog",
|
||||
display_name="Datadog",
|
||||
base_url="https://api.datadoghq.com/api/v1",
|
||||
auth_type=AuthType.HEADER,
|
||||
auth_header="DD-API-KEY",
|
||||
auth_prefix="",
|
||||
env_vars=["DD_API_KEY", "DATADOG_API_KEY"],
|
||||
),
|
||||
"sentry": ProviderConfig(
|
||||
name="sentry",
|
||||
display_name="Sentry",
|
||||
base_url="https://sentry.io/api/0",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["SENTRY_AUTH_TOKEN"],
|
||||
),
|
||||
# Search
|
||||
"algolia": ProviderConfig(
|
||||
name="algolia",
|
||||
display_name="Algolia",
|
||||
base_url="https://api.algolia.com",
|
||||
auth_type=AuthType.HEADER,
|
||||
auth_header="X-Algolia-API-Key",
|
||||
auth_prefix="",
|
||||
env_vars=["ALGOLIA_API_KEY"],
|
||||
),
|
||||
# Hanzo
|
||||
"hanzo": ProviderConfig(
|
||||
name="hanzo",
|
||||
display_name="Hanzo AI",
|
||||
base_url="https://api.hanzo.ai/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
env_vars=["HANZO_API_KEY", "HANZO_TOKEN"],
|
||||
),
|
||||
# Popular APIs from handmade-openapis
|
||||
"notion": ProviderConfig(
|
||||
name="notion",
|
||||
display_name="Notion",
|
||||
base_url="https://api.notion.com/v1",
|
||||
auth_type=AuthType.BEARER,
|
||||
spec_url="https://raw.githubusercontent.com/janwilmake/handmade-openapis/main/notion.json",
|
||||
env_vars=["NOTION_API_KEY", "NOTION_TOKEN"],
|
||||
extra_headers={"Notion-Version": "2022-06-28"},
|
||||
),
|
||||
"hackernews": ProviderConfig(
|
||||
name="hackernews",
|
||||
display_name="Hacker News",
|
||||
base_url="https://hacker-news.firebaseio.com/v0",
|
||||
auth_type=AuthType.BEARER, # No auth needed but required by system
|
||||
spec_url="https://raw.githubusercontent.com/janwilmake/handmade-openapis/main/hackernews.json",
|
||||
env_vars=[],
|
||||
),
|
||||
"serper": ProviderConfig(
|
||||
name="serper",
|
||||
display_name="Serper (Google Search)",
|
||||
base_url="https://google.serper.dev",
|
||||
auth_type=AuthType.HEADER,
|
||||
auth_header="X-API-KEY",
|
||||
auth_prefix="",
|
||||
spec_url="https://raw.githubusercontent.com/janwilmake/handmade-openapis/main/serper.json",
|
||||
env_vars=["SERPER_API_KEY"],
|
||||
),
|
||||
"jina": ProviderConfig(
|
||||
name="jina",
|
||||
display_name="Jina Reader",
|
||||
base_url="https://r.jina.ai",
|
||||
auth_type=AuthType.BEARER,
|
||||
spec_url="https://raw.githubusercontent.com/janwilmake/handmade-openapis/main/jina-reader.json",
|
||||
env_vars=["JINA_API_KEY"],
|
||||
),
|
||||
"upstash-redis": ProviderConfig(
|
||||
name="upstash-redis",
|
||||
display_name="Upstash Redis",
|
||||
base_url="https://global.upstash.io",
|
||||
auth_type=AuthType.BEARER,
|
||||
spec_url="https://raw.githubusercontent.com/janwilmake/handmade-openapis/main/upstash-redis.json",
|
||||
env_vars=["UPSTASH_REDIS_REST_TOKEN"],
|
||||
),
|
||||
"devto": ProviderConfig(
|
||||
name="devto",
|
||||
display_name="DEV.to",
|
||||
base_url="https://dev.to/api",
|
||||
auth_type=AuthType.HEADER,
|
||||
auth_header="api-key",
|
||||
auth_prefix="",
|
||||
spec_url="https://raw.githubusercontent.com/janwilmake/handmade-openapis/main/devto.json",
|
||||
env_vars=["DEV_API_KEY", "DEVTO_API_KEY"],
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def get_provider_config(provider: str) -> ProviderConfig | None:
|
||||
"""Get configuration for a provider.
|
||||
|
||||
Checks built-in configs first, then falls back to APIs.guru.
|
||||
"""
|
||||
# Built-in configs take priority
|
||||
if provider in PROVIDER_CONFIGS:
|
||||
return PROVIDER_CONFIGS[provider]
|
||||
|
||||
# Check APIs.guru auto-generated configs
|
||||
if provider in APIS_GURU_PROVIDERS:
|
||||
guru = APIS_GURU_PROVIDERS[provider]
|
||||
return ProviderConfig(
|
||||
name=provider,
|
||||
display_name=guru.get("display_name", provider),
|
||||
base_url=guru.get("base_url", ""),
|
||||
auth_type=AuthType.BEARER,
|
||||
spec_url=guru.get("spec_url"),
|
||||
env_vars=guru.get("env_vars", [f"{provider.upper()}_API_KEY"]),
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_env_vars(provider: str) -> list[str]:
|
||||
"""Get environment variable names for a provider."""
|
||||
# Check provider config first
|
||||
config = PROVIDER_CONFIGS.get(provider)
|
||||
if config:
|
||||
return config.env_vars
|
||||
|
||||
# Check APIs.guru configs
|
||||
guru = APIS_GURU_PROVIDERS.get(provider)
|
||||
if guru:
|
||||
return guru.get("env_vars", [])
|
||||
|
||||
# Fall back to env mappings
|
||||
return ENV_VAR_MAPPINGS.get(provider, [])
|
||||
|
||||
|
||||
def list_providers() -> list[str]:
|
||||
"""List all known provider names (1100+ providers)."""
|
||||
return sorted(
|
||||
set(PROVIDER_CONFIGS.keys())
|
||||
| set(ENV_VAR_MAPPINGS.keys())
|
||||
| set(APIS_GURU_PROVIDERS.keys())
|
||||
)
|
||||
|
||||
|
||||
def list_providers_with_specs() -> list[str]:
|
||||
"""List providers that have OpenAPI spec URLs."""
|
||||
providers = []
|
||||
for name, config in PROVIDER_CONFIGS.items():
|
||||
if config.spec_url:
|
||||
providers.append(name)
|
||||
for name, guru in APIS_GURU_PROVIDERS.items():
|
||||
if guru.get("spec_url") and name not in providers:
|
||||
providers.append(name)
|
||||
return sorted(providers)
|
||||
@@ -0,0 +1,323 @@
|
||||
"""Pluggable credential storage interface.
|
||||
|
||||
Provides abstractions for storing and retrieving credentials,
|
||||
with built-in implementations for file storage and environment variables.
|
||||
|
||||
Custom implementations can integrate with KMS, 1Password, OS keychain, etc.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import aiofiles
|
||||
|
||||
from .models import Credential
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CredentialStorage(ABC):
|
||||
"""Abstract interface for credential storage.
|
||||
|
||||
Implement this interface to integrate custom secret stores
|
||||
(e.g., AWS Secrets Manager, HashiCorp Vault, 1Password, OS keychain).
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def get(self, provider: str) -> Optional[Credential]:
|
||||
"""Retrieve credential for a provider.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
Credential if found, None otherwise
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def set(self, credential: Credential) -> None:
|
||||
"""Store a credential.
|
||||
|
||||
Args:
|
||||
credential: Credential to store
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, provider: str) -> bool:
|
||||
"""Delete a stored credential.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
True if deleted, False if not found
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def list(self) -> list[str]:
|
||||
"""List all stored provider names.
|
||||
|
||||
Returns:
|
||||
List of provider names with stored credentials
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class FileCredentialStorage(CredentialStorage):
|
||||
"""File-based credential storage with obfuscation.
|
||||
|
||||
Stores credentials in a JSON file with basic base64 obfuscation
|
||||
and restrictive file permissions (0600).
|
||||
|
||||
Note: This is NOT encryption. For production use with sensitive
|
||||
credentials, consider using a proper secrets manager.
|
||||
"""
|
||||
|
||||
def __init__(self, path: Optional[Path] = None):
|
||||
"""Initialize file storage.
|
||||
|
||||
Args:
|
||||
path: Path to credentials file. Defaults to ~/.hanzo/api/credentials.json
|
||||
"""
|
||||
self.path = path or Path.home() / ".hanzo" / "api" / "credentials.json"
|
||||
self._cache: dict[str, Credential] = {}
|
||||
self._loaded = False
|
||||
|
||||
def _ensure_dir(self) -> None:
|
||||
"""Ensure parent directory exists."""
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@staticmethod
|
||||
def _obfuscate(value: str) -> str:
|
||||
"""Basic obfuscation for at-rest storage."""
|
||||
return base64.b64encode(value.encode()).decode()
|
||||
|
||||
@staticmethod
|
||||
def _deobfuscate(value: str) -> str:
|
||||
"""Reverse obfuscation."""
|
||||
try:
|
||||
return base64.b64decode(value.encode()).decode()
|
||||
except Exception:
|
||||
return value
|
||||
|
||||
async def _load(self) -> None:
|
||||
"""Load credentials from file."""
|
||||
if self._loaded:
|
||||
return
|
||||
|
||||
self._ensure_dir()
|
||||
|
||||
if self.path.exists():
|
||||
try:
|
||||
async with aiofiles.open(self.path, "r") as f:
|
||||
content = await f.read()
|
||||
data = json.loads(content)
|
||||
|
||||
for name, cred_data in data.items():
|
||||
# Deobfuscate sensitive fields
|
||||
if cred_data.get("api_key"):
|
||||
cred_data["api_key"] = self._deobfuscate(cred_data["api_key"])
|
||||
if cred_data.get("api_secret"):
|
||||
cred_data["api_secret"] = self._deobfuscate(cred_data["api_secret"])
|
||||
self._cache[name] = Credential(**cred_data)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load credentials: {e}")
|
||||
|
||||
self._loaded = True
|
||||
|
||||
async def _save(self) -> None:
|
||||
"""Save credentials to file."""
|
||||
self._ensure_dir()
|
||||
|
||||
data = {}
|
||||
for name, cred in self._cache.items():
|
||||
cred_dict = cred.model_dump()
|
||||
# Obfuscate sensitive fields
|
||||
if cred_dict.get("api_key"):
|
||||
cred_dict["api_key"] = self._obfuscate(cred_dict["api_key"])
|
||||
if cred_dict.get("api_secret"):
|
||||
cred_dict["api_secret"] = self._obfuscate(cred_dict["api_secret"])
|
||||
data[name] = cred_dict
|
||||
|
||||
async with aiofiles.open(self.path, "w") as f:
|
||||
await f.write(json.dumps(data, indent=2))
|
||||
|
||||
# Set restrictive permissions
|
||||
self.path.chmod(0o600)
|
||||
|
||||
async def get(self, provider: str) -> Optional[Credential]:
|
||||
await self._load()
|
||||
return self._cache.get(provider)
|
||||
|
||||
async def set(self, credential: Credential) -> None:
|
||||
await self._load()
|
||||
self._cache[credential.provider] = credential
|
||||
await self._save()
|
||||
|
||||
async def delete(self, provider: str) -> bool:
|
||||
await self._load()
|
||||
if provider in self._cache:
|
||||
del self._cache[provider]
|
||||
await self._save()
|
||||
return True
|
||||
return False
|
||||
|
||||
async def list(self) -> list[str]:
|
||||
await self._load()
|
||||
return list(self._cache.keys())
|
||||
|
||||
|
||||
class EnvironmentCredentialStorage(CredentialStorage):
|
||||
"""Environment variable credential storage (read-only).
|
||||
|
||||
Looks up credentials from environment variables based on
|
||||
provider-specific mappings.
|
||||
"""
|
||||
|
||||
def __init__(self, env_mappings: Optional[dict[str, list[str]]] = None):
|
||||
"""Initialize environment storage.
|
||||
|
||||
Args:
|
||||
env_mappings: Mapping of provider names to environment variable names
|
||||
"""
|
||||
from .providers import ENV_VAR_MAPPINGS
|
||||
|
||||
self.env_mappings = env_mappings or ENV_VAR_MAPPINGS
|
||||
|
||||
async def get(self, provider: str) -> Optional[Credential]:
|
||||
env_vars = self.env_mappings.get(provider, [])
|
||||
|
||||
for var in env_vars:
|
||||
value = os.environ.get(var)
|
||||
if value:
|
||||
logger.debug(f"Found credential for {provider} from env var {var}")
|
||||
return Credential(provider=provider, api_key=value)
|
||||
|
||||
return None
|
||||
|
||||
async def set(self, credential: Credential) -> None:
|
||||
"""Environment storage is read-only."""
|
||||
raise NotImplementedError("Cannot write to environment storage")
|
||||
|
||||
async def delete(self, provider: str) -> bool:
|
||||
"""Environment storage is read-only."""
|
||||
raise NotImplementedError("Cannot delete from environment storage")
|
||||
|
||||
async def list(self) -> list[str]:
|
||||
"""List providers with credentials in environment."""
|
||||
result = []
|
||||
for provider, env_vars in self.env_mappings.items():
|
||||
for var in env_vars:
|
||||
if os.environ.get(var):
|
||||
result.append(provider)
|
||||
break
|
||||
return result
|
||||
|
||||
|
||||
class MemoryCredentialStorage(CredentialStorage):
|
||||
"""In-memory credential storage.
|
||||
|
||||
Useful for per-call overrides and testing.
|
||||
Credentials are lost when the process exits.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._credentials: dict[str, Credential] = {}
|
||||
|
||||
async def get(self, provider: str) -> Optional[Credential]:
|
||||
return self._credentials.get(provider)
|
||||
|
||||
async def set(self, credential: Credential) -> None:
|
||||
self._credentials[credential.provider] = credential
|
||||
|
||||
async def delete(self, provider: str) -> bool:
|
||||
if provider in self._credentials:
|
||||
del self._credentials[provider]
|
||||
return True
|
||||
return False
|
||||
|
||||
async def list(self) -> list[str]:
|
||||
return list(self._credentials.keys())
|
||||
|
||||
|
||||
class ChainedCredentialStorage(CredentialStorage):
|
||||
"""Chains multiple credential stores with priority.
|
||||
|
||||
Checks stores in order and returns the first match.
|
||||
Writes go to the first writable store.
|
||||
|
||||
Default order:
|
||||
1. Memory (per-call overrides)
|
||||
2. File storage
|
||||
3. Environment variables
|
||||
"""
|
||||
|
||||
def __init__(self, stores: Optional[list[CredentialStorage]] = None):
|
||||
"""Initialize chained storage.
|
||||
|
||||
Args:
|
||||
stores: List of storage backends in priority order.
|
||||
Defaults to [Memory, File, Environment].
|
||||
"""
|
||||
if stores is None:
|
||||
stores = [
|
||||
MemoryCredentialStorage(),
|
||||
FileCredentialStorage(),
|
||||
EnvironmentCredentialStorage(),
|
||||
]
|
||||
self.stores = stores
|
||||
|
||||
async def get(self, provider: str) -> Optional[Credential]:
|
||||
"""Get credential from first store that has it."""
|
||||
for store in self.stores:
|
||||
cred = await store.get(provider)
|
||||
if cred and cred.has_credentials:
|
||||
return cred
|
||||
return None
|
||||
|
||||
async def get_with_source(self, provider: str) -> tuple[Optional[Credential], Optional[CredentialStorage]]:
|
||||
"""Get credential and its source store."""
|
||||
for store in self.stores:
|
||||
cred = await store.get(provider)
|
||||
if cred and cred.has_credentials:
|
||||
return cred, store
|
||||
return None, None
|
||||
|
||||
async def set(self, credential: Credential) -> None:
|
||||
"""Store credential in first writable store."""
|
||||
for store in self.stores:
|
||||
try:
|
||||
await store.set(credential)
|
||||
return
|
||||
except NotImplementedError:
|
||||
continue
|
||||
raise RuntimeError("No writable credential store available")
|
||||
|
||||
async def delete(self, provider: str) -> bool:
|
||||
"""Delete from all writable stores."""
|
||||
deleted = False
|
||||
for store in self.stores:
|
||||
try:
|
||||
if await store.delete(provider):
|
||||
deleted = True
|
||||
except NotImplementedError:
|
||||
continue
|
||||
return deleted
|
||||
|
||||
async def list(self) -> list[str]:
|
||||
"""List providers from all stores."""
|
||||
providers = set()
|
||||
for store in self.stores:
|
||||
providers.update(await store.list())
|
||||
return list(providers)
|
||||
@@ -0,0 +1,53 @@
|
||||
[project]
|
||||
name = "hanzo-tools-api"
|
||||
version = "0.3.1"
|
||||
description = "Generic API tool for calling any REST API via OpenAPI specs - search, explore, and dynamically use ANY API"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
license = { text = "MIT" }
|
||||
authors = [{ name = "Hanzo AI", email = "[email protected]" }]
|
||||
keywords = ["hanzo", "mcp", "api", "openapi", "tools"]
|
||||
classifiers = [
|
||||
"Development Status :: 4 - Beta",
|
||||
"Intended Audience :: Developers",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
]
|
||||
|
||||
dependencies = [
|
||||
"hanzo-tools-core>=0.1.0",
|
||||
"httpx>=0.27.0",
|
||||
"pydantic>=2.0",
|
||||
"pyyaml>=6.0",
|
||||
"aiofiles>=24.1.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=8.0",
|
||||
"pytest-asyncio>=0.24.0",
|
||||
"respx>=0.22.0", # Mock httpx requests
|
||||
]
|
||||
|
||||
[project.entry-points."hanzo.tools"]
|
||||
api = "hanzo_tools.api:TOOLS"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["hanzo_tools"]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
target-version = "py311"
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "B", "UP"]
|
||||
ignore = ["E501"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
@@ -0,0 +1,385 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate provider configurations from APIs.guru with LLM-optimized descriptions from oapis.org."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import urllib.request
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
# Common env var patterns
|
||||
ENV_VAR_PATTERNS = {
|
||||
'stripe': ['STRIPE_API_KEY', 'STRIPE_SECRET_KEY'],
|
||||
'twilio': ['TWILIO_ACCOUNT_SID', 'TWILIO_AUTH_TOKEN'],
|
||||
'github': ['GITHUB_TOKEN', 'GH_TOKEN'],
|
||||
'gitlab': ['GITLAB_TOKEN'],
|
||||
'slack': ['SLACK_TOKEN', 'SLACK_BOT_TOKEN'],
|
||||
'discord': ['DISCORD_TOKEN', 'DISCORD_BOT_TOKEN'],
|
||||
'shopify': ['SHOPIFY_API_KEY', 'SHOPIFY_ACCESS_TOKEN'],
|
||||
'hubspot': ['HUBSPOT_API_KEY', 'HUBSPOT_ACCESS_TOKEN'],
|
||||
'mailchimp': ['MAILCHIMP_API_KEY'],
|
||||
'sendgrid': ['SENDGRID_API_KEY'],
|
||||
'mailgun': ['MAILGUN_API_KEY'],
|
||||
'contentful': ['CONTENTFUL_ACCESS_TOKEN'],
|
||||
'airtable': ['AIRTABLE_API_KEY'],
|
||||
'notion': ['NOTION_API_KEY', 'NOTION_TOKEN'],
|
||||
'figma': ['FIGMA_TOKEN', 'FIGMA_ACCESS_TOKEN'],
|
||||
'asana': ['ASANA_TOKEN', 'ASANA_ACCESS_TOKEN'],
|
||||
'trello': ['TRELLO_API_KEY', 'TRELLO_TOKEN'],
|
||||
'jira': ['JIRA_TOKEN', 'JIRA_API_TOKEN'],
|
||||
'confluence': ['CONFLUENCE_TOKEN'],
|
||||
'bitbucket': ['BITBUCKET_TOKEN'],
|
||||
'dropbox': ['DROPBOX_ACCESS_TOKEN'],
|
||||
'box': ['BOX_ACCESS_TOKEN'],
|
||||
'spotify': ['SPOTIFY_CLIENT_ID', 'SPOTIFY_CLIENT_SECRET'],
|
||||
'youtube': ['YOUTUBE_API_KEY'],
|
||||
'twitter': ['TWITTER_API_KEY', 'TWITTER_BEARER_TOKEN'],
|
||||
'linkedin': ['LINKEDIN_ACCESS_TOKEN'],
|
||||
'facebook': ['FACEBOOK_ACCESS_TOKEN'],
|
||||
'instagram': ['INSTAGRAM_ACCESS_TOKEN'],
|
||||
'pinterest': ['PINTEREST_ACCESS_TOKEN'],
|
||||
'reddit': ['REDDIT_CLIENT_ID', 'REDDIT_CLIENT_SECRET'],
|
||||
'zendesk': ['ZENDESK_API_TOKEN'],
|
||||
'freshdesk': ['FRESHDESK_API_KEY'],
|
||||
'intercom': ['INTERCOM_ACCESS_TOKEN'],
|
||||
'mixpanel': ['MIXPANEL_TOKEN'],
|
||||
'segment': ['SEGMENT_WRITE_KEY'],
|
||||
'amplitude': ['AMPLITUDE_API_KEY'],
|
||||
'plaid': ['PLAID_CLIENT_ID', 'PLAID_SECRET'],
|
||||
'square': ['SQUARE_ACCESS_TOKEN'],
|
||||
'paypal': ['PAYPAL_CLIENT_ID', 'PAYPAL_CLIENT_SECRET'],
|
||||
'braintree': ['BRAINTREE_MERCHANT_ID'],
|
||||
'quickbooks': ['QUICKBOOKS_CLIENT_ID'],
|
||||
'xero': ['XERO_CLIENT_ID'],
|
||||
'salesforce': ['SALESFORCE_ACCESS_TOKEN'],
|
||||
'pipedrive': ['PIPEDRIVE_API_TOKEN'],
|
||||
'zoho': ['ZOHO_ACCESS_TOKEN'],
|
||||
'monday': ['MONDAY_API_KEY'],
|
||||
'clickup': ['CLICKUP_API_KEY'],
|
||||
'linear': ['LINEAR_API_KEY'],
|
||||
'vercel': ['VERCEL_TOKEN'],
|
||||
'netlify': ['NETLIFY_AUTH_TOKEN'],
|
||||
'heroku': ['HEROKU_API_KEY'],
|
||||
'digitalocean': ['DIGITALOCEAN_TOKEN', 'DO_TOKEN'],
|
||||
'linode': ['LINODE_TOKEN'],
|
||||
'vultr': ['VULTR_API_KEY'],
|
||||
'cloudflare': ['CLOUDFLARE_API_TOKEN', 'CF_API_TOKEN'],
|
||||
'datadog': ['DD_API_KEY', 'DATADOG_API_KEY'],
|
||||
'newrelic': ['NEW_RELIC_API_KEY'],
|
||||
'sentry': ['SENTRY_AUTH_TOKEN'],
|
||||
'pagerduty': ['PAGERDUTY_API_KEY'],
|
||||
'opsgenie': ['OPSGENIE_API_KEY'],
|
||||
'splunk': ['SPLUNK_TOKEN'],
|
||||
'elasticsearch': ['ELASTIC_API_KEY'],
|
||||
'algolia': ['ALGOLIA_API_KEY'],
|
||||
'meilisearch': ['MEILISEARCH_API_KEY'],
|
||||
'typesense': ['TYPESENSE_API_KEY'],
|
||||
'twitch': ['TWITCH_CLIENT_ID', 'TWITCH_CLIENT_SECRET'],
|
||||
'ebay': ['EBAY_APP_ID', 'EBAY_DEV_ID'],
|
||||
'etsy': ['ETSY_API_KEY'],
|
||||
'walmart': ['WALMART_CLIENT_ID'],
|
||||
'bestbuy': ['BESTBUY_API_KEY'],
|
||||
'yelp': ['YELP_API_KEY'],
|
||||
'foursquare': ['FOURSQUARE_API_KEY'],
|
||||
'tripadvisor': ['TRIPADVISOR_API_KEY'],
|
||||
'airbnb': ['AIRBNB_API_KEY'],
|
||||
'uber': ['UBER_ACCESS_TOKEN'],
|
||||
'lyft': ['LYFT_ACCESS_TOKEN'],
|
||||
'doordash': ['DOORDASH_API_KEY'],
|
||||
'postmates': ['POSTMATES_API_KEY'],
|
||||
'mapbox': ['MAPBOX_ACCESS_TOKEN'],
|
||||
'here': ['HERE_API_KEY'],
|
||||
'tomtom': ['TOMTOM_API_KEY'],
|
||||
'openweathermap': ['OPENWEATHERMAP_API_KEY'],
|
||||
'weatherapi': ['WEATHERAPI_KEY'],
|
||||
'newsapi': ['NEWSAPI_KEY'],
|
||||
'nytimes': ['NYTIMES_API_KEY'],
|
||||
'guardian': ['GUARDIAN_API_KEY'],
|
||||
'giphy': ['GIPHY_API_KEY'],
|
||||
'unsplash': ['UNSPLASH_ACCESS_KEY'],
|
||||
'pexels': ['PEXELS_API_KEY'],
|
||||
'cloudinary': ['CLOUDINARY_API_KEY'],
|
||||
'imgix': ['IMGIX_API_KEY'],
|
||||
'uploadcare': ['UPLOADCARE_PUBLIC_KEY'],
|
||||
'filestack': ['FILESTACK_API_KEY'],
|
||||
'agora': ['AGORA_APP_ID'],
|
||||
'vonage': ['VONAGE_API_KEY', 'NEXMO_API_KEY'],
|
||||
'messagebird': ['MESSAGEBIRD_API_KEY'],
|
||||
'bandwidth': ['BANDWIDTH_API_TOKEN'],
|
||||
'telnyx': ['TELNYX_API_KEY'],
|
||||
'apilayer': ['APILAYER_API_KEY'],
|
||||
'currencylayer': ['CURRENCYLAYER_API_KEY'],
|
||||
'exchangerate': ['EXCHANGERATE_API_KEY'],
|
||||
'coinbase': ['COINBASE_API_KEY'],
|
||||
'binance': ['BINANCE_API_KEY'],
|
||||
'kraken': ['KRAKEN_API_KEY'],
|
||||
'alchemy': ['ALCHEMY_API_KEY'],
|
||||
'infura': ['INFURA_PROJECT_ID'],
|
||||
'moralis': ['MORALIS_API_KEY'],
|
||||
'thegraph': ['THEGRAPH_API_KEY'],
|
||||
'openai': ['OPENAI_API_KEY'],
|
||||
'anthropic': ['ANTHROPIC_API_KEY'],
|
||||
'cohere': ['COHERE_API_KEY'],
|
||||
'huggingface': ['HF_TOKEN', 'HUGGINGFACE_TOKEN'],
|
||||
'replicate': ['REPLICATE_API_TOKEN'],
|
||||
'stability': ['STABILITY_API_KEY'],
|
||||
'deepl': ['DEEPL_API_KEY'],
|
||||
'ably': ['ABLY_API_KEY'],
|
||||
'pusher': ['PUSHER_APP_KEY'],
|
||||
'pubnub': ['PUBNUB_SUBSCRIBE_KEY'],
|
||||
'firebase': ['FIREBASE_API_KEY'],
|
||||
'supabase': ['SUPABASE_API_KEY'],
|
||||
'mongodb': ['MONGODB_API_KEY'],
|
||||
'redis': ['REDIS_API_KEY'],
|
||||
'cockroachdb': ['COCKROACHDB_API_KEY'],
|
||||
'planetscale': ['PLANETSCALE_TOKEN'],
|
||||
'neon': ['NEON_API_KEY'],
|
||||
'upstash': ['UPSTASH_API_KEY'],
|
||||
'fauna': ['FAUNA_SECRET'],
|
||||
'adyen': ['ADYEN_API_KEY'],
|
||||
'sportsdata': ['SPORTSDATA_API_KEY'],
|
||||
'amadeus': ['AMADEUS_API_KEY', 'AMADEUS_API_SECRET'],
|
||||
'nexmo': ['NEXMO_API_KEY', 'NEXMO_API_SECRET'],
|
||||
'mastercard': ['MASTERCARD_API_KEY'],
|
||||
'hubapi': ['HUBSPOT_API_KEY'],
|
||||
'apideck': ['APIDECK_API_KEY'],
|
||||
'codat': ['CODAT_API_KEY'],
|
||||
'deutschebahn': ['DEUTSCHEBAHN_API_KEY'],
|
||||
'rapidapi': ['RAPIDAPI_KEY'],
|
||||
'whapi': ['WHAPI_TOKEN'],
|
||||
'vtex': ['VTEX_APP_KEY', 'VTEX_APP_TOKEN'],
|
||||
'interzoid': ['INTERZOID_API_KEY'],
|
||||
}
|
||||
|
||||
# Skip cloud-specific APIs that need special auth handling
|
||||
SKIP_PREFIXES = [
|
||||
'amazonaws', 'azure', 'googleapis', 'google.',
|
||||
'apisetu', 'parliament', 'gov.', 'opto22',
|
||||
'windows', 'o365', 'microsofthealth',
|
||||
]
|
||||
|
||||
# Popular APIs with oapis.org LLM-friendly descriptions
|
||||
# These get priority loading and better descriptions
|
||||
OAPIS_POPULAR = [
|
||||
'github', 'stripe', 'twilio', 'slack', 'notion', 'discord',
|
||||
'shopify', 'spotify', 'twitter', 'openai', 'anthropic',
|
||||
]
|
||||
|
||||
|
||||
def get_env_vars(name: str) -> list[str]:
|
||||
"""Get env vars for a provider name."""
|
||||
name_lower = name.lower()
|
||||
for key, vars in ENV_VAR_PATTERNS.items():
|
||||
if key in name_lower:
|
||||
return vars
|
||||
# Generate default pattern
|
||||
clean = re.sub(r'[^a-z0-9]', '_', name_lower).upper()
|
||||
return [f'{clean}_API_KEY', f'{clean}_TOKEN']
|
||||
|
||||
|
||||
def clean_name(name: str) -> str:
|
||||
"""Clean API name for use as provider ID."""
|
||||
# Remove version suffixes
|
||||
name = re.sub(r':\d+.*$', '', name)
|
||||
# Replace dots and colons with hyphens
|
||||
name = name.replace('.', '-').replace(':', '-')
|
||||
# Remove common suffixes
|
||||
name = re.sub(r'-com$|-io$|-net$|-org$|-local$', '', name)
|
||||
# Clean up
|
||||
name = re.sub(r'[^a-z0-9-]', '', name.lower())
|
||||
name = re.sub(r'-+', '-', name).strip('-')
|
||||
return name
|
||||
|
||||
|
||||
def fetch_oapis_slop(name: str) -> dict | None:
|
||||
"""Fetch LLM-optimized description from oapis.org/slop/{name}.
|
||||
|
||||
Returns dict with base_url, endpoints_count, description if available.
|
||||
"""
|
||||
try:
|
||||
url = f"https://oapis.org/slop/{name}"
|
||||
req = urllib.request.Request(url, headers={'User-Agent': 'hanzo-tools-api/0.2'})
|
||||
with urllib.request.urlopen(req, timeout=10) as response:
|
||||
content = response.read().decode()
|
||||
|
||||
# Parse the slop format
|
||||
result = {}
|
||||
|
||||
# Extract base URL (usually on a line with https://)
|
||||
base_match = re.search(r'Base URL:\s*(https?://[^\s]+)', content, re.IGNORECASE)
|
||||
if not base_match:
|
||||
base_match = re.search(r'\*\*Base URL\*\*:\s*(https?://[^\s]+)', content, re.IGNORECASE)
|
||||
if base_match:
|
||||
result['base_url'] = base_match.group(1).rstrip('/')
|
||||
|
||||
# Extract endpoint count
|
||||
count_match = re.search(r'(\d+)\s+endpoints?', content, re.IGNORECASE)
|
||||
if count_match:
|
||||
result['endpoints_count'] = int(count_match.group(1))
|
||||
|
||||
# Extract description (first paragraph after title)
|
||||
desc_match = re.search(r'^#[^\n]+\n+([^#\n][^\n]+)', content, re.MULTILINE)
|
||||
if desc_match:
|
||||
result['description'] = desc_match.group(1).strip()[:200]
|
||||
|
||||
return result if result else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def fetch_apis_guru() -> dict:
|
||||
"""Fetch API list from APIs.guru."""
|
||||
url = "https://api.apis.guru/v2/list.json"
|
||||
with urllib.request.urlopen(url, timeout=60) as response:
|
||||
return json.loads(response.read().decode())
|
||||
|
||||
|
||||
def fetch_oapis_batch(names: list[str], max_workers: int = 10) -> dict[str, dict]:
|
||||
"""Fetch oapis.org descriptions for multiple APIs concurrently."""
|
||||
results = {}
|
||||
|
||||
def fetch_one(name: str) -> tuple[str, dict | None]:
|
||||
return name, fetch_oapis_slop(name)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
for name, data in executor.map(lambda n: fetch_one(n), names):
|
||||
if data:
|
||||
results[name] = data
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def main():
|
||||
print("Fetching APIs from APIs.guru...", file=sys.stderr)
|
||||
data = fetch_apis_guru()
|
||||
print(f"Found {len(data)} APIs", file=sys.stderr)
|
||||
|
||||
# Fetch oapis.org LLM descriptions for popular APIs
|
||||
print(f"Fetching LLM descriptions from oapis.org for {len(OAPIS_POPULAR)} popular APIs...", file=sys.stderr)
|
||||
oapis_data = fetch_oapis_batch(OAPIS_POPULAR)
|
||||
print(f"Got {len(oapis_data)} oapis.org descriptions", file=sys.stderr)
|
||||
|
||||
apis = []
|
||||
seen_names = set()
|
||||
|
||||
for name, info in data.items():
|
||||
# Skip cloud-specific
|
||||
skip = False
|
||||
for prefix in SKIP_PREFIXES:
|
||||
if name.lower().startswith(prefix.lower()):
|
||||
skip = True
|
||||
break
|
||||
if skip:
|
||||
continue
|
||||
|
||||
preferred = info.get('preferred', '')
|
||||
if not preferred or preferred not in info.get('versions', {}):
|
||||
continue
|
||||
|
||||
version_info = info['versions'][preferred]
|
||||
spec_url = version_info.get('swaggerUrl', '')
|
||||
if not spec_url:
|
||||
continue
|
||||
|
||||
api_info = version_info.get('info', {})
|
||||
title = api_info.get('title', name)
|
||||
|
||||
clean = clean_name(name)
|
||||
if clean in seen_names or len(clean) < 2:
|
||||
continue
|
||||
seen_names.add(clean)
|
||||
|
||||
# Get base URL from x-origin
|
||||
base_url = ''
|
||||
if 'x-origin' in api_info:
|
||||
x_origin = api_info['x-origin']
|
||||
if isinstance(x_origin, list) and x_origin:
|
||||
base_url = x_origin[0].get('url', '')
|
||||
|
||||
# Check for oapis.org enhanced data
|
||||
oapis_key = None
|
||||
for popular in OAPIS_POPULAR:
|
||||
if popular in clean:
|
||||
oapis_key = popular
|
||||
break
|
||||
|
||||
description = api_info.get('description', '')[:200] if api_info.get('description') else ''
|
||||
endpoints_count = None
|
||||
|
||||
if oapis_key and oapis_key in oapis_data:
|
||||
oapis_info = oapis_data[oapis_key]
|
||||
if oapis_info.get('base_url'):
|
||||
base_url = oapis_info['base_url']
|
||||
if oapis_info.get('description'):
|
||||
description = oapis_info['description']
|
||||
if oapis_info.get('endpoints_count'):
|
||||
endpoints_count = oapis_info['endpoints_count']
|
||||
|
||||
# Clean strings - remove newlines, escape quotes, limit length
|
||||
def clean_str(s: str, max_len: int = 150) -> str:
|
||||
if not s:
|
||||
return ""
|
||||
# Replace various whitespace with single space
|
||||
s = ' '.join(s.split())
|
||||
# Escape backslashes first, then quotes
|
||||
s = s.replace('\\', '\\\\').replace('"', '\\"')
|
||||
return s[:max_len]
|
||||
|
||||
apis.append({
|
||||
'name': clean,
|
||||
'title': clean_str(title, 60),
|
||||
'description': clean_str(description, 150),
|
||||
'spec_url': spec_url,
|
||||
'base_url': base_url,
|
||||
'env_vars': get_env_vars(clean),
|
||||
'endpoints_count': endpoints_count,
|
||||
})
|
||||
|
||||
# Sort by name
|
||||
apis.sort(key=lambda x: x['name'])
|
||||
|
||||
# Generate Python code
|
||||
print('"""Auto-generated provider configurations from APIs.guru + oapis.org.')
|
||||
print()
|
||||
print(f'Generated {len(apis)} provider configurations.')
|
||||
print('Includes LLM-optimized descriptions from oapis.org for popular APIs.')
|
||||
print()
|
||||
print('Regenerate with: python scripts/generate_providers.py > hanzo_tools/api/apis_guru_providers.py')
|
||||
print('"""')
|
||||
print()
|
||||
print('from typing import Any')
|
||||
print()
|
||||
print()
|
||||
print('# =============================================================================')
|
||||
print('# APIs.guru Provider Configurations (with oapis.org enhancements)')
|
||||
print('# =============================================================================')
|
||||
print()
|
||||
print('APIS_GURU_PROVIDERS: dict[str, dict[str, Any]] = {')
|
||||
|
||||
for api in apis:
|
||||
env_str = ', '.join(f'"{v}"' for v in api['env_vars'][:2])
|
||||
print(f' "{api["name"]}": {{')
|
||||
print(f' "display_name": "{api["title"]}",')
|
||||
if api['description']:
|
||||
print(f' "description": "{api["description"][:150]}",')
|
||||
print(f' "spec_url": "{api["spec_url"]}",')
|
||||
if api['base_url']:
|
||||
safe_url = api['base_url'].replace('"', '\\"')
|
||||
print(f' "base_url": "{safe_url}",')
|
||||
print(f' "env_vars": [{env_str}],')
|
||||
if api['endpoints_count']:
|
||||
print(f' "endpoints_count": {api["endpoints_count"]},')
|
||||
print(' },')
|
||||
|
||||
print('}')
|
||||
print()
|
||||
print()
|
||||
print(f'# Total: {len(apis)} providers')
|
||||
|
||||
print(f"\nGenerated {len(apis)} providers", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,140 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Pre-download and cache popular OpenAPI specs for faster startup.
|
||||
|
||||
Usage:
|
||||
python scripts/preload_specs.py # Download top 50 specs
|
||||
python scripts/preload_specs.py --all # Download all specs with URLs
|
||||
python scripts/preload_specs.py --list # List available specs
|
||||
python scripts/preload_specs.py github stripe # Download specific specs
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add parent to path for local development
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from hanzo_tools.api.client import APIClient
|
||||
from hanzo_tools.api.providers import list_providers_with_specs, get_provider_config
|
||||
|
||||
# Top 50 most popular APIs to pre-cache
|
||||
POPULAR_APIS = [
|
||||
# AI/ML
|
||||
"openai", "anthropic", "together", "groq", "mistral", "cohere", "replicate",
|
||||
# Developer
|
||||
"github", "gitlab", "bitbucket", "vercel", "netlify", "fly", "railway", "render",
|
||||
# Cloud
|
||||
"cloudflare", "digitalocean", "hetzner", "linode", "vultr",
|
||||
# Payment
|
||||
"stripe", "paypal", "adyen-com-accountservice", "square",
|
||||
# Communication
|
||||
"twilio", "sendgrid", "resend", "slack", "discord",
|
||||
# Databases
|
||||
"supabase", "neon", "planetscale", "mongodb", "redis",
|
||||
# Search
|
||||
"algolia", "elasticsearch", "meilisearch", "typesense",
|
||||
# Monitoring
|
||||
"datadog", "sentry", "newrelic",
|
||||
# Other popular
|
||||
"notion", "airtable", "hubspot", "shopify", "jira",
|
||||
]
|
||||
|
||||
|
||||
async def preload_specs(
|
||||
providers: list[str] | None = None,
|
||||
concurrent: int = 10,
|
||||
verbose: bool = True,
|
||||
) -> dict[str, bool]:
|
||||
"""Pre-download and cache OpenAPI specs.
|
||||
|
||||
Args:
|
||||
providers: List of provider names, or None for POPULAR_APIS
|
||||
concurrent: Number of concurrent downloads
|
||||
verbose: Print progress
|
||||
|
||||
Returns:
|
||||
Dict mapping provider name to success status
|
||||
"""
|
||||
client = APIClient()
|
||||
|
||||
if providers is None:
|
||||
providers = POPULAR_APIS
|
||||
|
||||
# Filter to only those with spec URLs
|
||||
valid_providers = []
|
||||
for p in providers:
|
||||
config = get_provider_config(p)
|
||||
if config and config.spec_url:
|
||||
valid_providers.append(p)
|
||||
elif verbose:
|
||||
print(f" ⚠ {p}: no spec URL configured")
|
||||
|
||||
if verbose:
|
||||
print(f"Pre-loading {len(valid_providers)} OpenAPI specs...")
|
||||
print()
|
||||
|
||||
results: dict[str, bool] = {}
|
||||
semaphore = asyncio.Semaphore(concurrent)
|
||||
|
||||
async def fetch_one(provider: str) -> tuple[str, bool]:
|
||||
async with semaphore:
|
||||
try:
|
||||
await client.spec(provider)
|
||||
if verbose:
|
||||
print(f" ✓ {provider}")
|
||||
return provider, True
|
||||
except Exception as e:
|
||||
if verbose:
|
||||
print(f" ✗ {provider}: {e}")
|
||||
return provider, False
|
||||
|
||||
tasks = [fetch_one(p) for p in valid_providers]
|
||||
for coro in asyncio.as_completed(tasks):
|
||||
provider, success = await coro
|
||||
results[provider] = success
|
||||
|
||||
if verbose:
|
||||
success_count = sum(1 for v in results.values() if v)
|
||||
print()
|
||||
print(f"Downloaded {success_count}/{len(results)} specs")
|
||||
print(f"Cached at: {client._spec_cache.cache_dir}")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
async def main():
|
||||
parser = argparse.ArgumentParser(description="Pre-download OpenAPI specs")
|
||||
parser.add_argument("providers", nargs="*", help="Specific providers to download")
|
||||
parser.add_argument("--all", action="store_true", help="Download all available specs")
|
||||
parser.add_argument("--list", action="store_true", help="List available specs")
|
||||
parser.add_argument("-n", "--concurrent", type=int, default=10, help="Concurrent downloads")
|
||||
parser.add_argument("-q", "--quiet", action="store_true", help="Quiet mode")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.list:
|
||||
specs = list_providers_with_specs()
|
||||
print(f"Available specs ({len(specs)}):")
|
||||
for i, name in enumerate(specs, 1):
|
||||
config = get_provider_config(name)
|
||||
display = config.display_name if config else name
|
||||
print(f" {i:4}. {name}: {display}")
|
||||
return
|
||||
|
||||
providers = None
|
||||
if args.providers:
|
||||
providers = args.providers
|
||||
elif args.all:
|
||||
providers = list_providers_with_specs()
|
||||
|
||||
await preload_specs(
|
||||
providers=providers,
|
||||
concurrent=args.concurrent,
|
||||
verbose=not args.quiet,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1 @@
|
||||
# Tests for hanzo-tools-api
|
||||
@@ -0,0 +1,437 @@
|
||||
"""Tests for hanzo-tools-api v0.2.0."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from hanzo_tools.api import (
|
||||
APIClient,
|
||||
APITool,
|
||||
Credential,
|
||||
CredentialManager,
|
||||
OpenAPIClient,
|
||||
SpecCache,
|
||||
get_credential_manager,
|
||||
reset_credential_manager,
|
||||
PROVIDER_CONFIGS,
|
||||
ENV_VAR_MAPPINGS,
|
||||
)
|
||||
from hanzo_tools.api.models import AuthType
|
||||
from hanzo_tools.api.storage import (
|
||||
MemoryCredentialStorage,
|
||||
ChainedCredentialStorage,
|
||||
)
|
||||
|
||||
|
||||
class TestCredentialManager:
|
||||
"""Tests for CredentialManager."""
|
||||
|
||||
@pytest.fixture
|
||||
def temp_config_dir(self):
|
||||
"""Create a temporary config directory."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
yield Path(tmpdir)
|
||||
|
||||
@pytest.fixture
|
||||
def cred_manager(self, temp_config_dir):
|
||||
"""Create a CredentialManager with temp directory."""
|
||||
return CredentialManager(config_dir=temp_config_dir)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_and_get_credential(self, cred_manager):
|
||||
"""Test storing and retrieving credentials."""
|
||||
await cred_manager.set_credential(
|
||||
provider="test-provider",
|
||||
api_key="test-key",
|
||||
api_secret="test-secret",
|
||||
account_id="test-account",
|
||||
)
|
||||
|
||||
cred = await cred_manager.get_credential("test-provider")
|
||||
assert cred.provider == "test-provider"
|
||||
assert cred.api_key == "test-key"
|
||||
assert cred.api_secret == "test-secret"
|
||||
assert cred.account_id == "test-account"
|
||||
assert cred.has_credentials
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credentials_persist(self, temp_config_dir):
|
||||
"""Test that credentials persist across instances."""
|
||||
# Set credential
|
||||
cm1 = CredentialManager(config_dir=temp_config_dir)
|
||||
await cm1.set_credential("test", api_key="my-key")
|
||||
|
||||
# Read with new instance
|
||||
cm2 = CredentialManager(config_dir=temp_config_dir)
|
||||
cred = await cm2.get_credential("test")
|
||||
assert cred.api_key == "my-key"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_credential(self, cred_manager):
|
||||
"""Test deleting credentials."""
|
||||
await cred_manager.set_credential("test", api_key="key")
|
||||
assert await cred_manager.delete_credential("test")
|
||||
assert not await cred_manager.delete_credential("test") # Already deleted
|
||||
|
||||
cred = await cred_manager.get_credential("test")
|
||||
assert not cred.has_credentials
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_env_var_fallback(self, temp_config_dir):
|
||||
"""Test environment variable fallback."""
|
||||
# Use a unique provider with custom env var mapping
|
||||
with patch.dict(os.environ, {"TESTPROV_API_KEY": "env-token"}, clear=False):
|
||||
# Create custom storage chain with our test env var
|
||||
from hanzo_tools.api.storage import EnvironmentCredentialStorage
|
||||
|
||||
env_mappings = {"testprov": ["TESTPROV_API_KEY"]}
|
||||
storage = ChainedCredentialStorage([
|
||||
MemoryCredentialStorage(),
|
||||
EnvironmentCredentialStorage(env_mappings),
|
||||
])
|
||||
|
||||
cred_manager = CredentialManager(
|
||||
config_dir=temp_config_dir,
|
||||
storage=storage,
|
||||
)
|
||||
cred = await cred_manager.get_credential("testprov")
|
||||
assert cred.api_key == "env-token"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_providers(self, cred_manager):
|
||||
"""Test listing providers."""
|
||||
await cred_manager.set_credential("custom", api_key="key")
|
||||
|
||||
providers = await cred_manager.list_providers()
|
||||
provider_names = [p.name for p in providers]
|
||||
|
||||
# Should include configured providers
|
||||
assert "cloudflare" in provider_names
|
||||
assert "github" in provider_names
|
||||
assert "custom" in provider_names
|
||||
|
||||
# Custom should be configured
|
||||
custom = next(p for p in providers if p.name == "custom")
|
||||
assert custom.configured
|
||||
|
||||
def test_get_provider_config(self, cred_manager):
|
||||
"""Test getting provider configuration."""
|
||||
config = cred_manager.get_provider_config("cloudflare")
|
||||
assert config is not None
|
||||
assert config.name == "cloudflare"
|
||||
assert config.base_url == "https://api.cloudflare.com/client/v4"
|
||||
assert config.auth_type == AuthType.BEARER
|
||||
|
||||
|
||||
class TestProviderConfigs:
|
||||
"""Tests for built-in provider configurations."""
|
||||
|
||||
def test_all_providers_have_required_fields(self):
|
||||
"""Test that all providers have required fields."""
|
||||
for name, config in PROVIDER_CONFIGS.items():
|
||||
assert config.name == name
|
||||
assert config.display_name
|
||||
assert config.base_url
|
||||
assert config.auth_type in AuthType
|
||||
|
||||
def test_env_var_mappings_exist(self):
|
||||
"""Test that env var mappings cover common providers."""
|
||||
expected_providers = ["cloudflare", "github", "openai", "anthropic", "stripe"]
|
||||
for provider in expected_providers:
|
||||
assert provider in ENV_VAR_MAPPINGS
|
||||
assert len(ENV_VAR_MAPPINGS[provider]) > 0
|
||||
|
||||
|
||||
class TestCredential:
|
||||
"""Tests for Credential model."""
|
||||
|
||||
def test_has_credentials(self):
|
||||
"""Test has_credentials property."""
|
||||
cred = Credential(provider="test")
|
||||
assert not cred.has_credentials
|
||||
|
||||
cred = Credential(provider="test", api_key="key")
|
||||
assert cred.has_credentials
|
||||
|
||||
|
||||
class TestOpenAPIClient:
|
||||
"""Tests for OpenAPIClient."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_spec(self):
|
||||
"""Sample OpenAPI spec for testing."""
|
||||
return {
|
||||
"openapi": "3.0.0",
|
||||
"info": {"title": "Test API", "version": "1.0"},
|
||||
"servers": [{"url": "https://api.test.com/v1"}],
|
||||
"paths": {
|
||||
"/users": {
|
||||
"get": {
|
||||
"operationId": "listUsers",
|
||||
"summary": "List all users",
|
||||
"tags": ["users"],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "limit",
|
||||
"in": "query",
|
||||
"schema": {"type": "integer"},
|
||||
}
|
||||
],
|
||||
"responses": {"200": {"description": "Success"}},
|
||||
},
|
||||
"post": {
|
||||
"operationId": "createUser",
|
||||
"summary": "Create a user",
|
||||
"tags": ["users"],
|
||||
"requestBody": {
|
||||
"required": True,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {"name": {"type": "string"}},
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
"responses": {"201": {"description": "Created"}},
|
||||
},
|
||||
},
|
||||
"/users/{id}": {
|
||||
"get": {
|
||||
"operationId": "getUser",
|
||||
"summary": "Get a user",
|
||||
"tags": ["users"],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": True,
|
||||
"schema": {"type": "string"},
|
||||
}
|
||||
],
|
||||
"responses": {"200": {"description": "Success"}},
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def temp_config_dir(self):
|
||||
"""Create a temporary config directory."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
yield Path(tmpdir)
|
||||
|
||||
@pytest.fixture
|
||||
async def client(self, sample_spec, temp_config_dir):
|
||||
"""Create an OpenAPIClient with sample spec."""
|
||||
# Create credential manager
|
||||
cred_manager = CredentialManager(config_dir=temp_config_dir)
|
||||
await cred_manager.set_credential("test", api_key="test-key")
|
||||
|
||||
# Create spec cache and seed it with sample spec
|
||||
spec_cache = SpecCache(temp_config_dir / "specs")
|
||||
await spec_cache.set("test", sample_spec)
|
||||
|
||||
# Create client
|
||||
client = OpenAPIClient(
|
||||
"test",
|
||||
credential_manager=cred_manager,
|
||||
spec_cache=spec_cache,
|
||||
)
|
||||
|
||||
# Load the spec
|
||||
await client.load_spec()
|
||||
|
||||
return client
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_url_from_spec(self, client):
|
||||
"""Test base URL extraction from spec."""
|
||||
assert client.base_url == "https://api.test.com/v1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_operations(self, client):
|
||||
"""Test listing operations."""
|
||||
result = client.list_operations()
|
||||
assert result.total_count == 3
|
||||
op_ids = [op.operation_id for op in result.operations]
|
||||
assert "listUsers" in op_ids
|
||||
assert "createUser" in op_ids
|
||||
assert "getUser" in op_ids
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_operations_by_tag(self, client):
|
||||
"""Test filtering operations by tag."""
|
||||
result = client.list_operations(tag="users")
|
||||
assert result.total_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_operations_by_search(self, client):
|
||||
"""Test searching operations."""
|
||||
result = client.list_operations(search="list")
|
||||
assert result.total_count == 1
|
||||
assert result.operations[0].operation_id == "listUsers"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_operation(self, client):
|
||||
"""Test getting a specific operation."""
|
||||
op = client.get_operation("getUser")
|
||||
assert op is not None
|
||||
assert op.operation_id == "getUser"
|
||||
assert op.method == "GET"
|
||||
assert op.path == "/users/{id}"
|
||||
assert len(op.parameters) == 1
|
||||
assert op.parameters[0].name == "id"
|
||||
assert op.parameters[0].required
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_path_param_substitution(self, client):
|
||||
"""Test path parameter substitution."""
|
||||
path, remaining = client._substitute_path_params(
|
||||
"/users/{id}/posts/{post_id}",
|
||||
{"id": "123", "post_id": "456", "extra": "value"},
|
||||
)
|
||||
assert path == "/users/123/posts/456"
|
||||
assert remaining == {"extra": "value"}
|
||||
|
||||
|
||||
class TestAPIClient:
|
||||
"""Tests for the main APIClient."""
|
||||
|
||||
@pytest.fixture
|
||||
def temp_config_dir(self):
|
||||
"""Create a temporary config directory."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
yield Path(tmpdir)
|
||||
|
||||
@pytest.fixture
|
||||
def api_client(self, temp_config_dir):
|
||||
"""Create APIClient with temp directory."""
|
||||
return APIClient(config_dir=temp_config_dir)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_providers(self, api_client):
|
||||
"""Test listing providers."""
|
||||
result = await api_client.list_providers()
|
||||
assert result.total_count > 0
|
||||
|
||||
provider_names = [p.name for p in result.providers]
|
||||
assert "cloudflare" in provider_names
|
||||
assert "github" in provider_names
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_and_get_credentials(self, api_client):
|
||||
"""Test configuring and retrieving credentials."""
|
||||
await api_client.config("test-api", api_key="my-key")
|
||||
|
||||
effective = await api_client.get_effective_credentials("test-api")
|
||||
assert effective.has_credentials
|
||||
assert effective.credential.api_key == "my-key"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_config(self, api_client):
|
||||
"""Test deleting configuration."""
|
||||
await api_client.config("test-api", api_key="my-key")
|
||||
assert await api_client.delete_config("test-api")
|
||||
assert not await api_client.delete_config("test-api") # Already deleted
|
||||
|
||||
|
||||
class TestAPITool:
|
||||
"""Tests for APITool."""
|
||||
|
||||
@pytest.fixture
|
||||
def temp_config_dir(self):
|
||||
"""Create a temporary config directory."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
yield Path(tmpdir)
|
||||
|
||||
@pytest.fixture
|
||||
def tool(self, temp_config_dir):
|
||||
"""Create APITool with mocked client."""
|
||||
client = APIClient(config_dir=temp_config_dir)
|
||||
return APITool(client=client)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_action(self, tool):
|
||||
"""Test list action."""
|
||||
ctx = AsyncMock()
|
||||
result = await tool.call(ctx, action="list")
|
||||
assert "API Providers" in result
|
||||
assert "cloudflare" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_action(self, tool):
|
||||
"""Test config action."""
|
||||
ctx = AsyncMock()
|
||||
|
||||
# Configure
|
||||
result = await tool.call(
|
||||
ctx,
|
||||
action="config",
|
||||
provider="test-api",
|
||||
api_key="my-key",
|
||||
)
|
||||
assert "Configured" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_requires_provider(self, tool):
|
||||
"""Test that config action requires provider."""
|
||||
ctx = AsyncMock()
|
||||
result = await tool.call(ctx, action="config", api_key="key")
|
||||
assert "Error" in result
|
||||
assert "provider" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_requires_api_key(self, tool):
|
||||
"""Test that config action requires api_key."""
|
||||
ctx = AsyncMock()
|
||||
result = await tool.call(ctx, action="config", provider="test")
|
||||
assert "Error" in result
|
||||
assert "api_key" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_action(self, tool):
|
||||
"""Test unknown action."""
|
||||
ctx = AsyncMock()
|
||||
result = await tool.call(ctx, action="invalid")
|
||||
assert "Unknown action" in result
|
||||
|
||||
def test_tool_name(self, tool):
|
||||
"""Test tool name."""
|
||||
assert tool.name == "api"
|
||||
|
||||
def test_tool_description(self, tool):
|
||||
"""Test tool description."""
|
||||
desc = tool.description
|
||||
assert "API" in desc
|
||||
assert "OpenAPI" in desc
|
||||
assert "credential" in desc.lower()
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests (require network, skip by default)."""
|
||||
|
||||
@pytest.mark.skip(reason="Requires network and API keys")
|
||||
@pytest.mark.asyncio
|
||||
async def test_github_api(self):
|
||||
"""Test calling GitHub API."""
|
||||
# This would require GITHUB_TOKEN to be set
|
||||
async with APIClient() as client:
|
||||
await client.spec("github")
|
||||
result = await client.ops("github", search="user")
|
||||
assert result.total_count > 0
|
||||
|
||||
@pytest.mark.skip(reason="Requires network and API keys")
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloudflare_api(self):
|
||||
"""Test calling Cloudflare API."""
|
||||
# This would require CLOUDFLARE_API_TOKEN to be set
|
||||
async with APIClient() as client:
|
||||
await client.spec("cloudflare")
|
||||
result = await client.call("cloudflare", "listZones")
|
||||
assert result.success
|
||||
Generated
+1581
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "hanzo"
|
||||
version = "0.3.31"
|
||||
version = "0.3.32"
|
||||
description = "Hanzo AI - Complete AI Infrastructure Platform with CLI, Router, MCP, and Agent Runtime"
|
||||
authors = [
|
||||
{name = "Hanzo AI", email = "[email protected]"},
|
||||
@@ -73,6 +73,40 @@ repl = [
|
||||
# "hanzo-router>=1.74.3", # TODO: For AI chat functionality - publish to PyPI
|
||||
]
|
||||
|
||||
# Infrastructure service dependencies
|
||||
infra = [
|
||||
# Vector database (Qdrant)
|
||||
"qdrant-client>=1.7.0",
|
||||
# Key-value store (Redis/Valkey)
|
||||
"redis>=5.0.0",
|
||||
# Document database (MongoDB)
|
||||
"motor>=3.3.0",
|
||||
"pymongo>=4.6.0",
|
||||
# Object storage (S3/MinIO)
|
||||
"aiobotocore>=2.9.0",
|
||||
"botocore>=1.34.0",
|
||||
# Full-text search (Meilisearch)
|
||||
"meilisearch-python-sdk>=3.0.0",
|
||||
# Messaging (NATS)
|
||||
"nats-py>=2.6.0",
|
||||
# Workflows (Temporal)
|
||||
"temporalio>=1.4.0",
|
||||
# Scheduled jobs (cron parsing)
|
||||
"croniter>=2.0.0",
|
||||
# HTTP client for Nuclio
|
||||
"httpx>=0.23.0",
|
||||
]
|
||||
infra-vector = ["qdrant-client>=1.7.0"]
|
||||
infra-kv = ["redis>=5.0.0"]
|
||||
infra-documentdb = ["motor>=3.3.0", "pymongo>=4.6.0"]
|
||||
infra-storage = ["aiobotocore>=2.9.0", "botocore>=1.34.0"]
|
||||
infra-search = ["meilisearch-python-sdk>=3.0.0"]
|
||||
infra-pubsub = ["nats-py>=2.6.0"]
|
||||
infra-tasks = ["temporalio>=1.4.0"]
|
||||
infra-queues = ["redis>=5.0.0"]
|
||||
infra-cron = ["redis>=5.0.0", "croniter>=2.0.0"]
|
||||
infra-functions = ["httpx>=0.23.0"]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://hanzo.ai"
|
||||
Repository = "https://github.com/hanzoai/python-sdk"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Hanzo - Complete AI Infrastructure Platform with CLI, Router, MCP, and Agent Runtime."""
|
||||
|
||||
__version__ = "0.3.31"
|
||||
__version__ = "0.3.32"
|
||||
__all__ = ["main", "cli", "__version__"]
|
||||
|
||||
from .cli import cli, main
|
||||
|
||||
@@ -0,0 +1,408 @@
|
||||
"""Hanzo Infrastructure SDK - Unified client for Hanzo OSS infrastructure.
|
||||
|
||||
Provides async clients for all Hanzo infrastructure services:
|
||||
- Vector: Qdrant vector database
|
||||
- KV: Redis/Valkey key-value store
|
||||
- DocumentDB: MongoDB document database
|
||||
- Storage: S3/MinIO object storage
|
||||
- Search: Meilisearch full-text search
|
||||
- PubSub: NATS messaging
|
||||
- Tasks: Temporal workflows
|
||||
- Queues: Distributed work queues
|
||||
- Cron: Scheduled jobs
|
||||
- Functions: Nuclio serverless functions
|
||||
|
||||
Example:
|
||||
```python
|
||||
from hanzo.infra import HanzoInfra
|
||||
|
||||
async def main():
|
||||
# Initialize from environment variables
|
||||
infra = await HanzoInfra.from_env()
|
||||
|
||||
# Use individual clients
|
||||
await infra.kv.set("key", "value")
|
||||
await infra.vector.upsert("embeddings", points)
|
||||
|
||||
# Or with context manager
|
||||
async with await HanzoInfra.from_env() as infra:
|
||||
results = await infra.search.search("products", "query")
|
||||
|
||||
# Selective initialization
|
||||
infra = await HanzoInfra.from_env(services=["kv", "vector"])
|
||||
```
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any, Optional, Sequence
|
||||
|
||||
from .vector import VectorClient, VectorConfig, VectorPoint, ScoredPoint
|
||||
from .kv import KVClient, KVConfig
|
||||
from .documentdb import DocumentDBClient, DocumentDBConfig, Document, UpdateResult, DeleteResult
|
||||
from .storage import StorageClient, StorageConfig, ObjectInfo, UploadResult, PresignedUrl
|
||||
from .search import SearchClient, SearchConfig, SearchHit, SearchResult
|
||||
from .pubsub import PubSubClient, PubSubConfig, Message, Subscription
|
||||
from .tasks import TasksClient, TasksConfig, WorkflowHandle, WorkflowExecution
|
||||
from .queues import QueuesClient, QueuesConfig, Job, JobStatus, QueueStats
|
||||
from .cron import CronClient, CronConfig, CronJob, CronExecution
|
||||
from .functions import FunctionsClient, FunctionsConfig, FunctionSpec, FunctionStatus, InvokeResult
|
||||
|
||||
|
||||
__all__ = [
|
||||
# Main class
|
||||
"HanzoInfra",
|
||||
# Vector (Qdrant)
|
||||
"VectorClient",
|
||||
"VectorConfig",
|
||||
"VectorPoint",
|
||||
"ScoredPoint",
|
||||
# KV (Redis/Valkey)
|
||||
"KVClient",
|
||||
"KVConfig",
|
||||
# DocumentDB (MongoDB)
|
||||
"DocumentDBClient",
|
||||
"DocumentDBConfig",
|
||||
"Document",
|
||||
"UpdateResult",
|
||||
"DeleteResult",
|
||||
# Storage (S3/MinIO)
|
||||
"StorageClient",
|
||||
"StorageConfig",
|
||||
"ObjectInfo",
|
||||
"UploadResult",
|
||||
"PresignedUrl",
|
||||
# Search (Meilisearch)
|
||||
"SearchClient",
|
||||
"SearchConfig",
|
||||
"SearchHit",
|
||||
"SearchResult",
|
||||
# PubSub (NATS)
|
||||
"PubSubClient",
|
||||
"PubSubConfig",
|
||||
"Message",
|
||||
"Subscription",
|
||||
# Tasks (Temporal)
|
||||
"TasksClient",
|
||||
"TasksConfig",
|
||||
"WorkflowHandle",
|
||||
"WorkflowExecution",
|
||||
# Queues
|
||||
"QueuesClient",
|
||||
"QueuesConfig",
|
||||
"Job",
|
||||
"JobStatus",
|
||||
"QueueStats",
|
||||
# Cron
|
||||
"CronClient",
|
||||
"CronConfig",
|
||||
"CronJob",
|
||||
"CronExecution",
|
||||
# Functions (Nuclio)
|
||||
"FunctionsClient",
|
||||
"FunctionsConfig",
|
||||
"FunctionSpec",
|
||||
"FunctionStatus",
|
||||
"InvokeResult",
|
||||
]
|
||||
|
||||
|
||||
ALL_SERVICES = [
|
||||
"vector",
|
||||
"kv",
|
||||
"documentdb",
|
||||
"storage",
|
||||
"search",
|
||||
"pubsub",
|
||||
"tasks",
|
||||
"queues",
|
||||
"cron",
|
||||
"functions",
|
||||
]
|
||||
|
||||
|
||||
class HanzoInfra:
|
||||
"""Unified client for Hanzo infrastructure services.
|
||||
|
||||
Provides a single entry point for all infrastructure services,
|
||||
with lazy initialization and configurable service selection.
|
||||
|
||||
Attributes:
|
||||
vector: Qdrant vector database client.
|
||||
kv: Redis/Valkey key-value client.
|
||||
documentdb: MongoDB document database client.
|
||||
storage: S3/MinIO object storage client.
|
||||
search: Meilisearch full-text search client.
|
||||
pubsub: NATS messaging client.
|
||||
tasks: Temporal workflow client.
|
||||
queues: Distributed work queue client.
|
||||
cron: Scheduled jobs client.
|
||||
functions: Nuclio serverless functions client.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vector: Optional[VectorClient] = None,
|
||||
kv: Optional[KVClient] = None,
|
||||
documentdb: Optional[DocumentDBClient] = None,
|
||||
storage: Optional[StorageClient] = None,
|
||||
search: Optional[SearchClient] = None,
|
||||
pubsub: Optional[PubSubClient] = None,
|
||||
tasks: Optional[TasksClient] = None,
|
||||
queues: Optional[QueuesClient] = None,
|
||||
cron: Optional[CronClient] = None,
|
||||
functions: Optional[FunctionsClient] = None,
|
||||
) -> None:
|
||||
"""Initialize HanzoInfra with individual clients.
|
||||
|
||||
Args:
|
||||
vector: Qdrant vector client.
|
||||
kv: Redis/Valkey client.
|
||||
documentdb: MongoDB client.
|
||||
storage: S3/MinIO client.
|
||||
search: Meilisearch client.
|
||||
pubsub: NATS client.
|
||||
tasks: Temporal client.
|
||||
queues: Work queue client.
|
||||
cron: Cron scheduler client.
|
||||
functions: Nuclio functions client.
|
||||
"""
|
||||
self._vector = vector
|
||||
self._kv = kv
|
||||
self._documentdb = documentdb
|
||||
self._storage = storage
|
||||
self._search = search
|
||||
self._pubsub = pubsub
|
||||
self._tasks = tasks
|
||||
self._queues = queues
|
||||
self._cron = cron
|
||||
self._functions = functions
|
||||
|
||||
@property
|
||||
def vector(self) -> VectorClient:
|
||||
"""Get vector client (Qdrant)."""
|
||||
if self._vector is None:
|
||||
raise RuntimeError("Vector client not initialized. Initialize with services=['vector']")
|
||||
return self._vector
|
||||
|
||||
@property
|
||||
def kv(self) -> KVClient:
|
||||
"""Get key-value client (Redis/Valkey)."""
|
||||
if self._kv is None:
|
||||
raise RuntimeError("KV client not initialized. Initialize with services=['kv']")
|
||||
return self._kv
|
||||
|
||||
@property
|
||||
def documentdb(self) -> DocumentDBClient:
|
||||
"""Get document database client (MongoDB)."""
|
||||
if self._documentdb is None:
|
||||
raise RuntimeError("DocumentDB client not initialized. Initialize with services=['documentdb']")
|
||||
return self._documentdb
|
||||
|
||||
@property
|
||||
def storage(self) -> StorageClient:
|
||||
"""Get object storage client (S3/MinIO)."""
|
||||
if self._storage is None:
|
||||
raise RuntimeError("Storage client not initialized. Initialize with services=['storage']")
|
||||
return self._storage
|
||||
|
||||
@property
|
||||
def search(self) -> SearchClient:
|
||||
"""Get search client (Meilisearch)."""
|
||||
if self._search is None:
|
||||
raise RuntimeError("Search client not initialized. Initialize with services=['search']")
|
||||
return self._search
|
||||
|
||||
@property
|
||||
def pubsub(self) -> PubSubClient:
|
||||
"""Get pub/sub client (NATS)."""
|
||||
if self._pubsub is None:
|
||||
raise RuntimeError("PubSub client not initialized. Initialize with services=['pubsub']")
|
||||
return self._pubsub
|
||||
|
||||
@property
|
||||
def tasks(self) -> TasksClient:
|
||||
"""Get tasks/workflow client (Temporal)."""
|
||||
if self._tasks is None:
|
||||
raise RuntimeError("Tasks client not initialized. Initialize with services=['tasks']")
|
||||
return self._tasks
|
||||
|
||||
@property
|
||||
def queues(self) -> QueuesClient:
|
||||
"""Get work queue client."""
|
||||
if self._queues is None:
|
||||
raise RuntimeError("Queues client not initialized. Initialize with services=['queues']")
|
||||
return self._queues
|
||||
|
||||
@property
|
||||
def cron(self) -> CronClient:
|
||||
"""Get cron scheduler client."""
|
||||
if self._cron is None:
|
||||
raise RuntimeError("Cron client not initialized. Initialize with services=['cron']")
|
||||
return self._cron
|
||||
|
||||
@property
|
||||
def functions(self) -> FunctionsClient:
|
||||
"""Get serverless functions client (Nuclio)."""
|
||||
if self._functions is None:
|
||||
raise RuntimeError("Functions client not initialized. Initialize with services=['functions']")
|
||||
return self._functions
|
||||
|
||||
@classmethod
|
||||
async def from_env(
|
||||
cls,
|
||||
services: Optional[Sequence[str]] = None,
|
||||
connect: bool = True,
|
||||
) -> HanzoInfra:
|
||||
"""Create HanzoInfra from environment variables.
|
||||
|
||||
Args:
|
||||
services: List of services to initialize. If None, initializes all.
|
||||
Valid values: vector, kv, documentdb, storage, search, pubsub,
|
||||
tasks, queues, cron, functions.
|
||||
connect: Automatically connect to services.
|
||||
|
||||
Returns:
|
||||
Configured HanzoInfra instance.
|
||||
|
||||
Example:
|
||||
```python
|
||||
# Initialize all services
|
||||
infra = await HanzoInfra.from_env()
|
||||
|
||||
# Initialize only specific services
|
||||
infra = await HanzoInfra.from_env(services=["kv", "vector"])
|
||||
```
|
||||
"""
|
||||
services_to_init = set(services) if services else set(ALL_SERVICES)
|
||||
|
||||
# Validate service names
|
||||
invalid = services_to_init - set(ALL_SERVICES)
|
||||
if invalid:
|
||||
raise ValueError(f"Invalid services: {invalid}. Valid: {ALL_SERVICES}")
|
||||
|
||||
clients: dict[str, Any] = {}
|
||||
|
||||
# Initialize requested services
|
||||
if "vector" in services_to_init:
|
||||
clients["vector"] = VectorClient(VectorConfig.from_env())
|
||||
|
||||
if "kv" in services_to_init:
|
||||
clients["kv"] = KVClient(KVConfig.from_env())
|
||||
|
||||
if "documentdb" in services_to_init:
|
||||
clients["documentdb"] = DocumentDBClient(DocumentDBConfig.from_env())
|
||||
|
||||
if "storage" in services_to_init:
|
||||
clients["storage"] = StorageClient(StorageConfig.from_env())
|
||||
|
||||
if "search" in services_to_init:
|
||||
clients["search"] = SearchClient(SearchConfig.from_env())
|
||||
|
||||
if "pubsub" in services_to_init:
|
||||
clients["pubsub"] = PubSubClient(PubSubConfig.from_env())
|
||||
|
||||
if "tasks" in services_to_init:
|
||||
clients["tasks"] = TasksClient(TasksConfig.from_env())
|
||||
|
||||
if "queues" in services_to_init:
|
||||
clients["queues"] = QueuesClient(QueuesConfig.from_env())
|
||||
|
||||
if "cron" in services_to_init:
|
||||
clients["cron"] = CronClient(CronConfig.from_env())
|
||||
|
||||
if "functions" in services_to_init:
|
||||
clients["functions"] = FunctionsClient(FunctionsConfig.from_env())
|
||||
|
||||
infra = cls(**clients)
|
||||
|
||||
if connect:
|
||||
await infra.connect()
|
||||
|
||||
return infra
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Connect all initialized services."""
|
||||
if self._vector:
|
||||
await self._vector.connect()
|
||||
if self._kv:
|
||||
await self._kv.connect()
|
||||
if self._documentdb:
|
||||
await self._documentdb.connect()
|
||||
if self._storage:
|
||||
await self._storage.connect()
|
||||
if self._search:
|
||||
await self._search.connect()
|
||||
if self._pubsub:
|
||||
await self._pubsub.connect()
|
||||
if self._tasks:
|
||||
await self._tasks.connect()
|
||||
if self._queues:
|
||||
await self._queues.connect()
|
||||
if self._cron:
|
||||
await self._cron.connect()
|
||||
if self._functions:
|
||||
await self._functions.connect()
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close all service connections."""
|
||||
if self._vector:
|
||||
await self._vector.close()
|
||||
if self._kv:
|
||||
await self._kv.close()
|
||||
if self._documentdb:
|
||||
await self._documentdb.close()
|
||||
if self._storage:
|
||||
await self._storage.close()
|
||||
if self._search:
|
||||
await self._search.close()
|
||||
if self._pubsub:
|
||||
await self._pubsub.close()
|
||||
if self._tasks:
|
||||
await self._tasks.close()
|
||||
if self._queues:
|
||||
await self._queues.close()
|
||||
if self._cron:
|
||||
await self._cron.close()
|
||||
if self._functions:
|
||||
await self._functions.close()
|
||||
|
||||
async def health_check(self) -> dict[str, bool]:
|
||||
"""Check health of all initialized services.
|
||||
|
||||
Returns:
|
||||
Dict mapping service name to health status.
|
||||
"""
|
||||
results = {}
|
||||
|
||||
if self._vector:
|
||||
results["vector"] = await self._vector.health_check()
|
||||
if self._kv:
|
||||
results["kv"] = await self._kv.health_check()
|
||||
if self._documentdb:
|
||||
results["documentdb"] = await self._documentdb.health_check()
|
||||
if self._storage:
|
||||
results["storage"] = await self._storage.health_check()
|
||||
if self._search:
|
||||
results["search"] = await self._search.health_check()
|
||||
if self._pubsub:
|
||||
results["pubsub"] = await self._pubsub.health_check()
|
||||
if self._tasks:
|
||||
results["tasks"] = await self._tasks.health_check()
|
||||
if self._queues:
|
||||
results["queues"] = await self._queues.health_check()
|
||||
if self._cron:
|
||||
results["cron"] = await self._cron.health_check()
|
||||
if self._functions:
|
||||
results["functions"] = await self._functions.health_check()
|
||||
|
||||
return results
|
||||
|
||||
async def __aenter__(self) -> HanzoInfra:
|
||||
"""Async context manager entry."""
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,631 @@
|
||||
"""Scheduled jobs (cron) client for Hanzo infrastructure.
|
||||
|
||||
Provides async interface for scheduled/recurring jobs,
|
||||
built on top of Redis for distributed scheduling.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, Awaitable, Callable, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class CronConfig(BaseModel):
|
||||
"""Configuration for cron scheduler connection."""
|
||||
|
||||
host: str = Field(default="localhost", description="Redis server host")
|
||||
port: int = Field(default=6379, description="Redis server port")
|
||||
password: Optional[str] = Field(default=None, description="Redis password")
|
||||
db: int = Field(default=2, description="Redis database number")
|
||||
url: Optional[str] = Field(default=None, description="Full Redis URL")
|
||||
prefix: str = Field(default="hanzo:cron", description="Key prefix for cron")
|
||||
lock_timeout: int = Field(default=60, description="Lock timeout in seconds")
|
||||
timezone: str = Field(default="UTC", description="Default timezone")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> CronConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
CRON_REDIS_HOST: Server host (default: localhost)
|
||||
CRON_REDIS_PORT: Server port (default: 6379)
|
||||
CRON_REDIS_PASSWORD: Authentication password
|
||||
CRON_REDIS_DB: Database number (default: 2)
|
||||
CRON_REDIS_URL: Full URL (overrides host/port)
|
||||
CRON_PREFIX: Key prefix (default: hanzo:cron)
|
||||
CRON_TIMEZONE: Default timezone (default: UTC)
|
||||
"""
|
||||
return cls(
|
||||
host=os.getenv("CRON_REDIS_HOST") or os.getenv("REDIS_HOST", "localhost"),
|
||||
port=int(os.getenv("CRON_REDIS_PORT") or os.getenv("REDIS_PORT", "6379")),
|
||||
password=os.getenv("CRON_REDIS_PASSWORD") or os.getenv("REDIS_PASSWORD"),
|
||||
db=int(os.getenv("CRON_REDIS_DB", "2")),
|
||||
url=os.getenv("CRON_REDIS_URL") or os.getenv("REDIS_URL"),
|
||||
prefix=os.getenv("CRON_PREFIX", "hanzo:cron"),
|
||||
timezone=os.getenv("CRON_TIMEZONE", "UTC"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CronJob:
|
||||
"""A scheduled job definition."""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
schedule: str # Cron expression or interval
|
||||
handler: str # Handler name/identifier
|
||||
args: dict[str, Any] = field(default_factory=dict)
|
||||
enabled: bool = True
|
||||
timezone: str = "UTC"
|
||||
last_run: Optional[datetime] = None
|
||||
next_run: Optional[datetime] = None
|
||||
run_count: int = 0
|
||||
error_count: int = 0
|
||||
last_error: Optional[str] = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dictionary for storage."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"name": self.name,
|
||||
"schedule": self.schedule,
|
||||
"handler": self.handler,
|
||||
"args": self.args,
|
||||
"enabled": self.enabled,
|
||||
"timezone": self.timezone,
|
||||
"last_run": self.last_run.isoformat() if self.last_run else None,
|
||||
"next_run": self.next_run.isoformat() if self.next_run else None,
|
||||
"run_count": self.run_count,
|
||||
"error_count": self.error_count,
|
||||
"last_error": self.last_error,
|
||||
"metadata": self.metadata,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> CronJob:
|
||||
"""Create from dictionary."""
|
||||
return cls(
|
||||
id=data["id"],
|
||||
name=data["name"],
|
||||
schedule=data["schedule"],
|
||||
handler=data["handler"],
|
||||
args=data.get("args", {}),
|
||||
enabled=data.get("enabled", True),
|
||||
timezone=data.get("timezone", "UTC"),
|
||||
last_run=datetime.fromisoformat(data["last_run"]) if data.get("last_run") else None,
|
||||
next_run=datetime.fromisoformat(data["next_run"]) if data.get("next_run") else None,
|
||||
run_count=data.get("run_count", 0),
|
||||
error_count=data.get("error_count", 0),
|
||||
last_error=data.get("last_error"),
|
||||
metadata=data.get("metadata", {}),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CronExecution:
|
||||
"""Record of a job execution."""
|
||||
|
||||
job_id: str
|
||||
started_at: datetime
|
||||
completed_at: Optional[datetime] = None
|
||||
success: bool = False
|
||||
result: Any = None
|
||||
error: Optional[str] = None
|
||||
duration_ms: int = 0
|
||||
|
||||
|
||||
CronHandler = Callable[[CronJob], Awaitable[Any]]
|
||||
|
||||
|
||||
class CronClient:
|
||||
"""Async client for scheduled job management.
|
||||
|
||||
Implements distributed cron scheduling on top of Redis,
|
||||
with support for cron expressions, intervals, and distributed locking.
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = CronClient(CronConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Register a job
|
||||
await client.register(
|
||||
"cleanup",
|
||||
"daily-cleanup",
|
||||
schedule="0 2 * * *", # 2 AM daily
|
||||
handler="cleanup_old_files",
|
||||
args={"days": 30},
|
||||
)
|
||||
|
||||
# Or with interval
|
||||
await client.register(
|
||||
"heartbeat",
|
||||
"service-heartbeat",
|
||||
schedule="@every 5m", # Every 5 minutes
|
||||
handler="send_heartbeat",
|
||||
)
|
||||
|
||||
# Run the scheduler (in a worker)
|
||||
handlers = {
|
||||
"cleanup_old_files": cleanup_handler,
|
||||
"send_heartbeat": heartbeat_handler,
|
||||
}
|
||||
await client.run_scheduler(handlers)
|
||||
|
||||
# Manually trigger a job
|
||||
await client.trigger("cleanup")
|
||||
|
||||
# List jobs
|
||||
jobs = await client.list_jobs()
|
||||
for job in jobs:
|
||||
print(f"{job.name}: next run at {job.next_run}")
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[CronConfig] = None) -> None:
|
||||
"""Initialize cron client.
|
||||
|
||||
Args:
|
||||
config: Cron configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or CronConfig.from_env()
|
||||
self._redis: Any = None
|
||||
self._running = False
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection to Redis."""
|
||||
try:
|
||||
import redis.asyncio as redis
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"redis is required for CronClient. "
|
||||
"Install with: pip install redis"
|
||||
) from e
|
||||
|
||||
if self.config.url:
|
||||
self._redis = redis.from_url(
|
||||
self.config.url,
|
||||
decode_responses=True,
|
||||
)
|
||||
else:
|
||||
self._redis = redis.Redis(
|
||||
host=self.config.host,
|
||||
port=self.config.port,
|
||||
password=self.config.password,
|
||||
db=self.config.db,
|
||||
decode_responses=True,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
self._running = False
|
||||
if self._redis:
|
||||
await self._redis.aclose()
|
||||
self._redis = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if Redis is healthy.
|
||||
|
||||
Returns:
|
||||
True if Redis is responding.
|
||||
"""
|
||||
if not self._redis:
|
||||
return False
|
||||
try:
|
||||
await self._redis.ping()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _key(self, *parts: str) -> str:
|
||||
"""Build a Redis key with prefix."""
|
||||
return ":".join([self.config.prefix, *parts])
|
||||
|
||||
def _parse_schedule(self, schedule: str) -> datetime:
|
||||
"""Parse a schedule and return the next run time.
|
||||
|
||||
Supports:
|
||||
- Cron expressions: "* * * * *" (min hour day month weekday)
|
||||
- Intervals: "@every 5m", "@every 1h", "@every 30s"
|
||||
- Named schedules: "@hourly", "@daily", "@weekly", "@monthly"
|
||||
"""
|
||||
now = datetime.utcnow()
|
||||
|
||||
# Handle interval syntax
|
||||
if schedule.startswith("@every "):
|
||||
interval_str = schedule[7:]
|
||||
return now + self._parse_interval(interval_str)
|
||||
|
||||
# Handle named schedules
|
||||
named = {
|
||||
"@hourly": "0 * * * *",
|
||||
"@daily": "0 0 * * *",
|
||||
"@weekly": "0 0 * * 0",
|
||||
"@monthly": "0 0 1 * *",
|
||||
"@yearly": "0 0 1 1 *",
|
||||
"@annually": "0 0 1 1 *",
|
||||
}
|
||||
if schedule in named:
|
||||
schedule = named[schedule]
|
||||
|
||||
# Parse cron expression
|
||||
try:
|
||||
from croniter import croniter
|
||||
except ImportError:
|
||||
# Fallback: simple interval-based scheduling
|
||||
# Default to hourly if croniter not available
|
||||
return now + timedelta(hours=1)
|
||||
|
||||
cron = croniter(schedule, now)
|
||||
return cron.get_next(datetime)
|
||||
|
||||
def _parse_interval(self, interval: str) -> timedelta:
|
||||
"""Parse an interval string like '5m', '1h', '30s'."""
|
||||
units = {
|
||||
"s": 1,
|
||||
"m": 60,
|
||||
"h": 3600,
|
||||
"d": 86400,
|
||||
"w": 604800,
|
||||
}
|
||||
|
||||
if not interval:
|
||||
return timedelta(hours=1)
|
||||
|
||||
unit = interval[-1].lower()
|
||||
if unit not in units:
|
||||
raise ValueError(f"Unknown interval unit: {unit}")
|
||||
|
||||
try:
|
||||
value = int(interval[:-1])
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid interval value: {interval[:-1]}")
|
||||
|
||||
return timedelta(seconds=value * units[unit])
|
||||
|
||||
# Job management
|
||||
|
||||
async def register(
|
||||
self,
|
||||
job_id: str,
|
||||
name: str,
|
||||
schedule: str,
|
||||
handler: str,
|
||||
args: Optional[dict[str, Any]] = None,
|
||||
timezone: Optional[str] = None,
|
||||
enabled: bool = True,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
) -> CronJob:
|
||||
"""Register a scheduled job.
|
||||
|
||||
Args:
|
||||
job_id: Unique job identifier.
|
||||
name: Human-readable job name.
|
||||
schedule: Cron expression or interval.
|
||||
handler: Handler name/identifier.
|
||||
args: Arguments to pass to handler.
|
||||
timezone: Job timezone.
|
||||
enabled: Whether job is enabled.
|
||||
metadata: Additional metadata.
|
||||
|
||||
Returns:
|
||||
Created CronJob.
|
||||
"""
|
||||
next_run = self._parse_schedule(schedule) if enabled else None
|
||||
|
||||
job = CronJob(
|
||||
id=job_id,
|
||||
name=name,
|
||||
schedule=schedule,
|
||||
handler=handler,
|
||||
args=args or {},
|
||||
enabled=enabled,
|
||||
timezone=timezone or self.config.timezone,
|
||||
next_run=next_run,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
|
||||
# Store job
|
||||
job_key = self._key("job", job_id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
|
||||
# Add to jobs set
|
||||
jobs_key = self._key("jobs")
|
||||
await self._redis.sadd(jobs_key, job_id)
|
||||
|
||||
# Add to schedule
|
||||
if enabled and next_run:
|
||||
schedule_key = self._key("schedule")
|
||||
await self._redis.zadd(schedule_key, {job_id: next_run.timestamp()})
|
||||
|
||||
return job
|
||||
|
||||
async def unregister(self, job_id: str) -> bool:
|
||||
"""Unregister a job.
|
||||
|
||||
Args:
|
||||
job_id: Job ID to remove.
|
||||
|
||||
Returns:
|
||||
True if job was removed.
|
||||
"""
|
||||
job_key = self._key("job", job_id)
|
||||
jobs_key = self._key("jobs")
|
||||
schedule_key = self._key("schedule")
|
||||
|
||||
existed = await self._redis.delete(job_key)
|
||||
await self._redis.srem(jobs_key, job_id)
|
||||
await self._redis.zrem(schedule_key, job_id)
|
||||
|
||||
return existed > 0
|
||||
|
||||
async def get_job(self, job_id: str) -> Optional[CronJob]:
|
||||
"""Get a job by ID.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
|
||||
Returns:
|
||||
CronJob or None.
|
||||
"""
|
||||
job_key = self._key("job", job_id)
|
||||
data = await self._redis.get(job_key)
|
||||
if not data:
|
||||
return None
|
||||
return CronJob.from_dict(json.loads(data))
|
||||
|
||||
async def list_jobs(self) -> list[CronJob]:
|
||||
"""List all registered jobs.
|
||||
|
||||
Returns:
|
||||
List of CronJob objects.
|
||||
"""
|
||||
jobs_key = self._key("jobs")
|
||||
job_ids = await self._redis.smembers(jobs_key)
|
||||
|
||||
jobs = []
|
||||
for job_id in job_ids:
|
||||
job = await self.get_job(job_id)
|
||||
if job:
|
||||
jobs.append(job)
|
||||
|
||||
return sorted(jobs, key=lambda j: j.name)
|
||||
|
||||
async def enable_job(self, job_id: str) -> bool:
|
||||
"""Enable a job.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
|
||||
Returns:
|
||||
True if job was enabled.
|
||||
"""
|
||||
job = await self.get_job(job_id)
|
||||
if not job:
|
||||
return False
|
||||
|
||||
job.enabled = True
|
||||
job.next_run = self._parse_schedule(job.schedule)
|
||||
|
||||
# Update job
|
||||
job_key = self._key("job", job_id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
|
||||
# Add to schedule
|
||||
schedule_key = self._key("schedule")
|
||||
await self._redis.zadd(schedule_key, {job_id: job.next_run.timestamp()})
|
||||
|
||||
return True
|
||||
|
||||
async def disable_job(self, job_id: str) -> bool:
|
||||
"""Disable a job.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
|
||||
Returns:
|
||||
True if job was disabled.
|
||||
"""
|
||||
job = await self.get_job(job_id)
|
||||
if not job:
|
||||
return False
|
||||
|
||||
job.enabled = False
|
||||
job.next_run = None
|
||||
|
||||
# Update job
|
||||
job_key = self._key("job", job_id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
|
||||
# Remove from schedule
|
||||
schedule_key = self._key("schedule")
|
||||
await self._redis.zrem(schedule_key, job_id)
|
||||
|
||||
return True
|
||||
|
||||
async def trigger(self, job_id: str) -> Optional[CronExecution]:
|
||||
"""Manually trigger a job execution.
|
||||
|
||||
Args:
|
||||
job_id: Job ID to trigger.
|
||||
|
||||
Returns:
|
||||
Execution result or None if job not found.
|
||||
"""
|
||||
job = await self.get_job(job_id)
|
||||
if not job:
|
||||
return None
|
||||
|
||||
# Schedule for immediate execution
|
||||
schedule_key = self._key("schedule")
|
||||
await self._redis.zadd(schedule_key, {job_id: datetime.utcnow().timestamp()})
|
||||
|
||||
return CronExecution(
|
||||
job_id=job_id,
|
||||
started_at=datetime.utcnow(),
|
||||
)
|
||||
|
||||
# Scheduler
|
||||
|
||||
async def _acquire_lock(self, job_id: str) -> bool:
|
||||
"""Acquire a distributed lock for a job."""
|
||||
lock_key = self._key("lock", job_id)
|
||||
lock_value = hashlib.md5(f"{os.getpid()}-{datetime.utcnow().isoformat()}".encode()).hexdigest()
|
||||
|
||||
acquired = await self._redis.set(
|
||||
lock_key,
|
||||
lock_value,
|
||||
ex=self.config.lock_timeout,
|
||||
nx=True,
|
||||
)
|
||||
return acquired is not None
|
||||
|
||||
async def _release_lock(self, job_id: str) -> None:
|
||||
"""Release a distributed lock."""
|
||||
lock_key = self._key("lock", job_id)
|
||||
await self._redis.delete(lock_key)
|
||||
|
||||
async def _execute_job(
|
||||
self,
|
||||
job: CronJob,
|
||||
handlers: dict[str, CronHandler],
|
||||
) -> CronExecution:
|
||||
"""Execute a job."""
|
||||
execution = CronExecution(
|
||||
job_id=job.id,
|
||||
started_at=datetime.utcnow(),
|
||||
)
|
||||
|
||||
handler = handlers.get(job.handler)
|
||||
if not handler:
|
||||
execution.error = f"Handler not found: {job.handler}"
|
||||
return execution
|
||||
|
||||
try:
|
||||
result = await handler(job)
|
||||
execution.success = True
|
||||
execution.result = result
|
||||
except Exception as e:
|
||||
execution.success = False
|
||||
execution.error = str(e)
|
||||
|
||||
execution.completed_at = datetime.utcnow()
|
||||
execution.duration_ms = int(
|
||||
(execution.completed_at - execution.started_at).total_seconds() * 1000
|
||||
)
|
||||
|
||||
# Update job stats
|
||||
job.last_run = execution.started_at
|
||||
job.run_count += 1
|
||||
if not execution.success:
|
||||
job.error_count += 1
|
||||
job.last_error = execution.error
|
||||
|
||||
# Schedule next run
|
||||
job.next_run = self._parse_schedule(job.schedule)
|
||||
|
||||
# Save job
|
||||
job_key = self._key("job", job.id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
|
||||
# Update schedule
|
||||
schedule_key = self._key("schedule")
|
||||
await self._redis.zadd(schedule_key, {job.id: job.next_run.timestamp()})
|
||||
|
||||
return execution
|
||||
|
||||
async def run_scheduler(
|
||||
self,
|
||||
handlers: dict[str, CronHandler],
|
||||
poll_interval: float = 1.0,
|
||||
) -> None:
|
||||
"""Run the cron scheduler (blocking).
|
||||
|
||||
Args:
|
||||
handlers: Map of handler names to functions.
|
||||
poll_interval: Seconds between schedule checks.
|
||||
"""
|
||||
self._running = True
|
||||
schedule_key = self._key("schedule")
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
now = datetime.utcnow().timestamp()
|
||||
|
||||
# Get jobs due to run
|
||||
due_jobs = await self._redis.zrangebyscore(schedule_key, "-inf", now)
|
||||
|
||||
for job_id in due_jobs:
|
||||
# Try to acquire lock
|
||||
if not await self._acquire_lock(job_id):
|
||||
continue # Another worker has this job
|
||||
|
||||
try:
|
||||
job = await self.get_job(job_id)
|
||||
if job and job.enabled:
|
||||
await self._execute_job(job, handlers)
|
||||
finally:
|
||||
await self._release_lock(job_id)
|
||||
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception:
|
||||
# Log error but keep running
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
def stop_scheduler(self) -> None:
|
||||
"""Stop the scheduler."""
|
||||
self._running = False
|
||||
|
||||
# History (optional - for debugging)
|
||||
|
||||
async def get_recent_executions(
|
||||
self,
|
||||
job_id: str,
|
||||
limit: int = 10,
|
||||
) -> list[CronExecution]:
|
||||
"""Get recent executions for a job.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
limit: Maximum executions to return.
|
||||
|
||||
Returns:
|
||||
List of recent executions (most recent first).
|
||||
"""
|
||||
history_key = self._key("history", job_id)
|
||||
entries = await self._redis.lrange(history_key, 0, limit - 1)
|
||||
|
||||
executions = []
|
||||
for entry in entries:
|
||||
data = json.loads(entry)
|
||||
executions.append(CronExecution(
|
||||
job_id=data["job_id"],
|
||||
started_at=datetime.fromisoformat(data["started_at"]),
|
||||
completed_at=datetime.fromisoformat(data["completed_at"]) if data.get("completed_at") else None,
|
||||
success=data.get("success", False),
|
||||
result=data.get("result"),
|
||||
error=data.get("error"),
|
||||
duration_ms=data.get("duration_ms", 0),
|
||||
))
|
||||
|
||||
return executions
|
||||
|
||||
async def __aenter__(self) -> CronClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,565 @@
|
||||
"""MongoDB document database client wrapper for Hanzo infrastructure.
|
||||
|
||||
Provides async interface to MongoDB for document storage and querying,
|
||||
supporting both MongoDB Atlas and self-hosted deployments.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Optional, Sequence, TypeVar
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
T = TypeVar("T", bound=dict[str, Any])
|
||||
|
||||
|
||||
class DocumentDBConfig(BaseModel):
|
||||
"""Configuration for MongoDB connection."""
|
||||
|
||||
host: str = Field(default="localhost", description="MongoDB server host")
|
||||
port: int = Field(default=27017, description="MongoDB server port")
|
||||
username: Optional[str] = Field(default=None, description="Authentication username")
|
||||
password: Optional[str] = Field(default=None, description="Authentication password")
|
||||
database: str = Field(default="hanzo", description="Default database name")
|
||||
uri: Optional[str] = Field(default=None, description="Full connection URI (overrides host/port)")
|
||||
auth_source: str = Field(default="admin", description="Authentication database")
|
||||
replica_set: Optional[str] = Field(default=None, description="Replica set name")
|
||||
tls: bool = Field(default=False, description="Use TLS/SSL")
|
||||
server_selection_timeout_ms: int = Field(default=5000, description="Server selection timeout")
|
||||
connect_timeout_ms: int = Field(default=5000, description="Connection timeout")
|
||||
socket_timeout_ms: int = Field(default=30000, description="Socket timeout")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> DocumentDBConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
MONGODB_HOST: Server host (default: localhost)
|
||||
MONGODB_PORT: Server port (default: 27017)
|
||||
MONGODB_USERNAME: Authentication username
|
||||
MONGODB_PASSWORD: Authentication password
|
||||
MONGODB_DATABASE: Database name (default: hanzo)
|
||||
MONGODB_URI: Full connection URI (overrides host/port)
|
||||
MONGODB_TLS: Use TLS (default: false)
|
||||
"""
|
||||
return cls(
|
||||
host=os.getenv("MONGODB_HOST", "localhost"),
|
||||
port=int(os.getenv("MONGODB_PORT", "27017")),
|
||||
username=os.getenv("MONGODB_USERNAME"),
|
||||
password=os.getenv("MONGODB_PASSWORD"),
|
||||
database=os.getenv("MONGODB_DATABASE", "hanzo"),
|
||||
uri=os.getenv("MONGODB_URI"),
|
||||
tls=os.getenv("MONGODB_TLS", "").lower() in ("true", "1", "yes"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""A MongoDB document with metadata."""
|
||||
|
||||
id: str
|
||||
data: dict[str, Any]
|
||||
collection: str
|
||||
created_at: Optional[datetime] = None
|
||||
updated_at: Optional[datetime] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueryResult:
|
||||
"""Result of a query operation."""
|
||||
|
||||
documents: list[dict[str, Any]]
|
||||
count: int
|
||||
has_more: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class UpdateResult:
|
||||
"""Result of an update operation."""
|
||||
|
||||
matched_count: int
|
||||
modified_count: int
|
||||
upserted_id: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeleteResult:
|
||||
"""Result of a delete operation."""
|
||||
|
||||
deleted_count: int
|
||||
|
||||
|
||||
class DocumentDBClient:
|
||||
"""Async client for MongoDB document database.
|
||||
|
||||
Wraps motor (async MongoDB driver) with Hanzo conventions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = DocumentDBClient(DocumentDBConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Insert documents
|
||||
doc_id = await client.insert_one("users", {"name": "Alice", "email": "[email protected]"})
|
||||
|
||||
# Find documents
|
||||
users = await client.find("users", {"name": "Alice"})
|
||||
|
||||
# Update documents
|
||||
result = await client.update_one("users", {"_id": doc_id}, {"$set": {"status": "active"}})
|
||||
|
||||
# Aggregate
|
||||
pipeline = [{"$group": {"_id": "$status", "count": {"$sum": 1}}}]
|
||||
stats = await client.aggregate("users", pipeline)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[DocumentDBConfig] = None) -> None:
|
||||
"""Initialize document database client.
|
||||
|
||||
Args:
|
||||
config: MongoDB configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or DocumentDBConfig.from_env()
|
||||
self._client: Any = None
|
||||
self._db: Any = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection to MongoDB server."""
|
||||
try:
|
||||
from motor.motor_asyncio import AsyncIOMotorClient
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"motor is required for DocumentDBClient. "
|
||||
"Install with: pip install motor"
|
||||
) from e
|
||||
|
||||
if self.config.uri:
|
||||
self._client = AsyncIOMotorClient(
|
||||
self.config.uri,
|
||||
serverSelectionTimeoutMS=self.config.server_selection_timeout_ms,
|
||||
connectTimeoutMS=self.config.connect_timeout_ms,
|
||||
socketTimeoutMS=self.config.socket_timeout_ms,
|
||||
)
|
||||
else:
|
||||
self._client = AsyncIOMotorClient(
|
||||
host=self.config.host,
|
||||
port=self.config.port,
|
||||
username=self.config.username,
|
||||
password=self.config.password,
|
||||
authSource=self.config.auth_source,
|
||||
replicaSet=self.config.replica_set,
|
||||
tls=self.config.tls,
|
||||
serverSelectionTimeoutMS=self.config.server_selection_timeout_ms,
|
||||
connectTimeoutMS=self.config.connect_timeout_ms,
|
||||
socketTimeoutMS=self.config.socket_timeout_ms,
|
||||
)
|
||||
|
||||
self._db = self._client[self.config.database]
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
if self._client:
|
||||
self._client.close()
|
||||
self._client = None
|
||||
self._db = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if MongoDB server is healthy.
|
||||
|
||||
Returns:
|
||||
True if server is reachable and responding.
|
||||
"""
|
||||
if not self._client:
|
||||
return False
|
||||
try:
|
||||
await self._client.admin.command("ping")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _get_collection(self, collection: str) -> Any:
|
||||
"""Get a collection object."""
|
||||
return self._db[collection]
|
||||
|
||||
# Insert operations
|
||||
|
||||
async def insert_one(self, collection: str, document: dict[str, Any]) -> str:
|
||||
"""Insert a single document.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
document: Document to insert.
|
||||
|
||||
Returns:
|
||||
Inserted document ID as string.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
result = await coll.insert_one(document)
|
||||
return str(result.inserted_id)
|
||||
|
||||
async def insert_many(self, collection: str, documents: Sequence[dict[str, Any]]) -> list[str]:
|
||||
"""Insert multiple documents.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
documents: Documents to insert.
|
||||
|
||||
Returns:
|
||||
List of inserted document IDs.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
result = await coll.insert_many(list(documents))
|
||||
return [str(id_) for id_ in result.inserted_ids]
|
||||
|
||||
# Find operations
|
||||
|
||||
async def find_one(
|
||||
self,
|
||||
collection: str,
|
||||
filter: dict[str, Any],
|
||||
projection: Optional[dict[str, Any]] = None,
|
||||
) -> Optional[dict[str, Any]]:
|
||||
"""Find a single document.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
filter: Query filter.
|
||||
projection: Fields to include/exclude.
|
||||
|
||||
Returns:
|
||||
Document or None if not found.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
return await coll.find_one(filter, projection)
|
||||
|
||||
async def find(
|
||||
self,
|
||||
collection: str,
|
||||
filter: dict[str, Any],
|
||||
projection: Optional[dict[str, Any]] = None,
|
||||
sort: Optional[list[tuple[str, int]]] = None,
|
||||
skip: int = 0,
|
||||
limit: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Find documents matching a filter.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
filter: Query filter.
|
||||
projection: Fields to include/exclude.
|
||||
sort: Sort specification [(field, direction), ...].
|
||||
skip: Number of documents to skip.
|
||||
limit: Maximum documents to return (0 = no limit).
|
||||
|
||||
Returns:
|
||||
List of matching documents.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
cursor = coll.find(filter, projection)
|
||||
|
||||
if sort:
|
||||
cursor = cursor.sort(sort)
|
||||
if skip:
|
||||
cursor = cursor.skip(skip)
|
||||
if limit:
|
||||
cursor = cursor.limit(limit)
|
||||
|
||||
return await cursor.to_list(length=limit if limit else None)
|
||||
|
||||
async def find_by_id(
|
||||
self,
|
||||
collection: str,
|
||||
id: str,
|
||||
projection: Optional[dict[str, Any]] = None,
|
||||
) -> Optional[dict[str, Any]]:
|
||||
"""Find a document by ID.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
id: Document ID.
|
||||
projection: Fields to include/exclude.
|
||||
|
||||
Returns:
|
||||
Document or None if not found.
|
||||
"""
|
||||
from bson import ObjectId
|
||||
|
||||
return await self.find_one(collection, {"_id": ObjectId(id)}, projection)
|
||||
|
||||
# Update operations
|
||||
|
||||
async def update_one(
|
||||
self,
|
||||
collection: str,
|
||||
filter: dict[str, Any],
|
||||
update: dict[str, Any],
|
||||
upsert: bool = False,
|
||||
) -> UpdateResult:
|
||||
"""Update a single document.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
filter: Query filter.
|
||||
update: Update operations.
|
||||
upsert: Insert if not found.
|
||||
|
||||
Returns:
|
||||
Update result with counts.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
result = await coll.update_one(filter, update, upsert=upsert)
|
||||
return UpdateResult(
|
||||
matched_count=result.matched_count,
|
||||
modified_count=result.modified_count,
|
||||
upserted_id=str(result.upserted_id) if result.upserted_id else None,
|
||||
)
|
||||
|
||||
async def update_many(
|
||||
self,
|
||||
collection: str,
|
||||
filter: dict[str, Any],
|
||||
update: dict[str, Any],
|
||||
upsert: bool = False,
|
||||
) -> UpdateResult:
|
||||
"""Update multiple documents.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
filter: Query filter.
|
||||
update: Update operations.
|
||||
upsert: Insert if not found.
|
||||
|
||||
Returns:
|
||||
Update result with counts.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
result = await coll.update_many(filter, update, upsert=upsert)
|
||||
return UpdateResult(
|
||||
matched_count=result.matched_count,
|
||||
modified_count=result.modified_count,
|
||||
upserted_id=str(result.upserted_id) if result.upserted_id else None,
|
||||
)
|
||||
|
||||
async def replace_one(
|
||||
self,
|
||||
collection: str,
|
||||
filter: dict[str, Any],
|
||||
replacement: dict[str, Any],
|
||||
upsert: bool = False,
|
||||
) -> UpdateResult:
|
||||
"""Replace a single document.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
filter: Query filter.
|
||||
replacement: New document.
|
||||
upsert: Insert if not found.
|
||||
|
||||
Returns:
|
||||
Update result with counts.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
result = await coll.replace_one(filter, replacement, upsert=upsert)
|
||||
return UpdateResult(
|
||||
matched_count=result.matched_count,
|
||||
modified_count=result.modified_count,
|
||||
upserted_id=str(result.upserted_id) if result.upserted_id else None,
|
||||
)
|
||||
|
||||
# Delete operations
|
||||
|
||||
async def delete_one(self, collection: str, filter: dict[str, Any]) -> DeleteResult:
|
||||
"""Delete a single document.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
filter: Query filter.
|
||||
|
||||
Returns:
|
||||
Delete result with count.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
result = await coll.delete_one(filter)
|
||||
return DeleteResult(deleted_count=result.deleted_count)
|
||||
|
||||
async def delete_many(self, collection: str, filter: dict[str, Any]) -> DeleteResult:
|
||||
"""Delete multiple documents.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
filter: Query filter.
|
||||
|
||||
Returns:
|
||||
Delete result with count.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
result = await coll.delete_many(filter)
|
||||
return DeleteResult(deleted_count=result.deleted_count)
|
||||
|
||||
async def delete_by_id(self, collection: str, id: str) -> DeleteResult:
|
||||
"""Delete a document by ID.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
id: Document ID.
|
||||
|
||||
Returns:
|
||||
Delete result with count.
|
||||
"""
|
||||
from bson import ObjectId
|
||||
|
||||
return await self.delete_one(collection, {"_id": ObjectId(id)})
|
||||
|
||||
# Aggregation
|
||||
|
||||
async def aggregate(
|
||||
self,
|
||||
collection: str,
|
||||
pipeline: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Run an aggregation pipeline.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
pipeline: Aggregation pipeline stages.
|
||||
|
||||
Returns:
|
||||
List of aggregation results.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
cursor = coll.aggregate(pipeline)
|
||||
return await cursor.to_list(length=None)
|
||||
|
||||
# Count operations
|
||||
|
||||
async def count_documents(
|
||||
self,
|
||||
collection: str,
|
||||
filter: Optional[dict[str, Any]] = None,
|
||||
) -> int:
|
||||
"""Count documents matching a filter.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
filter: Query filter (empty for all).
|
||||
|
||||
Returns:
|
||||
Number of matching documents.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
return await coll.count_documents(filter or {})
|
||||
|
||||
async def estimated_document_count(self, collection: str) -> int:
|
||||
"""Get estimated document count (faster than count_documents).
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
|
||||
Returns:
|
||||
Estimated number of documents.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
return await coll.estimated_document_count()
|
||||
|
||||
# Index operations
|
||||
|
||||
async def create_index(
|
||||
self,
|
||||
collection: str,
|
||||
keys: list[tuple[str, int]],
|
||||
unique: bool = False,
|
||||
name: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Create an index.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
keys: Index keys [(field, direction), ...].
|
||||
unique: Unique index.
|
||||
name: Index name.
|
||||
|
||||
Returns:
|
||||
Created index name.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
return await coll.create_index(keys, unique=unique, name=name)
|
||||
|
||||
async def drop_index(self, collection: str, name: str) -> None:
|
||||
"""Drop an index.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
name: Index name.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
await coll.drop_index(name)
|
||||
|
||||
async def list_indexes(self, collection: str) -> list[dict[str, Any]]:
|
||||
"""List all indexes on a collection.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
|
||||
Returns:
|
||||
List of index specifications.
|
||||
"""
|
||||
coll = self._get_collection(collection)
|
||||
cursor = coll.list_indexes()
|
||||
return await cursor.to_list(length=None)
|
||||
|
||||
# Collection operations
|
||||
|
||||
async def list_collections(self) -> list[str]:
|
||||
"""List all collections in the database.
|
||||
|
||||
Returns:
|
||||
List of collection names.
|
||||
"""
|
||||
return await self._db.list_collection_names()
|
||||
|
||||
async def create_collection(self, name: str) -> None:
|
||||
"""Create a new collection.
|
||||
|
||||
Args:
|
||||
name: Collection name.
|
||||
"""
|
||||
await self._db.create_collection(name)
|
||||
|
||||
async def drop_collection(self, name: str) -> None:
|
||||
"""Drop a collection.
|
||||
|
||||
Args:
|
||||
name: Collection name.
|
||||
"""
|
||||
await self._db.drop_collection(name)
|
||||
|
||||
# Database operations
|
||||
|
||||
def use_database(self, name: str) -> None:
|
||||
"""Switch to a different database.
|
||||
|
||||
Args:
|
||||
name: Database name.
|
||||
"""
|
||||
self._db = self._client[name]
|
||||
|
||||
async def list_databases(self) -> list[str]:
|
||||
"""List all databases.
|
||||
|
||||
Returns:
|
||||
List of database names.
|
||||
"""
|
||||
return await self._client.list_database_names()
|
||||
|
||||
async def __aenter__(self) -> DocumentDBClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,539 @@
|
||||
"""Nuclio serverless functions client wrapper for Hanzo infrastructure.
|
||||
|
||||
Provides async interface to Nuclio for deploying and invoking
|
||||
serverless functions with auto-scaling and GPU support.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class FunctionsConfig(BaseModel):
|
||||
"""Configuration for Nuclio connection."""
|
||||
|
||||
dashboard_url: str = Field(default="http://localhost:8070", description="Nuclio dashboard URL")
|
||||
api_key: Optional[str] = Field(default=None, description="API key for authentication")
|
||||
namespace: str = Field(default="nuclio", description="Kubernetes namespace")
|
||||
default_runtime: str = Field(default="python:3.11", description="Default function runtime")
|
||||
default_handler: str = Field(default="main:handler", description="Default function handler")
|
||||
registry: Optional[str] = Field(default=None, description="Container registry URL")
|
||||
timeout: float = Field(default=30.0, description="Request timeout in seconds")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> FunctionsConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
NUCLIO_DASHBOARD_URL: Dashboard URL (default: http://localhost:8070)
|
||||
NUCLIO_API_KEY: API key for authentication
|
||||
NUCLIO_NAMESPACE: Kubernetes namespace (default: nuclio)
|
||||
NUCLIO_RUNTIME: Default runtime (default: python:3.11)
|
||||
NUCLIO_REGISTRY: Container registry URL
|
||||
"""
|
||||
return cls(
|
||||
dashboard_url=os.getenv("NUCLIO_DASHBOARD_URL", "http://localhost:8070"),
|
||||
api_key=os.getenv("NUCLIO_API_KEY"),
|
||||
namespace=os.getenv("NUCLIO_NAMESPACE", "nuclio"),
|
||||
default_runtime=os.getenv("NUCLIO_RUNTIME", "python:3.11"),
|
||||
registry=os.getenv("NUCLIO_REGISTRY"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionSpec:
|
||||
"""Specification for a Nuclio function."""
|
||||
|
||||
name: str
|
||||
handler: str = "main:handler"
|
||||
runtime: str = "python:3.11"
|
||||
code: Optional[str] = None # Inline code
|
||||
code_path: Optional[str] = None # Path to code directory
|
||||
image: Optional[str] = None # Pre-built image
|
||||
env: dict[str, str] = field(default_factory=dict)
|
||||
min_replicas: int = 0
|
||||
max_replicas: int = 10
|
||||
target_cpu: int = 75 # CPU utilization target for scaling
|
||||
triggers: dict[str, Any] = field(default_factory=dict)
|
||||
resources: dict[str, Any] = field(default_factory=dict)
|
||||
build_commands: list[str] = field(default_factory=list)
|
||||
requirements: list[str] = field(default_factory=list) # pip requirements
|
||||
labels: dict[str, str] = field(default_factory=dict)
|
||||
annotations: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
def to_nuclio_spec(self) -> dict[str, Any]:
|
||||
"""Convert to Nuclio API spec format."""
|
||||
spec: dict[str, Any] = {
|
||||
"spec": {
|
||||
"handler": self.handler,
|
||||
"runtime": self.runtime,
|
||||
"minReplicas": self.min_replicas,
|
||||
"maxReplicas": self.max_replicas,
|
||||
"targetCPU": self.target_cpu,
|
||||
"env": [{"name": k, "value": v} for k, v in self.env.items()],
|
||||
},
|
||||
"metadata": {
|
||||
"name": self.name,
|
||||
"labels": self.labels,
|
||||
"annotations": self.annotations,
|
||||
},
|
||||
}
|
||||
|
||||
if self.code:
|
||||
spec["spec"]["build"] = {
|
||||
"functionSourceCode": self.code,
|
||||
}
|
||||
elif self.code_path:
|
||||
spec["spec"]["build"] = {
|
||||
"path": self.code_path,
|
||||
}
|
||||
elif self.image:
|
||||
spec["spec"]["image"] = self.image
|
||||
|
||||
if self.build_commands:
|
||||
spec["spec"]["build"] = spec["spec"].get("build", {})
|
||||
spec["spec"]["build"]["commands"] = self.build_commands
|
||||
|
||||
if self.requirements:
|
||||
# Add pip install to build commands
|
||||
pip_cmd = f"pip install {' '.join(self.requirements)}"
|
||||
commands = spec["spec"].get("build", {}).get("commands", [])
|
||||
commands.insert(0, pip_cmd)
|
||||
spec["spec"]["build"] = spec["spec"].get("build", {})
|
||||
spec["spec"]["build"]["commands"] = commands
|
||||
|
||||
if self.triggers:
|
||||
spec["spec"]["triggers"] = self.triggers
|
||||
|
||||
if self.resources:
|
||||
spec["spec"]["resources"] = self.resources
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionStatus:
|
||||
"""Status of a deployed function."""
|
||||
|
||||
name: str
|
||||
state: str # ready, building, error, etc.
|
||||
replicas: int = 0
|
||||
version: str = ""
|
||||
invoke_url: Optional[str] = None
|
||||
internal_url: Optional[str] = None
|
||||
created_at: Optional[datetime] = None
|
||||
updated_at: Optional[datetime] = None
|
||||
message: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class InvokeResult:
|
||||
"""Result of a function invocation."""
|
||||
|
||||
status_code: int
|
||||
body: Any
|
||||
headers: dict[str, str] = field(default_factory=dict)
|
||||
duration_ms: int = 0
|
||||
|
||||
|
||||
class FunctionsClient:
|
||||
"""Async client for Nuclio serverless functions.
|
||||
|
||||
Wraps Nuclio HTTP API for deploying, managing, and invoking
|
||||
serverless functions with Hanzo conventions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = FunctionsClient(FunctionsConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Deploy a function
|
||||
spec = FunctionSpec(
|
||||
name="hello",
|
||||
code='''
|
||||
def handler(context, event):
|
||||
return "Hello, " + event.body.decode()
|
||||
''',
|
||||
runtime="python:3.11",
|
||||
)
|
||||
await client.deploy(spec)
|
||||
|
||||
# Invoke the function
|
||||
result = await client.invoke("hello", body=b"World")
|
||||
print(result.body) # "Hello, World"
|
||||
|
||||
# List functions
|
||||
functions = await client.list_functions()
|
||||
for fn in functions:
|
||||
print(f"{fn.name}: {fn.state}")
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[FunctionsConfig] = None) -> None:
|
||||
"""Initialize functions client.
|
||||
|
||||
Args:
|
||||
config: Nuclio configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or FunctionsConfig.from_env()
|
||||
self._client: Any = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection (validates connectivity)."""
|
||||
try:
|
||||
import httpx
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"httpx is required for FunctionsClient. "
|
||||
"Install with: pip install httpx"
|
||||
) from e
|
||||
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=self.config.dashboard_url,
|
||||
timeout=self.config.timeout,
|
||||
headers=self._get_headers(),
|
||||
)
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
"""Get request headers."""
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if self.config.api_key:
|
||||
headers["X-nuclio-project-name"] = "default"
|
||||
headers["X-v3io-session-key"] = self.config.api_key
|
||||
return headers
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if Nuclio is healthy.
|
||||
|
||||
Returns:
|
||||
True if Nuclio dashboard is reachable.
|
||||
"""
|
||||
if not self._client:
|
||||
return False
|
||||
try:
|
||||
response = await self._client.get("/api/projects")
|
||||
return response.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# Function management
|
||||
|
||||
async def deploy(
|
||||
self,
|
||||
spec: FunctionSpec,
|
||||
wait: bool = True,
|
||||
timeout: float = 300.0,
|
||||
) -> FunctionStatus:
|
||||
"""Deploy a function.
|
||||
|
||||
Args:
|
||||
spec: Function specification.
|
||||
wait: Wait for deployment to complete.
|
||||
timeout: Deployment timeout in seconds.
|
||||
|
||||
Returns:
|
||||
Function status after deployment.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Create/update function
|
||||
nuclio_spec = spec.to_nuclio_spec()
|
||||
response = await self._client.post(
|
||||
f"/api/functions/{spec.name}",
|
||||
json=nuclio_spec,
|
||||
params={"namespace": self.config.namespace},
|
||||
)
|
||||
|
||||
if response.status_code not in (200, 201, 202):
|
||||
raise Exception(f"Failed to deploy function: {response.text}")
|
||||
|
||||
if wait:
|
||||
# Poll until ready or timeout
|
||||
start_time = datetime.utcnow()
|
||||
while (datetime.utcnow() - start_time).total_seconds() < timeout:
|
||||
status = await self.get_function(spec.name)
|
||||
if status.state in ("ready", "imported"):
|
||||
return status
|
||||
if status.state == "error":
|
||||
raise Exception(f"Function deployment failed: {status.message}")
|
||||
await asyncio.sleep(2)
|
||||
|
||||
raise TimeoutError(f"Function deployment timed out after {timeout}s")
|
||||
|
||||
return await self.get_function(spec.name)
|
||||
|
||||
async def delete_function(self, name: str) -> bool:
|
||||
"""Delete a function.
|
||||
|
||||
Args:
|
||||
name: Function name.
|
||||
|
||||
Returns:
|
||||
True if function was deleted.
|
||||
"""
|
||||
response = await self._client.delete(
|
||||
f"/api/functions/{name}",
|
||||
params={"namespace": self.config.namespace},
|
||||
)
|
||||
return response.status_code in (200, 204)
|
||||
|
||||
async def get_function(self, name: str) -> FunctionStatus:
|
||||
"""Get function status.
|
||||
|
||||
Args:
|
||||
name: Function name.
|
||||
|
||||
Returns:
|
||||
Function status.
|
||||
"""
|
||||
response = await self._client.get(
|
||||
f"/api/functions/{name}",
|
||||
params={"namespace": self.config.namespace},
|
||||
)
|
||||
|
||||
if response.status_code == 404:
|
||||
raise Exception(f"Function not found: {name}")
|
||||
|
||||
data = response.json()
|
||||
status = data.get("status", {})
|
||||
metadata = data.get("metadata", {})
|
||||
|
||||
return FunctionStatus(
|
||||
name=name,
|
||||
state=status.get("state", "unknown"),
|
||||
replicas=status.get("replicas", 0),
|
||||
version=metadata.get("version", ""),
|
||||
invoke_url=status.get("httpPort"),
|
||||
internal_url=status.get("internalInvocationUrls", [None])[0] if status.get("internalInvocationUrls") else None,
|
||||
message=status.get("message"),
|
||||
)
|
||||
|
||||
async def list_functions(self) -> list[FunctionStatus]:
|
||||
"""List all functions.
|
||||
|
||||
Returns:
|
||||
List of function statuses.
|
||||
"""
|
||||
response = await self._client.get(
|
||||
"/api/functions",
|
||||
params={"namespace": self.config.namespace},
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to list functions: {response.text}")
|
||||
|
||||
functions = []
|
||||
for name, data in response.json().items():
|
||||
status = data.get("status", {})
|
||||
metadata = data.get("metadata", {})
|
||||
functions.append(FunctionStatus(
|
||||
name=name,
|
||||
state=status.get("state", "unknown"),
|
||||
replicas=status.get("replicas", 0),
|
||||
version=metadata.get("version", ""),
|
||||
invoke_url=status.get("httpPort"),
|
||||
message=status.get("message"),
|
||||
))
|
||||
|
||||
return functions
|
||||
|
||||
# Function invocation
|
||||
|
||||
async def invoke(
|
||||
self,
|
||||
name: str,
|
||||
body: bytes = b"",
|
||||
method: str = "POST",
|
||||
path: str = "",
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
query: Optional[dict[str, str]] = None,
|
||||
) -> InvokeResult:
|
||||
"""Invoke a function.
|
||||
|
||||
Args:
|
||||
name: Function name.
|
||||
body: Request body.
|
||||
method: HTTP method.
|
||||
path: Request path within function.
|
||||
headers: Additional headers.
|
||||
query: Query parameters.
|
||||
|
||||
Returns:
|
||||
Invocation result.
|
||||
"""
|
||||
import time
|
||||
|
||||
# Get function invoke URL
|
||||
status = await self.get_function(name)
|
||||
if not status.invoke_url and not status.internal_url:
|
||||
raise Exception(f"Function {name} has no invoke URL")
|
||||
|
||||
# Determine invoke URL
|
||||
if status.internal_url:
|
||||
invoke_url = status.internal_url
|
||||
else:
|
||||
# Construct URL from dashboard URL and port
|
||||
base = self.config.dashboard_url.rsplit(":", 1)[0]
|
||||
invoke_url = f"{base}:{status.invoke_url}"
|
||||
|
||||
# Build request URL
|
||||
url = f"{invoke_url}/{path.lstrip('/')}" if path else invoke_url
|
||||
|
||||
# Prepare headers
|
||||
req_headers = headers or {}
|
||||
|
||||
start = time.monotonic()
|
||||
response = await self._client.request(
|
||||
method=method,
|
||||
url=url,
|
||||
content=body,
|
||||
headers=req_headers,
|
||||
params=query,
|
||||
)
|
||||
duration = int((time.monotonic() - start) * 1000)
|
||||
|
||||
# Parse response
|
||||
try:
|
||||
response_body = response.json()
|
||||
except Exception:
|
||||
response_body = response.text
|
||||
|
||||
return InvokeResult(
|
||||
status_code=response.status_code,
|
||||
body=response_body,
|
||||
headers=dict(response.headers),
|
||||
duration_ms=duration,
|
||||
)
|
||||
|
||||
async def invoke_async(
|
||||
self,
|
||||
name: str,
|
||||
body: bytes = b"",
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
) -> str:
|
||||
"""Invoke a function asynchronously.
|
||||
|
||||
Args:
|
||||
name: Function name.
|
||||
body: Request body.
|
||||
headers: Additional headers.
|
||||
|
||||
Returns:
|
||||
Invocation ID for checking status.
|
||||
"""
|
||||
req_headers = headers or {}
|
||||
req_headers["X-Nuclio-Function-Async"] = "true"
|
||||
|
||||
result = await self.invoke(name, body, headers=req_headers)
|
||||
|
||||
# Return invocation ID from response
|
||||
return result.headers.get("X-Nuclio-Invoke-Id", "")
|
||||
|
||||
# Scaling
|
||||
|
||||
async def scale_function(
|
||||
self,
|
||||
name: str,
|
||||
replicas: Optional[int] = None,
|
||||
min_replicas: Optional[int] = None,
|
||||
max_replicas: Optional[int] = None,
|
||||
) -> FunctionStatus:
|
||||
"""Scale a function.
|
||||
|
||||
Args:
|
||||
name: Function name.
|
||||
replicas: Fixed replica count (overrides auto-scaling).
|
||||
min_replicas: Minimum replicas for auto-scaling.
|
||||
max_replicas: Maximum replicas for auto-scaling.
|
||||
|
||||
Returns:
|
||||
Updated function status.
|
||||
"""
|
||||
# Get current function
|
||||
response = await self._client.get(
|
||||
f"/api/functions/{name}",
|
||||
params={"namespace": self.config.namespace},
|
||||
)
|
||||
|
||||
if response.status_code == 404:
|
||||
raise Exception(f"Function not found: {name}")
|
||||
|
||||
data = response.json()
|
||||
spec = data.get("spec", {})
|
||||
|
||||
# Update scaling config
|
||||
if replicas is not None:
|
||||
spec["replicas"] = replicas
|
||||
if min_replicas is not None:
|
||||
spec["minReplicas"] = min_replicas
|
||||
if max_replicas is not None:
|
||||
spec["maxReplicas"] = max_replicas
|
||||
|
||||
# Update function
|
||||
data["spec"] = spec
|
||||
response = await self._client.put(
|
||||
f"/api/functions/{name}",
|
||||
json=data,
|
||||
params={"namespace": self.config.namespace},
|
||||
)
|
||||
|
||||
if response.status_code not in (200, 202):
|
||||
raise Exception(f"Failed to scale function: {response.text}")
|
||||
|
||||
return await self.get_function(name)
|
||||
|
||||
# Logs
|
||||
|
||||
async def get_logs(
|
||||
self,
|
||||
name: str,
|
||||
since: Optional[datetime] = None,
|
||||
follow: bool = False,
|
||||
):
|
||||
"""Get function logs.
|
||||
|
||||
Args:
|
||||
name: Function name.
|
||||
since: Get logs since this time.
|
||||
follow: Stream logs (not implemented in basic version).
|
||||
|
||||
Yields:
|
||||
Log lines.
|
||||
"""
|
||||
params: dict[str, Any] = {
|
||||
"namespace": self.config.namespace,
|
||||
"function": name,
|
||||
}
|
||||
if since:
|
||||
params["since"] = since.isoformat()
|
||||
|
||||
response = await self._client.get("/api/logs", params=params)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to get logs: {response.text}")
|
||||
|
||||
for line in response.text.split("\n"):
|
||||
if line.strip():
|
||||
yield line
|
||||
|
||||
async def __aenter__(self) -> FunctionsClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,525 @@
|
||||
"""Redis/Valkey key-value store client wrapper for Hanzo infrastructure.
|
||||
|
||||
Provides async interface to Redis/Valkey for caching, sessions,
|
||||
and general key-value storage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class KVConfig(BaseModel):
|
||||
"""Configuration for Redis/Valkey connection."""
|
||||
|
||||
host: str = Field(default="localhost", description="Redis server host")
|
||||
port: int = Field(default=6379, description="Redis server port")
|
||||
password: Optional[str] = Field(default=None, description="Redis password")
|
||||
db: int = Field(default=0, description="Redis database number")
|
||||
url: Optional[str] = Field(default=None, description="Full Redis URL (overrides host/port)")
|
||||
ssl: bool = Field(default=False, description="Use SSL/TLS")
|
||||
socket_timeout: float = Field(default=5.0, description="Socket timeout in seconds")
|
||||
socket_connect_timeout: float = Field(default=5.0, description="Connection timeout")
|
||||
max_connections: int = Field(default=10, description="Max pool connections")
|
||||
decode_responses: bool = Field(default=True, description="Decode responses as strings")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> KVConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
REDIS_HOST: Server host (default: localhost)
|
||||
REDIS_PORT: Server port (default: 6379)
|
||||
REDIS_PASSWORD: Authentication password
|
||||
REDIS_DB: Database number (default: 0)
|
||||
REDIS_URL: Full URL (overrides host/port)
|
||||
REDIS_SSL: Use SSL (default: false)
|
||||
"""
|
||||
return cls(
|
||||
host=os.getenv("REDIS_HOST", "localhost"),
|
||||
port=int(os.getenv("REDIS_PORT", "6379")),
|
||||
password=os.getenv("REDIS_PASSWORD"),
|
||||
db=int(os.getenv("REDIS_DB", "0")),
|
||||
url=os.getenv("REDIS_URL"),
|
||||
ssl=os.getenv("REDIS_SSL", "").lower() in ("true", "1", "yes"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class KVResult:
|
||||
"""Result of a KV operation with metadata."""
|
||||
|
||||
key: str
|
||||
value: Any
|
||||
ttl: Optional[int] = None # TTL in seconds, None if no expiry
|
||||
|
||||
|
||||
class KVClient:
|
||||
"""Async client for Redis/Valkey key-value store.
|
||||
|
||||
Wraps redis-py with async methods and Hanzo conventions.
|
||||
Supports both Redis and Valkey (Redis-compatible).
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = KVClient(KVConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Basic operations
|
||||
await client.set("key", "value", ttl=3600)
|
||||
value = await client.get("key")
|
||||
|
||||
# JSON operations
|
||||
await client.set_json("config", {"theme": "dark"})
|
||||
config = await client.get_json("config")
|
||||
|
||||
# Hash operations
|
||||
await client.hset("user:1", {"name": "Alice", "email": "[email protected]"})
|
||||
user = await client.hgetall("user:1")
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[KVConfig] = None) -> None:
|
||||
"""Initialize KV client.
|
||||
|
||||
Args:
|
||||
config: Redis configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or KVConfig.from_env()
|
||||
self._client: Any = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection to Redis server."""
|
||||
try:
|
||||
import redis.asyncio as redis
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"redis is required for KVClient. " "Install with: pip install redis"
|
||||
) from e
|
||||
|
||||
if self.config.url:
|
||||
self._client = redis.from_url(
|
||||
self.config.url,
|
||||
decode_responses=self.config.decode_responses,
|
||||
socket_timeout=self.config.socket_timeout,
|
||||
socket_connect_timeout=self.config.socket_connect_timeout,
|
||||
max_connections=self.config.max_connections,
|
||||
)
|
||||
else:
|
||||
self._client = redis.Redis(
|
||||
host=self.config.host,
|
||||
port=self.config.port,
|
||||
password=self.config.password,
|
||||
db=self.config.db,
|
||||
ssl=self.config.ssl,
|
||||
decode_responses=self.config.decode_responses,
|
||||
socket_timeout=self.config.socket_timeout,
|
||||
socket_connect_timeout=self.config.socket_connect_timeout,
|
||||
max_connections=self.config.max_connections,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if Redis server is healthy.
|
||||
|
||||
Returns:
|
||||
True if server is reachable and responding.
|
||||
"""
|
||||
if not self._client:
|
||||
return False
|
||||
try:
|
||||
await self._client.ping()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# Basic string operations
|
||||
|
||||
async def get(self, key: str) -> Optional[str]:
|
||||
"""Get a string value.
|
||||
|
||||
Args:
|
||||
key: Key to retrieve.
|
||||
|
||||
Returns:
|
||||
Value or None if not found.
|
||||
"""
|
||||
return await self._client.get(key)
|
||||
|
||||
async def set(
|
||||
self,
|
||||
key: str,
|
||||
value: str,
|
||||
ttl: Optional[Union[int, timedelta]] = None,
|
||||
nx: bool = False,
|
||||
xx: bool = False,
|
||||
) -> bool:
|
||||
"""Set a string value.
|
||||
|
||||
Args:
|
||||
key: Key to set.
|
||||
value: Value to store.
|
||||
ttl: Time to live in seconds or timedelta.
|
||||
nx: Only set if key does not exist.
|
||||
xx: Only set if key already exists.
|
||||
|
||||
Returns:
|
||||
True if successful.
|
||||
"""
|
||||
ex = ttl if isinstance(ttl, int) else None
|
||||
px = int(ttl.total_seconds() * 1000) if isinstance(ttl, timedelta) else None
|
||||
|
||||
result = await self._client.set(key, value, ex=ex, px=px, nx=nx, xx=xx)
|
||||
return result is not None
|
||||
|
||||
async def delete(self, *keys: str) -> int:
|
||||
"""Delete one or more keys.
|
||||
|
||||
Args:
|
||||
keys: Keys to delete.
|
||||
|
||||
Returns:
|
||||
Number of keys deleted.
|
||||
"""
|
||||
return await self._client.delete(*keys)
|
||||
|
||||
async def exists(self, *keys: str) -> int:
|
||||
"""Check if keys exist.
|
||||
|
||||
Args:
|
||||
keys: Keys to check.
|
||||
|
||||
Returns:
|
||||
Number of keys that exist.
|
||||
"""
|
||||
return await self._client.exists(*keys)
|
||||
|
||||
async def expire(self, key: str, ttl: Union[int, timedelta]) -> bool:
|
||||
"""Set expiration on a key.
|
||||
|
||||
Args:
|
||||
key: Key to expire.
|
||||
ttl: Time to live in seconds or timedelta.
|
||||
|
||||
Returns:
|
||||
True if expiration was set.
|
||||
"""
|
||||
seconds = ttl if isinstance(ttl, int) else int(ttl.total_seconds())
|
||||
return await self._client.expire(key, seconds)
|
||||
|
||||
async def ttl(self, key: str) -> int:
|
||||
"""Get time to live for a key.
|
||||
|
||||
Args:
|
||||
key: Key to check.
|
||||
|
||||
Returns:
|
||||
TTL in seconds, -1 if no expiry, -2 if key doesn't exist.
|
||||
"""
|
||||
return await self._client.ttl(key)
|
||||
|
||||
# JSON operations
|
||||
|
||||
async def get_json(self, key: str) -> Optional[Any]:
|
||||
"""Get a JSON value.
|
||||
|
||||
Args:
|
||||
key: Key to retrieve.
|
||||
|
||||
Returns:
|
||||
Parsed JSON value or None if not found.
|
||||
"""
|
||||
value = await self.get(key)
|
||||
if value is None:
|
||||
return None
|
||||
return json.loads(value)
|
||||
|
||||
async def set_json(
|
||||
self,
|
||||
key: str,
|
||||
value: Any,
|
||||
ttl: Optional[Union[int, timedelta]] = None,
|
||||
) -> bool:
|
||||
"""Set a JSON value.
|
||||
|
||||
Args:
|
||||
key: Key to set.
|
||||
value: Value to serialize as JSON.
|
||||
ttl: Time to live.
|
||||
|
||||
Returns:
|
||||
True if successful.
|
||||
"""
|
||||
return await self.set(key, json.dumps(value), ttl=ttl)
|
||||
|
||||
# Hash operations
|
||||
|
||||
async def hget(self, name: str, key: str) -> Optional[str]:
|
||||
"""Get a hash field value.
|
||||
|
||||
Args:
|
||||
name: Hash name.
|
||||
key: Field key.
|
||||
|
||||
Returns:
|
||||
Field value or None.
|
||||
"""
|
||||
return await self._client.hget(name, key)
|
||||
|
||||
async def hset(self, name: str, mapping: dict[str, Any]) -> int:
|
||||
"""Set hash fields.
|
||||
|
||||
Args:
|
||||
name: Hash name.
|
||||
mapping: Field-value pairs.
|
||||
|
||||
Returns:
|
||||
Number of fields added.
|
||||
"""
|
||||
return await self._client.hset(name, mapping=mapping)
|
||||
|
||||
async def hgetall(self, name: str) -> dict[str, str]:
|
||||
"""Get all hash fields.
|
||||
|
||||
Args:
|
||||
name: Hash name.
|
||||
|
||||
Returns:
|
||||
Dictionary of field-value pairs.
|
||||
"""
|
||||
return await self._client.hgetall(name)
|
||||
|
||||
async def hdel(self, name: str, *keys: str) -> int:
|
||||
"""Delete hash fields.
|
||||
|
||||
Args:
|
||||
name: Hash name.
|
||||
keys: Fields to delete.
|
||||
|
||||
Returns:
|
||||
Number of fields deleted.
|
||||
"""
|
||||
return await self._client.hdel(name, *keys)
|
||||
|
||||
# List operations
|
||||
|
||||
async def lpush(self, key: str, *values: str) -> int:
|
||||
"""Push values to the head of a list.
|
||||
|
||||
Args:
|
||||
key: List key.
|
||||
values: Values to push.
|
||||
|
||||
Returns:
|
||||
Length of list after push.
|
||||
"""
|
||||
return await self._client.lpush(key, *values)
|
||||
|
||||
async def rpush(self, key: str, *values: str) -> int:
|
||||
"""Push values to the tail of a list.
|
||||
|
||||
Args:
|
||||
key: List key.
|
||||
values: Values to push.
|
||||
|
||||
Returns:
|
||||
Length of list after push.
|
||||
"""
|
||||
return await self._client.rpush(key, *values)
|
||||
|
||||
async def lpop(self, key: str, count: Optional[int] = None) -> Optional[Union[str, list[str]]]:
|
||||
"""Pop values from the head of a list.
|
||||
|
||||
Args:
|
||||
key: List key.
|
||||
count: Number of values to pop.
|
||||
|
||||
Returns:
|
||||
Popped value(s) or None.
|
||||
"""
|
||||
return await self._client.lpop(key, count)
|
||||
|
||||
async def rpop(self, key: str, count: Optional[int] = None) -> Optional[Union[str, list[str]]]:
|
||||
"""Pop values from the tail of a list.
|
||||
|
||||
Args:
|
||||
key: List key.
|
||||
count: Number of values to pop.
|
||||
|
||||
Returns:
|
||||
Popped value(s) or None.
|
||||
"""
|
||||
return await self._client.rpop(key, count)
|
||||
|
||||
async def lrange(self, key: str, start: int, end: int) -> list[str]:
|
||||
"""Get a range of list elements.
|
||||
|
||||
Args:
|
||||
key: List key.
|
||||
start: Start index.
|
||||
end: End index (-1 for last).
|
||||
|
||||
Returns:
|
||||
List of elements.
|
||||
"""
|
||||
return await self._client.lrange(key, start, end)
|
||||
|
||||
async def llen(self, key: str) -> int:
|
||||
"""Get list length.
|
||||
|
||||
Args:
|
||||
key: List key.
|
||||
|
||||
Returns:
|
||||
Length of list.
|
||||
"""
|
||||
return await self._client.llen(key)
|
||||
|
||||
# Set operations
|
||||
|
||||
async def sadd(self, key: str, *members: str) -> int:
|
||||
"""Add members to a set.
|
||||
|
||||
Args:
|
||||
key: Set key.
|
||||
members: Members to add.
|
||||
|
||||
Returns:
|
||||
Number of members added.
|
||||
"""
|
||||
return await self._client.sadd(key, *members)
|
||||
|
||||
async def srem(self, key: str, *members: str) -> int:
|
||||
"""Remove members from a set.
|
||||
|
||||
Args:
|
||||
key: Set key.
|
||||
members: Members to remove.
|
||||
|
||||
Returns:
|
||||
Number of members removed.
|
||||
"""
|
||||
return await self._client.srem(key, *members)
|
||||
|
||||
async def smembers(self, key: str) -> set[str]:
|
||||
"""Get all members of a set.
|
||||
|
||||
Args:
|
||||
key: Set key.
|
||||
|
||||
Returns:
|
||||
Set of members.
|
||||
"""
|
||||
return await self._client.smembers(key)
|
||||
|
||||
async def sismember(self, key: str, member: str) -> bool:
|
||||
"""Check if value is a set member.
|
||||
|
||||
Args:
|
||||
key: Set key.
|
||||
member: Member to check.
|
||||
|
||||
Returns:
|
||||
True if member exists.
|
||||
"""
|
||||
return await self._client.sismember(key, member)
|
||||
|
||||
# Atomic operations
|
||||
|
||||
async def incr(self, key: str, amount: int = 1) -> int:
|
||||
"""Increment a value.
|
||||
|
||||
Args:
|
||||
key: Key to increment.
|
||||
amount: Amount to increment by.
|
||||
|
||||
Returns:
|
||||
New value.
|
||||
"""
|
||||
return await self._client.incrby(key, amount)
|
||||
|
||||
async def decr(self, key: str, amount: int = 1) -> int:
|
||||
"""Decrement a value.
|
||||
|
||||
Args:
|
||||
key: Key to decrement.
|
||||
amount: Amount to decrement by.
|
||||
|
||||
Returns:
|
||||
New value.
|
||||
"""
|
||||
return await self._client.decrby(key, amount)
|
||||
|
||||
# Pub/Sub (basic - for full pub/sub use PubSubClient)
|
||||
|
||||
async def publish(self, channel: str, message: str) -> int:
|
||||
"""Publish a message to a channel.
|
||||
|
||||
Args:
|
||||
channel: Channel name.
|
||||
message: Message to publish.
|
||||
|
||||
Returns:
|
||||
Number of subscribers that received the message.
|
||||
"""
|
||||
return await self._client.publish(channel, message)
|
||||
|
||||
# Keys operations
|
||||
|
||||
async def keys(self, pattern: str = "*") -> list[str]:
|
||||
"""Find keys matching a pattern.
|
||||
|
||||
Args:
|
||||
pattern: Glob-style pattern.
|
||||
|
||||
Returns:
|
||||
List of matching keys.
|
||||
"""
|
||||
return await self._client.keys(pattern)
|
||||
|
||||
async def scan(
|
||||
self,
|
||||
cursor: int = 0,
|
||||
match: Optional[str] = None,
|
||||
count: int = 100,
|
||||
) -> tuple[int, list[str]]:
|
||||
"""Incrementally iterate keys.
|
||||
|
||||
Args:
|
||||
cursor: Cursor position.
|
||||
match: Pattern to match.
|
||||
count: Approximate number per iteration.
|
||||
|
||||
Returns:
|
||||
Tuple of (next_cursor, keys).
|
||||
"""
|
||||
return await self._client.scan(cursor, match=match, count=count)
|
||||
|
||||
async def flushdb(self) -> bool:
|
||||
"""Flush the current database.
|
||||
|
||||
Returns:
|
||||
True if successful.
|
||||
"""
|
||||
await self._client.flushdb()
|
||||
return True
|
||||
|
||||
async def __aenter__(self) -> KVClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,534 @@
|
||||
"""NATS pub/sub client wrapper for Hanzo infrastructure.
|
||||
|
||||
Provides async interface to NATS for messaging, pub/sub,
|
||||
request/reply, and streaming with JetStream.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, AsyncIterator, Awaitable, Callable, Optional, Sequence
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class PubSubConfig(BaseModel):
|
||||
"""Configuration for NATS connection."""
|
||||
|
||||
servers: list[str] = Field(default=["nats://localhost:4222"], description="NATS server URLs")
|
||||
user: Optional[str] = Field(default=None, description="Username for authentication")
|
||||
password: Optional[str] = Field(default=None, description="Password for authentication")
|
||||
token: Optional[str] = Field(default=None, description="Token for authentication")
|
||||
nkey_seed: Optional[str] = Field(default=None, description="NKey seed for authentication")
|
||||
tls: bool = Field(default=False, description="Use TLS")
|
||||
name: str = Field(default="hanzo-client", description="Client name")
|
||||
connect_timeout: float = Field(default=5.0, description="Connection timeout in seconds")
|
||||
reconnect_time_wait: float = Field(default=2.0, description="Time between reconnect attempts")
|
||||
max_reconnect_attempts: int = Field(default=60, description="Max reconnection attempts")
|
||||
ping_interval: float = Field(default=120.0, description="Ping interval in seconds")
|
||||
max_outstanding_pings: int = Field(default=2, description="Max outstanding pings")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> PubSubConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
NATS_SERVERS: Comma-separated list of NATS URLs (default: nats://localhost:4222)
|
||||
NATS_USER: Username for authentication
|
||||
NATS_PASSWORD: Password for authentication
|
||||
NATS_TOKEN: Token for authentication
|
||||
NATS_NKEY_SEED: NKey seed for authentication
|
||||
NATS_TLS: Use TLS (default: false)
|
||||
"""
|
||||
servers_str = os.getenv("NATS_SERVERS", "nats://localhost:4222")
|
||||
servers = [s.strip() for s in servers_str.split(",")]
|
||||
|
||||
return cls(
|
||||
servers=servers,
|
||||
user=os.getenv("NATS_USER"),
|
||||
password=os.getenv("NATS_PASSWORD"),
|
||||
token=os.getenv("NATS_TOKEN"),
|
||||
nkey_seed=os.getenv("NATS_NKEY_SEED"),
|
||||
tls=os.getenv("NATS_TLS", "").lower() in ("true", "1", "yes"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Message:
|
||||
"""A NATS message."""
|
||||
|
||||
subject: str
|
||||
data: bytes
|
||||
reply: Optional[str] = None
|
||||
headers: dict[str, str] = field(default_factory=dict)
|
||||
timestamp: Optional[datetime] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Subscription:
|
||||
"""A NATS subscription handle."""
|
||||
|
||||
subject: str
|
||||
queue: Optional[str] = None
|
||||
_sub: Any = None
|
||||
|
||||
async def unsubscribe(self) -> None:
|
||||
"""Unsubscribe from the subject."""
|
||||
if self._sub:
|
||||
await self._sub.unsubscribe()
|
||||
|
||||
async def drain(self) -> None:
|
||||
"""Drain the subscription before unsubscribing."""
|
||||
if self._sub:
|
||||
await self._sub.drain()
|
||||
|
||||
|
||||
MessageHandler = Callable[[Message], Awaitable[None]]
|
||||
|
||||
|
||||
class PubSubClient:
|
||||
"""Async client for NATS messaging.
|
||||
|
||||
Wraps nats-py with Hanzo conventions for pub/sub, request/reply,
|
||||
and JetStream operations.
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = PubSubClient(PubSubConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Simple pub/sub
|
||||
async def handler(msg: Message):
|
||||
print(f"Received: {msg.data.decode()}")
|
||||
|
||||
sub = await client.subscribe("events.>", handler)
|
||||
await client.publish("events.user.created", b"user123")
|
||||
|
||||
# Request/reply
|
||||
response = await client.request("api.users.get", b"user123", timeout=5.0)
|
||||
|
||||
# JetStream (persistent messaging)
|
||||
await client.js_create_stream("EVENTS", subjects=["events.*"])
|
||||
await client.js_publish("events.user", b"data")
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[PubSubConfig] = None) -> None:
|
||||
"""Initialize pub/sub client.
|
||||
|
||||
Args:
|
||||
config: NATS configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or PubSubConfig.from_env()
|
||||
self._nc: Any = None
|
||||
self._js: Any = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection to NATS server."""
|
||||
try:
|
||||
import nats
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"nats-py is required for PubSubClient. "
|
||||
"Install with: pip install nats-py"
|
||||
) from e
|
||||
|
||||
connect_opts: dict[str, Any] = {
|
||||
"servers": self.config.servers,
|
||||
"name": self.config.name,
|
||||
"connect_timeout": self.config.connect_timeout,
|
||||
"reconnect_time_wait": self.config.reconnect_time_wait,
|
||||
"max_reconnect_attempts": self.config.max_reconnect_attempts,
|
||||
"ping_interval": self.config.ping_interval,
|
||||
"max_outstanding_pings": self.config.max_outstanding_pings,
|
||||
}
|
||||
|
||||
if self.config.user and self.config.password:
|
||||
connect_opts["user"] = self.config.user
|
||||
connect_opts["password"] = self.config.password
|
||||
elif self.config.token:
|
||||
connect_opts["token"] = self.config.token
|
||||
elif self.config.nkey_seed:
|
||||
connect_opts["nkeys_seed"] = self.config.nkey_seed
|
||||
|
||||
self._nc = await nats.connect(**connect_opts)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
if self._nc:
|
||||
await self._nc.drain()
|
||||
self._nc = None
|
||||
self._js = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if NATS connection is healthy.
|
||||
|
||||
Returns:
|
||||
True if connected and responding.
|
||||
"""
|
||||
if not self._nc:
|
||||
return False
|
||||
return self._nc.is_connected
|
||||
|
||||
# Core pub/sub operations
|
||||
|
||||
async def publish(
|
||||
self,
|
||||
subject: str,
|
||||
data: bytes,
|
||||
reply: Optional[str] = None,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
) -> None:
|
||||
"""Publish a message to a subject.
|
||||
|
||||
Args:
|
||||
subject: Subject to publish to.
|
||||
data: Message payload.
|
||||
reply: Optional reply subject.
|
||||
headers: Optional message headers.
|
||||
"""
|
||||
await self._nc.publish(subject, data, reply=reply, headers=headers)
|
||||
|
||||
async def subscribe(
|
||||
self,
|
||||
subject: str,
|
||||
handler: MessageHandler,
|
||||
queue: Optional[str] = None,
|
||||
) -> Subscription:
|
||||
"""Subscribe to a subject.
|
||||
|
||||
Args:
|
||||
subject: Subject pattern (supports wildcards * and >).
|
||||
handler: Async callback for received messages.
|
||||
queue: Queue group for load balancing.
|
||||
|
||||
Returns:
|
||||
Subscription handle for unsubscribing.
|
||||
"""
|
||||
async def _wrapper(msg: Any) -> None:
|
||||
wrapped = Message(
|
||||
subject=msg.subject,
|
||||
data=msg.data,
|
||||
reply=msg.reply,
|
||||
headers=dict(msg.headers) if msg.headers else {},
|
||||
)
|
||||
await handler(wrapped)
|
||||
|
||||
sub = await self._nc.subscribe(subject, cb=_wrapper, queue=queue or "")
|
||||
return Subscription(subject=subject, queue=queue, _sub=sub)
|
||||
|
||||
async def subscribe_iter(
|
||||
self,
|
||||
subject: str,
|
||||
queue: Optional[str] = None,
|
||||
) -> AsyncIterator[Message]:
|
||||
"""Subscribe to a subject with async iteration.
|
||||
|
||||
Args:
|
||||
subject: Subject pattern.
|
||||
queue: Queue group for load balancing.
|
||||
|
||||
Yields:
|
||||
Messages as they arrive.
|
||||
"""
|
||||
sub = await self._nc.subscribe(subject, queue=queue or "")
|
||||
try:
|
||||
async for msg in sub.messages:
|
||||
yield Message(
|
||||
subject=msg.subject,
|
||||
data=msg.data,
|
||||
reply=msg.reply,
|
||||
headers=dict(msg.headers) if msg.headers else {},
|
||||
)
|
||||
finally:
|
||||
await sub.unsubscribe()
|
||||
|
||||
# Request/reply pattern
|
||||
|
||||
async def request(
|
||||
self,
|
||||
subject: str,
|
||||
data: bytes,
|
||||
timeout: float = 5.0,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
) -> Message:
|
||||
"""Send a request and wait for a response.
|
||||
|
||||
Args:
|
||||
subject: Subject to send request to.
|
||||
data: Request payload.
|
||||
timeout: Response timeout in seconds.
|
||||
headers: Optional request headers.
|
||||
|
||||
Returns:
|
||||
Response message.
|
||||
|
||||
Raises:
|
||||
TimeoutError: If no response within timeout.
|
||||
"""
|
||||
response = await self._nc.request(subject, data, timeout=timeout, headers=headers)
|
||||
return Message(
|
||||
subject=response.subject,
|
||||
data=response.data,
|
||||
reply=response.reply,
|
||||
headers=dict(response.headers) if response.headers else {},
|
||||
)
|
||||
|
||||
async def respond(
|
||||
self,
|
||||
request: Message,
|
||||
data: bytes,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
) -> None:
|
||||
"""Send a response to a request message.
|
||||
|
||||
Args:
|
||||
request: Original request message.
|
||||
data: Response payload.
|
||||
headers: Optional response headers.
|
||||
"""
|
||||
if request.reply:
|
||||
await self.publish(request.reply, data, headers=headers)
|
||||
|
||||
# JetStream operations
|
||||
|
||||
def _get_js(self) -> Any:
|
||||
"""Get JetStream context, creating if needed."""
|
||||
if self._js is None:
|
||||
self._js = self._nc.jetstream()
|
||||
return self._js
|
||||
|
||||
async def js_create_stream(
|
||||
self,
|
||||
name: str,
|
||||
subjects: Sequence[str],
|
||||
storage: str = "file",
|
||||
retention: str = "limits",
|
||||
max_msgs: int = -1,
|
||||
max_bytes: int = -1,
|
||||
max_age: int = 0, # seconds, 0 = unlimited
|
||||
max_msg_size: int = -1,
|
||||
duplicate_window: int = 120, # seconds
|
||||
replicas: int = 1,
|
||||
) -> None:
|
||||
"""Create a JetStream stream.
|
||||
|
||||
Args:
|
||||
name: Stream name.
|
||||
subjects: Subjects to capture.
|
||||
storage: Storage type (file, memory).
|
||||
retention: Retention policy (limits, interest, workqueue).
|
||||
max_msgs: Max messages (-1 = unlimited).
|
||||
max_bytes: Max bytes (-1 = unlimited).
|
||||
max_age: Max age in seconds (0 = unlimited).
|
||||
max_msg_size: Max message size (-1 = unlimited).
|
||||
duplicate_window: Duplicate detection window in seconds.
|
||||
replicas: Number of replicas.
|
||||
"""
|
||||
from nats.js.api import (
|
||||
StreamConfig,
|
||||
StorageType,
|
||||
RetentionPolicy,
|
||||
)
|
||||
|
||||
storage_map = {
|
||||
"file": StorageType.FILE,
|
||||
"memory": StorageType.MEMORY,
|
||||
}
|
||||
retention_map = {
|
||||
"limits": RetentionPolicy.LIMITS,
|
||||
"interest": RetentionPolicy.INTEREST,
|
||||
"workqueue": RetentionPolicy.WORK_QUEUE,
|
||||
}
|
||||
|
||||
config = StreamConfig(
|
||||
name=name,
|
||||
subjects=list(subjects),
|
||||
storage=storage_map.get(storage, StorageType.FILE),
|
||||
retention=retention_map.get(retention, RetentionPolicy.LIMITS),
|
||||
max_msgs=max_msgs,
|
||||
max_bytes=max_bytes,
|
||||
max_age=max_age * 1_000_000_000 if max_age else 0, # nanoseconds
|
||||
max_msg_size=max_msg_size,
|
||||
duplicate_window=duplicate_window * 1_000_000_000, # nanoseconds
|
||||
num_replicas=replicas,
|
||||
)
|
||||
|
||||
js = self._get_js()
|
||||
await js.add_stream(config)
|
||||
|
||||
async def js_delete_stream(self, name: str) -> None:
|
||||
"""Delete a JetStream stream.
|
||||
|
||||
Args:
|
||||
name: Stream name.
|
||||
"""
|
||||
js = self._get_js()
|
||||
await js.delete_stream(name)
|
||||
|
||||
async def js_stream_info(self, name: str) -> dict[str, Any]:
|
||||
"""Get stream information.
|
||||
|
||||
Args:
|
||||
name: Stream name.
|
||||
|
||||
Returns:
|
||||
Stream info dict.
|
||||
"""
|
||||
js = self._get_js()
|
||||
info = await js.stream_info(name)
|
||||
return {
|
||||
"name": info.config.name,
|
||||
"subjects": info.config.subjects,
|
||||
"messages": info.state.messages,
|
||||
"bytes": info.state.bytes,
|
||||
"first_seq": info.state.first_seq,
|
||||
"last_seq": info.state.last_seq,
|
||||
}
|
||||
|
||||
async def js_publish(
|
||||
self,
|
||||
subject: str,
|
||||
data: bytes,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
msg_id: Optional[str] = None,
|
||||
expect_stream: Optional[str] = None,
|
||||
) -> int:
|
||||
"""Publish a message to JetStream.
|
||||
|
||||
Args:
|
||||
subject: Subject to publish to.
|
||||
data: Message payload.
|
||||
headers: Optional message headers.
|
||||
msg_id: Message ID for deduplication.
|
||||
expect_stream: Expected stream name.
|
||||
|
||||
Returns:
|
||||
Sequence number of published message.
|
||||
"""
|
||||
js = self._get_js()
|
||||
ack = await js.publish(
|
||||
subject,
|
||||
data,
|
||||
headers=headers,
|
||||
)
|
||||
return ack.seq
|
||||
|
||||
async def js_create_consumer(
|
||||
self,
|
||||
stream: str,
|
||||
name: str,
|
||||
durable: bool = True,
|
||||
filter_subjects: Optional[Sequence[str]] = None,
|
||||
ack_policy: str = "explicit",
|
||||
max_deliver: int = -1,
|
||||
ack_wait: int = 30, # seconds
|
||||
) -> None:
|
||||
"""Create a JetStream consumer.
|
||||
|
||||
Args:
|
||||
stream: Stream name.
|
||||
name: Consumer name.
|
||||
durable: Durable consumer (survives disconnects).
|
||||
filter_subjects: Filter to specific subjects.
|
||||
ack_policy: Ack policy (none, all, explicit).
|
||||
max_deliver: Max redelivery attempts (-1 = unlimited).
|
||||
ack_wait: Ack wait time in seconds.
|
||||
"""
|
||||
from nats.js.api import ConsumerConfig, AckPolicy
|
||||
|
||||
ack_map = {
|
||||
"none": AckPolicy.NONE,
|
||||
"all": AckPolicy.ALL,
|
||||
"explicit": AckPolicy.EXPLICIT,
|
||||
}
|
||||
|
||||
config = ConsumerConfig(
|
||||
durable_name=name if durable else None,
|
||||
name=name,
|
||||
filter_subjects=list(filter_subjects) if filter_subjects else None,
|
||||
ack_policy=ack_map.get(ack_policy, AckPolicy.EXPLICIT),
|
||||
max_deliver=max_deliver,
|
||||
ack_wait=ack_wait * 1_000_000_000, # nanoseconds
|
||||
)
|
||||
|
||||
js = self._get_js()
|
||||
await js.add_consumer(stream, config)
|
||||
|
||||
async def js_subscribe(
|
||||
self,
|
||||
stream: str,
|
||||
consumer: str,
|
||||
handler: MessageHandler,
|
||||
) -> Subscription:
|
||||
"""Subscribe to a JetStream consumer.
|
||||
|
||||
Args:
|
||||
stream: Stream name.
|
||||
consumer: Consumer name.
|
||||
handler: Message handler.
|
||||
|
||||
Returns:
|
||||
Subscription handle.
|
||||
"""
|
||||
async def _wrapper(msg: Any) -> None:
|
||||
wrapped = Message(
|
||||
subject=msg.subject,
|
||||
data=msg.data,
|
||||
reply=msg.reply,
|
||||
headers=dict(msg.headers) if msg.headers else {},
|
||||
)
|
||||
await handler(wrapped)
|
||||
await msg.ack()
|
||||
|
||||
js = self._get_js()
|
||||
sub = await js.pull_subscribe(durable=consumer, stream=stream)
|
||||
# Note: Pull subscribe works differently, using fetch
|
||||
return Subscription(subject=f"{stream}.{consumer}", _sub=sub)
|
||||
|
||||
async def js_fetch(
|
||||
self,
|
||||
stream: str,
|
||||
consumer: str,
|
||||
batch: int = 1,
|
||||
timeout: float = 5.0,
|
||||
) -> list[Message]:
|
||||
"""Fetch messages from a JetStream consumer.
|
||||
|
||||
Args:
|
||||
stream: Stream name.
|
||||
consumer: Consumer name.
|
||||
batch: Number of messages to fetch.
|
||||
timeout: Fetch timeout in seconds.
|
||||
|
||||
Returns:
|
||||
List of messages.
|
||||
"""
|
||||
js = self._get_js()
|
||||
sub = await js.pull_subscribe(durable=consumer, stream=stream)
|
||||
|
||||
try:
|
||||
msgs = await sub.fetch(batch, timeout=timeout)
|
||||
result = []
|
||||
for msg in msgs:
|
||||
result.append(Message(
|
||||
subject=msg.subject,
|
||||
data=msg.data,
|
||||
reply=msg.reply,
|
||||
headers=dict(msg.headers) if msg.headers else {},
|
||||
))
|
||||
await msg.ack()
|
||||
return result
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
async def __aenter__(self) -> PubSubClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,637 @@
|
||||
"""Work queue client for Hanzo infrastructure.
|
||||
|
||||
Provides async interface for distributed work queues,
|
||||
built on top of Redis/Valkey for simplicity and reliability.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
from enum import Enum
|
||||
from typing import Any, Awaitable, Callable, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class QueuesConfig(BaseModel):
|
||||
"""Configuration for work queue connection."""
|
||||
|
||||
host: str = Field(default="localhost", description="Redis server host")
|
||||
port: int = Field(default=6379, description="Redis server port")
|
||||
password: Optional[str] = Field(default=None, description="Redis password")
|
||||
db: int = Field(default=1, description="Redis database number")
|
||||
url: Optional[str] = Field(default=None, description="Full Redis URL")
|
||||
prefix: str = Field(default="hanzo:queue", description="Key prefix for queues")
|
||||
default_timeout: int = Field(default=300, description="Default job timeout in seconds")
|
||||
default_ttl: int = Field(default=86400, description="Default result TTL in seconds")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> QueuesConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
QUEUE_REDIS_HOST: Server host (default: localhost)
|
||||
QUEUE_REDIS_PORT: Server port (default: 6379)
|
||||
QUEUE_REDIS_PASSWORD: Authentication password
|
||||
QUEUE_REDIS_DB: Database number (default: 1)
|
||||
QUEUE_REDIS_URL: Full URL (overrides host/port)
|
||||
QUEUE_PREFIX: Key prefix (default: hanzo:queue)
|
||||
"""
|
||||
return cls(
|
||||
host=os.getenv("QUEUE_REDIS_HOST") or os.getenv("REDIS_HOST", "localhost"),
|
||||
port=int(os.getenv("QUEUE_REDIS_PORT") or os.getenv("REDIS_PORT", "6379")),
|
||||
password=os.getenv("QUEUE_REDIS_PASSWORD") or os.getenv("REDIS_PASSWORD"),
|
||||
db=int(os.getenv("QUEUE_REDIS_DB", "1")),
|
||||
url=os.getenv("QUEUE_REDIS_URL") or os.getenv("REDIS_URL"),
|
||||
prefix=os.getenv("QUEUE_PREFIX", "hanzo:queue"),
|
||||
)
|
||||
|
||||
|
||||
class JobStatus(str, Enum):
|
||||
"""Status of a queued job."""
|
||||
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
RETRYING = "retrying"
|
||||
CANCELLED = "cancelled"
|
||||
EXPIRED = "expired"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Job:
|
||||
"""A work queue job."""
|
||||
|
||||
id: str
|
||||
queue: str
|
||||
name: str
|
||||
args: dict[str, Any] = field(default_factory=dict)
|
||||
status: JobStatus = JobStatus.PENDING
|
||||
result: Any = None
|
||||
error: Optional[str] = None
|
||||
created_at: Optional[datetime] = None
|
||||
started_at: Optional[datetime] = None
|
||||
completed_at: Optional[datetime] = None
|
||||
attempts: int = 0
|
||||
max_attempts: int = 3
|
||||
timeout: int = 300 # seconds
|
||||
ttl: int = 86400 # result TTL in seconds
|
||||
priority: int = 0 # higher = more priority
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dictionary for storage."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"queue": self.queue,
|
||||
"name": self.name,
|
||||
"args": self.args,
|
||||
"status": self.status.value,
|
||||
"result": self.result,
|
||||
"error": self.error,
|
||||
"created_at": self.created_at.isoformat() if self.created_at else None,
|
||||
"started_at": self.started_at.isoformat() if self.started_at else None,
|
||||
"completed_at": self.completed_at.isoformat() if self.completed_at else None,
|
||||
"attempts": self.attempts,
|
||||
"max_attempts": self.max_attempts,
|
||||
"timeout": self.timeout,
|
||||
"ttl": self.ttl,
|
||||
"priority": self.priority,
|
||||
"metadata": self.metadata,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> Job:
|
||||
"""Create from dictionary."""
|
||||
return cls(
|
||||
id=data["id"],
|
||||
queue=data["queue"],
|
||||
name=data["name"],
|
||||
args=data.get("args", {}),
|
||||
status=JobStatus(data.get("status", "pending")),
|
||||
result=data.get("result"),
|
||||
error=data.get("error"),
|
||||
created_at=datetime.fromisoformat(data["created_at"]) if data.get("created_at") else None,
|
||||
started_at=datetime.fromisoformat(data["started_at"]) if data.get("started_at") else None,
|
||||
completed_at=datetime.fromisoformat(data["completed_at"]) if data.get("completed_at") else None,
|
||||
attempts=data.get("attempts", 0),
|
||||
max_attempts=data.get("max_attempts", 3),
|
||||
timeout=data.get("timeout", 300),
|
||||
ttl=data.get("ttl", 86400),
|
||||
priority=data.get("priority", 0),
|
||||
metadata=data.get("metadata", {}),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueueStats:
|
||||
"""Statistics for a queue."""
|
||||
|
||||
name: str
|
||||
pending: int = 0
|
||||
running: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
|
||||
|
||||
JobHandler = Callable[[Job], Awaitable[Any]]
|
||||
|
||||
|
||||
class QueuesClient:
|
||||
"""Async client for distributed work queues.
|
||||
|
||||
Implements a simple but reliable work queue on top of Redis,
|
||||
with support for priorities, retries, and result storage.
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = QueuesClient(QueuesConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Enqueue jobs
|
||||
job_id = await client.enqueue(
|
||||
"emails",
|
||||
"send_welcome",
|
||||
args={"user_id": "123", "email": "[email protected]"},
|
||||
priority=10,
|
||||
)
|
||||
|
||||
# Get job status
|
||||
job = await client.get_job(job_id)
|
||||
print(f"Status: {job.status}")
|
||||
|
||||
# Process jobs (in a worker)
|
||||
async def handler(job: Job) -> str:
|
||||
# Do work...
|
||||
return "sent"
|
||||
|
||||
await client.process("emails", handler)
|
||||
|
||||
# Get result
|
||||
result = await client.get_result(job_id)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[QueuesConfig] = None) -> None:
|
||||
"""Initialize queues client.
|
||||
|
||||
Args:
|
||||
config: Queue configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or QueuesConfig.from_env()
|
||||
self._redis: Any = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection to Redis."""
|
||||
try:
|
||||
import redis.asyncio as redis
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"redis is required for QueuesClient. "
|
||||
"Install with: pip install redis"
|
||||
) from e
|
||||
|
||||
if self.config.url:
|
||||
self._redis = redis.from_url(
|
||||
self.config.url,
|
||||
decode_responses=True,
|
||||
)
|
||||
else:
|
||||
self._redis = redis.Redis(
|
||||
host=self.config.host,
|
||||
port=self.config.port,
|
||||
password=self.config.password,
|
||||
db=self.config.db,
|
||||
decode_responses=True,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
if self._redis:
|
||||
await self._redis.aclose()
|
||||
self._redis = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if Redis is healthy.
|
||||
|
||||
Returns:
|
||||
True if Redis is responding.
|
||||
"""
|
||||
if not self._redis:
|
||||
return False
|
||||
try:
|
||||
await self._redis.ping()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _key(self, *parts: str) -> str:
|
||||
"""Build a Redis key with prefix."""
|
||||
return ":".join([self.config.prefix, *parts])
|
||||
|
||||
# Job operations
|
||||
|
||||
async def enqueue(
|
||||
self,
|
||||
queue: str,
|
||||
name: str,
|
||||
args: Optional[dict[str, Any]] = None,
|
||||
job_id: Optional[str] = None,
|
||||
priority: int = 0,
|
||||
delay: Optional[timedelta] = None,
|
||||
timeout: Optional[int] = None,
|
||||
max_attempts: int = 3,
|
||||
ttl: Optional[int] = None,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
) -> str:
|
||||
"""Enqueue a job.
|
||||
|
||||
Args:
|
||||
queue: Queue name.
|
||||
name: Job name/type.
|
||||
args: Job arguments.
|
||||
job_id: Custom job ID (auto-generated if not specified).
|
||||
priority: Job priority (higher = more priority).
|
||||
delay: Delay before job becomes available.
|
||||
timeout: Job timeout in seconds.
|
||||
max_attempts: Maximum retry attempts.
|
||||
ttl: Result TTL in seconds.
|
||||
metadata: Additional metadata.
|
||||
|
||||
Returns:
|
||||
Job ID.
|
||||
"""
|
||||
job = Job(
|
||||
id=job_id or f"{name}-{uuid.uuid4().hex[:12]}",
|
||||
queue=queue,
|
||||
name=name,
|
||||
args=args or {},
|
||||
status=JobStatus.PENDING,
|
||||
created_at=datetime.utcnow(),
|
||||
max_attempts=max_attempts,
|
||||
timeout=timeout or self.config.default_timeout,
|
||||
ttl=ttl or self.config.default_ttl,
|
||||
priority=priority,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
|
||||
# Store job data
|
||||
job_key = self._key("job", job.id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
|
||||
# Add to queue
|
||||
queue_key = self._key("queue", queue)
|
||||
score = -priority # Redis ZSET is ascending, we want descending priority
|
||||
|
||||
if delay:
|
||||
# Delayed job - use scheduled set
|
||||
scheduled_key = self._key("scheduled", queue)
|
||||
execute_at = datetime.utcnow() + delay
|
||||
await self._redis.zadd(scheduled_key, {job.id: execute_at.timestamp()})
|
||||
else:
|
||||
# Immediate job
|
||||
await self._redis.zadd(queue_key, {job.id: score})
|
||||
|
||||
return job.id
|
||||
|
||||
async def get_job(self, job_id: str) -> Optional[Job]:
|
||||
"""Get a job by ID.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
|
||||
Returns:
|
||||
Job or None if not found.
|
||||
"""
|
||||
job_key = self._key("job", job_id)
|
||||
data = await self._redis.get(job_key)
|
||||
if not data:
|
||||
return None
|
||||
return Job.from_dict(json.loads(data))
|
||||
|
||||
async def get_result(self, job_id: str, timeout: float = 0) -> Any:
|
||||
"""Get a job's result.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
timeout: Seconds to wait for result (0 = don't wait).
|
||||
|
||||
Returns:
|
||||
Job result or None.
|
||||
"""
|
||||
job = await self.get_job(job_id)
|
||||
if not job:
|
||||
return None
|
||||
|
||||
if job.status == JobStatus.COMPLETED:
|
||||
return job.result
|
||||
|
||||
if timeout > 0 and job.status in (JobStatus.PENDING, JobStatus.RUNNING):
|
||||
# Poll for completion
|
||||
end_time = datetime.utcnow() + timedelta(seconds=timeout)
|
||||
while datetime.utcnow() < end_time:
|
||||
await asyncio.sleep(0.5)
|
||||
job = await self.get_job(job_id)
|
||||
if job and job.status == JobStatus.COMPLETED:
|
||||
return job.result
|
||||
if job and job.status == JobStatus.FAILED:
|
||||
raise Exception(f"Job failed: {job.error}")
|
||||
|
||||
return None
|
||||
|
||||
async def cancel_job(self, job_id: str) -> bool:
|
||||
"""Cancel a pending job.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
|
||||
Returns:
|
||||
True if cancelled, False if job not found or already processed.
|
||||
"""
|
||||
job = await self.get_job(job_id)
|
||||
if not job:
|
||||
return False
|
||||
|
||||
if job.status not in (JobStatus.PENDING, JobStatus.RETRYING):
|
||||
return False
|
||||
|
||||
# Remove from queue
|
||||
queue_key = self._key("queue", job.queue)
|
||||
await self._redis.zrem(queue_key, job_id)
|
||||
|
||||
# Update status
|
||||
job.status = JobStatus.CANCELLED
|
||||
job_key = self._key("job", job_id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
|
||||
return True
|
||||
|
||||
async def retry_job(self, job_id: str) -> bool:
|
||||
"""Retry a failed job.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
|
||||
Returns:
|
||||
True if job was requeued.
|
||||
"""
|
||||
job = await self.get_job(job_id)
|
||||
if not job:
|
||||
return False
|
||||
|
||||
if job.status != JobStatus.FAILED:
|
||||
return False
|
||||
|
||||
# Reset and requeue
|
||||
job.status = JobStatus.RETRYING
|
||||
job.error = None
|
||||
job.started_at = None
|
||||
job.completed_at = None
|
||||
|
||||
job_key = self._key("job", job_id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
|
||||
queue_key = self._key("queue", job.queue)
|
||||
await self._redis.zadd(queue_key, {job_id: -job.priority})
|
||||
|
||||
return True
|
||||
|
||||
# Queue operations
|
||||
|
||||
async def get_queue_stats(self, queue: str) -> QueueStats:
|
||||
"""Get queue statistics.
|
||||
|
||||
Args:
|
||||
queue: Queue name.
|
||||
|
||||
Returns:
|
||||
Queue statistics.
|
||||
"""
|
||||
queue_key = self._key("queue", queue)
|
||||
running_key = self._key("running", queue)
|
||||
|
||||
pending = await self._redis.zcard(queue_key)
|
||||
running = await self._redis.scard(running_key)
|
||||
|
||||
# Count completed/failed from recent jobs (expensive, use sparingly)
|
||||
completed = 0
|
||||
failed = 0
|
||||
|
||||
return QueueStats(
|
||||
name=queue,
|
||||
pending=pending,
|
||||
running=running,
|
||||
completed=completed,
|
||||
failed=failed,
|
||||
)
|
||||
|
||||
async def list_queues(self) -> list[str]:
|
||||
"""List all queues.
|
||||
|
||||
Returns:
|
||||
List of queue names.
|
||||
"""
|
||||
pattern = self._key("queue", "*")
|
||||
keys = await self._redis.keys(pattern)
|
||||
prefix_len = len(self._key("queue", ""))
|
||||
return [k[prefix_len:] for k in keys]
|
||||
|
||||
async def purge_queue(self, queue: str) -> int:
|
||||
"""Remove all jobs from a queue.
|
||||
|
||||
Args:
|
||||
queue: Queue name.
|
||||
|
||||
Returns:
|
||||
Number of jobs removed.
|
||||
"""
|
||||
queue_key = self._key("queue", queue)
|
||||
count = await self._redis.zcard(queue_key)
|
||||
await self._redis.delete(queue_key)
|
||||
return count
|
||||
|
||||
# Worker operations
|
||||
|
||||
async def dequeue(
|
||||
self,
|
||||
queue: str,
|
||||
timeout: float = 0,
|
||||
) -> Optional[Job]:
|
||||
"""Dequeue a job for processing.
|
||||
|
||||
Args:
|
||||
queue: Queue name.
|
||||
timeout: Seconds to wait for a job (0 = don't wait).
|
||||
|
||||
Returns:
|
||||
Job or None if queue is empty.
|
||||
"""
|
||||
queue_key = self._key("queue", queue)
|
||||
running_key = self._key("running", queue)
|
||||
|
||||
# Move scheduled jobs that are ready
|
||||
await self._move_scheduled_jobs(queue)
|
||||
|
||||
# Try to get a job
|
||||
result = await self._redis.zpopmin(queue_key, 1)
|
||||
if not result:
|
||||
if timeout > 0:
|
||||
# Wait for job with blocking pop
|
||||
result = await self._redis.bzpopmin(queue_key, timeout)
|
||||
if result:
|
||||
_, job_id, _ = result
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
job_id = result[0][0]
|
||||
|
||||
# Get job data
|
||||
job = await self.get_job(job_id)
|
||||
if not job:
|
||||
return None
|
||||
|
||||
# Mark as running
|
||||
job.status = JobStatus.RUNNING
|
||||
job.started_at = datetime.utcnow()
|
||||
job.attempts += 1
|
||||
|
||||
job_key = self._key("job", job_id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
|
||||
# Add to running set
|
||||
await self._redis.sadd(running_key, job_id)
|
||||
|
||||
return job
|
||||
|
||||
async def _move_scheduled_jobs(self, queue: str) -> None:
|
||||
"""Move scheduled jobs that are ready to the main queue."""
|
||||
scheduled_key = self._key("scheduled", queue)
|
||||
queue_key = self._key("queue", queue)
|
||||
|
||||
now = datetime.utcnow().timestamp()
|
||||
|
||||
# Get jobs that are ready
|
||||
ready = await self._redis.zrangebyscore(scheduled_key, "-inf", now)
|
||||
|
||||
for job_id in ready:
|
||||
job = await self.get_job(job_id)
|
||||
if job:
|
||||
# Move to main queue
|
||||
await self._redis.zrem(scheduled_key, job_id)
|
||||
await self._redis.zadd(queue_key, {job_id: -job.priority})
|
||||
|
||||
async def complete_job(
|
||||
self,
|
||||
job_id: str,
|
||||
result: Any = None,
|
||||
) -> None:
|
||||
"""Mark a job as completed.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
result: Job result.
|
||||
"""
|
||||
job = await self.get_job(job_id)
|
||||
if not job:
|
||||
return
|
||||
|
||||
job.status = JobStatus.COMPLETED
|
||||
job.result = result
|
||||
job.completed_at = datetime.utcnow()
|
||||
|
||||
# Update job
|
||||
job_key = self._key("job", job_id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
await self._redis.expire(job_key, job.ttl)
|
||||
|
||||
# Remove from running
|
||||
running_key = self._key("running", job.queue)
|
||||
await self._redis.srem(running_key, job_id)
|
||||
|
||||
async def fail_job(
|
||||
self,
|
||||
job_id: str,
|
||||
error: str,
|
||||
) -> None:
|
||||
"""Mark a job as failed.
|
||||
|
||||
Args:
|
||||
job_id: Job ID.
|
||||
error: Error message.
|
||||
"""
|
||||
job = await self.get_job(job_id)
|
||||
if not job:
|
||||
return
|
||||
|
||||
# Check if should retry
|
||||
if job.attempts < job.max_attempts:
|
||||
job.status = JobStatus.RETRYING
|
||||
job.error = error
|
||||
|
||||
# Requeue with backoff
|
||||
delay = min(60 * (2 ** job.attempts), 3600) # Exponential backoff, max 1 hour
|
||||
queue_key = self._key("queue", job.queue)
|
||||
scheduled_key = self._key("scheduled", job.queue)
|
||||
execute_at = datetime.utcnow() + timedelta(seconds=delay)
|
||||
await self._redis.zadd(scheduled_key, {job_id: execute_at.timestamp()})
|
||||
else:
|
||||
job.status = JobStatus.FAILED
|
||||
job.error = error
|
||||
job.completed_at = datetime.utcnow()
|
||||
|
||||
# Update job
|
||||
job_key = self._key("job", job_id)
|
||||
await self._redis.set(job_key, json.dumps(job.to_dict()))
|
||||
|
||||
# Remove from running
|
||||
running_key = self._key("running", job.queue)
|
||||
await self._redis.srem(running_key, job_id)
|
||||
|
||||
async def process(
|
||||
self,
|
||||
queue: str,
|
||||
handler: JobHandler,
|
||||
batch_size: int = 1,
|
||||
poll_interval: float = 1.0,
|
||||
) -> None:
|
||||
"""Process jobs from a queue (blocking).
|
||||
|
||||
Args:
|
||||
queue: Queue name.
|
||||
handler: Async job handler function.
|
||||
batch_size: Jobs to process before checking for more.
|
||||
poll_interval: Seconds between queue checks.
|
||||
"""
|
||||
while True:
|
||||
for _ in range(batch_size):
|
||||
job = await self.dequeue(queue, timeout=poll_interval)
|
||||
if not job:
|
||||
break
|
||||
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
handler(job),
|
||||
timeout=job.timeout,
|
||||
)
|
||||
await self.complete_job(job.id, result)
|
||||
except asyncio.TimeoutError:
|
||||
await self.fail_job(job.id, "Job timed out")
|
||||
except Exception as e:
|
||||
await self.fail_job(job.id, str(e))
|
||||
|
||||
async def __aenter__(self) -> QueuesClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,510 @@
|
||||
"""Meilisearch client wrapper for Hanzo infrastructure.
|
||||
|
||||
Provides async interface to Meilisearch for full-text search
|
||||
with typo tolerance, filtering, and faceted search.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional, Sequence
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SearchConfig(BaseModel):
|
||||
"""Configuration for Meilisearch connection."""
|
||||
|
||||
host: str = Field(default="localhost", description="Meilisearch server host")
|
||||
port: int = Field(default=7700, description="Meilisearch server port")
|
||||
api_key: Optional[str] = Field(default=None, description="Master or API key")
|
||||
url: Optional[str] = Field(default=None, description="Full URL (overrides host/port)")
|
||||
timeout: float = Field(default=30.0, description="Request timeout in seconds")
|
||||
ssl: bool = Field(default=False, description="Use HTTPS")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> SearchConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
MEILISEARCH_HOST: Server host (default: localhost)
|
||||
MEILISEARCH_PORT: Server port (default: 7700)
|
||||
MEILISEARCH_API_KEY: API key for authentication
|
||||
MEILISEARCH_URL: Full URL (overrides host/port)
|
||||
MEILISEARCH_SSL: Use HTTPS (default: false)
|
||||
"""
|
||||
return cls(
|
||||
host=os.getenv("MEILISEARCH_HOST", "localhost"),
|
||||
port=int(os.getenv("MEILISEARCH_PORT", "7700")),
|
||||
api_key=os.getenv("MEILISEARCH_API_KEY") or os.getenv("MEILI_MASTER_KEY"),
|
||||
url=os.getenv("MEILISEARCH_URL"),
|
||||
ssl=os.getenv("MEILISEARCH_SSL", "").lower() in ("true", "1", "yes"),
|
||||
)
|
||||
|
||||
@property
|
||||
def effective_url(self) -> str:
|
||||
"""Get the effective URL to connect to."""
|
||||
if self.url:
|
||||
return self.url
|
||||
protocol = "https" if self.ssl else "http"
|
||||
return f"{protocol}://{self.host}:{self.port}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchHit:
|
||||
"""A single search result."""
|
||||
|
||||
id: str
|
||||
document: dict[str, Any]
|
||||
score: Optional[float] = None
|
||||
highlights: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
"""Result of a search query."""
|
||||
|
||||
hits: list[SearchHit]
|
||||
query: str
|
||||
processing_time_ms: int
|
||||
total_hits: int
|
||||
offset: int = 0
|
||||
limit: int = 20
|
||||
facet_distribution: dict[str, dict[str, int]] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class IndexStats:
|
||||
"""Statistics for an index."""
|
||||
|
||||
number_of_documents: int
|
||||
is_indexing: bool
|
||||
field_distribution: dict[str, int] = field(default_factory=dict)
|
||||
|
||||
|
||||
class SearchClient:
|
||||
"""Async client for Meilisearch full-text search.
|
||||
|
||||
Wraps meilisearch-python-sdk with async methods and Hanzo conventions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = SearchClient(SearchConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Create index
|
||||
await client.create_index("products", primary_key="id")
|
||||
|
||||
# Add documents
|
||||
docs = [{"id": "1", "name": "Widget", "price": 99}]
|
||||
await client.add_documents("products", docs)
|
||||
|
||||
# Search
|
||||
results = await client.search("products", "widget", filters="price > 50")
|
||||
for hit in results.hits:
|
||||
print(hit.document["name"])
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[SearchConfig] = None) -> None:
|
||||
"""Initialize search client.
|
||||
|
||||
Args:
|
||||
config: Meilisearch configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or SearchConfig.from_env()
|
||||
self._client: Any = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection to Meilisearch server."""
|
||||
try:
|
||||
from meilisearch_python_sdk import AsyncClient
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"meilisearch-python-sdk is required for SearchClient. "
|
||||
"Install with: pip install meilisearch-python-sdk"
|
||||
) from e
|
||||
|
||||
self._client = AsyncClient(
|
||||
self.config.effective_url,
|
||||
api_key=self.config.api_key,
|
||||
timeout=int(self.config.timeout),
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if Meilisearch server is healthy.
|
||||
|
||||
Returns:
|
||||
True if server is reachable and healthy.
|
||||
"""
|
||||
if not self._client:
|
||||
return False
|
||||
try:
|
||||
health = await self._client.health()
|
||||
return health.status == "available"
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# Index operations
|
||||
|
||||
async def create_index(
|
||||
self,
|
||||
name: str,
|
||||
primary_key: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Create a new index.
|
||||
|
||||
Args:
|
||||
name: Index name (uid).
|
||||
primary_key: Document primary key field.
|
||||
"""
|
||||
task = await self._client.create_index(uid=name, primary_key=primary_key)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def delete_index(self, name: str) -> None:
|
||||
"""Delete an index.
|
||||
|
||||
Args:
|
||||
name: Index name to delete.
|
||||
"""
|
||||
task = await self._client.index(name).delete()
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def index_exists(self, name: str) -> bool:
|
||||
"""Check if an index exists.
|
||||
|
||||
Args:
|
||||
name: Index name.
|
||||
|
||||
Returns:
|
||||
True if index exists.
|
||||
"""
|
||||
try:
|
||||
await self._client.get_index(name)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def list_indexes(self) -> list[str]:
|
||||
"""List all indexes.
|
||||
|
||||
Returns:
|
||||
List of index names.
|
||||
"""
|
||||
indexes = await self._client.get_indexes()
|
||||
return [idx.uid for idx in indexes]
|
||||
|
||||
async def get_index_stats(self, name: str) -> IndexStats:
|
||||
"""Get statistics for an index.
|
||||
|
||||
Args:
|
||||
name: Index name.
|
||||
|
||||
Returns:
|
||||
Index statistics.
|
||||
"""
|
||||
index = self._client.index(name)
|
||||
stats = await index.get_stats()
|
||||
return IndexStats(
|
||||
number_of_documents=stats.number_of_documents,
|
||||
is_indexing=stats.is_indexing,
|
||||
field_distribution=stats.field_distribution or {},
|
||||
)
|
||||
|
||||
# Document operations
|
||||
|
||||
async def add_documents(
|
||||
self,
|
||||
index: str,
|
||||
documents: Sequence[dict[str, Any]],
|
||||
primary_key: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Add or update documents.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
documents: Documents to add.
|
||||
primary_key: Override primary key for this operation.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.add_documents(list(documents), primary_key=primary_key)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def update_documents(
|
||||
self,
|
||||
index: str,
|
||||
documents: Sequence[dict[str, Any]],
|
||||
primary_key: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Update documents (partial update).
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
documents: Documents with updates.
|
||||
primary_key: Override primary key for this operation.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.update_documents(list(documents), primary_key=primary_key)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def delete_document(self, index: str, document_id: str) -> None:
|
||||
"""Delete a document by ID.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
document_id: Document ID to delete.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.delete_document(document_id)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def delete_documents(self, index: str, document_ids: list[str]) -> None:
|
||||
"""Delete multiple documents by ID.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
document_ids: Document IDs to delete.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.delete_documents(document_ids)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def delete_all_documents(self, index: str) -> None:
|
||||
"""Delete all documents in an index.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.delete_all_documents()
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def get_document(
|
||||
self,
|
||||
index: str,
|
||||
document_id: str,
|
||||
fields: Optional[list[str]] = None,
|
||||
) -> Optional[dict[str, Any]]:
|
||||
"""Get a document by ID.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
document_id: Document ID.
|
||||
fields: Fields to retrieve (all if not specified).
|
||||
|
||||
Returns:
|
||||
Document or None if not found.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
try:
|
||||
return await idx.get_document(document_id, fields=fields)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
async def get_documents(
|
||||
self,
|
||||
index: str,
|
||||
offset: int = 0,
|
||||
limit: int = 20,
|
||||
fields: Optional[list[str]] = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Get documents from an index.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
offset: Number of documents to skip.
|
||||
limit: Maximum documents to return.
|
||||
fields: Fields to retrieve (all if not specified).
|
||||
|
||||
Returns:
|
||||
List of documents.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
result = await idx.get_documents(offset=offset, limit=limit, fields=fields)
|
||||
return result.results
|
||||
|
||||
# Search operations
|
||||
|
||||
async def search(
|
||||
self,
|
||||
index: str,
|
||||
query: str,
|
||||
offset: int = 0,
|
||||
limit: int = 20,
|
||||
filters: Optional[str] = None,
|
||||
facets: Optional[list[str]] = None,
|
||||
attributes_to_retrieve: Optional[list[str]] = None,
|
||||
attributes_to_highlight: Optional[list[str]] = None,
|
||||
sort: Optional[list[str]] = None,
|
||||
show_matches_position: bool = False,
|
||||
show_ranking_score: bool = False,
|
||||
) -> SearchResult:
|
||||
"""Search documents in an index.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
query: Search query.
|
||||
offset: Number of results to skip.
|
||||
limit: Maximum results to return.
|
||||
filters: Filter expression (e.g., "price > 50 AND category = 'tech'").
|
||||
facets: Fields to get facet distribution for.
|
||||
attributes_to_retrieve: Fields to include in results.
|
||||
attributes_to_highlight: Fields to highlight matches in.
|
||||
sort: Sort expressions (e.g., ["price:asc", "name:desc"]).
|
||||
show_matches_position: Include match positions.
|
||||
show_ranking_score: Include ranking score.
|
||||
|
||||
Returns:
|
||||
Search results with hits and metadata.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
result = await idx.search(
|
||||
query,
|
||||
offset=offset,
|
||||
limit=limit,
|
||||
filter=filters,
|
||||
facets=facets,
|
||||
attributes_to_retrieve=attributes_to_retrieve,
|
||||
attributes_to_highlight=attributes_to_highlight,
|
||||
sort=sort,
|
||||
show_matches_position=show_matches_position,
|
||||
show_ranking_score=show_ranking_score,
|
||||
)
|
||||
|
||||
hits = []
|
||||
for hit in result.hits:
|
||||
# Handle both dict and object responses
|
||||
if isinstance(hit, dict):
|
||||
doc = hit
|
||||
formatted = hit.get("_formatted", {})
|
||||
score = hit.get("_rankingScore")
|
||||
else:
|
||||
doc = hit.__dict__ if hasattr(hit, "__dict__") else {}
|
||||
formatted = getattr(hit, "_formatted", {}) or {}
|
||||
score = getattr(hit, "_rankingScore", None)
|
||||
|
||||
# Extract ID from document
|
||||
doc_id = doc.get("id") or doc.get("_id") or str(hash(str(doc)))
|
||||
|
||||
hits.append(SearchHit(
|
||||
id=str(doc_id),
|
||||
document=doc,
|
||||
score=score,
|
||||
highlights=formatted,
|
||||
))
|
||||
|
||||
return SearchResult(
|
||||
hits=hits,
|
||||
query=query,
|
||||
processing_time_ms=result.processing_time_ms,
|
||||
total_hits=result.estimated_total_hits or len(hits),
|
||||
offset=offset,
|
||||
limit=limit,
|
||||
facet_distribution=result.facet_distribution or {},
|
||||
)
|
||||
|
||||
# Settings operations
|
||||
|
||||
async def get_settings(self, index: str) -> dict[str, Any]:
|
||||
"""Get index settings.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
|
||||
Returns:
|
||||
Current settings.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
settings = await idx.get_settings()
|
||||
return settings.__dict__ if hasattr(settings, "__dict__") else dict(settings)
|
||||
|
||||
async def update_settings(
|
||||
self,
|
||||
index: str,
|
||||
settings: dict[str, Any],
|
||||
) -> None:
|
||||
"""Update index settings.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
settings: Settings to update.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.update_settings(settings)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def update_searchable_attributes(
|
||||
self,
|
||||
index: str,
|
||||
attributes: list[str],
|
||||
) -> None:
|
||||
"""Set searchable attributes.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
attributes: Ordered list of searchable attributes.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.update_searchable_attributes(attributes)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def update_filterable_attributes(
|
||||
self,
|
||||
index: str,
|
||||
attributes: list[str],
|
||||
) -> None:
|
||||
"""Set filterable attributes.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
attributes: List of filterable attributes.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.update_filterable_attributes(attributes)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def update_sortable_attributes(
|
||||
self,
|
||||
index: str,
|
||||
attributes: list[str],
|
||||
) -> None:
|
||||
"""Set sortable attributes.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
attributes: List of sortable attributes.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.update_sortable_attributes(attributes)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def update_ranking_rules(
|
||||
self,
|
||||
index: str,
|
||||
rules: list[str],
|
||||
) -> None:
|
||||
"""Set ranking rules.
|
||||
|
||||
Args:
|
||||
index: Index name.
|
||||
rules: Ordered list of ranking rules.
|
||||
"""
|
||||
idx = self._client.index(index)
|
||||
task = await idx.update_ranking_rules(rules)
|
||||
await self._client.wait_for_task(task.task_uid)
|
||||
|
||||
async def __aenter__(self) -> SearchClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,535 @@
|
||||
"""S3/MinIO object storage client wrapper for Hanzo infrastructure.
|
||||
|
||||
Provides async interface to S3-compatible object storage including
|
||||
AWS S3, MinIO, and other compatible services.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from io import BytesIO
|
||||
from typing import Any, AsyncIterator, BinaryIO, Optional, Union
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class StorageConfig(BaseModel):
|
||||
"""Configuration for S3/MinIO connection."""
|
||||
|
||||
endpoint_url: Optional[str] = Field(default=None, description="Endpoint URL (required for MinIO)")
|
||||
region: str = Field(default="us-east-1", description="AWS region")
|
||||
access_key: Optional[str] = Field(default=None, description="Access key ID")
|
||||
secret_key: Optional[str] = Field(default=None, description="Secret access key")
|
||||
bucket: str = Field(default="hanzo", description="Default bucket name")
|
||||
use_ssl: bool = Field(default=True, description="Use SSL/TLS")
|
||||
verify_ssl: bool = Field(default=True, description="Verify SSL certificates")
|
||||
session_token: Optional[str] = Field(default=None, description="Session token (for temporary credentials)")
|
||||
addressing_style: str = Field(default="auto", description="Path or virtual addressing style")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> StorageConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
S3_ENDPOINT_URL: Endpoint URL (for MinIO/custom S3)
|
||||
S3_REGION / AWS_REGION: AWS region (default: us-east-1)
|
||||
S3_ACCESS_KEY / AWS_ACCESS_KEY_ID: Access key
|
||||
S3_SECRET_KEY / AWS_SECRET_ACCESS_KEY: Secret key
|
||||
S3_BUCKET: Default bucket name (default: hanzo)
|
||||
S3_USE_SSL: Use SSL (default: true)
|
||||
"""
|
||||
return cls(
|
||||
endpoint_url=os.getenv("S3_ENDPOINT_URL") or os.getenv("MINIO_ENDPOINT"),
|
||||
region=os.getenv("S3_REGION") or os.getenv("AWS_REGION", "us-east-1"),
|
||||
access_key=os.getenv("S3_ACCESS_KEY") or os.getenv("AWS_ACCESS_KEY_ID"),
|
||||
secret_key=os.getenv("S3_SECRET_KEY") or os.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
bucket=os.getenv("S3_BUCKET", "hanzo"),
|
||||
use_ssl=os.getenv("S3_USE_SSL", "true").lower() in ("true", "1", "yes"),
|
||||
session_token=os.getenv("AWS_SESSION_TOKEN"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ObjectInfo:
|
||||
"""Metadata about a stored object."""
|
||||
|
||||
key: str
|
||||
size: int
|
||||
etag: str
|
||||
last_modified: datetime
|
||||
content_type: Optional[str] = None
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class UploadResult:
|
||||
"""Result of an upload operation."""
|
||||
|
||||
key: str
|
||||
etag: str
|
||||
version_id: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PresignedUrl:
|
||||
"""A presigned URL for temporary access."""
|
||||
|
||||
url: str
|
||||
expires_in: int # seconds
|
||||
|
||||
|
||||
class StorageClient:
|
||||
"""Async client for S3-compatible object storage.
|
||||
|
||||
Wraps aiobotocore/boto3 for async S3 operations with Hanzo conventions.
|
||||
Supports AWS S3, MinIO, and other S3-compatible services.
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = StorageClient(StorageConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Upload
|
||||
await client.put_object("files/doc.txt", b"Hello, World!")
|
||||
|
||||
# Download
|
||||
data = await client.get_object("files/doc.txt")
|
||||
|
||||
# List objects
|
||||
async for obj in client.list_objects("files/"):
|
||||
print(obj.key, obj.size)
|
||||
|
||||
# Presigned URLs
|
||||
url = await client.presign_get("files/doc.txt", expires_in=3600)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[StorageConfig] = None) -> None:
|
||||
"""Initialize storage client.
|
||||
|
||||
Args:
|
||||
config: S3 configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or StorageConfig.from_env()
|
||||
self._session: Any = None
|
||||
self._client: Any = None
|
||||
self._exit_stack: Any = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection to S3 service."""
|
||||
try:
|
||||
from aiobotocore.session import get_session
|
||||
from contextlib import AsyncExitStack
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"aiobotocore is required for StorageClient. "
|
||||
"Install with: pip install aiobotocore"
|
||||
) from e
|
||||
|
||||
self._session = get_session()
|
||||
self._exit_stack = AsyncExitStack()
|
||||
|
||||
client_kwargs: dict[str, Any] = {
|
||||
"service_name": "s3",
|
||||
"region_name": self.config.region,
|
||||
}
|
||||
|
||||
if self.config.endpoint_url:
|
||||
client_kwargs["endpoint_url"] = self.config.endpoint_url
|
||||
|
||||
if self.config.access_key and self.config.secret_key:
|
||||
client_kwargs["aws_access_key_id"] = self.config.access_key
|
||||
client_kwargs["aws_secret_access_key"] = self.config.secret_key
|
||||
|
||||
if self.config.session_token:
|
||||
client_kwargs["aws_session_token"] = self.config.session_token
|
||||
|
||||
client_kwargs["config"] = self._get_client_config()
|
||||
|
||||
ctx_manager = self._session.create_client(**client_kwargs)
|
||||
self._client = await self._exit_stack.enter_async_context(ctx_manager)
|
||||
|
||||
def _get_client_config(self) -> Any:
|
||||
"""Get boto client config."""
|
||||
from botocore.config import Config
|
||||
|
||||
return Config(
|
||||
signature_version="s3v4",
|
||||
s3={"addressing_style": self.config.addressing_style},
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
if self._exit_stack:
|
||||
await self._exit_stack.aclose()
|
||||
self._exit_stack = None
|
||||
self._client = None
|
||||
self._session = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if S3 service is healthy.
|
||||
|
||||
Returns:
|
||||
True if service is reachable.
|
||||
"""
|
||||
if not self._client:
|
||||
return False
|
||||
try:
|
||||
await self._client.list_buckets()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# Bucket operations
|
||||
|
||||
async def create_bucket(self, bucket: Optional[str] = None) -> None:
|
||||
"""Create a bucket.
|
||||
|
||||
Args:
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
create_config = {}
|
||||
if self.config.region != "us-east-1":
|
||||
create_config["CreateBucketConfiguration"] = {
|
||||
"LocationConstraint": self.config.region
|
||||
}
|
||||
await self._client.create_bucket(Bucket=bucket_name, **create_config)
|
||||
|
||||
async def delete_bucket(self, bucket: Optional[str] = None) -> None:
|
||||
"""Delete a bucket.
|
||||
|
||||
Args:
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
await self._client.delete_bucket(Bucket=bucket_name)
|
||||
|
||||
async def bucket_exists(self, bucket: Optional[str] = None) -> bool:
|
||||
"""Check if a bucket exists.
|
||||
|
||||
Args:
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
True if bucket exists.
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
try:
|
||||
await self._client.head_bucket(Bucket=bucket_name)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def list_buckets(self) -> list[str]:
|
||||
"""List all buckets.
|
||||
|
||||
Returns:
|
||||
List of bucket names.
|
||||
"""
|
||||
response = await self._client.list_buckets()
|
||||
return [b["Name"] for b in response.get("Buckets", [])]
|
||||
|
||||
# Object operations
|
||||
|
||||
async def put_object(
|
||||
self,
|
||||
key: str,
|
||||
data: Union[bytes, BinaryIO, BytesIO],
|
||||
content_type: Optional[str] = None,
|
||||
metadata: Optional[dict[str, str]] = None,
|
||||
bucket: Optional[str] = None,
|
||||
) -> UploadResult:
|
||||
"""Upload an object.
|
||||
|
||||
Args:
|
||||
key: Object key.
|
||||
data: Object data as bytes or file-like object.
|
||||
content_type: Content type (auto-detected if not specified).
|
||||
metadata: Custom metadata.
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
Upload result with etag.
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
|
||||
put_kwargs: dict[str, Any] = {
|
||||
"Bucket": bucket_name,
|
||||
"Key": key,
|
||||
"Body": data,
|
||||
}
|
||||
|
||||
if content_type:
|
||||
put_kwargs["ContentType"] = content_type
|
||||
if metadata:
|
||||
put_kwargs["Metadata"] = metadata
|
||||
|
||||
response = await self._client.put_object(**put_kwargs)
|
||||
return UploadResult(
|
||||
key=key,
|
||||
etag=response.get("ETag", "").strip('"'),
|
||||
version_id=response.get("VersionId"),
|
||||
)
|
||||
|
||||
async def get_object(
|
||||
self,
|
||||
key: str,
|
||||
bucket: Optional[str] = None,
|
||||
) -> bytes:
|
||||
"""Download an object.
|
||||
|
||||
Args:
|
||||
key: Object key.
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
Object data as bytes.
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
response = await self._client.get_object(Bucket=bucket_name, Key=key)
|
||||
async with response["Body"] as stream:
|
||||
return await stream.read()
|
||||
|
||||
async def get_object_stream(
|
||||
self,
|
||||
key: str,
|
||||
bucket: Optional[str] = None,
|
||||
) -> AsyncIterator[bytes]:
|
||||
"""Stream an object in chunks.
|
||||
|
||||
Args:
|
||||
key: Object key.
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
|
||||
Yields:
|
||||
Chunks of object data.
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
response = await self._client.get_object(Bucket=bucket_name, Key=key)
|
||||
async with response["Body"] as stream:
|
||||
async for chunk in stream.iter_chunks():
|
||||
yield chunk[0] # iter_chunks returns (chunk, final)
|
||||
|
||||
async def delete_object(
|
||||
self,
|
||||
key: str,
|
||||
bucket: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Delete an object.
|
||||
|
||||
Args:
|
||||
key: Object key.
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
await self._client.delete_object(Bucket=bucket_name, Key=key)
|
||||
|
||||
async def delete_objects(
|
||||
self,
|
||||
keys: list[str],
|
||||
bucket: Optional[str] = None,
|
||||
) -> list[str]:
|
||||
"""Delete multiple objects.
|
||||
|
||||
Args:
|
||||
keys: Object keys to delete.
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
List of deleted keys.
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
response = await self._client.delete_objects(
|
||||
Bucket=bucket_name,
|
||||
Delete={"Objects": [{"Key": k} for k in keys]},
|
||||
)
|
||||
return [d["Key"] for d in response.get("Deleted", [])]
|
||||
|
||||
async def head_object(
|
||||
self,
|
||||
key: str,
|
||||
bucket: Optional[str] = None,
|
||||
) -> ObjectInfo:
|
||||
"""Get object metadata without downloading.
|
||||
|
||||
Args:
|
||||
key: Object key.
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
Object info with metadata.
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
response = await self._client.head_object(Bucket=bucket_name, Key=key)
|
||||
return ObjectInfo(
|
||||
key=key,
|
||||
size=response.get("ContentLength", 0),
|
||||
etag=response.get("ETag", "").strip('"'),
|
||||
last_modified=response.get("LastModified", datetime.now()),
|
||||
content_type=response.get("ContentType"),
|
||||
metadata=response.get("Metadata", {}),
|
||||
)
|
||||
|
||||
async def object_exists(
|
||||
self,
|
||||
key: str,
|
||||
bucket: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""Check if an object exists.
|
||||
|
||||
Args:
|
||||
key: Object key.
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
True if object exists.
|
||||
"""
|
||||
try:
|
||||
await self.head_object(key, bucket)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def list_objects(
|
||||
self,
|
||||
prefix: str = "",
|
||||
bucket: Optional[str] = None,
|
||||
max_keys: int = 1000,
|
||||
) -> AsyncIterator[ObjectInfo]:
|
||||
"""List objects with a prefix.
|
||||
|
||||
Args:
|
||||
prefix: Key prefix filter.
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
max_keys: Maximum keys per request.
|
||||
|
||||
Yields:
|
||||
Object info for each matching object.
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
continuation_token = None
|
||||
|
||||
while True:
|
||||
kwargs: dict[str, Any] = {
|
||||
"Bucket": bucket_name,
|
||||
"Prefix": prefix,
|
||||
"MaxKeys": max_keys,
|
||||
}
|
||||
if continuation_token:
|
||||
kwargs["ContinuationToken"] = continuation_token
|
||||
|
||||
response = await self._client.list_objects_v2(**kwargs)
|
||||
|
||||
for obj in response.get("Contents", []):
|
||||
yield ObjectInfo(
|
||||
key=obj["Key"],
|
||||
size=obj["Size"],
|
||||
etag=obj["ETag"].strip('"'),
|
||||
last_modified=obj["LastModified"],
|
||||
)
|
||||
|
||||
if not response.get("IsTruncated"):
|
||||
break
|
||||
continuation_token = response.get("NextContinuationToken")
|
||||
|
||||
# Copy operations
|
||||
|
||||
async def copy_object(
|
||||
self,
|
||||
source_key: str,
|
||||
dest_key: str,
|
||||
source_bucket: Optional[str] = None,
|
||||
dest_bucket: Optional[str] = None,
|
||||
) -> UploadResult:
|
||||
"""Copy an object.
|
||||
|
||||
Args:
|
||||
source_key: Source object key.
|
||||
dest_key: Destination object key.
|
||||
source_bucket: Source bucket (uses default if not specified).
|
||||
dest_bucket: Destination bucket (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
Upload result for the copy.
|
||||
"""
|
||||
src_bucket = source_bucket or self.config.bucket
|
||||
dst_bucket = dest_bucket or self.config.bucket
|
||||
|
||||
response = await self._client.copy_object(
|
||||
Bucket=dst_bucket,
|
||||
Key=dest_key,
|
||||
CopySource={"Bucket": src_bucket, "Key": source_key},
|
||||
)
|
||||
return UploadResult(
|
||||
key=dest_key,
|
||||
etag=response.get("CopyObjectResult", {}).get("ETag", "").strip('"'),
|
||||
version_id=response.get("VersionId"),
|
||||
)
|
||||
|
||||
# Presigned URLs
|
||||
|
||||
async def presign_get(
|
||||
self,
|
||||
key: str,
|
||||
expires_in: int = 3600,
|
||||
bucket: Optional[str] = None,
|
||||
) -> PresignedUrl:
|
||||
"""Generate a presigned URL for download.
|
||||
|
||||
Args:
|
||||
key: Object key.
|
||||
expires_in: URL validity in seconds (default: 1 hour).
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
Presigned URL.
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
url = await self._client.generate_presigned_url(
|
||||
"get_object",
|
||||
Params={"Bucket": bucket_name, "Key": key},
|
||||
ExpiresIn=expires_in,
|
||||
)
|
||||
return PresignedUrl(url=url, expires_in=expires_in)
|
||||
|
||||
async def presign_put(
|
||||
self,
|
||||
key: str,
|
||||
expires_in: int = 3600,
|
||||
content_type: Optional[str] = None,
|
||||
bucket: Optional[str] = None,
|
||||
) -> PresignedUrl:
|
||||
"""Generate a presigned URL for upload.
|
||||
|
||||
Args:
|
||||
key: Object key.
|
||||
expires_in: URL validity in seconds (default: 1 hour).
|
||||
content_type: Required content type.
|
||||
bucket: Bucket name (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
Presigned URL.
|
||||
"""
|
||||
bucket_name = bucket or self.config.bucket
|
||||
params: dict[str, Any] = {"Bucket": bucket_name, "Key": key}
|
||||
if content_type:
|
||||
params["ContentType"] = content_type
|
||||
|
||||
url = await self._client.generate_presigned_url(
|
||||
"put_object",
|
||||
Params=params,
|
||||
ExpiresIn=expires_in,
|
||||
)
|
||||
return PresignedUrl(url=url, expires_in=expires_in)
|
||||
|
||||
async def __aenter__(self) -> StorageClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,518 @@
|
||||
"""Temporal workflow client wrapper for Hanzo infrastructure.
|
||||
|
||||
Provides async interface to Temporal for durable workflows,
|
||||
activities, and task orchestration.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, Callable, Optional, Sequence, TypeVar
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class TasksConfig(BaseModel):
|
||||
"""Configuration for Temporal connection."""
|
||||
|
||||
host: str = Field(default="localhost", description="Temporal server host")
|
||||
port: int = Field(default=7233, description="Temporal server port")
|
||||
namespace: str = Field(default="default", description="Temporal namespace")
|
||||
target: Optional[str] = Field(default=None, description="Full target address (overrides host/port)")
|
||||
tls: bool = Field(default=False, description="Use TLS")
|
||||
tls_cert_path: Optional[str] = Field(default=None, description="Path to TLS certificate")
|
||||
tls_key_path: Optional[str] = Field(default=None, description="Path to TLS key")
|
||||
api_key: Optional[str] = Field(default=None, description="API key for Temporal Cloud")
|
||||
identity: str = Field(default="hanzo-client", description="Client identity")
|
||||
data_converter: Optional[str] = Field(default=None, description="Custom data converter class")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> TasksConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
TEMPORAL_HOST: Server host (default: localhost)
|
||||
TEMPORAL_PORT: Server port (default: 7233)
|
||||
TEMPORAL_NAMESPACE: Namespace (default: default)
|
||||
TEMPORAL_TARGET: Full target address (overrides host/port)
|
||||
TEMPORAL_TLS: Use TLS (default: false)
|
||||
TEMPORAL_TLS_CERT: Path to TLS certificate
|
||||
TEMPORAL_TLS_KEY: Path to TLS key
|
||||
TEMPORAL_API_KEY: API key for Temporal Cloud
|
||||
"""
|
||||
return cls(
|
||||
host=os.getenv("TEMPORAL_HOST", "localhost"),
|
||||
port=int(os.getenv("TEMPORAL_PORT", "7233")),
|
||||
namespace=os.getenv("TEMPORAL_NAMESPACE", "default"),
|
||||
target=os.getenv("TEMPORAL_TARGET"),
|
||||
tls=os.getenv("TEMPORAL_TLS", "").lower() in ("true", "1", "yes"),
|
||||
tls_cert_path=os.getenv("TEMPORAL_TLS_CERT"),
|
||||
tls_key_path=os.getenv("TEMPORAL_TLS_KEY"),
|
||||
api_key=os.getenv("TEMPORAL_API_KEY"),
|
||||
)
|
||||
|
||||
@property
|
||||
def effective_target(self) -> str:
|
||||
"""Get the effective target address."""
|
||||
return self.target or f"{self.host}:{self.port}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkflowExecution:
|
||||
"""Information about a workflow execution."""
|
||||
|
||||
workflow_id: str
|
||||
run_id: str
|
||||
workflow_type: str
|
||||
status: str
|
||||
start_time: Optional[datetime] = None
|
||||
close_time: Optional[datetime] = None
|
||||
history_length: int = 0
|
||||
memo: dict[str, Any] = field(default_factory=dict)
|
||||
search_attributes: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkflowHandle:
|
||||
"""Handle to a running workflow."""
|
||||
|
||||
workflow_id: str
|
||||
run_id: Optional[str] = None
|
||||
_handle: Any = None
|
||||
|
||||
async def result(self) -> Any:
|
||||
"""Wait for and return the workflow result."""
|
||||
if self._handle:
|
||||
return await self._handle.result()
|
||||
return None
|
||||
|
||||
async def cancel(self) -> None:
|
||||
"""Request cancellation of the workflow."""
|
||||
if self._handle:
|
||||
await self._handle.cancel()
|
||||
|
||||
async def terminate(self, reason: str = "") -> None:
|
||||
"""Terminate the workflow."""
|
||||
if self._handle:
|
||||
await self._handle.terminate(reason)
|
||||
|
||||
async def signal(self, name: str, *args: Any) -> None:
|
||||
"""Send a signal to the workflow."""
|
||||
if self._handle:
|
||||
await self._handle.signal(name, *args)
|
||||
|
||||
async def query(self, name: str, *args: Any) -> Any:
|
||||
"""Query the workflow state."""
|
||||
if self._handle:
|
||||
return await self._handle.query(name, *args)
|
||||
return None
|
||||
|
||||
async def describe(self) -> WorkflowExecution:
|
||||
"""Get workflow execution details."""
|
||||
if self._handle:
|
||||
desc = await self._handle.describe()
|
||||
return WorkflowExecution(
|
||||
workflow_id=desc.id,
|
||||
run_id=desc.run_id,
|
||||
workflow_type=desc.workflow_type,
|
||||
status=str(desc.status),
|
||||
start_time=desc.start_time,
|
||||
close_time=desc.close_time,
|
||||
)
|
||||
return WorkflowExecution(
|
||||
workflow_id=self.workflow_id,
|
||||
run_id=self.run_id or "",
|
||||
workflow_type="",
|
||||
status="UNKNOWN",
|
||||
)
|
||||
|
||||
|
||||
class TasksClient:
|
||||
"""Async client for Temporal workflow orchestration.
|
||||
|
||||
Wraps temporalio SDK with Hanzo conventions for workflow
|
||||
execution, activities, and task management.
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = TasksClient(TasksConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Start a workflow
|
||||
handle = await client.start_workflow(
|
||||
"ProcessOrder",
|
||||
args={"order_id": "123"},
|
||||
id="order-123",
|
||||
task_queue="orders",
|
||||
)
|
||||
|
||||
# Wait for result
|
||||
result = await handle.result()
|
||||
|
||||
# Query workflow state
|
||||
state = await handle.query("get_status")
|
||||
|
||||
# Send signal
|
||||
await handle.signal("approve", {"user": "admin"})
|
||||
|
||||
# List workflows
|
||||
async for wf in client.list_workflows("WorkflowType = 'ProcessOrder'"):
|
||||
print(wf.workflow_id, wf.status)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[TasksConfig] = None) -> None:
|
||||
"""Initialize tasks client.
|
||||
|
||||
Args:
|
||||
config: Temporal configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or TasksConfig.from_env()
|
||||
self._client: Any = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection to Temporal server."""
|
||||
try:
|
||||
from temporalio.client import Client, TLSConfig
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"temporalio is required for TasksClient. "
|
||||
"Install with: pip install temporalio"
|
||||
) from e
|
||||
|
||||
connect_kwargs: dict[str, Any] = {
|
||||
"target_host": self.config.effective_target,
|
||||
"namespace": self.config.namespace,
|
||||
"identity": self.config.identity,
|
||||
}
|
||||
|
||||
if self.config.tls:
|
||||
tls_config = TLSConfig()
|
||||
if self.config.tls_cert_path and self.config.tls_key_path:
|
||||
with open(self.config.tls_cert_path, "rb") as f:
|
||||
cert = f.read()
|
||||
with open(self.config.tls_key_path, "rb") as f:
|
||||
key = f.read()
|
||||
tls_config = TLSConfig(client_cert=cert, client_private_key=key)
|
||||
connect_kwargs["tls"] = tls_config
|
||||
|
||||
if self.config.api_key:
|
||||
connect_kwargs["api_key"] = self.config.api_key
|
||||
|
||||
self._client = await Client.connect(**connect_kwargs)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
# Temporal client doesn't require explicit close
|
||||
self._client = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if Temporal server is healthy.
|
||||
|
||||
Returns:
|
||||
True if server is reachable.
|
||||
"""
|
||||
if not self._client:
|
||||
return False
|
||||
try:
|
||||
# Try to get system info
|
||||
await self._client.service_client.check_health()
|
||||
return True
|
||||
except Exception:
|
||||
# Fallback: try listing workflows with limit 1
|
||||
try:
|
||||
async for _ in self._client.list_workflows(query="", page_size=1):
|
||||
break
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# Workflow operations
|
||||
|
||||
async def start_workflow(
|
||||
self,
|
||||
workflow: str,
|
||||
args: Any = None,
|
||||
id: Optional[str] = None,
|
||||
task_queue: str = "default",
|
||||
execution_timeout: Optional[timedelta] = None,
|
||||
run_timeout: Optional[timedelta] = None,
|
||||
task_timeout: Optional[timedelta] = None,
|
||||
id_reuse_policy: str = "allow_duplicate",
|
||||
retry_policy: Optional[dict[str, Any]] = None,
|
||||
cron_schedule: Optional[str] = None,
|
||||
memo: Optional[dict[str, Any]] = None,
|
||||
search_attributes: Optional[dict[str, Any]] = None,
|
||||
) -> WorkflowHandle:
|
||||
"""Start a workflow execution.
|
||||
|
||||
Args:
|
||||
workflow: Workflow type name.
|
||||
args: Workflow arguments.
|
||||
id: Workflow ID (auto-generated if not specified).
|
||||
task_queue: Task queue for the workflow.
|
||||
execution_timeout: Max workflow execution time.
|
||||
run_timeout: Max single run time.
|
||||
task_timeout: Max workflow task time.
|
||||
id_reuse_policy: ID reuse policy (allow_duplicate, allow_duplicate_failed_only, reject_duplicate, terminate_if_running).
|
||||
retry_policy: Retry configuration.
|
||||
cron_schedule: Cron schedule expression.
|
||||
memo: Workflow memo fields.
|
||||
search_attributes: Custom search attributes.
|
||||
|
||||
Returns:
|
||||
Handle to the started workflow.
|
||||
"""
|
||||
from temporalio.client import WorkflowIDReusePolicy
|
||||
from temporalio.common import RetryPolicy
|
||||
import uuid
|
||||
|
||||
policy_map = {
|
||||
"allow_duplicate": WorkflowIDReusePolicy.ALLOW_DUPLICATE,
|
||||
"allow_duplicate_failed_only": WorkflowIDReusePolicy.ALLOW_DUPLICATE_FAILED_ONLY,
|
||||
"reject_duplicate": WorkflowIDReusePolicy.REJECT_DUPLICATE,
|
||||
"terminate_if_running": WorkflowIDReusePolicy.TERMINATE_IF_RUNNING,
|
||||
}
|
||||
|
||||
start_kwargs: dict[str, Any] = {
|
||||
"workflow": workflow,
|
||||
"arg": args,
|
||||
"id": id or f"{workflow}-{uuid.uuid4().hex[:8]}",
|
||||
"task_queue": task_queue,
|
||||
"id_reuse_policy": policy_map.get(id_reuse_policy, WorkflowIDReusePolicy.ALLOW_DUPLICATE),
|
||||
}
|
||||
|
||||
if execution_timeout:
|
||||
start_kwargs["execution_timeout"] = execution_timeout
|
||||
if run_timeout:
|
||||
start_kwargs["run_timeout"] = run_timeout
|
||||
if task_timeout:
|
||||
start_kwargs["task_timeout"] = task_timeout
|
||||
if cron_schedule:
|
||||
start_kwargs["cron_schedule"] = cron_schedule
|
||||
if memo:
|
||||
start_kwargs["memo"] = memo
|
||||
if search_attributes:
|
||||
start_kwargs["search_attributes"] = search_attributes
|
||||
|
||||
if retry_policy:
|
||||
start_kwargs["retry_policy"] = RetryPolicy(
|
||||
initial_interval=timedelta(seconds=retry_policy.get("initial_interval", 1)),
|
||||
maximum_interval=timedelta(seconds=retry_policy.get("maximum_interval", 100)),
|
||||
backoff_coefficient=retry_policy.get("backoff_coefficient", 2.0),
|
||||
maximum_attempts=retry_policy.get("maximum_attempts", 0),
|
||||
)
|
||||
|
||||
handle = await self._client.start_workflow(**start_kwargs)
|
||||
return WorkflowHandle(
|
||||
workflow_id=handle.id,
|
||||
run_id=handle.result_run_id,
|
||||
_handle=handle,
|
||||
)
|
||||
|
||||
async def execute_workflow(
|
||||
self,
|
||||
workflow: str,
|
||||
args: Any = None,
|
||||
id: Optional[str] = None,
|
||||
task_queue: str = "default",
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Start a workflow and wait for its result.
|
||||
|
||||
Args:
|
||||
workflow: Workflow type name.
|
||||
args: Workflow arguments.
|
||||
id: Workflow ID.
|
||||
task_queue: Task queue.
|
||||
**kwargs: Additional start_workflow arguments.
|
||||
|
||||
Returns:
|
||||
Workflow result.
|
||||
"""
|
||||
handle = await self.start_workflow(workflow, args, id, task_queue, **kwargs)
|
||||
return await handle.result()
|
||||
|
||||
async def get_workflow_handle(
|
||||
self,
|
||||
workflow_id: str,
|
||||
run_id: Optional[str] = None,
|
||||
) -> WorkflowHandle:
|
||||
"""Get a handle to an existing workflow.
|
||||
|
||||
Args:
|
||||
workflow_id: Workflow ID.
|
||||
run_id: Optional specific run ID.
|
||||
|
||||
Returns:
|
||||
Workflow handle.
|
||||
"""
|
||||
handle = self._client.get_workflow_handle(workflow_id, run_id=run_id)
|
||||
return WorkflowHandle(
|
||||
workflow_id=workflow_id,
|
||||
run_id=run_id,
|
||||
_handle=handle,
|
||||
)
|
||||
|
||||
async def list_workflows(
|
||||
self,
|
||||
query: str = "",
|
||||
page_size: int = 100,
|
||||
):
|
||||
"""List workflow executions.
|
||||
|
||||
Args:
|
||||
query: Temporal list filter query.
|
||||
page_size: Results per page.
|
||||
|
||||
Yields:
|
||||
WorkflowExecution for each matching workflow.
|
||||
"""
|
||||
async for wf in self._client.list_workflows(query=query, page_size=page_size):
|
||||
yield WorkflowExecution(
|
||||
workflow_id=wf.id,
|
||||
run_id=wf.run_id,
|
||||
workflow_type=wf.workflow_type,
|
||||
status=str(wf.status),
|
||||
start_time=wf.start_time,
|
||||
close_time=wf.close_time,
|
||||
memo=dict(wf.memo) if wf.memo else {},
|
||||
search_attributes=dict(wf.search_attributes) if wf.search_attributes else {},
|
||||
)
|
||||
|
||||
async def count_workflows(self, query: str = "") -> int:
|
||||
"""Count workflow executions matching a query.
|
||||
|
||||
Args:
|
||||
query: Temporal list filter query.
|
||||
|
||||
Returns:
|
||||
Number of matching workflows.
|
||||
"""
|
||||
count = 0
|
||||
async for _ in self._client.list_workflows(query=query):
|
||||
count += 1
|
||||
return count
|
||||
|
||||
# Schedule operations
|
||||
|
||||
async def create_schedule(
|
||||
self,
|
||||
schedule_id: str,
|
||||
workflow: str,
|
||||
args: Any = None,
|
||||
task_queue: str = "default",
|
||||
cron: Optional[str] = None,
|
||||
interval: Optional[timedelta] = None,
|
||||
calendar: Optional[dict[str, Any]] = None,
|
||||
start_at: Optional[datetime] = None,
|
||||
end_at: Optional[datetime] = None,
|
||||
jitter: Optional[timedelta] = None,
|
||||
memo: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Create a workflow schedule.
|
||||
|
||||
Args:
|
||||
schedule_id: Unique schedule ID.
|
||||
workflow: Workflow type name.
|
||||
args: Workflow arguments.
|
||||
task_queue: Task queue.
|
||||
cron: Cron expression (e.g., "0 * * * *").
|
||||
interval: Interval between runs.
|
||||
calendar: Calendar-based schedule spec.
|
||||
start_at: Schedule start time.
|
||||
end_at: Schedule end time.
|
||||
jitter: Random jitter to add to scheduled times.
|
||||
memo: Schedule memo fields.
|
||||
"""
|
||||
from temporalio.client import (
|
||||
Schedule,
|
||||
ScheduleSpec,
|
||||
ScheduleActionStartWorkflow,
|
||||
ScheduleIntervalSpec,
|
||||
ScheduleCalendarSpec,
|
||||
)
|
||||
|
||||
specs = []
|
||||
if cron:
|
||||
specs.append(ScheduleSpec(cron_expressions=[cron]))
|
||||
if interval:
|
||||
specs.append(ScheduleSpec(intervals=[ScheduleIntervalSpec(every=interval)]))
|
||||
if calendar:
|
||||
specs.append(ScheduleSpec(calendars=[ScheduleCalendarSpec(**calendar)]))
|
||||
|
||||
if not specs:
|
||||
raise ValueError("At least one of cron, interval, or calendar must be specified")
|
||||
|
||||
schedule_spec = specs[0] if len(specs) == 1 else ScheduleSpec(
|
||||
cron_expressions=[cron] if cron else None,
|
||||
intervals=[ScheduleIntervalSpec(every=interval)] if interval else None,
|
||||
calendars=[ScheduleCalendarSpec(**calendar)] if calendar else None,
|
||||
start_at=start_at,
|
||||
end_at=end_at,
|
||||
jitter=jitter,
|
||||
)
|
||||
|
||||
await self._client.create_schedule(
|
||||
schedule_id,
|
||||
Schedule(
|
||||
action=ScheduleActionStartWorkflow(
|
||||
workflow,
|
||||
arg=args,
|
||||
task_queue=task_queue,
|
||||
),
|
||||
spec=schedule_spec,
|
||||
),
|
||||
memo=memo,
|
||||
)
|
||||
|
||||
async def delete_schedule(self, schedule_id: str) -> None:
|
||||
"""Delete a schedule.
|
||||
|
||||
Args:
|
||||
schedule_id: Schedule ID to delete.
|
||||
"""
|
||||
handle = self._client.get_schedule_handle(schedule_id)
|
||||
await handle.delete()
|
||||
|
||||
async def pause_schedule(self, schedule_id: str, note: str = "") -> None:
|
||||
"""Pause a schedule.
|
||||
|
||||
Args:
|
||||
schedule_id: Schedule ID to pause.
|
||||
note: Optional note about why paused.
|
||||
"""
|
||||
handle = self._client.get_schedule_handle(schedule_id)
|
||||
await handle.pause(note=note)
|
||||
|
||||
async def unpause_schedule(self, schedule_id: str, note: str = "") -> None:
|
||||
"""Unpause a schedule.
|
||||
|
||||
Args:
|
||||
schedule_id: Schedule ID to unpause.
|
||||
note: Optional note.
|
||||
"""
|
||||
handle = self._client.get_schedule_handle(schedule_id)
|
||||
await handle.unpause(note=note)
|
||||
|
||||
async def trigger_schedule(self, schedule_id: str) -> None:
|
||||
"""Trigger a schedule immediately.
|
||||
|
||||
Args:
|
||||
schedule_id: Schedule ID to trigger.
|
||||
"""
|
||||
handle = self._client.get_schedule_handle(schedule_id)
|
||||
await handle.trigger()
|
||||
|
||||
async def __aenter__(self) -> TasksClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
@@ -0,0 +1,350 @@
|
||||
"""Qdrant vector database client wrapper for Hanzo infrastructure.
|
||||
|
||||
Provides async interface to Qdrant for vector similarity search,
|
||||
supporting both cloud and local deployments.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional, Sequence
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class VectorConfig(BaseModel):
|
||||
"""Configuration for Qdrant vector database connection."""
|
||||
|
||||
host: str = Field(default="localhost", description="Qdrant server host")
|
||||
port: int = Field(default=6333, description="Qdrant REST API port")
|
||||
grpc_port: int = Field(default=6334, description="Qdrant gRPC port")
|
||||
api_key: Optional[str] = Field(default=None, description="API key for Qdrant Cloud")
|
||||
url: Optional[str] = Field(default=None, description="Full URL (overrides host/port)")
|
||||
https: bool = Field(default=False, description="Use HTTPS")
|
||||
timeout: float = Field(default=30.0, description="Request timeout in seconds")
|
||||
prefer_grpc: bool = Field(default=True, description="Prefer gRPC over REST")
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> VectorConfig:
|
||||
"""Create config from environment variables.
|
||||
|
||||
Environment variables:
|
||||
QDRANT_HOST: Server host (default: localhost)
|
||||
QDRANT_PORT: REST API port (default: 6333)
|
||||
QDRANT_GRPC_PORT: gRPC port (default: 6334)
|
||||
QDRANT_API_KEY: API key for authentication
|
||||
QDRANT_URL: Full URL (overrides host/port)
|
||||
QDRANT_HTTPS: Use HTTPS (default: false)
|
||||
"""
|
||||
return cls(
|
||||
host=os.getenv("QDRANT_HOST", "localhost"),
|
||||
port=int(os.getenv("QDRANT_PORT", "6333")),
|
||||
grpc_port=int(os.getenv("QDRANT_GRPC_PORT", "6334")),
|
||||
api_key=os.getenv("QDRANT_API_KEY"),
|
||||
url=os.getenv("QDRANT_URL"),
|
||||
https=os.getenv("QDRANT_HTTPS", "").lower() in ("true", "1", "yes"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VectorPoint:
|
||||
"""A point in vector space with optional payload."""
|
||||
|
||||
id: str | int
|
||||
vector: list[float]
|
||||
payload: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScoredPoint:
|
||||
"""A search result with similarity score."""
|
||||
|
||||
id: str | int
|
||||
score: float
|
||||
payload: dict[str, Any] = field(default_factory=dict)
|
||||
vector: Optional[list[float]] = None
|
||||
|
||||
|
||||
class VectorClient:
|
||||
"""Async client for Qdrant vector database.
|
||||
|
||||
Wraps qdrant-client with async methods and Hanzo conventions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = VectorClient(VectorConfig.from_env())
|
||||
await client.connect()
|
||||
|
||||
# Create collection
|
||||
await client.create_collection("embeddings", vector_size=1536)
|
||||
|
||||
# Upsert vectors
|
||||
points = [VectorPoint(id="doc1", vector=[0.1] * 1536, payload={"text": "hello"})]
|
||||
await client.upsert("embeddings", points)
|
||||
|
||||
# Search
|
||||
results = await client.search("embeddings", query_vector=[0.1] * 1536, limit=10)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[VectorConfig] = None) -> None:
|
||||
"""Initialize vector client.
|
||||
|
||||
Args:
|
||||
config: Qdrant configuration. If None, loads from environment.
|
||||
"""
|
||||
self.config = config or VectorConfig.from_env()
|
||||
self._client: Any = None
|
||||
self._async_client: Any = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Establish connection to Qdrant server."""
|
||||
try:
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"qdrant-client is required for VectorClient. "
|
||||
"Install with: pip install qdrant-client"
|
||||
) from e
|
||||
|
||||
if self.config.url:
|
||||
self._async_client = AsyncQdrantClient(
|
||||
url=self.config.url,
|
||||
api_key=self.config.api_key,
|
||||
timeout=self.config.timeout,
|
||||
prefer_grpc=self.config.prefer_grpc,
|
||||
)
|
||||
else:
|
||||
self._async_client = AsyncQdrantClient(
|
||||
host=self.config.host,
|
||||
port=self.config.port,
|
||||
grpc_port=self.config.grpc_port,
|
||||
api_key=self.config.api_key,
|
||||
https=self.config.https,
|
||||
timeout=self.config.timeout,
|
||||
prefer_grpc=self.config.prefer_grpc,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection."""
|
||||
if self._async_client:
|
||||
await self._async_client.close()
|
||||
self._async_client = None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if Qdrant server is healthy.
|
||||
|
||||
Returns:
|
||||
True if server is reachable and healthy.
|
||||
"""
|
||||
if not self._async_client:
|
||||
return False
|
||||
try:
|
||||
await self._async_client.get_collections()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def create_collection(
|
||||
self,
|
||||
name: str,
|
||||
vector_size: int,
|
||||
distance: str = "Cosine",
|
||||
on_disk: bool = False,
|
||||
) -> None:
|
||||
"""Create a new collection.
|
||||
|
||||
Args:
|
||||
name: Collection name.
|
||||
vector_size: Dimension of vectors.
|
||||
distance: Distance metric (Cosine, Euclid, Dot).
|
||||
on_disk: Store vectors on disk instead of RAM.
|
||||
"""
|
||||
from qdrant_client.models import Distance, VectorParams
|
||||
|
||||
distance_map = {
|
||||
"Cosine": Distance.COSINE,
|
||||
"Euclid": Distance.EUCLID,
|
||||
"Dot": Distance.DOT,
|
||||
}
|
||||
|
||||
await self._async_client.create_collection(
|
||||
collection_name=name,
|
||||
vectors_config=VectorParams(
|
||||
size=vector_size,
|
||||
distance=distance_map.get(distance, Distance.COSINE),
|
||||
on_disk=on_disk,
|
||||
),
|
||||
)
|
||||
|
||||
async def delete_collection(self, name: str) -> None:
|
||||
"""Delete a collection.
|
||||
|
||||
Args:
|
||||
name: Collection name to delete.
|
||||
"""
|
||||
await self._async_client.delete_collection(collection_name=name)
|
||||
|
||||
async def collection_exists(self, name: str) -> bool:
|
||||
"""Check if a collection exists.
|
||||
|
||||
Args:
|
||||
name: Collection name.
|
||||
|
||||
Returns:
|
||||
True if collection exists.
|
||||
"""
|
||||
return await self._async_client.collection_exists(collection_name=name)
|
||||
|
||||
async def list_collections(self) -> list[str]:
|
||||
"""List all collections.
|
||||
|
||||
Returns:
|
||||
List of collection names.
|
||||
"""
|
||||
result = await self._async_client.get_collections()
|
||||
return [c.name for c in result.collections]
|
||||
|
||||
async def upsert(
|
||||
self,
|
||||
collection: str,
|
||||
points: Sequence[VectorPoint],
|
||||
) -> None:
|
||||
"""Insert or update vectors.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
points: Points to upsert.
|
||||
"""
|
||||
from qdrant_client.models import PointStruct
|
||||
|
||||
qdrant_points = [
|
||||
PointStruct(id=p.id, vector=p.vector, payload=p.payload) for p in points
|
||||
]
|
||||
await self._async_client.upsert(collection_name=collection, points=qdrant_points)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
collection: str,
|
||||
query_vector: list[float],
|
||||
limit: int = 10,
|
||||
score_threshold: Optional[float] = None,
|
||||
filter_conditions: Optional[dict[str, Any]] = None,
|
||||
with_vectors: bool = False,
|
||||
) -> list[ScoredPoint]:
|
||||
"""Search for similar vectors.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
query_vector: Query vector.
|
||||
limit: Maximum results to return.
|
||||
score_threshold: Minimum similarity score.
|
||||
filter_conditions: Qdrant filter conditions.
|
||||
with_vectors: Include vectors in results.
|
||||
|
||||
Returns:
|
||||
List of scored points sorted by similarity.
|
||||
"""
|
||||
from qdrant_client.models import Filter
|
||||
|
||||
qdrant_filter = None
|
||||
if filter_conditions:
|
||||
qdrant_filter = Filter(**filter_conditions)
|
||||
|
||||
results = await self._async_client.search(
|
||||
collection_name=collection,
|
||||
query_vector=query_vector,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
query_filter=qdrant_filter,
|
||||
with_vectors=with_vectors,
|
||||
)
|
||||
|
||||
return [
|
||||
ScoredPoint(
|
||||
id=r.id,
|
||||
score=r.score,
|
||||
payload=r.payload or {},
|
||||
vector=r.vector if with_vectors else None,
|
||||
)
|
||||
for r in results
|
||||
]
|
||||
|
||||
async def delete(
|
||||
self,
|
||||
collection: str,
|
||||
ids: Optional[Sequence[str | int]] = None,
|
||||
filter_conditions: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Delete vectors by ID or filter.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
ids: Point IDs to delete.
|
||||
filter_conditions: Filter conditions for deletion.
|
||||
"""
|
||||
from qdrant_client.models import Filter, PointIdsList
|
||||
|
||||
if ids:
|
||||
await self._async_client.delete(
|
||||
collection_name=collection,
|
||||
points_selector=PointIdsList(points=list(ids)),
|
||||
)
|
||||
elif filter_conditions:
|
||||
await self._async_client.delete(
|
||||
collection_name=collection,
|
||||
points_selector=Filter(**filter_conditions),
|
||||
)
|
||||
|
||||
async def get(
|
||||
self,
|
||||
collection: str,
|
||||
ids: Sequence[str | int],
|
||||
with_vectors: bool = False,
|
||||
) -> list[VectorPoint]:
|
||||
"""Retrieve vectors by ID.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
ids: Point IDs to retrieve.
|
||||
with_vectors: Include vectors in results.
|
||||
|
||||
Returns:
|
||||
List of retrieved points.
|
||||
"""
|
||||
results = await self._async_client.retrieve(
|
||||
collection_name=collection,
|
||||
ids=list(ids),
|
||||
with_vectors=with_vectors,
|
||||
)
|
||||
|
||||
return [
|
||||
VectorPoint(
|
||||
id=r.id,
|
||||
vector=r.vector if with_vectors and r.vector else [],
|
||||
payload=r.payload or {},
|
||||
)
|
||||
for r in results
|
||||
]
|
||||
|
||||
async def count(self, collection: str) -> int:
|
||||
"""Count points in a collection.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
|
||||
Returns:
|
||||
Number of points.
|
||||
"""
|
||||
info = await self._async_client.get_collection(collection_name=collection)
|
||||
return info.points_count or 0
|
||||
|
||||
async def __aenter__(self) -> VectorClient:
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
Reference in New Issue
Block a user