beamlit 0.0.51__py3-none-any.whl → 0.0.53__py3-none-any.whl

beamlit/agents/chain.py

@@ -16,7 +16,7 @@ from beamlit.run import RunClient
 
 class ChainTool(BaseTool):
     """
-    Chain tool
+    A tool that allows chaining of agent actions. Extends LangChain's BaseTool.
     """
 
     client: RunClient
@@ -24,6 +24,16 @@ class ChainTool(BaseTool):
 
     @t.override
     def _run(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        """
+        Executes the tool synchronously.
+
+        Parameters:
+            *args (Any): Positional arguments.
+            **kwargs (Any): Keyword arguments.
+
+        Returns:
+            Any: The result of the tool execution.
+        """
         warnings.warn(
             "Invoke this tool asynchronousely using `ainvoke`. This method exists only to satisfy standard tests.",
             stacklevel=1,
@@ -32,6 +42,16 @@ class ChainTool(BaseTool):
 
     @t.override
     async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        """
+        Executes the tool asynchronously.
+
+        Parameters:
+            *args (Any): Positional arguments.
+            **kwargs (Any): Keyword arguments.
+
+        Returns:
+            Any: The result of the asynchronous tool execution.
+        """
         settings = get_settings()
         result = self.client.run(
             "agent",
@@ -45,17 +65,30 @@ class ChainTool(BaseTool):
     @t.override
     @property
     def tool_call_schema(self) -> type[pydantic.BaseModel]:
+        """
+        Defines the schema for tool calls based on the provided argument schema.
+
+        Returns:
+            type[pydantic.BaseModel]: The Pydantic model representing the tool call schema.
+        """
         assert self.args_schema is not None  # noqa: S101
         return self.args_schema
 
+
 class ChainInput(pydantic.BaseModel):
+    """
+    A Pydantic model representing the input structure for a chain.
+    """
+
     inputs: str
 
+
 @dataclass
 class ChainToolkit:
     """
-    Remote toolkit
+    A toolkit for managing and initializing a chain of agents.
     """
+
     client: AuthenticatedClient
     chain: list[AgentChain]
     _chain: list[Agent] | None = None
@@ -63,6 +96,12 @@ class ChainToolkit:
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
     def initialize(self) -> None:
+        """
+        Initializes the toolkit by retrieving and configuring the list of agents based on the provided chains.
+
+        Raises:
+            RuntimeError: If initialization fails due to missing agents.
+        """
         """Initialize the session and retrieve tools list"""
         if self._chain is None:
             agents = list_agents.sync_detailed(
@@ -77,8 +116,16 @@ class ChainToolkit:
                     agents_chain.append(agent[0])
             self._chain = agents_chain
 
-    @t.override
     def get_tools(self) -> list[BaseTool]:
+        """
+        Retrieves a list of tools corresponding to the initialized agents.
+
+        Returns:
+            list[BaseTool]: A list of initialized `ChainTool` instances.
+
+        Raises:
+            RuntimeError: If the toolkit has not been initialized.
+        """
         if self._chain is None:
             raise RuntimeError("Must initialize the toolkit first")
 

beamlit/agents/chat.py

@@ -7,39 +7,96 @@ from beamlit.api.models import get_model
 from beamlit.authentication import get_authentication_headers, new_client
 from beamlit.common.settings import get_settings
 from beamlit.models import Model
+
 from .voice.openai import OpenAIVoiceReactAgent
 
 logger = getLogger(__name__)
 
 
-def get_base_url(name: str):
+def get_base_url(name: str) -> str:
+    """
+    Constructs the base URL for a given model name based on the current settings.
+
+    Parameters:
+        name (str): The name of the model.
+
+    Returns:
+        str: The constructed base URL.
+    """
     settings = get_settings()
     return f"{settings.run_url}/{settings.workspace}/models/{name}/v1"
 
 
-def get_mistral_chat_model(**kwargs):
+def get_mistral_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns a MistralAI chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatMistralAI` model.
+
+    Returns:
+        ChatMistralAI: An instance of the MistralAI chat model.
+    """
     from langchain_mistralai.chat_models import ChatMistralAI  # type: ignore
 
     return ChatMistralAI(**kwargs)
 
 
-def get_openai_chat_model(**kwargs):
+def get_openai_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns an OpenAI chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatOpenAI` model.
+
+    Returns:
+        ChatOpenAI: An instance of the OpenAI chat model.
+    """
     from langchain_openai import ChatOpenAI  # type: ignore
 
     return ChatOpenAI(**kwargs)
 
 
-def get_anthropic_chat_model(**kwargs):
+def get_anthropic_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns an Anthropic chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatAnthropic` model.
+
+    Returns:
+        ChatAnthropic: An instance of the Anthropic chat model.
+    """
     from langchain_anthropic import ChatAnthropic  # type: ignore
 
     return ChatAnthropic(**kwargs)
 
-def get_xai_chat_model(**kwargs):
+
+def get_xai_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns an XAI chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatXAI` model.
+
+    Returns:
+        ChatXAI: An instance of the XAI chat model.
+    """
     from langchain_xai import ChatXAI  # type: ignore
 
     return ChatXAI(**kwargs)
 
-def get_cohere_chat_model(**kwargs):
+
+def get_cohere_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns a Cohere chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatCohere` model.
+
+    Returns:
+        ChatCohere: An instance of the Cohere chat model.
+    """
     from langchain_cohere import ChatCohere  # type: ignore
 
     return ChatCohere(**kwargs)
@@ -64,10 +121,35 @@ def get_azure_marketplace_chat_model(**kwargs):
     )  # It seems to use a compatible endpoint, so we can use the classic OpenAI interface
 
 def get_chat_model(name: str, agent_model: Union[Model, None] = None) -> BaseChatModel:
+    """
+    Gets a chat model instance for the specified model name.
+
+    Parameters:
+        name (str): The name of the model to retrieve.
+        agent_model (Union[Model, None], optional): A pre-fetched model instance.
+            If None, the model will be fetched from the API. Defaults to None.
+
+    Returns:
+        BaseChatModel: An instance of the appropriate chat model.
+    """
     [chat_model, _, __] = get_chat_model_full(name, agent_model)
     return chat_model
 
 def get_chat_model_full(name: str, agent_model: Union[Model, None] = None) -> Tuple[BaseChatModel, str, str]:
+    """
+    Gets a chat model instance along with provider and model information.
+
+    Parameters:
+        name (str): The name of the model to retrieve.
+        agent_model (Union[Model, None], optional): A pre-fetched model instance.
+            If None, the model will be fetched from the API. Defaults to None.
+
+    Returns:
+        Tuple[BaseChatModel, str, str]: A tuple containing:
+            - The chat model instance
+            - The provider name (e.g., 'openai', 'anthropic', etc.)
+            - The specific model name (e.g., 'gpt-4o-mini')
+    """
     settings = get_settings()
     client = new_client()
 
@@ -77,7 +159,10 @@ def get_chat_model_full(name: str, agent_model: Union[Model, None] = None) -> Tu
         except Exception:
             logger.warning(f"Model {name} not found, defaulting to gpt-4o-mini")
 
-    environment = (agent_model and agent_model.metadata and agent_model.metadata.environment) or settings.environment
+    environment = (
+        (agent_model and agent_model.metadata and agent_model.metadata.environment)
+        or settings.environment
+    )
     headers = get_authentication_headers(settings)
     headers["X-Beamlit-Environment"] = environment
 

beamlit/agents/decorator.py

@@ -1,7 +1,13 @@
+"""Module: decorator
+
+Defines decorators for agent functionalities.
+"""
+
 # Import necessary modules
 import functools
 import inspect
 from logging import getLogger
+from typing import Callable
 
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.prebuilt import create_react_agent
@@ -12,8 +18,9 @@ from beamlit.common.settings import init
 from beamlit.errors import UnexpectedStatus
 from beamlit.functions import get_functions
 from beamlit.models import Agent, AgentSpec, EnvironmentMetadata
-from .voice.openai import OpenAIVoiceReactAgent
+
 from .chat import get_chat_model_full
+from .voice.openai import OpenAIVoiceReactAgent
 
 
 def agent(
@@ -22,7 +29,32 @@ def agent(
     override_agent=None,
     override_functions=None,
     remote_functions=None,
-):
+) -> Callable:
+    """
+    A decorator factory that configures and wraps functions to integrate with Beamlit agents.
+    Handles model initialization, function retrieval, and agent setup.
+
+    Parameters:
+        agent (Agent | dict, optional): An `Agent` instance or a dictionary containing agent metadata and specifications.
+        override_model (Any, optional): An optional model to override the default agent model.
+        override_agent (Any, optional): An optional agent instance to override the default agent.
+        mcp_hub (Any, optional): An optional MCP hub configuration.
+        remote_functions (Any, optional): An optional list of remote functions to be integrated.
+
+    Returns:
+        Callable: A decorator that wraps the target function, injecting agent-related configurations and dependencies.
+
+    Behavior:
+        - Validates and initializes the agent configuration.
+        - Retrieves and sets up the appropriate chat model based on the agent's specifications.
+        - Retrieves functions from the specified directories or remote sources.
+        - Wraps the target function, injecting agent, model, and function dependencies as needed.
+        - Logs relevant information and handles exceptions during the setup process.
+
+    Raises:
+        ValueError: If required configurations such as the model are missing.
+        Re-raises exceptions encountered during model retrieval and agent setup.
+    """
     logger = getLogger(__name__)
     try:
         if agent is not None and not isinstance(agent, dict):
@@ -47,6 +79,7 @@ def agent(
                 param.name == "functions"
                 for param in inspect.signature(func).parameters.values()
             )
+
             @functools.wraps(func)
             def wrapped(*args, **kwargs):
                 if agent_kwargs:
@@ -111,21 +144,22 @@ def agent(
                         environment=settings.environment, client=client
                     )
                     models = ", ".join([model.metadata.name for model in models.parsed])
-                    models_select = f"You can select one from the your models: {models}"
+                    models_select = f"You can select one from your models: {models}"
                 except Exception:
                     pass
 
-                raise ValueError(f"You must provide a model.\n"
+                raise ValueError(
+                    f"You must provide a model.\n"
                     f"{models_select}\n"
                     f"You can create one at {settings.app_url}/{settings.workspace}/global-inference-network/models/create\n"
                     "Add it to your agent spec\n"
                     "agent={\n"
-                    "    \"metadata\": {\n"
-                    f"        \"name\": \"{agent.metadata.name}\",\n"
+                    '    "metadata": {\n'
+                    f'        "name": "{agent.metadata.name}",\n'
                     "    },\n"
-                    "    \"spec\": {\n"
-                    "        \"model\": \"MODEL_NAME\",\n"
-                    f"        \"description\": \"{agent.spec.description}\",\n"
+                    '    "spec": {\n'
+                    '        "model": "MODEL_NAME",\n'
+                    f'        "description": "{agent.spec.description}",\n'
                     "    },\n"
                     "}")
             if isinstance(chat_model, OpenAIVoiceReactAgent):

beamlit/agents/thread.py

@@ -1,9 +1,23 @@
+"""Module: thread
 
+Defines threading capabilities for agents.
+"""
 import jwt
 from fastapi import Request
 
 
 def get_default_thread(request: Request) -> str:
+    """
+    Extracts the default thread identifier from an incoming HTTP request.
+    Prioritizes the `X-Beamlit-Sub` header and falls back to decoding the JWT
+    from the `Authorization` or `X-Beamlit-Authorization` headers.
+
+    Parameters:
+        request (Request): The incoming HTTP request object.
+
+    Returns:
+        str: The extracted thread identifier. Returns an empty string if no valid identifier is found.
+    """
     if request.headers.get("X-Beamlit-Sub"):
         return request.headers.get("X-Beamlit-Sub")
     authorization = request.headers.get("Authorization", request.headers.get("X-Beamlit-Authorization"))

beamlit/agents/voice/openai.py

@@ -1,16 +1,14 @@
 import asyncio
 import json
-import websockets
-
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, AsyncIterator, Any, Callable, Coroutine
-from .utils import amerge
+from typing import Any, AsyncGenerator, AsyncIterator, Callable, Coroutine
 
-from langchain_core.tools import BaseTool
+import websockets
 from langchain_core._api import beta
-from langchain_core.utils import secret_from_env
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field, PrivateAttr
 
-from pydantic import BaseModel, Field, SecretStr, PrivateAttr
+from .utils import amerge
 
 EVENTS_TO_IGNORE = {
     "response.function_call_arguments.delta",

beamlit/authentication/apikey.py

@@ -1,20 +1,50 @@
+"""
+This module provides the ApiKeyProvider class, which handles API key-based authentication for Beamlit.
+"""
+
 from typing import Generator
 
 from httpx import Auth, Request, Response
 
 
 class ApiKeyProvider(Auth):
+    """
+    A provider that authenticates requests using an API key.
+    """
+
     def __init__(self, credentials, workspace_name: str):
+        """
+        Initializes the ApiKeyProvider with the given credentials and workspace name.
+
+        Parameters:
+            credentials: Credentials containing the API key.
+            workspace_name (str): The name of the workspace.
+        """
         self.credentials = credentials
         self.workspace_name = workspace_name
 
     def get_headers(self):
+        """
+        Retrieves the authentication headers containing the API key and workspace information.
+
+        Returns:
+            dict: A dictionary of headers with API key and workspace.
+        """
         return {
             "X-Beamlit-Api-Key": self.credentials.apiKey,
             "X-Beamlit-Workspace": self.workspace_name,
         }
 
     def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
+        """
+        Authenticates the request by adding API key and workspace headers.
+
+        Parameters:
+            request (Request): The HTTP request to authenticate.
+
+        Yields:
+            Request: The authenticated request.
+        """
         request.headers["X-Beamlit-Api-Key"] = self.credentials.apiKey
         request.headers["X-Beamlit-Workspace"] = self.workspace_name
         yield request

beamlit/authentication/authentication.py

@@ -1,3 +1,9 @@
+"""
+This module provides classes and functions for handling various authentication methods,
+including public access, API key authentication, client credentials, and bearer tokens.
+It also includes utilities for creating authenticated clients and managing authentication headers.
+"""
+
 from dataclasses import dataclass
 from typing import Dict, Generator
 
@@ -18,18 +24,43 @@ from .device_mode import BearerToken
 
 
 class PublicProvider(Auth):
+    """
+    A provider that allows public access without any authentication.
+    """
+
     def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
+        """
+        Processes the authentication flow for public access by yielding the request as-is.
+
+        Parameters:
+            request (Request): The HTTP request to process.
+
+        Yields:
+            Request: The unmodified request.
+        """
         yield request
 
 
 @dataclass
 class RunClientWithCredentials:
+    """
+    A dataclass that holds credentials and workspace information for initializing an authenticated client.
+
+    Attributes:
+        credentials (Credentials): The credentials used for authentication.
+        workspace (str): The name of the workspace.
+        api_url (str): The base API URL.
+        run_url (str): The run-specific URL.
+    """
     credentials: Credentials
     workspace: str
     api_url: str = ""
     run_url: str = ""
 
     def __post_init__(self):
+        """
+        Post-initialization to set the API and run URLs from settings.
+        """
         from ..common.settings import get_settings
 
         settings = get_settings()
@@ -38,6 +69,15 @@ class RunClientWithCredentials:
 
 
 def new_client_from_settings(settings: Settings):
+    """
+    Creates a new authenticated client using the provided settings.
+
+    Parameters:
+        settings (Settings): The settings containing authentication and workspace information.
+
+    Returns:
+        AuthenticatedClient: An instance of AuthenticatedClient configured with the provided settings.
+    """
     credentials = load_credentials_from_settings(settings)
 
     client_config = RunClientWithCredentials(
@@ -48,6 +88,12 @@ def new_client_from_settings(settings: Settings):
 
 
 def new_client():
+    """
+    Creates a new authenticated client based on the current context and settings.
+
+    Returns:
+        AuthenticatedClient: An instance of AuthenticatedClient configured with the current context or settings.
+    """
     settings = get_settings()
     context = current_context()
     if context.workspace and not settings.authentication.client.credentials:
@@ -66,6 +112,15 @@ def new_client():
 
 
 def new_client_with_credentials(config: RunClientWithCredentials):
+    """
+    Creates a new authenticated client using the provided client configuration.
+
+    Parameters:
+        config (RunClientWithCredentials): The client configuration containing credentials and workspace information.
+
+    Returns:
+        AuthenticatedClient: An instance of AuthenticatedClient configured with the provided credentials.
+    """
     provider: Auth = None
     if config.credentials.apiKey:
         provider = ApiKeyProvider(config.credentials, config.workspace)
@@ -80,6 +135,15 @@ def new_client_with_credentials(config: RunClientWithCredentials):
 
 
 def get_authentication_headers(settings: Settings) -> Dict[str, str]:
+    """
+    Retrieves authentication headers based on the current context and settings.
+
+    Parameters:
+        settings (Settings): The settings containing authentication and workspace information.
+
+    Returns:
+        Dict[str, str]: A dictionary of authentication headers.
+    """
     context = current_context()
     if context.workspace and not settings.authentication.client.credentials:
         credentials = load_credentials(context.workspace)

beamlit/authentication/clientcredentials.py

@@ -1,3 +1,9 @@
+"""
+This module provides the ClientCredentials class, which handles client credentials-based
+authentication for Beamlit. It manages token refreshing and authentication flows using
+client credentials and refresh tokens.
+"""
+
 import base64
 import json
 from dataclasses import dataclass
@@ -19,12 +25,33 @@ class DeviceLoginFinalizeResponse:
 
 
 class ClientCredentials(Auth):
+    """
+    A provider that authenticates requests using client credentials.
+    """
+
     def __init__(self, credentials, workspace_name: str, base_url: str):
+        """
+        Initializes the ClientCredentials provider with the given credentials, workspace name, and base URL.
+
+        Parameters:
+            credentials: Credentials containing access and refresh tokens.
+            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):
+        """
+        Retrieves the authentication headers after ensuring tokens are valid.
+
+        Returns:
+            dict: A dictionary of headers with Bearer token and workspace.
+
+        Raises:
+            Exception: If token refresh fails.
+        """
         err = self.refresh_if_needed()
         if err:
             raise err
@@ -35,6 +62,12 @@ class ClientCredentials(Auth):
         }
 
     def refresh_if_needed(self) -> Optional[Exception]:
+        """
+        Checks if the access token needs to be refreshed and performs the refresh if necessary.
+
+        Returns:
+            Optional[Exception]: An exception if refreshing fails, otherwise None.
+        """
         settings = get_settings()
         if self.credentials.client_credentials and not self.credentials.refresh_token:
             headers = {"Authorization": f"Basic {self.credentials.client_credentials}", "Content-Type": "application/json"}
@@ -61,10 +94,21 @@ class ClientCredentials(Auth):
         if current_time + timedelta(minutes=10) > exp_time:
             return self.do_refresh()
 
-
         return None
 
     def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
+        """
+        Processes the authentication flow by ensuring tokens are 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
@@ -74,6 +118,12 @@ class ClientCredentials(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")