zaldivards's picture
Update agent prompt
fba8c96
import os
from smolagents import OpenAIServerModel, ToolCallingAgent
from tools import (
AudioTranscriberTool,
txt_reader,
pdf_reader,
excel_reader,
math_calculator,
search_engine,
image_analyzer,
YoutubeTranscriberTool,
YoutubeVideoDescriptorTool,
)
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.
Key points:
- Do not add any prefix or suffix to YOUR FINAL ANSWER.
- Answer shortly and concisely, following the template strictly.
- Again, YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
Example:
[Final Answer: 42]
Good example: "42"
Bad example: "The answer is 42"
"""
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,
AudioTranscriberTool(),
txt_reader,
pdf_reader,
excel_reader,
image_analyzer,
YoutubeTranscriberTool(),
YoutubeVideoDescriptorTool(),
],
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)