graph

app.py

@@ -21,13 +21,8 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition  ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
-class BasicAgent:
-    def __init__(self):
-        self.model = ChatOpenAI(model="gpt-4o", temperature=0)
 
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        prompt = f"""
+PROMPT = """
 You are a general AI assistant. I will ask you a question.
 Report your thoughts,always use at least 2 sources to double check your answer, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER].
 YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
@@ -56,16 +51,70 @@ The part after `FINAL ANSWER: ` **MUST** contain **ONLY** the final answer, form
 
 Question: {question}
 """
-        messages = [HumanMessage(content=prompt)]
-        response = self.model.invoke(messages)
 
+
+class AnswerState(TypedDict):
+    question: str
+    answer: Optional[str]
+    messages: List[Dict[str, Any]]
+
+
+class BasicAgent:
+    def __init__(self):
+        self.model = ChatOpenAI(model="gpt-4o", temperature=0)
+        self.graph = StateGraph(AnswerState)
+        self.graph.add_node("log_question", self.log_question)
+        self.graph.add_node("try_answer", self.try_answer)
+        self.graph.add_node("final_answer", self.final_answer)
+
+        self.graph.add_edge(START, "log_question")
+        self.graph.add_edge("log_question", "try_answer")
+
+        # Add conditional edges
+        self.graph.add_conditional_edges(
+            "try_answer",
+            self.try_answer,
+            {"FINAL_ANSWER": "final_answer", "SOME_TOOL": "final_answer"},
+        )
+
+        self.graph.add_edge("final_answer", END)
+        self.compiled_graph = self.graph.compile()
+
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        res = self.compiled_graph.invoke(
+            {
+                "question": question,
+                "answer": None,
+                "messages": [],
+            }
+        )
+        return str(res)
+
+    def log_question(self, state: AnswerState) -> None:
+        print(f"Agent received question (first 50 chars): {state['question']}...")
+
+    def final_answer(self, state: AnswerState) -> str:
+        print(f"Agent returning answer: {state['answer']}")
+        return state["answer"] if state["answer"] else "I don't know"
+
+    def try_answer(self, state: AnswerState) -> str:
+
+        messages = [HumanMessage(content=PROMPT.format(question=state["question"]))]
+        response = self.model.invoke(messages)
         txt = response.text().split("FINAL ANSWER:")
         print(f"Agent returning answer: {txt}")
-
-        if len(txt) > 1:
-            return txt[-1].strip()
-        else:
-            return "I don't know."
+        answer = txt[-1].strip() if len(txt) > 1 else "I don't know."
+        print(f"Agent returning answer: {state['answer']}")
+        state["messages"] = state.get("messages", []) + [
+            {
+                "role": "assistant",
+                "content": f"FINAL ANSWER: {state['answer']}",
+            }
+        ]
+        state["answer"] = answer
+
+        return "FINAL_ANSWER"
 
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):