Update app.py

app.py

@@ -155,29 +155,41 @@ def safe_calculator(expression: str) -> str:
 class HfAgentWrapper:
     def __init__(self):
         print("Initializing HfAgentWrapper...")
-        model_id_or_path = "bigcode/starcoderbase-1b"
+        model_id_or_path = "bigcode/starcoderbase-1b" # Compatible model
 
         try:
-            print(f"Attempting to create LLM pipeline for model: {model_id_or_path}")
-            
-            # Retrieve the token from Space Secrets
-            hf_auth_token = os.getenv("HF_AUTH_TOKEN") # Use the name you gave the secret
+            print(f"Attempting to initialize HfAgent with model: {model_id_or_path}")
 
+            hf_auth_token = os.getenv("HF_AUTH_TOKEN")
             if not hf_auth_token:
-                print("WARNING: HF_AUTH_TOKEN secret not found. This will likely fail for gated models like Gemma.")
-                # You might want to raise an error here if the token is absolutely essential
-                # raise ValueError("HF_AUTH_TOKEN secret is missing. Please add it to your Space secrets.")
+                print("WARNING: HF_AUTH_TOKEN secret not found. Starcoder is gated and requires a token.")
+                # For Starcoder, token is necessary
+                raise ValueError("HF_AUTH_TOKEN secret is missing and is required for bigcode/starcoderbase-1b.")
             else:
-                print("HF_AUTH_TOKEN found, will use it for pipeline creation.")
-
-            # 1. Create the pipeline first, with trust_remote_code=True AND the token
-            llm_pipeline = pipeline(
-                task="text-generation", 
-                model=model_id_or_path,
-                trust_remote_code=True,
-                token=hf_auth_token  # <<< --- PASS THE TOKEN HERE
+                print(f"HF_AUTH_TOKEN found (length: {len(hf_auth_token)}), will pass it to HfAgent.")
+
+            # --- Ensure Tool names are correct ---
+            if not get_current_time_in_timezone.__doc__: raise ValueError("Tool 'get_current_time_in_timezone' is missing a docstring.")
+            if not web_search.__doc__: raise ValueError("Tool 'web_search' is missing a docstring.")
+            if not safe_calculator.__doc__: raise ValueError("Tool 'safe_calculator' is missing a docstring.")
+
+            time_tool_obj = Tool(name=get_current_time_in_timezone.__name__, func=get_current_time_in_timezone, description=get_current_time_in_timezone.__doc__)
+            search_tool_obj = Tool(name=web_search.__name__, func=web_search, description=web_search.__doc__)
+            calculator_tool_obj = Tool(name=safe_calculator.__name__, func=safe_calculator, description=safe_calculator.__doc__)
+            self.actual_tools_for_agent = [time_tool_obj, search_tool_obj, calculator_tool_obj]
+            tool_names_for_log = [tool.name for tool in self.actual_tools_for_agent]
+            print(f"Prepared Tool objects with names: {tool_names_for_log} and descriptions.") # Verify this prints correct names!
+            # --- End Tool Name Check ---
+
+            # Let HfAgent create its own pipeline internally.
+            # Pass the model ID string and the hf_token.
+            # trust_remote_code=True is not needed for HfAgent constructor and not for starcoder with 4.36.0.
+            self.agent = HfAgent(
+                llm_endpoint=model_id_or_path,  # Pass model ID string
+                additional_tools=self.actual_tools_for_agent,
+                hf_token=hf_auth_token  # HfAgent uses this for its internal LLMProxy
             )
-            print(f"Successfully created LLM pipeline for {model_id_or_path}.")
+            print(f"HfAgent successfully instantiated for model {model_id_or_path}.")
 
             # ... (rest of your Tool object creations - these were correct)
             if not get_current_time_in_timezone.__doc__: raise ValueError("Tool 'get_current_time_in_timezone' is missing a docstring.")
@@ -203,11 +215,10 @@ class HfAgentWrapper:
             # ...
 
         except Exception as e:
-            print(f"CRITICAL ERROR: Failed to initialize HfAgent or Pipeline: {e}")
-            # import traceback # Already imported at the top of your file
-            print("Full traceback of HfAgent/Pipeline initialization error:")
+            print(f"CRITICAL ERROR: Failed to initialize HfAgent: {e}")
+            print("Full traceback of HfAgent initialization error:")
             traceback.print_exc() 
-            raise RuntimeError(f"HfAgent/Pipeline initialization failed: {e}") from e
+            raise RuntimeError(f"HfAgent initialization failed: {e}") from e
 
     # __call__ method remains the same
     def __call__(self, question: str) -> str: