beamlit 0.0.56rc107__py3-none-any.whl → 0.0.57__py3-none-any.whl

beamlit/agents/chain.py

@@ -55,11 +55,9 @@ class ChainTool(BaseTool):
         Returns:
             Any: The result of the asynchronous tool execution.
         """
-        settings = get_settings()
         result = self.client.run(
             "agent",
             self.name,
-            settings.environment,
             "POST",
             cloud=self._cloud,
             service_name=self._service_name,

beamlit/agents/chat.py

@@ -120,6 +120,13 @@ def get_azure_marketplace_chat_model(**kwargs):
         **kwargs
     )  # It seems to use a compatible endpoint, so we can use the classic OpenAI interface
 
+def get_gemini_chat_model(**kwargs):
+    from langchain_google_genai import ChatGoogleGenerativeAI  # type: ignore
+
+    return ChatGoogleGenerativeAI(
+        **kwargs,
+    )
+
 def get_chat_model(name: str, agent_model: Union[Model, None] = None) -> BaseChatModel:
     """
     Gets a chat model instance for the specified model name.
@@ -155,19 +162,13 @@ def get_chat_model_full(name: str, agent_model: Union[Model, None] = None) -> Tu
 
     if agent_model is None:
         try:
-            agent_model = get_model.sync(name, client=client, environment=settings.environment)
+            agent_model = get_model.sync(name, client=client)
         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
 
     jwt = headers.get("X-Beamlit-Authorization", "").replace("Bearer ", "")
-    params = {"environment": environment}
     chat_classes = {
         "openai": {
             "func": get_openai_chat_model,
@@ -221,6 +222,18 @@ def get_chat_model_full(name: str, agent_model: Union[Model, None] = None) -> Tu
             "func": get_azure_marketplace_chat_model,
             "kwargs": {},
         },
+        "gemini": {
+            "func": get_gemini_chat_model,
+            "kwargs": {
+                "api_key": "fake_api_key",
+                "client_options": {
+                    "api_endpoint": get_base_url(name).replace("/v1", "")
+                },
+                "transport": "rest",
+                "additional_headers": {"X-Beamlit-Authorization": f"Bearer {jwt}"},
+            },
+            "remove_kwargs": ["api_key", "default_headers"]
+        },
     }
 
     provider = (
@@ -257,7 +270,6 @@ def get_chat_model_full(name: str, agent_model: Union[Model, None] = None) -> Tu
     kwargs = {
         "model": model,
         "base_url": get_base_url(name),
-        "default_query": params,
         "default_headers": headers,
         "api_key": "fake_api_key",
         "temperature": 0,

beamlit/agents/decorator.py

@@ -4,6 +4,7 @@ Defines decorators for agent functionalities.
 """
 
 # Import necessary modules
+import asyncio
 import functools
 import inspect
 from logging import getLogger
@@ -14,15 +15,113 @@ from langgraph.prebuilt import create_react_agent
 
 from beamlit.api.models import get_model, list_models
 from beamlit.authentication import new_client
-from beamlit.common.settings import init
+from beamlit.common.settings import Settings, init
 from beamlit.errors import UnexpectedStatus
 from beamlit.functions import get_functions
-from beamlit.models import Agent, AgentSpec, EnvironmentMetadata
+from beamlit.models import Agent, AgentSpec, Metadata
 
 from .chat import get_chat_model_full
 from .voice.openai import OpenAIVoiceReactAgent
 
 
+async def initialize_agent(
+    settings: Settings,
+    agent: Agent | dict = None,
+    override_model=None,
+    override_agent=None,
+    override_functions=None,
+    remote_functions=None,
+    local_functions=None,
+):
+    logger = getLogger(__name__)
+    client = new_client()
+    chat_model = override_model or None
+
+    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:
+            try:
+                response = get_model.sync_detailed(
+                    agent.spec.model, client=client
+                )
+                settings.agent.model = response.parsed
+            except UnexpectedStatus as e:
+                if e.status_code == 404:
+                    if e.status_code == 404:
+                        raise ValueError(f"Model {agent.spec.model} not found")
+                raise e
+            except Exception as e:
+                raise e
+
+            if settings.agent.model:
+                chat_model, provider, model = get_chat_model_full(agent.spec.model, settings.agent.model)
+                settings.agent.chat_model = chat_model
+                logger.info(f"Chat model configured, using: {provider}:{model}")
+
+    if override_functions is not None:
+        functions = override_functions
+    else:
+        functions = await get_functions(
+            client=client,
+            dir=settings.agent.functions_directory,
+            remote_functions=remote_functions,
+            chain=agent.spec.agent_chain,
+            local_functions=local_functions,
+            remote_functions_empty=not remote_functions,
+            warning=chat_model is not None,
+        )
+    settings.agent.functions = functions
+
+    if override_agent is None:
+        if chat_model is None:
+            models_select = ""
+            try:
+                models = list_models.sync_detailed(
+                    client=client
+                )
+                models = ", ".join([model.metadata.name for model in models.parsed])
+                models_select = f"You can select one from 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'
+                f'        "prompt": "{agent.spec.prompt}",\n'
+                "    },\n"
+                "}")
+        if isinstance(chat_model, OpenAIVoiceReactAgent):
+            _agent = chat_model
+        else:
+            memory = MemorySaver()
+            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")
+            try:
+                _agent = create_react_agent(chat_model, functions, checkpointer=memory, state_modifier=agent.spec.prompt or "")
+            except AttributeError: # special case for azure-marketplace where it uses the old OpenAI interface (no tools)
+                logger.warning("Using the old OpenAI interface for Azure Marketplace, no tools available")
+                _agent = create_react_agent(chat_model, [], checkpointer=memory, state_modifier=(agent and agent.spec and agent.spec.prompt) or "")
+
+        settings.agent.agent = _agent
+    else:
+        settings.agent.agent = override_agent
+
 def agent(
     agent: Agent | dict = None,
     override_model=None,
@@ -58,17 +157,17 @@ def agent(
         Re-raises exceptions encountered during model retrieval and agent setup.
     """
     logger = getLogger(__name__)
+    settings = init()
+    _is_initialized = False
     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"}})'
             )
 
-        client = new_client()
-        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()
@@ -83,109 +182,23 @@ def agent(
             )
 
             @functools.wraps(func)
-            def wrapped(*args, **kwargs):
+            async def wrapped(*args, **kwargs):
+                nonlocal _is_initialized
+                if not _is_initialized:
+                    async with asyncio.Lock():
+                        if not _is_initialized:
+                            await initialize_agent(settings, agent, override_model, override_agent, override_functions, remote_functions, local_functions)
+                            _is_initialized = True
                 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 await func(*args, **kwargs)
 
             return wrapped
 
-        if agent is not None:
-            metadata = EnvironmentMetadata(**agent.get("metadata", {}))
-            spec = AgentSpec(**agent.get("spec", {}))
-            agent = Agent(metadata=metadata, spec=spec)
-            if agent.spec.model and chat_model is None:
-                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, provider, model = get_chat_model_full(agent.spec.model, settings.agent.model)
-                    settings.agent.chat_model = chat_model
-                    logger.info(f"Chat model configured, using: {provider}:{model}")
-
-        if override_functions is not None:
-            functions = override_functions
-        else:
-            functions = get_functions(
-                client=client,
-                dir=settings.agent.functions_directory,
-                remote_functions=remote_functions,
-                chain=agent.spec.agent_chain,
-                local_functions=local_functions,
-                remote_functions_empty=not remote_functions,
-                warning=chat_model is not None,
-            )
-
-        settings.agent.functions = functions
-
-        if override_agent is None:
-            if chat_model is None:
-                models_select = ""
-                try:
-                    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 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'
-                    f'        "prompt": "{agent.spec.prompt}",\n'
-                    "    },\n"
-                    "}")
-            if isinstance(chat_model, OpenAIVoiceReactAgent):
-                _agent = chat_model
-            else:
-                memory = MemorySaver()
-                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")
-                try:
-                    _agent = create_react_agent(chat_model, functions, checkpointer=memory, state_modifier=agent.spec.prompt or "")
-                except AttributeError: # special case for azure-marketplace where it uses the old OpenAI interface (no tools)
-                    logger.warning("Using the old OpenAI interface for Azure Marketplace, no tools available")
-                    _agent = create_react_agent(chat_model, [], checkpointer=memory, state_modifier=agent.spec.prompt or "")
-
-            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}")

beamlit/agents/thread.py

@@ -2,7 +2,6 @@
 
 Defines threading capabilities for agents.
 """
-import jwt
 from fastapi import Request
 
 
@@ -18,11 +17,8 @@ def get_default_thread(request: Request) -> str:
     Returns:
         str: The extracted thread identifier. Returns an empty string if no valid identifier is found.
     """
-    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 ""
+    if request.headers.get("X-Beamlit-Thread-Id"):
+        return request.headers.get("X-Beamlit-Thread-Id")
+    if request.headers.get("Thread-Id"):
+        return request.headers.get("Thread-Id")
+    return ""

beamlit/api/agents/delete_agent.py

@@ -6,24 +6,15 @@ import httpx
 from ... import errors
 from ...client import AuthenticatedClient, Client
 from ...models.agent import Agent
-from ...types import UNSET, Response
+from ...types import 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
@@ -53,13 +44,11 @@ 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.
@@ -71,7 +60,6 @@ def sync_detailed(
 
     kwargs = _get_kwargs(
         agent_name=agent_name,
-        environment=environment,
     )
 
     response = client.get_httpx_client().request(
@@ -85,13 +73,11 @@ 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.
@@ -104,7 +90,6 @@ def sync(
     return sync_detailed(
         agent_name=agent_name,
         client=client,
-        environment=environment,
     ).parsed
 
 
@@ -112,13 +97,11 @@ 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.
@@ -130,7 +113,6 @@ async def asyncio_detailed(
 
     kwargs = _get_kwargs(
         agent_name=agent_name,
-        environment=environment,
     )
 
     response = await client.get_async_httpx_client().request(**kwargs)
@@ -142,13 +124,11 @@ 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.
@@ -162,6 +142,5 @@ async def asyncio(
         await asyncio_detailed(
             agent_name=agent_name,
             client=client,
-            environment=environment,
         )
     ).parsed

beamlit/api/agents/get_agent.py

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