Update app.py

app.py

@@ -11,19 +11,19 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
+
+def assistant(state: AgentState):
+    return {
+        "messages": [chat_with_tools.invoke(state["messages"])],
+    }
+    
 class BasicAgent:
-    def __init__(self):
-        llm = HuggingFaceEndpoint(repo_id="Qwen/Qwen2.5-Coder-32B-Instruct")
+    def __init__(self, llm_checkpoint):
+        llm = HuggingFaceEndpoint(repo_id=llm_checkpoint)
         chat = ChatHuggingFace(llm=llm, verbose=True)
-        tools = [guest_info_tool]
+        self.tools = [guest_info_tool]
         chat_with_tools = chat.bind_tools(tools)
-        self.builder = StateGraph(AgentState)
-        self.builder.add_node("assistant", assistant)
-        self.builder.add_node("tools", ToolNode(tools))
-        self.builder.add_edge(START, "assistant")
-        self.builder.add_conditional_edges("assistant", tools_condition)
-        self.builder.add_edge("tools", "assistant")
-        self.builder.compile()
+        self.graph = self._build_graph()
         print("BasicAgent initialized.")
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
@@ -33,6 +33,15 @@ class BasicAgent:
         fixed_answer = "This is a default answer."
         print(f"Agent returning answer: {response}")
         
+    def _build_graph():
+        self.builder = StateGraph(AgentState)
+        self.builder.add_node("assistant", assistant)
+        self.builder.add_node("tools", ToolNode(self.tools))
+        self.builder.add_edge(START, "assistant")
+        self.builder.add_conditional_edges("assistant", tools_condition)
+        self.builder.add_edge("tools", "assistant")
+        self.builder.compile()
+        
         return fixed_answer
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
@@ -56,7 +65,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
 
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
-        agent = BasicAgent()
+        agent = BasicAgent(llm_checkpoint="Qwen/Qwen2.5-Coder-32B-Instruct")
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None