beamlit 0.0.24rc20__py3-none-any.whl → 0.0.24rc21__py3-none-any.whl

beamlit/common/generate.py

@@ -1,8 +1,7 @@
 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 import Agent, Function, FunctionMetadata, FunctionSpec
 from beamlit.models.function_kit import FunctionKit
 
 
@@ -10,55 +9,67 @@ 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]:
+def generate_kit_function_code(
+    settings: Settings, function: Function, 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,
+        fn = Function(
+            metadata=FunctionMetadata(
+                name=kit.name,
+                workspace=settings.workspace,
+                environment=settings.environment,
+            ),
+            spec=FunctionSpec(
+                parameters=kit.parameters,
+                description=kit.description,
+            ),
+        )
+        new_code, export = generate_function_code(
+            settings, fn, force_name_in_endpoint=function.metadata.name, kit=True
         )
-        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
+    settings: Settings,
+    function: Function,
+    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)
+    name = get_titles_name(function.metadata.name)
+    if function.spec.parameters and len(function.spec.parameters) > 0:
+        args_list = ", ".join(f"{param.name}: str" for param in function.spec.parameters)
         args_list += ", "
     else:
         args_list = ""
     args_schema = ""
-    if function.parameters:
-        for param in function.parameters:
+    if function.spec.parameters:
+        for param in function.spec.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
+    endpoint_name = force_name_in_endpoint or function.metadata.name
     body = "{}"
-    if function.parameters:
-        body = f'{", ".join(f'"{param.name}": {param.name}' for param in function.parameters)}'
+    if function.spec.parameters:
+        body = f'{", ".join(f'"{param.name}": {param.name}' for param in function.spec.parameters)}'
     if kit is True:
         has_name = False
-        if function.parameters:
-            for param in function.parameters:
+        if function.spec.parameters:
+            for param in function.spec.parameters:
                 if param.name == "name":
                     has_name = True
                     break
         if not has_name:
             if len(body) > 0:
                 body += ", "
-            body += f'"name": "{function.function}"'
+            body += f'"name": "{function.metadata.name}"'
     return (
         f'''
 
@@ -66,8 +77,8 @@ class Beamlit{name}Input(BaseModel):
     {args_schema}
 
 class Beamlit{name}(BaseTool):
-    name: str = "beamlit_{function.function.replace("-", "_")}"
-    description: str = """{function.description}"""
+    name: str = "beamlit_{function.metadata.name.replace("-", "_")}"
+    description: str = """{function.spec.description}"""
     args_schema: Type[BaseModel] = Beamlit{name}Input
 
     response_format: Literal["content_and_artifact"] = "content_and_artifact"
@@ -84,22 +95,22 @@ class Beamlit{name}(BaseTool):
         except Exception as e:
             return repr(e), {{}}
 ''',
-        f"Beamlit{get_titles_name(function.function)},",
+        f"Beamlit{get_titles_name(function.metadata.name)},",
     )
 
 
-def generate_chain_code(settings: Settings, agent: AgentDeployment) -> Tuple[str, str]:
-    name = get_titles_name(agent.agent)
+def generate_chain_code(settings: Settings, agent: Agent) -> Tuple[str, str]:
+    name = get_titles_name(agent.metadata.name)
     # TODO: add return direct in agent configuration
     return_direct = False
     return (
         f'''
 class BeamlitChain{name}Input(BaseModel):
-    input: str = Field(description='{agent.description}')
+    input: str = Field(description='{agent.spec.description}')
 
 class BeamlitChain{name}(BaseTool):
-    name: str = "beamlit_chain_{agent.agent.replace("-", "_")}"
-    description: str = """{agent.description}"""
+    name: str = "beamlit_chain_{agent.metadata.name.replace("-", "_")}"
+    description: str = """{agent.spec.description}"""
     args_schema: Type[BaseModel] = BeamlitChain{name}Input
 
     response_format: Literal["content_and_artifact"] = "content_and_artifact"
@@ -112,10 +123,10 @@ class BeamlitChain{name}(BaseTool):
     ) -> 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}})
+            response = run_client.run("agent", "{agent.metadata.name}", 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}}")
+                logger.error(f"Failed to run tool {agent.metadata.name}, {{response.status_code}}::{{response.text}}")
+                raise Exception(f"Failed to run tool {agent.metadata.name}, {{response.status_code}}::{{response.text}}")
             if response.headers.get("Content-Type") == "application/json":
                 return response.json(), {{}}
             else:
@@ -157,8 +168,10 @@ run_client = RunClient(client=client)
     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)
+            if function_config.spec.kit and len(function_config.spec.kit) > 0:
+                new_code, export = generate_kit_function_code(
+                    settings, function_config, function_config.spec.kit
+                )
                 code += new_code
                 export_code += export
             else:

beamlit/common/settings.py

@@ -1,14 +1,13 @@
 import os
 from logging import getLogger
-from typing import List, Tuple, Type, Union
+from typing import Tuple, Type, Union
 
-from beamlit.api.functions import get_function_deployment
-from beamlit.api.models import get_model_deployment
+from beamlit.api.agents import get_agent
+from beamlit.api.functions import get_function
+from beamlit.api.models import get_model
 from beamlit.client import AuthenticatedClient
 from beamlit.common.logger import init as init_logger
-from beamlit.models.agent_deployment import AgentDeployment
-from beamlit.models.function_deployment import FunctionDeployment
-from beamlit.models.model_deployment import ModelDeployment
+from beamlit.models import Agent, Function, Model
 from beamlit.types import UNSET, Unset
 from langchain_core.language_models.chat_models import BaseChatModel
 from langgraph.graph.graph import CompiledGraph
@@ -26,9 +25,9 @@ def get_settings():
 
 class SettingsAgent(BaseSettings):
     agent: Union[None, CompiledGraph, BaseChatModel] = None
-    chain: Union[Unset, List[AgentDeployment]] = UNSET
-    model: Union[Unset, ModelDeployment] = UNSET
-    functions: Union[Unset, List[FunctionDeployment]] = UNSET
+    chain: Union[Unset, list[Agent]] = UNSET
+    model: Union[Unset, Model] = UNSET
+    functions: Union[Unset, list[Function]] = UNSET
     functions_directory: str = Field(default="src/functions")
     chat_model: Union[None, BaseChatModel] = None
     module: str = Field(default="main.main")
@@ -106,26 +105,29 @@ def init_agent(
     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:
+    agent = get_agent.sync(name, environment=env, client=client)
+    if not agent:
+        raise ValueError(f"Agent {name} not found")
+    functions: list[Function] = []
+    agents_chain: list[Agent] = []
+    if agent.spec.functions:
+        for function in agent.spec.functions:
+            function = get_function.sync(function, environment=env, client=client)
+            if function:
+                functions.append(function)
+        settings.agent.functions = functions
+
+    if agent.spec.agent_chain:
+        for chain in agent.spec.agent_chain:
             if chain.enabled:
-                agent_deployment = get_agent_deployment.sync(chain.name, env, client=client)
+                agent_chain = get_agent.sync(chain.name, environment=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
+                    agent_chain.spec.description = chain.description
+                agents_chain.append(agent_chain)
+        settings.agent.chain = agents_chain
+    if agent.spec.model:
+        model = get_model.sync(agent.spec.model, environment=env, client=client)
+        settings.agent.model = model
 
     content_generate = generate(destination, dry_run=True)
     compared_content = None