Update app.py

app.py

@@ -9,7 +9,7 @@ import wikipediaapi
 import pandas as pd
 from transformers import HuggingFaceAgent
 from transformers import pipeline  # or HfAgent if you want the higher-level agent
-
+from huggingface_hub import HfAgent
 
 load_dotenv()
 
@@ -22,20 +22,23 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 class BasicAgent:
     def __init__(self, model="google/gemma-7b"):
-        # Using pipeline for text generation
-        self.agent = pipeline("text-generation", model=model)
+        # Initialize the HF Agent with your model endpoint
+        self.agent = HfAgent(
+            endpoint=f"https://api-inference.huggingface.co/models/{model}",
+            headers={"Authorization": f"Bearer {os.getenv('HF_API_KEY')}"}
+        )
         print("BasicAgent initialized.")
 
     def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        answer = self.agent(question, max_length=100)[0]['generated_text']
-        print(f"Agent returning answer (first 50 chars): {answer[:50]}...")
-        return answer
+        print(f"Question: {question[:50]}...")
+        return self.agent.run(question)  # For full agent workflow
+        # OR for simple chat:
+        # return self.agent.chat(question)
 
     def generate_response(self, prompt: str) -> str:
-        """Get response from model"""
+        """Get response from model using chat interface"""
         try:
-            response = self.agent(prompt, max_length=100)[0]['generated_text']
+            response = self.agent.chat(prompt)
             return response
         except Exception as e:
             return f"Error generating response: {str(e)}"