File size: 4,072 Bytes
476b409 bbccf68 476b409 ebda551 476b409 ebda551 476b409 ebda551 bbccf68 476b409 89326bd 476b409 ebda551 476b409 ebda551 476b409 |
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 |
import contextlib
import io
import logging
import os
logger = logging.getLogger(__name__)
from models import GoogleModelID, OpenRouterModelID
from settings import Settings
from smolagents import LiteLLMModel, CodeAgent
from smolagents import GoogleSearchTool, VisitWebpageTool, FinalAnswerTool
from smolagents.local_python_executor import BASE_PYTHON_TOOLS
from tools import GetTaskFileTool, VideoUnderstandingTool, AudioUnderstandingTool
from tools import ChessBoardFENTool, BestChessMoveTool, ConvertChessMoveTool
# Base tools may use these to process files
BASE_PYTHON_TOOLS["open"] = open
BASE_PYTHON_TOOLS["os"] = os
BASE_PYTHON_TOOLS["io"] = io
BASE_PYTHON_TOOLS["contextlib"] = contextlib
BASE_PYTHON_TOOLS["exec"] = exec
class ResearchAgent:
def __init__(self, settings: Settings):
self.agent = CodeAgent(
name="researcher",
description="Searches the web, works with files, and answers questions for you. Give it your query as an argument.",
add_base_tools=False,
tools=[GoogleSearchTool("serper"),
VisitWebpageTool(max_output_length=100000),
VideoUnderstandingTool(settings, GoogleModelID.GEMINI_2_0_FLASH),
AudioUnderstandingTool(settings, GoogleModelID.GEMINI_2_0_FLASH)
],
additional_authorized_imports=[
"unicodedata",
"stat",
"datetime",
"random",
"pandas",
"itertools",
"math",
"statistics",
"queue",
"time",
"collections",
"re",
"os"
],
max_steps=10,
verbosity_level=1,
model=LiteLLMModel(
model_id=OpenRouterModelID.GPT_O4_MINI,
api_key = settings.openrouter_api_key.get_secret_value(),
temperature=0.0, timeout=180
)
)
class ChessAgent:
def __init__(self, settings: Settings):
self.agent = CodeAgent(
name="chess_player",
description="Makes a chess move. Give it a query including board image filepath and player turn (black or white).",
add_base_tools=False,
tools=[ChessBoardFENTool(),
BestChessMoveTool(settings),
ConvertChessMoveTool(settings, OpenRouterModelID.GPT_O4_MINI),
],
additional_authorized_imports=[
"unicodedata",
"stat",
"datetime",
"random",
"pandas",
"itertools",
"math",
"statistics",
"queue",
"time",
"collections",
"re",
"os"
],
max_steps=10,
verbosity_level=1,
model=LiteLLMModel(
model_id=OpenRouterModelID.GPT_O4_MINI,
api_key = settings.openrouter_api_key.get_secret_value(),
temperature=0.0, timeout=180
)
)
class ManagerAgent:
def __init__(self, settings: Settings):
self.researcher = ResearchAgent(settings).agent
self.chess_player = ChessAgent(settings).agent
self.agent = CodeAgent(
tools=[GetTaskFileTool(settings), FinalAnswerTool()],
model=LiteLLMModel(
model_id=OpenRouterModelID.GPT_O4_MINI,
api_key = settings.openrouter_api_key.get_secret_value(),
temperature=0.0, timeout=180
),
managed_agents=[self.researcher, self.chess_player],
)
# print("BasicAgent initialized.")
def __call__(self, question: str) -> str:
logger.info(f"Agent received question (first 50 chars): {question[:50]}...")
final_answer = self.agent.run(question)
logger.info(f"Agent returning fixed answer: {final_answer}")
return final_answer |