blaxel 0.2.0rc1__py3-none-any.whl → 0.2.0rc3__py3-none-any.whl

blaxel/core/mcp/client.py

@@ -7,6 +7,7 @@ import anyio
 import mcp.types as types
 from anyio.abc import TaskStatus
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.shared.message import SessionMessage
 from websockets.asyncio.client import ClientConnection
 from websockets.asyncio.client import connect as ws_connect
 
@@ -28,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)
@@ -64,7 +65,9 @@ async def websocket_client(
                                         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)
@@ -77,10 +80,12 @@ async def websocket_client(
                     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,
                                         )

blaxel/core/mcp/server.py

@@ -2,71 +2,23 @@ import logging
 import traceback
 import uuid
 from contextlib import asynccontextmanager
-from typing import Dict, Literal
+from typing import Literal
 
 import anyio
 import mcp.types as types
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from mcp.server.fastmcp import FastMCP as FastMCPBase
+from mcp.shared.message import SessionMessage
 from websockets.asyncio.server import ServerConnection, serve
 
-try:
-    from opentelemetry.trace import Span, StatusCode
-
-    HAS_OPENTELEMETRY = True
-except ImportError:
-    HAS_OPENTELEMETRY = False
-
-    # Create dummy classes for when opentelemetry is not available
-    class Span:
-        def set_attributes(self, *args, **kwargs):
-            pass
-
-        def set_status(self, *args, **kwargs):
-            pass
-
-        def record_exception(self, *args, **kwargs):
-            pass
-
-        def end(self, *args, **kwargs):
-            pass
-
-        def add_event(self, *args, **kwargs):
-            pass
-
-        def get_span_context(self, *args, **kwargs):
-            return None
-
-        def is_recording(self, *args, **kwargs):
-            return False
-
-        def set_attribute(self, *args, **kwargs):
-            pass
-
-        def update_name(self, *args, **kwargs):
-            pass
-
-    class StatusCode:
-        ERROR = "ERROR"
-
-
 from ..common.env import env
 
 logger = logging.getLogger(__name__)
 
 
-class DummySpanManager:
-    """Dummy span manager for when opentelemetry is not available."""
-
-    def create_span(self, name: str, attributes: dict = None):
-        return Span()
-
-
 class BlaxelMcpServerTransport:
     """WebSocket server transport for MCP."""
 
-    spans: Dict[str, Span] = {}
-
     def __init__(self, port: int = 8080):
         """Initialize the WebSocket server transport.
 
@@ -80,21 +32,14 @@ class BlaxelMcpServerTransport:
         self.clients = {}
         self.server = None
 
-        # Initialize span manager
-        if HAS_OPENTELEMETRY:
-            # TODO: Implement proper OpenTelemetry span manager when telemetry is available
-            self.span_manager = DummySpanManager()
-        else:
-            self.span_manager = DummySpanManager()
-
     @asynccontextmanager
     async def websocket_server(self):
         """Create and run a WebSocket server for MCP communication."""
-        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)
@@ -106,30 +51,14 @@ class BlaxelMcpServerTransport:
 
             try:
                 async for message in websocket:
-                    span = self.span_manager.create_span("message", {"mcp.client.id": client_id})
                     try:
                         msg = types.JSONRPCMessage.model_validate_json(message)
                         # Modify message ID to include client ID
                         if hasattr(msg, "id") and msg.id is not None:
                             original_id = msg.id
                             msg.id = f"{client_id}:{original_id}"
-                            span.set_attributes(
-                                {
-                                    "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
-                                    ),
-                                    "span.type": "mcp.message",
-                                }
-                            )
-                            self.spans[client_id + ":" + msg.id] = span
-                        await read_stream_writer.send(msg)
+                        await read_stream_writer.send(SessionMessage(message=msg))
                     except Exception as exc:
-                        span.set_status(StatusCode.ERROR)
-                        span.record_exception(exc)
-                        span.end()
                         logger.error(f"Failed to parse message: {exc}\n{traceback.format_exc()}")
                         await read_stream_writer.send(exc)
             except Exception as e:
@@ -141,11 +70,11 @@ class BlaxelMcpServerTransport:
 
         async def message_sender():
             async with write_stream_reader:
-                async for message in write_stream_reader:
+                async for session_message in write_stream_reader:
                     # Extract client ID from message ID
                     client_id = None
                     msg_id = None
-
+                    message = session_message.message
                     if hasattr(message, "id") and message.id is not None:
                         parts = str(message.id).split(":", 1)
                         if len(parts) == 2:
@@ -158,25 +87,12 @@ class BlaxelMcpServerTransport:
                     if client_id and client_id in self.clients:
                         # Send to specific client
                         websocket = self.clients[client_id]
-                        span = self.spans.get(client_id + ":" + msg_id)
                         try:
                             await websocket.send(data)
-                            if span:
-                                span.set_attributes(
-                                    {
-                                        "mcp.message.response_sent": True,
-                                    }
-                                )
                         except Exception as e:
-                            if span:
-                                span.set_status(StatusCode.ERROR)
-                                span.record_exception(e)
                             logger.error(f"Failed to send message to client {client_id}: {e}")
                             if client_id in self.clients:
                                 del self.clients[client_id]
-                        finally:
-                            if span:
-                                span.end()
                     else:
                         # Broadcast to all clients
                         dead_clients = []

{blaxel-0.2.0rc1.dist-info → blaxel-0.2.0rc3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blaxel
-Version: 0.2.0rc1
+Version: 0.2.0rc3
 Summary: Blaxel - AI development platform SDK
 Project-URL: Homepage, https://blaxel.ai
 Project-URL: Documentation, https://docs.blaxel.ai
@@ -11,7 +11,7 @@ License-File: LICENSE
 Requires-Python: >=3.10
 Requires-Dist: attrs>=21.3.0
 Requires-Dist: httpx>=0.27.0
-Requires-Dist: mcp<=1.7.1
+Requires-Dist: mcp>=1.9.4
 Requires-Dist: pydantic>=2.0.0
 Requires-Dist: pyjwt>=2.0.0
 Requires-Dist: python-dateutil>=2.8.0
@@ -43,7 +43,7 @@ Requires-Dist: llama-index-llms-groq>=0.3.1; extra == 'all'
 Requires-Dist: llama-index-llms-mistralai>=0.4.0; extra == 'all'
 Requires-Dist: llama-index-llms-openai>=0.3.42; extra == 'all'
 Requires-Dist: llama-index>=0.12.37; extra == 'all'
-Requires-Dist: openai-agents>=0.0.7; extra == 'all'
+Requires-Dist: openai-agents>=0.0.14; extra == 'all'
 Requires-Dist: opentelemetry-exporter-otlp>=1.28.0; extra == 'all'
 Requires-Dist: opentelemetry-instrumentation-anthropic==0.40.6; extra == 'all'
 Requires-Dist: opentelemetry-instrumentation-cohere==0.40.6; extra == 'all'
@@ -95,7 +95,7 @@ Requires-Dist: llama-index-llms-openai>=0.3.42; extra == 'llamaindex'
 Requires-Dist: llama-index>=0.12.37; extra == 'llamaindex'
 Requires-Dist: opentelemetry-instrumentation-llamaindex>=0.40.7; extra == 'llamaindex'
 Provides-Extra: openai
-Requires-Dist: openai-agents>=0.0.7; extra == 'openai'
+Requires-Dist: openai-agents>=0.0.14; extra == 'openai'
 Provides-Extra: pydantic
 Requires-Dist: pydantic-ai>=0.0.48; extra == 'pydantic'
 Provides-Extra: telemetry

{blaxel-0.2.0rc1.dist-info → blaxel-0.2.0rc3.dist-info}/RECORD

@@ -301,8 +301,8 @@ blaxel/core/common/logger.py,sha256=Jt0MCJgYDPq36rl7UyKRDJH76a-AwYdfggNeNYJt6N0,
 blaxel/core/common/settings.py,sha256=7KTryuBdud0IfHqykX7xEEtpgq5M5h1Z8YEzYKsHB-Q,2327
 blaxel/core/jobs/__init__.py,sha256=tApZ6t6Savd5c9yH1RhuiPDGq19B464Fakd-YqCOEtI,5441
 blaxel/core/mcp/__init__.py,sha256=5VjkiQFb1QWW5QKRgwPHARlxZJ9Xqaz0diJTpM8LLF0,142
-blaxel/core/mcp/client.py,sha256=dFjPHirybqA3IMecSvdDss-iU5ferPUaunNrPyhFm8o,5195
-blaxel/core/mcp/server.py,sha256=OU1sYcwepnvGtWrxH5IxZdcZfMqMF8hwjvyAzssM6bo,8741
+blaxel/core/mcp/client.py,sha256=aK3wSnsO8DmT1BZqw4eiCMF71Jwvni6Qga0DhPP806Y,5437
+blaxel/core/mcp/server.py,sha256=tXySGZKgK3IllYOzYOecp58BixKBkmAIvQp_4nSM_Ww,5919
 blaxel/core/models/__init__.py,sha256=HbRDsMnUFHkPC-MMkzPXh4mUqkVjqO6p3j7m00N_XSo,1722
 blaxel/core/sandbox/__init__.py,sha256=eWTH7PWNQZe9efqParwgXzltc7L4OyPtT_REqCdQybU,551
 blaxel/core/sandbox/action.py,sha256=9Zjkco7YkLzBThD3N2Hr5SpeEiqU_-Ktk8HlKpkpiAg,2802
@@ -402,7 +402,7 @@ blaxel/telemetry/instrumentation/map.py,sha256=EVCw6ug1hemE-9kh77Vxs9cqr_p9TPzFw
 blaxel/telemetry/instrumentation/utils.py,sha256=KInMYZH-mu9_wvetmf0EmgrfN3Sw8IWk2Y95v2u90_U,1901
 blaxel/telemetry/log/log.py,sha256=RvQByRjZMoP_dRaAZu8oK6DTegsHs-xV4W-UIqis6CA,2461
 blaxel/telemetry/log/logger.py,sha256=NPAS3g82ryROjvc_DEZaTIfrcehoLEZoP-JkLxADxc0,4113
-blaxel-0.2.0rc1.dist-info/METADATA,sha256=n61p-dUYa8NGsZdSxAllJkrdiq2PzI1BvlI_UNXPRAQ,9876
-blaxel-0.2.0rc1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-blaxel-0.2.0rc1.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
-blaxel-0.2.0rc1.dist-info/RECORD,,
+blaxel-0.2.0rc3.dist-info/METADATA,sha256=6uAwLfyFgd98jymj7QVZMz6et4U1oAr3JgqyL0ZcMkQ,9878
+blaxel-0.2.0rc3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+blaxel-0.2.0rc3.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
+blaxel-0.2.0rc3.dist-info/RECORD,,