beamlit 0.0.19__py3-none-any.whl → 0.0.20rc1__py3-none-any.whl

beamlit/common/generate.py

@@ -0,0 +1,166 @@
+from typing import Dict, 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:
+        body = {"function": kit.name, "workspace": settings.workspace, **kit.to_dict()}
+        new_code, export = generate_function_code(body, 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_config,
+                                    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_config(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)
+                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,30 @@
+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
+    }
+
+    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("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]
+    )

beamlit/common/settings.py

@@ -0,0 +1,112 @@
+import os
+from logging import getLogger
+from typing import Dict, List, Tuple, Type, Union
+
+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(
+        cli_parse_args=True,
+        yaml_file="beamlit.yaml",
+        env_prefix="bl_",
+    )
+
+    workspace: str
+    type: str
+    name: str
+    environment: str = Field(default="production")
+    base_url: str = Field(default="https://api.beamlit.dev/v0")
+    run_url: str = Field(default="https://run.beamlit.dev")
+    port: int = Field(default=80)
+    host: str = Field(default="0.0.0.0")
+    log_level: str = Field(default="INFO")
+    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_functions: Union[Unset, List[FunctionDeployment]] = UNSET
+    agent_model: Union[Unset, ModelDeployment] = UNSET
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: Type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        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):
+    from beamlit.api.agents import get_agent_deployment
+    from beamlit.common.generate import generate
+
+    logger = getLogger(__name__)
+
+    # Init configuration from environment variables
+    if SETTINGS.agent_functions or SETTINGS.agent_chain:
+        return
+
+    # Init configuration from beamlit control plane
+    name = SETTINGS.name
+    env = SETTINGS.environment
+
+    agent_deployment = get_agent_deployment.sync(name, env, client=client)
+    function_deployments = []
+    agent_chain_deployments = []
+    if agent_deployment.functions:
+        for function in agent_deployment.functions:
+            function_deployment = get_function_deployment.sync(function, env, client=client)
+            function_deployments.append(function_deployment)
+        SETTINGS.agent_functions = function_deployments
+
+    if agent_deployment.agent_chain:
+        for chain in agent_deployment.agent_chain:
+            if chain.enabled:
+                agent_deployment = get_agent_deployment.sync(chain.name, env, client=client)
+                if chain.description:
+                    agent_deployment.description = chain.description
+                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)
+        SETTINGS.agent_model = model_deployment
+
+
+    destination = f"{os.getcwd()}/src/beamlit_generated.py"
+
+
+    content_generate = generate(destination, dry_run=True)
+    compared_content = None
+    if os.path.exists(destination):
+        compared_content = open(destination, "r").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
+
+    SETTINGS = Settings()
+    init_logger(SETTINGS.log_level)
+    return SETTINGS

beamlit/common/utils.py

@@ -0,0 +1,13 @@
+import filecmp
+import os
+import shutil
+
+
+def copy_folder(source_folder: str, destination_folder: str):
+    for file in os.listdir(source_folder):
+        if os.path.isdir(f"{source_folder}/{file}"):
+            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}")

beamlit/models/__init__.py

@@ -5,11 +5,12 @@ from .agent import Agent
 from .agent_chain import AgentChain
 from .agent_configuration import AgentConfiguration
 from .agent_deployment import AgentDeployment
-from .agent_deployment_configuration import AgentDeploymentConfiguration
+from .agent_deployment_configuration_type_0 import AgentDeploymentConfigurationType0
 from .agent_deployment_history import AgentDeploymentHistory
 from .agent_deployment_history_event import AgentDeploymentHistoryEvent
-from .agent_deployment_pod_template import AgentDeploymentPodTemplate
+from .agent_deployment_pod_template_type_0 import AgentDeploymentPodTemplateType0
 from .agent_release import AgentRelease
+from .agent_with_deployments import AgentWithDeployments
 from .api_key import ApiKey
 from .authentication_provider_model import AuthenticationProviderModel
 from .authentication_provider_organization import AuthenticationProviderOrganization
@@ -22,19 +23,28 @@ from .create_workspace_service_account_response_200 import CreateWorkspaceServic
 from .delete_workspace_service_account_response_200 import DeleteWorkspaceServiceAccountResponse200
 from .deployment_configuration import DeploymentConfiguration
 from .deployment_configurations import DeploymentConfigurations
-from .deployment_serverless_config import DeploymentServerlessConfig
+from .deployment_serverless_config_type_0 import DeploymentServerlessConfigType0
 from .environment import Environment
 from .environment_metrics import EnvironmentMetrics
 from .flavor import Flavor
 from .function import Function
 from .function_configuration import FunctionConfiguration
 from .function_deployment import FunctionDeployment
-from .function_deployment_configuration import FunctionDeploymentConfiguration
-from .function_deployment_pod_template import FunctionDeploymentPodTemplate
+from .function_deployment_configuration_type_0 import FunctionDeploymentConfigurationType0
+from .function_deployment_pod_template_type_0 import FunctionDeploymentPodTemplateType0
 from .function_kit import FunctionKit
 from .function_provider_ref import FunctionProviderRef
 from .function_release import FunctionRelease
+from .function_with_deployments import FunctionWithDeployments
 from .get_workspace_service_accounts_response_200_item import GetWorkspaceServiceAccountsResponse200Item
+from .increase_and_rate_metric import IncreaseAndRateMetric
+from .integration import Integration
+from .integration_config import IntegrationConfig
+from .integration_connection import IntegrationConnection
+from .integration_connection_config import IntegrationConnectionConfig
+from .integration_connection_secret import IntegrationConnectionSecret
+from .integration_model import IntegrationModel
+from .integration_secret import IntegrationSecret
 from .invite_workspace_user_body import InviteWorkspaceUserBody
 from .labels_type_0 import LabelsType0
 from .location import Location
@@ -49,7 +59,7 @@ from .model_deployment_metrics_inference_per_second_per_region import ModelDeplo
 from .model_deployment_metrics_query_per_second_per_region_per_code import (
     ModelDeploymentMetricsQueryPerSecondPerRegionPerCode,
 )
-from .model_deployment_pod_template import ModelDeploymentPodTemplate
+from .model_deployment_pod_template_type_0 import ModelDeploymentPodTemplateType0
 from .model_metrics import ModelMetrics
 from .model_provider import ModelProvider
 from .model_provider_ref import ModelProviderRef
@@ -67,27 +77,31 @@ from .provider_config import ProviderConfig
 from .qps import QPS
 from .resource_deployment_log import ResourceDeploymentLog
 from .resource_deployment_metrics import ResourceDeploymentMetrics
-from .resource_deployment_metrics_inference_per_second_per_region import (
-    ResourceDeploymentMetricsInferencePerSecondPerRegion,
+from .resource_deployment_metrics_inference_per_region_type_0 import ResourceDeploymentMetricsInferencePerRegionType0
+from .resource_deployment_metrics_inference_per_second_per_region_type_0 import (
+    ResourceDeploymentMetricsInferencePerSecondPerRegionType0,
 )
-from .resource_deployment_metrics_query_per_second_per_region_per_code import (
-    ResourceDeploymentMetricsQueryPerSecondPerRegionPerCode,
+from .resource_deployment_metrics_query_per_region_per_code_type_0 import (
+    ResourceDeploymentMetricsQueryPerRegionPerCodeType0,
+)
+from .resource_deployment_metrics_query_per_second_per_region_per_code_type_0 import (
+    ResourceDeploymentMetricsQueryPerSecondPerRegionPerCodeType0,
 )
 from .resource_metrics import ResourceMetrics
 from .runtime import Runtime
-from .runtime_readiness_probe import RuntimeReadinessProbe
+from .runtime_readiness_probe_type_0 import RuntimeReadinessProbeType0
 from .runtime_resources import RuntimeResources
 from .serverless_config import ServerlessConfig
 from .standard_fields_dynamo_db import StandardFieldsDynamoDb
 from .store_agent import StoreAgent
 from .store_agent_configuration import StoreAgentConfiguration
-from .store_agent_labels import StoreAgentLabels
+from .store_agent_labels_type_0 import StoreAgentLabelsType0
 from .store_configuration import StoreConfiguration
 from .store_configuration_option import StoreConfigurationOption
 from .store_function import StoreFunction
 from .store_function_configuration import StoreFunctionConfiguration
 from .store_function_kit import StoreFunctionKit
-from .store_function_labels import StoreFunctionLabels
+from .store_function_labels_type_0 import StoreFunctionLabelsType0
 from .store_function_parameter import StoreFunctionParameter
 from .update_workspace_service_account_body import UpdateWorkspaceServiceAccountBody
 from .update_workspace_service_account_response_200 import UpdateWorkspaceServiceAccountResponse200
@@ -102,11 +116,12 @@ __all__ = (
     "AgentChain",
     "AgentConfiguration",
     "AgentDeployment",
-    "AgentDeploymentConfiguration",
+    "AgentDeploymentConfigurationType0",
     "AgentDeploymentHistory",
     "AgentDeploymentHistoryEvent",
-    "AgentDeploymentPodTemplate",
+    "AgentDeploymentPodTemplateType0",
     "AgentRelease",
+    "AgentWithDeployments",
     "ApiKey",
     "AuthenticationProviderModel",
     "AuthenticationProviderOrganization",
@@ -119,19 +134,28 @@ __all__ = (
     "DeleteWorkspaceServiceAccountResponse200",
     "DeploymentConfiguration",
     "DeploymentConfigurations",
-    "DeploymentServerlessConfig",
+    "DeploymentServerlessConfigType0",
     "Environment",
     "EnvironmentMetrics",
     "Flavor",
     "Function",
     "FunctionConfiguration",
     "FunctionDeployment",
-    "FunctionDeploymentConfiguration",
-    "FunctionDeploymentPodTemplate",
+    "FunctionDeploymentConfigurationType0",
+    "FunctionDeploymentPodTemplateType0",
     "FunctionKit",
     "FunctionProviderRef",
     "FunctionRelease",
+    "FunctionWithDeployments",
     "GetWorkspaceServiceAccountsResponse200Item",
+    "IncreaseAndRateMetric",
+    "Integration",
+    "IntegrationConfig",
+    "IntegrationConnection",
+    "IntegrationConnectionConfig",
+    "IntegrationConnectionSecret",
+    "IntegrationModel",
+    "IntegrationSecret",
     "InviteWorkspaceUserBody",
     "LabelsType0",
     "Location",
@@ -144,7 +168,7 @@ __all__ = (
     "ModelDeploymentMetrics",
     "ModelDeploymentMetricsInferencePerSecondPerRegion",
     "ModelDeploymentMetricsQueryPerSecondPerRegionPerCode",
-    "ModelDeploymentPodTemplate",
+    "ModelDeploymentPodTemplateType0",
     "ModelMetrics",
     "ModelProvider",
     "ModelProviderRef",
@@ -162,23 +186,25 @@ __all__ = (
     "QPS",
     "ResourceDeploymentLog",
     "ResourceDeploymentMetrics",
-    "ResourceDeploymentMetricsInferencePerSecondPerRegion",
-    "ResourceDeploymentMetricsQueryPerSecondPerRegionPerCode",
+    "ResourceDeploymentMetricsInferencePerRegionType0",
+    "ResourceDeploymentMetricsInferencePerSecondPerRegionType0",
+    "ResourceDeploymentMetricsQueryPerRegionPerCodeType0",
+    "ResourceDeploymentMetricsQueryPerSecondPerRegionPerCodeType0",
     "ResourceMetrics",
     "Runtime",
-    "RuntimeReadinessProbe",
+    "RuntimeReadinessProbeType0",
     "RuntimeResources",
     "ServerlessConfig",
     "StandardFieldsDynamoDb",
     "StoreAgent",
     "StoreAgentConfiguration",
-    "StoreAgentLabels",
+    "StoreAgentLabelsType0",
     "StoreConfiguration",
     "StoreConfigurationOption",
     "StoreFunction",
     "StoreFunctionConfiguration",
     "StoreFunctionKit",
-    "StoreFunctionLabels",
+    "StoreFunctionLabelsType0",
     "StoreFunctionParameter",
     "UpdateWorkspaceServiceAccountBody",
     "UpdateWorkspaceServiceAccountResponse200",