agent

agent_utils.py

@@ -16,5 +16,5 @@ def wiki_ret(question: str) -> str:
         ]
     )
 
-    return formatted_search_docs
+    return {"context": formatted_search_docs}
 

app.py

@@ -29,7 +29,7 @@ llm = ChatGroq(api_key=groq_api_key, model="llama-3.3-70b-versatile")
 
 class GraphState(TypedDict):
     messages: str #Annotated[list, operator.add]
-    answer: str
+    context: str
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
@@ -42,13 +42,14 @@ class BasicAgent:
         response = (self.graph).invoke({"messages": question})
         return response
 
-    def agent(self, question: str) -> str:
+    def agent(self, state: GraphState):
 
         # 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 = state.get("context", "")
+        question = state.get("messages", "")
         prompt = f"""
         You are a general AI assistant. I will ask you a question. 
         YOUR FINAL ANSWER should be a number OR 
@@ -58,7 +59,8 @@ 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}"""
+        in the list is a number or a string. Question: {question}
+        For answering the question use this context: {context}"""
         # 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"}}