changed model to .env gemini-flash-2.0

cereberas_langchain_agent.py

@@ -16,15 +16,22 @@ import datetime as dt
 matplotlib.use('Agg')
 
 load_dotenv()
-model_name = os.getenv("CEREBRAS_LLM_MODEL")  # Specify your Cerebras model name
-cerebras_api_key = os.getenv("CEREBRAS_API_KEY")
-cerebras_base_url = os.getenv("CEREBRAS_BASE_URL")
+api_keys = os.getenv("CEREBRAS_API_KEYS", "").split(",")
+base_url = os.getenv("CEREBRAS_BASE_URL")
+model_name = os.getenv("CEREBRAS_MODEL")
 
-# Initialize ChatCerebras LLM
-llm = ChatCerebras(model=model_name, api_key=cerebras_api_key)
+current_key_index = 0  # Track which key is being used
+
+def get_next_llm():
+    """Return a ChatCerebras instance using the next available API key"""
+    global current_key_index
+    if current_key_index >= len(api_keys):
+        raise ValueError("All Cerebras API keys exhausted.")
+    key = api_keys[current_key_index]
+    print(f"Using Cerebras API key index: {current_key_index}")
+    return ChatCerebras(model=model_name, api_key=key, base_url=base_url)
 
 def create_agent(llm, data, tools):
-    """Create agent with tool names"""
     return create_pandas_dataframe_agent(
         llm,
         data,
@@ -37,7 +44,6 @@ def create_agent(llm, data, tools):
 
 def _prompt_generator(question: str, chart_required: bool, csv_url: str):
     chat_prompt = f"""You are a senior data analyst working with CSV data. Adhere strictly to the following guidelines:  
-
                      1. **Data Verification:** Always inspect the data with `.sample(5).to_dict()` before performing any analysis.  
                      2. **Data Integrity:** Ensure proper handling of null values to maintain accuracy and reliability.  
                      3. **Communication:** Provide concise, professional, and well-structured responses.
@@ -48,7 +54,6 @@ def _prompt_generator(question: str, chart_required: bool, csv_url: str):
                  """               
         
     chart_prompt = f"""You are a senior data analyst working with CSV data. Follow these rules STRICTLY:
-
                    1. Generate ONE unique identifier FIRST using: unique_id = uuid.uuid4().hex
                    2. Visualization requirements:
                    - Adjust font sizes, rotate labels (45° if needed), truncate for readability
@@ -80,15 +85,12 @@ def _prompt_generator(question: str, chart_required: bool, csv_url: str):
                    - Always use pd.read_csv({csv_url}) to read the CSV file
                    """
 
-    if chart_required:
-        return ChatPromptTemplate.from_template(chart_prompt)
-    else:
-        return ChatPromptTemplate.from_template(chat_prompt)
+    return ChatPromptTemplate.from_template(chart_prompt if chart_required else chat_prompt)
 
 def cerebras_csv_handler(csv_url: str, question: str, chart_required: bool):
-    """Process CSV using ChatCerebras"""
+    """Process CSV using ChatCerebras with key rotation"""
+    global current_key_index
     data = pd.read_csv(csv_url)
-    
     tool = PythonAstREPLTool(
         locals={
             "df": data,
@@ -101,13 +103,19 @@ def cerebras_csv_handler(csv_url: str, question: str, chart_required: bool):
             "dt": dt
         },
     )
-
-    agent = create_agent(llm, data, [tool])
-    prompt = _prompt_generator(question, chart_required, csv_url)
-    result = agent.invoke({"input": prompt})
-    output = result.get("output")
     
-    if output is None:
-        raise ValueError("Received None response from agent")
-        
-    return output
+    while current_key_index < len(api_keys):
+        try:
+            llm = get_next_llm()
+            agent = create_agent(llm, data, [tool])
+            prompt = _prompt_generator(question, chart_required, csv_url)
+            result = agent.invoke({"input": prompt})
+            output = result.get("output")
+            if output is None:
+                raise ValueError("Received None response from agent")
+            return output
+        except Exception as e:
+            print(f"Error with key index {current_key_index}: {e}")
+            current_key_index += 1
+    
+    raise ValueError("All Cerebras API keys exhausted.")

controller.py

@@ -16,6 +16,7 @@ from pandasai import SmartDataframe
 from langchain_groq.chat_models import ChatGroq
 from dotenv import load_dotenv
 from pydantic import BaseModel, Field
+from cerebras_report_generator import generate_csv_report_cerebras
 from csv_service import clean_data, extract_chart_filenames, generate_csv_data, get_csv_basic_info
 from urllib.parse import unquote
 from langchain_groq import ChatGroq
@@ -325,7 +326,45 @@ 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 def handle_detailed_answer(decoded_url, query, conversation_history, chat_id):
+    """
+    Try CSV processing first with Cerebras orchestrator, then fallback to Gemini if needed.
+    """
+    orchestrator_answer = None
+
+    # Step 1: Try Cerebras
+    try:
+        logger.info("Processing detailed answer with Cerebras orchestrator...")
+        orchestrator_answer = await asyncio.to_thread(
+            csv_orchestrator_chat_cerebras, decoded_url, query, conversation_history, chat_id
+        )
+        if orchestrator_answer is not None:
+            logger.info(f"Cerebras answer successful: {str(orchestrator_answer)[:200]}...")
+            return {"answer": jsonable_encoder(orchestrator_answer)}
+        else:
+            logger.warning("Cerebras orchestrator returned None")
+    except Exception as e:
+        logger.error(f"Cerebras orchestrator failed: {str(e)}")
+
+    # Step 2: Fallback to Gemini
+    try:
+        logger.info("Falling back to Gemini orchestrator...")
+        orchestrator_answer = await asyncio.to_thread(
+            csv_orchestrator_chat_gemini, decoded_url, query, conversation_history, chat_id
+        )
+        if orchestrator_answer is not None:
+            logger.info(f"Gemini answer successful: {str(orchestrator_answer)[:200]}...")
+            return {"answer": jsonable_encoder(orchestrator_answer)}
+        else:
+            logger.warning("Gemini orchestrator returned None")
+    except Exception as e:
+        logger.error(f"Gemini orchestrator failed: {str(e)}")
+
+    # Step 3: Both failed
+    logger.error("Both Cerebras and Gemini orchestrators failed or returned None")
+    return {"answer": None}
 
 # Async endpoint with non-blocking execution
 @app.post("/api/csv-chat")
@@ -400,17 +439,7 @@ async def csv_chat(request: Dict, authorization: str = Header(None)):
         # Handle detailed answers with orchestrator
         if detailed_answer is True:
             logger.info("Processing detailed answer with orchestrator...")
-            try:
-                orchestrator_answer = await asyncio.to_thread(
-                    csv_orchestrator_chat_cerebras, 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)}")
+            return await handle_detailed_answer(decoded_url, query, conversation_history, chat_id)
 
         # Process with standard CSV agent (not Cerebras)
         logger.info("Processing with standard CSV agent...")