blaxel 0.2.31rc120__py3-none-any.whl → 0.2.32__py3-none-any.whl

blaxel/langgraph/model.py

@@ -23,70 +23,69 @@ logger = getLogger(__name__)
 
 class TokenRefreshingWrapper:
     """Base wrapper class that refreshes token before each call."""
-    
+
     def __init__(self, model_config: dict):
         self.model_config = model_config
         self.wrapped_model = self._create_model()
-    
+
     def _create_model(self):
         """Create the model instance with current token."""
         config = self.model_config
-        model_type = config['type']
-        model = config['model']
-        url = config['url']
-        kwargs = config.get('kwargs', {})
-        
-        
-        if model_type == 'mistral':
+        model_type = config["type"]
+        model = config["model"]
+        url = config["url"]
+        kwargs = config.get("kwargs", {})
+
+        if model_type == "mistral":
             return ChatOpenAI(
                 api_key=settings.auth.token,
                 model=model,
                 base_url=f"{url}/v1",
-                **kwargs
+                **kwargs,
             )
-        elif model_type == 'cohere':
+        elif model_type == "cohere":
             return ChatCohere(
                 cohere_api_key=settings.auth.token,
                 model=model,
                 base_url=url,
-                **kwargs
+                **kwargs,
             )
-        elif model_type == 'xai':
+        elif model_type == "xai":
             return ChatXAI(
                 model=model,
                 api_key=settings.auth.token,
                 xai_api_base=f"{url}/v1",
-                **kwargs
+                **kwargs,
             )
-        elif model_type == 'deepseek':
+        elif model_type == "deepseek":
             return ChatDeepSeek(
                 api_key=settings.auth.token,
                 model=model,
                 api_base=f"{url}/v1",
-                **kwargs
+                **kwargs,
             )
-        elif model_type == 'anthropic':
+        elif model_type == "anthropic":
             return ChatAnthropic(
                 api_key=settings.auth.token,
                 anthropic_api_url=url,
                 model=model,
                 default_headers=settings.auth.get_headers(),
-                **kwargs
+                **kwargs,
             )
-        elif model_type == 'gemini':
+        elif model_type == "gemini":
             return ChatGoogleGenerativeAI(
                 model=model,
                 client_options={"api_endpoint": url},
                 additional_headers=settings.auth.get_headers(),
                 transport="rest",
-                **kwargs
+                **kwargs,
             )
         elif model_type == "cerebras":
             return ChatCerebras(
                 api_key=settings.auth.token,
                 model=model,
                 base_url=f"{url}/v1",
-                **kwargs
+                **kwargs,
             )
         else:
             if model_type != "openai":
@@ -95,24 +94,24 @@ class TokenRefreshingWrapper:
                 api_key=settings.auth.token,
                 model=model,
                 base_url=f"{url}/v1",
-                **kwargs
+                **kwargs,
             )
-    
+
     def _refresh_token(self):
         """Refresh the token and recreate the model if needed."""
         # Only refresh if using ClientCredentials (which has get_token method)
         current_token = settings.auth.token
-        
-        if hasattr(settings.auth, 'get_token'):
+
+        if hasattr(settings.auth, "get_token"):
             # This will trigger token refresh if needed
             settings.auth.get_token()
-        
+
         new_token = settings.auth.token
-        
+
         # If token changed, recreate the model
         if current_token != new_token:
             self.wrapped_model = self._create_model()
-    
+
     def __getattr__(self, name):
         """Delegate attribute access to wrapped model."""
         return getattr(self.wrapped_model, name)
@@ -120,7 +119,7 @@ class TokenRefreshingWrapper:
 
 class TokenRefreshingChatModel(TokenRefreshingWrapper):
     """Wrapper for chat models that refreshes token before each call."""
-    
+
     async def ainvoke(
         self,
         input: LanguageModelInput,
@@ -132,7 +131,7 @@ class TokenRefreshingChatModel(TokenRefreshingWrapper):
         """Async invoke with token refresh."""
         self._refresh_token()
         return await self.wrapped_model.ainvoke(input, config, stop=stop, **kwargs)
-    
+
     def invoke(
         self,
         input: LanguageModelInput,
@@ -144,7 +143,7 @@ class TokenRefreshingChatModel(TokenRefreshingWrapper):
         """Sync invoke with token refresh."""
         self._refresh_token()
         return self.wrapped_model.invoke(input, config, stop=stop, **kwargs)
-    
+
     async def astream(
         self,
         input: LanguageModelInput,
@@ -157,7 +156,7 @@ class TokenRefreshingChatModel(TokenRefreshingWrapper):
         self._refresh_token()
         async for chunk in self.wrapped_model.astream(input, config, stop=stop, **kwargs):
             yield chunk
-    
+
     def stream(
         self,
         input: LanguageModelInput,
@@ -170,7 +169,7 @@ class TokenRefreshingChatModel(TokenRefreshingWrapper):
         self._refresh_token()
         for chunk in self.wrapped_model.stream(input, config, stop=stop, **kwargs):
             yield chunk
-    
+
     async def agenerate(
         self,
         messages: List[List[BaseMessage]],
@@ -185,10 +184,15 @@ class TokenRefreshingChatModel(TokenRefreshingWrapper):
         """Async generate with token refresh."""
         self._refresh_token()
         return await self.wrapped_model.agenerate(
-            messages, stop=stop, callbacks=callbacks, tags=tags,
-            metadata=metadata, run_name=run_name, **kwargs
+            messages,
+            stop=stop,
+            callbacks=callbacks,
+            tags=tags,
+            metadata=metadata,
+            run_name=run_name,
+            **kwargs,
         )
-    
+
     def generate(
         self,
         messages: List[List[BaseMessage]],
@@ -203,21 +207,26 @@ class TokenRefreshingChatModel(TokenRefreshingWrapper):
         """Sync generate with token refresh."""
         self._refresh_token()
         return self.wrapped_model.generate(
-            messages, stop=stop, callbacks=callbacks, tags=tags,
-            metadata=metadata, run_name=run_name, **kwargs
+            messages,
+            stop=stop,
+            callbacks=callbacks,
+            tags=tags,
+            metadata=metadata,
+            run_name=run_name,
+            **kwargs,
         )
-    
+
     async def astream_events(self, *args, **kwargs):
         """Async stream events with token refresh."""
         self._refresh_token()
         async for event in self.wrapped_model.astream_events(*args, **kwargs):
             yield event
-    
+
     def batch(self, *args, **kwargs):
         """Batch with token refresh."""
         self._refresh_token()
         return self.wrapped_model.batch(*args, **kwargs)
-    
+
     async def abatch(self, *args, **kwargs):
         """Async batch with token refresh."""
         self._refresh_token()
@@ -226,14 +235,9 @@ class TokenRefreshingChatModel(TokenRefreshingWrapper):
 
 async def bl_model(name: str, **kwargs):
     url, type, model = await bl_model_core(name).get_parameters()
-    
+
     # Store model configuration for recreation
-    model_config = {
-        'type': type,
-        'model': model,
-        'url': url,
-        'kwargs': kwargs
-    }
-    
+    model_config = {"type": type, "model": model, "url": url, "kwargs": kwargs}
+
     # Create and return the wrapper
-    return TokenRefreshingChatModel(model_config)
+    return TokenRefreshingChatModel(model_config)

blaxel/langgraph/tools.py

@@ -1,16 +1,22 @@
 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 blaxel.core.tools import bl_tools as bl_tools_core
 from blaxel.core.tools.types import Tool, ToolException
 
 NonTextContent = ImageContent | EmbeddedResource
 
+
 def get_langchain_tool(tool: Tool) -> StructuredTool:
     async def langchain_coroutine(
-       **arguments: dict[str, Any],
+        **arguments: dict[str, Any],
     ) -> tuple[str | list[str], list[NonTextContent] | None]:
         result: CallToolResult = await tool.coroutine(**arguments)
         text_contents: list[TextContent] = []
@@ -42,4 +48,4 @@ def get_langchain_tool(tool: Tool) -> StructuredTool:
 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()]
+    return [get_langchain_tool(tool) for tool in tools.get_tools()]

blaxel/llamaindex/custom/cohere.py

@@ -17,7 +17,10 @@ from llama_index.core.base.llms.types import (
 )
 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
@@ -56,7 +59,8 @@ class Cohere(FunctionCallingLLM):
     )
     max_retries: int = Field(default=10, description="The maximum number of API retries.")
     additional_kwargs: Dict[str, Any] = Field(
-        default_factory=dict, description="Additional kwargs for the Cohere API."
+        default_factory=dict,
+        description="Additional kwargs for the Cohere API.",
     )
     max_tokens: int = Field(description="The maximum number of tokens to generate.")
 
@@ -213,9 +217,9 @@ class Cohere(FunctionCallingLLM):
 
         messages, documents = remove_documents_from_messages(messages)
 
-        tool_results: List[Dict[str, Any]] | None = (
-            _messages_to_cohere_tool_results_curr_chat_turn(messages) or kwargs.get("tool_results")
-        )
+        tool_results: List[Dict[str, Any]] | None = _messages_to_cohere_tool_results_curr_chat_turn(
+            messages
+        ) or kwargs.get("tool_results")
         if not tool_results:
             tool_results = None
 
@@ -285,12 +289,14 @@ class Cohere(FunctionCallingLLM):
 
         if "stream" in all_kwargs:
             warnings.warn(
-                "Parameter `stream` is not supported by the `chat` method."
-                "Use the `stream_chat` method instead"
+                "Parameter `stream` is not supported by the `chat` method.Use the `stream_chat` method instead"
             )
 
         response = completion_with_retry(
-            client=self._client, max_retries=self.max_retries, chat=True, **chat_request
+            client=self._client,
+            max_retries=self.max_retries,
+            chat=True,
+            **chat_request,
         )
         if not isinstance(response, cohere.NonStreamedChatResponse):
             tool_calls = response.get("tool_calls")
@@ -316,8 +322,7 @@ class Cohere(FunctionCallingLLM):
         all_kwargs = self._get_all_kwargs(**kwargs)
         if "stream" in all_kwargs:
             warnings.warn(
-                "Parameter `stream` is not supported by the `chat` method."
-                "Use the `stream_chat` method instead"
+                "Parameter `stream` is not supported by the `chat` method.Use the `stream_chat` method instead"
             )
 
         response = completion_with_retry(
@@ -343,7 +348,10 @@ class Cohere(FunctionCallingLLM):
         chat_request = self.get_cohere_chat_request(messages=messages, **all_kwargs)
 
         response = completion_with_retry(
-            client=self._client, max_retries=self.max_retries, chat=True, **chat_request
+            client=self._client,
+            max_retries=self.max_retries,
+            chat=True,
+            **chat_request,
         )
 
         def gen() -> ChatResponseGen:
@@ -394,8 +402,7 @@ class Cohere(FunctionCallingLLM):
             raise ValueError(f"{all_kwargs['model']} not supported for chat")
         if "stream" in all_kwargs:
             warnings.warn(
-                "Parameter `stream` is not supported by the `chat` method."
-                "Use the `stream_chat` method instead"
+                "Parameter `stream` is not supported by the `chat` method.Use the `stream_chat` method instead"
             )
 
         chat_request = self.get_cohere_chat_request(messages=messages, **all_kwargs)
@@ -443,8 +450,7 @@ class Cohere(FunctionCallingLLM):
         all_kwargs = self._get_all_kwargs(**kwargs)
         if "stream" in all_kwargs:
             warnings.warn(
-                "Parameter `stream` is not supported by the `chat` method."
-                "Use the `stream_chat` method instead"
+                "Parameter `stream` is not supported by the `chat` method.Use the `stream_chat` method instead"
             )
 
         response = await acompletion_with_retry(
@@ -472,7 +478,10 @@ class Cohere(FunctionCallingLLM):
         chat_request = self.get_cohere_chat_request(messages, **all_kwargs)
 
         response = completion_with_retry(
-            client=self._client, max_retries=self.max_retries, chat=True, **chat_request
+            client=self._client,
+            max_retries=self.max_retries,
+            chat=True,
+            **chat_request,
         )
 
         async def gen() -> ChatResponseAsyncGen:

blaxel/llamaindex/model.py

@@ -34,35 +34,35 @@ logger = getLogger(__name__)
 
 class TokenRefreshingWrapper:
     """Base wrapper class that refreshes token before each call."""
-    
+
     def __init__(self, model_config: dict):
         self.model_config = model_config
         self.wrapped_model = self._create_model()
-    
+
     def _create_model(self):
         """Create the model instance with current token."""
         config = self.model_config
-        model_type = config['type']
-        model = config['model']
-        url = config['url']
-        kwargs = config.get('kwargs', {})
-        
-        if model_type == 'anthropic':
+        model_type = config["type"]
+        model = config["model"]
+        url = config["url"]
+        kwargs = config.get("kwargs", {})
+
+        if model_type == "anthropic":
             return Anthropic(
                 model=model,
                 api_key=settings.auth.token,
                 base_url=url,
                 default_headers=settings.auth.get_headers(),
-                **kwargs
+                **kwargs,
             )
-        elif model_type == 'xai':
+        elif model_type == "xai":
             return Groq(
                 model=model,
                 api_key=settings.auth.token,
                 api_base=f"{url}/v1",
-                **kwargs
+                **kwargs,
             )
-        elif model_type == 'gemini':
+        elif model_type == "gemini":
             return GoogleGenAI(
                 api_key=settings.auth.token,
                 model=model,
@@ -71,62 +71,54 @@ class TokenRefreshingWrapper:
                     base_url=url,
                     headers=settings.auth.get_headers(),
                 ),
-                **kwargs
+                **kwargs,
             )
-        elif model_type == 'cohere':
-            return Cohere(
-                model=model,
-                api_key=settings.auth.token,
-                api_base=url,
-                **kwargs
-            )
-        elif model_type == 'deepseek':
+        elif model_type == "cohere":
+            return Cohere(model=model, api_key=settings.auth.token, api_base=url, **kwargs)
+        elif model_type == "deepseek":
             return DeepSeek(
                 model=model,
                 api_key=settings.auth.token,
                 api_base=f"{url}/v1",
-                **kwargs
+                **kwargs,
             )
-        elif model_type == 'mistral':
-            return MistralAI(
-                model=model,
-                api_key=settings.auth.token,
-                endpoint=url,
-                **kwargs
-            )
-        elif model_type == 'cerebras':
+        elif model_type == "mistral":
+            return MistralAI(model=model, api_key=settings.auth.token, endpoint=url, **kwargs)
+        elif model_type == "cerebras":
             return Cerebras(
                 model=model,
                 api_key=settings.auth.token,
                 api_base=f"{url}/v1",
-                **kwargs
+                **kwargs,
             )
         else:
             if model_type != "openai":
-                logger.warning(f"Model {model} is not supported by LlamaIndex, defaulting to OpenAI")
-            
+                logger.warning(
+                    f"Model {model} is not supported by LlamaIndex, defaulting to OpenAI"
+                )
+
             return OpenAI(
                 model=model,
                 api_key=settings.auth.token,
                 api_base=f"{url}/v1",
-                **kwargs
+                **kwargs,
             )
-    
+
     def _refresh_token(self):
         """Refresh the token and recreate the model if needed."""
         # Only refresh if using ClientCredentials (which has get_token method)
         current_token = settings.auth.token
-        
-        if hasattr(settings.auth, 'get_token'):
+
+        if hasattr(settings.auth, "get_token"):
             # This will trigger token refresh if needed
             settings.auth.get_token()
-        
+
         new_token = settings.auth.token
-        
+
         # If token changed, recreate the model
         if current_token != new_token:
             self.wrapped_model = self._create_model()
-    
+
     def __getattr__(self, name):
         """Delegate attribute access to wrapped model."""
         return getattr(self.wrapped_model, name)
@@ -134,7 +126,7 @@ class TokenRefreshingWrapper:
 
 class TokenRefreshingLLM(TokenRefreshingWrapper):
     """Wrapper for LlamaIndex LLMs that refreshes token before each call."""
-    
+
     async def achat(
         self,
         messages: Sequence[ChatMessage],
@@ -143,7 +135,7 @@ class TokenRefreshingLLM(TokenRefreshingWrapper):
         """Async chat with token refresh."""
         self._refresh_token()
         return await self.wrapped_model.achat(messages, **kwargs)
-    
+
     def chat(
         self,
         messages: Sequence[ChatMessage],
@@ -152,7 +144,7 @@ class TokenRefreshingLLM(TokenRefreshingWrapper):
         """Sync chat with token refresh."""
         self._refresh_token()
         return self.wrapped_model.chat(messages, **kwargs)
-    
+
     async def astream_chat(
         self,
         messages: Sequence[ChatMessage],
@@ -162,7 +154,7 @@ class TokenRefreshingLLM(TokenRefreshingWrapper):
         self._refresh_token()
         async for chunk in self.wrapped_model.astream_chat(messages, **kwargs):
             yield chunk
-    
+
     def stream_chat(
         self,
         messages: Sequence[ChatMessage],
@@ -172,7 +164,7 @@ class TokenRefreshingLLM(TokenRefreshingWrapper):
         self._refresh_token()
         for chunk in self.wrapped_model.stream_chat(messages, **kwargs):
             yield chunk
-    
+
     async def acomplete(
         self,
         prompt: str,
@@ -181,7 +173,7 @@ class TokenRefreshingLLM(TokenRefreshingWrapper):
         """Async complete with token refresh."""
         self._refresh_token()
         return await self.wrapped_model.acomplete(prompt, **kwargs)
-    
+
     def complete(
         self,
         prompt: str,
@@ -190,7 +182,7 @@ class TokenRefreshingLLM(TokenRefreshingWrapper):
         """Sync complete with token refresh."""
         self._refresh_token()
         return self.wrapped_model.complete(prompt, **kwargs)
-    
+
     async def astream_complete(
         self,
         prompt: str,
@@ -200,7 +192,7 @@ class TokenRefreshingLLM(TokenRefreshingWrapper):
         self._refresh_token()
         async for chunk in self.wrapped_model.astream_complete(prompt, **kwargs):
             yield chunk
-    
+
     def stream_complete(
         self,
         prompt: str,
@@ -214,14 +206,9 @@ class TokenRefreshingLLM(TokenRefreshingWrapper):
 
 async def bl_model(name, **kwargs):
     url, type, model = await bl_model_core(name).get_parameters()
-    
+
     # Store model configuration for recreation
-    model_config = {
-        'type': type,
-        'model': model,
-        'url': url,
-        'kwargs': kwargs
-    }
-    
+    model_config = {"type": type, "model": model, "url": url, "kwargs": kwargs}
+
     # Create and return the wrapper
-    return TokenRefreshingLLM(model_config)
+    return TokenRefreshingLLM(model_config)

blaxel/llamaindex/tools.py

@@ -1,5 +1,3 @@
-
-
 from llama_index.core.tools import FunctionTool
 from llama_index.core.tools.types import ToolMetadata
 
@@ -22,7 +20,8 @@ def get_llamaindex_tool(tool: Tool) -> FunctionTool:
         ),
     )
 
+
 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()]
+    return [get_llamaindex_tool(tool) for tool in tools.get_tools()]

blaxel/pydantic/custom/gemini.py

@@ -8,11 +8,11 @@ class GoogleGLAProvider(Provider[AsyncClient]):
 
     @property
     def name(self):
-        return 'google-gla'
+        return "google-gla"
 
     @property
     def base_url(self) -> str:
-        return 'https://generativelanguage.googleapis.com/v1beta/models/'
+        return "https://generativelanguage.googleapis.com/v1beta/models/"
 
     @property
     def client(self) -> httpx.AsyncClient:
@@ -28,4 +28,4 @@ class GoogleGLAProvider(Provider[AsyncClient]):
         """
         self._client = http_client
         # https://cloud.google.com/docs/authentication/api-keys-use#using-with-rest
-        self._client.headers['X-Goog-Api-Key'] = api_key
+        self._client.headers["X-Goog-Api-Key"] = api_key

blaxel/pydantic/tools.py

@@ -23,9 +23,7 @@ def get_pydantic_tool(tool: Tool) -> PydanticTool:
 
     # Define the prepare function dynamically to capture the 'tool' object
     # Assuming RunContext and ToolDefinition types based on the pydantic_ai example
-    async def prepare_tool(
-        ctx: RunContext, tool_def: ToolDefinition
-    ) -> ToolDefinition | None:
+    async def prepare_tool(ctx: RunContext, tool_def: ToolDefinition) -> ToolDefinition | None:
         """Dynamically prepares the ToolDefinition using the custom Tool's attributes."""
         tool_def.name = tool.name  # Override inferred name
         tool_def.description = tool.description  # Override inferred description
@@ -48,4 +46,4 @@ def get_pydantic_tool(tool: Tool) -> PydanticTool:
 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()]
+    return [get_pydantic_tool(tool) for tool in tools.get_tools()]