agv 0.1.0__py3-none-any.whl

agentenv/__init__.py

@@ -0,0 +1,26 @@
+"""AgentEnv Cloud CLI - Manage containers, sandboxes, and more."""
+
+from ._version import __version__
+
+# SDK exports
+from .sdk import Client, connect, SandboxResource
+from .sdk.ai import AIClient
+from .sdk.cluster import ray_init, spark_init, remote
+from .sdk.function import function, RemoteFunctionError
+from .sdk.image import ImageBuilder, image, py
+
+__all__ = [
+    "Client",
+    "connect",
+    "SandboxResource",
+    "AIClient",
+    "ImageBuilder",
+    "image",
+    "py",
+    "function",
+    "RemoteFunctionError",
+    "ray_init",
+    "spark_init",
+    "remote",
+    "__version__",
+]

agentenv/__main__.py

@@ -0,0 +1,6 @@
+"""Entry point for `python -m agentenv`."""
+
+from .cli.app import main_cli
+
+if __name__ == "__main__":
+    main_cli()

agentenv/_version.py

@@ -0,0 +1,3 @@
+"""Single source of truth for the package version."""
+
+__version__ = "0.1.0"

agentenv/api/__init__.py

@@ -0,0 +1,23 @@
+"""API client for AgentEnv Cloud."""
+
+from .analytics import AnalyticsAPI
+from .models import (
+    AnalyticsInstallation,
+    AnalyticsFunnel,
+    AnalyticsFunnelStep,
+    AnalyticsUserProfile,
+    AnalyticsEvent,
+    CreateAnalyticsInstallationRequest,
+    CreateAnalyticsFunnelRequest,
+)
+
+__all__ = [
+    "AnalyticsAPI",
+    "AnalyticsInstallation",
+    "AnalyticsFunnel",
+    "AnalyticsFunnelStep",
+    "AnalyticsUserProfile",
+    "AnalyticsEvent",
+    "CreateAnalyticsInstallationRequest",
+    "CreateAnalyticsFunnelRequest",
+]

agentenv/api/ai_gateway.py

@@ -0,0 +1,322 @@
+"""AI Gateway API client."""
+
+from typing import Any, Iterator
+
+from .client import Client
+from .models import (
+    GatewayUpstream,
+    GatewayPool,
+    GatewayApiKey,
+    CreateGatewayApiKeyRequest,
+    GatewayUsageReport,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionChunk,
+    EmbeddingRequest,
+    EmbeddingResponse,
+)
+
+
+class AIGatewayAPI:
+    """API client for AI Gateway resources."""
+
+    def __init__(self, client: Client, gateway_key: str | None = None):
+        """Initialize AI Gateway API client.
+
+        Args:
+            client: HTTP client instance (for management APIs with JWT)
+            gateway_key: Gateway API key (for proxy APIs)
+        """
+        self.client = client
+        self.gateway_key = gateway_key
+        self.chat = ChatCompletionsAPI(client, gateway_key)
+        self.embeddings = EmbeddingsAPI(client, gateway_key)
+
+    # -------------------------------------------------------------------------
+    # Upstreams
+    # -------------------------------------------------------------------------
+
+    def list_upstreams(self) -> list[GatewayUpstream]:
+        """List all upstream providers.
+
+        Returns:
+            List of upstream configurations
+        """
+        response = self.client.get("/ai-gateway/upstreams")
+        data = response.json()
+
+        if isinstance(data, list):
+            return [GatewayUpstream(**item) for item in data]
+        return []
+
+    def create_upstream(
+        self,
+        name: str,
+        provider: str,
+        api_key: str,
+        base_url: str | None = None,
+        rate_limit_rpm: int | None = None,
+        rate_limit_tpm: int | None = None,
+        priority: int = 0,
+        weight: int = 1,
+    ) -> GatewayUpstream:
+        """Create a new upstream provider."""
+        data: dict[str, Any] = {
+            "name": name,
+            "provider": provider,
+            "apiKey": api_key,
+            "priority": priority,
+            "weight": weight,
+        }
+        if base_url:
+            data["baseUrl"] = base_url
+        if rate_limit_rpm:
+            data["rateLimitRpm"] = rate_limit_rpm
+        if rate_limit_tpm:
+            data["rateLimitTpm"] = rate_limit_tpm
+
+        response = self.client.post("/ai-gateway/upstreams", json=data)
+        return GatewayUpstream(**response.json())
+
+    def get_upstream(self, upstream_id: str) -> GatewayUpstream:
+        """Get upstream details."""
+        response = self.client.get(f"/ai-gateway/upstreams/{upstream_id}")
+        return GatewayUpstream(**response.json())
+
+    def delete_upstream(self, upstream_id: str) -> None:
+        """Delete an upstream provider."""
+        self.client.delete(f"/ai-gateway/upstreams/{upstream_id}")
+
+    def test_upstream(self, upstream_id: str) -> dict[str, Any]:
+        """Test upstream connectivity."""
+        response = self.client.post(f"/ai-gateway/upstreams/{upstream_id}/test")
+        return response.json()
+
+    # -------------------------------------------------------------------------
+    # Pools
+    # -------------------------------------------------------------------------
+
+    def list_pools(self) -> list[GatewayPool]:
+        """List all pools."""
+        response = self.client.get("/ai-gateway/pools")
+        data = response.json()
+
+        if isinstance(data, list):
+            return [GatewayPool(**item) for item in data]
+        return []
+
+    def create_pool(
+        self,
+        name: str,
+        strategy: str = "round-robin",
+        model_mappings: dict[str, str] | None = None,
+    ) -> GatewayPool:
+        """Create a new pool."""
+        data: dict[str, Any] = {
+            "name": name,
+            "strategy": strategy,
+        }
+        if model_mappings:
+            data["modelMappings"] = model_mappings
+
+        response = self.client.post("/ai-gateway/pools", json=data)
+        return GatewayPool(**response.json())
+
+    def delete_pool(self, pool_id: str) -> None:
+        """Delete a pool."""
+        self.client.delete(f"/ai-gateway/pools/{pool_id}")
+
+    def add_pool_member(
+        self,
+        pool_id: str,
+        upstream_id: str,
+        priority: int = 0,
+        weight: int = 1,
+    ) -> GatewayPool:
+        """Add member to pool."""
+        data = {
+            "upstreamId": upstream_id,
+            "priority": priority,
+            "weight": weight,
+        }
+        response = self.client.post(f"/ai-gateway/pools/{pool_id}/members", json=data)
+        return GatewayPool(**response.json())
+
+    # -------------------------------------------------------------------------
+    # API Keys
+    # -------------------------------------------------------------------------
+
+    def list_keys(self) -> list[GatewayApiKey]:
+        """List all API keys."""
+        response = self.client.get("/ai-gateway/keys")
+        data = response.json()
+
+        if isinstance(data, list):
+            return [GatewayApiKey(**item) for item in data]
+        return []
+
+    def create_key(self, request: CreateGatewayApiKeyRequest) -> GatewayApiKey:
+        """Create a new API key."""
+        response = self.client.post(
+            "/ai-gateway/keys",
+            json=request.model_dump(exclude_none=True, by_alias=True),
+        )
+        return GatewayApiKey(**response.json())
+
+    def revoke_key(self, key_id: str) -> None:
+        """Revoke an API key."""
+        self.client.delete(f"/ai-gateway/keys/{key_id}")
+
+    # -------------------------------------------------------------------------
+    # Discovery
+    # -------------------------------------------------------------------------
+
+    def list_providers(self) -> dict[str, Any]:
+        """List supported gateway providers."""
+        response = self.client.get("/ai-gateway/providers")
+        return response.json()
+
+    def list_platform_models(self) -> dict[str, Any]:
+        """List platform models grouped by provider."""
+        response = self.client.get("/ai-gateway/platform-models")
+        return response.json()
+
+    # -------------------------------------------------------------------------
+    # Usage
+    # -------------------------------------------------------------------------
+
+    def get_usage(
+        self,
+        start_date: str | None = None,
+        end_date: str | None = None,
+        upstream_id: str | None = None,
+    ) -> GatewayUsageReport:
+        """Get usage report."""
+        params: dict[str, Any] = {}
+        if start_date:
+            params["startDate"] = start_date
+        if end_date:
+            params["endDate"] = end_date
+        if upstream_id:
+            params["upstreamId"] = upstream_id
+
+        response = self.client.get("/ai-gateway/usage", params=params)
+        return GatewayUsageReport(**response.json())
+
+
+class ChatCompletionsAPI:
+    """Chat completions API (OpenAI-compatible endpoint)."""
+
+    def __init__(self, client: Client, gateway_key: str | None = None):
+        self.client = client
+        self.gateway_key = gateway_key
+
+    def _auth_headers(self) -> dict[str, str]:
+        """Get authentication headers for proxy API."""
+        if self.gateway_key:
+            return {"x-gateway-key": self.gateway_key}
+        return {}
+
+    def create(
+        self,
+        request: ChatCompletionRequest | dict[str, Any],
+    ) -> ChatCompletionResponse | Iterator[ChatCompletionChunk]:
+        """Create chat completion."""
+        if isinstance(request, ChatCompletionRequest):
+            data = request.model_dump(exclude_none=True, by_alias=True)
+        else:
+            data = request
+
+        headers = self._auth_headers()
+
+        if data.get("stream"):
+            return self._stream_create(data, headers)
+
+        response = self.client.post(
+            "/ai-gateway/openai/v1/chat/completions",
+            json=data,
+            headers=headers,
+        )
+        return ChatCompletionResponse(**response.json())
+
+    def _stream_create(
+        self,
+        data: dict[str, Any],
+        headers: dict[str, str],
+    ) -> Iterator[ChatCompletionChunk]:
+        """Create streaming chat completion."""
+        # Use stream_request from StreamingClient if available
+        if hasattr(self.client, "stream_request"):
+            response = self.client.stream_request(
+                "POST",
+                "/ai-gateway/openai/v1/chat/completions",
+                json=data,
+                headers=headers,
+            )
+        else:
+            # Fallback to regular client with streaming
+            import httpx
+
+            with httpx.Client() as http_client:
+                url = f"{self.client.base_url}/ai-gateway/openai/v1/chat/completions"
+                headers = {**headers, "Content-Type": "application/json"}
+
+                with http_client.stream("POST", url, json=data, headers=headers) as response:
+                    for line in response.iter_lines():
+                        if line.startswith("data: "):
+                            data_str = line[6:]
+                            if data_str == "[DONE]":
+                                break
+                            try:
+                                chunk = ChatCompletionChunk.model_validate_json(data_str)
+                                yield chunk
+                            except Exception:  # noqa: S110 - Skip invalid chunks
+                                continue
+                    return
+
+        # Handle response from stream_request
+        for line in response.iter_lines():
+            if isinstance(line, bytes):
+                line = line.decode("utf-8")
+            if line.startswith("data: "):
+                data_str = line[6:]
+                if data_str == "[DONE]":
+                    break
+                try:
+                    chunk = ChatCompletionChunk.model_validate_json(data_str)
+                    yield chunk
+                except Exception:  # noqa: S110 - Skip invalid chunks
+                    continue
+
+
+class EmbeddingsAPI:
+    """Embeddings API (OpenAI-compatible endpoint)."""
+
+    def __init__(self, client: Client, gateway_key: str | None = None):
+        self.client = client
+        self.gateway_key = gateway_key
+
+    def _auth_headers(self) -> dict[str, str]:
+        """Get authentication headers for proxy API."""
+        if self.gateway_key:
+            return {"x-gateway-key": self.gateway_key}
+        return {}
+
+    def create(
+        self,
+        request: EmbeddingRequest | dict[str, Any],
+    ) -> EmbeddingResponse:
+        """Create embeddings."""
+        if isinstance(request, EmbeddingRequest):
+            data = request.model_dump(exclude_none=True, by_alias=True)
+        else:
+            data = request
+
+        headers = self._auth_headers()
+
+        response = self.client.post(
+            "/ai-gateway/openai/v1/embeddings",
+            json=data,
+            headers=headers,
+        )
+        return EmbeddingResponse(**response.json())