Update openai_agent.py

openai_agent.py

@@ -1,4 +1,32 @@
-import openai
+from openai import OpenAI
+
+class OpenAIAgent:
+    def __init__(self):
+        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."
+
+    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,
+        )
+        
+        response = completion.choices[0].message.content
+        return response
+
+'''import openai
 
 class OpenAIAgent:
     def __init__(self):
@@ -20,4 +48,5 @@ class OpenAIAgent:
             max_tokens=150
         )
         
-        return response.choices[0].message['content']
+        return response.choices[0].message['content']
+'''