feat: add LiteLLM as AI gateway engine

gui_agents/s2/core/engine.py

@@ -451,6 +451,33 @@ def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
         )
 
 
+class LMMEngineLiteLLM(LMMEngine):
+    def __init__(self, api_key=None, model=None, rate_limit=-1, **kwargs):
+        assert model is not None, "model must be provided"
+        self.model = model
+        self.api_key = api_key
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        import litellm
+
+        params = {
+            "model": self.model,
+            "messages": messages,
+            "max_tokens": max_new_tokens if max_new_tokens else 4096,
+            "temperature": temperature,
+            "drop_params": True,
+            **kwargs,
+        }
+        if self.api_key:
+            params["api_key"] = self.api_key
+        response = litellm.completion(**params)
+        return response.choices[0].message.content
+
+
 class LMMEngineParasail(LMMEngine):
     def __init__(self, api_key=None, model=None, rate_limit=-1, **kwargs):
         assert model is not None, "Parasail model id must be provided"

gui_agents/s2/core/mllm.py

@@ -5,12 +5,13 @@
 from gui_agents.s2.core.engine import (
     LMMEngineAnthropic,
     LMMEngineAzureOpenAI,
+    LMMEngineGemini,
     LMMEngineHuggingFace,
+    LMMEngineLiteLLM,
     LMMEngineOpenAI,
     LMMEngineOpenRouter,
     LMMEngineParasail,
     LMMEnginevLLM,
-    LMMEngineGemini,
 )
 
 
@@ -35,6 +36,8 @@ def __init__(self, engine_params=None, system_prompt=None, engine=None):
                     self.engine = LMMEngineOpenRouter(**engine_params)
                 elif engine_type == "parasail":
                     self.engine = LMMEngineParasail(**engine_params)
+                elif engine_type == "litellm":
+                    self.engine = LMMEngineLiteLLM(**engine_params)
                 else:
                     raise ValueError("engine_type is not supported")
             else:
@@ -127,6 +130,7 @@ def add_message(
                 LMMEngineAzureOpenAI,
                 LMMEngineHuggingFace,
                 LMMEngineGemini,
+                LMMEngineLiteLLM,
                 LMMEngineOpenRouter,
                 LMMEngineParasail,
             ),

gui_agents/s2_5/core/engine.py

@@ -398,6 +398,33 @@ def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
         )
 
 
+class LMMEngineLiteLLM(LMMEngine):
+    def __init__(self, api_key=None, model=None, rate_limit=-1, **kwargs):
+        assert model is not None, "model must be provided"
+        self.model = model
+        self.api_key = api_key
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        import litellm
+
+        params = {
+            "model": self.model,
+            "messages": messages,
+            "max_tokens": max_new_tokens if max_new_tokens else 4096,
+            "temperature": temperature,
+            "drop_params": True,
+            **kwargs,
+        }
+        if self.api_key:
+            params["api_key"] = self.api_key
+        response = litellm.completion(**params)
+        return response.choices[0].message.content
+
+
 class LMMEngineParasail(LMMEngine):
     def __init__(
         self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs

gui_agents/s2_5/core/mllm.py

@@ -5,12 +5,13 @@
 from gui_agents.s2_5.core.engine import (
     LMMEngineAnthropic,
     LMMEngineAzureOpenAI,
+    LMMEngineGemini,
     LMMEngineHuggingFace,
+    LMMEngineLiteLLM,
     LMMEngineOpenAI,
     LMMEngineOpenRouter,
     LMMEngineParasail,
     LMMEnginevLLM,
-    LMMEngineGemini,
 )
 
 
@@ -35,6 +36,8 @@ def __init__(self, engine_params=None, system_prompt=None, engine=None):
                     self.engine = LMMEngineOpenRouter(**engine_params)
                 elif engine_type == "parasail":
                     self.engine = LMMEngineParasail(**engine_params)
+                elif engine_type == "litellm":
+                    self.engine = LMMEngineLiteLLM(**engine_params)
                 else:
                     raise ValueError("engine_type is not supported")
             else:
@@ -127,6 +130,7 @@ def add_message(
                 LMMEngineAzureOpenAI,
                 LMMEngineHuggingFace,
                 LMMEngineGemini,
+                LMMEngineLiteLLM,
                 LMMEngineOpenRouter,
                 LMMEngineParasail,
             ),

gui_agents/s3/core/engine.py

@@ -402,6 +402,37 @@ def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
         )
 
 
+class LMMEngineLiteLLM(LMMEngine):
+    def __init__(
+        self, api_key=None, model=None, rate_limit=-1, temperature=None, **kwargs
+    ):
+        assert model is not None, "model must be provided"
+        self.model = model
+        self.api_key = api_key
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+        self.temperature = temperature
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        import litellm
+
+        temp = self.temperature if self.temperature is not None else temperature
+        params = {
+            "model": self.model,
+            "messages": messages,
+            "max_tokens": max_new_tokens if max_new_tokens else 4096,
+            "temperature": temp,
+            "drop_params": True,
+            **kwargs,
+        }
+        if self.api_key:
+            params["api_key"] = self.api_key
+        response = litellm.completion(**params)
+        return response.choices[0].message.content
+
+
 class LMMEngineParasail(LMMEngine):
     def __init__(
         self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs

gui_agents/s3/core/mllm.py

@@ -5,12 +5,13 @@
 from gui_agents.s3.core.engine import (
     LMMEngineAnthropic,
     LMMEngineAzureOpenAI,
+    LMMEngineGemini,
     LMMEngineHuggingFace,
+    LMMEngineLiteLLM,
     LMMEngineOpenAI,
     LMMEngineOpenRouter,
     LMMEngineParasail,
     LMMEnginevLLM,
-    LMMEngineGemini,
 )
 
 
@@ -35,6 +36,8 @@ def __init__(self, engine_params=None, system_prompt=None, engine=None):
                     self.engine = LMMEngineOpenRouter(**engine_params)
                 elif engine_type == "parasail":
                     self.engine = LMMEngineParasail(**engine_params)
+                elif engine_type == "litellm":
+                    self.engine = LMMEngineLiteLLM(**engine_params)
                 else:
                     raise ValueError(f"engine_type '{engine_type}' is not supported")
             else:
@@ -127,6 +130,7 @@ def add_message(
                 LMMEngineAzureOpenAI,
                 LMMEngineHuggingFace,
                 LMMEngineGemini,
+                LMMEngineLiteLLM,
                 LMMEngineOpenRouter,
                 LMMEngineParasail,
             ),

requirements.txt

@@ -16,6 +16,7 @@ toml
 black
 pytesseract
 google-genai
+litellm>=1.60.0,<2.0.0
 
 # Platform-specific dependencies
 pyobjc; platform_system == "Darwin"

setup.py

@@ -29,6 +29,7 @@
         "toml",
         "pytesseract",
         "google-genai",
+        "litellm>=1.60.0,<2.0.0",
         'pywinauto; platform_system == "Windows"',  # Only for Windows
         'pywin32; platform_system == "Windows"',  # Only for Windows
     ],