Update analytics/ai_assistant.py

analytics/ai_assistant.py

@@ -11,7 +11,7 @@ except ImportError as exc:
 
 HF_MODEL_ID = os.environ.get("HF_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")
 HF_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN")
-HF_MAX_TOKENS = int(os.environ.get("HF_MAX_TOKENS", 256))
+HF_MAX_TOKENS = int(os.environ.get("HF_MAX_TOKENS", 512))
 HF_TEMPERATURE = float(os.environ.get("HF_TEMPERATURE", 0.3))
 
 
@@ -164,10 +164,14 @@ def generate_ai_answer(question: str, df: pd.DataFrame, as_of_month: str = None,
     #    top_k= 50,
     #    top_p= 0.95
     #)
-    print(response)
+    print(response.choices[0].message.content)
 
-    if isinstance(response, dict):
-        return response.get("generated_text") or response.get("text") or str(response)
-    if isinstance(response, list) and len(response) > 0:
-        return response[0].get("generated_text", str(response[0]))
-    return str(response)
+
+
+    output_text = response.choices[0].message.content if hasattr(response, 'choices') else (response[0].get('generated_text') if isinstance(response, list) else response)
+    return (output_text)
+    #if isinstance(response, dict):
+    #    return response.get("generated_text") or response.get("text") or str(response)
+    #if isinstance(response, list) and len(response) > 0:
+    #    return response[0].get("generated_text", str(response[0]))
+    #return str(response)