blaxel 0.1.13__py3-none-any.whl → 0.1.14.dev3__py3-none-any.whl

blaxel/authentication/clientcredentials.py

@@ -7,6 +7,7 @@ client credentials and refresh tokens.
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from typing import Generator, Optional
+import time
 
 import requests
 from httpx import Request, Response
@@ -46,14 +47,17 @@ class ClientCredentials(BlaxelAuth):
             "X-Blaxel-Workspace": self.workspace_name,
         }
 
-    def get_token(self) -> Optional[Exception]:
+    def _request_token(self, remaining_retries: int = 3) -> Optional[Exception]:
         """
-        Checks if the access token needs to be refreshed and performs the refresh if necessary.
+        Makes the token request with recursive retry logic.
+
+        Args:
+            remaining_retries (int): Number of retry attempts remaining.
 
         Returns:
-            Optional[Exception]: An exception if refreshing fails, otherwise None.
+            Optional[Exception]: An exception if refreshing fails after all retries, otherwise None.
         """
-        if self.need_token():
+        try:
             headers = {"Authorization": f"Basic {self.credentials.client_credentials}", "Content-Type": "application/json"}
             body = {"grant_type": "client_credentials"}
             response = requests.post(f"{self.base_url}/oauth/token", headers=headers, json=body)
@@ -63,6 +67,23 @@ class ClientCredentials(BlaxelAuth):
             self.credentials.refresh_token = creds["refresh_token"]
             self.credentials.expires_in = creds["expires_in"]
             self.expires_at = datetime.now() + timedelta(seconds=self.credentials.expires_in)
+            return None
+        except Exception as e:
+            if remaining_retries > 0:
+                time.sleep(1)
+                return self._request_token(remaining_retries - 1)
+            return e
+
+    def get_token(self) -> Optional[Exception]:
+        """
+        Checks if the access token needs to be refreshed and performs the refresh if necessary.
+        Uses recursive retry logic for up to 3 attempts.
+
+        Returns:
+            Optional[Exception]: An exception if refreshing fails after all retries, otherwise None.
+        """
+        if self.need_token():
+            return self._request_token()
         return None
 
     def need_token(self):

blaxel/common/internal.py

@@ -7,21 +7,25 @@ from typing import Optional
 
 logger = getLogger(__name__)
 
-def get_alphanumeric_limited_hash(input_str, max_size):
-    # Create SHA-256 hash of the input string
-    hash_obj = hashlib.sha256(input_str.encode('utf-8'))
-
-    # Get the hash digest in base64 format
-    hash_base64 = base64.b64encode(hash_obj.digest()).decode('utf-8')
-
-    # Remove non-alphanumeric characters and convert to lowercase
-    alphanumeric = re.sub(r'[^a-zA-Z0-9]', '', hash_base64).lower()
-
-    # Skip the first character to match the Node.js crypto output
-    alphanumeric = alphanumeric[1:]
-
-    # Limit to max_size characters
-    return alphanumeric[:max_size] if len(alphanumeric) > max_size else alphanumeric
+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()
+    
+    # Limit to max_size
+    if len(hash_hex) > max_size:
+        return hash_hex[:max_size]
+    
+    return hash_hex
 
 
 def get_global_unique_hash(workspace: str, type: str, name: str) -> str:

blaxel/common/logger.py

@@ -2,8 +2,55 @@
 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
+from opentelemetry import trace
+
+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: {}
+        }
+
+        # Add trace context if available
+        current_span = trace.get_current_span()
+        if current_span.is_recording():
+            span_context = current_span.get_span_context()
+            log_entry[self.trace_id_name] = f"{self.trace_id_prefix}{span_context.trace_id}"
+            log_entry[self.span_id_name] = f"{self.span_id_prefix}{span_context.span_id}"
+
+        # 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):
@@ -51,5 +98,10 @@ def init_logger(log_level: str):
     logging.getLogger('urllib3').setLevel(logging.CRITICAL)
     logging.getLogger("httpx").setLevel(logging.CRITICAL)
     handler = logging.StreamHandler()
-    handler.setFormatter(ColoredFormatter("%(levelname)s %(name)s - %(message)s"))
+
+    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/common/settings.py

@@ -62,9 +62,14 @@ class Settings:
     @property
     def run_internal_hostname(self) -> str:
         """Get the run internal hostname."""
-        return ''
-        # Not working in python
-        # return os.environ.get("BL_RUN_INTERNAL_HOSTNAME", "")
+        if self.generation == "":
+            return ""
+        return os.environ.get("BL_RUN_INTERNAL_HOST", "")
+    
+    @property
+    def generation(self) -> str:
+        """Get the generation."""
+        return os.environ.get("BL_GENERATION", "")
 
     @property
     def bl_cloud(self) -> bool:

blaxel/instrumentation/manager.py

@@ -7,6 +7,7 @@ import importlib
 import logging
 import signal
 import time
+import os
 from typing import Any, Dict, List, Optional, Type
 
 from opentelemetry import metrics, trace
@@ -164,14 +165,15 @@ class TelemetryManager:
         metrics.set_meter_provider(meter_provider)
         self.meter = meter_provider.get_meter(__name__)
 
-        # Set up the LoggerProvider
-        self.logger_provider = LoggerProvider(resource=resource)
-        set_logger_provider(self.logger_provider)
-        self.logger_provider.add_log_record_processor(
-            AsyncLogRecordProcessor(self.get_log_exporter())
-        )
-        handler = LoggingHandler(level=logging.NOTSET, logger_provider=self.logger_provider)
-        logging.getLogger().addHandler(handler)
+        logger_type = os.environ.get("BL_LOGGER", "http")
+        if logger_type == "http":
+            self.logger_provider = LoggerProvider(resource=resource)
+            set_logger_provider(self.logger_provider)
+            self.logger_provider.add_log_record_processor(
+                AsyncLogRecordProcessor(self.get_log_exporter())
+            )
+            handler = LoggingHandler(level=logging.NOTSET, logger_provider=self.logger_provider)
+            logging.getLogger().addHandler(handler)
 
         # Load and enable instrumentations
         for name, mapping in MAPPINGS.items():

blaxel/jobs/__init__.py

@@ -0,0 +1,187 @@
+import argparse
+import os
+import sys
+import asyncio
+import json
+import requests
+
+
+from typing import Any, Dict, Callable, Awaitable
+from logging import getLogger
+from ..client import client
+from ..common.env import env
+from ..common.internal import 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'):
+            parser = argparse.ArgumentParser()
+            # Parse known args, ignore unknown
+            args, unknown = parser.parse_known_args()
+            # Convert to dict and include unknown args
+            args_dict = vars(args)
+            # Add unknown args to dict
+            for i in range(0, len(unknown), 2):
+                if i + 1 < len(unknown):
+                    key = unknown[i].lstrip('-')
+                    args_dict[key] = unknown[i + 1]
+            return args_dict
+
+        response = requests.get(os.getenv('BL_EXECUTION_DATA_URL'))
+        data = response.json()
+        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')
+
+    @property
+    def index(self) -> int:
+        index_value = os.getenv(self.index_key)
+        return int(index_value) if index_value else 0
+
+    def start(self, func: Callable):
+        """
+        Run a job defined in a function, it's run in the current process.
+        Handles both async and sync functions.
+        Arguments are passed as keyword arguments to the function.
+        """
+        try:
+            parsed_args = self.get_arguments()
+            if asyncio.iscoroutinefunction(func):
+                asyncio.run(func(**parsed_args))
+            else:
+                func(**parsed_args)
+            sys.exit(0)
+        except Exception as error:
+            print('Job execution failed:', error, file=sys.stderr)
+            sys.exit(1)
+
+
+
+
+logger = getLogger(__name__)
+
+class BlJob:
+    def __init__(self, name: str):
+        self.name = name
+
+    @property
+    def internal_url(self):
+        """Get the internal URL for the job using a hash of workspace and job name."""
+        hash = get_global_unique_hash(settings.workspace, "job", self.name)
+        return f"{settings.run_internal_protocol}://bl-{settings.env}-{hash}.{settings.run_internal_hostname}"
+
+    @property
+    def forced_url(self):
+        """Get the forced URL from environment variables if set."""
+        env_var = self.name.replace("-", "_").upper()
+        if env[f"BL_JOB_{env_var}_URL"]:
+            return env[f"BL_JOB_{env_var}_URL"]
+        return None
+
+    @property
+    def external_url(self):
+        return f"{settings.run_url}/{settings.workspace}/jobs/{self.name}"
+
+    @property
+    def fallback_url(self):
+        if self.external_url != self.url:
+            return self.external_url
+        return None
+
+    @property
+    def url(self):
+        if self.forced_url:
+            return self.forced_url
+        if settings.run_internal_hostname:
+            return self.internal_url
+        return self.external_url
+
+    def call(self, url, input_data, headers: dict = {}, params: dict = {}):
+        body = {
+            "tasks": input_data
+        }
+
+        return client.get_httpx_client().post(
+            url+"/executions",
+            headers={
+                'Content-Type': 'application/json',
+                **headers
+            },
+            json=body,
+            params=params
+        )
+
+    async def acall(self, url, input_data, headers: dict = {}, params: dict = {}):
+        logger.debug(f"Job Calling: {self.name}")
+        body = {
+            "tasks": input_data
+        }
+
+        return await client.get_async_httpx_client().post(
+            url+"/executions",
+            headers={
+                'Content-Type': 'application/json',
+                **headers
+            },
+            json=body,
+            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)
+            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
+
+    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)
+            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
+
+    def __str__(self):
+        return f"Job {self.name}"
+
+    def __repr__(self):
+        return self.__str__()
+
+
+
+def bl_job(name: str):
+    return BlJob(name)
+
+# Create a singleton instance
+bl_start_job = BlJobWrapper()

{blaxel-0.1.13.dist-info → blaxel-0.1.14.dev3.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: blaxel
-Version: 0.1.13
+Version: 0.1.14.dev3
 Summary: Add your description here
 Author-email: cploujoux <cploujoux@blaxel.ai>
 License-File: LICENSE
 Requires-Python: >=3.10
 Requires-Dist: attrs>=21.3.0
 Requires-Dist: httpx>=0.27.0
-Requires-Dist: mcp>=1.2.1
+Requires-Dist: mcp<=1.7.1
 Requires-Dist: opentelemetry-api>=1.28.0
 Requires-Dist: opentelemetry-exporter-otlp>=1.28.0
 Requires-Dist: opentelemetry-instrumentation-anthropic==0.35.0

{blaxel-0.1.13.dist-info → blaxel-0.1.14.dev3.dist-info}/RECORD

@@ -2,7 +2,7 @@ blaxel/__init__.py,sha256=qmuJKjl5oGnjj4TbqHcJqUkKoxk4PvCsMb6-8rp67pE,159
 blaxel/agents/__init__.py,sha256=RDWkvfICIXXaQxJuuSu63jsFj_F8NBAL4U752hfN4AE,5262
 blaxel/authentication/__init__.py,sha256=tL9XKNCek5ixszTqjlKRBvidXMg4Nj6ODlBKlxxA9uk,3283
 blaxel/authentication/apikey.py,sha256=nOgLVba7EfVk3V-qm7cj-30LAL-BT7NOMIlGL9Ni1jY,1249
-blaxel/authentication/clientcredentials.py,sha256=SfWNuZVHZw6jjMqoBMMB4ZawmpyKbVPbOpv-JDFqzw8,3080
+blaxel/authentication/clientcredentials.py,sha256=7ZrqsXTQTBmCzzrWWUl9qRqq7ZPjucBWQ6dzFCB7JT4,3804
 blaxel/authentication/devicemode.py,sha256=kWbArs4okIIDqW-ql5oV2eQRE_LpRwfadCB6LG83irw,5986
 blaxel/authentication/oauth.py,sha256=Q5J0taIK1JrvGB6BC-zz3hM77HPCNu01DPGf4l7xjPQ,1417
 blaxel/authentication/types.py,sha256=E3lmfbmZuJ4Bc_kGA0Kc0GZC02Sjha1_2CbabP7z6oo,1603
@@ -262,14 +262,15 @@ blaxel/client/models/workspace_runtime.py,sha256=dxEpmwCFPOCRKHRKhY-iW7j6TbtL5qU
 blaxel/client/models/workspace_user.py,sha256=70CcifQWYbeWG7TDui4pblTzUe5sVK0AS19vNCzKE8g,3423
 blaxel/common/autoload.py,sha256=NFuK71-IHOY2JQyEBSjDCVfUaQ8D8PJsEUEryIdG4AU,263
 blaxel/common/env.py,sha256=wTbzPDdNgz4HMJiS2NCZmQlN0qpxy1PQEYBaZgtvhoc,1247
-blaxel/common/internal.py,sha256=J-etgnBzelb-ln8uBR9WIWzrEUyDds8rdT9FImjes9g,2390
-blaxel/common/logger.py,sha256=emqgonfZMBIaQPowpngWOOZxWQieKP-yj_OzGZT8Oe8,1918
-blaxel/common/settings.py,sha256=uVM6pjcovw_JPjiNZQb9WbtvLu3RdVGy7PNBcDGzoWY,2190
+blaxel/common/internal.py,sha256=PExgeKfJEmjINKreNb3r2nB5GAfG7uJhbfqHxuxBED8,2395
+blaxel/common/logger.py,sha256=7oWvrZ4fg7qUfrXe7oFAeH-1pTxadrvWQmPDn7bvbmQ,4111
+blaxel/common/settings.py,sha256=7KTryuBdud0IfHqykX7xEEtpgq5M5h1Z8YEzYKsHB-Q,2327
 blaxel/instrumentation/exporters.py,sha256=EoX3uaBVku1Rg49pSNXKFyHhgY5OV3Ih6UlqgjF5epw,1670
 blaxel/instrumentation/log.py,sha256=4tGyvLg6r4DbjqJfajYbbZ1toUzF4Q4H7kHVqYWFAEA,2537
-blaxel/instrumentation/manager.py,sha256=jVwVasyMI6tu0zv27DVYOaN57nXYuHFJ8QzJQKaNoK4,8880
+blaxel/instrumentation/manager.py,sha256=vX8RT84upjzgCUeiULp9QpDSSNVnPNFxLq0sMVz4Pjs,8974
 blaxel/instrumentation/map.py,sha256=zZoiUiQHmik5WQZ4VCWNARSa6ppMi0r7D6hlb41N-Mg,1589
 blaxel/instrumentation/span.py,sha256=X2lwfu_dyxwQTMQJT2vbXOrbVSChEhjRLc413QOxQJM,3244
+blaxel/jobs/__init__.py,sha256=4Mk-S_eT-bV-bI9Y_TOfisv_rxme9O1Z69AgmDoAKTY,6510
 blaxel/mcp/__init__.py,sha256=KednMrtuc4Y0O3lv7u1Lla54FCk8UX9c1k0USjL3Ahk,69
 blaxel/mcp/client.py,sha256=cFFXfpKXoMu8qTUly2ejF0pX2iBQkSNAxqwvDV1V6xY,4979
 blaxel/mcp/server.py,sha256=GIldtA_NgIc2dzd7ZpPvpbhpIt_7AfKu5yS_YJ0bDGg,7310
@@ -340,7 +341,7 @@ blaxel/tools/llamaindex.py,sha256=-gQ-C9V_h9a11J4ItsbWjXrCJOg0lRKsb98v9rVsNak,71
 blaxel/tools/openai.py,sha256=GuFXkj6bXEwldyVr89jEsRAi5ihZUVEVe327QuWiGNs,653
 blaxel/tools/pydantic.py,sha256=CvnNbAG_J4yBtA-XFI4lQrq3FYKjNd39hu841vZT004,1801
 blaxel/tools/types.py,sha256=YPCGJ4vZDhqR0X2H_TWtc5chQScsC32nGTQdRKJlO8Y,707
-blaxel-0.1.13.dist-info/METADATA,sha256=Tj-4s4B9NAhuurkfhmtgcPorlt3y7b-UrNO4DAMPhPE,11772
-blaxel-0.1.13.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-blaxel-0.1.13.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
-blaxel-0.1.13.dist-info/RECORD,,
+blaxel-0.1.14.dev3.dist-info/METADATA,sha256=wwb37Z03Dg-wMo0jN0MhdZUf-uA5lLlKrdbywdH8hMM,11777
+blaxel-0.1.14.dev3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+blaxel-0.1.14.dev3.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
+blaxel-0.1.14.dev3.dist-info/RECORD,,