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

blaxel/langgraph/model.py

@@ -23,69 +23,70 @@ 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":
@@ -94,24 +95,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)
@@ -119,7 +120,7 @@ class TokenRefreshingWrapper:
 
 class TokenRefreshingChatModel(TokenRefreshingWrapper):
     """Wrapper for chat models that refreshes token before each call."""
-
+    
     async def ainvoke(
         self,
         input: LanguageModelInput,
@@ -131,7 +132,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,
@@ -143,7 +144,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,
@@ -156,7 +157,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,
@@ -169,7 +170,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]],
@@ -184,15 +185,10 @@ 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]],
@@ -207,26 +203,21 @@ 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()
@@ -235,9 +226,14 @@ 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,22 +1,16 @@
 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] = []
@@ -48,4 +42,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,10 +17,7 @@ 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
@@ -59,8 +56,7 @@ 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.")
 
@@ -217,9 +213,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
 
@@ -289,14 +285,12 @@ 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")
@@ -322,7 +316,8 @@ 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(
@@ -348,10 +343,7 @@ 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:
@@ -402,7 +394,8 @@ 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)
@@ -450,7 +443,8 @@ 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(
@@ -478,10 +472,7 @@ 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,54 +71,62 @@ 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)
@@ -126,7 +134,7 @@ class TokenRefreshingWrapper:
 
 class TokenRefreshingLLM(TokenRefreshingWrapper):
     """Wrapper for LlamaIndex LLMs that refreshes token before each call."""
-
+    
     async def achat(
         self,
         messages: Sequence[ChatMessage],
@@ -135,7 +143,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],
@@ -144,7 +152,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],
@@ -154,7 +162,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],
@@ -164,7 +172,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,
@@ -173,7 +181,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,
@@ -182,7 +190,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,
@@ -192,7 +200,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,
@@ -206,9 +214,14 @@ 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,3 +1,5 @@
+
+
 from llama_index.core.tools import FunctionTool
 from llama_index.core.tools.types import ToolMetadata
 
@@ -20,8 +22,7 @@ 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,7 +23,9 @@ 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
@@ -46,4 +48,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()]