updates

llm_be/chat_backend/services/data_analysis_service.py

@@ -8,6 +8,8 @@ from typing import AsyncGenerator
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_ollama import OllamaLLM
 from langchain_core.output_parsers import StrOutputParser
+import docx
+import pypdf
 
 
 class AsyncDataAnalysisService:
@@ -25,14 +27,14 @@ 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. 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.
+        template = """You are an expert data analyst. Your role is to directly answer a user's question about a dataset or document they have provided.
+You will be given a summary and a sample of the dataset, or the content of the document.
 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:
+Data/Document Content:
 {data_summary}
 ---
 
@@ -77,6 +79,22 @@ Answer:"""
 
         return "\n".join(summary_lines)
 
+    def _read_docx(self, file_bytes: bytes) -> str:
+        """Reads text from a DOCX file."""
+        doc = docx.Document(io.BytesIO(file_bytes))
+        full_text = []
+        for para in doc.paragraphs:
+            full_text.append(para.text)
+        return "\n".join(full_text)
+
+    def _read_pdf(self, file_bytes: bytes) -> str:
+        """Reads text from a PDF file."""
+        pdf_reader = pypdf.PdfReader(io.BytesIO(file_bytes))
+        full_text = []
+        for page in pdf_reader.pages:
+            full_text.append(page.extract_text())
+        return "\n".join(full_text)
+
     def _generate_plot(self, query: str, df: pd.DataFrame) -> str:
         """
         Generates a plot from a DataFrame based on a natural language query,
@@ -132,28 +150,40 @@ Answer:"""
         This can be a text analysis or a plot visualization.
         """
         try:
-            if file_type == "csv":
+            df = None
+            data_summary = ""
+            file_type = file_type.lower()
+            print(file_type)
+
+            if "csv" in file_type:
                 df = pd.read_csv(io.BytesIO(decoded_file))
-            elif file_type == "xlsx":
+                data_summary = self._get_dataframe_summary(df)
+            elif "xlsx" in file_type or "spreadsheet" in file_type:
                 df = pd.read_excel(io.BytesIO(decoded_file))
+                data_summary = self._get_dataframe_summary(df)
+            elif "word" in file_type or "docx" in file_type:
+                data_summary = self._read_docx(decoded_file)
+            elif "pdf" in file_type:
+                data_summary = self._read_pdf(decoded_file)
             else:
-                yield json.dumps({"type": "error", "content": "I can only analyze CSV and XLSX files."})
+                yield json.dumps({"type": "error", "content": f"Unsupported file type: {file_type}. I can analyze CSV, XLSX, DOCX, and PDF files."})
                 return
 
-            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
+            # Only attempt plotting if we have a DataFrame
+            if df is not None:
+                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
 
-            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):

llm_be/chat_backend/services/prompt_classifier/prompt_classifier.py

@@ -10,6 +10,7 @@ class PromptType(Enum):
     RAG = auto()
     IMAGE_GENERATION = auto()
     DATA_ANALYSIS = auto()
+    SEARCH = auto()
     UNKNOWN = auto()
 
 
@@ -36,8 +37,9 @@ 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. 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
+4. DATA_ANALYSIS - When a user wants to read an uploaded document (PDF, Word, etc.) and generate an index, summary, or extract structured information. Includes prompts like "Please read this document and make me an index for it".
+5. SEARCH - When the user is seeking specific, up-to-date information (e.g., current events, celebrity news, sports scores).
+6. UNKNOWN - If none of the above fit
 
 1. IMAGE_GENERATION - ONLY if:
     - Explicitly contains: "generate/create/draw/make an image/picture/photo/art/illustration"
@@ -50,15 +52,22 @@ class PromptClassifier(BaseService):
     - Example: "What does contracts.pdf say?" → RAG
 
 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
+    - The user provides or references an uploaded document (PDF, Word, etc.) and asks for an index, summary, extraction, or analysis of its contents.
+    - Example: "Please read this document and make me an index for it" → DATA_ANALYSIS
+    - Example: "Here is the file content. What is the sum of all 'Sales'?" → DATA_ANALYSIS
 
-4. GENERAL_CHAT - DEFAULT category when:
+5. SEARCH - ONLY if:
+    - User asks for current information (news, weather, sports, stock prices)
+    - User asks for specific facts that might change or require lookup (e.g. "Who won the 2024 election?")
+    - Example: "What is the latest news on X?" → SEARCH
+    - Example: "Who won the Super Bowl this year?" → SEARCH
+
+6. GENERAL_CHAT - DEFAULT category when:
     - Doesn't meet above criteria
-    - Conversational/general knowledge
-    - Uncertain cases
+    - Conversational/general knowledge (that doesn't require live search)
+    - Creative writing (poems, jokes)
     - Example: "Tell me a joke" → GENERAL_CHAT
+    - Example: "Write a poem about cats" → GENERAL_CHAT
 
 Examples:
 [Definitely RAG]
@@ -66,9 +75,15 @@ Examples:
 - "Search our documents for the 2023 marketing strategy"
 
 [Definitely DATA_ANALYSIS]
+- "Please read this document and make me an index for it"
 - "Here is the file content. What is the sum of all 'Sales'?"
 - "Based on this CSV data, show me the top 5 customers."
 
+[Definitely SEARCH]
+- "Who won the 2024 presidential race?"
+- "What is the latest celebrity news?"
+- "Current stock price of Apple"
+
 [Definitely GENERAL_CHAT]
 - "How does photosynthesis work?" (General knowledge)
 - "Tell me a joke"
@@ -77,7 +92,7 @@ Examples:
 [Borderline -> GENERAL_CHAT]
 - "What's our company policy on X?" (No doc reference -> general)
 
-Return ONLY the label, no explanations.""",
+Return ONLY the exact Enum label (e.g. "GENERAL_CHAT"), no explanations."""
                 ),
                 ("human", "{prompt}"),
             ]
@@ -85,8 +100,40 @@ Return ONLY the label, no explanations.""",
 
         self.chain = self.classification_prompt | self.llm
 
+    def _quick_check(self, prompt: str) -> PromptType | None:
+        """
+        Performs a quick, rule-based classification before involving the LLM.
+        Returns a PromptType if a clear match is found, otherwise None.
+        """
+        lower_prompt = prompt.lower()
+
+        # IMAGE_GENERATION
+        if any(keyword in lower_prompt for keyword in ["generate image", "create picture", "draw an image", "make a photo", "generate an image", "create an illustration"]):
+            return PromptType.IMAGE_GENERATION
+        
+        # DATA_ANALYSIS (often involves uploaded documents)
+        if any(keyword in lower_prompt for keyword in ["read this document", "analyze this file", "summarize this pdf", "extract data from", "index this document", "based on this csv", "from this spreadsheet"]):
+            return PromptType.DATA_ANALYSIS
+
+        # RAG (explicitly asking to search within provided context/documents)
+        # This might overlap with DATA_ANALYSIS, but RAG is more about retrieval from a knowledge base.
+        # The prompt examples for RAG are "What does the uploaded PDF say about quarterly results?"
+        # "Search our documents for the 2023 marketing strategy"
+        if ("uploaded pdf" in lower_prompt or "our documents" in lower_prompt or "this document" in lower_prompt or "the document" in lower_prompt) and \
+           any(keyword in lower_prompt for keyword in ["say about", "search for", "find in", "lookup in"]):
+            return PromptType.RAG
+
+        # SEARCH
+        if any(keyword in lower_prompt for keyword in ["latest news", "current weather", "stock price", "who won", "what is the current", "breaking news", "real-time information", "up-to-date"]):
+            return PromptType.SEARCH
+
+        return None
+
     async def classify_async(self, prompt: str) -> PromptType:
         """Asynchronously classify the prompt"""
+        quick = self._quick_check(prompt)
+        if quick:
+            return quick
         try:
             response = await self.chain.ainvoke({"prompt": prompt})
             return self._parse_response(response.strip())
@@ -96,6 +143,9 @@ Return ONLY the label, no explanations.""",
 
     def classify(self, prompt: str) -> PromptType:
         """Synchronously classify the prompt"""
+        quick = self._quick_check(prompt)
+        if quick:
+            return quick
         try:
             response = self.chain.invoke({"prompt": prompt})
             return self._parse_response(response.strip())
@@ -105,10 +155,26 @@ Return ONLY the label, no explanations.""",
 
     def _parse_response(self, response: str) -> PromptType:
         """Convert string response to PromptType enum"""
-        response = response.upper()
+        response = response.upper().strip()
+        print(response)
+        
+        # Direct match
+        try:
+            return PromptType[response]
+        except KeyError:
+            pass
+            
+        # Handle missing underscores (e.g. GENERALCHAT)
+        normalized_response = response.replace("_", "")
+        for prompt_type in PromptType:
+            if prompt_type.name.replace("_", "") == normalized_response:
+                return prompt_type
+                
+        # Substring match as fallback
         for prompt_type in PromptType:
             if prompt_type.name in response:
                 return prompt_type
+                
         return PromptType.UNKNOWN