blaxel 0.64.0__py3-none-any.whl

blaxel/deploy/parser.py

@@ -0,0 +1,192 @@
+"""
+This module provides classes and functions for parsing deployment resources within Blaxel.
+It includes the Resource dataclass for representing deployment resources and functions to extract and process resources
+decorated within Python files.
+"""
+
+import ast
+import importlib
+import os
+from dataclasses import dataclass
+from logging import getLogger
+from typing import Callable, Literal
+
+from blaxel.models import StoreFunctionParameter
+
+
+@dataclass
+class Resource:
+    """
+    A dataclass representing a deployment resource.
+
+    Attributes:
+        type (Literal["agent", "function"]): The type of deployment ("agent" or "function").
+        module (Callable): The module containing the deployment.
+        name (str): The name of the deployment.
+        decorator (ast.Call): The decorator AST node used on the deployment function.
+        func (Callable): The deployment function.
+    """
+    type: Literal["agent", "function"]
+    module: Callable
+    name: str
+    decorator: ast.Call
+    func: Callable
+
+
+def get_resources(from_decorator, dir) -> list[Resource]:
+    """
+    Scans through Python files in a directory to find functions decorated with a specific decorator.
+
+    Args:
+        from_decorator (str): The name of the decorator to search for
+        dir (str): The directory to scan, defaults to "src"
+
+    Returns:
+        list[Resource]: List of Resource objects containing information about decorated functions
+    """
+    resources = []
+    logger = getLogger(__name__)
+
+    # Walk through all Python files in resources directory and subdirectories
+    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 resources
+                        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
+
+                                # Get the decorated function
+                                if hasattr(module, func_name) and isinstance(decorator, ast.Call):
+                                    resources.append(
+                                        Resource(
+                                            type=decorator_name,
+                                            module=module,
+                                            name=func_name,
+                                            func=getattr(module, func_name),
+                                            decorator=decorator,
+                                        )
+                                    )
+                    except Exception as e:
+                        logger.warning(f"Error processing {file_path}: {e!s} at line {e.__traceback__.tb_lineno}")
+    return resources
+
+
+def get_parameters(resource: Resource) -> list[StoreFunctionParameter]:
+    """
+    Extracts parameter information from a function's signature and docstring.
+
+    Args:
+        resource (Resource): The resource object containing the function to analyze
+
+    Returns:
+        list[StoreFunctionParameter]: List of parameter objects with name, type, required status, and description
+    """
+    parameters = []
+    # Get function signature
+    import inspect
+
+    sig = inspect.signature(resource.func)
+    # Get docstring for parameter descriptions
+    docstring = inspect.getdoc(resource.func)
+    param_descriptions = {}
+    if docstring:
+        # Parse docstring for parameter descriptions
+        lines = docstring.split("\n")
+        for line in lines:
+            line = line.strip().lower()
+            if line.startswith(":param "):
+                # Extract parameter name and description
+                param_line = line[7:].split(":", 1)
+                if len(param_line) == 2:
+                    param_name = param_line[0].strip()
+                    param_desc = param_line[1].strip()
+                    param_descriptions[param_name] = param_desc
+    for name, param in sig.parameters.items():
+        # Skip *args and **kwargs parameters
+        if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
+            continue
+
+        param_type = "string"  # Default type
+        type_mapping = {
+            "str": "string",
+            "int": "integer",
+            "float": "number",
+            "bool": "boolean",
+            "list": "array",
+            "dict": "object",
+            "none": "null",
+        }
+        if param.annotation != inspect.Parameter.empty:
+            # Map Python types to OpenAPI types
+            if hasattr(param.annotation, "__name__"):
+                param_type = param.annotation.__name__.lower()
+            else:
+                # Handle special types like Union, Optional etc
+                param_type = str(param.annotation).lower()
+        parameter = StoreFunctionParameter(
+            name=name,
+            type_=type_mapping.get(param_type, "string"),
+            required=param.default == inspect.Parameter.empty,
+            description=param_descriptions.get(name, f"Parameter {name}"),
+        )
+        parameters.append(parameter)
+
+    return parameters
+
+
+def get_description(description: str | None, resource: Resource) -> str:
+    """
+    Gets the description of a function from either a provided description or the function's docstring.
+
+    Args:
+        description (str | None): Optional explicit description
+        resource (Resource): The resource object containing the function
+
+    Returns:
+        str: The function description
+    """
+    if description:
+        return description
+    doc = resource.func.__doc__
+    if doc:
+        # Split docstring into sections and get only the description part
+        doc_lines = doc.split("\n")
+        description_lines = []
+        for line in doc_lines:
+            line = line.strip()
+            # Stop when we hit param/return sections
+            if line.startswith(":param") or line.startswith(":return"):
+                break
+            if line:
+                description_lines.append(line)
+        return " ".join(description_lines).strip()
+    return ""

blaxel/errors.py

@@ -0,0 +1,16 @@
+"""Contains shared errors types that can be raised from API functions"""
+
+
+class UnexpectedStatus(Exception):
+    """Raised by api functions when the response status an undocumented status and Client.raise_on_unexpected_status is True"""
+
+    def __init__(self, status_code: int, content: bytes):
+        self.status_code = status_code
+        self.content = content
+
+        super().__init__(
+            f"Unexpected status code: {status_code}\n\nResponse content:\n{content.decode(errors='ignore')}"
+        )
+
+
+__all__ = ["UnexpectedStatus"]

blaxel/functions/__init__.py

@@ -0,0 +1,7 @@
+"""Functions package providing function decorators and utilities for Blaxel integration.
+It includes decorators for creating function tools and utilities for managing and retrieving functions."""
+
+from .common import get_functions
+from .decorator import function, kit
+
+__all__ = ["function", "kit", "get_functions"]

blaxel/functions/common.py

@@ -0,0 +1,228 @@
+"""Decorators for creating function tools with Blaxel and LangChain integration.
+
+This module provides functionality to discover and load function tools from Python files,
+supporting both local and remote function execution.
+
+Key Features:
+- Automatic function discovery in specified directories
+- Support for both synchronous and asynchronous functions
+- Integration with LangChain's StructuredTool system
+- Remote function toolkit handling
+- Chain toolkit integration
+
+Main Components:
+- get_functions(): Core function that discovers and loads function tools
+"""
+import ast
+import asyncio
+import importlib.util
+import os
+import traceback
+from logging import getLogger
+from typing import Union
+
+from langchain_core.tools import StructuredTool
+from langchain_core.tools.base import create_schema_from_function
+
+from blaxel.aimon import slugify
+from blaxel.aimon.settings import get_settings
+from blaxel.authentication import new_client
+from blaxel.client import AuthenticatedClient
+from blaxel.functions.local.local import LocalToolKit
+from blaxel.functions.remote.remote import RemoteToolkit
+from blaxel.models import AgentChain
+
+logger = getLogger(__name__)
+
+MAX_RETRIES = 10
+RETRY_DELAY = 1  # 1 second delay between retries
+
+async def initialize_with_retry(toolkit, function_name: str, max_retries: int):
+    for attempt in range(1, max_retries + 1):
+        try:
+            await toolkit.initialize()
+            return await toolkit.get_tools()
+        except Exception as e:
+            if attempt == max_retries:
+                logger.warn(f"Failed to initialize function {function_name} after {max_retries} attempts: {e!s}")
+                raise
+            logger.info(f"Attempt {attempt} failed for {function_name}, retrying...")
+            await asyncio.sleep(RETRY_DELAY)
+    return []
+
+async def get_functions(
+    remote_functions: Union[list[str], None] = None,
+    local_functions: Union[list[dict], None] = None,
+    client: Union[AuthenticatedClient, None] = None,
+    dir: Union[str, None] = None,
+    chain: Union[list[AgentChain], None] = None,
+    remote_functions_empty: bool = True,
+    local_functions_empty: bool = True,
+    from_decorator: str = "function",
+    warning: bool = True,
+):
+    """Discovers and loads function tools from Python files and remote sources.
+
+    This function walks through Python files in a specified directory, looking for
+    decorated functions to convert into LangChain tools. It also handles remote
+    functions and chain toolkits.
+
+    Args:
+        remote_functions (Union[list[str], None]): List of remote function names to load
+        client (Union[AuthenticatedClient, None]): Authenticated client instance for API calls
+        dir (Union[str, None]): Directory to search for Python files containing functions
+        chain (Union[list[AgentChain], None]): List of agent chains to include
+        remote_functions_empty (bool): Whether to allow empty remote functions
+        from_decorator (str): Name of the decorator to look for (default: "function")
+        warning (bool): Whether to show warning messages
+
+    Returns:
+        list: List of discovered and loaded function tools
+
+    The function performs the following steps:
+    1. Walks through Python files in the specified directory
+    2. Parses each file to find decorated functions
+    3. Converts found functions into LangChain StructuredTools
+    4. Handles both synchronous and asynchronous functions
+    5. Processes remote functions if specified
+    6. Integrates chain toolkits if provided
+
+    Example:
+        ```python
+        tools = get_functions(
+            dir="./functions",
+            from_decorator="function",
+            warning=True
+        )
+        ```
+    """
+    from blaxel.agents.chain import ChainToolkit
+
+    settings = get_settings()
+    if client is None:
+        client = new_client()
+    if dir is None:
+        dir = settings.agent.functions_directory
+
+    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 and warning:
+            logger.warn(f"Functions directory {dir} not found")
+    if os.path.exists(dir):
+        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 = await get_functions(
+                                            client=client,
+                                            dir=os.path.join(root),
+                                            remote_functions_empty=remote_functions_empty,
+                                            local_functions_empty=local_functions_empty,
+                                            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 settings.remote:
+                                            toolkit = RemoteToolkit(client, slugify(func.__name__))
+                                            await toolkit.initialize()
+                                            functions.extend(await toolkit.get_tools())
+                                        else:
+                                            if asyncio.iscoroutinefunction(func):
+                                                functions.append(
+                                                    StructuredTool(
+                                                        name=func.__name__,
+                                                        description=func.__doc__,
+                                                        func=func,
+                                                        coroutine=func,
+                                                        args_schema=create_schema_from_function(func.__name__, func)
+                                                    )
+                                                )
+                                            else:
+
+                                                functions.append(
+                                                    StructuredTool(
+                                                        name=func.__name__,
+                                                        description=func.__doc__,
+                                                        func=func,
+                                                        args_schema=create_schema_from_function(func.__name__, func)
+                                                    )
+                                                )
+                        except Exception as e:
+                            logger.warning(f"Error processing {file_path}: {e!s}")
+    if remote_functions:
+        for function in remote_functions:
+            try:
+                toolkit = RemoteToolkit(client, function)
+                tools = await initialize_with_retry(toolkit, function, MAX_RETRIES)
+                functions.extend(tools)
+            except Exception as e:
+                if not isinstance(e, RuntimeError):
+                    logger.debug(
+                        f"Failed to initialize remote function {function}: {e!s}\n"
+                        f"Traceback:\n{traceback.format_exc()}"
+                    )
+                    logger.warn(f"Failed to initialize remote function {function}: {e!s}")
+    if local_functions:
+        for function in local_functions:
+            try:
+                toolkit = LocalToolKit(client, function)
+                tools = await initialize_with_retry(toolkit, function["name"], MAX_RETRIES)
+                functions.extend(tools)
+            except Exception as e:
+                logger.debug(
+                    f"Failed to initialize local function {function}: {e!s}\n"
+                    f"Traceback:\n{traceback.format_exc()}"
+                )
+                logger.warn(f"Failed to initialize local function {function}: {e!s}")
+
+    if chain:
+        toolkit = ChainToolkit(client, chain)
+        await toolkit.initialize()
+        functions.extend(await toolkit.get_tools())
+    return functions
+

blaxel/functions/decorator.py

@@ -0,0 +1,64 @@
+"""Decorators for creating function tools with Blaxel and LangChain integration."""
+import asyncio
+import functools
+from collections.abc import Callable
+from logging import getLogger
+
+from fastapi import Request
+
+from blaxel.models import Function, FunctionKit
+
+logger = getLogger(__name__)
+
+def kit(bl_kit: FunctionKit = None, **kwargs: dict) -> Callable:
+    """
+    Decorator to create function tools with Blaxel and LangChain integration.
+
+    Args:
+        bl_kit (FunctionKit | None): Optional FunctionKit to associate with the function.
+        **kwargs (dict): Additional keyword arguments for function configuration.
+
+    Returns:
+        Callable: The decorated function.
+    """
+
+    def wrapper(func: Callable) -> Callable:
+        if bl_kit and not func.__doc__ and bl_kit.description:
+            func.__doc__ = bl_kit.description
+        return func
+
+    return wrapper
+
+
+def function(*args, function: Function | dict = None, kit=False, **kwargs: dict) -> Callable:
+    """
+    Decorator to create function tools with Blaxel and LangChain integration.
+
+    Args:
+        function (Function | dict): Function metadata or a dictionary representing it.
+        kit (bool): Whether to associate a function kit.
+        **kwargs (dict): Additional keyword arguments for function configuration.
+
+    Returns:
+        Callable: The decorated function.
+    """
+    if function is not None and not isinstance(function, dict):
+        raise Exception(
+            'function must be a dictionary, example: @function(function={"metadata": {"name": "my_function"}})'
+        )
+    if isinstance(function, dict):
+        function = Function(**function)
+
+    def wrapper(func: Callable) -> Callable:
+        if function and not func.__doc__ and function.spec and function.spec.description:
+            func.__doc__ = function.spec.description
+
+        @functools.wraps(func)
+        async def wrapped(*args, **kwargs):
+            if len(args) > 0 and isinstance(args[0], Request):
+                body = await args[0].json()
+                args = [body.get(param) for param in func.__code__.co_varnames[:func.__code__.co_argcount]]
+            return await func(*args, **kwargs) if asyncio.iscoroutinefunction(func) else func(*args, **kwargs)
+        return wrapped
+
+    return wrapper

blaxel/functions/local/local.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+
+import pydantic
+from langchain_core.tools.base import BaseTool
+
+from blaxel.authentication.authentication import AuthenticatedClient
+from blaxel.functions.mcp.mcp import MCPClient, MCPToolkit
+from blaxel.models import Function
+
+
+@dataclass
+class LocalToolKit:
+    """
+    Toolkit for managing local tools.
+
+    Attributes:
+        client (AuthenticatedClient): The authenticated client instance.
+        function (str): The name of the local function to integrate.
+        _function (Function | None): Cached Function object after initialization.
+    """
+    client: AuthenticatedClient
+    local_function: dict
+    _function: Function | None = None
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    async def initialize(self) -> None:
+        """Initialize the session and retrieve the local function details."""
+        if self._function is None:
+            try:
+                # For local functions, we directly create the Function object
+                # based on the local function name
+                self._function = Function(
+                    metadata={"name": self.local_function['name']},
+                    spec={
+                        "configurations": {
+                            "url": self.local_function['url'],
+                        },
+                        "description": self.local_function['description'] or "",
+                    }
+                )
+            except Exception as e:
+                raise RuntimeError(f"Failed to initialize local function: {e}")
+
+    async def get_tools(self) -> list[BaseTool]:
+        mcp_client = MCPClient(self.client, self._function.spec["configurations"]["url"], None)
+        mcp_toolkit = MCPToolkit(client=mcp_client, url=self._function.spec["configurations"]["url"])
+        await mcp_toolkit.initialize()
+        return await mcp_toolkit.get_tools()

blaxel/functions/mcp/client.py

@@ -0,0 +1,96 @@
+import logging
+from contextlib import asynccontextmanager
+from typing import Any
+from urllib.parse import urljoin, urlparse
+
+import anyio
+import mcp.types as types
+from anyio.abc import TaskStatus
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from websockets.client import WebSocketClientProtocol
+from websockets.client import connect as ws_connect
+
+logger = logging.getLogger(__name__)
+
+
+def remove_request_params(url: str) -> str:
+    return urljoin(url, urlparse(url).path)
+
+
+@asynccontextmanager
+async def websocket_client(
+    url: str,
+    headers: dict[str, Any] | None = None,
+    timeout: float = 5,
+):
+    """
+    Client transport for WebSocket.
+
+    The `timeout` parameter controls connection timeout.
+    """
+    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    async with anyio.create_task_group() as tg:
+        try:
+            # Convert http(s):// to ws(s)://
+            ws_url = url.replace("http://", "ws://").replace("https://", "wss://")
+            logger.debug(f"Connecting to WebSocket endpoint: {remove_request_params(ws_url)}")
+            
+            async with ws_connect(ws_url, extra_headers=headers, open_timeout=timeout) as websocket:
+                logger.debug("WebSocket connection established")
+
+                async def ws_reader(
+                    websocket: WebSocketClientProtocol,
+                    task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
+                ):
+                    try:
+                        task_status.started()
+                        async for message in websocket:
+                            logger.debug(f"Received WebSocket message: {message}")
+                            try:
+                                parsed_message = types.JSONRPCMessage.model_validate_json(message)
+                                logger.debug(f"Received server message: {parsed_message}")
+                                await read_stream_writer.send(parsed_message)
+                            except Exception as exc:
+                                logger.error(f"Error parsing server message: {exc}")
+                                await read_stream_writer.send(exc)
+                    except Exception as exc:
+                        logger.error(f"Error in ws_reader: {exc}")
+                        await read_stream_writer.send(exc)
+                    finally:
+                        await read_stream_writer.aclose()
+
+                async def ws_writer(websocket: WebSocketClientProtocol):
+                    try:
+                        async with write_stream_reader:
+                            async for message in write_stream_reader:
+                                logger.debug(f"Sending client message: {message}")
+                                await websocket.send(
+                                    message.model_dump_json(
+                                        by_alias=True,
+                                        exclude_none=True,
+                                    )
+                                )
+                                logger.debug("Client message sent successfully")
+                    except Exception as exc:
+                        logger.error(f"Error in ws_writer: {exc}")
+                    finally:
+                        await write_stream.aclose()
+
+                await tg.start(ws_reader, websocket)
+                tg.start_soon(ws_writer, websocket)
+
+                try:
+                    yield read_stream, write_stream
+                finally:
+                    tg.cancel_scope.cancel()
+        finally:
+            await read_stream_writer.aclose()
+            await write_stream.aclose()