Updated data analysis to generate images to perform data analysis

llm_be/chat_backend/services/data_analysis_service.py

@@ -1,10 +1,15 @@
 import pandas as pd
 import io
+import re
+import json
+import base64
+import matplotlib.pyplot as plt
 from typing import AsyncGenerator
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_ollama import OllamaLLM
 from langchain_core.output_parsers import StrOutputParser
 
+
 class AsyncDataAnalysisService:
     """Asynchronous service for performing data analysis with an LLM."""
 
@@ -20,8 +25,11 @@ class AsyncDataAnalysisService:
 
     def _setup_chain(self):
         """Set up the LLM chain with a prompt tailored for data analysis."""
-        template = """You are an expert data analyst. A user has provided a summary and sample of a dataset and is asking a question about it.
-Analyze the provided information and answer the user's question. If a calculation is requested, perform it based on the summary statistics provided. If the data is not suitable for the request, explain why.
+        template = """You are an expert data analyst. Your role is to directly answer a user's question about a dataset they have provided.
+You will be given a summary and a sample of the dataset.
+Based on this information, provide a clear and concise answer to the user's question.
+Do not provide Python code or any other code. The user is not a developer and wants a direct answer.
+Even if you don't think the data provides enough evidence for the query, still provide a response
 
 ---
 Data Summary:
@@ -69,35 +77,87 @@ Answer:"""
 
         return "\n".join(summary_lines)
 
+    def _generate_plot(self, query: str, df: pd.DataFrame) -> str:
+        """
+        Generates a plot from a DataFrame based on a natural language query,
+        encodes it in Base64, and returns it.
+
+        If columns are specified (e.g., "plot X vs Y"), it uses them.
+        If not, it automatically picks the first two numerical columns.
+        """
+        col1, col2 = None, None
+        title = "Scatter Plot"
+
+        # Attempt to find explicitly mentioned columns, e.g., "plot Column1 vs Column2"
+        match = re.search(r"(?:plot|scatter|visualize)\s+(.*?)\s+(?:vs|versus|and)\s+(.*)", query, re.IGNORECASE)
+        if match:
+            potential_col1 = match.group(1).strip()
+            potential_col2 = match.group(2).strip()
+            if potential_col1 in df.columns and potential_col2 in df.columns:
+                col1, col2 = potential_col1, potential_col2
+                title = f"Scatterplot of {col1} vs {col2}"
+
+        # If no valid columns were explicitly found, auto-detect
+        if not col1 or not col2:
+            numeric_cols = df.select_dtypes(include=['number']).columns.tolist()
+            if len(numeric_cols) >= 2:
+                col1, col2 = numeric_cols[0], numeric_cols[1]
+                title = f"Scatterplot of {col1} vs {col2} (Auto-selected)"
+            else:
+                raise ValueError("I couldn't find two numerical columns to plot automatically. Please specify columns, like 'plot column_A vs column_B'.")
+
+        fig, ax = plt.subplots()
+        ax.scatter(df[col1], df[col2])
+        ax.set_xlabel(col1)
+        ax.set_ylabel(col2)
+        ax.set_title(title)
+        ax.grid(True)
+
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', bbox_inches='tight')
+        plt.close(fig)
+        buf.seek(0)
+
+        image_base64 = base64.b64encode(buf.read()).decode('utf-8')
+        return image_base64
+
     async def generate_response(
         self,
         query: str,
         decoded_file: bytes,
         file_type: str,
     ) -> AsyncGenerator[str, None]:
-        """Generate a response based on the uploaded data and user query."""
-
+        """
+        Generate a response based on the uploaded data and user query.
+        This can be a text analysis or a plot visualization.
+        """
         try:
-            # Read the file content into a DataFrame
             if file_type == "csv":
                 df = pd.read_csv(io.BytesIO(decoded_file))
             elif file_type == "xlsx":
                 df = pd.read_excel(io.BytesIO(decoded_file))
             else:
-                yield "I can only analyze CSV and XLSX files at this time."
+                yield json.dumps({"type": "error", "content": "I can only analyze CSV and XLSX files."})
                 return
 
-            # Get the structured summary instead of the full data
-            data_summary = self._get_dataframe_summary(df)
+            plot_keywords = ["plot", "graph", "scatter", "visualize"]
+            if any(keyword in query.lower() for keyword in plot_keywords):
+                try:
+                    image_base64 = self._generate_plot(query, df)
+                    yield json.dumps({
+                        "type": "plot",
+                        "format": "png",
+                        "image": image_base64
+                    })
+                except ValueError as e:
+                    yield json.dumps({"type": "error", "content": str(e)})
+                return
 
-            # Prepare the input for the LLM chain
-            chain_input = {
-                "data_summary": data_summary,
-                "query": query,
-            }
+            data_summary = self._get_dataframe_summary(df)
+            chain_input = {"data_summary": data_summary, "query": query}
 
             async for chunk in self.analysis_chain.astream(chain_input):
-                yield chunk
+                yield chunk #json.dumps({"type": "text", "content": chunk})
 
         except Exception as e:
-            yield f"An error occurred while processing the file: {e}"
+            yield json.dumps({"type": "error", "content": f"An error occurred: {e}"})

llm_be/chat_backend/services/moderation_classifier.py

@@ -29,7 +29,7 @@ class ModerationClassifier(BaseService):
                 (
                     "system",
                     """You are a strict content moderator. Classify the following prompt as either NSFW or FINE.
-             
+
 NSFW includes:
 - Sexual content
 - Violence/gore
@@ -50,6 +50,7 @@ Examples:
 - "Write a love poem" → FINE
 - "Explicit sex scene" → NSFW
 - "Python tutorial" → FINE
+- "Please analyze this file and project the next 12 months for me. Add a graph visual of the data as well" → FINE
 
 Return ONLY "NSFW" or "FINE", nothing else.""",
                 ),