ABVM commited on
Commit
b6a56ef
·
verified ·
1 Parent(s): 41b14b9

Upload 3 files

Browse files
Files changed (3) hide show
  1. multi_agent.py +187 -0
  2. requirements.txt +15 -0
  3. vision_tool.py +70 -0
multi_agent.py ADDED
@@ -0,0 +1,187 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from smolagents import (
2
+ CodeAgent,
3
+ VisitWebpageTool,
4
+ WebSearchTool,
5
+ WikipediaSearchTool,
6
+ PythonInterpreterTool,
7
+ FinalAnswerTool,
8
+ )
9
+ from groq import Groq
10
+ from vision_tool import image_reasoning_tool
11
+ import os
12
+ import time
13
+
14
+
15
+ # ---- TOOLS ----
16
+
17
+
18
+ # ---- GROQ MODEL WRAPPER ----
19
+ class GroqModel:
20
+ def __init__(self, model_name=""):
21
+ self.model_name = model_name
22
+ self.client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
23
+
24
+ def __call__(self, prompt, max_tokens=8096):
25
+ if isinstance(prompt, str):
26
+ messages = [{"role": "user", "content": prompt}]
27
+ else:
28
+ messages = prompt
29
+
30
+ response = None
31
+ for attempt in range(3):
32
+ try:
33
+ response = self.client.chat.completions.create(
34
+ messages=messages,
35
+ model=self.model_name,
36
+ stream=False,
37
+ max_tokens=max_tokens,
38
+ )
39
+ break
40
+ except Exception as e:
41
+ msg = str(e).lower()
42
+ if "rate limit" in msg and attempt < 2:
43
+ wait = 10 * (attempt + 1)
44
+ time.sleep(wait)
45
+ continue
46
+ raise
47
+
48
+ if response is None:
49
+ response = self.client.chat.completions.create(
50
+ messages=messages,
51
+ model=self.model_name,
52
+ stream=False,
53
+ max_tokens=max_tokens,
54
+ )
55
+
56
+ content = response.choices[0].message.content
57
+ # token usage is calculated but currently unused
58
+ if hasattr(response, "usage") and response.usage is not None:
59
+ _ = response.usage.total_tokens
60
+
61
+ return content
62
+
63
+ def generate(self, prompt, max_tokens=8096, **kwargs):
64
+ # For compatibility with agent frameworks
65
+ return self.__call__(prompt, max_tokens=max_tokens)
66
+
67
+
68
+ # ---- MULTI-AGENT SYSTEM ----
69
+ class MultyAgentSystem:
70
+ def __init__(self):
71
+ self.primary_model_name = "deepseek-r1-distill-llama-70b"
72
+ self.fallback_model_name = "llama3-70b-8k"
73
+
74
+ self.deepseek_model = GroqModel(self.primary_model_name)
75
+ qwen_model = GroqModel("qwen-qwq-32b")
76
+ self.verification_limit = int(os.getenv("VERIFY_WORD_LIMIT", "75"))
77
+
78
+ # --- Web agent definition ---
79
+ self.web_agent = CodeAgent(
80
+ model=qwen_model,
81
+ tools=[WebSearchTool(), VisitWebpageTool(), WikipediaSearchTool()],
82
+ name="web_agent",
83
+ description=(
84
+ "You are a web browsing agent. Whenever the given {task} involves browsing "
85
+ "the web or a specific website such as Wikipedia or YouTube, you will use "
86
+ "the provided tools. For web-based factual and retrieval tasks, be as precise and source-reliable as possible."
87
+ ),
88
+ additional_authorized_imports=[
89
+ "markdownify",
90
+ "json",
91
+ "requests",
92
+ "urllib.request",
93
+ "urllib.parse",
94
+ "wikipedia-api",
95
+ ],
96
+ verbosity_level=0,
97
+ max_steps=10,
98
+ )
99
+
100
+ # --- Info agent definition ---
101
+ self.info_agent = CodeAgent(
102
+ model=qwen_model,
103
+ tools=[PythonInterpreterTool(), image_reasoning_tool],
104
+ name="info_agent",
105
+ description=(
106
+ "You are an agent tasked with cleaning, parsing, calculating information, and performing OCR if images are provided in the {task}. "
107
+ "You can also analyze images using a vision model. You handle all math, code, and data manipulation. Use numpy, math, and available libraries. "
108
+ "For image or chess tasks, use pytesseract, PIL, chess, or the image_reasoning_tool as required."
109
+ ),
110
+ additional_authorized_imports=[
111
+ "numpy",
112
+ "math",
113
+ "pytesseract",
114
+ "PIL",
115
+ "chess",
116
+ ],
117
+ )
118
+
119
+ # --- Manager agent definition ---
120
+ manager_planning_interval = int(os.getenv("MANAGER_PLANNING_INTERVAL", "3"))
121
+ manager_max_steps = int(os.getenv("MANAGER_MAX_STEPS", "8"))
122
+
123
+ self.manager_agent = CodeAgent(
124
+ model=qwen_model,
125
+ tools=[FinalAnswerTool()],
126
+ managed_agents=[self.web_agent, self.info_agent],
127
+ name="manager_agent",
128
+ description=(
129
+ "You are the manager. Given a {task}, plan which agent to use: "
130
+ "If web data is needed, delegate to web_agent. If math, parsing, image reasoning, or code is needed, use info_agent. "
131
+ "After collecting outputs, optionally cross-validate and check correctness, then finalize and submit the best answer using FinalAnswerTool. "
132
+ "For each task, explicitly explain your planning steps and reasons for choosing which agent, and always prefer the most accurate and complete answer possible."
133
+ ),
134
+ additional_authorized_imports=[
135
+ "json",
136
+ "pandas",
137
+ "numpy",
138
+ ],
139
+ planning_interval=manager_planning_interval,
140
+ verbosity_level=2,
141
+ max_steps=manager_max_steps,
142
+ )
143
+
144
+ # runtime tracking for fallback switching
145
+ self.total_runtime = 0.0
146
+ self.first_call_duration = None
147
+ self.model_switched = False
148
+
149
+ def _switch_to_fallback(self):
150
+ if self.model_switched:
151
+ return
152
+ self.manager_agent.model = GroqModel(self.fallback_model_name)
153
+ self.model_switched = True
154
+
155
+ def run(self, question, high_stakes: bool = False, **kwargs):
156
+ start_time = time.time()
157
+ print("Generating initial answer with Qwen-32B")
158
+ initial_answer = self.manager_agent(question, **kwargs)
159
+ call_duration = time.time() - start_time
160
+
161
+ answer = initial_answer
162
+ if high_stakes or len(initial_answer.split()) > self.verification_limit:
163
+ print("Verifying answer using DeepSeek-70B")
164
+ verification_prompt = (
165
+ "Review the following answer for accuracy and rewrite if needed:"
166
+ f"\n\n{initial_answer}"
167
+ )
168
+ try:
169
+ answer = self.deepseek_model(verification_prompt)
170
+ except Exception as e:
171
+ print(f"Verification failed: {e}. Using initial answer.")
172
+ answer = initial_answer
173
+
174
+ if self.first_call_duration is None:
175
+ self.first_call_duration = call_duration
176
+ if self.first_call_duration > 30:
177
+ self._switch_to_fallback()
178
+
179
+ self.total_runtime += call_duration
180
+ if self.total_runtime > 300 and not self.model_switched:
181
+ self._switch_to_fallback()
182
+
183
+ return answer
184
+
185
+ def __call__(self, question, high_stakes: bool = False, **kwargs):
186
+
187
+ return self.run(question, high_stakes=high_stakes, **kwargs)
requirements.txt ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ gradio
2
+ requests
3
+ wikipedia-api
4
+ smolagents
5
+ huggingface_hub
6
+ pandas==2.3.0
7
+ numpy==2.2.6
8
+ plotly==6.1.2
9
+ kaleido
10
+ pillow==11.2.1
11
+ groq
12
+ pytesseract
13
+ pytesseract-ocr
14
+ python-chess
15
+ markdownify
vision_tool.py ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Vision tool using Groq's Meta-Llama Scout model
2
+ from smolagents import tool
3
+ from groq import Groq
4
+
5
+ import os
6
+
7
+
8
+ def _llama_analyze(image_b64: str, prompt: str) -> str:
9
+ """Internal helper to query the Llama vision model."""
10
+ client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
11
+ messages = [
12
+ {
13
+ "role": "user",
14
+ "content": [
15
+ {"type": "text", "text": prompt},
16
+ {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}},
17
+ ],
18
+ }
19
+ ]
20
+ response = client.chat.completions.create(
21
+ model="meta-llama/llama-4-scout-17b-16e-instruct",
22
+ messages=messages,
23
+ stream=False,
24
+ max_tokens=512,
25
+ )
26
+ return response.choices[0].message.content
27
+
28
+
29
+ @tool
30
+ def image_reasoning_tool(image_file: str, prompt: str | None = None) -> dict:
31
+ """Perform OCR and optional vision analysis on an image.
32
+
33
+ This single entry point unifies OCR extraction and Llama vision reasoning so
34
+ the planner only sees one image tool.
35
+
36
+ Args:
37
+ image_file: Path to the image file to analyze.
38
+ prompt: Optional instruction for the vision model. If omitted, only OCR
39
+ is performed.
40
+
41
+ Returns:
42
+ Dictionary with OCR text, base64 image data and optional vision model
43
+ response.
44
+ """
45
+ try:
46
+ from PIL import Image
47
+ from smolagents.utils import encode_image_base64
48
+ import pytesseract
49
+
50
+ image = Image.open(image_file)
51
+ b64 = encode_image_base64(image)
52
+ ocr_text = pytesseract.image_to_string(image)
53
+
54
+ vision_text = ""
55
+ if prompt:
56
+ try:
57
+ vision_text = _llama_analyze(b64, prompt)
58
+ except Exception as e: # vision errors shouldn't break OCR result
59
+ vision_text = f"Error processing image with vision model: {e}"
60
+
61
+ return {"ocr_text": ocr_text, "vision_text": vision_text, "base64_image": b64}
62
+ except Exception as e:
63
+ return {
64
+ "ocr_text": "",
65
+ "vision_text": "",
66
+ "base64_image": "",
67
+ "error": str(e),
68
+ }
69
+
70
+