File size: 2,610 Bytes
7bfc491
d1b9537
7bfc491
 
 
7c53c52
7bfc491
 
 
 
 
fba8c96
7c53c52
fba8c96
7bfc491
 
d1b9537
 
 
 
 
 
 
 
 
 
 
 
7c53c52
 
 
 
fba8c96
 
 
 
 
 
d1b9537
 
 
 
 
 
 
7bfc491
 
 
 
d1b9537
7bfc491
 
7c53c52
d1b9537
 
 
fba8c96
7c53c52
fba8c96
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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
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)