add parasail engine

gui_agents/s2/core/engine.py

@@ -457,3 +457,37 @@ class LMMEngineHuggingFace(LMMEngine):
             .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"
+        self.model = model
+
+        api_key = api_key or os.getenv("PARASAIL_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "A Parasail API key needs to be provided in either the api_key parameter or as an environment variable named PARASAIL_API_KEY"
+            )
+        
+        self.api_key = api_key
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+
+        self.llm_client = OpenAI(base_url="https://api.parasail.io/v1", api_key=self.api_key)
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        """Generate the next message based on previous messages"""
+        return (
+            self.llm_client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                max_tokens=max_new_tokens if max_new_tokens else 4096,
+                temperature=temperature,
+                **kwargs
+            )
+            .choices[0].
+            message.content
+        )

gui_agents/s2/core/mllm.py

@@ -8,6 +8,7 @@ from gui_agents.s2.core.engine import (
     LMMEngineHuggingFace,
     LMMEngineOpenAI,
     LMMEngineOpenRouter,
+    LMMEngineParasail,
     LMMEnginevLLM,
     LMMEngineGemini,
 )
@@ -32,6 +33,8 @@ class LMMAgent:
                     self.engine = LMMEngineGemini(**engine_params)
                 elif engine_type == "open_router":
                     self.engine = LMMEngineOpenRouter(**engine_params)
+                elif engine_type == "parasail":
+                    self.engine = LMMEngineParasail(**engine_params)
                 else:
                     raise ValueError("engine_type is not supported")
             else:
@@ -125,6 +128,7 @@ class LMMAgent:
                 LMMEngineHuggingFace,
                 LMMEngineGemini,
                 LMMEngineOpenRouter,
+                LMMEngineParasail
             ),
         ):
             # infer role from previous message