llamactl 0.2.7a1__py3-none-any.whl → 0.3.0__py3-none-any.whl

llama_deploy/cli/__init__.py

@@ -1,26 +1,10 @@
-import click
-from .commands import projects, deployments, profile, health_check, serve
-from .options import global_options
+from llama_deploy.cli.commands.auth import auth
+from llama_deploy.cli.commands.deployment import deployments
+from llama_deploy.cli.commands.env import env_group
+from llama_deploy.cli.commands.init import init
+from llama_deploy.cli.commands.serve import serve
 
-
-# Main CLI application
-@click.group(help="LlamaDeploy CLI - Manage projects and deployments")
-@global_options
-def app():
-    """LlamaDeploy CLI - Manage projects and deployments"""
-    pass
-
-
-# Add sub-commands
-app.add_command(profile, name="profile")
-app.add_command(projects, name="project")
-app.add_command(deployments, name="deployment")
-
-# Add health check at root level
-app.add_command(health_check, name="health")
-
-# Add serve command at root level
-app.add_command(serve, name="serve")
+from .app import app
 
 
 # Main entry point function (called by the script)
@@ -28,5 +12,8 @@ def main() -> None:
     app()
 
 
+__all__ = ["app", "deployments", "auth", "serve", "init", "env_group"]
+
+
 if __name__ == "__main__":
     app()

llama_deploy/cli/app.py

@@ -0,0 +1,69 @@
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as pkg_version
+
+import click
+from llama_deploy.cli.commands.aliased_group import AliasedGroup
+from llama_deploy.cli.config.env_service import service
+from llama_deploy.cli.options import global_options
+from rich import print as rprint
+from rich.console import Console
+from rich.text import Text
+
+console = Console(highlight=False)
+
+
+def print_version(ctx: click.Context, param: click.Option, value: bool) -> None:
+    """Print the version of llama_deploy"""
+    if not value or ctx.resilient_parsing:
+        return
+    try:
+        ver = pkg_version("llamactl")
+        console.print(Text.assemble("client version: ", (ver, "green")))
+
+        # If there is an active profile, attempt to query server version
+        auth_service = service.current_auth_service()
+        if auth_service:
+            try:
+                data = auth_service.fetch_server_version()
+                server_ver = data.version
+                console.print(
+                    Text.assemble(
+                        "server version: ",
+                        (
+                            server_ver or "unknown",
+                            "bright_yellow" if server_ver is None else "green",
+                        ),
+                    )
+                )
+            except Exception as e:
+                console.print(
+                    Text.assemble(
+                        "server version: ",
+                        ("unavailable", "bright_yellow"),
+                        (f" - {e}", "dim"),
+                    )
+                )
+    except PackageNotFoundError:
+        rprint("[red]Package 'llamactl' not found[/red]")
+        raise click.Abort()
+    except Exception as e:
+        rprint(f"[red]Error: {e}[/red]")
+        raise click.Abort()
+    ctx.exit()
+
+
+# Main CLI application
+@click.group(
+    help="Create, develop, and deploy LlamaIndex workflow based apps", cls=AliasedGroup
+)
+@click.option(
+    "--version",
+    is_flag=True,
+    callback=print_version,
+    expose_value=False,
+    is_eager=True,
+    help="Print client and server versions of LlamaDeploy",
+)
+@global_options
+def app():
+    pass

llama_deploy/cli/auth/client.py

@@ -0,0 +1,362 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from types import TracebackType
+from typing import Any, AsyncContextManager, AsyncGenerator, Awaitable, Callable, Self
+
+import httpx
+import jwt
+from jwt.algorithms import RSAAlgorithm  # type: ignore[possibly-unbound-import]
+from llama_deploy.cli.config.schema import DeviceOIDC
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class OidcDiscoveryResponse(BaseModel):
+    discovery_url: str
+    client_ids: dict[str, str] | None = None
+
+
+class OidcProviderConfiguration(BaseModel):
+    device_authorization_endpoint: str | None = None
+    token_endpoint: str | None = None
+    scopes_supported: list[str] | None = None
+    jwks_uri: str | None = None
+
+
+class JsonWebKey(BaseModel):
+    kty: str
+    kid: str | None = None
+    use: str | None = None
+    alg: str | None = None
+    n: str | None = None
+    e: str | None = None
+    x5c: list[str] | None = None
+    x5t: str | None = None
+    x5t_s256: str | None = None
+
+
+class JsonWebKeySet(BaseModel):
+    keys: list[JsonWebKey]
+
+
+class AuthMeResponse(BaseModel):
+    id: str
+    email: str | None = None
+    last_login_provider: str | None = None
+    name: str | None = None
+    first_name: str | None = None
+    last_name: str | None = None
+    claims: dict[str, Any] | None = None
+    restrict: Any | None = None
+    created_at: str | None = None
+
+
+class ClientContextManager(AsyncContextManager):
+    def __init__(self, base_url: str | None, auth: httpx.Auth | None = None) -> None:
+        self.base_url = base_url.rstrip("/") if base_url else None
+        if self.base_url:
+            self.client = httpx.AsyncClient(base_url=self.base_url, auth=auth)
+        else:
+            self.client = httpx.AsyncClient(auth=auth)
+
+    async def close(self) -> None:
+        try:
+            await self.client.aclose()
+        except Exception:
+            pass
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await self.close()
+
+
+class PlatformAuthDiscoveryClient(ClientContextManager):
+    """Client for ad hoc auth endpoints under /api/v1/auth."""
+
+    def __init__(self, base_url: str) -> None:
+        super().__init__(base_url)
+
+    async def oidc_discovery(self) -> OidcDiscoveryResponse:
+        resp = await self.client.get("/api/v1/auth/oidc/discovery", timeout=10.0)
+        resp.raise_for_status()
+        return OidcDiscoveryResponse.model_validate(resp.json())
+
+
+class APIToken(BaseModel):
+    token: str
+    id: str
+
+
+class PlatformAuthClient(ClientContextManager):
+    """Client for user introspection under /api/v1/auth/me."""
+
+    def __init__(
+        self, base_url: str, id_token: str | None = None, auth: httpx.Auth | None = None
+    ) -> None:
+        self.id_token = id_token
+        super().__init__(base_url, auth=auth)
+
+    async def me(self) -> AuthMeResponse:
+        headers = (
+            {"Authorization": f"Bearer {self.id_token}"} if self.id_token else None
+        )
+        resp = await self.client.get("/api/v1/auth/me", headers=headers, timeout=10.0)
+        resp.raise_for_status()
+        return AuthMeResponse.model_validate(resp.json())
+
+    async def create_agent_api_key(self, name: str) -> APIToken:
+        resp = await self.client.post(
+            "/api/v1/api-keys",
+            json={"name": name, "project_id": None},
+        )
+        resp.raise_for_status()
+        json = resp.json()
+        token = json["redacted_api_key"]
+        id = json["id"]
+        return APIToken(token=token, id=id)
+
+    async def delete_api_key(self, id: str) -> None:
+        response = await self.client.delete(f"/api/v1/api-keys/{id}")
+        response.raise_for_status()
+
+
+class RefreshMiddleware(httpx.Auth):
+    def __init__(
+        self,
+        device_oidc: DeviceOIDC,
+        on_refresh: Callable[[DeviceOIDC], Awaitable[None]],
+    ) -> None:
+        self.device_oidc = device_oidc
+        self.on_refresh = on_refresh
+        self.lock = asyncio.Lock()
+
+    async def _refresh_and_update(self) -> None:
+        new_device_oidc = await refresh(self.device_oidc)
+        self.device_oidc = new_device_oidc
+        try:
+            await self.on_refresh(new_device_oidc)
+        except Exception:
+            logger.exception("Error in on_refresh callback")
+
+    async def async_auth_flow(
+        self, request: httpx.Request
+    ) -> AsyncGenerator[httpx.Request, httpx.Response]:
+        token = self.device_oidc.device_access_token
+        request.headers["Authorization"] = f"Bearer {token}"
+
+        response = yield request
+        if response.status_code == 401:
+            async with self.lock:
+                if token == self.device_oidc.device_access_token:
+                    await self._refresh_and_update()
+                    request.headers["Authorization"] = (
+                        f"Bearer {self.device_oidc.device_access_token}"
+                    )
+                yield request
+
+
+class DeviceAuthorizationRequest(BaseModel):
+    client_id: str
+    scope: str
+
+
+class DeviceAuthorizationResponse(BaseModel):
+    device_code: str
+    user_code: str
+    verification_uri: str
+    verification_uri_complete: str | None = None
+    expires_in: int
+    interval: int | None = None
+
+
+class TokenRequestDeviceCode(BaseModel):
+    grant_type: str = "urn:ietf:params:oauth:grant-type:device_code"
+    device_code: str
+    client_id: str
+
+
+class TokenResponse(BaseModel):
+    # Success fields
+    id_token: str | None = None
+    access_token: str | None = None
+    refresh_token: str | None = None
+    expires_in: int | None = None
+    token_type: str | None = None
+    scope: str | None = None
+    # Error fields
+    error: str | None = None
+    error_description: str | None = None
+
+
+class TokenRequestRefresh(BaseModel):
+    grant_type: str = "refresh_token"
+    refresh_token: str
+    client_id: str
+
+
+class OIDCClient(ClientContextManager):
+    def __init__(self) -> None:
+        super().__init__(None)
+
+    async def fetch_provider_configuration(
+        self, discovery_url: str
+    ) -> OidcProviderConfiguration:
+        resp = await self.client.get(discovery_url, timeout=10.0)
+        resp.raise_for_status()
+        return OidcProviderConfiguration.model_validate(resp.json())
+
+    async def device_authorization(
+        self, device_endpoint: str, request: DeviceAuthorizationRequest
+    ) -> DeviceAuthorizationResponse:
+        resp = await self.client.post(
+            device_endpoint,
+            data=request.model_dump(),
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
+            timeout=10.0,
+        )
+        resp.raise_for_status()
+        return DeviceAuthorizationResponse.model_validate(resp.json())
+
+    async def token_with_device_code(
+        self, token_endpoint: str, request: TokenRequestDeviceCode
+    ) -> TokenResponse:
+        resp = await self.client.post(
+            token_endpoint,
+            data=request.model_dump(),
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
+            timeout=10.0,
+        )
+        # Do not raise for status; callers inspect error payloads during polling
+        try:
+            payload = resp.json()
+        except Exception:
+            # Fall back to minimal error information
+            return TokenResponse(error="invalid_response", error_description=resp.text)
+        return TokenResponse.model_validate(payload)
+
+    async def token_with_refresh(
+        self, token_endpoint: str, request: TokenRequestRefresh
+    ) -> TokenResponse:
+        resp = await self.client.post(
+            token_endpoint,
+            data=request.model_dump(),
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
+            timeout=10.0,
+        )
+        try:
+            payload = resp.json()
+        except Exception:
+            return TokenResponse(error="invalid_response", error_description=resp.text)
+        return TokenResponse.model_validate(payload)
+
+    async def get_jwks(self, jwks_uri: str) -> JsonWebKeySet:
+        resp = await self.client.get(jwks_uri, timeout=10.0)
+        resp.raise_for_status()
+        return JsonWebKeySet.model_validate(resp.json())
+
+
+async def decode_jwt_claims_from_device_oidc(
+    oidc_device: DeviceOIDC,
+    verify_audience: bool = False,
+    verify_expiration: bool = False,
+    audience: str | None = None,
+) -> dict[str, Any]:
+    """Decode JWT claims by discovering provider and verifying via JWKS.
+
+    Assumes RSA signing. Audience verification can be toggled and, when enabled,
+    an audience value can be provided.
+    """
+    if not oidc_device.device_id_token:
+        raise ValueError("Device ID token is missing. Cannot decode claims.")
+    async with OIDCClient() as oidc:
+        provider = await oidc.fetch_provider_configuration(oidc_device.discovery_url)
+        jwks_uri = provider.jwks_uri
+        if not jwks_uri:
+            raise ValueError("Provider does not expose jwks_uri")
+    return await decode_jwt_claims(
+        oidc_device.device_id_token,
+        jwks_uri,
+        verify_audience,
+        verify_expiration,
+        audience,
+    )
+
+
+async def decode_jwt_claims(
+    token: str,
+    jwks_uri: str,
+    verify_audience: bool = False,
+    verify_expiration: bool = False,
+    audience: str | None = None,
+) -> dict[str, Any]:
+    async with OIDCClient() as oidc:
+        jwks = await oidc.get_jwks(jwks_uri)
+
+    # Select key
+    header = jwt.get_unverified_header(token)
+    kid = header.get("kid")
+    alg = header.get("alg", "RS256")
+    keys = jwks.keys
+    key = next((k for k in keys if k.kid == kid), None) or next(iter(keys), None)
+    if not key:
+        raise ValueError("Signing key not found in JWKS")
+
+    # Build public key (RSA-only)
+    if key.kty != "RSA":
+        raise ValueError("Unsupported JWK kty; only RSA is supported")
+    key_json = key.model_dump_json()
+    public_key = RSAAlgorithm.from_jwk(key_json)
+
+    return jwt.decode(
+        token,
+        public_key,
+        algorithms=[alg],
+        options={"verify_aud": verify_audience, "verify_exp": verify_expiration},
+        audience=audience,
+    )
+
+
+async def refresh(device_oidc: DeviceOIDC) -> DeviceOIDC:
+    """
+    Run a refresh on the access token, storing updated tokens in a new DeviceOIDC.
+    """
+    async with OIDCClient() as oidc:
+        provider = await oidc.fetch_provider_configuration(device_oidc.discovery_url)
+        token_endpoint = provider.token_endpoint
+        if not token_endpoint:
+            raise ValueError("Provider does not expose token_endpoint")
+        if not device_oidc.device_refresh_token:
+            raise ValueError("Device refresh token is missing. Cannot refresh.")
+        token = await oidc.token_with_refresh(
+            token_endpoint,
+            TokenRequestRefresh(
+                refresh_token=device_oidc.device_refresh_token,
+                client_id=device_oidc.client_id,
+            ),
+        )
+        copy = device_oidc.model_copy()
+        if not token.access_token:
+            raise ValueError("Refresh failed: token response missing access_token")
+        copy.device_access_token = token.access_token
+        copy.device_refresh_token = token.refresh_token or copy.device_refresh_token
+        copy.device_id_token = token.id_token or copy.device_id_token
+        return copy

llama_deploy/cli/client.py

@@ -1,173 +1,50 @@
-import logging
-from typing import List, Optional
-
-import httpx
-from llama_deploy.core.schema.deployments import (
-    DeploymentCreate,
-    DeploymentResponse,
-    DeploymentsListResponse,
-    DeploymentUpdate,
-)
-from llama_deploy.core.schema.git_validation import (
-    RepositoryValidationRequest,
-    RepositoryValidationResponse,
-)
-from llama_deploy.core.schema.projects import ProjectSummary, ProjectsListResponse
-from rich.console import Console
-
-from .config import config_manager
-
-
-class LlamaDeployClient:
-    """HTTP client for communicating with the LlamaDeploy control plane API"""
-
-    def __init__(
-        self, base_url: Optional[str] = None, project_id: Optional[str] = None
-    ):
-        """Initialize the client with a configured profile"""
-        self.console = Console()
-
-        # Get profile data
-        profile = config_manager.get_current_profile()
-        if not profile:
-            self.console.print("\n[bold red]No profile configured![/bold red]")
-            self.console.print("\nTo get started, create a profile with:")
-            self.console.print("[cyan]llamactl profile create[/cyan]")
-            raise SystemExit(1)
-
-        # Use profile data with optional overrides
-        self.base_url = base_url or profile.api_url
-        self.project_id = project_id or profile.active_project_id
-
-        if not self.base_url:
-            raise ValueError("API URL is required")
-
-        if not self.project_id:
-            raise ValueError("Project ID is required")
-
-        self.base_url = self.base_url.rstrip("/")
-
-        # Create persistent client with event hooks
-        self.client = httpx.Client(
-            base_url=self.base_url, event_hooks={"response": [self._handle_response]}
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+
+from llama_deploy.cli.config.env_service import service
+from llama_deploy.core.client.manage_client import ControlPlaneClient, ProjectClient
+from rich import print as rprint
+
+
+def get_control_plane_client() -> ControlPlaneClient:
+    auth_svc = service.current_auth_service()
+    profile = service.current_auth_service().get_current_profile()
+    if profile:
+        resolved_base_url = profile.api_url.rstrip("/")
+        resolved_api_key = profile.api_key
+        return ControlPlaneClient(
+            resolved_base_url, resolved_api_key, auth_svc.auth_middleware()
         )
 
-    def _handle_response(self, response: httpx.Response) -> None:
-        """Handle response middleware - warnings and error conversion"""
-        # Check for warnings in response headers
-        if "X-Warning" in response.headers:
-            self.console.print(
-                f"[yellow]Warning: {response.headers['X-Warning']}[/yellow]"
-            )
-
-        # Convert httpx errors to our current exception format
+    # Fallback: allow env-scoped client construction for env operations
+    env = service.get_current_environment()
+    resolved_base_url = env.api_url.rstrip("/")
+    return ControlPlaneClient(resolved_base_url)
+
+
+def get_project_client() -> ProjectClient:
+    auth_svc = service.current_auth_service()
+    profile = auth_svc.get_current_profile()
+    if not profile:
+        rprint("\n[bold red]No profile configured![/bold red]")
+        rprint("\nTo get started, create a profile with:")
+        if auth_svc.env.requires_auth:
+            rprint("[cyan]llamactl auth login[/cyan]")
+        else:
+            rprint("[cyan]llamactl auth token[/cyan]")
+        raise SystemExit(1)
+    return ProjectClient(
+        profile.api_url, profile.project_id, profile.api_key, auth_svc.auth_middleware()
+    )
+
+
+@asynccontextmanager
+async def project_client_context() -> AsyncGenerator[ProjectClient, None]:
+    client = get_project_client()
+    try:
+        yield client
+    finally:
         try:
-            response.raise_for_status()
-        except httpx.HTTPStatusError as e:
-            # Try to parse JSON error response
-            try:
-                response.read()  # need to collect streaming data before calling json
-                error_data = e.response.json()
-                if isinstance(error_data, dict) and "detail" in error_data:
-                    error_message = error_data["detail"]
-                else:
-                    error_message = str(error_data)
-            except (ValueError, KeyError):
-                # Fallback to raw response text
-                error_message = e.response.text
-
-            raise Exception(f"HTTP {e.response.status_code}: {error_message}") from e
-        except httpx.RequestError as e:
-            raise Exception(f"Request failed: {e}") from e
-
-    # Health check
-    def health_check(self) -> dict:
-        """Check if the API server is healthy"""
-        response = self.client.get("/health")
-        return response.json()
-
-    # Projects
-    def list_projects(self) -> List[ProjectSummary]:
-        """List all projects with deployment counts"""
-        response = self.client.get("/projects/")
-        projects_response = ProjectsListResponse.model_validate(response.json())
-        return [project for project in projects_response.projects]
-
-    # Deployments
-    def list_deployments(self) -> List[DeploymentResponse]:
-        """List deployments for the configured project"""
-        response = self.client.get(f"/{self.project_id}/deployments/")
-        deployments_response = DeploymentsListResponse.model_validate(response.json())
-        return [deployment for deployment in deployments_response.deployments]
-
-    def get_deployment(self, deployment_id: str) -> DeploymentResponse:
-        """Get a specific deployment"""
-        response = self.client.get(f"/{self.project_id}/deployments/{deployment_id}")
-        deployment = DeploymentResponse.model_validate(response.json())
-        return deployment
-
-    def create_deployment(
-        self,
-        deployment_data: DeploymentCreate,
-    ) -> DeploymentResponse:
-        """Create a new deployment"""
-
-        response = self.client.post(
-            f"/{self.project_id}/deployments/",
-            json=deployment_data.model_dump(exclude_none=True),
-        )
-        deployment = DeploymentResponse.model_validate(response.json())
-        return deployment
-
-    def delete_deployment(self, deployment_id: str) -> None:
-        """Delete a deployment"""
-        self.client.delete(f"/{self.project_id}/deployments/{deployment_id}")
-
-    def update_deployment(
-        self,
-        deployment_id: str,
-        update_data: DeploymentUpdate,
-        force_git_sha_update: bool = False,
-    ) -> DeploymentResponse:
-        """Update an existing deployment"""
-
-        params = {}
-        if force_git_sha_update:
-            params["force_git_sha_update"] = True
-
-        response = self.client.patch(
-            f"/{self.project_id}/deployments/{deployment_id}",
-            json=update_data.model_dump(),
-            params=params,
-        )
-        deployment = DeploymentResponse.model_validate(response.json())
-        return deployment
-
-    def validate_repository(
-        self,
-        repo_url: str,
-        deployment_id: str | None = None,
-        pat: str | None = None,
-    ) -> RepositoryValidationResponse:
-        """Validate a repository URL"""
-        logging.info(
-            f"Validating repository with params: {repo_url}, {deployment_id}, {pat}"
-        )
-        response = self.client.post(
-            f"/{self.project_id}/deployments/validate-repository",
-            json=RepositoryValidationRequest(
-                repository_url=repo_url,
-                deployment_id=deployment_id,
-                pat=pat,
-            ).model_dump(),
-        )
-        logging.info(f"Response: {response.json()}")
-        return RepositoryValidationResponse.model_validate(response.json())
-
-
-# Global client factory function
-def get_client(
-    base_url: Optional[str] = None, project_id: Optional[str] = None
-) -> LlamaDeployClient:
-    """Get a client instance with optional overrides"""
-    return LlamaDeployClient(base_url=base_url, project_id=project_id)
+            await client.aclose()
+        except Exception:
+            pass