Update agent.py

agent.py

@@ -12,12 +12,14 @@ from langchain_community.document_loaders import WikipediaLoader
 from langchain_google_genai import ChatGoogleGenerativeAI
 from pathlib import Path
 import tempfile
+from dotenv import load_dotenv
 
 # constants
 API_URL = "https://agents-course-unit4-scoring.hf.space"
 QUESTIONS_URL = f"{API_URL}/questions"
 FILES_URL = f"{API_URL}/files"
 SUBMIT_URL = f"{API_URL}/submit"
+load_dotenv()
 
 class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
@@ -28,7 +30,7 @@ class AgentState(TypedDict):
 def build_gemini_llm():
     if not os.environ.get("GOOGLE_API_KEY"):
         raise ValueError("GOOGLE_API_KEY environment variable is not set.")
-    return ChatGoogleGenerativeAI(model = "gemini-3.7-flash", temperature = 0, max_output_tokens = 1024)
+    return ChatGoogleGenerativeAI(model = "gemini-3.7-flash", temperature = 0, max_output_tokens = 1025, include_thoughts=True)
 
 @tool
 def add_numbers(a: int, b: int) -> int:
@@ -228,15 +230,16 @@ tools = [
     multiply_numbers,
     divide_numbers,
     extract_text_from_image,
-    wikipedia_search,
-    download_and_read_file,
     search_web
 ]
 model_with_tools = model.bind_tools(tools)
 
 def assistant(state: AgentState):
+    response = model_with_tools.invoke(state["messages"])
+    reasoning_tokens = response.usage_metadata["output_token_details"]["reasoning"]
+    print("Reasoning tokens used:", reasoning_tokens)
     return {
-        "messages": state["messages"],
+        "messages": [response],
         "file_path": state["file_path"],
         "task_id": state["task_id"],
         "url": state["url"]
@@ -251,4 +254,4 @@ builder.add_edge(START, "assistant")
 builder.add_conditional_edges("assistant", tools_condition)
 builder.add_edge("tools", "assistant")
 
-graph = builder.compile()
+graph = builder.compile()