File size: 8,013 Bytes
4d89d4a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c49b8ba
4d89d4a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
715a773
 
4d89d4a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4e188ba
4d89d4a
 
c49b8ba
 
4d89d4a
 
 
4e188ba
4d89d4a
 
 
 
 
 
 
 
 
 
 
 
4e188ba
4d89d4a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
import logging
import os

import tempfile

from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_google_community import GoogleSearchResults
from langchain_google_community import GoogleSearchAPIWrapper
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import Tool
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_experimental.utilities import PythonREPL

from tools_audio import transcribe_audio

from tools_doc import (
    analyze_csv_file,
    analyze_excel_file,
    download_file_from_url,
    extract_text_from_image,
    read_file,
)
from tools_video import (
    review_youtube_video,
    use_vision_model,
    transcribe_youtube,
    video_frames_to_images,
)

from tools_browser import website_scrape, web_search

from answers import create_final_answer_graph, validate_answer

logger = logging.getLogger(__name__)
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"


class BasicAgent:
    def __init__(self):
        try:
            logger.info("Initializing BasicAgent")

            # Create the prompt template
            prompt = ChatPromptTemplate.from_messages(
                [
                    (
                        "system",
                        """You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [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.
                """,
                    ),
                    ("placeholder", "{chat_history}"),
                    ("human", "{input}"),
                    ("placeholder", "{agent_scratchpad}"),
                ]
            )
            logger.info("Created prompt template")

            # Initialize Gemini model
            logger.info("Creating Gemini model...")
            llm = ChatGoogleGenerativeAI(
                # model="models/gemini-2.5-flash-preview-04-17",
                model="gemini-2.5-flash-preview-04-17",
                google_api_key=os.getenv("GEMINI_KEY"),
                temperature=0.2,
            )
            logger.info("Created Gemini model successfully")

            # Define available tools
            tools = [
                # GoogleSearchResults(
                #     api_wrapper=GoogleSearchAPIWrapper(
                #         google_api_key=os.getenv("GOOGLE_SEARCH_API_KEY"),
                #         google_cse_id=os.getenv("GOOGLE_CSE_ID"),
                #         k=5,  # Number of results to return
                #     )
                # ),
                web_search,
                analyze_csv_file,
                analyze_excel_file,
                download_file_from_url,
                extract_text_from_image,
                read_file,
                review_youtube_video,
                transcribe_audio,
                transcribe_youtube,
                use_vision_model,
                video_frames_to_images,
                website_scrape,
                Tool(
                    name="python_repl",
                    description="A Python shell. Use this to execute python commands. Input # should be a valid python command. If you want to see the output of a value, # you should print it out with `print(...)`.",
                    func=PythonREPL().run,
                ),
            ]
            logger.info("Tools: %s", tools)

            # Create the agent
            agent = create_tool_calling_agent(llm, tools, prompt)
            logger.info("Created tool calling agent")

            # Create the agent executor
            self.agent_executor = AgentExecutor(
                agent=agent,
                tools=tools,
                return_intermediate_steps=True,
                verbose=True,
            )
            logger.info("Created agent executor")

            # Create the graph
            self.validation_graph = create_final_answer_graph()

        except Exception as e:
            logger.error("Error initializing agent: %s", e, exc_info=True)
            raise

    def __call__(self, question: str, task_id: str) -> str:
        """Execute the agent with the given question and optional file.
        Args:
            question (str): The question to answer
            task_id (str): The task ID to fetch the file
        """
        max_retries = 3
        attempt = 0

        # Create a temporary directory that will be automatically cleaned up
        print("HELLO")
        with tempfile.TemporaryDirectory() as temp_dir:
            while attempt < max_retries:
                # default_api_url = os.getenv("DEFAULT_API_URL")
                default_api_url = DEFAULT_API_URL
                file_url = f"{default_api_url}/files/{task_id}"

                try:
                    print("HELLO-A")
                    # Download file to temporary directory
                    file = download_file_from_url.invoke(
                        {
                            "url": file_url,
                            "directory": temp_dir,
                        }
                    )
                except Exception as e:
                    logger.error(f"Error downloading file: {e}")
                    file = None

                try:
                    print("HELLO-B")
                    attempt += 1
                    logger.info(f"Attempt {attempt} of {max_retries}")

                    # Prepare input with file information
                    if file and file.get("type") != "error":
                        input_data = {
                            "input": question
                            + f" [File: type={file.get('type', 'None')}, path={file.get('path', 'None')}]",
                        }
                    else:
                        input_data = {
                            "input": question,
                        }

                    # Run the agent to get the answer
                    result = self.agent_executor.invoke(input_data)
                    answer = result.get("output", "")

                    logger.info(f"Attempt {attempt} result: {result}")

                    # Run validation
                    validation_result = validate_answer(
                        self.validation_graph,
                        answer,
                        [result.get("intermediate_steps", [])],
                    )

                    valid_answer = validation_result.get("valid_answer", False)
                    final_answer = validation_result.get("final_answer", "")

                    if valid_answer:
                        logger.info(f"Valid answer found on attempt {attempt}")
                        return final_answer

                    logger.warning(
                        f"Validation failed on attempt {attempt}: {final_answer}"
                    )
                    if attempt >= max_retries:
                        raise Exception(
                            f"Failed to get valid answer after {max_retries} attempts. Last error: {final_answer}"
                        )

                except Exception as e:
                    logger.error(
                        f"Error in attempt {attempt}: {e}", exc_info=True
                    )
                    if attempt >= max_retries:
                        raise Exception(
                            f"Failed after {max_retries} attempts. Last error: {str(e)}"
                        )
                    continue