Update app.py

app.py

@@ -62,32 +62,41 @@ def safe_calculator_func(expression: str) -> str:
         print(f"Error during calculation for '{expression}': {e}")
         return f"Error calculating '{expression}': Invalid expression or calculation error ({e})."
 
+# --- Custom Pipeline to Fix LangChain Integration ---
+class FlanT5Pipeline(transformers.Pipeline):
+    def _call(self, *args, **kwargs):
+        # The HuggingFacePipeline class in LangChain might not pass the input
+        # with the 'inputs' keyword. This custom _call method ensures that
+        # whatever is passed as the first argument is correctly forwarded.
+        if args and len(args) > 0:
+            return super()._call(args[0], **kwargs)
+        else:
+            # Fallback in case no positional arguments are provided
+            return super()._call(kwargs)
 
 class LangChainAgentWrapper:
     def __init__(self):
         print("Initializing LangChainAgentWrapper...")
         
+        # Switched to a smaller, CPU-friendly instruction-tuned model
         model_id = "google/flan-t5-base"
         
         try:
+            hf_auth_token = os.getenv("HF_TOKEN")
             print(f"Loading model pipeline for: {model_id}")
             
-            # --- MODIFICATION: Use the custom pipeline class ---
-            # Load the tokenizer first
+            # We load the model and tokenizer objects first
             tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
-            # Load the model
             model = transformers.AutoModelForSeq2SeqLM.from_pretrained(model_id)
 
-            # Create an instance of our custom pipeline
+            # Now we use our custom FlanT5Pipeline class
             llm_pipeline = FlanT5Pipeline(
                 task="text2text-generation",
                 model=model,
                 tokenizer=tokenizer,
                 device_map="auto",
-                max_new_tokens=512 # Add max_new_tokens to control output length
+                max_new_tokens=512
             )
-            # --- END MODIFICATION ---
-
             print("Model pipeline loaded successfully.")
             
             # Wrap the pipeline in a LangChain LLM object