beamlit 0.0.24rc19__tar.gz → 0.0.24rc21__tar.gz

{beamlit-0.0.24rc19 → beamlit-0.0.24rc21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: beamlit
-Version: 0.0.24rc19
+Version: 0.0.24rc21
 Summary: Add your description here
 Author-email: cploujoux <ch.ploujoux@gmail.com>
 Requires-Python: >=3.12

{beamlit-0.0.24rc19 → beamlit-0.0.24rc21}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "beamlit"
-version = "0.0.24.pre-19"
+version = "0.0.24.pre-21"
 description = "Add your description here"
 readme = "README.md"
 authors = [

{beamlit-0.0.24rc19 → beamlit-0.0.24rc21}/src/beamlit/agents/chat.py

@@ -2,14 +2,14 @@ from logging import getLogger
 
 from beamlit.authentication import get_authentication_headers, new_client
 from beamlit.common.settings import get_settings
-from beamlit.models import AgentDeployment
+from beamlit.models import Model
 
 logger = getLogger(__name__)
 
 
-def get_base_url(agent_model: AgentDeployment):
+def get_base_url(agent_model: Model):
     settings = get_settings()
-    return f"{settings.run_url}/{settings.workspace}/models/{agent_model.model}/v1"
+    return f"{settings.run_url}/{settings.workspace}/models/{agent_model.metadata.name}/v1"
 
 
 def get_mistral_chat_model(**kwargs):
@@ -29,16 +29,25 @@ def get_anthropic_chat_model(**kwargs):
 
     return ChatAnthropic(**kwargs)
 
+def get_xai_chat_model(**kwargs):
+    from langchain_xai import ChatXAI  # type: ignore
 
-def get_chat_model(agent_model: AgentDeployment):
+    return ChatXAI(**kwargs)
+
+def get_cohere_chat_model(**kwargs):
+    from langchain_cohere import ChatCohere  # type: ignore
+
+    return ChatCohere(**kwargs)
+
+def get_chat_model(agent_model: Model):
     settings = get_settings()
     client = new_client()
 
     headers = get_authentication_headers(settings)
-    headers["X-Beamlit-Environment"] = agent_model.environment
+    headers["X-Beamlit-Environment"] = agent_model.metadata.environment
 
     jwt = headers.get("X-Beamlit-Authorization", "").replace("Bearer ", "")
-    params = {"environment": agent_model.environment}
+    params = {"environment": agent_model.metadata.environment}
     chat_classes = {
         "openai": {
             "func": get_openai_chat_model,
@@ -57,19 +66,33 @@ def get_chat_model(agent_model: AgentDeployment):
                 "api_key": jwt,
             },
         },
+        "xai": {
+            "func": get_xai_chat_model,
+            "kwargs": {
+                "api_key": jwt,
+                "xai_api_base": get_base_url(),
+            },
+            "remove_kwargs": ["base_url"],
+        },
+        "cohere": {
+            "func": get_cohere_chat_model,
+            "kwargs": {
+                "cohere_api_key": jwt,
+            },
+        },
     }
 
     if agent_model is None:
         raise ValueError("agent_model not found in configuration")
-    if agent_model.runtime is None:
+    if agent_model.spec.runtime is None:
         raise ValueError("runtime not found in agent model")
-    if agent_model.runtime.type_ is None:
+    if agent_model.spec.runtime.type_ is None:
         raise ValueError("type not found in runtime")
-    if agent_model.runtime.model is None:
+    if agent_model.spec.runtime.model is None:
         raise ValueError("model not found in runtime")
 
-    provider = agent_model.runtime.type_
-    model = agent_model.runtime.model
+    provider = agent_model.spec.runtime.type_
+    model = agent_model.spec.runtime.model
 
     kwargs = {
         "model": model,
@@ -85,4 +108,7 @@ def get_chat_model(agent_model: AgentDeployment):
         chat_class = chat_classes["openai"]
     if "kwargs" in chat_class:
         kwargs.update(chat_class["kwargs"])
+    if "remove_kwargs" in chat_class:
+        for key in chat_class["remove_kwargs"]:
+            kwargs.pop(key, None)
     return chat_class["func"](**kwargs)

beamlit-0.0.24rc21/src/beamlit/agents/decorator.py

@@ -0,0 +1,185 @@
+# Import necessary modules
+import ast
+import asyncio
+import functools
+import importlib
+import os
+from logging import getLogger
+
+from beamlit.api.models import get_model
+from beamlit.authentication import new_client
+from beamlit.common.settings import get_settings, init
+from beamlit.errors import UnexpectedStatus
+from beamlit.models import Agent, AgentSpec, Metadata
+from langchain_core.tools import Tool
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.prebuilt import create_react_agent
+
+from .chat import get_chat_model
+
+
+def get_functions(dir="src/functions", from_decorator="function"):
+    functions = []
+    logger = getLogger(__name__)
+
+    # Walk through all Python files in functions directory and subdirectories
+    if not os.path.exists(dir):
+        return []
+    for root, _, files in os.walk(dir):
+        for file in files:
+            if file.endswith(".py"):
+                file_path = os.path.join(root, file)
+                # Read and compile the file content
+                with open(file_path) as f:
+                    try:
+                        file_content = f.read()
+                        # Parse the file content to find decorated functions
+                        tree = ast.parse(file_content)
+
+                        # Look for function definitions with decorators
+                        for node in ast.walk(tree):
+                            if (
+                                not isinstance(node, ast.FunctionDef)
+                                and not isinstance(node, ast.AsyncFunctionDef)
+                            ) or len(node.decorator_list) == 0:
+                                continue
+                            decorator = node.decorator_list[0]
+
+                            decorator_name = ""
+                            if isinstance(decorator, ast.Call):
+                                decorator_name = decorator.func.id
+                            if isinstance(decorator, ast.Name):
+                                decorator_name = decorator.id
+
+                            if decorator_name == from_decorator:
+                                # Get the function name and decorator name
+                                func_name = node.name
+
+                                # Import the module to get the actual function
+                                spec = importlib.util.spec_from_file_location(func_name, file_path)
+                                module = importlib.util.module_from_spec(spec)
+                                spec.loader.exec_module(module)
+                                # Check if kit=True in the decorator arguments
+                                is_kit = False
+                                if isinstance(decorator, ast.Call):
+                                    for keyword in decorator.keywords:
+                                        if keyword.arg == "kit" and isinstance(
+                                            keyword.value, ast.Constant
+                                        ):
+                                            is_kit = keyword.value.value
+                                if is_kit:
+                                    kit_functions = get_functions(
+                                        dir=os.path.join(root),
+                                        from_decorator="kit",
+                                    )
+                                    functions.extend(kit_functions)
+
+                                # Get the decorated function
+                                if not is_kit and hasattr(module, func_name):
+                                    func = getattr(module, func_name)
+                                    if asyncio.iscoroutinefunction(func):
+                                        functions.append(
+                                            Tool(
+                                                name=func.__name__,
+                                                description=func.__doc__,
+                                                func=func,
+                                                coroutine=func,
+                                            )
+                                        )
+                                    else:
+                                        functions.append(
+                                            Tool(
+                                                name=func.__name__,
+                                                description=func.__doc__,
+                                                func=func,
+                                            )
+                                        )
+                    except Exception as e:
+                        logger.warning(f"Error processing {file_path}: {e!s}")
+    return functions
+
+
+def agent(
+    agent: Agent | dict = None,
+    override_chat_model=None,
+    override_agent=None,
+):
+    logger = getLogger(__name__)
+    try:
+        if agent is not None and not isinstance(agent, dict):
+            raise Exception(
+                'agent must be a dictionary, example: @agent(agent={"metadata": {"name": "my_agent"}})'
+            )
+
+        chat_model = override_chat_model or None
+        settings = init()
+
+        def wrapper(func):
+            @functools.wraps(func)
+            def wrapped(*args, **kwargs):
+                return func(
+                    settings.agent.agent,
+                    settings.agent.chat_model,
+                    settings.agent.functions,
+                    *args,
+                    **kwargs,
+                )
+
+            return wrapped
+
+        # Initialize functions array to store decorated functions
+        functions = get_functions(dir=settings.agent.functions_directory)
+        settings.agent.functions = functions
+
+        if agent is not None:
+            metadata = Metadata(**agent.get("metadata", {}))
+            spec = AgentSpec(**agent.get("spec", {}))
+            agent = Agent(metadata=metadata, spec=spec)
+            if agent.spec.model and chat_model is None:
+                client = new_client()
+                try:
+                    response = get_model.sync_detailed(
+                        agent.spec.model, environment=settings.environment, client=client
+                    )
+                    settings.agent.model = response.parsed
+                except UnexpectedStatus as e:
+                    if e.status_code == 404 and settings.environment != "production":
+                        try:
+                            response = get_model.sync_detailed(
+                                agent.spec.model, environment="production", client=client
+                            )
+                            settings.agent.model = response.parsed
+                        except UnexpectedStatus as e:
+                            if e.status_code == 404:
+                                raise ValueError(f"Model {agent.spec.model} not found")
+                    else:
+                        raise e
+                except Exception as e:
+                    raise e
+
+                if settings.agent.model:
+                    chat_model = get_chat_model(settings.agent.model)
+                    settings.agent.chat_model = chat_model
+                    runtime = settings.agent.model.spec.runtime
+                    logger.info(f"Chat model configured, using: {runtime.type_}:{runtime.model}")
+
+        if override_agent is None and len(functions) == 0:
+            raise ValueError(
+                "You must define at least one function, you can define this function in directory "
+                f'"{settings.agent.functions_directory}". Here is a sample function you can use:\n\n'
+                "from beamlit.functions import function\n\n"
+                "@function()\n"
+                "def hello_world(query: str):\n"
+                "    return 'Hello, world!'\n"
+            )
+
+        if override_agent is None and chat_model is not None:
+            memory = MemorySaver()
+            agent = create_react_agent(chat_model, functions, checkpointer=memory)
+            settings.agent.agent = agent
+        else:
+            settings.agent.agent = override_agent
+        return wrapper
+    except Exception as e:
+        logger.error(f"Error in agent decorator: {e!s} at line {e.__traceback__.tb_lineno}")
+        raise e

{beamlit-0.0.24rc19 → beamlit-0.0.24rc21}/src/beamlit/api/agents/create_agent.py

@@ -6,13 +6,12 @@ import httpx
 from ... import errors
 from ...client import AuthenticatedClient, Client
 from ...models.agent import Agent
-from ...models.agent_with_deployments import AgentWithDeployments
 from ...types import Response
 
 
 def _get_kwargs(
     *,
-    body: AgentWithDeployments,
+    body: Agent,
 ) -> dict[str, Any]:
     headers: dict[str, Any] = {}
 
@@ -53,13 +52,12 @@ def _build_response(*, client: Union[AuthenticatedClient, Client], response: htt
 def sync_detailed(
     *,
     client: AuthenticatedClient,
-    body: AgentWithDeployments,
+    body: Agent,
 ) -> Response[Agent]:
     """Create agent by name
 
     Args:
-        body (AgentWithDeployments): Logical object representing an agent but with deployment
-            definition inside
+        body (Agent): Agent
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -83,13 +81,12 @@ def sync_detailed(
 def sync(
     *,
     client: AuthenticatedClient,
-    body: AgentWithDeployments,
+    body: Agent,
 ) -> Optional[Agent]:
     """Create agent by name
 
     Args:
-        body (AgentWithDeployments): Logical object representing an agent but with deployment
-            definition inside
+        body (Agent): Agent
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -108,13 +105,12 @@ def sync(
 async def asyncio_detailed(
     *,
     client: AuthenticatedClient,
-    body: AgentWithDeployments,
+    body: Agent,
 ) -> Response[Agent]:
     """Create agent by name
 
     Args:
-        body (AgentWithDeployments): Logical object representing an agent but with deployment
-            definition inside
+        body (Agent): Agent
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -136,13 +132,12 @@ async def asyncio_detailed(
 async def asyncio(
     *,
     client: AuthenticatedClient,
-    body: AgentWithDeployments,
+    body: Agent,
 ) -> Optional[Agent]:
     """Create agent by name
 
     Args:
-        body (AgentWithDeployments): Logical object representing an agent but with deployment
-            definition inside
+        body (Agent): Agent
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.

{beamlit-0.0.24rc19 → beamlit-0.0.24rc21}/src/beamlit/api/agents/delete_agent.py

@@ -6,15 +6,24 @@ import httpx
 from ... import errors
 from ...client import AuthenticatedClient, Client
 from ...models.agent import Agent
-from ...types import Response
+from ...types import UNSET, Response
 
 
 def _get_kwargs(
     agent_name: str,
+    *,
+    environment: str,
 ) -> dict[str, Any]:
+    params: dict[str, Any] = {}
+
+    params["environment"] = environment
+
+    params = {k: v for k, v in params.items() if v is not UNSET and v is not None}
+
     _kwargs: dict[str, Any] = {
         "method": "delete",
         "url": f"/agents/{agent_name}",
+        "params": params,
     }
 
     return _kwargs
@@ -44,11 +53,13 @@ def sync_detailed(
     agent_name: str,
     *,
     client: AuthenticatedClient,
+    environment: str,
 ) -> Response[Agent]:
     """Delete agent by name
 
     Args:
         agent_name (str):
+        environment (str):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -60,6 +71,7 @@ def sync_detailed(
 
     kwargs = _get_kwargs(
         agent_name=agent_name,
+        environment=environment,
     )
 
     response = client.get_httpx_client().request(
@@ -73,11 +85,13 @@ def sync(
     agent_name: str,
     *,
     client: AuthenticatedClient,
+    environment: str,
 ) -> Optional[Agent]:
     """Delete agent by name
 
     Args:
         agent_name (str):
+        environment (str):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -90,6 +104,7 @@ def sync(
     return sync_detailed(
         agent_name=agent_name,
         client=client,
+        environment=environment,
     ).parsed
 
 
@@ -97,11 +112,13 @@ async def asyncio_detailed(
     agent_name: str,
     *,
     client: AuthenticatedClient,
+    environment: str,
 ) -> Response[Agent]:
     """Delete agent by name
 
     Args:
         agent_name (str):
+        environment (str):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -113,6 +130,7 @@ async def asyncio_detailed(
 
     kwargs = _get_kwargs(
         agent_name=agent_name,
+        environment=environment,
     )
 
     response = await client.get_async_httpx_client().request(**kwargs)
@@ -124,11 +142,13 @@ async def asyncio(
     agent_name: str,
     *,
     client: AuthenticatedClient,
+    environment: str,
 ) -> Optional[Agent]:
     """Delete agent by name
 
     Args:
         agent_name (str):
+        environment (str):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -142,5 +162,6 @@ async def asyncio(
         await asyncio_detailed(
             agent_name=agent_name,
             client=client,
+            environment=environment,
         )
     ).parsed

{beamlit-0.0.24rc19 → beamlit-0.0.24rc21}/src/beamlit/api/agents/delete_agent_history.py

@@ -21,9 +21,7 @@ def _get_kwargs(
     return _kwargs
 
 
-def _parse_response(
-    *, client: Union[AuthenticatedClient, Client], response: httpx.Response
-) -> Optional[AgentHistory]:
+def _parse_response(*, client: Union[AuthenticatedClient, Client], response: httpx.Response) -> Optional[AgentHistory]:
     if response.status_code == 200:
         response_200 = AgentHistory.from_dict(response.json())
 
@@ -34,9 +32,7 @@ def _parse_response(
         return None
 
 
-def _build_response(
-    *, client: Union[AuthenticatedClient, Client], response: httpx.Response
-) -> Response[AgentHistory]:
+def _build_response(*, client: Union[AuthenticatedClient, Client], response: httpx.Response) -> Response[AgentHistory]:
     return Response(
         status_code=HTTPStatus(response.status_code),
         content=response.content,

beamlit-0.0.24rc19/src/beamlit/api/agents/update_agent.py → beamlit-0.0.24rc21/src/beamlit/api/agents/get_agent.py

@@ -6,28 +6,26 @@ import httpx
 from ... import errors
 from ...client import AuthenticatedClient, Client
 from ...models.agent import Agent
-from ...models.agent_with_deployments import AgentWithDeployments
-from ...types import Response
+from ...types import UNSET, Response, Unset
 
 
 def _get_kwargs(
     agent_name: str,
     *,
-    body: AgentWithDeployments,
+    environment: Union[Unset, str] = UNSET,
 ) -> dict[str, Any]:
-    headers: dict[str, Any] = {}
+    params: dict[str, Any] = {}
+
+    params["environment"] = environment
+
+    params = {k: v for k, v in params.items() if v is not UNSET and v is not None}
 
     _kwargs: dict[str, Any] = {
-        "method": "put",
+        "method": "get",
         "url": f"/agents/{agent_name}",
+        "params": params,
     }
 
-    _body = body.to_dict()
-
-    _kwargs["json"] = _body
-    headers["Content-Type"] = "application/json"
-
-    _kwargs["headers"] = headers
     return _kwargs
 
 
@@ -55,14 +53,13 @@ def sync_detailed(
     agent_name: str,
     *,
     client: AuthenticatedClient,
-    body: AgentWithDeployments,
+    environment: Union[Unset, str] = UNSET,
 ) -> Response[Agent]:
-    """Update agent by name
+    """Get agent by name
 
     Args:
         agent_name (str):
-        body (AgentWithDeployments): Logical object representing an agent but with deployment
-            definition inside
+        environment (Union[Unset, str]):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -74,7 +71,7 @@ def sync_detailed(
 
     kwargs = _get_kwargs(
         agent_name=agent_name,
-        body=body,
+        environment=environment,
     )
 
     response = client.get_httpx_client().request(
@@ -88,14 +85,13 @@ def sync(
     agent_name: str,
     *,
     client: AuthenticatedClient,
-    body: AgentWithDeployments,
+    environment: Union[Unset, str] = UNSET,
 ) -> Optional[Agent]:
-    """Update agent by name
+    """Get agent by name
 
     Args:
         agent_name (str):
-        body (AgentWithDeployments): Logical object representing an agent but with deployment
-            definition inside
+        environment (Union[Unset, str]):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -108,7 +104,7 @@ def sync(
     return sync_detailed(
         agent_name=agent_name,
         client=client,
-        body=body,
+        environment=environment,
     ).parsed
 
 
@@ -116,14 +112,13 @@ async def asyncio_detailed(
     agent_name: str,
     *,
     client: AuthenticatedClient,
-    body: AgentWithDeployments,
+    environment: Union[Unset, str] = UNSET,
 ) -> Response[Agent]:
-    """Update agent by name
+    """Get agent by name
 
     Args:
         agent_name (str):
-        body (AgentWithDeployments): Logical object representing an agent but with deployment
-            definition inside
+        environment (Union[Unset, str]):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -135,7 +130,7 @@ async def asyncio_detailed(
 
     kwargs = _get_kwargs(
         agent_name=agent_name,
-        body=body,
+        environment=environment,
     )
 
     response = await client.get_async_httpx_client().request(**kwargs)
@@ -147,14 +142,13 @@ async def asyncio(
     agent_name: str,
     *,
     client: AuthenticatedClient,
-    body: AgentWithDeployments,
+    environment: Union[Unset, str] = UNSET,
 ) -> Optional[Agent]:
-    """Update agent by name
+    """Get agent by name
 
     Args:
         agent_name (str):
-        body (AgentWithDeployments): Logical object representing an agent but with deployment
-            definition inside
+        environment (Union[Unset, str]):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -168,6 +162,6 @@ async def asyncio(
         await asyncio_detailed(
             agent_name=agent_name,
             client=client,
-            body=body,
+            environment=environment,
         )
     ).parsed