Update app.py

app.py

@@ -9,6 +9,8 @@ from smolagents import DuckDuckGoSearchTool
 from smolagents import Tool
 import traceback
 from openai import OpenAI
+from transformers import pipeline
+
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
@@ -54,42 +56,36 @@ CalculatorTool=Tool(
 )
 
 class BasicAgent:
-    def __init__(self):
+    def __init__(self, use_local_model=True):
         try:
-            # Load your HF token
-            hf_token = os.getenv("chatbotagenthf")
-            if not hf_token:
-                raise ValueError("HF token not found in environment variable 'chatbotagenthf'.")
-
-            # --- Connect to Hugging Face router ---
-            self.client = OpenAI(
-                base_url="https://router.huggingface.co/v1",
-                api_key=hf_token
-            )
-
-            # --- Choose any accessible instruct model ---
-            self.model_id = "moonshotai/Kimi-K2-Thinking:novita"  
-            # (You can change to: mistralai/Mistral-7B-Instruct-v0.3, zephyr-7b-beta, etc.)
 
-            # --- Define your ToolCallingAgent as before ---
+            print("🚀 Loading local Hugging Face model...")
+            self.pipe = pipeline("text-generation", model="WeiboAI/VibeThinker-1.5B")
+            self.model_type = "local"
             self.agent = ToolCallingAgent(
                 tools=[DuckDuckGoSearchTool()],
-                model=None  # Model handled manually via OpenAI client
+                model=None
             )
-
-            print("✅ Agent initialized successfully using Hugging Face Router.")
+            print("✅ Agent initialized successfully.")
 
         except Exception as e:
             print(f"❌ Error initializing agent: {e}")
             traceback.print_exc()
             raise e
 
-    def __call__(self, question: str) -> str:
-        """Run a chat completion on the HF model."""
-        try:
-            print(f"\n➡️ Agent received question: {question[:80]}")
 
-            # Send the prompt to the model
+   def __call__(self, question: str) -> str:
+    """Run either local HF pipeline or remote API chat."""
+    try:
+        print(f"\n➡️ Agent received question: {question[:80]}")
+
+        if getattr(self, "model_type", None) == "local":
+            # Use local model
+            messages = [{"role": "user", "content": question}]
+            response = self.pipe(messages, max_new_tokens=200)
+            answer = response[0]["generated_text"]
+        else:
+            # Use Hugging Face router (OpenAI-style)
             response = self.client.chat.completions.create(
                 model=self.model_id,
                 messages=[
@@ -97,15 +93,16 @@ class BasicAgent:
                     {"role": "user", "content": question}
                 ],
             )
-
             answer = response.choices[0].message["content"].strip()
-            print(f"✅ Agent returning answer: {answer}")
-            return answer
 
-        except Exception as e:
-            print(f"❌ Agent encountered an error: {e}")
-            traceback.print_exc()
-            return f"Error generating answer: {e}"
+        print(f"✅ Agent returning answer: {answer[:100]}...")
+        return answer
+
+    except Exception as e:
+        print(f"❌ Agent encountered an error: {e}")
+        traceback.print_exc()
+        return f"Error generating answer: {e}"
+