blaxel 0.64.0__py3-none-any.whl

blaxel/__init__.py

@@ -0,0 +1,8 @@
+"""A client library for accessing Blaxel Control Plane"""
+
+from .client import AuthenticatedClient, Client
+
+__all__ = (
+    "AuthenticatedClient",
+    "Client",
+)

blaxel/agents/__init__.py

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

blaxel/agents/chain.py

@@ -0,0 +1,153 @@
+import asyncio
+import os
+import warnings
+from dataclasses import dataclass
+from typing import Callable
+
+import pydantic
+import typing_extensions as t
+from langchain_core.tools.base import BaseTool, ToolException
+
+from blaxel.aimon.settings import get_settings
+from blaxel.api.agents import list_agents
+from blaxel.authentication.authentication import AuthenticatedClient
+from blaxel.models import Agent, AgentChain
+from blaxel.run import RunClient
+
+
+class ChainTool(BaseTool):
+    """
+    A tool that allows chaining of agent actions. Extends LangChain's BaseTool.
+    """
+
+    client: RunClient
+    handle_tool_error: bool | str | Callable[[ToolException], str] | None = True
+    _cloud: bool = False
+    _service_name: str | None = None
+
+    @t.override
+    def _run(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        """
+        Executes the tool synchronously.
+
+        Parameters:
+            *args (Any): Positional arguments.
+            **kwargs (Any): Keyword arguments.
+
+        Returns:
+            Any: The result of the tool execution.
+        """
+        warnings.warn(
+            "Invoke this tool asynchronousely using `ainvoke`. This method exists only to satisfy standard tests.",
+            stacklevel=1,
+        )
+        return asyncio.run(self._arun(*args, **kwargs))
+
+    @t.override
+    async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        """
+        Executes the tool asynchronously.
+
+        Parameters:
+            *args (Any): Positional arguments.
+            **kwargs (Any): Keyword arguments.
+
+        Returns:
+            Any: The result of the asynchronous tool execution.
+        """
+        result = self.client.run(
+            "agent",
+            self.name,
+            "POST",
+            cloud=self._cloud,
+            service_name=self._service_name,
+            json=kwargs,
+        )
+        return result.text
+
+    @t.override
+    @property
+    def tool_call_schema(self) -> type[pydantic.BaseModel]:
+        """
+        Defines the schema for tool calls based on the provided argument schema.
+
+        Returns:
+            type[pydantic.BaseModel]: The Pydantic model representing the tool call schema.
+        """
+        assert self.args_schema is not None  # noqa: S101
+        return self.args_schema
+
+
+class ChainInput(pydantic.BaseModel):
+    """
+    A Pydantic model representing the input structure for a chain.
+    """
+
+    inputs: str
+
+
+@dataclass
+class ChainToolkit:
+    """
+    A toolkit for managing and initializing a chain of agents.
+    """
+
+    client: AuthenticatedClient
+    chain: list[AgentChain]
+    _cloud: bool = False
+    _service_name: str | None = None
+    _chain: list[Agent] | None = None
+
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    def initialize(self) -> None:
+        """
+        Initializes the toolkit by retrieving and configuring the list of agents based on the provided chains.
+
+        Raises:
+            RuntimeError: If initialization fails due to missing agents.
+        """
+        """Initialize the session and retrieve tools list"""
+        settings = get_settings()
+        self._cloud = settings.cloud
+        if self._chain is None:
+            agents = list_agents.sync_detailed(
+                client=self.client,
+            ).parsed
+            chain_enabled = [chain for chain in self.chain if chain.enabled]
+            agents_chain = []
+            for chain in chain_enabled:
+                agent = [agent for agent in agents if agent.metadata.name == chain.name]
+                agent_name = agent[0].metadata.name.upper().replace("-", "_")
+                if os.getenv(f"BL_AGENT_{agent_name}_SERVICE_NAME"):
+                    self._service_name = os.getenv(f"BL_AGENT_{agent_name}_SERVICE_NAME")
+                if agent:
+                    agent[0].spec.prompt = chain.prompt or agent[0].spec.prompt
+                    agent[0].spec.description = chain.description or agent[0].spec.description
+                    agents_chain.append(agent[0])
+            self._chain = agents_chain
+
+    def get_tools(self) -> list[BaseTool]:
+        """
+        Retrieves a list of tools corresponding to the initialized agents.
+
+        Returns:
+            list[BaseTool]: A list of initialized `ChainTool` instances.
+
+        Raises:
+            RuntimeError: If the toolkit has not been initialized.
+        """
+        if self._chain is None:
+            raise RuntimeError("Must initialize the toolkit first")
+
+        return [
+            ChainTool(
+                client=RunClient(self.client),
+                name=agent.metadata.name,
+                description=agent.spec.description or agent.spec.prompt or "",
+                args_schema=ChainInput,
+                cloud=self._cloud,
+                service_name=self._service_name,
+            )
+            for agent in self._chain
+        ]

blaxel/agents/chat.py

@@ -0,0 +1,286 @@
+from logging import getLogger
+from typing import Tuple, Union
+
+from langchain_core.language_models import BaseChatModel
+
+from blaxel.aimon.settings import get_settings
+from blaxel.api.models import get_model
+from blaxel.authentication import get_authentication_headers, new_client
+from blaxel.models import Model
+
+from .voice.openai import OpenAIVoiceReactAgent
+
+logger = getLogger(__name__)
+
+
+def get_base_url(name: str) -> str:
+    """
+    Constructs the base URL for a given model name based on the current settings.
+
+    Parameters:
+        name (str): The name of the model.
+
+    Returns:
+        str: The constructed base URL.
+    """
+    settings = get_settings()
+    return f"{settings.run_url}/{settings.workspace}/models/{name}/v1"
+
+
+def get_mistral_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns a MistralAI chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatMistralAI` model.
+
+    Returns:
+        ChatMistralAI: An instance of the MistralAI chat model.
+    """
+    from langchain_mistralai.chat_models import ChatMistralAI  # type: ignore
+
+    return ChatMistralAI(**kwargs)
+
+
+def get_openai_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns an OpenAI chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatOpenAI` model.
+
+    Returns:
+        ChatOpenAI: An instance of the OpenAI chat model.
+    """
+    from langchain_openai import ChatOpenAI  # type: ignore
+
+    return ChatOpenAI(**kwargs)
+
+
+def get_anthropic_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns an Anthropic chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatAnthropic` model.
+
+    Returns:
+        ChatAnthropic: An instance of the Anthropic chat model.
+    """
+    from langchain_anthropic import ChatAnthropic  # type: ignore
+
+    return ChatAnthropic(**kwargs)
+
+
+def get_xai_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns an XAI chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatXAI` model.
+
+    Returns:
+        ChatXAI: An instance of the XAI chat model.
+    """
+    from langchain_xai import ChatXAI  # type: ignore
+
+    return ChatXAI(**kwargs)
+
+
+def get_cohere_chat_model(**kwargs) -> BaseChatModel:
+    """
+    Initializes and returns a Cohere chat model with the provided keyword arguments.
+
+    Parameters:
+        **kwargs: Arbitrary keyword arguments for configuring the `ChatCohere` model.
+
+    Returns:
+        ChatCohere: An instance of the Cohere chat model.
+    """
+    from langchain_cohere import ChatCohere  # type: ignore
+
+    return ChatCohere(**kwargs)
+
+def get_deepseek_chat_model(**kwargs):
+    from langchain_deepseek import ChatDeepSeek  # type: ignore
+
+    return ChatDeepSeek(**kwargs)
+
+def get_azure_ai_inference_chat_model(**kwargs):
+    from langchain_openai import ChatOpenAI  # type: ignore
+
+    return ChatOpenAI(
+        **kwargs
+    )  # It uses a compatible endpoint, so we can use the ChatOpenAI interface
+
+def get_azure_marketplace_chat_model(**kwargs):
+    from langchain_openai import OpenAI  # type: ignore
+
+    return OpenAI(
+        **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.
+
+    Parameters:
+        name (str): The name of the model to retrieve.
+        agent_model (Union[Model, None], optional): A pre-fetched model instance.
+            If None, the model will be fetched from the API. Defaults to None.
+
+    Returns:
+        BaseChatModel: An instance of the appropriate chat model.
+    """
+    [chat_model, _, __] = get_chat_model_full(name, agent_model)
+    return chat_model
+
+def get_chat_model_full(name: str, agent_model: Union[Model, None] = None) -> Tuple[BaseChatModel, str, str]:
+    """
+    Gets a chat model instance along with provider and model information.
+
+    Parameters:
+        name (str): The name of the model to retrieve.
+        agent_model (Union[Model, None], optional): A pre-fetched model instance.
+            If None, the model will be fetched from the API. Defaults to None.
+
+    Returns:
+        Tuple[BaseChatModel, str, str]: A tuple containing:
+            - The chat model instance
+            - The provider name (e.g., 'openai', 'anthropic', etc.)
+            - The specific model name (e.g., 'gpt-4o-mini')
+    """
+    settings = get_settings()
+    client = new_client()
+
+    if agent_model is None:
+        try:
+            agent_model = get_model.sync(name, client=client)
+        except Exception:
+            logger.warning(f"Model {name} not found, defaulting to gpt-4o-mini")
+
+    headers = get_authentication_headers(settings)
+
+    jwt = headers.get("X-Blaxel-Authorization", "").replace("Bearer ", "")
+    chat_classes = {
+        "openai": {
+            "func": get_openai_chat_model,
+            "kwargs": {
+                "http_async_client": client.get_async_httpx_client(),
+                "http_client": client.get_httpx_client(),
+            },
+        },
+        "anthropic": {
+            "func": get_anthropic_chat_model,
+            "kwargs": {
+                "base_url": get_base_url(name).replace("/v1", ""),
+            },
+            "remove_kwargs": ["default_query"]
+        },
+        "mistral": {
+            "func": get_mistral_chat_model,
+            "kwargs": {
+                "api_key": jwt,
+            },
+        },
+        "xai": {
+            "func": get_xai_chat_model,
+            "kwargs": {
+                "api_key": jwt,
+                "xai_api_base": get_base_url(name),
+            },
+            "remove_kwargs": ["base_url"],
+        },
+        "cohere": {
+            "func": get_cohere_chat_model,
+            "kwargs": {
+                "cohere_api_key": jwt,
+                "base_url": get_base_url(name).replace("/v1", ""),
+            },
+        },
+        "deepseek": {
+            "func": get_deepseek_chat_model,
+            "kwargs": {
+                "api_key": jwt,
+                "api_base": get_base_url(name),
+            },
+        },
+        "azure-ai-inference": {
+            "func": get_azure_ai_inference_chat_model,
+            "kwargs": {
+                "base_url": get_base_url(name).replace("/v1", ""),
+            },
+        },
+        "azure-marketplace": {
+            "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-Blaxel-Authorization": f"Bearer {jwt}"},
+            },
+            "remove_kwargs": ["api_key", "default_headers"]
+        },
+    }
+
+    provider = (
+        agent_model
+        and agent_model.spec
+        and agent_model.spec.runtime
+        and agent_model.spec.runtime.type_
+    )
+    if not provider:
+        logger.warning("Provider not found in agent model, defaulting to OpenAI")
+        provider = "openai"
+
+    model = (
+        agent_model
+        and agent_model.spec
+        and agent_model.spec.runtime
+        and agent_model.spec.runtime.model
+    )
+    if not model:
+        logger.warning("Model not found in agent model, defaulting to gpt-4o-mini")
+        model = "gpt-4o-mini"
+
+    if provider == "openai" and "realtime" in model:
+        logger.info("Starting OpenAI Realtime Agent")
+        return (
+            OpenAIVoiceReactAgent(
+                url=get_base_url(name),
+                model=model,
+                headers=headers
+            ),
+            provider,
+            model
+        )
+    kwargs = {
+        "model": model,
+        "base_url": get_base_url(name),
+        "default_headers": headers,
+        "api_key": "fake_api_key",
+        "temperature": 0,
+    }
+    chat_class = chat_classes.get(provider)
+    if not chat_class:
+        logger.warning(f"Provider {provider} not currently supported, defaulting to OpenAI")
+        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), provider, model

blaxel/agents/decorator.py

@@ -0,0 +1,208 @@
+"""Module: decorator
+
+Defines decorators for agent functionalities.
+"""
+
+# Import necessary modules
+import asyncio
+import functools
+import inspect
+from logging import getLogger
+from typing import Any, Callable
+
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.prebuilt import create_react_agent
+
+from blaxel.aimon.settings import Settings, init
+from blaxel.api.models import get_model, list_models
+from blaxel.authentication import new_client
+from blaxel.errors import UnexpectedStatus
+from blaxel.functions import get_functions
+from blaxel.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 asyncio.iscoroutinefunction(override_agent):
+        override_agent = await override_agent()
+
+    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 blaxel.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,
+    override_agent: Any | Callable = None,
+    override_functions=None,
+    remote_functions=None,
+    local_functions=None,
+) -> Callable:
+    """
+    A decorator factory that configures and wraps functions to integrate with Blaxel agents.
+    Handles model initialization, function retrieval, and agent setup.
+
+    Parameters:
+        agent (Agent | dict, optional): An `Agent` instance or a dictionary containing agent metadata and specifications.
+        override_model (Any, optional): An optional model to override the default agent model.
+        override_agent (Any, Callable, optional): An optional agent instance to override the default agent, or a function which returns an agent instance.
+        mcp_hub (Any, optional): An optional MCP hub configuration.
+        remote_functions (Any, optional): An optional list of remote functions to be integrated.
+        local_functions (Any, optional): An optional list of local functions to be integrated.
+
+    Returns:
+        Callable: A decorator that wraps the target function, injecting agent-related configurations and dependencies.
+
+    Behavior:
+        - Validates and initializes the agent configuration.
+        - Retrieves and sets up the appropriate chat model based on the agent's specifications.
+        - Retrieves functions from the specified directories or remote sources.
+        - Wraps the target function, injecting agent, model, and function dependencies as needed.
+        - Logs relevant information and handles exceptions during the setup process.
+
+    Raises:
+        ValueError: If required configurations such as the model are missing.
+        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"}})'
+            )
+
+
+        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)
+            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 await func(*args, **kwargs)
+
+            return wrapped
+
+        return wrapper
+    except Exception as e:
+        logger.error(f"Error in agent decorator: {e!s} at line {e.__traceback__.tb_lineno}")
+        raise e

blaxel/agents/thread.py

@@ -0,0 +1,24 @@
+"""Module: thread
+
+Defines threading capabilities for agents.
+"""
+from fastapi import Request
+
+
+def get_default_thread(request: Request) -> str:
+    """
+    Extracts the default thread identifier from an incoming HTTP request.
+    Prioritizes the `X-Blaxel-Sub` header and falls back to decoding the JWT
+    from the `Authorization` or `X-Blaxel-Authorization` headers.
+
+    Parameters:
+        request (Request): The incoming HTTP request object.
+
+    Returns:
+        str: The extracted thread identifier. Returns an empty string if no valid identifier is found.
+    """
+    if request.headers.get("X-Blaxel-Thread-Id"):
+        return request.headers.get("X-Blaxel-Thread-Id")
+    if request.headers.get("Thread-Id"):
+        return request.headers.get("Thread-Id")
+    return ""