beamlit 0.0.33rc49__py3-none-any.whl → 0.0.34__py3-none-any.whl

beamlit/agents/__init__.py

@@ -1,4 +1,5 @@
 from .chat import get_chat_model
 from .decorator import agent
+from .thread import get_default_thread
 
-__all__ = ["agent", "get_chat_model"]
+__all__ = ["agent", "get_chat_model", "get_default_thread"]

beamlit/agents/chat.py

@@ -1,5 +1,9 @@
 from logging import getLogger
+from typing import Tuple, Union
 
+from langchain_core.language_models import BaseChatModel
+
+from beamlit.api.models import get_model
 from beamlit.authentication import get_authentication_headers, new_client
 from beamlit.common.settings import get_settings
 from beamlit.models import Model
@@ -39,11 +43,17 @@ def get_cohere_chat_model(**kwargs):
 
     return ChatCohere(**kwargs)
 
-def get_chat_model(name: str, agent_model: Model):
+def get_chat_model(name: str, agent_model: Union[Model, None] = None) -> Tuple[BaseChatModel, str, str]:
     settings = get_settings()
     client = new_client()
 
-    environment = (agent_model.metadata and agent_model.metadata.environment) or settings.environment
+    if agent_model is None:
+        try:
+            agent_model = get_model.sync(name, client=client, environment=settings.environment)
+        except Exception:
+            logger.warning(f"Model {name} not found, defaulting to gpt-4o-mini")
+
+    environment = (agent_model and agent_model.metadata and agent_model.metadata.environment) or settings.environment
     headers = get_authentication_headers(settings)
     headers["X-Beamlit-Environment"] = environment
 
@@ -84,7 +94,8 @@ def get_chat_model(name: str, agent_model: Model):
     }
 
     provider = (
-        agent_model.spec
+        agent_model
+        and agent_model.spec
         and agent_model.spec.runtime
         and agent_model.spec.runtime.type_
     )
@@ -93,7 +104,8 @@ def get_chat_model(name: str, agent_model: Model):
         provider = "openai"
 
     model = (
-        agent_model.spec
+        agent_model
+        and agent_model.spec
         and agent_model.spec.runtime
         and agent_model.spec.runtime.model
     )

beamlit/agents/decorator.py

@@ -1,122 +1,24 @@
 # Import necessary modules
-import ast
-import asyncio
 import functools
-import importlib
-import os
+import inspect
 from logging import getLogger
 
-from langchain_core.tools import Tool
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.prebuilt import create_react_agent
 
-from beamlit.api.models import get_model
+from beamlit.api.models import get_model, list_models
 from beamlit.authentication import new_client
-from beamlit.common import slugify
-from beamlit.common.settings import get_settings, init
+from beamlit.common.settings import init
 from beamlit.errors import UnexpectedStatus
-from beamlit.functions.mcp.mcp import MCPClient, MCPToolkit
-from beamlit.functions.remote.remote import RemoteToolkit
+from beamlit.functions import get_functions
 from beamlit.models import Agent, AgentMetadata, AgentSpec
 
-from .chain import ChainToolkit
 from .chat import get_chat_model
 
 
-def get_functions(client, dir="src/functions", from_decorator="function", remote_functions_empty=True):
-    functions = []
-    logger = getLogger(__name__)
-    settings = get_settings()
-
-    # Walk through all Python files in functions directory and subdirectories
-    if not os.path.exists(dir):
-        if remote_functions_empty:
-            logger.warn(f"Functions directory {dir} not found")
-        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 and not settings.remote:
-                                    kit_functions = get_functions(
-                                        client,
-                                        dir=os.path.join(root),
-                                        from_decorator="kit",
-                                        remote_functions_empty=remote_functions_empty,
-                                    )
-                                    functions.extend(kit_functions)
-
-                                # Get the decorated function
-                                if not is_kit and hasattr(module, func_name):
-                                    func = getattr(module, func_name)
-                                    if settings.remote:
-                                        toolkit = RemoteToolkit(client, slugify(func.__name__))
-                                        toolkit.initialize()
-                                        functions.extend(toolkit.get_tools())
-                                    else:
-                                        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_model=None,
     override_agent=None,
     mcp_hub=None,
     remote_functions=None,
@@ -129,31 +31,34 @@ def agent(
             )
 
         client = new_client()
-        chat_model = override_chat_model or None
+        chat_model = override_model or None
         settings = init()
 
         def wrapper(func):
+            agent_kwargs = any(
+                param.name == "agent"
+                for param in inspect.signature(func).parameters.values()
+            )
+            model_kwargs = any(
+                param.name == "model"
+                for param in inspect.signature(func).parameters.values()
+            )
+            functions_kwargs = any(
+                param.name == "functions"
+                for param in inspect.signature(func).parameters.values()
+            )
             @functools.wraps(func)
             def wrapped(*args, **kwargs):
-                return func(
-                    settings.agent.agent,
-                    settings.agent.chat_model,
-                    settings.agent.functions,
-                    *args,
-                    **kwargs,
-                )
+                if agent_kwargs:
+                    kwargs["agent"] = settings.agent.agent
+                if model_kwargs:
+                    kwargs["model"] = settings.agent.chat_model
+                if functions_kwargs:
+                    kwargs["functions"] = settings.agent.functions
+                return func(*args, **kwargs)
 
             return wrapped
 
-        # Initialize functions array to store decorated functions
-        functions = get_functions(
-            client,
-            dir=settings.agent.functions_directory,
-            remote_functions_empty=not remote_functions,
-        )
-
-        settings.agent.functions = functions
-
         if agent is not None:
             metadata = AgentMetadata(**agent.get("metadata", {}))
             spec = AgentSpec(**agent.get("spec", {}))
@@ -184,44 +89,52 @@ def agent(
                     settings.agent.chat_model = chat_model
                     logger.info(f"Chat model configured, using: {provider}:{model}")
 
-        if mcp_hub:
-            for server in mcp_hub:
-                try:
-                    mcp_client = MCPClient(client, server)
-                    toolkit = MCPToolkit(client=mcp_client)
-                    toolkit.initialize()
-                    functions.extend(toolkit.get_tools())
-                except Exception as e:
-                    logger.warn(f"Failed to initialize MCP server {server}: {e!s}")
+        functions = get_functions(
+            client=client,
+            dir=settings.agent.functions_directory,
+            mcp_hub=mcp_hub,
+            remote_functions=remote_functions,
+            chain=agent.spec.agent_chain,
+            remote_functions_empty=not remote_functions,
+            warning=chat_model is not None,
+        )
+        settings.agent.functions = functions
 
-        if remote_functions:
-            for function in remote_functions:
+        if override_agent is None:
+            if chat_model is None:
+                models_select = ""
                 try:
-                    toolkit = RemoteToolkit(client, function)
-                    toolkit.initialize()
-                    functions.extend(toolkit.get_tools())
-                except Exception as e:
-                    logger.warn(f"Failed to initialize remote function {function}: {e!s}")
-
-        if agent.spec.agent_chain:
-            toolkit = ChainToolkit(client, agent.spec.agent_chain)
-            toolkit.initialize()
-            functions.extend(toolkit.get_tools())
-
-        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:
+                    models = list_models.sync_detailed(
+                        environment=settings.environment, client=client
+                    )
+                    models = ", ".join([model.metadata.name for model in models.parsed])
+                    models_select = f"You can select one from the your models: {models}"
+                except Exception:
+                    pass
+                    
+                raise ValueError(f"You must provide a model.\n"
+                    f"{models_select}\n"
+                    f"You can create one at {settings.app_url}/{settings.workspace}/global-inference-network/models/create\n"
+                    "Add it to your agent spec\n"
+                    "agent={\n"
+                    "    \"metadata\": {\n"
+                    f"        \"name\": \"{agent.metadata.name}\",\n"
+                    "    },\n"
+                    "    \"spec\": {\n"
+                    "        \"model\": \"MODEL_NAME\",\n"
+                    f"        \"description\": \"{agent.spec.description}\",\n"
+                    "    },\n"
+                    "}")
             memory = MemorySaver()
-            agent = create_react_agent(chat_model, functions, checkpointer=memory)
-            settings.agent.agent = agent
+            if len(functions) == 0:
+                raise ValueError("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")
+            _agent = create_react_agent(chat_model, functions, checkpointer=memory)
+            settings.agent.agent = _agent
         else:
             settings.agent.agent = override_agent
         return wrapper

beamlit/agents/thread.py

@@ -0,0 +1,14 @@
+
+import jwt
+from fastapi import Request
+
+
+def get_default_thread(request: Request) -> str:
+    if request.headers.get("X-Beamlit-Sub"):
+        return request.headers.get("X-Beamlit-Sub")
+    authorization = request.headers.get("Authorization", request.headers.get("X-Beamlit-Authorization"))
+    if authorization and len(authorization.split("Bearer ")) > 1:
+        token = authorization.split(" ")[1]
+        decoded = jwt.decode(token, options={"verify_signature": False})
+        return decoded["sub"]
+    return ""

beamlit/api/workspaces/workspace_quotas_request.py

@@ -0,0 +1,97 @@
+from http import HTTPStatus
+from typing import Any, Optional, Union
+
+import httpx
+
+from ... import errors
+from ...client import AuthenticatedClient, Client
+from ...types import Response
+
+
+def _get_kwargs(
+    workspace_name: str,
+) -> dict[str, Any]:
+    _kwargs: dict[str, Any] = {
+        "method": "put",
+        "url": f"/workspaces/{workspace_name}/quotas",
+    }
+
+    return _kwargs
+
+
+def _parse_response(*, client: Union[AuthenticatedClient, Client], response: httpx.Response) -> Optional[Any]:
+    if response.status_code == 200:
+        return None
+    if client.raise_on_unexpected_status:
+        raise errors.UnexpectedStatus(response.status_code, response.content)
+    else:
+        return None
+
+
+def _build_response(*, client: Union[AuthenticatedClient, Client], response: httpx.Response) -> Response[Any]:
+    return Response(
+        status_code=HTTPStatus(response.status_code),
+        content=response.content,
+        headers=response.headers,
+        parsed=_parse_response(client=client, response=response),
+    )
+
+
+def sync_detailed(
+    workspace_name: str,
+    *,
+    client: AuthenticatedClient,
+) -> Response[Any]:
+    """Request workspace quotas
+
+     Send a request to update quotas for a workspace.
+
+    Args:
+        workspace_name (str):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[Any]
+    """
+
+    kwargs = _get_kwargs(
+        workspace_name=workspace_name,
+    )
+
+    response = client.get_httpx_client().request(
+        **kwargs,
+    )
+
+    return _build_response(client=client, response=response)
+
+
+async def asyncio_detailed(
+    workspace_name: str,
+    *,
+    client: AuthenticatedClient,
+) -> Response[Any]:
+    """Request workspace quotas
+
+     Send a request to update quotas for a workspace.
+
+    Args:
+        workspace_name (str):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[Any]
+    """
+
+    kwargs = _get_kwargs(
+        workspace_name=workspace_name,
+    )
+
+    response = await client.get_async_httpx_client().request(**kwargs)
+
+    return _build_response(client=client, response=response)

beamlit/authentication/clientcredentials.py

@@ -1,7 +1,7 @@
 import base64
 import json
-import time
 from dataclasses import dataclass
+from datetime import datetime, timedelta
 from typing import Generator, Optional
 
 import requests
@@ -55,11 +55,13 @@ class ClientCredentials(Auth):
         except Exception as e:
             return Exception(f"Failed to decode/parse JWT claims: {str(e)}")
 
-        exp_time = time.gmtime(claims["exp"])
+        exp_time = datetime.fromtimestamp(claims["exp"])
+        current_time = datetime.now()
         # Refresh if token expires in less than 10 minutes
-        if time.time() + (10 * 60) > time.mktime(exp_time):
+        if current_time + timedelta(minutes=10) > exp_time:
             return self.do_refresh()
 
+
         return None
 
     def auth_flow(self, request: Request) -> Generator[Request, Response, None]:

beamlit/authentication/device_mode.py

@@ -1,7 +1,7 @@
 import base64
 import json
-import time
 from dataclasses import dataclass
+from datetime import datetime, timedelta
 from typing import Dict, Generator, Optional
 
 from httpx import Auth, Request, Response, post
@@ -65,10 +65,10 @@ class BearerToken(Auth):
             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 = datetime.fromtimestamp(claims["exp"])
+        current_time = datetime.now()
         # Refresh if token expires in less than 10 minutes
-        if time.time() + (10 * 60) > time.mktime(exp_time):
+        if current_time + timedelta(minutes=10) > exp_time:
             return self.do_refresh()
 
         return None

beamlit/common/__init__.py

@@ -1,7 +1,7 @@
 from .error import HTTPError
 from .logger import init as init_logger
 from .secrets import Secret
-from .settings import Settings, get_settings, init, init_agent
+from .settings import Settings, get_settings, init
 from .slugify import slugify
 from .utils import copy_folder
 
@@ -9,7 +9,6 @@ __all__ = [
     "Secret",
     "Settings",
     "get_settings",
-    "init_agent",
     "init",
     "copy_folder",
     "init_logger",