beamlit 0.0.20rc5__py3-none-any.whl → 0.0.20rc7__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:
@@ -47,8 +72,15 @@ def new_client_with_credentials(config: RunClientWithCredentials):
 
     return AuthenticatedClient(base_url=config.api_url, provider=provider)
 
+
 def get_authentication_headers(settings: Settings) -> Dict[str, str]:
-    credentials = load_credentials_from_settings(settings)
+    context = current_context()
+    if context.workspace:
+        credentials = load_credentials(context.workspace)
+    else:
+        settings = get_settings()
+        credentials = load_credentials_from_settings(settings)
+
     config = RunClientWithCredentials(
         credentials=credentials,
         workspace=settings.workspace,
@@ -63,5 +95,4 @@ def get_authentication_headers(settings: Settings) -> Dict[str, str]:
 
     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

@@ -1,10 +1,14 @@
 from dataclasses import dataclass
+from logging import getLogger
 from pathlib import Path
 from typing import List
 
 import yaml
+
 from beamlit.common.settings import Settings
 
+logger = getLogger(__name__)
+
 
 @dataclass
 class Credentials:
@@ -15,6 +19,7 @@ class Credentials:
     device_code: str = ""
     client_credentials: str = ""
 
+
 @dataclass
 class WorkspaceConfig:
     name: str
@@ -69,15 +74,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,34 +122,36 @@ def load_credentials(workspace_name: str) -> Credentials:
             return workspace.credentials
     return Credentials()
 
-def load_credentials_from_settings(config: Settings) -> Credentials:
+
+def load_credentials_from_settings(settings: Settings) -> Credentials:
     return Credentials(
-        api_key=config.api_key,
-        client_credentials=config.client_credentials
+        api_key=settings.authentication.api_key,
+        client_credentials=settings.authentication.client_credentials,
     )
 
+
 def create_home_dir_if_missing():
     home_dir = Path.home()
     if not home_dir:
-        print("Error getting home directory")
+        logger.error("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.")
+        logger.warning("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}")
+            logger.error(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")
+        logger.info("No credentials to save, error")
         return
 
     config = load_config()

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,13 +150,13 @@ 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:
+    if settings.agent.functions and len(settings.agent.functions) > 0:
+        for function_config in settings.agent.functions:
             if function_config.kit and len(function_config.kit) > 0:
                 new_code, export = generate_kit_function_code(settings, function_config, function_config.kit)
                 code += new_code
@@ -153,19 +165,19 @@ run_client = RunClient(client=client)
                 new_code, export = generate_function_code(settings, function_config)
                 code += new_code
                 export_code += export
-    if settings.agent_chain and len(settings.agent_chain) > 0:
-        for agent in settings.agent_chain:
+    if settings.agent.chain and len(settings.agent.chain) > 0:
+        for agent in settings.agent.chain:
             new_code, export = generate_chain_code(settings, agent)
             code += new_code
             export_chain += export
-    if settings.agent_functions and len(settings.agent_functions) > 0:
+    if settings.agent.functions and len(settings.agent.functions) > 0:
         export_code = export_code[:-1]
-    export_code += ']'
-    if settings.agent_chain and len(settings.agent_chain) > 0:
+    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