changed model to .env gemini-flash-2.0

cerebras_report_generator.py

@@ -0,0 +1,385 @@
+import json
+import numpy as np
+import pandas as pd
+import re
+import os
+import uuid
+import logging
+from io import StringIO
+import sys
+import traceback
+from typing import Optional, Dict, Any, List
+from pydantic import BaseModel, Field
+from openai import OpenAI
+from dotenv import load_dotenv
+import seaborn as sns 
+import datetime as dt
+
+from supabase_service import upload_file_to_supabase
+
+pd.set_option('display.max_columns', None)
+pd.set_option('display.max_rows', None)
+pd.set_option('display.max_colwidth', None)
+
+load_dotenv()
+
+API_KEYS = os.getenv("CEREBRAS_API_KEYS", "").split(",")[::-1]
+MODEL_NAME = os.getenv("CEREBRAS_MODEL")  # Default Cerebras model
+CEREBRAS_BASE_URL = os.getenv("CEREBRAS_BASE_URL")
+
+class FileProps(BaseModel):
+    fileName: str
+    filePath: str
+    fileType: str  # 'csv' | 'image'
+
+class Files(BaseModel):
+    csv_files: List[FileProps]
+    image_files: List[FileProps]
+
+class FileBoxProps(BaseModel):
+    files: Files
+
+os.environ['MPLBACKEND'] = 'agg'
+import matplotlib.pyplot as plt
+plt.show = lambda: None
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+class CerebrasKeyManager:
+    """Manage multiple Cerebras API keys with failover"""
+    
+    def __init__(self, api_keys: List[str], base_url: str):
+        self.original_keys = api_keys.copy()
+        self.available_keys = api_keys.copy()
+        self.base_url = base_url
+        self.active_key = None
+        self.failed_keys = {}
+        self.client = None
+        
+    def configure(self) -> bool:
+        while self.available_keys:
+            key = self.available_keys.pop(0)
+            try:
+                self.client = OpenAI(
+                    api_key=key,
+                    base_url=self.base_url
+                )
+                # Test the connection with a simple request
+                response = self.client.models.list()
+                self.active_key = key
+                logger.info(f"Configured with key: {self._mask_key(key)}")
+                return True
+            except Exception as e:
+                self.failed_keys[key] = str(e)
+                logger.error(f"Key failed: {self._mask_key(key)}. Error: {str(e)}")
+        logger.critical("All API keys failed")
+        return False
+    
+    def _mask_key(self, key: str) -> str:
+        return f"{key[:8]}...{key[-4:]}" if key else ""
+
+class PythonREPL:
+    """Secure Python REPL with file generation tracking"""
+    
+    def __init__(self, df: pd.DataFrame):
+        self.df = df
+        self.output_dir = os.path.abspath(f'generated_outputs/{uuid.uuid4()}')
+        os.makedirs(self.output_dir, exist_ok=True)
+        self.local_env = {
+            "pd": pd,
+            "df": self.df.copy(),
+            "plt": plt,
+            "os": os,
+            "uuid": uuid,
+            "sns": sns,
+            "json": json,
+            "dt": dt,
+            "output_dir": self.output_dir
+        }
+        
+    def execute(self, code: str) -> Dict[str, Any]:
+        print('Executing code...', code)
+        old_stdout = sys.stdout
+        sys.stdout = mystdout = StringIO()
+        file_tracker = {
+            'csv_files': set(),
+            'image_files': set()
+        }
+
+        try:
+            code = f"""
+import matplotlib.pyplot as plt
+plt.switch_backend('agg')
+{code}
+plt.close('all')
+            """
+            exec(code, self.local_env)
+            self.df = self.local_env.get('df', self.df)
+            
+            # Track generated files
+            for fname in os.listdir(self.output_dir):
+                if fname.endswith('.csv'):
+                    file_tracker['csv_files'].add(fname)
+                elif fname.lower().endswith(('.png', '.jpg', '.jpeg')):
+                    file_tracker['image_files'].add(fname)
+                    
+            error = False
+        except Exception as e:
+            error_msg = traceback.format_exc()
+            error = True
+        finally:
+            sys.stdout = old_stdout
+            
+        return {
+            "output": mystdout.getvalue(),
+            "error": error,
+            "error_message": error_msg if error else None,
+            "df": self.local_env.get('df', self.df),
+            "output_dir": self.output_dir,
+            "files": {
+                "csv": [os.path.join(self.output_dir, f) for f in file_tracker['csv_files']],
+                "images": [os.path.join(self.output_dir, f) for f in file_tracker['image_files']]
+            }
+        }
+
+class RethinkAgent(BaseModel):
+    df: pd.DataFrame
+    max_retries: int = Field(default=5, ge=1)
+    cerebras_client: Optional[OpenAI] = None
+    model_name: str = Field(default="llama3.1-8b")
+    current_retry: int = Field(default=0, ge=0)
+    repl: Optional[PythonREPL] = None
+    key_manager: Optional[CerebrasKeyManager] = None
+    conversation: List[Dict[str, Any]] = []
+    
+    class Config:
+        arbitrary_types_allowed = True
+        
+    def _extract_code(self, response: str) -> str:
+        code_match = re.search(r'```python(.*?)```', response, re.DOTALL)
+        return code_match.group(1).strip() if code_match else response.strip()
+    
+    def _generate_initial_prompt(self, query: str) -> str:
+        initial_prompt = f"""Generate DIRECT EXECUTION CODE (no functions, no explanations) following STRICT RULES:
+    
+        CONVERSATION HISTORY:
+        {self.conversation}
+        
+        MANDATORY REQUIREMENTS:
+        1. Operate directly on existing 'df' variable
+        2. Save ALL final DataFrames to CSV using: df.to_csv(f'{{output_dir}}/descriptive_name.csv')
+        3. For visualizations: plt.savefig(f'{{output_dir}}/chart_name.png')
+        4. Use EXACTLY this structure:
+           # Data processing
+           df_processed = df[...]  # filtering/grouping
+           # Save results
+           df_processed.to_csv(f'{{output_dir}}/result.csv')
+           # Visualizations (if needed)
+           plt.figure()
+           ... plotting code ...
+           plt.savefig(f'{{output_dir}}/chart.png')
+           plt.close()
+
+        FORBIDDEN:
+        - Function definitions
+        - Dummy data creation
+        - Any code blocks besides pandas operations and matplotlib
+        - Print statements showing dataframes
+        - Using any visualization library other than matplotlib or seaborn
+
+        DATAFRAME COLUMNS: {', '.join(self.df.columns)}
+        DATAFRAME'S FIRST FIVE ROWS: {self.df.head().to_dict('records')}
+        USER QUERY: {query}
+
+        EXAMPLE RESPONSE FOR "Sales by region":
+        # Data processing
+        sales_by_region = df.groupby('region')['sales'].sum().reset_index()
+        # Save results
+        sales_by_region.to_csv(f'{{output_dir}}/sales_by_region.csv')
+        """
+        logger.info('Conversation history:', self.conversation)
+        return initial_prompt
+
+    def _generate_retry_prompt(self, query: str, error: str, code: str) -> str:
+        return f"""FIX THIS CODE (failed with: {error}) by STRICTLY FOLLOWING:
+        
+        1. REMOVE ALL FUNCTION DEFINITIONS
+        2. ENSURE DIRECT DF OPERATIONS
+        3. USE EXPLICIT output_dir PATHS
+        4. ADD NECESSARY IMPORTS IF MISSING
+        5. VALIDATE COLUMN NAMES EXIST
+
+        BAD CODE:
+        {code}
+
+        CORRECTED CODE:"""
+
+    def initialize_model(self, api_keys: List[str], base_url: str) -> bool:
+        self.key_manager = CerebrasKeyManager(api_keys, base_url)
+        if not self.key_manager.configure():
+            raise RuntimeError("API key initialization failed")
+        try:
+            self.cerebras_client = self.key_manager.client
+            return True
+        except Exception as e:
+            logger.error(f"Model init failed: {str(e)}")
+            return False
+    
+    def generate_code(self, query: str, error: Optional[str] = None, previous_code: Optional[str] = None) -> str:
+        prompt = self._generate_retry_prompt(query, error, previous_code) if error else self._generate_initial_prompt(query)
+        try:
+            response = self.cerebras_client.chat.completions.create(
+                model=self.model_name,
+                messages=[
+                    {"role": "system", "content": "You are a Python code generator. Generate only executable Python code without explanations."},
+                    {"role": "user", "content": prompt}
+                ],
+                max_tokens=2048,
+                temperature=0.1
+            )
+            return self._extract_code(response.choices[0].message.content)
+        except Exception as e:
+            if self.key_manager.available_keys and self.key_manager.configure():
+                self.cerebras_client = self.key_manager.client
+                return self.generate_code(query, error, previous_code)
+            raise
+    
+    def execute_query(self, query: str) -> Dict[str, Any]:
+        self.repl = PythonREPL(self.df)
+        result = None
+        
+        while self.current_retry < self.max_retries:
+            try:
+                code = self.generate_code(query, 
+                                         result["error_message"] if result else None,
+                                         result["code"] if result else None)
+                execution_result = self.repl.execute(code)
+                
+                if execution_result["error"]:
+                    self.current_retry += 1
+                    result = {
+                        "error_message": execution_result["error_message"],
+                        "code": code
+                    }
+                else:
+                    return {
+                        "text": execution_result["output"],
+                        "csv_files": execution_result["files"]["csv"],
+                        "image_files": execution_result["files"]["images"]
+                    }
+            except Exception as e:
+                return {
+                    "error": f"Critical failure: {str(e)}",
+                    "csv_files": [],
+                    "image_files": []
+                }
+        
+        return {
+            "error": f"Failed after {self.max_retries} retries",
+            "csv_files": [],
+            "image_files": []
+        }
+
+def cerebras_llm_chat(csv_url: str, query: str, conversation_history: List[Dict[str, Any]]) -> Dict[str, Any]:
+    try:
+        df = pd.read_csv(csv_url)
+        agent = RethinkAgent(df=df, conversation=conversation_history, model_name=MODEL_NAME)
+        
+        if not agent.initialize_model(API_KEYS, CEREBRAS_BASE_URL):
+            return {"error": "API configuration failed"}
+        
+        result = agent.execute_query(query)
+        
+        if "error" in result:
+            return result
+        
+        return {
+            "message": result["text"],
+            "csv_files": result["csv_files"],
+            "image_files": result["image_files"]
+        }
+    except Exception as e:
+        logger.error(f"Processing failed: {str(e)}")
+        return {
+            "error": f"Processing error: {str(e)}",
+            "csv_files": [],
+            "image_files": []
+        }
+        
+
+async def generate_csv_report_cerebras(csv_url: str, query: str, chat_id: str, conversation_history: List[Dict[str, Any]]) -> FileBoxProps:
+    try:
+        result = cerebras_llm_chat(csv_url, query, conversation_history)
+        logger.info(f"Raw result from cerebras_llm_chat: {result}")
+        
+        csv_files = []
+        image_files = []
+        
+        # Check if we got the expected response structure
+        if isinstance(result, dict) and 'csv_files' in result and 'image_files' in result:
+            # Process CSV files
+            for csv_path in result['csv_files']:
+                if os.path.exists(csv_path):
+                    file_name = os.path.basename(csv_path)
+                    try:
+                        unique_file_name = f"{uuid.uuid4()}_{file_name}"
+                        public_url = await upload_file_to_supabase(
+                            file_path=csv_path,
+                            file_name=unique_file_name,
+                            chat_id=chat_id
+                        )
+                        csv_files.append(FileProps(
+                            fileName=file_name,
+                            filePath=public_url,
+                            fileType="csv"
+                        ))
+                        os.remove(csv_path)  # Clean up
+                    except Exception as upload_error:
+                        logger.error(f"Failed to upload CSV {file_name}: {str(upload_error)}")
+                        continue
+            
+            # Process image files
+            for img_path in result['image_files']:
+                if os.path.exists(img_path):
+                    file_name = os.path.basename(img_path)
+                    try:
+                        unique_file_name = f"{uuid.uuid4()}_{file_name}"
+                        public_url = await upload_file_to_supabase(
+                            file_path=img_path,
+                            file_name=unique_file_name,
+                            chat_id=chat_id
+                        )
+                        image_files.append(FileProps(
+                            fileName=file_name,
+                            filePath=public_url,
+                            fileType="image"
+                        ))
+                        os.remove(img_path)  # Clean up
+                    except Exception as upload_error:
+                        logger.error(f"Failed to upload image {file_name}: {str(upload_error)}")
+                        continue
+            
+            return FileBoxProps(
+                files=Files(
+                    csv_files=csv_files,
+                    image_files=image_files
+                )
+            )
+        else:
+            raise ValueError("Unexpected response format from cerebras_llm_chat")
+            
+    except Exception as e:
+        logger.error(f"Report generation failed: {str(e)}")
+        # Return empty response but log the files we found
+        if 'csv_files' in locals() and 'image_files' in locals():
+            logger.info(f"Files that were generated but not processed: CSV: {result.get('csv_files', [])}, Images: {result.get('image_files', [])}")
+        return FileBoxProps(
+            files=Files(
+                csv_files=[],
+                image_files=[]
+            )
+        )

controller.py

@@ -4,6 +4,7 @@ import logging
 import os
 import asyncio
 import threading
+import traceback
 import uuid
 from fastapi import FastAPI, HTTPException, Header
 from fastapi.encoders import jsonable_encoder
@@ -25,7 +26,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib
 import seaborn as sns
-from gemini_report_generator import generate_csv_report
+from gemini_report_generator import generate_csv_report_gemini
 from intitial_q_handler import if_initial_chart_question, if_initial_chat_question
 from orchestrator_agent import csv_orchestrator_chat
 from python_code_executor_service import CsvChatResult, PythonExecutor
@@ -33,6 +34,7 @@ from supabase_service import upload_file_to_supabase
 from cerebras_csv_agent import query_csv_agent
 from util_service import _prompt_generator, process_answer
 from fastapi.middleware.cors import CORSMiddleware
+
 import matplotlib
 matplotlib.use('Agg')
 
@@ -110,8 +112,6 @@ async def root():
 async def root():
     return {"message": "Pong !!"}
 
-
-
 # BASIC KNOWLEDGE BASED ON CSV
 
 # Remove trailing slash from the URL otherwise it will redirect to GET method
@@ -324,19 +324,25 @@ def langchain_csv_chat(csv_url: str, question: str, chart_required: bool):
                 return {"error": error_message}
 
     return {"error": "All API keys exhausted"}
+from cerebras_report_generator import generate_csv_report_cerebras
 
 # Async endpoint with non-blocking execution
 @app.post("/api/csv-chat")
 async def csv_chat(request: Dict, authorization: str = Header(None)):
     # Authorization checks
     if not authorization or not authorization.startswith("Bearer "):
+        logger.error("Authorization failed: Missing or invalid authorization header")
         raise HTTPException(status_code=401, detail="Invalid authorization")
     
     token = authorization.split(" ")[1]
     if token != os.getenv("AUTH_TOKEN"):
+        logger.error("Authorization failed: Invalid token")
         raise HTTPException(status_code=403, detail="Invalid token")
 
+    logger.info("Authorization successful")
+
     try:
+        # Extract request parameters
         query = request.get("query")
         csv_url = request.get("csv_url")
         decoded_url = unquote(csv_url)
@@ -345,57 +351,112 @@ async def csv_chat(request: Dict, authorization: str = Header(None)):
         generate_report = request.get("generate_report")
         chat_id = request.get("chat_id")
         
+        logger.info(f"Request parameters: query='{query[:100]}...', csv_url='{csv_url}', detailed_answer={detailed_answer}, generate_report={generate_report}, chat_id={chat_id}")
+        
+        # Handle report generation with Cerebras first, then Gemini fallback
         if generate_report is True:
-            report_files = await generate_csv_report(csv_url, query, chat_id, conversation_history)
-            if report_files is not None:
-                return {"answer": jsonable_encoder(report_files)}
+            logger.info("Starting report generation process...")
+            
+            # Try Cerebras first for report generation
+            logger.info("Attempting report generation with Cerebras...")
+            try:
+                report_files = await generate_csv_report_cerebras(csv_url, query, chat_id, conversation_history)
+                if report_files is not None and (report_files.files.csv_files or report_files.files.image_files):
+                    logger.info(f"Cerebras report generation successful: {len(report_files.files.csv_files)} CSV files, {len(report_files.files.image_files)} image files")
+                    return {"answer": jsonable_encoder(report_files)}
+                else:
+                    logger.warning("Cerebras report generation returned empty or None result")
+            except Exception as cerebras_error:
+                logger.error(f"Cerebras report generation failed: {str(cerebras_error)}")
+            
+            # Fallback to Gemini for report generation
+            logger.info("Falling back to Gemini for report generation...")
+            try:
+                report_files = await generate_csv_report_gemini(csv_url, query, chat_id, conversation_history)
+                if report_files is not None and (report_files.files.csv_files or report_files.files.image_files):
+                    logger.info(f"Gemini report generation successful: {len(report_files.files.csv_files)} CSV files, {len(report_files.files.image_files)} image files")
+                    return {"answer": jsonable_encoder(report_files)}
+                else:
+                    logger.warning("Gemini report generation returned empty or None result")
+            except Exception as gemini_error:
+                logger.error(f"Gemini report generation failed: {str(gemini_error)}")
+            
+            logger.error("Both Cerebras and Gemini report generation failed")
+            return {"answer": "error"}
 
+        # Handle initial chat questions with langchain
         if if_initial_chat_question(query):
-            answer = await asyncio.to_thread(
-                langchain_csv_chat, decoded_url, query, False
-            )
-            logger.info("langchain_answer:", answer)
-            return {"answer": jsonable_encoder(answer)}
+            logger.info("Processing as initial chat question with langchain...")
+            try:
+                answer = await asyncio.to_thread(
+                    langchain_csv_chat, decoded_url, query, False
+                )
+                logger.info(f"Langchain initial chat answer: {str(answer)[:200]}...")
+                return {"answer": jsonable_encoder(answer)}
+            except Exception as e:
+                logger.error(f"Langchain initial chat failed: {str(e)}")
         
-        # Orchestrate the execution
+        # Handle detailed answers with orchestrator
         if detailed_answer is True:
-           orchestrator_answer = await asyncio.to_thread(
-            csv_orchestrator_chat, decoded_url, query, conversation_history, chat_id
-           )
-           if orchestrator_answer is not None:
-              return {"answer": jsonable_encoder(orchestrator_answer)}
+            logger.info("Processing detailed answer with orchestrator...")
+            try:
+                orchestrator_answer = await asyncio.to_thread(
+                    csv_orchestrator_chat, decoded_url, query, conversation_history, chat_id
+                )
+                if orchestrator_answer is not None:
+                    logger.info(f"Orchestrator answer successful: {str(orchestrator_answer)[:200]}...")
+                    return {"answer": jsonable_encoder(orchestrator_answer)}
+                else:
+                    logger.warning("Orchestrator returned None result")
+            except Exception as e:
+                logger.error(f"Orchestrator processing failed: {str(e)}")
+
+        # Process with standard CSV agent (not Cerebras)
+        logger.info("Processing with standard CSV agent...")
+        try:
+            result = await query_csv_agent(decoded_url, query, chat_id)
+            logger.info(f"Standard CSV agent result: {str(result)[:200]}...")
+            if result is not None and result != "":
+                return {"answer": result}
+            else:
+                logger.warning("Standard CSV agent returned empty or None result")
+        except Exception as e:
+            logger.error(f"Standard CSV agent failed: {str(e)}")
 
-        # Process with groq_chat first
-        # groq_answer = await asyncio.to_thread(groq_chat, decoded_url, query)
-        # logger.info("groq_answer:", groq_answer)
-        
-        result = await query_csv_agent(decoded_url, query, chat_id)
-        logger.info("cerebras csv answer == >", result)
-        if result is not None or result == "":
-            return {"answer": result}
-     
-        # if process_answer(groq_answer) == "Empty response received.":
-        #     return {"answer": "Sorry, I couldn't find relevant data..."}
-
-        # if process_answer(groq_answer):
-        lang_answer = await asyncio.to_thread(
-            langchain_csv_chat, decoded_url, query, False
-        )
-        if process_answer(lang_answer):
-            return {"answer": "error"}
-        return {"answer": jsonable_encoder(lang_answer)}
+        # Fallback to langchain
+        logger.info("Falling back to langchain CSV chat...")
+        try:
+            lang_answer = await asyncio.to_thread(
+                langchain_csv_chat, decoded_url, query, False
+            )
+            logger.info(f"Langchain fallback result: {str(lang_answer)[:200]}...")
+            
+            if process_answer(lang_answer):
+                logger.error("Langchain fallback produced error response")
+                return {"answer": "error"}
+            
+            logger.info("Langchain fallback successful")
+            return {"answer": jsonable_encoder(lang_answer)}
+        except Exception as e:
+            logger.error(f"Langchain fallback failed: {str(e)}")
 
-        # return {"answer": jsonable_encoder(groq_answer)}
+        # If all methods fail
+        logger.error("All processing methods failed")
+        return {"answer": "error"}
 
     except Exception as e:
-        logger.error(f"Error processing request: {str(e)}")
+        logger.error(f"Critical error processing request: {str(e)}")
+        logger.error(f"Error traceback: {traceback.format_exc()}")
         return {"answer": "error"}
 
 def handle_out_of_range_float(value):
+    """Handle out of range float values for JSON serialization"""
     if isinstance(value, float):
         if np.isnan(value):
+            logger.debug("Converting NaN to None")
             return None
         elif np.isinf(value):
+            logger.debug("Converting Infinity to string")
             return "Infinity"
     return value
 
@@ -404,7 +465,6 @@ def handle_out_of_range_float(value):
 
 
 
-
 # CHART CODING STARTS FROM HERE
 
 instructions = """
@@ -593,7 +653,7 @@ async def csv_chart(request: dict, authorization: str = Header(None)):
         chat_id = request.get("chat_id", "")
         
         if generate_report is True:
-            report_files = await generate_csv_report(csv_url, query, chat_id, conversation_history)
+            report_files = await generate_csv_report_gemini(csv_url, query, chat_id, conversation_history)
             if report_files is not None:
                 return {"orchestrator_response": jsonable_encoder(report_files)}
 

gemini_report_generator.py

@@ -294,7 +294,7 @@ def gemini_llm_chat(csv_url: str, query: str, conversation_history: List[Dict[st
         }
         
 
-async def generate_csv_report(csv_url: str, query: str, chat_id: str, conversation_history: List[Dict[str, Any]]) -> FileBoxProps:
+async def generate_csv_report_gemini(csv_url: str, query: str, chat_id: str, conversation_history: List[Dict[str, Any]]) -> FileBoxProps:
     try:
         result = gemini_llm_chat(csv_url, query, conversation_history)
         logger.info(f"Raw result from gemini_llm_chat: {result}")