Spaces:
Sleeping
Sleeping
File size: 4,717 Bytes
7f5a92e 31d8082 7f5a92e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
from smolagents import (
CodeAgent,
InferenceClientModel,
VisitWebpageTool,
WebSearchTool,
WikipediaSearchTool,
PythonInterpreterTool,
FinalAnswerTool,
tool
)
from groq import Groq
from typing import Dict, Any
import os
# ---- 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) # <<<< CORRECT VARIABLE
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 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 = self.client.chat.completions.create(
messages=messages,
model=self.model_name,
stream=False,
max_tokens=max_tokens,
)
return response.choices[0].message.content
# ---- 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)
|