beamlit 0.0.11__py3-none-any.whl

beamlit/api/workspaces/update_workspace_user_role.py

@@ -0,0 +1,187 @@
+from http import HTTPStatus
+from typing import Any, Optional, Union, cast
+
+import httpx
+
+from ... import errors
+from ...client import AuthenticatedClient, Client
+from ...models.update_workspace_user_role_body import UpdateWorkspaceUserRoleBody
+from ...models.workspace_user import WorkspaceUser
+from ...types import Response
+
+
+def _get_kwargs(
+    sub_or_email: str,
+    *,
+    body: UpdateWorkspaceUserRoleBody,
+) -> dict[str, Any]:
+    headers: dict[str, Any] = {}
+
+    _kwargs: dict[str, Any] = {
+        "method": "put",
+        "url": f"/users/{sub_or_email}",
+    }
+
+    _body = body.to_dict()
+
+    _kwargs["json"] = _body
+    headers["Content-Type"] = "application/json"
+
+    _kwargs["headers"] = headers
+    return _kwargs
+
+
+def _parse_response(
+    *, client: Union[AuthenticatedClient, Client], response: httpx.Response
+) -> Optional[Union[Any, WorkspaceUser]]:
+    if response.status_code == 200:
+        response_200 = WorkspaceUser.from_dict(response.json())
+
+        return response_200
+    if response.status_code == 400:
+        response_400 = cast(Any, None)
+        return response_400
+    if response.status_code == 404:
+        response_404 = cast(Any, None)
+        return response_404
+    if client.raise_on_unexpected_status:
+        raise errors.UnexpectedStatus(response.status_code, response.content)
+    else:
+        return None
+
+
+def _build_response(
+    *, client: Union[AuthenticatedClient, Client], response: httpx.Response
+) -> Response[Union[Any, WorkspaceUser]]:
+    return Response(
+        status_code=HTTPStatus(response.status_code),
+        content=response.content,
+        headers=response.headers,
+        parsed=_parse_response(client=client, response=response),
+    )
+
+
+def sync_detailed(
+    sub_or_email: str,
+    *,
+    client: AuthenticatedClient,
+    body: UpdateWorkspaceUserRoleBody,
+) -> Response[Union[Any, WorkspaceUser]]:
+    """Update user role in workspace
+
+     Updates the role of a user in the workspace.
+
+    Args:
+        sub_or_email (str):
+        body (UpdateWorkspaceUserRoleBody):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[Union[Any, WorkspaceUser]]
+    """
+
+    kwargs = _get_kwargs(
+        sub_or_email=sub_or_email,
+        body=body,
+    )
+
+    response = client.get_httpx_client().request(
+        **kwargs,
+    )
+
+    return _build_response(client=client, response=response)
+
+
+def sync(
+    sub_or_email: str,
+    *,
+    client: AuthenticatedClient,
+    body: UpdateWorkspaceUserRoleBody,
+) -> Optional[Union[Any, WorkspaceUser]]:
+    """Update user role in workspace
+
+     Updates the role of a user in the workspace.
+
+    Args:
+        sub_or_email (str):
+        body (UpdateWorkspaceUserRoleBody):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Union[Any, WorkspaceUser]
+    """
+
+    return sync_detailed(
+        sub_or_email=sub_or_email,
+        client=client,
+        body=body,
+    ).parsed
+
+
+async def asyncio_detailed(
+    sub_or_email: str,
+    *,
+    client: AuthenticatedClient,
+    body: UpdateWorkspaceUserRoleBody,
+) -> Response[Union[Any, WorkspaceUser]]:
+    """Update user role in workspace
+
+     Updates the role of a user in the workspace.
+
+    Args:
+        sub_or_email (str):
+        body (UpdateWorkspaceUserRoleBody):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[Union[Any, WorkspaceUser]]
+    """
+
+    kwargs = _get_kwargs(
+        sub_or_email=sub_or_email,
+        body=body,
+    )
+
+    response = await client.get_async_httpx_client().request(**kwargs)
+
+    return _build_response(client=client, response=response)
+
+
+async def asyncio(
+    sub_or_email: str,
+    *,
+    client: AuthenticatedClient,
+    body: UpdateWorkspaceUserRoleBody,
+) -> Optional[Union[Any, WorkspaceUser]]:
+    """Update user role in workspace
+
+     Updates the role of a user in the workspace.
+
+    Args:
+        sub_or_email (str):
+        body (UpdateWorkspaceUserRoleBody):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Union[Any, WorkspaceUser]
+    """
+
+    return (
+        await asyncio_detailed(
+            sub_or_email=sub_or_email,
+            client=client,
+            body=body,
+        )
+    ).parsed

beamlit/authentication/__init__.py

@@ -0,0 +1,24 @@
+from .apikey import ApiKeyProvider
+from .authentication import (PublicProvider, RunClientWithCredentials,
+                             new_client_with_credentials)
+from .credentials import (Config, ContextConfig, Credentials, WorkspaceConfig,
+                          load_credentials)
+from .device_mode import (BearerToken, DeviceLogin, DeviceLoginFinalizeRequest,
+                          DeviceLoginFinalizeResponse, DeviceLoginResponse)
+
+__all__ = (
+    "ApiKeyProvider",
+    "PublicProvider",
+    "RunClientWithCredentials",
+    "new_client_with_credentials",
+    "Config",
+    "ContextConfig",
+    "Credentials",
+    "WorkspaceConfig",
+    "load_credentials",
+    "BearerToken",
+    "DeviceLogin",
+    "DeviceLoginFinalizeRequest",
+    "DeviceLoginFinalizeResponse",
+    "DeviceLoginResponse"
+)

beamlit/authentication/apikey.py

@@ -0,0 +1,14 @@
+from typing import Generator
+
+from httpx import Auth, Request, Response
+
+
+class ApiKeyProvider(Auth):
+    def __init__(self, credentials, workspace_name: str):
+        self.credentials = credentials
+        self.workspace_name = workspace_name
+
+    def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
+        request.headers['X-Beamlit-Api-Key'] = self.credentials.api_key
+        request.headers['X-Beamlit-Workspace'] = self.workspace_name
+        yield request

beamlit/authentication/authentication.py

@@ -0,0 +1,35 @@
+from dataclasses import dataclass
+from typing import Generator
+
+from httpx import Auth, Request, Response
+
+from ..client import AuthenticatedClient
+from .apikey import ApiKeyProvider
+from .credentials import Credentials
+from .device_mode import BearerToken
+
+
+class PublicProvider(Auth):
+    def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
+        yield request
+
+
+
+@dataclass
+class RunClientWithCredentials:
+    credentials: Credentials
+    workspace: str
+    api_url: str = "https://api.beamlit.dev/v0"
+    run_url: str = "https://run.beamlit.dev/v0"
+
+
+def new_client_with_credentials(config: RunClientWithCredentials):
+    provider: Auth = None
+    if config.credentials.api_key:
+        provider = ApiKeyProvider(config.credentials, config.workspace)
+    elif config.credentials.access_token:
+        provider = BearerToken(config.credentials, config.workspace, config.api_url)
+    else:
+        provider = PublicProvider()
+
+    return AuthenticatedClient(base_url=config.api_url, provider=provider)

beamlit/authentication/credentials.py

@@ -0,0 +1,165 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List
+
+import yaml
+
+
+@dataclass
+class Credentials:
+    api_key: str = ""
+    access_token: str = ""
+    refresh_token: str = ""
+    expires_in: int = 0
+    device_code: str = ""
+
+@dataclass
+class WorkspaceConfig:
+    name: str
+    credentials: Credentials
+
+
+@dataclass
+class ContextConfig:
+    workspace: str = ""
+    environment: str = ""
+
+
+@dataclass
+class Config:
+    workspaces: List[WorkspaceConfig] = None
+    context: ContextConfig = None
+
+    def __post_init__(self):
+        if self.workspaces is None:
+            self.workspaces = []
+        if self.context is None:
+            self.context = ContextConfig()
+
+
+def load_config() -> Config:
+    config = Config()
+    home_dir = Path.home()
+    if home_dir:
+        config_path = home_dir / ".beamlit" / "config.yaml"
+        if config_path.exists():
+            try:
+                with open(config_path) as f:
+                    data = yaml.safe_load(f)
+                    if data:
+                        workspaces = []
+                        for ws in data.get("workspaces", []):
+                            creds = Credentials(**ws.get("credentials", {}))
+                            workspaces.append(WorkspaceConfig(name=ws["name"], credentials=creds))
+                        config.workspaces = workspaces
+                        if "context" in data:
+                            config.context = ContextConfig(**data["context"])
+            except yaml.YAMLError:
+                # Invalid YAML, use empty config
+                pass
+    return config
+
+
+def save_config(config: Config):
+    data = {
+        "workspaces": [
+            {
+                "name": ws.name,
+                "credentials": {
+                    "access_token": ws.credentials.access_token,
+                    "api_key": ws.credentials.api_key
+                }
+            }
+            for ws in config.workspaces
+        ],
+        "context": {
+            "workspace": config.context.workspace,
+            "environment": config.context.environment
+        }
+    }
+
+    home_dir = Path.home()
+    if not home_dir:
+        raise RuntimeError("Could not determine home directory")
+
+    config_dir = home_dir / ".beamlit"
+    config_file = config_dir / "config.yaml"
+
+    config_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
+
+    with open(config_file, "w", encoding="utf-8") as f:
+        yaml.dump(data, f)
+
+
+def list_workspaces() -> List[str]:
+    config = load_config()
+    return [workspace.name for workspace in config.workspaces]
+
+
+def current_context() -> ContextConfig:
+    config = load_config()
+    return config.context
+
+
+def set_current_workspace(workspace_name: str, environment: str):
+    config = load_config()
+    config.context.workspace = workspace_name
+    config.context.environment = environment
+    save_config(config)
+
+
+def load_credentials(workspace_name: str) -> Credentials:
+    config = load_config()
+    for workspace in config.workspaces:
+        if workspace.name == workspace_name:
+            return workspace.credentials
+    return Credentials()
+
+
+def create_home_dir_if_missing():
+    home_dir = Path.home()
+    if not home_dir:
+        print("Error getting home directory")
+        return
+
+    credentials_dir = home_dir / ".beamlit"
+    credentials_file = credentials_dir / "credentials.json"
+
+    if credentials_file.exists():
+        print("You are already logged in. Enter a new API key to overwrite it.")
+    else:
+        try:
+            credentials_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
+        except Exception as e:
+            print(f"Error creating credentials directory: {e}")
+
+
+def save_credentials(workspace_name: str, credentials: Credentials):
+    create_home_dir_if_missing()
+    if not credentials.access_token and not credentials.api_key:
+        print("No credentials to save, error")
+        return
+
+    config = load_config()
+    found = False
+
+    for i, workspace in enumerate(config.workspaces):
+        if workspace.name == workspace_name:
+            config.workspaces[i].credentials = credentials
+            found = True
+            break
+
+    if not found:
+        config.workspaces.append(WorkspaceConfig(name=workspace_name, credentials=credentials))
+
+    save_config(config)
+
+
+def clear_credentials(workspace_name: str):
+    config = load_config()
+    config.workspaces = [ws for ws in config.workspaces if ws.name != workspace_name]
+
+    if config.context.workspace == workspace_name:
+        config.context.workspace = ""
+
+    save_config(config)

beamlit/authentication/device_mode.py

@@ -0,0 +1,115 @@
+import base64
+import json
+import time
+from dataclasses import dataclass
+from typing import Generator, Optional
+
+from httpx import Auth, Request, Response, post
+
+
+@dataclass
+class DeviceLogin:
+    client_id: str
+    scope: str
+
+
+@dataclass
+class DeviceLoginResponse:
+    client_id: str
+    device_code: str
+    user_code: str
+    expires_in: int
+    interval: int
+    verification_uri: str
+    verification_uri_complete: str
+
+
+@dataclass
+class DeviceLoginFinalizeRequest:
+    grant_type: str
+    client_id: str
+    device_code: str
+
+
+@dataclass
+class DeviceLoginFinalizeResponse:
+    access_token: str
+    expires_in: int
+    refresh_token: str
+    token_type: str
+
+
+class BearerToken(Auth):
+    def __init__(self, credentials, workspace_name: str, base_url: str):
+        self.credentials = credentials
+        self.workspace_name = workspace_name
+        self.base_url = base_url
+
+    def refresh_if_needed(self) -> Optional[Exception]:
+        # Need to refresh token if expires in less than 10 minutes
+        parts = self.credentials.access_token.split('.')
+        if len(parts) != 3:
+            return Exception("Invalid JWT token format")
+
+        try:
+            claims_bytes = base64.urlsafe_b64decode(parts[1] + '=' * (-len(parts[1]) % 4))
+            claims = json.loads(claims_bytes)
+        except Exception as e:
+            return Exception(f"Failed to decode/parse JWT claims: {str(e)}")
+
+        exp_time = time.gmtime(claims['exp'])
+        # Refresh if token expires in less than 10 minutes
+        if time.time() + (10 * 60) > time.mktime(exp_time):
+            return self.do_refresh()
+
+        return None
+
+    def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
+        err = self.refresh_if_needed()
+        if err:
+            return err
+
+        request.headers['X-Beamlit-Authorization'] = f'Bearer {self.credentials.access_token}'
+        request.headers['X-Beamlit-Workspace'] = self.workspace_name
+        yield request
+
+    def do_refresh(self) -> Optional[Exception]:
+        if not self.credentials.refresh_token:
+            return Exception("No refresh token to refresh")
+
+        url = f"{self.base_url}/oauth/token"
+        refresh_data = {
+            "grant_type": "refresh_token",
+            "refresh_token": self.credentials.refresh_token,
+            "device_code": self.credentials.device_code,
+            "client_id": "beamlit"
+        }
+
+        try:
+            print(refresh_data)
+            response = post(
+                url,
+                json=refresh_data,
+                headers={"Content-Type": "application/json"}
+            )
+            print(response.text)
+            response.raise_for_status()
+            finalize_response = DeviceLoginFinalizeResponse(**response.json())
+
+            if not finalize_response.refresh_token:
+                finalize_response.refresh_token = self.credentials.refresh_token
+
+            from .credentials import Credentials, save_credentials
+            creds = Credentials(
+                access_token=finalize_response.access_token,
+                refresh_token=finalize_response.refresh_token,
+                expires_in=finalize_response.expires_in,
+                device_code=self.credentials.device_code
+            )
+
+            self.credentials = creds
+            save_credentials(self.workspace_name, creds)
+            return None
+
+        except Exception as e:
+            return Exception(f"Failed to refresh token: {str(e)}")