Update multiagents.py

multiagents.py

@@ -42,22 +42,54 @@ dotenv.load_dotenv()
 #     api_key=os.environ["OPENAI_API_KEY"],
 # )
 # --- Agent wrappers ---
+# Fixed GrokApiAgent class
 class GrokApiAgent:
     def __init__(self):
         pass  # No setup needed for API-based agent
 
-    def __call__(self, prompt):
+    def __call__(self, prompt, **kwargs):
         # You could keep this for backward compatibility if needed
-        return self.generate(prompt)
-
-    def generate(self, prompt):
-        response = GrokApi(
-            system_prompt="You are a helpful assistant.",
-            user_input=prompt
-        )
-        # Assuming GrokApi returns a string or has a property with the answer text
-        # Adjust this depending on the actual return type of GrokApi
-        return str(response)
+        return self.generate(prompt, **kwargs)
+
+    def generate(self, prompt, stop_sequences=None, max_tokens=None, temperature=None, **kwargs):
+        """
+        Generate method compatible with smolagents framework.
+        
+        Args:
+            prompt: The input prompt
+            stop_sequences: List of strings where generation should stop (optional)
+            max_tokens: Maximum number of tokens to generate (optional)
+            temperature: Sampling temperature (optional)
+            **kwargs: Other parameters that might be passed by the framework
+        """
+        try:
+            # Call your GrokApi with the prompt
+            # Note: You might need to adjust this based on how GrokApi actually works
+            # and whether it supports the additional parameters
+            response = GrokApi(
+                system_prompt="You are a helpful assistant.",
+                user_input=prompt
+                # Add other parameters here if GrokApi supports them:
+                # max_tokens=max_tokens,
+                # temperature=temperature,
+                # etc.
+            )
+            
+            # Convert response to string and handle stop_sequences if needed
+            response_text = str(response)
+            
+            # If stop_sequences are provided, truncate the response at the first occurrence
+            if stop_sequences:
+                for stop_seq in stop_sequences:
+                    if stop_seq in response_text:
+                        response_text = response_text.split(stop_seq)[0]
+                        break
+            
+            return response_text
+            
+        except Exception as e:
+            print(f"Error in GrokApiAgent.generate: {e}")
+            return f"Error generating response: {e}"
 
 # Pick your agent here (just change one line to test different ones)
 # My_Agent = BasicAgent