litellm Integration

os_computer_use/config.py

@@ -5,15 +5,19 @@
 grounding_model = providers.OSAtlasProvider()
 # grounding_model = providers.ShowUIProvider()
 
-# vision_model = providers.FireworksProvider("llama3.2")
-# vision_model = providers.OpenAIProvider("gpt-4o")
-# vision_model = providers.AnthropicProvider("claude-3.5-sonnet")
-vision_model = providers.GroqProvider("llama3.2")
-# vision_model = providers.MistralProvider("pixtral")  # pixtral-large-latest has vision capabilities
 
+# Vision models using LiteLLM:
+vision_model = providers.LiteLLMProvider("pixtral")  # Mistral
+# vision_model = providers.LiteLLMProvider("llama3.2", provider="fireworks")  # Fireworks
+# vision_model = providers.LiteLLMProvider("gpt-4-vision")  # OpenAI
+# vision_model = providers.LiteLLMProvider("llama3.2", provider="groq")  # Groq
+# vision_model = providers.LiteLLMProvider("claude-3-5-sonnet")  # Anthropic
+# vision_model = providers.LiteLLMProvider("gemini-2.0-flash", provider="gemini")  # Gemini
 
-# action_model = providers.FireworksProvider("llama3.3")
-# action_model = providers.OpenAIProvider("gpt-4o")
-# action_model = providers.AnthropicProvider("claude-3.5-sonnet")
-action_model = providers.GroqProvider("llama3.3")
-# action_model = providers.MistralProvider("large")  # mistral-large-latest for non-vision tasks
+# Action models using LiteLLM:
+action_model = providers.LiteLLMProvider("large")  # Mistral
+# action_model = providers.LiteLLMProvider("llama3.3", provider="fireworks")  # Fireworks
+# action_model = providers.LiteLLMProvider("llama3.3", provider="groq")  # Groq
+# action_model = providers.LiteLLMProvider("gpt-4")  # OpenAI
+# action_model = providers.LiteLLMProvider("claude-3-5-sonnet")  # Anthropic
+# action_model = providers.LiteLLMProvider("gemini-2.0-flash", provider="gemini")  # Gemini

os_computer_use/llm_provider.py

@@ -4,6 +4,7 @@
 import json
 import re
 import base64
+import imghdr
 
 
 def Message(content, role="assistant"):
@@ -22,6 +23,29 @@ def parse_json(s):
         return None
 
 
+def extract_json_objects(s):
+    """Extract all balanced JSON objects from a string."""
+    objects = []
+    brace_level = 0
+    start_index = None
+    for i, char in enumerate(s):
+        if char == "{":
+            if brace_level == 0:
+                start_index = i
+            brace_level += 1
+        elif char == "}":
+            brace_level -= 1
+            if brace_level == 0 and start_index is not None:
+                candidate = s[start_index : i + 1]
+                try:
+                    obj = json.loads(candidate)
+                    objects.append(obj)
+                except json.JSONDecodeError:
+                    pass
+                start_index = None
+    return objects
+
+
 class LLMProvider:
     """
     The LLM provider is used to make calls to an LLM given a provider and model name, with optional tool use support
@@ -52,6 +76,13 @@ def create_function_schema(self, definitions):
                 properties[param_name] = {"type": "string", "description": param_desc}
                 required.append(param_name)
 
+            # Add a dummy property if no parameters are provided, because providers like Gemini require a non-empty "properties" object.
+            if not properties:
+                properties["noop"] = {
+                    "type": "string",
+                    "description": "Dummy parameter for function with no parameters.",
+                }
+
             function_def = self.create_function_def(name, details, properties, required)
             functions.append(function_def)
 
@@ -68,8 +99,7 @@ def create_tool_call(self, name, parameters):
     # Wrap a content block in a text or an image object
     def wrap_block(self, block):
         if isinstance(block, bytes):
-            encoded_image = base64.b64encode(block).decode("utf-8")
-            return self.create_image_block(encoded_image)
+            return self.create_image_block(block)
         else:
             return Text(block)
 
@@ -117,10 +147,17 @@ def create_function_def(self, name, details, properties, required):
             },
         }
 
-    def create_image_block(self, base64_image):
+    def create_image_block(self, image_data):
+        # Detect the image type using imghdr.
+        image_type = imghdr.what(None, image_data)
+        if image_type is None:
+            image_type = "png"  # fallback if type cannot be detected
+
+        # Base64-encode the raw image bytes.
+        encoded = base64.b64encode(image_data).decode("utf-8")
         return {
             "type": "image_url",
-            "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
+            "image_url": {"url": f"data:image/{image_type};base64,{encoded}"},
         }
 
     def call(self, messages, functions=None):
@@ -140,18 +177,17 @@ def call(self, messages, functions=None):
                 if parse_json(tool_call.function.arguments) is not None
             ]
 
-            # Sometimes, function calls are returned unparsed by the inference provider. This code parses them manually.
+            # Sometimes, function calls are returned unparsed by the inference provider.
             if message.content and not tool_calls:
-                tool_call_matches = re.search(r"\{.*\}", message.content)
-                if tool_call_matches:
-                    tool_call = parse_json(tool_call_matches.group(0))
-                    # Some models use "arguments" as the key instead of "parameters"
-                    parameters = tool_call.get("parameters", tool_call.get("arguments"))
-                    if tool_call.get("name") and parameters:
+                json_objs = extract_json_objects(message.content)
+                for obj in json_objs:
+                    parameters = obj.get("parameters", obj.get("arguments"))
+                    if obj.get("name") and parameters is not None:
                         combined_tool_calls.append(
-                            self.create_tool_call(tool_call.get("name"), parameters)
+                            self.create_tool_call(obj.get("name"), parameters)
                         )
-                        return None, combined_tool_calls
+                if combined_tool_calls:
+                    return None, combined_tool_calls
 
             return message.content, combined_tool_calls
 
@@ -160,75 +196,52 @@ def call(self, messages, functions=None):
             return message.content
 
 
-class AnthropicBaseProvider(LLMProvider):
+class LiteLLMBaseProvider(OpenAIBaseProvider):
+    """Base provider using LiteLLM"""
 
     def create_client(self):
-        return Anthropic(api_key=self.api_key).messages
+        from litellm import completion
 
-    def create_function_def(self, name, details, properties, required):
-        return {
-            "name": name,
-            "description": details["description"],
-            "input_schema": {
-                "type": "object",
-                "properties": properties,
-                "required": required,
-            },
-        }
+        import litellm
 
-    def create_image_block(self, base64_image):
-        return {
-            "type": "image",
-            "source": {
-                "type": "base64",
-                "media_type": "image/png",
-                "data": base64_image,
-            },
-        }
-
-    def call(self, messages, functions=None):
-        tools = self.create_function_schema(functions) if functions else None
-
-        # Move all messages with the system role to a system parameter
-        system = "\n".join(
-            msg.get("content") for msg in messages if msg.get("role") == "system"
-        )
-        messages = [msg for msg in messages if msg.get("role") != "system"]
-
-        # Call the Anthropic API
-        completion = self.completion(
-            messages, system=system, tools=tools, max_tokens=4096
-        )
-        text = "".join(getattr(block, "text", "") for block in completion.content)
+        # Enable dropping unsupported params globally
+        litellm.drop_params = True
+        litellm.modify_params = True
+        # Enable debug mode for better error messages
+        # litellm._turn_on_debug()
+        return completion
 
-        # Return response text and tool calls separately
-        if functions:
-            tool_calls = [
-                self.create_tool_call(block.name, block.input)
-                for block in completion.content
-                if block.type == "tool_use"
-            ]
-            return text, tool_calls
+    def completion(self, messages, **kwargs):
+        # Skip the tools parameter if it's None
+        filtered_kwargs = {k: v for k, v in kwargs.items() if v is not None}
 
-        # Only return response text
-        else:
-            return text
+        # No need to remove tools; pass tools so that function calling works with Claude.
 
+        # Wrap content blocks in image or text objects if necessary
+        new_messages = [self.transform_message(message) for message in messages]
 
-class MistralBaseProvider(OpenAIBaseProvider):
-    def create_function_def(self, name, details, properties, required):
-        # If description is wrapped in a dict, extract the inner string
-        if isinstance(details.get("description"), dict):
-            details["description"] = details["description"].get("description", "")
-        return super().create_function_def(name, details, properties, required)
+        # Call LiteLLM completion
+        completion_response = self.client(
+            model=self.model,
+            messages=new_messages,
+            api_key=self.api_key,
+            **filtered_kwargs,
+        )
+        return completion_response
 
+    # Added method to adjust the final message role for Mistral-based models only
     def call(self, messages, functions=None):
-        if messages and messages[-1].get("role") == "assistant":
+        if (
+            "mistral" in self.model.lower()
+            and messages
+            and messages[-1].get("role") == "assistant"
+        ):
             prefix = messages.pop()["content"]
             if messages and messages[-1].get("role") == "user":
                 messages[-1]["content"] = (
                     prefix + "\n" + messages[-1].get("content", "")
                 )
             else:
                 messages.append({"role": "user", "content": prefix})
+
         return super().call(messages, functions)