DB Chat for pro users only

orchestrator_functions.py

@@ -3,6 +3,7 @@ import asyncio
 import logging
 import os
 import threading
+from typing import Dict
 import uuid
 from fastapi.encoders import jsonable_encoder
 import numpy as np
@@ -344,164 +345,19 @@ def langchain_csv_chart(csv_url: str, question: str, chart_required: bool):
 
 
 
-###########################################################################################################################
+####################################### Orchestrator Function for Chat #######################################
 
-
-
-
-# async def csv_chart(csv_url: str, query: str):
-#     """
-#     Generate a chart based on the provided CSV URL and query.
-#     Parameters:
-#     - csv_url (str): The URL of the CSV file.
-#     - query (str): The query for generating the chart.
-#     Returns:
-#     - dict: A dictionary containing the generated chart image URL.
-#     Example:
-#     - csv_url: "https://example.com/data.csv"
-#     - query: "Generate a bar chart showing sales by region."
-#     Returns:
-#     - dict: {"image_url": "https://example.com/chart.png"}.
-
-#     """
-    
-#     try:
-#         # First try Groq-based chart generation
-#         try:
-#             groq_result = await asyncio.to_thread(groq_chart, csv_url, query)
-#             logger.info(f"Generated Chart (Groq): {groq_result}")
-            
-#             if groq_result != 'Chart not generated':
-#                 unique_file_name = f'{str(uuid.uuid4())}.png'
-#                 image_public_url = await upload_file_to_supabase(groq_result, unique_file_name)
-#                 logger.info(f"Image uploaded to Supabase: {image_public_url}")
-#                 return {"image_url": image_public_url}
-                
-#         except Exception as groq_error:
-#             logger.info(f"Groq chart generation failed, falling back to Langchain: {str(groq_error)}")
-
-#         # Fallback to Langchain if Groq fails
-#         try:
-#             langchain_paths = await asyncio.to_thread(langchain_csv_chart, csv_url, query, True)
-#             logger.info("Fallback langchain chart result:", langchain_paths)
-            
-#             if isinstance(langchain_paths, list) and len(langchain_paths) > 0:
-#                 unique_file_name = f'{str(uuid.uuid4())}.png'
-#                 logger.info("Uploading the chart to supabase...")
-#                 image_public_url = await upload_file_to_supabase(langchain_paths[0], unique_file_name)
-#                 logger.info("Image uploaded to Supabase and Image URL is... ", image_public_url)
-#                 return {"image_url": image_public_url}
-                
-#         except Exception as langchain_error:
-#             logger.info(f"Langchain chart generation also failed: {str(langchain_error)}")
-#             try:
-#                 # Last resort: Try with the gemini langchain agent
-#                 logger.info("Trying with the gemini langchain agent...")
-#                 lc_gemini_chart_result = await asyncio.to_thread(langchain_gemini_csv_handler, csv_url, query, True)
-#                 if lc_gemini_chart_result is not None:
-#                     clean_path = lc_gemini_chart_result.strip()
-#                     unique_file_name = f'{str(uuid.uuid4())}.png'
-#                     logger.info("Uploading the chart to supabase...")
-#                     image_public_url = await upload_file_to_supabase(clean_path, unique_file_name)
-#                     logger.info("Image uploaded to Supabase and Image URL is... ", image_public_url)
-#                     return {"image_url": image_public_url}
-#             except Exception as gemini_error:
-#                 logger.info(f"Gemini Langchain chart generation also failed: {str(gemini_error)}")
-
-#         # If both methods fail
-#         return {"error": "Could not generate the chart, please try again."}
-
-#     except Exception as e:
-#         logger.info(f"Critical chart error: {str(e)}")
-#         return {"error": "Internal system error"}
-    
-    
-    
-
-
-
-
-
-# async def csv_chat(csv_url: str, query: str):
-#     """
-#     Generate a response based on the provided CSV URL and query.
-#     Parameters:
-#     - csv_url (str): The URL of the CSV file.
-#     - query (str): The query for generating the response.
-#     Returns:
-#     - dict: A dictionary containing the generated response.
-#     Example:
-#     - csv_url: "https://example.com/data.csv"
-#     - query: "What is the total sales for the year 2022?"
-#     Returns:
-#     - dict: {"answer": "The total sales for 2022 is $100,000."}.
-#     """
-#     try:
-#         updated_query = f"{query} and Do not show any charts or graphs."
-        
-#         # Process with Groq first
-#         try:
-#             groq_answer = await asyncio.to_thread(groq_chat, csv_url, updated_query)
-#             logger.info("groq_answer:", groq_answer)
-            
-#             if process_answer(groq_answer) == "Empty response received." or groq_answer == None:
-#                 return {"answer": "Sorry, I couldn't find relevant data..."}
-            
-#             if process_answer(groq_answer) or groq_answer == None:
-#                 raise Exception("Groq response not usable, falling back to LangChain")
-            
-#             return {"answer": jsonable_encoder(groq_answer)}
-            
-#         except Exception as groq_error:
-#             logger.info(f"Groq error, falling back to LangChain: {str(groq_error)}")
-            
-#             # Process with LangChain if Groq fails
-#             try:
-#                 lang_answer = await asyncio.to_thread(
-#                     langchain_csv_chat, csv_url, query, False
-#                 )
-#                 if not process_answer(lang_answer):
-#                     return {"answer": jsonable_encoder(lang_answer)}
-#                 return {"answer": "Sorry, I couldn't find relevant data..."}
-#             except Exception as langchain_error:
-#                 logger.info(f"LangChain processing error: {str(langchain_error)}")
-                
-#                 # last resort: Try with the gemini langchain agent
-#                 try:
-#                     gemini_answer = await asyncio.to_thread(
-#                         langchain_gemini_csv_handler, csv_url, query, False
-#                     )
-#                     if not process_answer(gemini_answer):
-#                         return {"answer": jsonable_encoder(gemini_answer)}
-#                     return {"answer": "Sorry, I couldn't find relevant data..."}
-#                 except Exception as gemini_error:
-#                     logger.info(f"Gemini Langchain processing error: {str(gemini_error)}")
-#                     return {"answer": "error"}
-                
-#     except Exception as e:
-#         logger.info(f"Error processing request: {str(e)}")
-#         return {"answer": "error"}
-
-
-
-
-
-
-
-####################################### Start with lc_gemini #######################################
-
-
-async def csv_chat(csv_url: str, query: str):
+async def csv_chat(csv_url: str, query: str) -> dict:
     """
     Generate a response based on the provided CSV URL and query.
-    Prioritizes LangChain-Groq, then raw Groq, and finally LangChain-Gemini as fallback.
+    Prioritizes LangChain-Gemini first, then falls back to LangChain-Groq.
     
     Parameters:
     - csv_url (str): The URL of the CSV file.
     - query (str): The query for generating the response.
     
     Returns:
-    - dict: A dictionary containing the generated response.
+    - dict: A dictionary containing the generated response or error message.
     
     Example:
     - csv_url: "https://example.com/data.csv"
@@ -509,159 +365,136 @@ async def csv_chat(csv_url: str, query: str):
     Returns:
     - dict: {"answer": "The total sales for 2022 is $100,000."}
     """
+    updated_query = f"{query} and Do not show any charts or graphs."
+    error_messages = []
+    
+    # --- 1. First Attempt: LangChain Gemini ---
     try:
-        updated_query = f"{query} and Do not show any charts or graphs."
+        gemini_answer = await asyncio.to_thread(
+            langchain_gemini_csv_handler, csv_url, updated_query, False
+        )
+        logger.info(f"LangChain-Gemini answer: {gemini_answer}")
         
-        # --- 1. First Attempt: LangChain Groq ---
-        try:
-            lang_groq_answer = await asyncio.to_thread(
-                langchain_csv_chat, csv_url, updated_query, False
-            )
-            logger.info("LangChain-Groq answer:", lang_groq_answer)
-            
-            if lang_groq_answer is not None:
-                return {"answer": jsonable_encoder(lang_groq_answer)}
-            
-            raise Exception("LangChain-Groq response not usable, falling back to raw Groq")
-            
-        except Exception as lang_groq_error:
-            logger.info(f"LangChain-Groq error: {str(lang_groq_error)}")
-            
-            # --- 2. Second Attempt: Raw Groq Chat ---
-            try:
-                raw_groq_answer = await asyncio.to_thread(groq_chat, csv_url, updated_query)
-                logger.info("Raw Groq answer:", raw_groq_answer)
-                
-                if process_answer(raw_groq_answer) == "Empty response received." or raw_groq_answer is None:
-                    raise Exception("Raw Groq response not usable, falling back to LangChain-Gemini")
-                
-                if process_answer(raw_groq_answer):
-                    raise Exception("Raw Groq response not usable, falling back to LangChain-Gemini")
-                
-                return {"answer": jsonable_encoder(raw_groq_answer)}
-                
-            except Exception as raw_groq_error:
-                logger.info(f"Raw Groq error: {str(raw_groq_error)}")
-                
-                # --- 3. Final Attempt: LangChain Gemini ---
-                try:
-                    gemini_answer = await asyncio.to_thread(
-                        langchain_gemini_csv_handler, csv_url, updated_query, False
-                    )
-                    logger.info("LangChain-Gemini answer:", gemini_answer)
-                    
-                    if gemini_answer is not None:
-                        return {"answer": jsonable_encoder(gemini_answer)}
-                    
-                    raise Exception("All fallbacks exhausted")
-                    
-                except Exception as gemini_error:
-                    logger.info(f"LangChain-Gemini error: {str(gemini_error)}")
-                    return {"answer": "Sorry, I couldn't find relevant data..."}
-                
-    except Exception as e:
-        logger.info(f"Unexpected error: {str(e)}")
-        return {"answer": "error"}
+        if gemini_answer and is_valid_response(gemini_answer):
+            return {"answer": jsonable_encoder(gemini_answer)}
+        
+        error_messages.append("LangChain-Gemini response not usable")
+    except Exception as gemini_error:
+        error_messages.append(f"LangChain-Gemini error: {str(gemini_error)}")
+        logger.error(f"LangChain-Gemini failed: {str(gemini_error)}")
+    
+    # --- 2. Fallback Attempt: LangChain Groq ---
+    try:
+        lang_groq_answer = await asyncio.to_thread(
+            langchain_csv_chat, csv_url, updated_query, False
+        )
+        logger.info(f"LangChain-Groq answer: {lang_groq_answer}")
+        
+        if lang_groq_answer and is_valid_response(lang_groq_answer):
+            return {"answer": jsonable_encoder(lang_groq_answer)}
+        
+        error_messages.append("LangChain-Groq response not usable")
+    except Exception as lang_groq_error:
+        error_messages.append(f"LangChain-Groq error: {str(lang_groq_error)}")
+        logger.error(f"LangChain-Groq failed: {str(lang_groq_error)}")
+    
+    # --- Final Fallback when all attempts fail ---
+    logger.error(f"All attempts failed. Errors: {'; '.join(error_messages)}")
+    return {"answer": "Sorry, I couldn't process your request with the available data."}
 
+def is_valid_response(response) -> bool:
+    """Check if the response is valid and not empty."""
+    if not response:
+        return False
+    if isinstance(response, str) and response.strip() == "":
+        return False
+    if isinstance(response, dict) and not response.get("answer"):
+        return False
+    return True
 
 
 
 
+####################################### Orchestrator Function for Chart #######################################
 
 
-async def csv_chart(csv_url: str, query: str, chat_id: str):
+async def csv_chart(csv_url: str, query: str, chat_id: str) -> Dict[str, str]:
     """
     Generate a chart based on the provided CSV URL and query.
-    Prioritizes PandasAI Groq, then LangChain Gemini, and finally LangChain Groq as fallback.
+    Prioritizes LangChain-Gemini first, then falls back to LangChain-Groq.
     
     Parameters:
-    - csv_url (str): The URL of the CSV file.
-    - query (str): The query for generating the chart.
+    - csv_url (str): URL of the CSV file
+    - query (str): Query for generating the chart
+    - chat_id (str): Chat session ID for file storage
     
     Returns:
-    - dict: A dictionary containing either:
-        - {"image_url": "https://example.com/chart.png"} on success, or
-        - {"error": "error message"} on failure
+    - dict: Either {"image_url": "url"} on success or {"error": "message"} on failure
     
     Example:
     - csv_url: "https://example.com/data.csv"
-    - query: "Show sales trends as a line chart"
+    - query: "Show sales trends as line chart"
     Returns:
     - dict: {"image_url": "https://storage.example.com/chart_uuid.png"}
     """
+    error_messages = []
     
-    async def upload_and_return(image_path: str, chat_id: str) -> dict:
-        """Helper function to handle image uploads"""
+    async def upload_and_return(image_path: str) -> Dict[str, str]:
+        """Handle image upload and return public URL"""
         unique_name = f'{uuid.uuid4()}.png'
         public_url = await upload_file_to_supabase(image_path, unique_name, chat_id)
-        logger.info(f"Uploaded chart: {public_url}")
-        os.remove(image_path) # Remove the local image file after upload
+        logger.info(f"Uploaded chart to: {public_url}")
+        try:
+            os.remove(image_path)
+        except OSError as e:
+            logger.warning(f"Could not delete temp file {image_path}: {str(e)}")
         return {"image_url": public_url}
 
+    # --- 1. First Attempt: LangChain Gemini ---
     try:
-        # Commented out for now because aiml api is not working
-        # try:
-        #     # --- 1. First Attempt: OpenAI ---
-        #     openai_result = await asyncio.to_thread(openai_chart, csv_url, query)
-        #     logger.info(f"OpenAI chart result:", openai_result)
-
-        #     if openai_result and openai_result != 'Chart not generated':
-        #         return await upload_and_return(openai_result, chat_id)
-
-        #     raise Exception("OpenAI failed to generate chart")
-
-        # except Exception as openai_error:
-        #     logger.info(f"OpenAI failed ({str(openai_error)}), trying raw Groq...")
-            # --- 2. Second Attempt: Raw Groq ---
-            try:
-                groq_result = await asyncio.to_thread(groq_chart, csv_url, query)
-                logger.info(f"Raw Groq chart result:", groq_result)
-            
-                if groq_result and groq_result != 'Chart not generated':
-                    return await upload_and_return(groq_result, chat_id)
-                
-                raise Exception("Raw Groq failed to generate chart")
+        gemini_result = await asyncio.to_thread(
+            langchain_gemini_csv_handler, csv_url, query, True
+        )
+        logger.info(f"LangChain-Gemini chart result: {gemini_result}")
+        
+        if gemini_result:
+            # Handle string or list response
+            if isinstance(gemini_result, str):
+                clean_path = gemini_result.strip()
+                if os.path.exists(clean_path):
+                    return await upload_and_return(clean_path)
             
-            except Exception as groq_error:
-                logger.info(f"Raw Groq failed ({str(groq_error)}), trying LangChain Gemini...")
+            if isinstance(gemini_result, list) and gemini_result:
+                first_path = gemini_result[0]
+                if os.path.exists(first_path):
+                    return await upload_and_return(first_path)
+        
+        error_messages.append("LangChain-Gemini returned invalid result")
+    except Exception as gemini_error:
+        error_messages.append(f"LangChain-Gemini error: {str(gemini_error)}")
+        logger.error(f"Gemini chart failed: {str(gemini_error)}")
 
-                # --- 3. Third Attempt: LangChain Gemini ---
-                try:
-                    gemini_result = await asyncio.to_thread(
-                        langchain_gemini_csv_handler, csv_url, query, True
-                    )
-                    logger.info("LangChain Gemini chart result:", gemini_result)
-                
-                    # --- i) If Gemini result is a string, return it ---
-                    if gemini_result and isinstance(gemini_result, str):
-                        clean_path = gemini_result.strip()
-                        return await upload_and_return(clean_path, chat_id)
-                
-                    # --- ii) If Gemini result is a list, return the first element ---
-                    if gemini_result and isinstance(gemini_result, list) and len(gemini_result) > 0:
-                        return await upload_and_return(gemini_result[0], chat_id)
-                    
-                    raise Exception("LangChain Gemini returned empty result")
-                
-                except Exception as gemini_error:
-                    logger.info(f"LangChain Gemini failed ({str(gemini_error)}), trying LangChain Groq...")
-
-                    # --- 4. Final Attempt: LangChain Groq ---
-                    try:
-                        lc_groq_paths = await asyncio.to_thread(
-                            langchain_csv_chart, csv_url, query, True
-                        )
-                        logger.info("LangChain Groq chart result:", lc_groq_paths)
-                    
-                        if isinstance(lc_groq_paths, list) and lc_groq_paths:
-                            return await upload_and_return(lc_groq_paths[0], chat_id)
-                        
-                        return {"error": "All chart generation methods failed"}
-                    
-                    except Exception as lc_groq_error:
-                        logger.info(f"LangChain Groq failed: {str(lc_groq_error)}")
-                        return {"error": "Could not generate chart"}
-
-    except Exception as e:
-        logger.info(f"Critical error: {str(e)}")
-        return {"error": "Internal system error"}
+    # --- 2. Fallback Attempt: LangChain Groq ---
+    try:
+        lc_groq_paths = await asyncio.to_thread(
+            langchain_csv_chart, csv_url, query, True
+        )
+        logger.info(f"LangChain-Groq chart result: {lc_groq_paths}")
+        
+        if lc_groq_paths:
+            if isinstance(lc_groq_paths, list) and lc_groq_paths:
+                first_path = lc_groq_paths[0]
+                if os.path.exists(first_path):
+                    return await upload_and_return(first_path)
+            
+            if isinstance(lc_groq_paths, str) and os.path.exists(lc_groq_paths):
+                return await upload_and_return(lc_groq_paths)
+        
+        error_messages.append("LangChain-Groq returned invalid result")
+    except Exception as lc_groq_error:
+        error_messages.append(f"LangChain-Groq error: {str(lc_groq_error)}")
+        logger.error(f"Groq chart failed: {str(lc_groq_error)}")
+
+    # --- Final Error Handling ---
+    logger.error(f"All chart generation failed. Errors: {'; '.join(error_messages)}")
+    return {"error": "Could not generate chart from the provided data"}