Allow for data analysis

2025-09-08 12:29:20 -05:00
parent 951a58f2fa
commit 14d8211715
4 changed files with 236 additions and 83 deletions
@@ -0,0 +1,103 @@
+import pandas as pd
+import io
+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."""
+
+    def __init__(self):
+        # A model with a large context window and strong analytical skills is best
+        self.llm = OllamaLLM(
+            model="llama3.2",
+            temperature=0.3,
+            num_ctx=8192,
+        )
+        self.output_parser = StrOutputParser()
+        self._setup_chain()
+
+    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.
+
+---
+Data Summary:
+{data_summary}
+---
+
+User's Question: {query}
+Answer:"""
+
+        self.prompt = ChatPromptTemplate.from_template(template)
+
+        self.analysis_chain = (
+            {
+                "data_summary": lambda x: x["data_summary"],
+                "query": lambda x: x["query"],
+            }
+            | self.prompt
+            | self.llm
+            | self.output_parser
+        )
+
+    def _get_dataframe_summary(self, df: pd.DataFrame) -> str:
+        """Generates a structured summary of the DataFrame for the LLM."""
+
+        num_rows, num_cols = df.shape
+        summary_lines = [
+            f"DataFrame has {num_rows} rows and {num_cols} columns.",
+            "Column Information (Name, Dtype, Non-Null Count):",
+            "--------------------------------------------------",
+        ]
+
+        # Add a concise summary using df.info()
+        info_buffer = io.StringIO()
+        df.info(buf=info_buffer, verbose=True, show_counts=True)
+        summary_lines.append(info_buffer.getvalue())
+
+        summary_lines.append("\nDescriptive Statistics (for numerical columns):")
+        summary_lines.append("--------------------------------------------")
+        summary_lines.append(df.describe().to_string())
+
+        summary_lines.append("\nSample of Data:")
+        summary_lines.append("-----------------")
+        # Show the first 5 rows and a few random rows to give a feel for the data
+        summary_lines.append(df.head(5).to_string())
+
+        return "\n".join(summary_lines)
+
+    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."""
+
+        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."
+                return
+
+            # Get the structured summary instead of the full data
+            data_summary = self._get_dataframe_summary(df)
+
+            # Prepare the input for the LLM chain
+            chain_input = {
+                "data_summary": data_summary,
+                "query": query,
+            }
+
+            async for chunk in self.analysis_chain.astream(chain_input):
+                yield chunk
+
+        except Exception as e:
+            yield f"An error occurred while processing the file: {e}"
@@ -9,6 +9,7 @@ class PromptType(Enum):
    GENERAL_CHAT = auto()
    RAG = auto()
    IMAGE_GENERATION = auto()
+    DATA_ANALYSIS = auto()
    UNKNOWN = auto()


@@ -35,43 +36,45 @@ class PromptClassifier(BaseService):
 1. GENERAL_CHAT - Casual conversation, personal questions, or non-specific inquiries
 2. RAG - ONLY when explicitly requesting document/search-based knowledge
 3. IMAGE_GENERATION - Specific requests to create/modify images
-4. UNKNOWN - If none of the above fit
+4. DATA_ANALYSIS - When a user is asking questions about an uploaded spreadsheet or CSV file. The user's message contains the data from the file.
+5. UNKNOWN - If none of the above fit

 1. IMAGE_GENERATION - ONLY if:
-   - Explicitly contains: "generate/create/draw/make an image/picture/photo/art/illustration"
-   - Requests visual content creation
-   - Example: "Make a picture of a castle" → IMAGE_GENERATION
+    - Explicitly contains: "generate/create/draw/make an image/picture/photo/art/illustration"
+    - Requests visual content creation
+    - Example: "Make a picture of a castle" → IMAGE_GENERATION

 2. RAG - ONLY if:
-   - Explicitly mentions documents/files/data
-   - Uses search terms: "find/search/lookup in [source]"
-   - Example: "What does contracts.pdf say?" → RAG
+    - Explicitly mentions documents/files/data
+    - Uses search terms: "find/search/lookup in [source]"
+    - Example: "What does contracts.pdf say?" → RAG

-3. GENERAL_CHAT - DEFAULT category when:
-   - Doesn't meet above criteria
-   - Conversational/general knowledge
-   - Uncertain cases
-   - Example: "Tell me a joke" → GENERAL_CHAT
+3. DATA_ANALYSIS - ONLY if:
+    - The message explicitly contains structured data from a file (e.g., a DataFrame string)
+    - The user is asking to analyze, summarize, or plot the data
+    - Example: "Here is the sales data. What is the average revenue per product?" -> DATA_ANALYSIS
+
+4. GENERAL_CHAT - DEFAULT category when:
+    - Doesn't meet above criteria
+    - Conversational/general knowledge
+    - Uncertain cases
+    - Example: "Tell me a joke" → GENERAL_CHAT

 Examples:
 [Definitely RAG]
 - "What does the uploaded PDF say about quarterly results?"
 - "Search our documents for the 2023 marketing strategy"
- "Find the contract clause about termination"
+
+[Definitely DATA_ANALYSIS]
+- "Here is the file content. What is the sum of all 'Sales'?"
+- "Based on this CSV data, show me the top 5 customers."

 [Definitely GENERAL_CHAT]
 - "How does photosynthesis work?" (General knowledge)
 - "Tell me a joke"
- "What's your opinion on AI?"

-[Borderline → GENERAL_CHAT]
- "What's our company policy on X?" (No doc reference → general)
- "Explain quantum computing" (General knowledge)
- "Summarize the meeting" (No doc reference)
-
-[Definitely NOT IMAGE_GENERATION]
- "Great, can you make it about a duck now"
- "highlight the features of the backyard playset if they were to choose us and make the language more long form"
+[Borderline -> GENERAL_CHAT]
+- "What's our company policy on X?" (No doc reference -> general)

 Return ONLY the label, no explanations.""",
                ),