Update src.py

src.py

@@ -1,379 +1,548 @@
-import os
-import pandas as pd
-from pandasai import Agent, SmartDataframe
-from typing import Tuple
-from PIL import Image
-from pandasai.llm import HuggingFaceTextGen
-from dotenv import load_dotenv
-from langchain_groq import ChatGroq
-from langchain_google_genai import ChatGoogleGenerativeAI
-import matplotlib.pyplot as plt
-
-# FORCE reload environment variables
-load_dotenv(override=True)
-
-# Get API keys with explicit None handling and debugging
-Groq_Token = os.getenv("GROQ_API_KEY")
-hf_token = os.getenv("HF_TOKEN")
-gemini_token = os.getenv("GEMINI_TOKEN")
-
-# Debug print (remove in production)
-print(f"Debug - Groq Token: {'Present' if Groq_Token else 'Missing'}")
-print(f"Debug - Groq Token Value: {Groq_Token[:10] + '...' if Groq_Token else 'None'}")
-print(f"Debug - Gemini Token: {'Present' if gemini_token else 'Missing'}")
-
-models = {
-    "mistral": "mistral-saba-24b", 
-    "llama3.3": "llama-3.3-70b-versatile",
-    "llama3.1": "llama-3.1-8b-instant", 
-    "gemma2": "gemma2-9b-it", 
-    "gemini-pro": "gemini-1.5-pro"
-}
-
-def preprocess_and_load_df(path: str) -> pd.DataFrame:
-    """Load and preprocess the dataframe"""
-    try:
-        df = pd.read_csv(path)
-        df["Timestamp"] = pd.to_datetime(df["Timestamp"])
-        return df
-    except Exception as e:
-        raise Exception(f"Error loading dataframe: {e}")
-
-def load_agent(df: pd.DataFrame, context: str, inference_server: str, name="mistral") -> Agent:
-    """Load pandas AI agent with error handling"""
-    try:
-        if name == "gemini-pro":
-            if not gemini_token or gemini_token.strip() == "":
-                raise ValueError("Gemini API token not available or empty")
-            llm = ChatGoogleGenerativeAI(
-                model=models[name], 
-                google_api_key=gemini_token, 
-                temperature=0.1
-            )
-        else:
-            if not Groq_Token or Groq_Token.strip() == "":
-                raise ValueError("Groq API token not available or empty")
-            llm = ChatGroq(
-                model=models[name], 
-                api_key=Groq_Token,
-                temperature=0.1
-            )
-        
-        agent = Agent(df, config={"llm": llm, "enable_cache": False, "options": {"wait_for_model": True}})
-        if context:
-            agent.add_message(context)
-        return agent
-    except Exception as e:
-        raise Exception(f"Error loading agent: {e}")
-
-def load_smart_df(df: pd.DataFrame, inference_server: str, name="mistral") -> SmartDataframe:
-    """Load smart dataframe with error handling"""
-    try:
-        if name == "gemini-pro":
-            if not gemini_token or gemini_token.strip() == "":
-                raise ValueError("Gemini API token not available or empty")
-            llm = ChatGoogleGenerativeAI(
-                model=models[name], 
-                google_api_key=gemini_token, 
-                temperature=0.1
-            )
-        else:
-            if not Groq_Token or Groq_Token.strip() == "":
-                raise ValueError("Groq API token not available or empty")
-            llm = ChatGroq(
-                model=models[name], 
-                api_key=Groq_Token,
-                temperature=0.1
-            )
-        
-        df = SmartDataframe(df, config={"llm": llm, "max_retries": 5, "enable_cache": False})
-        return df
-    except Exception as e:
-        raise Exception(f"Error loading smart dataframe: {e}")
-
-def get_from_user(prompt):
-    """Format user prompt"""
-    return {"role": "user", "content": prompt}
-
-def ask_agent(agent: Agent, prompt: str) -> dict:
-    """Ask agent with comprehensive error handling"""
-    try:
-        response = agent.chat(prompt)
-        gen_code = getattr(agent, 'last_code_generated', '')
-        ex_code = getattr(agent, 'last_code_executed', '')
-        last_prompt = getattr(agent, 'last_prompt', prompt)
-        
-        return {
-            "role": "assistant", 
-            "content": response, 
-            "gen_code": gen_code, 
-            "ex_code": ex_code, 
-            "last_prompt": last_prompt,
-            "error": None
-        }
-    except Exception as e:
-        return {
-            "role": "assistant", 
-            "content": f"Error: {str(e)}", 
-            "gen_code": "", 
-            "ex_code": "", 
-            "last_prompt": prompt,
-            "error": str(e)
-        }
-
-def decorate_with_code(response: dict) -> str:
-    """Decorate response with code details"""
-    gen_code = response.get("gen_code", "No code generated")
-    last_prompt = response.get("last_prompt", "No prompt")
-    
-    return f"""<details>
-<summary>Generated Code</summary>
-    
-```python
-{gen_code}
-```
-</details>
-
-<details>
-<summary>Prompt</summary>
-
-{last_prompt}
-"""
-
-def show_response(st, response):
-    """Display response with error handling"""
-    try:
-        with st.chat_message(response["role"]):
-            content = response.get("content", "No content")
-            
-            try:
-                # Try to open as image
-                image = Image.open(content)
-                if response.get("gen_code"):
-                    st.markdown(decorate_with_code(response), unsafe_allow_html=True)
-                st.image(image)
-                return {"is_image": True}
-            except:
-                # Not an image, display as text
-                if response.get("gen_code"):
-                    display_content = decorate_with_code(response) + f"""</details>
-
-{content}"""
-                else:
-                    display_content = content
-                st.markdown(display_content, unsafe_allow_html=True)
-                return {"is_image": False}
-    except Exception as e:
-        st.error(f"Error displaying response: {e}")
-        return {"is_image": False}
-
-def ask_question(model_name, question):
-    """Ask question with comprehensive error handling"""
-    try:
-        # Reload environment variables to get fresh values
-        load_dotenv(override=True)
-        fresh_groq_token = os.getenv("GROQ_API_KEY")
-        fresh_gemini_token = os.getenv("GEMINI_TOKEN")
-        
-        print(f"ask_question - Fresh Groq Token: {'Present' if fresh_groq_token else 'Missing'}")
-        
-        # Check API availability with fresh tokens
-        if model_name == "gemini-pro":
-            if not fresh_gemini_token or fresh_gemini_token.strip() == "":
-                return {
-                    "role": "assistant", 
-                    "content": "❌ Gemini API token not available or empty. Please set GEMINI_TOKEN in your environment variables.", 
-                    "gen_code": "", 
-                    "ex_code": "", 
-                    "last_prompt": question,
-                    "error": "Missing or empty API token"
-                }
-            llm = ChatGoogleGenerativeAI(
-                model=models[model_name], 
-                google_api_key=fresh_gemini_token,
-                temperature=0
-            )
-        else:
-            if not fresh_groq_token or fresh_groq_token.strip() == "":
-                return {
-                    "role": "assistant", 
-                    "content": "❌ Groq API token not available or empty. Please set GROQ_API_KEY in your environment variables and restart the application.", 
-                    "gen_code": "", 
-                    "ex_code": "", 
-                    "last_prompt": question,
-                    "error": "Missing or empty API token"
-                }
-            
-            # Test the API key by trying to create the client
-            try:
-                llm = ChatGroq(
-                    model=models[model_name], 
-                    api_key=fresh_groq_token,
-                    temperature=0.1
-                )
-                # Test with a simple call to verify the API key works
-                test_response = llm.invoke("Test")
-                print("API key test successful")
-            except Exception as api_error:
-                error_msg = str(api_error).lower()
-                if "organization_restricted" in error_msg or "unauthorized" in error_msg:
-                    return {
-                        "role": "assistant", 
-                        "content": "❌ API Key Error: Your Groq API key appears to be invalid, expired, or restricted. Please check your API key in the .env file.", 
-                        "gen_code": "", 
-                        "ex_code": "", 
-                        "last_prompt": question,
-                        "error": f"API key validation failed: {str(api_error)}"
-                    }
-                else:
-                    return {
-                        "role": "assistant", 
-                        "content": f"❌ API Connection Error: {str(api_error)}", 
-                        "gen_code": "", 
-                        "ex_code": "", 
-                        "last_prompt": question,
-                        "error": str(api_error)
-                    }
-
-        # Check if data file exists
-        if not os.path.exists("Data.csv"):
-            return {
-                "role": "assistant", 
-                "content": "❌ Data.csv file not found. Please ensure the data file is in the correct location.", 
-                "gen_code": "", 
-                "ex_code": "", 
-                "last_prompt": question,
-                "error": "Data file not found"
-            }
-
-        df_check = pd.read_csv("Data.csv")
-        df_check["Timestamp"] = pd.to_datetime(df_check["Timestamp"])
-        df_check = df_check.head(5)
-
-        new_line = "\n"
-        parameters = {"font.size": 12, "figure.dpi": 600}
-
-        template = f"""```python
-import pandas as pd
-import matplotlib.pyplot as plt
-import uuid
-
-plt.rcParams.update({parameters})
-
-df = pd.read_csv("Data.csv")
-df["Timestamp"] = pd.to_datetime(df["Timestamp"])
-
-# Available columns and data types:
-{new_line.join(map(lambda x: '# '+x, str(df_check.dtypes).split(new_line)))}
-
-# Question: {question.strip()}
-# Generate code to answer the question and save result in 'answer' variable
-# If creating a plot, save it with a unique filename and store the filename in 'answer'
-# If returning text/numbers, store the result directly in 'answer'
-```"""
-
-        system_prompt = """You are a helpful assistant that generates Python code for data analysis.
-        
-        Rules:
-        1. Always save your final result in a variable called 'answer'
-        2. If creating a plot, save it with plt.savefig() and store the filename in 'answer'
-        3. If returning text/numbers, store the result directly in 'answer'
-        4. Use descriptive variable names and add comments
-        5. Handle potential errors gracefully
-        6. For plots, use unique filenames to avoid conflicts
-        """
-
-        query = f"""{system_prompt}
-
-Complete the following code to answer the user's question:
-
-{template}
-"""
-        
-        # Make API call
-        if model_name == "gemini-pro":
-            response = llm.invoke(query)
-            answer = response.content
-        else:
-            response = llm.invoke(query)
-            answer = response.content
-        
-        # Extract and execute code
-        try:
-            if "```python" in answer:
-                code_part = answer.split("```python")[1].split("```")[0]
-            else:
-                code_part = answer
-            
-            full_code = f"""
-{template.split("```python")[1].split("```")[0]}
-{code_part}
-"""
-            
-            # Execute code in a controlled environment
-            local_vars = {}
-            global_vars = {
-                'pd': pd,
-                'plt': plt,
-                'os': os,
-                'uuid': __import__('uuid')
-            }
-            
-            exec(full_code, global_vars, local_vars)
-            
-            # Get the answer
-            if 'answer' in local_vars:
-                answer_result = local_vars['answer']
-            else:
-                answer_result = "No answer variable found in generated code"
-                
-            return {
-                "role": "assistant", 
-                "content": answer_result, 
-                "gen_code": full_code, 
-                "ex_code": full_code, 
-                "last_prompt": question,
-                "error": None
-            }
-            
-        except Exception as code_error:
-            return {
-                "role": "assistant", 
-                "content": f"❌ Error executing generated code: {str(code_error)}", 
-                "gen_code": full_code if 'full_code' in locals() else "", 
-                "ex_code": full_code if 'full_code' in locals() else "", 
-                "last_prompt": question,
-                "error": str(code_error)
-            }
-            
-    except Exception as e:
-        error_msg = str(e)
-        
-        # Handle specific API errors
-        if "organization_restricted" in error_msg:
-            return {
-                "role": "assistant", 
-                "content": "❌ API Organization Restricted: Your API key access has been restricted. Please check your Groq API key or try generating a new one.", 
-                "gen_code": "", 
-                "ex_code": "", 
-                "last_prompt": question,
-                "error": "API access restricted"
-            }
-        elif "rate_limit" in error_msg.lower():
-            return {
-                "role": "assistant", 
-                "content": "❌ Rate limit exceeded. Please wait a moment and try again.", 
-                "gen_code": "", 
-                "ex_code": "", 
-                "last_prompt": question,
-                "error": "Rate limit exceeded"
-            }
-        else:
-            return {
-                "role": "assistant", 
-                "content": f"❌ Error: {error_msg}", 
-                "gen_code": "", 
-                "ex_code": "", 
-                "last_prompt": question,
-                "error": error_msg
-            }
+import os
+import pandas as pd
+from pandasai import Agent, SmartDataframe
+from typing import Tuple
+from PIL import Image
+from pandasai.llm import HuggingFaceTextGen
+from dotenv import load_dotenv
+from langchain_groq import ChatGroq
+from langchain_google_genai import ChatGoogleGenerativeAI
+import matplotlib.pyplot as plt
+import json
+from datetime import datetime
+from huggingface_hub import HfApi
+import uuid
+
+# FORCE reload environment variables
+load_dotenv(override=True)
+
+# Get API keys with explicit None handling and debugging
+Groq_Token = os.getenv("GROQ_API_KEY")
+hf_token = os.getenv("HF_TOKEN")
+gemini_token = os.getenv("GEMINI_TOKEN")
+
+# Debug print (remove in production)
+print(f"Debug - Groq Token: {'Present' if Groq_Token else 'Missing'}")
+print(f"Debug - Groq Token Value: {Groq_Token[:10] + '...' if Groq_Token else 'None'}")
+print(f"Debug - Gemini Token: {'Present' if gemini_token else 'Missing'}")
+
+models = {
+    "mistral": "mistral-saba-24b", 
+    "llama3.3": "llama-3.3-70b-versatile",
+    "llama3.1": "llama-3.1-8b-instant", 
+    "gemma2": "gemma2-9b-it", 
+    "gemini-pro": "gemini-1.5-pro"
+}
+
+def log_interaction(user_query, model_name, response_content, generated_code, execution_time, error_message=None, is_image=False):
+    """Log user interactions to Hugging Face dataset"""
+    try:
+        if not hf_token or hf_token.strip() == "":
+            print("Warning: HF_TOKEN not available, skipping logging")
+            return
+        
+        # Create log entry
+        log_entry = {
+            "timestamp": datetime.now().isoformat(),
+            "session_id": str(uuid.uuid4()),
+            "user_query": user_query,
+            "model_name": model_name,
+            "response_content": str(response_content),
+            "generated_code": generated_code or "",
+            "execution_time_seconds": execution_time,
+            "error_message": error_message or "",
+            "is_image_output": is_image,
+            "success": error_message is None
+        }
+        
+        # Create DataFrame
+        df = pd.DataFrame([log_entry])
+        
+        # Create unique filename with timestamp
+        timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
+        random_id = str(uuid.uuid4())[:8]
+        filename = f"interaction_log_{timestamp_str}_{random_id}.parquet"
+        
+        # Save locally first
+        local_path = f"/tmp/{filename}"
+        df.to_parquet(local_path, index=False)
+        
+        # Upload to Hugging Face
+        api = HfApi(token=hf_token)
+        api.upload_file(
+            path_or_fileobj=local_path,
+            path_in_repo=f"data/{filename}",
+            repo_id="SustainabilityLabIITGN/VayuBuddy_logs",
+            repo_type="dataset",
+        )
+        
+        # Clean up local file
+        if os.path.exists(local_path):
+            os.remove(local_path)
+            
+        print(f"Successfully logged interaction to HuggingFace: {filename}")
+        
+    except Exception as e:
+        print(f"Error logging interaction: {e}")
+
+def preprocess_and_load_df(path: str) -> pd.DataFrame:
+    """Load and preprocess the dataframe"""
+    try:
+        df = pd.read_csv(path)
+        df["Timestamp"] = pd.to_datetime(df["Timestamp"])
+        return df
+    except Exception as e:
+        raise Exception(f"Error loading dataframe: {e}")
+
+def load_agent(df: pd.DataFrame, context: str, inference_server: str, name="mistral") -> Agent:
+    """Load pandas AI agent with error handling"""
+    try:
+        if name == "gemini-pro":
+            if not gemini_token or gemini_token.strip() == "":
+                raise ValueError("Gemini API token not available or empty")
+            llm = ChatGoogleGenerativeAI(
+                model=models[name], 
+                google_api_key=gemini_token, 
+                temperature=0.1
+            )
+        else:
+            if not Groq_Token or Groq_Token.strip() == "":
+                raise ValueError("Groq API token not available or empty")
+            llm = ChatGroq(
+                model=models[name], 
+                api_key=Groq_Token,
+                temperature=0.1
+            )
+        
+        agent = Agent(df, config={"llm": llm, "enable_cache": False, "options": {"wait_for_model": True}})
+        if context:
+            agent.add_message(context)
+        return agent
+    except Exception as e:
+        raise Exception(f"Error loading agent: {e}")
+
+def load_smart_df(df: pd.DataFrame, inference_server: str, name="mistral") -> SmartDataframe:
+    """Load smart dataframe with error handling"""
+    try:
+        if name == "gemini-pro":
+            if not gemini_token or gemini_token.strip() == "":
+                raise ValueError("Gemini API token not available or empty")
+            llm = ChatGoogleGenerativeAI(
+                model=models[name], 
+                google_api_key=gemini_token, 
+                temperature=0.1
+            )
+        else:
+            if not Groq_Token or Groq_Token.strip() == "":
+                raise ValueError("Groq API token not available or empty")
+            llm = ChatGroq(
+                model=models[name], 
+                api_key=Groq_Token,
+                temperature=0.1
+            )
+        
+        df = SmartDataframe(df, config={"llm": llm, "max_retries": 5, "enable_cache": False})
+        return df
+    except Exception as e:
+        raise Exception(f"Error loading smart dataframe: {e}")
+
+def get_from_user(prompt):
+    """Format user prompt"""
+    return {"role": "user", "content": prompt}
+
+def ask_agent(agent: Agent, prompt: str) -> dict:
+    """Ask agent with comprehensive error handling"""
+    start_time = datetime.now()
+    try:
+        response = agent.chat(prompt)
+        execution_time = (datetime.now() - start_time).total_seconds()
+        
+        gen_code = getattr(agent, 'last_code_generated', '')
+        ex_code = getattr(agent, 'last_code_executed', '')
+        last_prompt = getattr(agent, 'last_prompt', prompt)
+        
+        # Log the interaction
+        log_interaction(
+            user_query=prompt,
+            model_name="pandas_ai_agent",
+            response_content=response,
+            generated_code=gen_code,
+            execution_time=execution_time,
+            error_message=None,
+            is_image=isinstance(response, str) and any(response.endswith(ext) for ext in ['.png', '.jpg', '.jpeg'])
+        )
+        
+        return {
+            "role": "assistant", 
+            "content": response, 
+            "gen_code": gen_code, 
+            "ex_code": ex_code, 
+            "last_prompt": last_prompt,
+            "error": None
+        }
+    except Exception as e:
+        execution_time = (datetime.now() - start_time).total_seconds()
+        error_msg = str(e)
+        
+        # Log the failed interaction
+        log_interaction(
+            user_query=prompt,
+            model_name="pandas_ai_agent",
+            response_content=f"Error: {error_msg}",
+            generated_code="",
+            execution_time=execution_time,
+            error_message=error_msg,
+            is_image=False
+        )
+        
+        return {
+            "role": "assistant", 
+            "content": f"Error: {error_msg}", 
+            "gen_code": "", 
+            "ex_code": "", 
+            "last_prompt": prompt,
+            "error": error_msg
+        }
+
+def decorate_with_code(response: dict) -> str:
+    """Decorate response with code details"""
+    gen_code = response.get("gen_code", "No code generated")
+    last_prompt = response.get("last_prompt", "No prompt")
+    
+    return f"""<details>
+<summary>Generated Code</summary>
+    
+```python
+{gen_code}
+```
+</details>
+
+<details>
+<summary>Prompt</summary>
+
+{last_prompt}
+"""
+
+def show_response(st, response):
+    """Display response with error handling"""
+    try:
+        with st.chat_message(response["role"]):
+            content = response.get("content", "No content")
+            
+            try:
+                # Try to open as image
+                image = Image.open(content)
+                if response.get("gen_code"):
+                    st.markdown(decorate_with_code(response), unsafe_allow_html=True)
+                st.image(image)
+                return {"is_image": True}
+            except:
+                # Not an image, display as text
+                if response.get("gen_code"):
+                    display_content = decorate_with_code(response) + f"""</details>
+
+{content}"""
+                else:
+                    display_content = content
+                st.markdown(display_content, unsafe_allow_html=True)
+                return {"is_image": False}
+    except Exception as e:
+        st.error(f"Error displaying response: {e}")
+        return {"is_image": False}
+
+def ask_question(model_name, question):
+    """Ask question with comprehensive error handling and logging"""
+    start_time = datetime.now()
+    try:
+        # Reload environment variables to get fresh values
+        load_dotenv(override=True)
+        fresh_groq_token = os.getenv("GROQ_API_KEY")
+        fresh_gemini_token = os.getenv("GEMINI_TOKEN")
+        
+        print(f"ask_question - Fresh Groq Token: {'Present' if fresh_groq_token else 'Missing'}")
+        
+        # Check API availability with fresh tokens
+        if model_name == "gemini-pro":
+            if not fresh_gemini_token or fresh_gemini_token.strip() == "":
+                execution_time = (datetime.now() - start_time).total_seconds()
+                error_msg = "Missing or empty API token"
+                
+                # Log the failed interaction
+                log_interaction(
+                    user_query=question,
+                    model_name=model_name,
+                    response_content="❌ Gemini API token not available or empty",
+                    generated_code="",
+                    execution_time=execution_time,
+                    error_message=error_msg,
+                    is_image=False
+                )
+                
+                return {
+                    "role": "assistant", 
+                    "content": "❌ Gemini API token not available or empty. Please set GEMINI_TOKEN in your environment variables.", 
+                    "gen_code": "", 
+                    "ex_code": "", 
+                    "last_prompt": question,
+                    "error": error_msg
+                }
+            llm = ChatGoogleGenerativeAI(
+                model=models[model_name], 
+                google_api_key=fresh_gemini_token,
+                temperature=0
+            )
+        else:
+            if not fresh_groq_token or fresh_groq_token.strip() == "":
+                execution_time = (datetime.now() - start_time).total_seconds()
+                error_msg = "Missing or empty API token"
+                
+                # Log the failed interaction
+                log_interaction(
+                    user_query=question,
+                    model_name=model_name,
+                    response_content="❌ Groq API token not available or empty",
+                    generated_code="",
+                    execution_time=execution_time,
+                    error_message=error_msg,
+                    is_image=False
+                )
+                
+                return {
+                    "role": "assistant", 
+                    "content": "❌ Groq API token not available or empty. Please set GROQ_API_KEY in your environment variables and restart the application.", 
+                    "gen_code": "", 
+                    "ex_code": "", 
+                    "last_prompt": question,
+                    "error": error_msg
+                }
+            
+            # Test the API key by trying to create the client
+            try:
+                llm = ChatGroq(
+                    model=models[model_name], 
+                    api_key=fresh_groq_token,
+                    temperature=0.1
+                )
+                # Test with a simple call to verify the API key works
+                test_response = llm.invoke("Test")
+                print("API key test successful")
+            except Exception as api_error:
+                execution_time = (datetime.now() - start_time).total_seconds()
+                error_msg = str(api_error)
+                
+                if "organization_restricted" in error_msg.lower() or "unauthorized" in error_msg.lower():
+                    response_content = "❌ API Key Error: Your Groq API key appears to be invalid, expired, or restricted. Please check your API key in the .env file."
+                    log_error_msg = f"API key validation failed: {error_msg}"
+                else:
+                    response_content = f"❌ API Connection Error: {error_msg}"
+                    log_error_msg = error_msg
+                
+                # Log the failed interaction
+                log_interaction(
+                    user_query=question,
+                    model_name=model_name,
+                    response_content=response_content,
+                    generated_code="",
+                    execution_time=execution_time,
+                    error_message=log_error_msg,
+                    is_image=False
+                )
+                
+                return {
+                    "role": "assistant", 
+                    "content": response_content, 
+                    "gen_code": "", 
+                    "ex_code": "", 
+                    "last_prompt": question,
+                    "error": log_error_msg
+                }
+
+        # Check if data file exists
+        if not os.path.exists("Data.csv"):
+            execution_time = (datetime.now() - start_time).total_seconds()
+            error_msg = "Data file not found"
+            
+            # Log the failed interaction
+            log_interaction(
+                user_query=question,
+                model_name=model_name,
+                response_content="❌ Data.csv file not found",
+                generated_code="",
+                execution_time=execution_time,
+                error_message=error_msg,
+                is_image=False
+            )
+            
+            return {
+                "role": "assistant", 
+                "content": "❌ Data.csv file not found. Please ensure the data file is in the correct location.", 
+                "gen_code": "", 
+                "ex_code": "", 
+                "last_prompt": question,
+                "error": error_msg
+            }
+
+        df_check = pd.read_csv("Data.csv")
+        df_check["Timestamp"] = pd.to_datetime(df_check["Timestamp"])
+        df_check = df_check.head(5)
+
+        new_line = "\n"
+        parameters = {"font.size": 12, "figure.dpi": 600}
+
+        template = f"""```python
+import pandas as pd
+import matplotlib.pyplot as plt
+import uuid
+
+plt.rcParams.update({parameters})
+
+df = pd.read_csv("Data.csv")
+df["Timestamp"] = pd.to_datetime(df["Timestamp"])
+
+# Available columns and data types:
+{new_line.join(map(lambda x: '# '+x, str(df_check.dtypes).split(new_line)))}
+
+# Question: {question.strip()}
+# Generate code to answer the question and save result in 'answer' variable
+# If creating a plot, save it with a unique filename and store the filename in 'answer'
+# If returning text/numbers, store the result directly in 'answer'
+```"""
+
+        system_prompt = """You are a helpful assistant that generates Python code for data analysis.
+        
+        Rules:
+        1. Always save your final result in a variable called 'answer'
+        2. If creating a plot, save it with plt.savefig() and store the filename in 'answer'
+        3. If returning text/numbers, store the result directly in 'answer'
+        4. Use descriptive variable names and add comments
+        5. Handle potential errors gracefully
+        6. For plots, use unique filenames to avoid conflicts
+        """
+
+        query = f"""{system_prompt}
+
+Complete the following code to answer the user's question:
+
+{template}
+"""
+        
+        # Make API call
+        if model_name == "gemini-pro":
+            response = llm.invoke(query)
+            answer = response.content
+        else:
+            response = llm.invoke(query)
+            answer = response.content
+        
+        # Extract and execute code
+        try:
+            if "```python" in answer:
+                code_part = answer.split("```python")[1].split("```")[0]
+            else:
+                code_part = answer
+            
+            full_code = f"""
+{template.split("```python")[1].split("```")[0]}
+{code_part}
+"""
+            
+            # Execute code in a controlled environment
+            local_vars = {}
+            global_vars = {
+                'pd': pd,
+                'plt': plt,
+                'os': os,
+                'uuid': __import__('uuid')
+            }
+            
+            exec(full_code, global_vars, local_vars)
+            
+            # Get the answer
+            if 'answer' in local_vars:
+                answer_result = local_vars['answer']
+            else:
+                answer_result = "No answer variable found in generated code"
+            
+            execution_time = (datetime.now() - start_time).total_seconds()
+            
+            # Determine if output is an image
+            is_image = isinstance(answer_result, str) and any(answer_result.endswith(ext) for ext in ['.png', '.jpg', '.jpeg'])
+            
+            # Log successful interaction
+            log_interaction(
+                user_query=question,
+                model_name=model_name,
+                response_content=str(answer_result),
+                generated_code=full_code,
+                execution_time=execution_time,
+                error_message=None,
+                is_image=is_image
+            )
+                
+            return {
+                "role": "assistant", 
+                "content": answer_result, 
+                "gen_code": full_code, 
+                "ex_code": full_code, 
+                "last_prompt": question,
+                "error": None
+            }
+            
+        except Exception as code_error:
+            execution_time = (datetime.now() - start_time).total_seconds()
+            error_msg = str(code_error)
+            
+            # Log the failed code execution
+            log_interaction(
+                user_query=question,
+                model_name=model_name,
+                response_content=f"❌ Error executing generated code: {error_msg}",
+                generated_code=full_code if 'full_code' in locals() else "",
+                execution_time=execution_time,
+                error_message=error_msg,
+                is_image=False
+            )
+            
+            return {
+                "role": "assistant", 
+                "content": f"❌ Error executing generated code: {error_msg}", 
+                "gen_code": full_code if 'full_code' in locals() else "", 
+                "ex_code": full_code if 'full_code' in locals() else "", 
+                "last_prompt": question,
+                "error": error_msg
+            }
+            
+    except Exception as e:
+        execution_time = (datetime.now() - start_time).total_seconds()
+        error_msg = str(e)
+        
+        # Handle specific API errors
+        if "organization_restricted" in error_msg:
+            response_content = "❌ API Organization Restricted: Your API key access has been restricted. Please check your Groq API key or try generating a new one."
+            log_error_msg = "API access restricted"
+        elif "rate_limit" in error_msg.lower():
+            response_content = "❌ Rate limit exceeded. Please wait a moment and try again."
+            log_error_msg = "Rate limit exceeded"
+        else:
+            response_content = f"❌ Error: {error_msg}"
+            log_error_msg = error_msg
+        
+        # Log the failed interaction
+        log_interaction(
+            user_query=question,
+            model_name=model_name,
+            response_content=response_content,
+            generated_code="",
+            execution_time=execution_time,
+            error_message=log_error_msg,
+            is_image=False
+        )
+        
+        return {
+            "role": "assistant", 
+            "content": response_content, 
+            "gen_code": "", 
+            "ex_code": "", 
+            "last_prompt": question,
+            "error": log_error_msg
+        }