File size: 2,119 Bytes
7bfc491 d1b9537 7bfc491 d1b9537 7bfc491 d1b9537 7bfc491 d1b9537 7bfc491 d1b9537 |
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 |
import os
from smolagents import OpenAIServerModel, ToolCallingAgent
from tools import (
AudioTranscriber,
txt_reader,
pdf_reader,
excel_reader,
math_calculator,
search_engine,
image_transcriber,
)
from utils import OPENAI_MODEL_ID
MAIN_PROMPT = """
You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template:
[YOUR FINAL ANSWER]
YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Each question will contain the question, the task ID, and if available, a file name.
"""
class MainAgent:
"""Main agent that orchestrates file handling and web search tasks."""
def __init__(self):
model = OpenAIServerModel(model_id=OPENAI_MODEL_ID, api_key=os.getenv("OPENAI_API_KEY"))
self.agent = ToolCallingAgent(
name="main_agent",
description="An agent that can search the web, visit webpages, perform calculations, and handle files.",
tools=[
search_engine,
math_calculator,
AudioTranscriber(),
txt_reader,
pdf_reader,
excel_reader,
image_transcriber,
],
max_steps=5,
model=model,
)
def run(self, question: str, task_id: str, file_name: str | None) -> str:
"""Run the agent with the provided question, task ID, and optional file name."""
query = f"""q: {question}
task_id: {task_id}
{'file_name: ' + file_name if file_name else ''}
"""
return self.agent.run(query)
|