blaxel 0.1.21rc69__py3-none-any.whl → 0.2.0__py3-none-any.whl

blaxel/{common → core/common}/internal.py

@@ -1,29 +1,30 @@
 import hashlib
-import os
 from logging import getLogger
 from typing import Optional
+
 from .env import env
 
 logger = getLogger(__name__)
 
+
 def get_alphanumeric_limited_hash(input_str, max_size=48):
     """
     Create an alphanumeric hash using MD5 that can be reproduced in Go, TypeScript, and Python.
-    
+
     Args:
         input_str (str): The input string to hash
         max_size (int): The maximum length of the returned hash
-        
+
     Returns:
         str: An alphanumeric hash of the input string, limited to max_size
     """
     # Calculate MD5 hash and convert to hexadecimal
-    hash_hex = hashlib.md5(input_str.encode('utf-8')).hexdigest()
-    
+    hash_hex = hashlib.md5(input_str.encode("utf-8")).hexdigest()
+
     # Limit to max_size
     if len(hash_hex) > max_size:
         return hash_hex[:max_size]
-    
+
     return hash_hex
 
 
@@ -43,103 +44,74 @@ def get_global_unique_hash(workspace: str, type: str, name: str) -> str:
     hash = get_alphanumeric_limited_hash(global_unique_name, 48)
     return hash
 
-class Agent:
-    def __init__(self, agent_name: str, workspace: str, run_internal_protocol: str, run_internal_hostname: str):
-        self.agent_name = agent_name
-        self.workspace = workspace
-        self.run_internal_protocol = run_internal_protocol
-        self.run_internal_hostname = run_internal_hostname
-
-    @property
-    def internal_url(self) -> str:
-        """
-        Generate the internal URL for the agent using a unique hash.
-
-        Returns:
-            The internal URL as a string
-        """
-        hash_value = get_global_unique_hash(
-            self.workspace,
-            "agent",
-            self.agent_name
-        )
-        return f"{self.run_internal_protocol}://{hash_value}.{self.run_internal_hostname}"
-
-    @property
-    def forced_url(self) -> Optional[str]:
-        """
-        Check for a forced URL in environment variables.
-
-        Returns:
-            The forced URL if found in environment variables, None otherwise
-        """
-        env_var = self.agent_name.replace("-", "_").upper()
-        env_key = f"BL_AGENT_{env_var}_URL"
-        return os.environ.get(env_key)
 
 def pluralize(type_str: str) -> str:
     """
     Convert a string to its plural form following English pluralization rules.
-    
+
     Args:
         type_str: The input string to pluralize
-        
+
     Returns:
         The pluralized form of the input string
     """
     word = type_str.lower()
 
     # Words ending in s, ss, sh, ch, x, z - add 'es'
-    if (word.endswith('s') or word.endswith('ss') or word.endswith('sh') or
-        word.endswith('ch') or word.endswith('x') or word.endswith('z')):
-        return type_str + 'es'
+    if (
+        word.endswith("s")
+        or word.endswith("ss")
+        or word.endswith("sh")
+        or word.endswith("ch")
+        or word.endswith("x")
+        or word.endswith("z")
+    ):
+        return type_str + "es"
 
     # Words ending in consonant + y - change y to ies
-    if word.endswith('y') and len(word) > 1:
+    if word.endswith("y") and len(word) > 1:
         before_y = word[-2]
-        if before_y not in 'aeiou':
-            return type_str[:-1] + 'ies'
+        if before_y not in "aeiou":
+            return type_str[:-1] + "ies"
 
     # Words ending in f or fe - change to ves
-    if word.endswith('f'):
-        return type_str[:-1] + 'ves'
-    if word.endswith('fe'):
-        return type_str[:-2] + 'ves'
+    if word.endswith("f"):
+        return type_str[:-1] + "ves"
+    if word.endswith("fe"):
+        return type_str[:-2] + "ves"
 
     # Words ending in consonant + o - add 'es'
-    if word.endswith('o') and len(word) > 1:
+    if word.endswith("o") and len(word) > 1:
         before_o = word[-2]
-        if before_o not in 'aeiou':
-            return type_str + 'es'
+        if before_o not in "aeiou":
+            return type_str + "es"
 
     # Default case - just add 's'
-    return type_str + 's'
+    return type_str + "s"
 
 
 def get_forced_url(type_str: str, name: str) -> Optional[str]:
     """
     Check for forced URLs in environment variables using both plural and singular forms.
-    
+
     Args:
         type_str: The type identifier
         name: The name identifier
-        
+
     Returns:
         The forced URL if found in environment variables, None otherwise
     """
     plural_type = pluralize(type_str)
     env_var = name.replace("-", "_").upper()
-    
+
     # BL_FUNCTIONS_NAME_URL (plural form)
     plural_env_key = f"BL_{plural_type.upper()}_{env_var}_URL"
     if env[plural_env_key] is not None:
         return env[plural_env_key]
-    
+
     # BL_FUNCTION_NAME_URL (singular form)
     singular_env_key = f"BL_{type_str.upper()}_{env_var}_URL"
     if env[singular_env_key] is not None:
         return env[singular_env_key]
-    
-    return None
-
 
+    return None

blaxel/core/common/logger.py

@@ -0,0 +1,131 @@
+"""
+This module provides a custom colored formatter for logging and an initialization function
+to set up logging configurations for Blaxel applications.
+"""
+
+import json
+import logging
+import os
+
+try:
+    from opentelemetry.trace import get_current_span
+
+    HAS_OPENTELEMETRY = True
+except ImportError:
+    HAS_OPENTELEMETRY = False
+
+    def get_current_span():
+        """Fallback function when opentelemetry is not available."""
+        return None
+
+
+class JsonFormatter(logging.Formatter):
+    """
+    A logger compatible with standard json logging.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.trace_id_name = os.environ.get("BL_LOGGER_TRACE_ID", "trace_id")
+        self.span_id_name = os.environ.get("BL_LOGGER_SPAN_ID", "span_id")
+        self.labels_name = os.environ.get("BL_LOGGER_LABELS", "labels")
+        self.trace_id_prefix = os.environ.get("BL_LOGGER_TRACE_ID_PREFIX", "")
+        self.span_id_prefix = os.environ.get("BL_LOGGER_SPAN_ID_PREFIX", "")
+        self.task_index = os.environ.get("BL_TASK_KEY", "TASK_INDEX")
+        self.task_prefix = os.environ.get("BL_TASK_PREFIX", "")
+        self.execution_key = os.environ.get("BL_EXECUTION_KEY", "BL_EXECUTION_ID")
+        self.execution_prefix = os.environ.get("BL_EXECUTION_PREFIX", "")
+
+    def format(self, record):
+        """
+        Formats the log record by converting it to a JSON object with trace context and environment variables.
+        """
+        log_entry = {
+            "message": record.getMessage(),
+            "severity": record.levelname,
+            self.labels_name: {},
+        }
+
+        # Get current active span - equivalent to trace.getActiveSpan() in JS
+        if HAS_OPENTELEMETRY:
+            current_span = get_current_span()
+
+            # Check if span exists and has valid context (equivalent to 'if (currentSpan)' in JS)
+            if current_span and current_span.get_span_context().is_valid:
+                span_context = current_span.get_span_context()
+                # Format trace_id and span_id as hex strings (like JS does)
+                trace_id_hex = format(span_context.trace_id, "032x")
+                span_id_hex = format(span_context.span_id, "016x")
+
+                log_entry[self.trace_id_name] = f"{self.trace_id_prefix}{trace_id_hex}"
+                log_entry[self.span_id_name] = f"{self.span_id_prefix}{span_id_hex}"
+
+        # Add task ID if available
+        task_id = os.environ.get(self.task_index)
+        if task_id:
+            log_entry[self.labels_name]["blaxel-task"] = f"{self.task_prefix}{task_id}"
+
+        # Add execution ID if available
+        execution_id = os.environ.get(self.execution_key)
+        if execution_id:
+            log_entry[self.labels_name]["blaxel-execution"] = (
+                f"{self.execution_prefix}{execution_id.split('-')[-1]}"
+            )
+
+        return json.dumps(log_entry)
+
+
+class ColoredFormatter(logging.Formatter):
+    """
+    A custom logging formatter that adds ANSI color codes to log levels for enhanced readability.
+
+    Attributes:
+        COLORS (dict): A mapping of log level names to their corresponding ANSI color codes.
+    """
+
+    COLORS = {
+        "DEBUG": "\033[1;36m",  # Cyan
+        "INFO": "\033[1;32m",  # Green
+        "WARNING": "\033[1;33m",  # Yellow
+        "ERROR": "\033[1;31m",  # Red
+        "CRITICAL": "\033[1;41m",  # Red background
+    }
+
+    def format(self, record):
+        """
+        Formats the log record by adding color codes based on the log level.
+
+        Parameters:
+            record (LogRecord): The log record to format.
+
+        Returns:
+            str: The formatted log message with appropriate color codes.
+        """
+        n_spaces = len("CRITICAL") - len(record.levelname)
+        tab = " " * n_spaces
+        color = self.COLORS.get(record.levelname, "\033[0m")
+        record.levelname = f"{color}{record.levelname}\033[0m:{tab}"
+        return super().format(record)
+
+
+def init_logger(log_level: str):
+    """
+    Initializes the logging configuration for Blaxel.
+
+    This function clears existing handlers for specific loggers, sets up a colored formatter,
+    and configures the root logger with the specified log level.
+
+    Parameters:
+        log_level (str): The logging level to set (e.g., "DEBUG", "INFO").
+    """
+    # Disable urllib3 logging
+    logging.getLogger("urllib3").setLevel(logging.CRITICAL)
+    logging.getLogger("httpx").setLevel(logging.CRITICAL)
+    handler = logging.StreamHandler()
+
+    logger_type = os.environ.get("BL_LOGGER", "http")
+    if logger_type == "json":
+        handler.setFormatter(JsonFormatter())
+    else:
+        handler.setFormatter(ColoredFormatter("%(levelname)s %(name)s - %(message)s"))
+    logging.basicConfig(level=log_level, handlers=[handler])

blaxel/{jobs → core/jobs}/__init__.py

@@ -8,14 +8,13 @@ from typing import Any, Awaitable, Callable, Dict
 import requests
 
 from ..client import client
-from ..common.internal import get_global_unique_hash, get_forced_url
+from ..common.internal import get_forced_url, get_global_unique_hash
 from ..common.settings import settings
-from ..instrumentation.span import SpanManager
 
 
 class BlJobWrapper:
     def get_arguments(self) -> Dict[str, Any]:
-        if not os.getenv('BL_EXECUTION_DATA_URL'):
+        if not os.getenv("BL_EXECUTION_DATA_URL"):
             parser = argparse.ArgumentParser()
             # Parse known args, ignore unknown
             args, unknown = parser.parse_known_args()
@@ -24,18 +23,18 @@ class BlJobWrapper:
             # Add unknown args to dict
             for i in range(0, len(unknown), 2):
                 if i + 1 < len(unknown):
-                    key = unknown[i].lstrip('-')
+                    key = unknown[i].lstrip("-")
                     args_dict[key] = unknown[i + 1]
             return args_dict
 
-        response = requests.get(os.getenv('BL_EXECUTION_DATA_URL'))
+        response = requests.get(os.getenv("BL_EXECUTION_DATA_URL"))
         data = response.json()
-        tasks = data.get('tasks', [])
+        tasks = data.get("tasks", [])
         return tasks[self.index] if self.index < len(tasks) else {}
 
     @property
     def index_key(self) -> str:
-        return os.getenv('BL_TASK_KEY', 'TASK_INDEX')
+        return os.getenv("BL_TASK_KEY", "TASK_INDEX")
 
     @property
     def index(self) -> int:
@@ -55,14 +54,13 @@ class BlJobWrapper:
             else:
                 func(**parsed_args)
         except Exception as error:
-            print('Job execution failed:', error, file=sys.stderr)
+            print("Job execution failed:", error, file=sys.stderr)
             sys.exit(1)
 
 
-
-
 logger = getLogger(__name__)
 
+
 class BlJob:
     def __init__(self, name: str):
         self.name = name
@@ -97,73 +95,55 @@ class BlJob:
         return self.external_url
 
     def call(self, url, input_data, headers: dict = {}, params: dict = {}):
-        body = {
-            "tasks": input_data
-        }
+        body = {"tasks": input_data}
 
         return client.get_httpx_client().post(
-            url+"/executions",
-            headers={
-                'Content-Type': 'application/json',
-                **headers
-            },
+            url + "/executions",
+            headers={"Content-Type": "application/json", **headers},
             json=body,
-            params=params
+            params=params,
         )
 
     async def acall(self, url, input_data, headers: dict = {}, params: dict = {}):
         logger.debug(f"Job Calling: {self.name}")
-        body = {
-            "tasks": input_data
-        }
+        body = {"tasks": input_data}
 
         return await client.get_async_httpx_client().post(
-            url+"/executions",
-            headers={
-                'Content-Type': 'application/json',
-                **headers
-            },
+            url + "/executions",
+            headers={"Content-Type": "application/json", **headers},
             json=body,
-            params=params
+            params=params,
         )
 
     def run(self, input: Any, headers: dict = {}, params: dict = {}) -> str:
-        attributes = {
-            "job.name": self.name,
-            "span.type": "job.run",
-        }
-        with SpanManager("blaxel-tracer").create_active_span(self.name, attributes) as span:
-            logger.debug(f"Job Calling: {self.name}")
-            response = self.call(self.url, input, headers, params)
+        logger.debug(f"Job Calling: {self.name}")
+        response = self.call(self.url, input, headers, params)
+        if response.status_code >= 400:
+            if not self.fallback_url:
+                raise Exception(
+                    f"Job {self.name} returned status code {response.status_code} with body {response.text}"
+                )
+            response = self.call(self.fallback_url, input, headers, params)
             if response.status_code >= 400:
-                if not self.fallback_url:
-                    span.set_attribute("job.run.error", response.text)
-                    raise Exception(f"Job {self.name} returned status code {response.status_code} with body {response.text}")
-                response = self.call(self.fallback_url, input, headers, params)
-                if response.status_code >= 400:
-                    span.set_attribute("job.run.error", response.text)
-                    raise Exception(f"Job {self.name} returned status code {response.status_code} with body {response.text}")
-            span.set_attribute("job.run.result", response.text)
-            return response.text
+                raise Exception(
+                    f"Job {self.name} returned status code {response.status_code} with body {response.text}"
+                )
+        return response.text
 
     async def arun(self, input: Any, headers: dict = {}, params: dict = {}) -> Awaitable[str]:
-        attributes = {
-            "job.name": self.name,
-            "span.type": "job.run",
-        }
-        with SpanManager("blaxel-tracer").create_active_span(self.name, attributes) as span:
-            logger.debug(f"Job Calling: {self.name}")
-            response = await self.acall(self.url, input, headers, params)
+        logger.debug(f"Job Calling: {self.name}")
+        response = await self.acall(self.url, input, headers, params)
+        if response.status_code >= 400:
+            if not self.fallback_url:
+                raise Exception(
+                    f"Job {self.name} returned status code {response.status_code} with body {response.text}"
+                )
+            response = await self.acall(self.fallback_url, input, headers, params)
             if response.status_code >= 400:
-                if not self.fallback_url:
-                    span.set_attribute("job.run.error", response.text)
-                    raise Exception(f"Job {self.name} returned status code {response.status_code} with body {response.text}")
-                response = await self.acall(self.fallback_url, input, headers, params)
-                if response.status_code >= 400:
-                    span.set_attribute("job.run.error", response.text)
-                    raise Exception(f"Job {self.name} returned status code {response.status_code} with body {response.text}")
-            span.set_attribute("job.run.result", response.text)
-            return response.text
+                raise Exception(
+                    f"Job {self.name} returned status code {response.status_code} with body {response.text}"
+                )
+        return response.text
 
     def __str__(self):
         return f"Job {self.name}"
@@ -172,9 +152,9 @@ class BlJob:
         return self.__str__()
 
 
-
 def bl_job(name: str):
     return BlJob(name)
 
+
 # Create a singleton instance
 bl_start_job = BlJobWrapper()

blaxel/core/mcp/__init__.py

@@ -0,0 +1,4 @@
+from .client import websocket_client
+from .server import BlaxelMcpServerTransport
+
+__all__ = ["websocket_client", "BlaxelMcpServerTransport"]

blaxel/{mcp → core/mcp}/client.py

@@ -7,11 +7,13 @@ 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
+from mcp.shared.message import SessionMessage
+from websockets.asyncio.client import ClientConnection
+from websockets.asyncio.client import connect as ws_connect
 
 logger = logging.getLogger(__name__)
 
+
 def remove_request_params(url: str) -> str:
     return urljoin(url, urlparse(url).path)
 
@@ -27,11 +29,11 @@ async def websocket_client(
 
     The `timeout` parameter controls connection timeout.
     """
-    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+    read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
 
-    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+    write_stream: MemoryObjectSendStream[SessionMessage]
+    write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -44,11 +46,14 @@ async def websocket_client(
                 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:
+                # Use different parameters based on websockets version
+                connection_kwargs = {"additional_headers": headers, "open_timeout": timeout}
+
+                async with ws_connect(ws_url, **connection_kwargs) as websocket:
                     logger.debug("WebSocket connection established")
 
                     async def ws_reader(
-                        websocket: WebSocketClientProtocol,
+                        websocket: ClientConnection,
                         task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
                     ):
                         try:
@@ -56,9 +61,13 @@ async def websocket_client(
                             async for message in websocket:
                                 logger.debug(f"Received WebSocket message: {message}")
                                 try:
-                                    parsed_message = types.JSONRPCMessage.model_validate_json(message)
+                                    parsed_message = types.JSONRPCMessage.model_validate_json(
+                                        message
+                                    )
                                     logger.debug(f"Received server message: {parsed_message}")
-                                    await read_stream_writer.send(parsed_message)
+                                    await read_stream_writer.send(
+                                        SessionMessage(message=parsed_message)
+                                    )
                                 except Exception as exc:
                                     logger.error(f"Error parsing server message: {exc}")
                                     await read_stream_writer.send(exc)
@@ -68,13 +77,15 @@ async def websocket_client(
                         finally:
                             await read_stream_writer.aclose()
 
-                    async def ws_writer(websocket: WebSocketClientProtocol):
+                    async def ws_writer(websocket: ClientConnection):
                         try:
                             async with write_stream_reader:
-                                async for message in write_stream_reader:
-                                    logger.debug(f"Sending client message: {message}")
+                                async for session_message in write_stream_reader:
+                                    logger.debug(
+                                        f"Sending client message: {session_message.message}"
+                                    )
                                     await websocket.send(
-                                        message.model_dump_json(
+                                        session_message.message.model_dump_json(
                                             by_alias=True,
                                             exclude_none=True,
                                         )
@@ -107,4 +118,4 @@ async def websocket_client(
     finally:
         # Ensure streams are closed even if task group setup fails or GeneratorExit is caught
         await read_stream_writer.aclose()
-        await write_stream.aclose()
+        await write_stream.aclose()