beamlit 0.0.20rc5__py3-none-any.whl → 0.0.20rc6__py3-none-any.whl

beamlit/authentication/authentication.py

@@ -1,13 +1,19 @@
 from dataclasses import dataclass
 from typing import Dict, Generator
 
-from beamlit.common.settings import Settings
 from httpx import Auth, Request, Response
 
+from beamlit.common.settings import Settings, get_settings
+
 from ..client import AuthenticatedClient
 from .apikey import ApiKeyProvider
 from .clientcredentials import ClientCredentials
-from .credentials import Credentials, load_credentials_from_settings
+from .credentials import (
+    Credentials,
+    current_context,
+    load_credentials,
+    load_credentials_from_settings,
+)
 from .device_mode import BearerToken
 
 
@@ -16,7 +22,6 @@ class PublicProvider(Auth):
         yield request
 
 
-
 @dataclass
 class RunClientWithCredentials:
     credentials: Credentials
@@ -34,6 +39,26 @@ def new_client_from_settings(settings: Settings):
     )
     return new_client_with_credentials(client_config)
 
+
+def new_client():
+    context = current_context()
+    if context.workspace:
+        credentials = load_credentials(context.workspace)
+        client_config = RunClientWithCredentials(
+            credentials=credentials,
+            workspace=context.workspace,
+        )
+    else:
+        settings = get_settings()
+        credentials = load_credentials_from_settings(settings)
+
+        client_config = RunClientWithCredentials(
+            credentials=credentials,
+            workspace=settings.workspace,
+        )
+    return new_client_with_credentials(client_config)
+
+
 def new_client_with_credentials(config: RunClientWithCredentials):
     provider: Auth = None
     if config.credentials.api_key:
@@ -41,12 +66,15 @@ def new_client_with_credentials(config: RunClientWithCredentials):
     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)
+        provider = ClientCredentials(
+            config.credentials, config.workspace, config.api_url
+        )
     else:
         provider = PublicProvider()
 
     return AuthenticatedClient(base_url=config.api_url, provider=provider)
 
+
 def get_authentication_headers(settings: Settings) -> Dict[str, str]:
     credentials = load_credentials_from_settings(settings)
     config = RunClientWithCredentials(
@@ -59,9 +87,10 @@ def get_authentication_headers(settings: Settings) -> Dict[str, str]:
     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)
+        provider = ClientCredentials(
+            config.credentials, config.workspace, config.api_url
+        )
 
     if provider is None:
         return None
-
     return provider.get_headers()

beamlit/authentication/clientcredentials.py

@@ -5,9 +5,10 @@ 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
 
+from beamlit.common.settings import get_settings
+
 
 @dataclass
 class DeviceLoginFinalizeResponse:
@@ -29,32 +30,32 @@ class ClientCredentials(Auth):
             raise err
 
         return {
-            'X-Beamlit-Authorization': f'Bearer {self.credentials.access_token}',
-            'X-Beamlit-Workspace': self.workspace_name
+            "X-Beamlit-Authorization": f"Bearer {self.credentials.access_token}",
+            "X-Beamlit-Workspace": self.workspace_name,
         }
 
     def refresh_if_needed(self) -> Optional[Exception]:
         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" }
+            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']
+            self.credentials.access_token = response.json()["access_token"]
+            self.credentials.refresh_token = response.json()["refresh_token"]
+            self.credentials.expires_in = response.json()["expires_in"]
 
         # Need to refresh token if expires in less than 10 minutes
-        parts = self.credentials.access_token.split('.')
+        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_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'])
+        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()
@@ -66,8 +67,8 @@ class ClientCredentials(Auth):
         if err:
             return err
 
-        request.headers['X-Beamlit-Authorization'] = f'Bearer {self.credentials.access_token}'
-        request.headers['X-Beamlit-Workspace'] = self.workspace_name
+        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]:
@@ -79,15 +80,11 @@ class ClientCredentials(Auth):
             "grant_type": "refresh_token",
             "refresh_token": self.credentials.refresh_token,
             "device_code": self.credentials.device_code,
-            "client_id": "beamlit"
+            "client_id": "beamlit",
         }
 
         try:
-            response = post(
-                url,
-                json=refresh_data,
-                headers={"Content-Type": "application/json"}
-            )
+            response = post(url, json=refresh_data, headers={"Content-Type": "application/json"})
             response.raise_for_status()
             finalize_response = DeviceLoginFinalizeResponse(**response.json())
 
@@ -95,11 +92,12 @@ class ClientCredentials(Auth):
                 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
+                device_code=self.credentials.device_code,
             )
 
             self.credentials = creds

beamlit/authentication/credentials.py

@@ -3,6 +3,7 @@ from pathlib import Path
 from typing import List
 
 import yaml
+
 from beamlit.common.settings import Settings
 
 
@@ -15,6 +16,7 @@ class Credentials:
     device_code: str = ""
     client_credentials: str = ""
 
+
 @dataclass
 class WorkspaceConfig:
     name: str
@@ -69,15 +71,15 @@ def save_config(config: Config):
                 "name": ws.name,
                 "credentials": {
                     "access_token": ws.credentials.access_token,
-                    "api_key": ws.credentials.api_key
-                }
+                    "api_key": ws.credentials.api_key,
+                },
             }
             for ws in config.workspaces
         ],
         "context": {
             "workspace": config.context.workspace,
-            "environment": config.context.environment
-        }
+            "environment": config.context.environment,
+        },
     }
 
     home_dir = Path.home()
@@ -117,11 +119,10 @@ def load_credentials(workspace_name: str) -> Credentials:
             return workspace.credentials
     return Credentials()
 
+
 def load_credentials_from_settings(config: Settings) -> Credentials:
-    return Credentials(
-        api_key=config.api_key,
-        client_credentials=config.client_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

@@ -50,23 +50,25 @@ class BearerToken(Auth):
         if err:
             raise err
         return {
-            'X-Beamlit-Authorization': f'Bearer {self.credentials.access_token}',
-            'X-Beamlit-Workspace': self.workspace_name
+            "X-Beamlit-Authorization": f"Bearer {self.credentials.access_token}",
+            "X-Beamlit-Workspace": self.workspace_name,
         }
 
     def refresh_if_needed(self) -> Optional[Exception]:
         # Need to refresh token if expires in less than 10 minutes
-        parts = self.credentials.access_token.split('.')
+        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_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'])
+        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()
@@ -78,8 +80,10 @@ class BearerToken(Auth):
         if err:
             return err
 
-        request.headers['X-Beamlit-Authorization'] = f'Bearer {self.credentials.access_token}'
-        request.headers['X-Beamlit-Workspace'] = self.workspace_name
+        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]:
@@ -91,14 +95,12 @@ class BearerToken(Auth):
             "grant_type": "refresh_token",
             "refresh_token": self.credentials.refresh_token,
             "device_code": self.credentials.device_code,
-            "client_id": "beamlit"
+            "client_id": "beamlit",
         }
 
         try:
             response = post(
-                url,
-                json=refresh_data,
-                headers={"Content-Type": "application/json"}
+                url, json=refresh_data, headers={"Content-Type": "application/json"}
             )
             response.raise_for_status()
             finalize_response = DeviceLoginFinalizeResponse(**response.json())
@@ -107,11 +109,12 @@ class BearerToken(Auth):
                 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
+                device_code=self.credentials.device_code,
             )
 
             self.credentials = creds

beamlit/common/__init__.py

@@ -0,0 +1,13 @@
+from .logger import init as init_logger
+from .secrets import Secret
+from .settings import Settings, get_settings, init_agent
+from .utils import copy_folder
+
+__all__ = [
+    "Secret",
+    "Settings",
+    "get_settings",
+    "init_agent",
+    "copy_folder",
+    "init_logger",
+]

beamlit/common/generate.py

@@ -1,4 +1,4 @@
-from typing import Dict, Tuple
+from typing import Tuple
 
 from beamlit.common.settings import Settings, get_settings
 from beamlit.models.agent_deployment import AgentDeployment
@@ -9,6 +9,7 @@ from beamlit.models.function_kit import FunctionKit
 def get_titles_name(name: str) -> str:
     return name.title().replace("-", "").replace("_", "")
 
+
 def generate_kit_function_code(settings: Settings, function: FunctionDeployment, kit: FunctionKit) -> Tuple[str, str]:
     export_code = ""
     code = ""
@@ -24,7 +25,10 @@ def generate_kit_function_code(settings: Settings, function: FunctionDeployment,
         export_code += export
     return code, export_code
 
-def generate_function_code(settings: Settings, function: FunctionDeployment, force_name_in_endpoint: str = "", kit: bool = False) -> Tuple[str, str]:
+
+def generate_function_code(
+    settings: Settings, function: FunctionDeployment, force_name_in_endpoint: str = "", kit: bool = False
+) -> Tuple[str, str]:
     name = get_titles_name(function.function)
     if function.parameters and len(function.parameters) > 0:
         args_list = ", ".join(f"{param.name}: str" for param in function.parameters)
@@ -55,7 +59,8 @@ def generate_function_code(settings: Settings, function: FunctionDeployment, for
             if len(body) > 0:
                 body += ", "
             body += f'"name": "{function.function}"'
-    return f'''
+    return (
+        f'''
 
 class Beamlit{name}Input(BaseModel):
     {args_schema}
@@ -78,13 +83,17 @@ class Beamlit{name}(BaseTool):
             return response.json(), {{}}
         except Exception as e:
             return repr(e), {{}}
-''', f'Beamlit{get_titles_name(function.function)},'
+''',
+        f"Beamlit{get_titles_name(function.function)},",
+    )
+
 
 def generate_chain_code(settings: Settings, agent: AgentDeployment) -> Tuple[str, str]:
     name = get_titles_name(agent.agent)
     # TODO: add return direct in agent configuration
     return_direct = False
-    return f'''
+    return (
+        f'''
 class BeamlitChain{name}Input(BaseModel):
     input: str = Field(description='{agent.description}')
 
@@ -113,10 +122,13 @@ class BeamlitChain{name}(BaseTool):
                 return response.text, {{}}
         except Exception as e:
             return repr(e), {{}}
-''', f'BeamlitChain{name},'
+''',
+        f"BeamlitChain{name},",
+    )
+
 
 def generate(destination: str, dry_run: bool = False):
-    imports = '''from logging import getLogger
+    imports = """from logging import getLogger
 from typing import Dict, List, Literal, Optional, Tuple, Type, Union
 
 from langchain_core.callbacks import CallbackManagerForToolRun
@@ -138,10 +150,10 @@ client_config = RunClientWithCredentials(
 )
 client = new_client_with_credentials(client_config)
 run_client = RunClient(client=client)
-'''
+"""
     settings = get_settings()
-    export_code = '\n\nfunctions = ['
-    export_chain = '\n\nchains = ['
+    export_code = "\n\nfunctions = ["
+    export_chain = "\n\nchains = ["
     code = imports
     if settings.agent_functions and len(settings.agent_functions) > 0:
         for function_config in settings.agent_functions:
@@ -160,12 +172,12 @@ run_client = RunClient(client=client)
             export_chain += export
     if settings.agent_functions and len(settings.agent_functions) > 0:
         export_code = export_code[:-1]
-    export_code += ']'
+    export_code += "]"
     if settings.agent_chain and len(settings.agent_chain) > 0:
         export_chain = export_chain[:-1]
-    export_chain += ']'
+    export_chain += "]"
     content = code + export_code + export_chain
     if not dry_run:
         with open(destination, "w") as f:
             f.write(content)
-    return content
+    return content

beamlit/common/logger.py

@@ -1,21 +1,21 @@
 import logging
-import os
 
 
 class ColoredFormatter(logging.Formatter):
     COLORS = {
-        'DEBUG': '\033[1;36m',    # Cyan
-        'INFO': '\033[1;32m',     # Green
-        'WARNING': '\033[1;33m',  # Yellow
-        'ERROR': '\033[1;31m',    # Red
-        'CRITICAL': '\033[1;41m'  # Red background
+        "DEBUG": "\033[1;36m",  # Cyan
+        "INFO": "\033[1;32m",  # Green
+        "WARNING": "\033[1;33m",  # Yellow
+        "ERROR": "\033[1;31m",  # Red
+        "CRITICAL": "\033[1;41m",  # Red background
     }
 
     def format(self, record):
-        color = self.COLORS.get(record.levelname, '\033[0m')
+        color = self.COLORS.get(record.levelname, "\033[0m")
         record.levelname = f"{color}{record.levelname}\033[0m"
         return super().format(record)
 
+
 def init(log_level: str):
     logging.getLogger("uvicorn.access").handlers.clear()
     logging.getLogger("uvicorn.access").propagate = False
@@ -25,8 +25,5 @@ def init(log_level: str):
     logging.getLogger("httpx").propagate = False
 
     handler = logging.StreamHandler()
-    handler.setFormatter(ColoredFormatter('%(levelname)s:\t  %(name)s - %(message)s'))
-    logging.basicConfig(
-        level=log_level,
-        handlers=[handler]
-    )
+    handler.setFormatter(ColoredFormatter("%(levelname)s:\t  %(name)s - %(message)s"))
+    logging.basicConfig(level=log_level, handlers=[handler])

beamlit/common/secrets.py

@@ -0,0 +1,11 @@
+import os
+
+
+class Secret:
+    @staticmethod
+    def get(name: str):
+        return os.getenv(name, os.getenv(f"bl_{name}"))
+
+    @staticmethod
+    def set(name: str, value: str):
+        os.environ[name] = value

beamlit/common/settings.py

@@ -1,26 +1,29 @@
 import os
 from logging import getLogger
-from typing import Dict, List, Tuple, Type, Union
+from typing import Any, List, Tuple, Type, Union
+
+from langchain_core.language_models.chat_models import BaseChatModel
+from langgraph.graph.graph import CompiledGraph
+from pydantic import Field
+from pydantic_settings import BaseSettings, PydanticBaseSettingsSource, SettingsConfigDict, YamlConfigSettingsSource
 
 from beamlit.api.functions import get_function_deployment
 from beamlit.api.models import get_model_deployment
 from beamlit.client import AuthenticatedClient
 from beamlit.common.logger import init as init_logger
-from beamlit.models.agent_chain import AgentChain
 from beamlit.models.agent_deployment import AgentDeployment
 from beamlit.models.function_deployment import FunctionDeployment
 from beamlit.models.model_deployment import ModelDeployment
 from beamlit.types import UNSET, Unset
-from pydantic import Field
-from pydantic_settings import (BaseSettings, PydanticBaseSettingsSource,
-                               SettingsConfigDict, YamlConfigSettingsSource)
 
 global SETTINGS
 SETTINGS = None
 
+
 def get_settings():
     return SETTINGS
 
+
 class Settings(BaseSettings):
     model_config = SettingsConfigDict(
         yaml_file="beamlit.yaml",
@@ -39,9 +42,13 @@ class Settings(BaseSettings):
     api_key: Union[None, str] = None
     jwt: Union[None, str] = None
     client_credentials: Union[None, str] = None
-    agent_chain:  Union[Unset, List[AgentDeployment]] = UNSET
+    agent_module: str = Field(default="main.main")
+    agent_chain: Union[Unset, List[AgentDeployment]] = UNSET
     agent_functions: Union[Unset, List[FunctionDeployment]] = UNSET
     agent_model: Union[Unset, ModelDeployment] = UNSET
+    agent: Union[None, CompiledGraph, BaseChatModel] = None
+    agent_chat_model: Union[None, BaseChatModel] = None
+    agent_functions: Union[None, List[Any]] = None
 
     @classmethod
     def settings_customise_sources(
@@ -52,9 +59,19 @@ class Settings(BaseSettings):
         dotenv_settings: PydanticBaseSettingsSource,
         file_secret_settings: PydanticBaseSettingsSource,
     ) -> Tuple[PydanticBaseSettingsSource, ...]:
-        return (env_settings, dotenv_settings, file_secret_settings, YamlConfigSettingsSource(settings_cls), init_settings, )
-
-def init_agent(client: AuthenticatedClient, destination: str = f"{os.getcwd()}/src/beamlit_generated.py"):
+        return (
+            env_settings,
+            dotenv_settings,
+            file_secret_settings,
+            YamlConfigSettingsSource(settings_cls),
+            init_settings,
+        )
+
+
+def init_agent(
+    client: AuthenticatedClient,
+    destination: str = f"{os.getcwd()}/src/beamlit_generated.py",
+):
     from beamlit.api.agents import get_agent_deployment
     from beamlit.common.generate import generate
 
@@ -86,18 +103,19 @@ def init_agent(client: AuthenticatedClient, destination: str = f"{os.getcwd()}/s
                 agent_chain_deployments.append(agent_deployment)
         SETTINGS.agent_chain = agent_chain_deployments
     if agent_deployment.model:
-        model_deployment = get_model_deployment.sync(agent_deployment.model, env , client=client)
+        model_deployment = get_model_deployment.sync(agent_deployment.model, env, client=client)
         SETTINGS.agent_model = model_deployment
 
     content_generate = generate(destination, dry_run=True)
     compared_content = None
     if os.path.exists(destination):
-        compared_content = open(destination, "r").read()
+        compared_content = open(destination).read()
 
     if not os.path.exists(destination) or (compared_content and content_generate != compared_content):
         logger.info("Generating agent code")
         generate(destination)
 
+
 def init() -> Settings:
     """Parse the beamlit.yaml file to get configurations."""
     global SETTINGS

beamlit/common/utils.py

@@ -9,5 +9,7 @@ def copy_folder(source_folder: str, destination_folder: str):
             if not os.path.exists(f"{destination_folder}/{file}"):
                 os.makedirs(f"{destination_folder}/{file}")
             copy_folder(f"{source_folder}/{file}", f"{destination_folder}/{file}")
-        elif not os.path.exists(f"{destination_folder}/{file}") or not filecmp.cmp(f"{source_folder}/{file}", f"{destination_folder}/{file}"):
-            shutil.copy(f"{source_folder}/{file}", f"{destination_folder}/{file}")
+        elif not os.path.exists(f"{destination_folder}/{file}") or not filecmp.cmp(
+            f"{source_folder}/{file}", f"{destination_folder}/{file}"
+        ):
+            shutil.copy(f"{source_folder}/{file}", f"{destination_folder}/{file}")

beamlit/functions/__init__.py

@@ -0,0 +1,5 @@
+"""Functions package providing function decorators and utilities."""
+
+from .decorator import function, kit
+
+__all__ = ["function", "kit"]