beamlit 0.0.18__py3-none-any.whl → 0.0.20__py3-none-any.whl

beamlit/authentication/__init__.py

@@ -1,24 +1,45 @@
 from .apikey import ApiKeyProvider
-from .authentication import (PublicProvider, RunClientWithCredentials,
-                             new_client_with_credentials)
-from .credentials import (Config, ContextConfig, Credentials, WorkspaceConfig,
-                          load_credentials)
-from .device_mode import (BearerToken, DeviceLogin, DeviceLoginFinalizeRequest,
-                          DeviceLoginFinalizeResponse, DeviceLoginResponse)
+from .authentication import (
+    PublicProvider,
+    RunClientWithCredentials,
+    get_authentication_headers,
+    new_client,
+    new_client_from_settings,
+    new_client_with_credentials,
+)
+from .credentials import (
+    Config,
+    ContextConfig,
+    Credentials,
+    WorkspaceConfig,
+    load_credentials,
+    load_credentials_from_settings,
+)
+from .device_mode import (
+    BearerToken,
+    DeviceLogin,
+    DeviceLoginFinalizeRequest,
+    DeviceLoginFinalizeResponse,
+    DeviceLoginResponse,
+)
 
 __all__ = (
     "ApiKeyProvider",
     "PublicProvider",
     "RunClientWithCredentials",
     "new_client_with_credentials",
+    "new_client_from_settings",
+    "new_client",
+    "get_authentication_headers",
     "Config",
     "ContextConfig",
     "Credentials",
     "WorkspaceConfig",
     "load_credentials",
+    "load_credentials_from_settings",
     "BearerToken",
     "DeviceLogin",
     "DeviceLoginFinalizeRequest",
     "DeviceLoginFinalizeResponse",
-    "DeviceLoginResponse"
-)
+    "DeviceLoginResponse",
+)

beamlit/authentication/apikey.py

@@ -8,7 +8,13 @@ class ApiKeyProvider(Auth):
         self.credentials = credentials
         self.workspace_name = workspace_name
 
+    def get_headers(self):
+        return {
+            "X-Beamlit-Api-Key": self.credentials.api_key,
+            "X-Beamlit-Workspace": self.workspace_name,
+        }
+
     def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
-        request.headers['X-Beamlit-Api-Key'] = self.credentials.api_key
-        request.headers['X-Beamlit-Workspace'] = self.workspace_name
+        request.headers["X-Beamlit-Api-Key"] = self.credentials.api_key
+        request.headers["X-Beamlit-Workspace"] = self.workspace_name
         yield request

beamlit/authentication/authentication.py

@@ -1,11 +1,19 @@
 from dataclasses import dataclass
-from typing import Generator
+from typing import Dict, Generator
 
 from httpx import Auth, Request, Response
 
+from beamlit.common.settings import Settings, get_settings
+
 from ..client import AuthenticatedClient
 from .apikey import ApiKeyProvider
-from .credentials import Credentials
+from .clientcredentials import ClientCredentials
+from .credentials import (
+    Credentials,
+    current_context,
+    load_credentials,
+    load_credentials_from_settings,
+)
 from .device_mode import BearerToken
 
 
@@ -14,7 +22,6 @@ class PublicProvider(Auth):
         yield request
 
 
-
 @dataclass
 class RunClientWithCredentials:
     credentials: Credentials
@@ -23,13 +30,69 @@ class RunClientWithCredentials:
     run_url: str = "https://run.beamlit.dev/v0"
 
 
+def new_client_from_settings(settings: 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():
+    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:
         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()
 
     return AuthenticatedClient(base_url=config.api_url, provider=provider)
+
+
+def get_authentication_headers(settings: Settings) -> Dict[str, str]:
+    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,
+    )
+    provider = None
+    if config.credentials.api_key:
+        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)
+
+    if provider is None:
+        return None
+    return provider.get_headers()

beamlit/authentication/clientcredentials.py

@@ -0,0 +1,108 @@
+import base64
+import json
+import time
+from dataclasses import dataclass
+from typing import Generator, Optional
+
+import requests
+from httpx import Auth, Request, Response, post
+
+from beamlit.common.settings import get_settings
+
+
+@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 get_headers(self):
+        err = self.refresh_if_needed()
+        if err:
+            raise err
+
+        return {
+            "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"}
+            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"]
+
+        # 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]:
+        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

@@ -1,9 +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:
@@ -12,6 +17,8 @@ class Credentials:
     refresh_token: str = ""
     expires_in: int = 0
     device_code: str = ""
+    client_credentials: str = ""
+
 
 @dataclass
 class WorkspaceConfig:
@@ -67,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()
@@ -116,28 +123,35 @@ def load_credentials(workspace_name: str) -> Credentials:
     return Credentials()
 
 
+def load_credentials_from_settings(settings: Settings) -> Credentials:
+    return 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

@@ -2,7 +2,7 @@ import base64
 import json
 import time
 from dataclasses import dataclass
-from typing import Generator, Optional
+from typing import Dict, Generator, Optional
 
 from httpx import Auth, Request, Response, post
 
@@ -45,19 +45,30 @@ class BearerToken(Auth):
         self.workspace_name = workspace_name
         self.base_url = base_url
 
+    def get_headers(self) -> Dict[str, str]:
+        err = self.refresh_if_needed()
+        if err:
+            raise err
+        return {
+            "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()
@@ -69,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]:
@@ -82,17 +95,13 @@ 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:
-            print(refresh_data)
             response = post(
-                url,
-                json=refresh_data,
-                headers={"Content-Type": "application/json"}
+                url, json=refresh_data, headers={"Content-Type": "application/json"}
             )
-            print(response.text)
             response.raise_for_status()
             finalize_response = DeviceLoginFinalizeResponse(**response.json())
 
@@ -100,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

@@ -0,0 +1,183 @@
+from typing import Tuple
+
+from beamlit.common.settings import Settings, get_settings
+from beamlit.models.agent_deployment import AgentDeployment
+from beamlit.models.function_deployment import FunctionDeployment
+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 = ""
+    for kit in kit:
+        fn = FunctionDeployment(
+            function=kit.name,
+            workspace=settings.workspace,
+            parameters=kit.parameters,
+            description=kit.description,
+        )
+        new_code, export = generate_function_code(settings, fn, force_name_in_endpoint=function.function, kit=True)
+        code += new_code
+        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]:
+    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)
+        args_list += ", "
+    else:
+        args_list = ""
+    args_schema = ""
+    if function.parameters:
+        for param in function.parameters:
+            args_schema += f'{param.name}: str = Field(description="""{param.description}""")\n    '
+    if len(args_schema) == 0:
+        args_schema = "pass"
+
+    # TODO: add return direct in function configuration
+    return_direct = False
+    endpoint_name = force_name_in_endpoint or function.function
+    body = "{}"
+    if function.parameters:
+        body = f'{", ".join(f'"{param.name}": {param.name}' for param in function.parameters)}'
+    if kit is True:
+        has_name = False
+        if function.parameters:
+            for param in function.parameters:
+                if param.name == "name":
+                    has_name = True
+                    break
+        if not has_name:
+            if len(body) > 0:
+                body += ", "
+            body += f'"name": "{function.function}"'
+    return (
+        f'''
+
+class Beamlit{name}Input(BaseModel):
+    {args_schema}
+
+class Beamlit{name}(BaseTool):
+    name: str = "beamlit_{function.function.replace("-", "_")}"
+    description: str = """{function.description}"""
+    args_schema: Type[BaseModel] = Beamlit{name}Input
+
+    response_format: Literal["content_and_artifact"] = "content_and_artifact"
+    return_direct: bool = {return_direct}
+
+    def _run(self, {args_list} run_manager: Optional[CallbackManagerForToolRun] = None) -> Tuple[Union[List[Dict[str, str]], str], Dict]:
+        try:
+            params = self.metadata.get("params", {{}})
+            response = run_client.run("function", "{endpoint_name}", settings.environment, "POST", json={{{body}}})
+            if response.status_code >= 400:
+                logger.error(f"Failed to run function {name}, {{response.status_code}}::{{response.text}}")
+                raise Exception(f"Failed to run function {name}, {{response.status_code}}::{{response.text}}")
+            return response.json(), {{}}
+        except Exception as e:
+            return repr(e), {{}}
+''',
+        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'''
+class BeamlitChain{name}Input(BaseModel):
+    input: str = Field(description='{agent.description}')
+
+class BeamlitChain{name}(BaseTool):
+    name: str = "beamlit_chain_{agent.agent.replace("-", "_")}"
+    description: str = """{agent.description}"""
+    args_schema: Type[BaseModel] = BeamlitChain{name}Input
+
+    response_format: Literal["content_and_artifact"] = "content_and_artifact"
+    return_direct: bool = {return_direct}
+
+    def _run(
+        self,
+        input: str,
+        run_manager: Optional[CallbackManagerForToolRun] = None,
+    ) -> Tuple[Union[List[Dict[str, str]], str], Dict]:
+        try:
+            params = self.metadata.get("params", {{}})
+            response = run_client.run("agent", "{agent.agent}", settings.environment, "POST", json={{"input": input}})
+            if response.status_code >= 400:
+                logger.error(f"Failed to run tool {agent.agent}, {{response.status_code}}::{{response.text}}")
+                raise Exception(f"Failed to run tool {agent.agent}, {{response.status_code}}::{{response.text}}")
+            if response.headers.get("Content-Type") == "application/json":
+                return response.json(), {{}}
+            else:
+                return response.text, {{}}
+        except Exception as e:
+            return repr(e), {{}}
+''',
+        f"BeamlitChain{name},",
+    )
+
+
+def generate(destination: str, dry_run: bool = False):
+    imports = """from logging import getLogger
+from typing import Dict, List, Literal, Optional, Tuple, Type, Union
+
+from langchain_core.callbacks import CallbackManagerForToolRun
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field
+from beamlit.authentication import (RunClientWithCredentials,
+                                    load_credentials_from_settings,
+                                    new_client_with_credentials)
+from beamlit.common.settings import get_settings
+from beamlit.run import RunClient
+
+logger = getLogger(__name__)
+settings = get_settings()
+credentials = load_credentials_from_settings(settings)
+
+client_config = RunClientWithCredentials(
+    credentials=credentials,
+    workspace=settings.workspace,
+)
+client = new_client_with_credentials(client_config)
+run_client = RunClient(client=client)
+"""
+    settings = get_settings()
+    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 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
+                export_code += export
+            else:
+                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:
+            new_code, export = generate_chain_code(settings, agent)
+            code += new_code
+            export_chain += export
+    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_chain = export_chain[:-1]
+    export_chain += "]"
+    content = code + export_code + export_chain
+    if not dry_run:
+        with open(destination, "w") as f:
+            f.write(content)
+    return content

beamlit/common/logger.py

@@ -0,0 +1,29 @@
+import logging
+
+
+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
+    }
+
+    def format(self, record):
+        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
+    logging.getLogger("uvicorn.error").handlers.clear()
+    logging.getLogger("uvicorn.error").propagate = False
+    logging.getLogger("httpx").handlers.clear()
+    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])