Update app.py

app.py

@@ -1,17 +1,50 @@
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import tools_condition, ToolNode
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import tool
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_community.tools.tavily_search import TavilySearchResults
+from groq import Groq
 import os
-import gradio as gr
-import requests
-import string
-import warnings
-import pandas as pd
 import re
-from huggingface_hub import login
-from groq import Groq
 
-# --- Constants ---
-DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+# Define the tools
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return up to 2 results."""
+    docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    return "\n\n".join([doc.page_content for doc in docs])
+
+@tool
+def web_search(query: str) -> str:
+    """Search the web using Tavily."""
+    docs = TavilySearchResults(max_results=3).invoke(query)
+    return "\n\n".join([doc.page_content for doc in docs])
+
+@tool
+def arvix_search(query: str) -> str:
+    """Search Arxiv and return up to 3 results."""
+    docs = ArxivLoader(query=query, load_max_docs=3).load()
+    return "\n\n".join([doc.page_content[:1000] for doc in docs])
+
+# Tool-based LangGraph builder
+def build_tool_graph(system_prompt):
+    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+    llm_with_tools = llm.bind_tools([wiki_search, web_search, arvix_search])
+
+    def assistant(state: MessagesState):
+        return {"messages": [llm_with_tools.invoke(state["messages"]) ]}
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode([wiki_search, web_search, arvix_search]))
+    builder.set_entry_point("assistant")
+    builder.set_finish_point("assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
+    return builder.compile()
 
-# --- Basic Agent Definition ---
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
@@ -28,11 +61,53 @@ class BasicAgent:
             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."""
         )
+        self.tool_chain = build_tool_graph(self.agent_prompt)
 
     def format_final_answer(self, answer: str) -> str:
         cleaned = " ".join(answer.split())
         return f"FINAL ANSWER: {cleaned}"
 
+    def query_groq(self, question: str) -> str:
+        full_prompt = f"{self.agent_prompt}\n\nQuestion: {question}"
+        try:
+            response = self.client.chat.completions.create(
+                model="llama3-8b-8192",
+                messages=[{"role": "user", "content": full_prompt}]
+            )
+            answer = response.choices[0].message.content
+            print(f"[Groq Raw Response]: {answer}")
+            return self.format_final_answer(answer).upper()
+        except Exception as e:
+            print(f"[Groq ERROR]: {e}")
+            return self.format_final_answer("GROQ_ERROR")
+
+    def query_tools(self, question: str) -> str:
+        try:
+            input_state = {
+                "messages": [
+                    SystemMessage(content=self.agent_prompt),
+                    HumanMessage(content=question)
+                ]
+            }
+            result = self.tool_chain.invoke(input_state)
+            final_msg = result["messages"][-1].content
+            print(f"[LangGraph Final Response]: {final_msg}")
+            return self.format_final_answer(final_msg)
+        except Exception as e:
+            print(f"[LangGraph ERROR]: {e}")
+            return self.format_final_answer("TOOL_ERROR")
+
+    def __call__(self, question: str) -> str:
+        print(f"Received question: {question[:50]}...")
+        if "commutative" in question.lower():
+            return self.check_commutativity()
+        if self.maybe_reversed(question):
+            print("Detected likely reversed riddle.")
+            return self.solve_riddle(question)
+        if "use tools" in question.lower():
+            return self.query_tools(question)
+        return self.query_groq(question)
+
     def check_commutativity(self):
         S = ['a', 'b', 'c', 'd', 'e']
         counter_example_elements = set()
@@ -75,126 +150,4 @@ class BasicAgent:
                 }
                 opposite = opposites.get(word, f"UNKNOWN_OPPOSITE_OF_{word}")
                 return f"FINAL ANSWER: {opposite.upper()}"
-        return self.format_final_answer("COULD_NOT_SOLVE")
-
-    def query_groq(self, question: str) -> str:
-        full_prompt = f"{self.agent_prompt}\n\nQuestion: {question}"
-        try:
-            response = self.client.chat.completions.create(
-                model="llama3-8b-8192",
-                messages=[{"role": "user", "content": full_prompt}]
-            )
-            answer = response.choices[0].message.content
-            print(f"[Groq Raw Response]: {answer}")
-            if "FINAL ANSWER: " in answer:
-                return answer.split("FINAL ANSWER: ")[-1].strip().upper()
-            else:
-                return self.format_final_answer(answer).upper()
-        except Exception as e:
-            print(f"[Groq ERROR]: {e}")
-            return self.format_final_answer("GROQ_ERROR")
-
-    def __call__(self, question: str) -> str:
-        print(f"Received question: {question[:50]}...")
-        if "commutative" in question.lower():
-            return self.check_commutativity()
-        if self.maybe_reversed(question):
-            print("Detected likely reversed riddle.")
-            return self.solve_riddle(question)
-        return self.query_groq(question)
-
-def run_and_submit_all(profile: gr.OAuthProfile | None):
-    space_id = os.getenv("SPACE_ID")
-    if profile:
-        username = f"{profile.username}"
-        print("User logged in.")
-    else:
-        print("User not logged in.")
-        return "Please Login to Hugging Face with the button.", None
-
-    api_url = DEFAULT_API_URL
-    questions_url = f"{api_url}/questions"
-    submit_url = f"{api_url}/submit"
-
-    try:
-        agent = BasicAgent()
-    except Exception as e:
-        return f"Error initializing agent: {e}", None
-
-    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
-    try:
-        response = requests.get(questions_url, timeout=15)
-        response.raise_for_status()
-        questions_data = response.json()
-        if not questions_data:
-            return "Fetched questions list is empty or invalid format.", None
-    except requests.exceptions.RequestException as e:
-        return f"Error fetching questions: {e}", None
-
-    results_log = []
-    answers_payload = []
-    for item in questions_data:
-        task_id = item.get("task_id")
-        question_text = item.get("question")
-
-        if not task_id or question_text is None:
-            continue
-        try:
-            submitted_answer = agent(question_text)
-            print(f"Q: {question_text}")
-            print(f"Predicted: {submitted_answer}")
-
-            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
-            results_log.append({
-                "Task ID": task_id,
-                "Question": question_text,
-                "Submitted Answer": submitted_answer
-            })
-        except Exception as e:
-            results_log.append({
-                "Task ID": task_id,
-                "Question": question_text,
-                "Submitted Answer": f"AGENT ERROR: {e}"
-            })
-
-    if not answers_payload:
-        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
-
-    submission_data = {
-        "username": username.strip(),
-        "agent_code": agent_code,
-        "answers": answers_payload
-    }
-
-    try:
-        response = requests.post(submit_url, json=submission_data, timeout=60)
-        response.raise_for_status()
-        result_data = response.json()
-        print(result_data)
-        final_status = (
-            f"Submission Successful!\n"
-            f"User: {result_data.get('username')}\n"
-            f"Overall Score: {result_data.get('score', '?')}% "
-            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
-            f"Message: {result_data.get('message', 'No message received.')}"
-        )
-        return final_status, pd.DataFrame(results_log)
-    except Exception as e:
-        return f"Submission Failed: {e}", pd.DataFrame(results_log)
-
-# --- Build Gradio Interface ---
-with gr.Blocks() as demo:
-    gr.Markdown("# Basic Agent Evaluation Runner")
-    gr.LoginButton()
-    run_button = gr.Button("Run Evaluation & Submit All Answers")
-    status_output = gr.Textbox(label="Run Status / Submission Result", max_lines=5, interactive=False, max_length=200)
-    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)
-
-    run_button.click(
-        fn=run_and_submit_all,
-        outputs=[status_output, results_table]
-    )
-
-if __name__ == "__main__":
-    print("Launching Gradio Interface for Basic Agent Evaluation...")
-    demo.launch(debug=True, share=False)
+    return self.format_final_answer("COULD_NOT_SOLVE")