Update app.py

app.py

@@ -1,52 +1,44 @@
 import os
 import gradio as gr
 import requests
-import inspect
 import pandas as pd
 from dotenv import load_dotenv
-from typing import List, Dict, Any, Tuple, Optional
+from typing import List, Dict, Any, Optional
 
 # LangChain imports
-from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
-from langchain_core.messages import BaseMessage
-from langchain.schema import Document
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
 from langchain_openai import ChatOpenAI
-# from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.tools.wikipedia.tool import WikipediaQueryRun
 from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
 from langchain_community.tools.arxiv.tool import ArxivQueryRun
-from langgraph.graph import StateGraph, START, END
+from langgraph.graph import StateGraph, END
 from langgraph.prebuilt import ToolNode, tools_condition
-from dataclasses import dataclass
-from typing import TypedDict, Annotated, Literal
+from typing import TypedDict
 
-# Constants
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 class MessagesState(TypedDict):
     messages: List[BaseMessage]
 
-# Load system prompt
 try:
     with open("system_prompt.txt", "r", encoding="utf-8") as f:
         system_prompt = f.read()
 except FileNotFoundError:
-    system_prompt = """You are a helpful AI assistant that uses tools to find information and answer questions.
-When you don't know something, use the available tools to look up information. Be concise, direct, and provide accurate responses.
-Always cite your sources when using information from searches or reference materials."""
+    system_prompt = (
+        "You are a helpful AI assistant that uses tools to find information and answer questions.\n"
+        "When you don't know something, use the available tools to look up information. Be concise, direct, and provide accurate responses.\n"
+        "Always cite your sources when using information from searches or reference materials."
+    )
 
-# Advanced agent using LangGraph
 class AdvancedAgent:
     def __init__(self):
         print("Initializing AdvancedAgent with LangGraph, Wikipedia, Arxiv, and Gemini 2.0 Flash")
-        load_dotenv()  # Load environment variables from .env file
+        load_dotenv()
         self.graph = self.build_graph()
         print("Graph successfully built")
         
     def build_graph(self):
-        """Build the LangGraph agent with necessary tools"""
-
         llm = ChatOpenAI(
             model="google/gemini-2.0-flash-001",
             temperature=0,
@@ -61,31 +53,25 @@ class AdvancedAgent:
         tools = [wikipedia_tool, arxiv_tool, tavily_search]
         print(f"Initialized {len(tools)} tools: Wikipedia, Arxiv, Tavily Search")
         
-        sys_msg = SystemMessage(content=system_prompt)
         llm_with_tools = llm.bind_tools(tools)
 
         def assistant(state: MessagesState):
-            """Assistant node that processes messages and generates responses"""
             messages = state["messages"]
-            response = llm_with_tools.invoke(messages)  # <-- messages, not dict!
-            return {"messages": messages + [response]}
+            response = llm_with_tools.invoke(messages)
+            return {"messages": messages + [response]}  # Always return dict
 
         tools_node = ToolNode(tools)
 
         builder = StateGraph(MessagesState)
         builder.add_node("assistant", assistant)
         builder.add_node("tools", tools_node)
-
-        builder.set_entry_point("assistant")
         builder.add_edge("assistant", "tools")
-        builder.add_edge("tools", "assistant")
-        builder.add_edge("assistant", END)
+        builder.set_entry_point("assistant")
         builder.add_conditional_edges(
             "assistant",
             tools_condition,
             {"tools": "tools", END: END}
         )
-
         return builder.compile()
     
     def __call__(self, question: str) -> str:
@@ -95,7 +81,8 @@ class AdvancedAgent:
             HumanMessage(content=question)
         ]
         try:
-            result = self.graph.invoke(messages)  # <-- messages, not dict!
+            # Initial state must be a dict with "messages" key!
+            result = self.graph.invoke({"messages": messages})
             final_messages = result["messages"]
             ai_messages = [msg for msg in final_messages if isinstance(msg, AIMessage)]
             if not ai_messages: