Update utils.py

utils.py

@@ -13,6 +13,7 @@ def summarize_dataframe(df: pd.DataFrame, max_rows: int = 30) -> str:
     summary += sample.to_string(index=False)
     return summary
 
+
 def query_agent(df: pd.DataFrame, query: str) -> str:
     query_lower = query.lower()
     try:
@@ -39,17 +40,32 @@ Question:
 Answer (with explanation):
 """
 
-    # Initialize the InferenceClient
     client = InferenceClient(token=os.environ.get("HUGGINGFACE_API_KEY"))
 
-    # Use the text_generation method with the correct parameters
-    response = client.text_generation(
-        model="google/gemma-2b-it",
-        inputs=prompt,
-        parameters={"max_new_tokens": 1024, "temperature": 0.7}
-    )
+    try:
+        # ✅ Works for huggingface-hub >=0.33
+        response = client.text_generation(
+            prompt,
+            model="google/gemma-2b-it",
+            max_new_tokens=1024,
+            temperature=0.7,
+        )
+    except TypeError:
+        # ✅ Fallback for older versions that use 'inputs' & 'parameters'
+        response = client.text_generation(
+            model="google/gemma-2b-it",
+            inputs=prompt,
+            parameters={"max_new_tokens": 1024, "temperature": 0.7},
+        )
 
-    # Extract the generated text from the response
-    answer = response[0]["generated_text"] if isinstance(response, list) and "generated_text" in response[0] else str(response)
+    # Extract text safely
+    if isinstance(response, str):
+        answer = response
+    elif isinstance(response, dict) and "generated_text" in response:
+        answer = response["generated_text"]
+    elif isinstance(response, list) and len(response) > 0 and "generated_text" in response[0]:
+        answer = response[0]["generated_text"]
+    else:
+        answer = str(response)
 
-    return answer
+    return answer