Update openai_agent.py

openai_agent.py

@@ -1,3 +1,4 @@
+import streamlit as st
 from openai import OpenAI
 
 class OpenAIAgent:
@@ -5,27 +6,46 @@ class OpenAIAgent:
         self.base_url = "https://api.aimlapi.com/v1"
         self.api_key = "c496d9094ba54ddb9d66eeeb35a6196f"  # Replace with your actual API key
         self.api = OpenAI(api_key=self.api_key, base_url=self.base_url)
-        self.system_prompt = "You are a data analyst assistant. Provide insights based on the given data."
+        self.system_prompt = "You are an AI assistant specialized in data analysis, visualization, and insights generation. Provide clear and concise responses."
 
     def get_insights(self, df, prompt):
-        # Prepare the data for the API call
-        data_sample = df.head(5).to_string()
-        
-        user_prompt = f"Here's a sample of my data:\n{data_sample}\n\nBased on this data, {prompt}"
-        
-        completion = self.api.chat.completions.create(
-            model="gpt-4o",
-            messages=[
-                {"role": "system", "content": self.system_prompt},
-                {"role": "user", "content": user_prompt},
-            ],
-            temperature=0.7,
-            max_tokens=256,
-        )
+        try:
+            data_sample = df.head(5).to_string()
+            data_info = f"Dataset shape: {df.shape}\nColumns: {', '.join(df.columns)}\n"
+            full_prompt = f"{data_info}\n{data_sample}\n\n{prompt}"
+            
+            completion = self.api.chat.completions.create(
+                model="mistralai/Mistral-7B-Instruct-v0.2",
+                messages=[
+                    {"role": "system", "content": self.system_prompt},
+                    {"role": "user", "content": full_prompt},
+                ],
+                temperature=0.7,
+                max_tokens=500,
+            )
+            
+            return completion.choices[0].message.content
+        except Exception as e:
+            return f"An error occurred while getting insights: {str(e)}"
+
+    def generate_insights(self, df):
+        insight_type = st.selectbox("Select insight type", 
+                                    ["General Insights", "Trend Analysis", "Anomaly Detection", "Predictive Analysis"])
         
-        response = completion.choices[0].message.content
-        return response
+        if st.button("Generate Insights"):
+            with st.spinner("Generating insights..."):
+                prompt = f"Analyze the following dataset and provide {insight_type}."
+                ai_response = self.get_insights(df, prompt)
+                st.write(ai_response)
 
+    def custom_ai_task(self, df):
+        custom_task = st.text_area("Describe your custom task:")
+        
+        if st.button("Execute Custom Task"):
+            with st.spinner("Processing custom task..."):
+                prompt = f"Perform the following task on the given dataset:\n{custom_task}"
+                ai_response = self.get_insights(df, prompt)
+                st.write(ai_response)
 '''import openai
 
 class OpenAIAgent: