beamlit 0.0.19__py3-none-any.whl → 0.0.20rc1__py3-none-any.whl

beamlit/api/{authentication_providers/list_models_for_authentication_provider.py → integrations/list_integration_connections.py}

@@ -5,16 +5,14 @@ import httpx
 
 from ... import errors
 from ...client import AuthenticatedClient, Client
-from ...models.authentication_provider_model import AuthenticationProviderModel
+from ...models.integration_connection import IntegrationConnection
 from ...types import Response
 
 
-def _get_kwargs(
-    authentication_provider_name: str,
-) -> dict[str, Any]:
+def _get_kwargs() -> dict[str, Any]:
     _kwargs: dict[str, Any] = {
         "method": "get",
-        "url": f"/authentication_providers/{authentication_provider_name}/models",
+        "url": "/integrations/connections",
     }
 
     return _kwargs
@@ -22,12 +20,12 @@ def _get_kwargs(
 
 def _parse_response(
     *, client: Union[AuthenticatedClient, Client], response: httpx.Response
-) -> Optional[List["AuthenticationProviderModel"]]:
+) -> Optional[List["IntegrationConnection"]]:
     if response.status_code == 200:
         response_200 = []
         _response_200 = response.json()
         for response_200_item_data in _response_200:
-            response_200_item = AuthenticationProviderModel.from_dict(response_200_item_data)
+            response_200_item = IntegrationConnection.from_dict(response_200_item_data)
 
             response_200.append(response_200_item)
 
@@ -40,7 +38,7 @@ def _parse_response(
 
 def _build_response(
     *, client: Union[AuthenticatedClient, Client], response: httpx.Response
-) -> Response[List["AuthenticationProviderModel"]]:
+) -> Response[List["IntegrationConnection"]]:
     return Response(
         status_code=HTTPStatus(response.status_code),
         content=response.content,
@@ -50,28 +48,22 @@ def _build_response(
 
 
 def sync_detailed(
-    authentication_provider_name: str,
     *,
     client: AuthenticatedClient,
-) -> Response[List["AuthenticationProviderModel"]]:
-    """List models for a authentication provider
+) -> Response[List["IntegrationConnection"]]:
+    """List integrations connections
 
-     Returns a list of all models for an integration by ID.
-
-    Args:
-        authentication_provider_name (str):
+     Returns a list of all connections integrations in the workspace.
 
     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[List['AuthenticationProviderModel']]
+        Response[List['IntegrationConnection']]
     """
 
-    kwargs = _get_kwargs(
-        authentication_provider_name=authentication_provider_name,
-    )
+    kwargs = _get_kwargs()
 
     response = client.get_httpx_client().request(
         **kwargs,
@@ -81,54 +73,43 @@ def sync_detailed(
 
 
 def sync(
-    authentication_provider_name: str,
     *,
     client: AuthenticatedClient,
-) -> Optional[List["AuthenticationProviderModel"]]:
-    """List models for a authentication provider
-
-     Returns a list of all models for an integration by ID.
+) -> Optional[List["IntegrationConnection"]]:
+    """List integrations connections
 
-    Args:
-        authentication_provider_name (str):
+     Returns a list of all connections integrations in the workspace.
 
     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:
-        List['AuthenticationProviderModel']
+        List['IntegrationConnection']
     """
 
     return sync_detailed(
-        authentication_provider_name=authentication_provider_name,
         client=client,
     ).parsed
 
 
 async def asyncio_detailed(
-    authentication_provider_name: str,
     *,
     client: AuthenticatedClient,
-) -> Response[List["AuthenticationProviderModel"]]:
-    """List models for a authentication provider
+) -> Response[List["IntegrationConnection"]]:
+    """List integrations connections
 
-     Returns a list of all models for an integration by ID.
-
-    Args:
-        authentication_provider_name (str):
+     Returns a list of all connections integrations in the workspace.
 
     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[List['AuthenticationProviderModel']]
+        Response[List['IntegrationConnection']]
     """
 
-    kwargs = _get_kwargs(
-        authentication_provider_name=authentication_provider_name,
-    )
+    kwargs = _get_kwargs()
 
     response = await client.get_async_httpx_client().request(**kwargs)
 
@@ -136,28 +117,23 @@ async def asyncio_detailed(
 
 
 async def asyncio(
-    authentication_provider_name: str,
     *,
     client: AuthenticatedClient,
-) -> Optional[List["AuthenticationProviderModel"]]:
-    """List models for a authentication provider
-
-     Returns a list of all models for an integration by ID.
+) -> Optional[List["IntegrationConnection"]]:
+    """List integrations connections
 
-    Args:
-        authentication_provider_name (str):
+     Returns a list of all connections integrations in the workspace.
 
     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:
-        List['AuthenticationProviderModel']
+        List['IntegrationConnection']
     """
 
     return (
         await asyncio_detailed(
-            authentication_provider_name=authentication_provider_name,
             client=client,
         )
     ).parsed

beamlit/api/integrations/list_integration_models.py

@@ -0,0 +1,97 @@
+from http import HTTPStatus
+from typing import Any, Optional, Union
+
+import httpx
+
+from ... import errors
+from ...client import AuthenticatedClient, Client
+from ...types import Response
+
+
+def _get_kwargs(
+    integration_name: str,
+) -> dict[str, Any]:
+    _kwargs: dict[str, Any] = {
+        "method": "get",
+        "url": f"/integrations/{integration_name}/models",
+    }
+
+    return _kwargs
+
+
+def _parse_response(*, client: Union[AuthenticatedClient, Client], response: httpx.Response) -> Optional[Any]:
+    if response.status_code == 200:
+        return None
+    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[Any]:
+    return Response(
+        status_code=HTTPStatus(response.status_code),
+        content=response.content,
+        headers=response.headers,
+        parsed=_parse_response(client=client, response=response),
+    )
+
+
+def sync_detailed(
+    integration_name: str,
+    *,
+    client: AuthenticatedClient,
+) -> Response[Any]:
+    """List integration models
+
+     Returns a list of all models for an integration.
+
+    Args:
+        integration_name (str):
+
+    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[Any]
+    """
+
+    kwargs = _get_kwargs(
+        integration_name=integration_name,
+    )
+
+    response = client.get_httpx_client().request(
+        **kwargs,
+    )
+
+    return _build_response(client=client, response=response)
+
+
+async def asyncio_detailed(
+    integration_name: str,
+    *,
+    client: AuthenticatedClient,
+) -> Response[Any]:
+    """List integration models
+
+     Returns a list of all models for an integration.
+
+    Args:
+        integration_name (str):
+
+    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[Any]
+    """
+
+    kwargs = _get_kwargs(
+        integration_name=integration_name,
+    )
+
+    response = await client.get_async_httpx_client().request(**kwargs)
+
+    return _build_response(client=client, response=response)

beamlit/api/integrations/update_integration_connection.py

@@ -0,0 +1,180 @@
+from http import HTTPStatus
+from typing import Any, Optional, Union
+
+import httpx
+
+from ... import errors
+from ...client import AuthenticatedClient, Client
+from ...models.integration_connection import IntegrationConnection
+from ...types import Response
+
+
+def _get_kwargs(
+    connection_name: str,
+    *,
+    body: IntegrationConnection,
+) -> dict[str, Any]:
+    headers: dict[str, Any] = {}
+
+    _kwargs: dict[str, Any] = {
+        "method": "put",
+        "url": f"/integrations/connections/{connection_name}",
+    }
+
+    _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[IntegrationConnection]:
+    if response.status_code == 200:
+        response_200 = IntegrationConnection.from_dict(response.json())
+
+        return response_200
+    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[IntegrationConnection]:
+    return Response(
+        status_code=HTTPStatus(response.status_code),
+        content=response.content,
+        headers=response.headers,
+        parsed=_parse_response(client=client, response=response),
+    )
+
+
+def sync_detailed(
+    connection_name: str,
+    *,
+    client: Union[AuthenticatedClient, Client],
+    body: IntegrationConnection,
+) -> Response[IntegrationConnection]:
+    """Update integration connection
+
+     Update an integration connection by integration name and connection name.
+
+    Args:
+        connection_name (str):
+        body (IntegrationConnection): Integration Connection
+
+    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[IntegrationConnection]
+    """
+
+    kwargs = _get_kwargs(
+        connection_name=connection_name,
+        body=body,
+    )
+
+    response = client.get_httpx_client().request(
+        **kwargs,
+    )
+
+    return _build_response(client=client, response=response)
+
+
+def sync(
+    connection_name: str,
+    *,
+    client: Union[AuthenticatedClient, Client],
+    body: IntegrationConnection,
+) -> Optional[IntegrationConnection]:
+    """Update integration connection
+
+     Update an integration connection by integration name and connection name.
+
+    Args:
+        connection_name (str):
+        body (IntegrationConnection): Integration Connection
+
+    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:
+        IntegrationConnection
+    """
+
+    return sync_detailed(
+        connection_name=connection_name,
+        client=client,
+        body=body,
+    ).parsed
+
+
+async def asyncio_detailed(
+    connection_name: str,
+    *,
+    client: Union[AuthenticatedClient, Client],
+    body: IntegrationConnection,
+) -> Response[IntegrationConnection]:
+    """Update integration connection
+
+     Update an integration connection by integration name and connection name.
+
+    Args:
+        connection_name (str):
+        body (IntegrationConnection): Integration Connection
+
+    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[IntegrationConnection]
+    """
+
+    kwargs = _get_kwargs(
+        connection_name=connection_name,
+        body=body,
+    )
+
+    response = await client.get_async_httpx_client().request(**kwargs)
+
+    return _build_response(client=client, response=response)
+
+
+async def asyncio(
+    connection_name: str,
+    *,
+    client: Union[AuthenticatedClient, Client],
+    body: IntegrationConnection,
+) -> Optional[IntegrationConnection]:
+    """Update integration connection
+
+     Update an integration connection by integration name and connection name.
+
+    Args:
+        connection_name (str):
+        body (IntegrationConnection): Integration Connection
+
+    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:
+        IntegrationConnection
+    """
+
+    return (
+        await asyncio_detailed(
+            connection_name=connection_name,
+            client=client,
+            body=body,
+        )
+    ).parsed

beamlit/authentication/__init__.py

@@ -2,7 +2,7 @@ from .apikey import ApiKeyProvider
 from .authentication import (PublicProvider, RunClientWithCredentials,
                              new_client_with_credentials)
 from .credentials import (Config, ContextConfig, Credentials, WorkspaceConfig,
-                          load_credentials)
+                          load_credentials, load_credentials_from_config)
 from .device_mode import (BearerToken, DeviceLogin, DeviceLoginFinalizeRequest,
                           DeviceLoginFinalizeResponse, DeviceLoginResponse)
 
@@ -16,6 +16,7 @@ __all__ = (
     "Credentials",
     "WorkspaceConfig",
     "load_credentials",
+    "load_credentials_from_config",
     "BearerToken",
     "DeviceLogin",
     "DeviceLoginFinalizeRequest",

beamlit/authentication/authentication.py

@@ -5,6 +5,7 @@ from httpx import Auth, Request, Response
 
 from ..client import AuthenticatedClient
 from .apikey import ApiKeyProvider
+from .clientcredentials import ClientCredentials
 from .credentials import Credentials
 from .device_mode import BearerToken
 
@@ -29,6 +30,8 @@ def new_client_with_credentials(config: RunClientWithCredentials):
         provider = ApiKeyProvider(config.credentials, config.workspace)
     elif config.credentials.access_token:
         provider = BearerToken(config.credentials, config.workspace, config.api_url)
+    elif config.credentials.client_credentials:
+        provider = ClientCredentials(config.credentials, config.workspace, config.api_url)
     else:
         provider = PublicProvider()
 

beamlit/authentication/clientcredentials.py

@@ -0,0 +1,99 @@
+import base64
+import json
+import time
+from dataclasses import dataclass
+from typing import Generator, Optional
+
+import requests
+from beamlit.common.settings import get_settings
+from httpx import Auth, Request, Response, post
+
+
+@dataclass
+class DeviceLoginFinalizeResponse:
+    access_token: str
+    expires_in: int
+    refresh_token: str
+    token_type: str
+
+
+class ClientCredentials(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]:
+        settings = get_settings()
+        if self.credentials.client_credentials and not self.credentials.refresh_token:
+            headers = { "Authorization": f"Basic {self.credentials.client_credentials}" }
+            body = { "grant_type": "client_credentials" }
+            response = requests.post(f"{settings.base_url}/oauth/token", headers=headers, json=body)
+            response.raise_for_status()
+            self.credentials.access_token = response.json()['access_token']
+            self.credentials.refresh_token = response.json()['refresh_token']
+            self.credentials.expires_in = response.json()['expires_in']
+        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:
+            response = post(
+                url,
+                json=refresh_data,
+                headers={"Content-Type": "application/json"}
+            )
+            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)}")

beamlit/authentication/credentials.py

@@ -3,6 +3,7 @@ from pathlib import Path
 from typing import List
 
 import yaml
+from beamlit.common.settings import Settings
 
 
 @dataclass
@@ -12,7 +13,7 @@ class Credentials:
     refresh_token: str = ""
     expires_in: int = 0
     device_code: str = ""
-
+    client_credentials: str = ""
 @dataclass
 class WorkspaceConfig:
     name: str
@@ -115,6 +116,11 @@ def load_credentials(workspace_name: str) -> Credentials:
             return workspace.credentials
     return Credentials()
 
+def load_credentials_from_config(config: Settings) -> Credentials:
+    return Credentials(
+        api_key=config.api_key,
+        client_credentials=config.client_credentials
+    )
 
 def create_home_dir_if_missing():
     home_dir = Path.home()

beamlit/authentication/device_mode.py

@@ -86,13 +86,11 @@ class BearerToken(Auth):
         }
 
         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())