blaxel 0.1.9rc35__py3-none-any.whl → 0.1.9rc37__py3-none-any.whl

blaxel/common/settings.py

@@ -21,7 +21,7 @@ class Settings:
     @property
     def log_level(self) -> str:
         """Get the log level."""
-        return os.environ.get("LOG_LEVEL", "INFO")
+        return os.environ.get("LOG_LEVEL", "INFO").upper()
 
     @property
     def base_url(self) -> str:
@@ -69,6 +69,11 @@ class Settings:
         """Is running on bl cloud."""
         return os.environ.get("BL_CLOUD", "") == "true"
 
+    @property
+    def run_internal_protocol(self) -> str:
+        """Get the run internal protocol."""
+        return os.environ.get("BL_RUN_INTERNAL_PROTOCOL", "https")
+
     @property
     def enable_opentelemetry(self) -> bool:
         """Get the enable opentelemetry."""

blaxel/instrumentation/exporters.py

@@ -1,11 +1,8 @@
 from typing import Callable, Dict, Sequence
 
-from opentelemetry.exporter.otlp.proto.http._log_exporter import \
-    OTLPLogExporter
-from opentelemetry.exporter.otlp.proto.http.metric_exporter import \
-    OTLPMetricExporter
-from opentelemetry.exporter.otlp.proto.http.trace_exporter import \
-    OTLPSpanExporter
+from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk._logs import LogData
 from opentelemetry.sdk.metrics.export import MetricExportResult, MetricsData
 

blaxel/instrumentation/manager.py

@@ -20,9 +20,11 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.trace import NoOpTracerProvider
 
-from blaxel.instrumentation.exporters import (DynamicHeadersLogExporter,
-                                              DynamicHeadersMetricExporter,
-                                              DynamicHeadersSpanExporter)
+from blaxel.instrumentation.exporters import (
+    DynamicHeadersLogExporter,
+    DynamicHeadersMetricExporter,
+    DynamicHeadersSpanExporter,
+)
 from blaxel.instrumentation.span import DefaultAttributesSpanProcessor
 
 from ..common.settings import Settings

blaxel/mcp/client.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from contextlib import asynccontextmanager
 from typing import Any
 from urllib.parse import urljoin, urlparse
@@ -7,8 +6,7 @@ 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 anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from websockets.client import WebSocketClientProtocol
 from websockets.client import connect as ws_connect
 

blaxel/mcp/server.py

@@ -1,19 +1,18 @@
 import logging
+import traceback
 import uuid
 from contextlib import asynccontextmanager
 from typing import Dict, Literal
 
 import anyio
 import mcp.types as types
-from anyio.streams.memory import (MemoryObjectReceiveStream,
-                                  MemoryObjectSendStream)
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from mcp.server.fastmcp import FastMCP as FastMCPBase
 from opentelemetry.trace import Span, StatusCode
 from websockets.server import WebSocketServerProtocol, serve
 
 from ..common.env import env
 from ..instrumentation.span import SpanManager
-import traceback
 
 logger = logging.getLogger(__name__)
 
@@ -66,7 +65,8 @@ class BlaxelMcpServerTransport:
                                 "mcp.message.parsed": True,
                                 "mcp.method": getattr(msg, "method", None),
                                 "mcp.messageId": getattr(msg, "id", None),
-                                "mcp.toolName": getattr(getattr(msg, "params", None), "name", None)
+                                "mcp.toolName": getattr(getattr(msg, "params", None), "name", None),
+                                "span.type": "mcp.message",
                             })
                             self.spans[client_id+":"+msg.id] =  span
                         await read_stream_writer.send(msg)

blaxel/models/__init__.py

@@ -4,6 +4,7 @@ from ..client import client
 from ..client.api.models import get_model
 from ..client.models import Model
 from ..common.settings import settings
+
 # This has to be here because livekit plugins must be registered on the main thread
 from .livekit import get_livekit_model
 
@@ -94,7 +95,7 @@ class BLModel:
 
         try:
             return await get_model.asyncio(client=client, model_name=self.model_name)
-        except Exception as e:
+        except Exception:
             return None
 
 def bl_model(model_name, **kwargs):

blaxel/models/custom/langchain/gemini.py

@@ -6,33 +6,56 @@ import logging
 import uuid
 import warnings
 from operator import itemgetter
-from typing import (Any, AsyncIterator, Callable, Dict, Iterator, List,
-                    Mapping, Optional, Sequence, Tuple, Type, Union, cast)
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
 
 import httpx
 import requests
-from langchain_core.callbacks.manager import (AsyncCallbackManagerForLLMRun,
-                                              CallbackManagerForLLMRun)
+from langchain_core.callbacks.manager import AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun
 from langchain_core.language_models import LanguageModelInput
-from langchain_core.language_models.chat_models import (BaseChatModel,
-                                                        LangSmithParams)
-from langchain_core.messages import (AIMessage, AIMessageChunk, BaseMessage,
-                                     FunctionMessage, HumanMessage,
-                                     SystemMessage, ToolMessage)
+from langchain_core.language_models.chat_models import BaseChatModel, LangSmithParams
+from langchain_core.messages import (
+    AIMessage,
+    AIMessageChunk,
+    BaseMessage,
+    FunctionMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolMessage,
+)
 from langchain_core.messages.ai import UsageMetadata
-from langchain_core.messages.tool import (invalid_tool_call, tool_call,
-                                          tool_call_chunk)
+from langchain_core.messages.tool import invalid_tool_call, tool_call, tool_call_chunk
 from langchain_core.output_parsers.openai_tools import (
-    JsonOutputKeyToolsParser, PydanticToolsParser, parse_tool_calls)
-from langchain_core.outputs import (ChatGeneration, ChatGenerationChunk,
-                                    ChatResult)
+    JsonOutputKeyToolsParser,
+    PydanticToolsParser,
+    parse_tool_calls,
+)
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from langchain_core.runnables import Runnable, RunnablePassthrough
 from langchain_core.tools import BaseTool
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from PIL import Image
 from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
-from tenacity import (before_sleep_log, retry, retry_if_exception_type,
-                      stop_after_attempt, wait_exponential)
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
 from typing_extensions import Self, is_typeddict
 
 WARNED_STRUCTURED_OUTPUT_JSON_MODE = False
@@ -189,7 +212,7 @@ def _chat_with_retry(generation_method: Callable, **kwargs: Any) -> Any:
         try:
             # Extract request parameters and other kwargs
             request = kwargs.pop('request', {})
-            metadata = kwargs.pop('metadata', None)
+            kwargs.pop('metadata', None)
 
             # Unpack request parameters into kwargs
             kwargs.update(request)
@@ -238,7 +261,7 @@ async def _achat_with_retry(generation_method: Callable, **kwargs: Any) -> Any:
         try:
             # Extract request parameters and other kwargs
             request = kwargs.pop('request', {})
-            metadata = kwargs.pop('metadata', None)
+            kwargs.pop('metadata', None)
 
             # Unpack request parameters into kwargs
             kwargs.update(request)

blaxel/models/custom/llamaindex/cohere.py

@@ -4,27 +4,36 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Union
 
 import cohere
 from cohere.types import ToolCall
-from llama_index.core.base.llms.types import (ChatMessage, ChatResponse,
-                                              ChatResponseAsyncGen,
-                                              ChatResponseGen,
-                                              CompletionResponse,
-                                              CompletionResponseAsyncGen,
-                                              CompletionResponseGen,
-                                              LLMMetadata, MessageRole)
+from llama_index.core.base.llms.types import (
+    ChatMessage,
+    ChatResponse,
+    ChatResponseAsyncGen,
+    ChatResponseGen,
+    CompletionResponse,
+    CompletionResponseAsyncGen,
+    CompletionResponseGen,
+    LLMMetadata,
+    MessageRole,
+)
 from llama_index.core.bridge.pydantic import Field, PrivateAttr
 from llama_index.core.callbacks import CallbackManager
-from llama_index.core.llms.callbacks import (llm_chat_callback,
-                                             llm_completion_callback)
+from llama_index.core.llms.callbacks import llm_chat_callback, llm_completion_callback
 from llama_index.core.llms.function_calling import FunctionCallingLLM
 from llama_index.core.llms.llm import ToolSelection
 from llama_index.core.tools.types import BaseTool
 from llama_index.core.types import BaseOutputParser, PydanticProgramMode
 from llama_index.llms.cohere.utils import (
-    CHAT_MODELS, _get_message_cohere_format, _message_to_cohere_tool_results,
-    _messages_to_cohere_tool_results_curr_chat_turn, acompletion_with_retry,
-    cohere_modelname_to_contextsize, completion_with_retry,
-    format_to_cohere_tools, is_cohere_function_calling_model,
-    remove_documents_from_messages)
+    CHAT_MODELS,
+    _get_message_cohere_format,
+    _message_to_cohere_tool_results,
+    _messages_to_cohere_tool_results_curr_chat_turn,
+    acompletion_with_retry,
+    cohere_modelname_to_contextsize,
+    completion_with_retry,
+    format_to_cohere_tools,
+    is_cohere_function_calling_model,
+    remove_documents_from_messages,
+)
 
 
 class Cohere(FunctionCallingLLM):
@@ -200,9 +209,9 @@ class Cohere(FunctionCallingLLM):
             The request for the Cohere chat API.
         """
         additional_kwargs = messages[-1].additional_kwargs
-
+        documents = additional_kwargs.pop("documents", [])
         # cohere SDK will fail loudly if both connectors and documents are provided
-        if additional_kwargs.get("documents", []) and documents and len(documents) > 0:
+        if documents and len(documents) > 0:
             raise ValueError(
                 "Received documents both as a keyword argument and as an prompt additional keyword argument. Please choose only one option."
             )

blaxel/models/custom/pydantic/gemini.py

@@ -1,4 +1,3 @@
-import os
 
 import httpx
 from httpx import AsyncClient

blaxel/models/livekit.py

@@ -1,6 +1,6 @@
 from logging import getLogger
 
-from livekit.plugins import anthropic, groq, openai
+from livekit.plugins import groq, openai
 
 from ..common.settings import settings
 

blaxel/tools/__init__.py

@@ -4,12 +4,14 @@ import os
 from contextlib import AsyncExitStack
 from logging import getLogger
 from typing import Any, cast
+import traceback
 
 from mcp import ClientSession
 from mcp.types import CallToolResult
 from mcp.types import Tool as MCPTool
 
 from ..common.env import env
+from ..common.internal import get_global_unique_hash
 from ..common.settings import settings
 from ..instrumentation.span import SpanManager
 from ..mcp.client import websocket_client
@@ -22,15 +24,19 @@ if os.getenv("BL_SERVER_PORT"):
     DEFAULT_TIMEOUT = 5
 
 class PersistentWebSocket:
-    def __init__(self, url: str, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True):
+    def __init__(self, url: str, name: str, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True):
         self.url = url
+        self.name = name
         self.timeout = timeout
-        self.timeout_enabled = timeout_enabled
         self.session_exit_stack = AsyncExitStack()
         self.client_exit_stack = AsyncExitStack()
         self.session: ClientSession = None
         self.timer_task = None
         self.tools_cache = []
+        if settings.bl_cloud:
+            self.timeout_enabled = False
+        else:
+            self.timeout_enabled = timeout_enabled
 
     def with_metas(self, metas: dict[str, Any]):
         self.metas = metas
@@ -46,19 +52,32 @@ class PersistentWebSocket:
         logger.debug(f"Tool {tool_name} returned {call_tool_result}")
         if self.timeout_enabled:
             self._reset_timer()
+        else:
+            await self._close()
         return call_tool_result
 
     async def list_tools(self):
-        await self._initialize()
-        if self.timeout_enabled:
-            self._remove_timer()
-        logger.debug("Listing tools")
-        list_tools_result = await self.session.list_tools()
-        self.tools_cache = list_tools_result.tools
-        logger.debug(f"Tools listed: {list_tools_result}")
-        if self.timeout_enabled:
-            self._reset_timer()
-        return list_tools_result
+        logger.debug(f"Listing tools for {self.name}")
+        span_attributes = {
+            "tool.server": self.url,
+            "tool.server_name": self.name,
+            "span.type": "tool.list",
+        }
+        with SpanManager("blaxel-tracer").create_active_span(self.name, span_attributes) as span:
+            await self._initialize()
+            logger.debug(f"Initialized websocket for {self.name}")
+            if self.timeout_enabled:
+                self._remove_timer()
+            logger.debug("Listing tools")
+            list_tools_result = await self.session.list_tools()
+            self.tools_cache = list_tools_result.tools
+            logger.debug(f"Tools listed: {list_tools_result}")
+            if self.timeout_enabled:
+                self._reset_timer()
+            else:
+                await self._close()
+            span.set_attribute("tool.list.result", list_tools_result.model_dump_json())
+            return list_tools_result
 
     def get_tools(self):
         return self.tools_cache
@@ -125,8 +144,9 @@ def convert_mcp_tool_to_blaxel_tool(
             "tool.args": json.dumps(arguments),
             "tool.server": url,
             "tool.server_name": name,
+            "span.type": "tool.call",
         }
-        with SpanManager("blaxel-tracer").create_active_span("blaxel-tool-call", span_attributes) as span:
+        with SpanManager("blaxel-tracer").create_active_span("blaxel-tool-call", span_attributes):
             logger.debug(f"Calling tool {tool.name} with arguments {arguments}")
             call_tool_result = await websocket_client.call_tool(tool.name, arguments)
             logger.debug(f"Tool {tool.name} returned {call_tool_result}")
@@ -163,22 +183,37 @@ class BlTools:
         self.timeout = timeout
         self.timeout_enabled = timeout_enabled
 
-    def _external_url(self, name: str) -> str:
-        return f"{settings.run_url}/{settings.auth.workspace_name}/functions/{name}"
+    def _internal_url(self, name: str):
+        """Get the internal URL for the agent using a hash of workspace and agent name."""
+        hash = get_global_unique_hash(settings.workspace, "function", name)
+        return f"{settings.run_internal_protocol}://bl-{settings.env}-{hash}.{settings.run_internal_hostname}"
 
-    def _url(self, name: str) -> str:
+    def _forced_url(self, name: str):
+        """Get the forced URL from environment variables if set."""
         env_var = name.replace("-", "_").upper()
         if env[f"BL_FUNCTION_{env_var}_URL"]:
             return env[f"BL_FUNCTION_{env_var}_URL"]
-        elif env[f"BL_FUNCTION_{env_var}_SERVICE_NAME"]:
-            return f"https://{env[f'BL_FUNCTION_{env_var}_SERVICE_NAME']}.{settings.run_internal_hostname}"
-        return self._external_url(name)
+        return None
+
+    def _external_url(self, name: str):
+        return f"{settings.run_url}/{settings.workspace}/functions/{name}"
 
-    def _fallback_url(self, name: str) -> str | None:
+    def _fallback_url(self, name: str):
         if self._external_url(name) != self._url(name):
             return self._external_url(name)
         return None
 
+    def _url(self, name: str):
+        logger.debug(f"Getting URL for {name}")
+        if self._forced_url(name):
+            logger.debug(f"Forced URL found for {name}: {self._forced_url(name)}")
+            return self._forced_url(name)
+        if settings.run_internal_hostname:
+            logger.debug(f"Internal hostname found for {name}: {self._internal_url(name)}")
+            return self._internal_url(name)
+        logger.debug(f"No URL found for {name}, using external URL")
+        return self._external_url(name)
+
     def get_tools(self) -> list[Tool]:
         """Get a list of all tools from all connected servers."""
         all_tools: list[Tool] = []
@@ -240,8 +275,13 @@ class BlTools:
         except Exception as e:
             if not self._fallback_url(name):
                 raise e
+            logger.warning(f"Error connecting to {name}: {e}\n{traceback.format_exc()}")
             url = self._fallback_url(name)
-            await self.connect_with_url(name, url)
+            try:
+                await self.connect_with_url(name, url)
+            except Exception as e:
+                logger.error(f"Error connecting to {name} with fallback URL: {e}\n{traceback.format_exc()}")
+                raise e
 
     async def connect_with_url(
         self, name: str, url: str
@@ -255,7 +295,8 @@ class BlTools:
         logger.debug(f"Initializing session and loading tools from {url}")
 
         if not toolPersistances.get(name):
-            toolPersistances[name] = PersistentWebSocket(url, timeout=self.timeout, timeout_enabled=self.timeout_enabled)
+            logger.debug(f"Creating new persistent websocket for {name}")
+            toolPersistances[name] = PersistentWebSocket(url, name, timeout=self.timeout, timeout_enabled=self.timeout_enabled)
             await toolPersistances[name].list_tools()
         logger.debug(f"Loaded {len(toolPersistances[name].get_tools())} tools from {url}")
         return toolPersistances[name].with_metas(self.metas)

blaxel/tools/langchain.py

@@ -1,8 +1,7 @@
 from typing import Any
 
 from langchain_core.tools import StructuredTool
-from mcp.types import (CallToolResult, EmbeddedResource, ImageContent,
-                       TextContent)
+from mcp.types import CallToolResult, EmbeddedResource, ImageContent, TextContent
 
 from .types import Tool, ToolException
 

{blaxel-0.1.9rc35.dist-info → blaxel-0.1.9rc37.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blaxel
-Version: 0.1.9rc35
+Version: 0.1.9rc37
 Summary: Add your description here
 Author-email: cploujoux <cploujoux@blaxel.ai>
 License-File: LICENSE
@@ -25,9 +25,6 @@ Requires-Dist: pyyaml<6.1.0,>=6.0.2
 Requires-Dist: requests<2.33.0,>=2.32.3
 Requires-Dist: tomli>=2.2.1
 Requires-Dist: websockets<15.0.0
-Provides-Extra: crewai
-Requires-Dist: crewai>=0.102.0; extra == 'crewai'
-Requires-Dist: opentelemetry-instrumentation-crewai>=0.39.0; extra == 'crewai'
 Provides-Extra: google-adk
 Requires-Dist: google-adk>=0.2.0; extra == 'google-adk'
 Requires-Dist: litellm>=1.63.11; extra == 'google-adk'