Upload multi_agent.py

multi_agent.py

@@ -0,0 +1,168 @@
+from smolagents import (
+    CodeAgent,
+    InferenceClientModel,
+    VisitWebpageTool,
+    WebSearchTool,
+    WikipediaSearchTool,
+    PythonInterpreterTool,
+    FinalAnswerTool,
+    tool,
+)
+from groq import Groq
+from typing import Dict, Any
+import os
+import time
+
+
+# ---- TOOLS ----
+@tool
+def image_process(image_file: str) -> Dict[str, str]:
+    """
+    Extract text from an image file using OCR and return OCR text and base64 encoding.
+
+    Args:
+        image_file: Path to the image file
+
+    Returns:
+        Dict with keys 'ocr_text' and 'base64_image'
+    """
+    try:
+        import pytesseract
+        from PIL import Image
+        from smolagents.utils import encode_image_base64
+
+        image = Image.open(image_file)
+        base64_img = encode_image_base64(image)
+        text = pytesseract.image_to_string(image)
+        return {"ocr_text": text, "base64_image": base64_img}
+    except Exception as e:
+        return {"ocr_text": "", "base64_image": "", "error": str(e)}
+
+
+# ---- GROQ MODEL WRAPPER ----
+class LLMResponse:
+    def __init__(self, content: str, token_usage: int | None = None):
+        self.content = content
+        self.token_usage = token_usage
+
+
+class GroqModel:
+    def __init__(self, model_name=""):
+        self.model_name = model_name
+        self.client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+
+    def __call__(self, prompt, max_tokens=8096):
+        if isinstance(prompt, str):
+            messages = [{"role": "user", "content": prompt}]
+        else:
+            messages = prompt
+
+        response = None
+        for attempt in range(3):
+            try:
+                response = self.client.chat.completions.create(
+                    messages=messages,
+                    model=self.model_name,
+                    stream=False,
+                    max_tokens=max_tokens,
+                )
+                break
+            except Exception as e:
+                msg = str(e).lower()
+                if "rate limit" in msg and attempt < 2:
+                    wait = 10 * (attempt + 1)
+                    time.sleep(wait)
+                    continue
+                raise
+
+        if response is None:
+            response = self.client.chat.completions.create(
+                messages=messages,
+                model=self.model_name,
+                stream=False,
+                max_tokens=max_tokens,
+            )
+
+        content = response.choices[0].message.content
+        token_usage = None
+        if hasattr(response, "usage") and response.usage is not None:
+            token_usage = response.usage.total_tokens
+
+        return LLMResponse(content, token_usage)
+
+    def generate(self, prompt, max_tokens=8096, **kwargs):
+        # For compatibility with agent frameworks
+        return self.__call__(prompt, max_tokens=max_tokens)
+
+
+# ---- MULTI-AGENT SYSTEM ----
+class MultyAgentSystem:
+    def __init__(self):
+        deepseek_model = GroqModel("deepseek-r1-distill-llama-70b")
+        qwen_model = GroqModel("qwen-qwq-32b")
+
+        # --- Web agent definition ---
+        self.web_agent = CodeAgent(
+            model=qwen_model,
+            tools=[WebSearchTool(), VisitWebpageTool(), WikipediaSearchTool()],
+            name="web_agent",
+            description=(
+                "You are a web browsing agent. Whenever the given {task} involves browsing "
+                "the web or a specific website such as Wikipedia or YouTube, you will use "
+                "the provided tools. For web-based factual and retrieval tasks, be as precise and source-reliable as possible."
+            ),
+            additional_authorized_imports=[
+                "markdownify",
+                "json",
+                "requests",
+                "urllib.request",
+                "urllib.parse",
+                "wikipedia-api",
+            ],
+            verbosity_level=0,
+            max_steps=10,
+        )
+
+        # --- Info agent definition ---
+        self.info_agent = CodeAgent(
+            model=qwen_model,
+            tools=[PythonInterpreterTool(), image_process],
+            name="info_agent",
+            description=(
+                "You are an agent tasked with cleaning, parsing, calculating information, and performing OCR if images are provided in the {task}. "
+                "You handle all math, code, and data manipulation. Use numpy, math, and available libraries. For image or chess tasks, use pytesseract, PIL, or chess as required."
+            ),
+            additional_authorized_imports=[
+                "numpy",
+                "math",
+                "pytesseract",
+                "PIL",
+                "chess",
+            ],
+        )
+
+        # --- Manager agent definition ---
+        self.manager_agent = CodeAgent(
+            model=deepseek_model,
+            tools=[FinalAnswerTool()],
+            managed_agents=[self.web_agent, self.info_agent],
+            name="manager_agent",
+            description=(
+                "You are the manager. Given a {task}, plan which agent to use: "
+                "If web data is needed, delegate to web_agent. If math, parsing, or code is needed, use info_agent. "
+                "After collecting outputs, optionally cross-validate and check correctness, then finalize and submit the best answer using FinalAnswerTool. "
+                "For each task, explicitly explain your planning steps and reasons for choosing which agent, and always prefer the most accurate and complete answer possible."
+            ),
+            additional_authorized_imports=[
+                "json",
+                "pandas",
+                "numpy",
+            ],
+            planning_interval=5,
+            verbosity_level=2,
+            max_steps=20,
+        )
+
+    def __call__(self, question, **kwargs):
+
+        return self.manager_agent(question, **kwargs)