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

blaxel/googleadk/tools.py

@@ -0,0 +1,72 @@
+import inspect
+from typing import Any, Optional, override
+
+from google.adk.tools import BaseTool, ToolContext
+from google.genai import types
+
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.types import Tool
+
+
+class GoogleADKTool(BaseTool):
+    _tool: Tool
+
+    def __init__(self, tool: Tool):
+        super().__init__(
+            name=tool.name,
+            description=tool.description,
+        )
+        self._tool = tool
+
+    def _clean_schema(self, schema: dict) -> dict:
+        if not isinstance(schema, dict):
+            return schema
+
+        # Create a copy of the schema
+        cleaned_schema = schema.copy()
+
+        # Remove $schema and additionalProperties at current level
+        if "$schema" in cleaned_schema:
+            del cleaned_schema["$schema"]
+        if "additionalProperties" in cleaned_schema:
+            del cleaned_schema["additionalProperties"]
+
+        # Recursively clean properties if they exist
+        if "properties" in cleaned_schema:
+            cleaned_schema["properties"] = {
+                k: self._clean_schema(v) for k, v in cleaned_schema["properties"].items()
+            }
+
+        return cleaned_schema
+
+    @override
+    def _get_declaration(self) -> Optional[types.FunctionDeclaration]:
+        # Clean the schema recursively
+        schema = self._clean_schema(self._tool.input_schema)
+
+        function_decl = types.FunctionDeclaration.model_validate(
+            types.FunctionDeclaration(
+                name=self._tool.name,
+                description=self._tool.description,
+                parameters=schema,
+            )
+        )
+
+        return function_decl
+
+    @override
+    async def run_async(
+        self, *, args: dict[str, Any], tool_context: ToolContext
+    ) -> Any:
+        args_to_call = args.copy()
+        signature = inspect.signature(self._tool.coroutine)
+        if 'tool_context' in signature.parameters:
+            args_to_call['tool_context'] = tool_context
+        return await self._tool.coroutine(**args_to_call) or {}
+
+
+async def bl_tools(tools_names: list[str], **kwargs: Any) -> list[BaseTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [GoogleADKTool(tool) for tool in tools.get_tools()]
+    return [GoogleADKTool(tool) for tool in tools.get_tools()]

blaxel/langgraph/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel LangGraph integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/langchain.py → langgraph/model.py}

@@ -7,12 +7,16 @@ from langchain_deepseek import ChatDeepSeek
 from langchain_openai import ChatOpenAI
 from langchain_xai import ChatXAI
 
-from ..common.settings import settings
-from .custom.langchain.gemini import ChatGoogleGenerativeAI
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
+
+from .custom.gemini import ChatGoogleGenerativeAI
 
 logger = getLogger(__name__)
 
-async def get_langchain_model(url: str,type: str, model: str, **kwargs):
+
+async def bl_model(name: str, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
     if type == 'mistral':
         return ChatOpenAI(
             api_key=settings.auth.token,
@@ -72,4 +76,4 @@ async def get_langchain_model(url: str,type: str, model: str, **kwargs):
             model=model,
             base_url=f"{url}/v1",
             **kwargs
-        )
+        )

blaxel/{tools/langchain.py → langgraph/tools.py}

@@ -3,7 +3,8 @@ from typing import Any
 from langchain_core.tools import StructuredTool
 from mcp.types import CallToolResult, EmbeddedResource, ImageContent, TextContent
 
-from .types import Tool, ToolException
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.types import Tool, ToolException
 
 NonTextContent = ImageContent | EmbeddedResource
 
@@ -37,5 +38,8 @@ def get_langchain_tool(tool: Tool) -> StructuredTool:
         sync_coroutine=tool.sync_coroutine,
     )
 
-def get_langchain_tools(tools: list[Tool]) -> list[StructuredTool]:
-    return [get_langchain_tool(tool) for tool in tools]
+
+async def bl_tools(tools_names: list[str], **kwargs) -> list[StructuredTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [get_langchain_tool(tool) for tool in tools.get_tools()]

blaxel/livekit/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel LiveKit integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/livekit.py → livekit/model.py}

@@ -2,7 +2,8 @@ from logging import getLogger
 
 from livekit.plugins import groq, openai
 
-from ..common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
 
 logger = getLogger(__name__)
 
@@ -23,3 +24,8 @@ async def get_livekit_model(url: str, type: str, model: str, **kwargs):
             base_url=f"{url}/v1",
             **kwargs
         )
+
+async def bl_model(name: str, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
+    return await get_livekit_model(url, type, model, **kwargs)
+    return await get_livekit_model(url, type, model, **kwargs)

blaxel/{tools/livekit.py → livekit/tools.py}

@@ -1,6 +1,7 @@
 from livekit.agents import function_tool, llm
 
-from .types import Tool
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.types import Tool
 
 
 def livekit_coroutine(tool: Tool):
@@ -21,3 +22,9 @@ def get_livekit_tools(tools: list[Tool]) -> list[llm.FunctionTool]:
             },
         ))
     return livekit_tools
+
+async def bl_tools(tools_names: list[str], **kwargs) -> list[llm.FunctionTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return get_livekit_tools(tools.get_tools())
+    return get_livekit_tools(tools.get_tools())

blaxel/llamaindex/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel LlamaIndex integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/llamaindex.py → llamaindex/model.py}

@@ -14,12 +14,15 @@ from llama_index.llms.groq import Groq
 from llama_index.llms.mistralai import MistralAI
 from llama_index.llms.openai import OpenAI
 
-from ..common.settings import settings
-from .custom.llamaindex.cohere import Cohere
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
+
+from .custom.cohere import Cohere
 
 logger = getLogger(__name__)
 
-async def get_llamaindex_model(url: str, type: str, model: str, **kwargs):
+async def bl_model(name, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
     if type == 'anthropic':
         return Anthropic(
             model=model,

blaxel/{tools/llamaindex.py → llamaindex/tools.py}

@@ -3,8 +3,9 @@
 from llama_index.core.tools import FunctionTool
 from llama_index.core.tools.types import ToolMetadata
 
-from .common import create_model_from_json_schema
-from .types import Tool
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.common import create_model_from_json_schema
+from blaxel.core.tools.types import Tool
 
 
 def get_llamaindex_tool(tool: Tool) -> FunctionTool:
@@ -21,5 +22,7 @@ def get_llamaindex_tool(tool: Tool) -> FunctionTool:
         ),
     )
 
-def get_llamaindex_tools(tools: list[Tool]) -> list[FunctionTool]:
-    return [get_llamaindex_tool(tool) for tool in tools]
+async def bl_tools(tools_names: list[str], **kwargs) -> list[FunctionTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [get_llamaindex_tool(tool) for tool in tools.get_tools()]

blaxel/openai/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel OpenAI integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/openai.py → openai/model.py}

@@ -1,9 +1,11 @@
 from agents import AsyncOpenAI, OpenAIChatCompletionsModel
 
-from ..common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
 
 
-async def get_openai_model(url: str, type: str, model: str, **kwargs):
+async def bl_model(name, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
     if type != "openai":
         raise ValueError(f"Invalid model type: {type}")
     external_client = AsyncOpenAI(

blaxel/{tools/openai.py → openai/tools.py}

@@ -3,7 +3,8 @@ from typing import Any
 
 from agents import FunctionTool, RunContextWrapper
 
-from .types import Tool
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.types import Tool
 
 
 def get_openai_tool(tool: Tool) -> FunctionTool:
@@ -21,5 +22,8 @@ def get_openai_tool(tool: Tool) -> FunctionTool:
         on_invoke_tool=openai_coroutine,
     )
 
-def get_openai_tools(tools: list[Tool]) -> list[FunctionTool]:
-    return [get_openai_tool(tool) for tool in tools]
+async def bl_tools(tools_names: list[str], **kwargs) -> list[FunctionTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [get_openai_tool(tool) for tool in tools.get_tools()]
+    return [get_openai_tool(tool) for tool in tools.get_tools()]

blaxel/pydantic/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel Pydantic AI integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/custom/pydantic → pydantic/custom}/gemini.py

@@ -1,4 +1,3 @@
-
 import httpx
 from httpx import AsyncClient
 from pydantic_ai.providers import Provider

blaxel/{models/pydantic.py → pydantic/model.py}

@@ -14,13 +14,15 @@ from pydantic_ai.providers.cohere import CohereProvider
 from pydantic_ai.providers.mistral import MistralProvider
 from pydantic_ai.providers.openai import OpenAIProvider
 
-from blaxel.common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
+from blaxel.core.client import client
 
-from ..client import client
-from .custom.pydantic.gemini import GoogleGLAProvider
+from .custom.gemini import GoogleGLAProvider
 
 logger = logging.getLogger(__name__)
-async def get_pydantic_model(url: str,type: str, model: str, **kwargs) -> Model:
+async def bl_model(name: str, **kwargs) -> Model:
+    url, type, model = await bl_model_core(name).get_parameters()
     if type == 'mistral':
         return MistralModel(
             model_name=model,

blaxel/{tools/pydantic.py → pydantic/tools.py}

@@ -2,7 +2,8 @@ from pydantic_ai import RunContext
 from pydantic_ai.tools import Tool as PydanticTool
 from pydantic_ai.tools import ToolDefinition
 
-from .types import Tool
+from blaxel.core.tools import Tool
+from blaxel.core.tools import bl_tools as bl_tools_core
 
 
 def get_pydantic_tool(tool: Tool) -> PydanticTool:
@@ -44,5 +45,7 @@ def get_pydantic_tool(tool: Tool) -> PydanticTool:
     )
 
 
-def get_pydantic_tools(tools: list[Tool]) -> list[PydanticTool]:
-    return [get_pydantic_tool(tool) for tool in tools]
+async def bl_tools(tools_names: list[str], **kwargs) -> list[PydanticTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [get_pydantic_tool(tool) for tool in tools.get_tools()]

blaxel/telemetry/__init__.py

@@ -0,0 +1,6 @@
+"""Blaxel telemetry module."""
+
+from .exporters import *  # noqa: F403, F401
+from .instrumentation import *  # noqa: F403, F401
+from .log import *  # noqa: F403, F401
+from .manager import *  # noqa: F403, F401

blaxel/telemetry/instrumentation/blaxel_core.py

@@ -0,0 +1,124 @@
+import json
+import logging
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import Tracer, get_tracer
+
+import blaxel.core.agents
+import blaxel.core.jobs
+from blaxel.telemetry.span import SpanManager
+
+logger = logging.getLogger(__name__)
+
+class BlaxelCoreInstrumentor(BaseInstrumentor):
+    def instrumentation_dependencies(self):
+        return ["blaxel.core"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "0.2.0", tracer_provider)
+        self._patch_agent(tracer)
+        self._patch_job(tracer)
+        self._patch_tools(tracer)
+
+    def _uninstrument(self, **kwargs):
+        # Optionally restore original methods
+        pass
+
+    def _patch_agent(self, tracer: Tracer):
+        orig_run = blaxel.core.agents.BlAgent.run
+
+        def traced_run(self, *args, **kwargs):
+            attributes = {
+                "agent.name": self.name,
+                "agent.args": str(args),
+                "span.type": "agent.run",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span(self.name, attributes=attributes) as span:
+                try:
+                    result = orig_run(self, *args, **kwargs)
+                    span.set_attribute("agent.run.result", result)
+                    return result
+                except Exception as e:
+                    span.set_attribute("agent.run.error", str(e))
+                    raise
+
+        blaxel.core.agents.BlAgent.run = traced_run
+
+    def _patch_job(self, tracer: Tracer):
+        orig_run = blaxel.core.jobs.BlJob.run
+
+        def traced_run(self, *args, **kwargs):
+            attributes = {
+                "job.name": self.name,
+                "span.type": "job.run",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span(self.name, attributes=attributes) as span:
+                try:
+                    result = orig_run(self, *args, **kwargs)
+                    span.set_attribute("job.run.result", result)
+                    return result
+                except Exception as e:
+                    span.set_attribute("job.run.error", str(e))
+                    raise
+
+        blaxel.core.jobs.BlJob.run = traced_run
+
+    def _patch_tools(self, tracer: Tracer):
+        # Patch PersistentWebSocket.list_tools
+        orig_list_tools = blaxel.core.tools.PersistentWebSocket.list_tools
+
+        async def traced_list_tools(self, *args, **kwargs):
+            span_attributes = {
+                "tool.server": self.url,
+                "tool.server_name": self.name,
+                "span.type": "tool.list",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span(self.name, attributes=span_attributes):
+                result = await orig_list_tools(self, *args, **kwargs)
+                # Optionally: span.set_attribute("tool.list.result", str(result))
+                return result
+
+        blaxel.core.tools.PersistentWebSocket.list_tools = traced_list_tools
+
+        # Patch convert_mcp_tool_to_blaxel_tool to wrap tool calls
+        orig_convert = blaxel.core.tools.convert_mcp_tool_to_blaxel_tool
+
+        def traced_convert_mcp_tool_to_blaxel_tool(websocket_client, name, url, tool):
+            Tool = orig_convert(websocket_client, name, url, tool)
+            orig_coroutine = Tool.coroutine
+            orig_sync_coroutine = Tool.sync_coroutine
+
+            async def traced_coroutine(*args, **kwargs):
+                span_attributes = {
+                    "tool.name": tool.name,
+                    "tool.args": json.dumps(kwargs),
+                    "tool.server": url,
+                    "tool.server_name": name,
+                    "span.type": "tool.call",
+                    **SpanManager.get_default_attributes(),
+                }
+                logger.info(span_attributes)
+                with tracer.start_span("blaxel-tool-call", attributes=span_attributes):
+                    return await orig_coroutine(*args, **kwargs)
+
+            def traced_sync_coroutine(*args, **kwargs):
+                span_attributes = {
+                    "tool.name": tool.name,
+                    "tool.args": json.dumps(kwargs),
+                    "tool.server": url,
+                    "tool.server_name": name,
+                    "span.type": "tool.call",
+                    **SpanManager.get_default_attributes(),
+                }
+                with tracer.start_span("blaxel-tool-call", attributes=span_attributes):
+                    return orig_sync_coroutine(*args, **kwargs)
+
+            Tool.coroutine = traced_coroutine
+            Tool.sync_coroutine = traced_sync_coroutine
+            return Tool
+
+        blaxel.core.tools.convert_mcp_tool_to_blaxel_tool = traced_convert_mcp_tool_to_blaxel_tool

blaxel/telemetry/instrumentation/blaxel_langgraph.py

@@ -0,0 +1,74 @@
+import logging
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import Tracer, get_tracer
+from wrapt import wrap_function_wrapper
+
+from blaxel.telemetry.span import SpanManager
+
+logger = logging.getLogger(__name__)
+
+class BlaxelLanggraphInstrumentor(BaseInstrumentor):
+    def instrumentation_dependencies(self):
+        return ["blaxel_langgraph", "blaxel.core"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "0.2.0", tracer_provider)
+        # self._patch_gemini(tracer)
+        self._patch_tool(tracer)
+
+    def _uninstrument(self, **kwargs):
+        import blaxel.langgraph.custom.gemini
+
+        # Restore original methods
+        if hasattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_generate_content"):
+            blaxel.langgraph.custom.gemini.GeminiRestClient.generate_content = blaxel.langgraph.custom.gemini.GeminiRestClient._blaxel_original_generate_content
+            delattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_generate_content")
+
+        if hasattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_generate_content_async"):
+            blaxel.langgraph.custom.gemini.GeminiRestClient.generate_content_async = blaxel.langgraph.custom.gemini.GeminiRestClient._blaxel_original_generate_content_async
+            delattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_generate_content_async")
+
+        if hasattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_stream_generate_content"):
+            blaxel.langgraph.custom.gemini.GeminiRestClient.stream_generate_content = blaxel.langgraph.custom.gemini.GeminiRestClient._blaxel_original_stream_generate_content
+            delattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_stream_generate_content")
+
+        if hasattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_stream_generate_content_async"):
+            blaxel.langgraph.custom.gemini.GeminiRestClient.stream_generate_content_async = blaxel.langgraph.custom.gemini.GeminiRestClient._blaxel_original_stream_generate_content_async
+            delattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_stream_generate_content_async")
+
+    def _patch_tool(self, tracer: Tracer):
+        def _traced_run(original, instance, args, kwargs):
+            attributes = {
+                "tool.type": "langgraph",
+                "span.type": "tool.call",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span("langgraph.tool", attributes=attributes) as span:
+                try:
+                    result = original(*args, **kwargs)
+                    span.set_attribute("tool.name", result.name)
+                    return result
+                except Exception as e:
+                    span.set_attribute("tool.error", str(e))
+                    raise
+
+        async def _traced_arun(original, instance, args, kwargs):
+            attributes = {
+                "tool.type": "langgraph",
+                "span.type": "tool.call",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span("langgraph.tool", attributes=attributes) as span:
+                try:
+                    result = await original(*args, **kwargs)
+                    span.set_attribute("tool.name", result.name)
+                    return result
+                except Exception as e:
+                    span.set_attribute("tool.error", str(e))
+                    raise
+
+        wrap_function_wrapper('langchain_core.tools', 'BaseTool.run', _traced_run)
+        wrap_function_wrapper('langchain_core.tools', 'BaseTool.arun', _traced_arun)
+        wrap_function_wrapper('langchain_core.tools', 'BaseTool.arun', _traced_arun)