Update app.py

app.py

@@ -7,10 +7,8 @@ from typing import List, Dict, Union
 import requests
 import wikipediaapi
 import pandas as pd
-from transformers import pipeline  # or HfAgent if you want the higher-level agent
-from huggingface_hub import InferenceClient  # Updated import
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from langchain_core.messages import HumanMessage
+from agent import build_graph
 
 load_dotenv()
 
@@ -23,19 +21,18 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 
 class BasicAgent:
-    def __init__(self, model="google/gemma-7b"):
-        self.pipe = pipeline("text-generation", model=model)
+    """A langgraph agent."""
+    def __init__(self):
         print("BasicAgent initialized.")
+        self.graph = build_graph()
 
     def __call__(self, question: str) -> str:
-        print(f"Question: {question[:50]}...")
-        response = self.pipe(question, max_new_tokens=100)
-        return response[0]['generated_text']
-            
-    def wikipedia_search(self, query: str) -> str:
-        """Get Wikipedia summary"""
-        page = self.wiki.page(query)
-        return page.summary if page.exists() else "No Wikipedia page found"
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        # Wrap the question in a HumanMessage from langchain_core
+        messages = [HumanMessage(content=question)]
+        messages = self.graph.invoke({"messages": messages})
+        answer = messages['messages'][-1].content
+        return answer[14:]