beamlit 0.0.52__py3-none-any.whl → 0.0.53rc99__py3-none-any.whl

beamlit/authentication/credentials.py

@@ -1,3 +1,9 @@
+"""
+This module provides classes and functions for managing credentials and workspace configurations.
+It includes functionalities to load, save, and manage authentication credentials, as well as to handle
+workspace contexts and configurations.
+"""
+
 from dataclasses import dataclass
 from logging import getLogger
 from pathlib import Path
@@ -12,6 +18,17 @@ logger = getLogger(__name__)
 
 @dataclass
 class Credentials:
+    """
+    A dataclass representing user credentials for authentication.
+
+    Attributes:
+        apiKey (str): The API key.
+        access_token (str): The access token.
+        refresh_token (str): The refresh token.
+        expires_in (int): Token expiration time in seconds.
+        device_code (str): The device code for device authentication.
+        client_credentials (str): The client credentials for authentication.
+    """
     apiKey: str = ""
     access_token: str = ""
     refresh_token: str = ""
@@ -22,28 +39,58 @@ class Credentials:
 
 @dataclass
 class WorkspaceConfig:
+    """
+    A dataclass representing the configuration for a workspace.
+
+    Attributes:
+        name (str): The name of the workspace.
+        credentials (Credentials): The credentials associated with the workspace.
+    """
     name: str
     credentials: Credentials
 
 
 @dataclass
 class ContextConfig:
+    """
+    A dataclass representing the current context configuration.
+
+    Attributes:
+        workspace (str): The name of the current workspace.
+        environment (str): The current environment (e.g., development, production).
+    """
     workspace: str = ""
     environment: str = ""
 
 
 @dataclass
 class Config:
+    """
+    A dataclass representing the overall configuration, including workspaces and context.
+
+    Attributes:
+        workspaces (List[WorkspaceConfig]): A list of workspace configurations.
+        context (ContextConfig): The current context configuration.
+    """
     workspaces: List[WorkspaceConfig] = None
     context: ContextConfig = None
 
     def __post_init__(self):
+        """
+        Post-initialization to ensure workspaces and context are initialized.
+        """
         if self.workspaces is None:
             self.workspaces = []
         if self.context is None:
             self.context = ContextConfig()
 
     def to_json(self) -> dict:
+        """
+        Converts the Config dataclass to a JSON-compatible dictionary.
+
+        Returns:
+            dict: The JSON representation of the configuration.
+        """
         return {
             "workspaces": [
                 {
@@ -67,6 +114,12 @@ class Config:
 
 
 def load_config() -> Config:
+    """
+    Loads the configuration from the user's home directory.
+
+    Returns:
+        Config: The loaded configuration.
+    """
     config = Config()
     home_dir = Path.home()
     if home_dir:
@@ -90,6 +143,15 @@ def load_config() -> Config:
 
 
 def save_config(config: Config):
+    """
+    Saves the provided configuration to the user's home directory.
+
+    Parameters:
+        config (Config): The configuration to save.
+
+    Raises:
+        RuntimeError: If the home directory cannot be determined.
+    """
     home_dir = Path.home()
     if not home_dir:
         raise RuntimeError("Could not determine home directory")
@@ -103,16 +165,35 @@ def save_config(config: Config):
 
 
 def list_workspaces() -> List[str]:
+    """
+    Lists all available workspace names from the configuration.
+
+    Returns:
+        List[str]: A list of workspace names.
+    """
     config = load_config()
     return [workspace.name for workspace in config.workspaces]
 
 
 def current_context() -> ContextConfig:
+    """
+    Retrieves the current context configuration.
+
+    Returns:
+        ContextConfig: The current context configuration.
+    """
     config = load_config()
     return config.context
 
 
 def set_current_workspace(workspace_name: str, environment: str):
+    """
+    Sets the current workspace and environment in the configuration.
+
+    Parameters:
+        workspace_name (str): The name of the workspace to set as current.
+        environment (str): The environment to set for the workspace.
+    """
     config = load_config()
     config.context.workspace = workspace_name
     config.context.environment = environment
@@ -120,6 +201,15 @@ def set_current_workspace(workspace_name: str, environment: str):
 
 
 def load_credentials(workspace_name: str) -> Credentials:
+    """
+    Loads credentials for the specified workspace.
+
+    Parameters:
+        workspace_name (str): The name of the workspace whose credentials are to be loaded.
+
+    Returns:
+        Credentials: The credentials associated with the workspace. Returns empty credentials if not found.
+    """
     config = load_config()
     for workspace in config.workspaces:
         if workspace.name == workspace_name:
@@ -128,6 +218,15 @@ def load_credentials(workspace_name: str) -> Credentials:
 
 
 def load_credentials_from_settings(settings: Settings) -> Credentials:
+    """
+    Loads credentials from the provided settings.
+
+    Parameters:
+        settings (Settings): The settings containing authentication information.
+
+    Returns:
+        Credentials: The loaded credentials from settings.
+    """
     return Credentials(
         apiKey=settings.authentication.apiKey,
         client_credentials=settings.authentication.client.credentials,
@@ -135,6 +234,11 @@ def load_credentials_from_settings(settings: Settings) -> Credentials:
 
 
 def create_home_dir_if_missing():
+    """
+    Creates the Beamlit home directory if it does not exist.
+
+    Logs a warning if credentials already exist or an error if directory creation fails.
+    """
     home_dir = Path.home()
     if not home_dir:
         logger.error("Error getting home directory")
@@ -153,6 +257,13 @@ def create_home_dir_if_missing():
 
 
 def save_credentials(workspace_name: str, credentials: Credentials):
+    """
+    Saves the provided credentials for the specified workspace.
+
+    Parameters:
+        workspace_name (str): The name of the workspace.
+        credentials (Credentials): The credentials to save.
+    """
     create_home_dir_if_missing()
     if not credentials.access_token and not credentials.apiKey:
         logger.info("No credentials to save, error")
@@ -174,6 +285,12 @@ def save_credentials(workspace_name: str, credentials: Credentials):
 
 
 def clear_credentials(workspace_name: str):
+    """
+    Clears the credentials for the specified workspace.
+
+    Parameters:
+        workspace_name (str): The name of the workspace whose credentials are to be cleared.
+    """
     config = load_config()
     config.workspaces = [ws for ws in config.workspaces if ws.name != workspace_name]
 

beamlit/authentication/device_mode.py

@@ -1,3 +1,9 @@
+"""
+This module provides classes for handling device-based authentication,
+including device login processes and bearer token management. It facilitates token refreshing
+and authentication flows using device codes and bearer tokens.
+"""
+
 import base64
 import json
 from dataclasses import dataclass
@@ -9,12 +15,31 @@ from httpx import Auth, Request, Response, post
 
 @dataclass
 class DeviceLogin:
+    """
+    A dataclass representing a device login request.
+
+    Attributes:
+        client_id (str): The client ID for the device.
+        scope (str): The scope of the authentication.
+    """
     client_id: str
     scope: str
 
 
 @dataclass
 class DeviceLoginResponse:
+    """
+    A dataclass representing the response from a device login request.
+
+    Attributes:
+        client_id (str): The client ID associated with the device login.
+        device_code (str): The device code for authentication.
+        user_code (str): The user code for completing authentication.
+        expires_in (int): Time in seconds until the device code expires.
+        interval (int): Polling interval in seconds.
+        verification_uri (str): URI for user to verify device login.
+        verification_uri_complete (str): Complete URI including the user code for verification.
+    """
     client_id: str
     device_code: str
     user_code: str
@@ -26,6 +51,14 @@ class DeviceLoginResponse:
 
 @dataclass
 class DeviceLoginFinalizeRequest:
+    """
+    A dataclass representing a device login finalize request.
+
+    Attributes:
+        grant_type (str): The type of grant being requested.
+        client_id (str): The client ID for finalizing the device login.
+        device_code (str): The device code to finalize login.
+    """
     grant_type: str
     client_id: str
     device_code: str
@@ -40,12 +73,33 @@ class DeviceLoginFinalizeResponse:
 
 
 class BearerToken(Auth):
+    """
+    A provider that authenticates requests using a Bearer token.
+    """
+
     def __init__(self, credentials, workspace_name: str, base_url: str):
+        """
+        Initializes the BearerToken provider with the given credentials, workspace name, and base URL.
+
+        Parameters:
+            credentials: Credentials containing the Bearer token and refresh token.
+            workspace_name (str): The name of the workspace.
+            base_url (str): The base URL for authentication.
+        """
         self.credentials = credentials
         self.workspace_name = workspace_name
         self.base_url = base_url
 
     def get_headers(self) -> Dict[str, str]:
+        """
+        Retrieves the authentication headers containing the Bearer token and workspace information.
+
+        Returns:
+            Dict[str, str]: A dictionary of headers with Bearer token and workspace.
+
+        Raises:
+            Exception: If token refresh fails.
+        """
         err = self.refresh_if_needed()
         if err:
             raise err
@@ -55,6 +109,12 @@ class BearerToken(Auth):
         }
 
     def refresh_if_needed(self) -> Optional[Exception]:
+        """
+        Checks if the Bearer token needs to be refreshed and performs the refresh if necessary.
+
+        Returns:
+            Optional[Exception]: An exception if refreshing fails, otherwise None.
+        """
         # Need to refresh token if expires in less than 10 minutes
         parts = self.credentials.access_token.split(".")
         if len(parts) != 3:
@@ -74,6 +134,18 @@ class BearerToken(Auth):
         return None
 
     def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
+        """
+        Processes the authentication flow by ensuring the Bearer token is valid and adding necessary headers.
+
+        Parameters:
+            request (Request): The HTTP request to authenticate.
+
+        Yields:
+            Request: The authenticated request.
+
+        Raises:
+            Exception: If token refresh fails.
+        """
         err = self.refresh_if_needed()
         if err:
             return err
@@ -83,6 +155,12 @@ class BearerToken(Auth):
         yield request
 
     def do_refresh(self) -> Optional[Exception]:
+        """
+        Performs the token refresh using the refresh token.
+
+        Returns:
+            Optional[Exception]: An exception if refreshing fails, otherwise None.
+        """
         if not self.credentials.refresh_token:
             return Exception("No refresh token to refresh")
 

beamlit/common/error.py

@@ -1,9 +1,27 @@
+"""
+This module defines custom exception classes used for handling HTTP-related errors within Beamlit.
+"""
+
 class HTTPError(Exception):
+    """
+    A custom exception class for HTTP errors.
+
+    Attributes:
+        status_code (int): The HTTP status code associated with the error.
+        message (str): A descriptive message explaining the error.
+    """
+
     def __init__(self, status_code: int, message: str):
         self.status_code = status_code
         self.message = message
         super().__init__(self.message)
 
     def __str__(self):
+        """
+        Returns a string representation of the HTTPError.
+
+        Returns:
+            str: A string in the format "status_code message".
+        """
         return f"{self.status_code} {self.message}"
 

beamlit/common/instrumentation.py

@@ -1,3 +1,8 @@
+"""
+This module provides utilities for setting up and managing OpenTelemetry instrumentation within Beamlit.
+It includes classes and functions for configuring tracers, meters, loggers, and integrating with FastAPI applications.
+"""
+
 import importlib
 import logging
 from typing import Any, Optional, Type
@@ -5,18 +10,10 @@ from typing import Any, Optional, Type
 from fastapi import FastAPI
 from opentelemetry import _logs, metrics, trace
 from opentelemetry._logs import set_logger_provider
-from opentelemetry.exporter.otlp.proto.http._log_exporter import (
-    OTLPLogExporter,
-)
-from opentelemetry.exporter.otlp.proto.http.metric_exporter import (
-    OTLPMetricExporter,
-)
-from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
-    OTLPSpanExporter,
-)
-from opentelemetry.instrumentation.fastapi import (  # type: ignore
-    FastAPIInstrumentor,
-)
+from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor  # type: ignore
 from opentelemetry.metrics import NoOpMeterProvider
 from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
 from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
@@ -40,6 +37,12 @@ log = logging.getLogger(__name__)
 
 
 def auth_headers() -> Dict[str, str]:
+    """
+    Retrieves authentication headers based on the current settings.
+
+    Returns:
+        Dict[str, str]: A dictionary containing authentication headers.
+    """
     settings = get_settings()
     headers = get_authentication_headers(settings)
     return {
@@ -49,6 +52,15 @@ def auth_headers() -> Dict[str, str]:
 
 
 def get_logger() -> LoggerProvider:
+    """
+    Retrieves the current logger provider.
+
+    Returns:
+        LoggerProvider: The active logger provider.
+
+    Raises:
+        Exception: If the logger has not been initialized.
+    """
     if logger is None:
         raise Exception("Logger is not initialized")
     return logger
@@ -202,6 +214,15 @@ def _is_package_installed(package_name: str) -> bool:
 
 
 def instrument_app(app: FastAPI):
+    """
+    Instruments the given FastAPI application with OpenTelemetry.
+
+    This includes setting up tracer and meter providers, configuring exporters, and instrumenting
+    various modules based on available packages.
+
+    Parameters:
+        app (FastAPI): The FastAPI application to instrument.
+    """
     global tracer
     global meter
     settings = get_settings()
@@ -281,6 +302,11 @@ def instrument_app(app: FastAPI):
 
 
 def shutdown_instrumentation():
+    """
+    Shuts down the OpenTelemetry instrumentation providers gracefully.
+
+    This ensures that all spans and metrics are properly exported before the application exits.
+    """
     if tracer is not None:
         trace_provider = trace.get_tracer_provider()
         if isinstance(trace_provider, TracerProvider):

beamlit/common/logger.py

@@ -1,7 +1,18 @@
+"""
+This module provides a custom colored formatter for logging and an initialization function
+to set up logging configurations for Beamlit applications.
+"""
+
 import logging
 
 
 class ColoredFormatter(logging.Formatter):
+    """
+    A custom logging formatter that adds ANSI color codes to log levels for enhanced readability.
+
+    Attributes:
+        COLORS (dict): A mapping of log level names to their corresponding ANSI color codes.
+    """
     COLORS = {
         "DEBUG": "\033[1;36m",  # Cyan
         "INFO": "\033[1;32m",  # Green
@@ -11,6 +22,15 @@ class ColoredFormatter(logging.Formatter):
     }
 
     def format(self, record):
+        """
+        Formats the log record by adding color codes based on the log level.
+
+        Parameters:
+            record (LogRecord): The log record to format.
+
+        Returns:
+            str: The formatted log message with appropriate color codes.
+        """
         n_spaces = len("CRITICAL") - len(record.levelname)
         tab = " " * n_spaces
         color = self.COLORS.get(record.levelname, "\033[0m")
@@ -19,6 +39,15 @@ class ColoredFormatter(logging.Formatter):
 
 
 def init(log_level: str):
+    """
+    Initializes the logging configuration for Beamlit.
+
+    This function clears existing handlers for specific loggers, sets up a colored formatter,
+    and configures the root logger with the specified log level.
+
+    Parameters:
+        log_level (str): The logging level to set (e.g., "DEBUG", "INFO").
+    """
     logging.getLogger("uvicorn.access").handlers.clear()
     logging.getLogger("uvicorn.access").propagate = False
     logging.getLogger("uvicorn.error").handlers.clear()

beamlit/common/secrets.py

@@ -2,10 +2,38 @@ import os
 
 
 class Secret:
+    """
+    A utility class for managing environment secrets.
+
+    Provides static methods to retrieve and set secret values, supporting both standard and
+    prefixed environment variable naming conventions.
+    """
+
     @staticmethod
     def get(name: str):
+        """
+        Retrieves the value of a secret environment variable.
+
+        This method first attempts to get the value using the standard name. If not found,
+        it tries with a "bl_" prefix.
+
+        Parameters:
+            name (str): The name of the environment variable to retrieve.
+
+        Returns:
+            str: The value of the environment variable if found, otherwise an empty string.
+        """
         return os.getenv(name, os.getenv(f"bl_{name}"))
 
     @staticmethod
     def set(name: str, value: str):
+        """
+        Sets the value of a secret environment variable.
+
+        This method sets the value using the standard name, allowing for consistent secret management.
+
+        Parameters:
+            name (str): The name of the environment variable to set.
+            value (str): The value to assign to the environment variable.
+        """
         os.environ[name] = value

beamlit/common/settings.py

@@ -1,3 +1,9 @@
+"""
+This module defines the configuration management system for Beamlit applications using Pydantic.
+It includes dataclasses for various configuration aspects, such as agents, authentication, and server settings.
+The module provides functions to initialize settings, load configurations from YAML files, and customize settings sources.
+"""
+
 import os
 from typing import Tuple, Type, Union
 
@@ -18,6 +24,18 @@ global SETTINGS
 SETTINGS = None
 
 class SettingsAgent(BaseSettings):
+    """
+    Configuration settings for agents within Beamlit.
+
+    Attributes:
+        agent (Union[None, CompiledGraph]): The compiled agent graph.
+        chain (Union[None, list[Agent]]): A list of agent chains.
+        model (Union[None, Model]): The model configuration.
+        functions (Union[None, list[Function]]): A list of functions available to agents.
+        functions_directory (str): The directory path where agent functions are located.
+        chat_model (Union[None, BaseChatModel]): The chat model used by agents.
+        module (str): The module path to the main application.
+    """
     agent: Union[None, CompiledGraph] = None
     chain: Union[None, list[Agent]] = None
     model: Union[None, Model] = None
@@ -28,6 +46,12 @@ class SettingsAgent(BaseSettings):
 
 
 class SettingsAuthenticationClient(BaseSettings):
+    """
+    Configuration settings for authentication clients.
+
+    Attributes:
+        credentials (Union[None, str]): Client credentials for authentication.
+    """
     credentials: Union[None, str] = None
 
 
@@ -92,10 +116,24 @@ class Settings(BaseSettings):
         )
 
 def get_settings() -> Settings:
+    """
+    Retrieves the current settings instance.
+
+    Returns:
+        Settings: The current settings configuration.
+    """
     return SETTINGS
 
 def init() -> Settings:
-    """Parse the beamlit.yaml file to get configurations."""
+    """
+    Initializes the settings by parsing the `beamlit.yaml` file and setting up logging.
+
+    This function reads workspace and environment configurations from the current context,
+    initializes the global SETTINGS variable, and configures the logger based on the log level.
+
+    Returns:
+        Settings: The initialized settings configuration.
+    """
     from beamlit.authentication.credentials import current_context
 
     global SETTINGS

beamlit/common/slugify.py

@@ -1,2 +1,18 @@
+"""
+This module provides utility functions for string manipulation, including slugification.
+The `slugify` function transforms a given string into a URL-friendly slug by replacing spaces and underscores with hyphens and converting the string to lowercase.
+"""
+
 def slugify(name: str) -> str:
+    """
+    Converts a given string into a URL-friendly slug.
+
+    This function transforms the input string by converting it to lowercase and replacing spaces and underscores with hyphens.
+
+    Parameters:
+        name (str): The string to slugify.
+
+    Returns:
+        str: The slugified version of the input string.
+    """
     return name.lower().replace(" ", "-").replace("_", "-")

beamlit/common/utils.py

@@ -1,9 +1,28 @@
+"""
+This module provides utility functions for file operations within Beamlit.
+It includes functions to copy folders and synchronize directory contents efficiently.
+"""
+
 import filecmp
 import os
 import shutil
 
 
 def copy_folder(source_folder: str, destination_folder: str):
+    """
+    Copies the contents of the source folder to the destination folder.
+
+    This function recursively copies all files and subdirectories from the `source_folder` to the `destination_folder`.
+    It ensures that existing files are only overwritten if they differ from the source.
+
+    Parameters:
+        source_folder (str): The path to the source directory.
+        destination_folder (str): The path to the destination directory.
+
+    Raises:
+        FileNotFoundError: If the source folder does not exist.
+        PermissionError: If the program lacks permissions to read from the source or write to the destination.
+    """
     for file in os.listdir(source_folder):
         if os.path.isdir(f"{source_folder}/{file}"):
             if not os.path.exists(f"{destination_folder}/{file}"):

beamlit/deploy/__init__.py

@@ -1,3 +1,8 @@
+"""
+This module provides functions and classes to handle deployment processes within Beamlit.
+It includes utilities to generate deployment configurations and manage deployment-related operations.
+"""
+
 from .deploy import generate_beamlit_deployment
 
 __all__ = ["generate_beamlit_deployment"]