simple solution

agent_utils.py

@@ -0,0 +1,20 @@
+from langchain_community.document_loaders import WikipediaLoader
+
+
+def wiki_ret(question: str) -> str:
+    """ Retrieve docs from wikipedia """
+
+    # Search
+    search_docs = WikipediaLoader(query=question, 
+                                load_max_docs=2).load()
+
+    # Format
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ]
+    )
+
+    return formatted_search_docs
+

app.py

@@ -3,12 +3,20 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-import re
-import json
+# import re
+# import json
+from typing_extensions import TypedDict
+from typing import list
+
+import operator
+from typing import  Annotated
 
 from langchain_groq import ChatGroq
 from langchain_core.messages import HumanMessage
-from langchain_community.document_loaders import WikipediaLoader
+
+from langgraph.graph import START, END, StateGraph, ToolNode
+
+from agent_utils import wiki_ret
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -18,7 +26,9 @@ groq_api_key = os.getenv("GROQ_API_KEY")
 
 llm = ChatGroq(api_key=groq_api_key, model="llama-3.3-70b-versatile")
 
-
+class GraphState(TypedDict):
+    messages: Annotated[list, operator.add]
+    answer: str
 
 
 # --- Basic Agent Definition ---
@@ -26,35 +36,14 @@ llm = ChatGroq(api_key=groq_api_key, model="llama-3.3-70b-versatile")
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
+        self.graph = self.create_graph()
 
-
-    def wiki_ret(self, question: str) -> str:
-        """ Retrieve docs from wikipedia """
-
-        # Search query
-        # structured_llm = llm.with_structured_output(SearchQuery)
-        # search_query = structured_llm.invoke([search_instructions]+state['messages'])
-        
-        # Search
-        search_docs = WikipediaLoader(query=question, 
-                                    load_max_docs=2).load()
-
-        # Format
-        formatted_search_docs = "\n\n---\n\n".join(
-            [
-                f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-                for doc in search_docs
-            ]
-        )
-
-        return formatted_search_docs
-
-    def __call__(self, question: str) -> str:
+    def agent(self, question: str) -> str:
 
         # print(f"Agent received question (first 50 chars): {question[:50]}...")
         # fixed_answer = "This is a default answer. --- 1"
         # print(f"Agent returning fixed answer: {fixed_answer}")
-        context = self.wiki_ret(question)
+        # context = self.wiki_ret(question)
 
         prompt = f"""
         You are a general AI assistant. I will ask you a question. 
@@ -65,8 +54,7 @@ class BasicAgent:
         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. Question: {question} 
-        For answering questions you can use context from wikiperdia: {context}"""
+        in the list is a number or a string. Question: {question}"""
         # Your answer must be in the following format:
 
         # {{"task_id": "task_id_1", "model_answer": "Answer 1 from your model", "reasoning_trace": "The different steps by which your model reached answer 1"}}
@@ -84,6 +72,20 @@ class BasicAgent:
         # res = json.loads(json_str)
         return response
 
+    def create_graph(self):
+        builder = StateGraph(GraphState)
+        builder.add_node("agent", self.agent)
+        builder.add_node("tools", ToolNode(tools = [wiki_ret]))
+
+        builder.add_edge(START, "agent")
+        builder.add_edge("agent", END)
+        return builder.compile()
+
+    def __call___(self, question):
+        response = (self.graph).invoke({"messages": question})
+        return response
+
+
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,

requirements.txt

@@ -3,4 +3,5 @@ requests
 langchain-groq
 langchain-community
 langchain-core
-wikipedia
+wikipedia
+langgraph