changed model to .env gemini-flash-2.0

controller.py

@@ -28,7 +28,8 @@ import matplotlib
 import seaborn as sns
 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 orc_agent_main_cerebras import csv_orchestrator_chat_cerebras
+from orchestrator_agent import csv_orchestrator_chat_gemini
 from python_code_executor_service import CsvChatResult, PythonExecutor
 from supabase_service import upload_file_to_supabase
 from cerebras_csv_agent import query_csv_agent
@@ -401,7 +402,7 @@ async def csv_chat(request: Dict, authorization: str = Header(None)):
             logger.info("Processing detailed answer with orchestrator...")
             try:
                 orchestrator_answer = await asyncio.to_thread(
-                    csv_orchestrator_chat, decoded_url, query, conversation_history, chat_id
+                    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]}...")
@@ -676,7 +677,7 @@ async def csv_chart(request: dict, authorization: str = Header(None)):
         # Use orchestrator to handle the user's chart query first
         if detailed_answer is True:
            orchestrator_answer = await asyncio.to_thread(
-             csv_orchestrator_chat, csv_url, query, conversation_history, chat_id
+             csv_orchestrator_chat_gemini, csv_url, query, conversation_history, chat_id
            )
         
            if orchestrator_answer is not None:

orc_agent_main_cerebras.py

@@ -0,0 +1,190 @@
+import os
+from typing import List, Any
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIChatModel
+from openai import RateLimitError, APIError  
+from csv_service import get_csv_basic_info
+from orchestrator_functions import csv_chart, csv_chat
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Load all API keys from the environment variable
+CEREBRAS_API_KEYS = os.getenv("CEREBRAS_API_KEYS", "").split(",")  # Expecting a comma-separated list of keys
+CEREBRAS_BASE_URL = os.getenv("CEREBRAS_BASE_URL")  # Cerebras API base URL
+CEREBRAS_MODEL = os.getenv("CEREBRAS_MODEL")  # Default Cerebras model
+
+# Function to initialize the model with a specific API key
+def initialize_model(api_key: str) -> OpenAIChatModel:
+    """Initialize Cerebras model using OpenAI-compatible interface"""
+    return OpenAIChatModel(
+        CEREBRAS_MODEL,
+        base_url=CEREBRAS_BASE_URL,
+        api_key=api_key
+    )
+
+# Define the tools
+async def generate_csv_answer(csv_url: str, user_questions: List[str]) -> Any:
+    """
+    This function generates answers for the given user questions using the CSV URL.
+    It uses the csv_chat function to process each question and return the answers.
+    
+    Args:
+        csv_url (str): The URL of the CSV file.
+        user_questions (List[str]): A list of user questions.
+        
+    Returns:
+        List[Dict[str, Any]]: A list of dictionaries containing the question and answer for each question.
+        
+    Example:
+        [
+            {"question": "What is the average age of the customers?", "answer": "The average age is 35."},
+            {"question": "What is the most common gender?", "answer": "The most common gender is Male."}
+        ]
+    """
+    
+    print("LLM using the csv chat function....")
+    print("CSV URL:", csv_url)
+    print("User question:", user_questions)
+
+    # Create an array to accumulate the answers
+    answers = []
+    # Loop through the user questions and generate answers for each
+    for question in user_questions:
+        answer = await csv_chat(csv_url, question)
+        answers.append(dict(question=question, answer=answer))
+    return answers
+
+async def generate_chart(csv_url: str, user_questions: List[str], chat_id: str) -> Any:
+    """
+    This function generates charts for the given user questions using the CSV URL.
+    It uses the csv_chart function to process each question and return the chart URLs.
+    It returns a list of dictionaries containing the question and chart URL for each question.
+    
+    Args:
+        csv_url (str): The URL of the CSV file.
+        user_questions (List[str]): A list of user questions.
+        chat_id (str): The chat ID for the session.
+
+    Returns:
+        List[Dict[str, Any]]: A list of dictionaries containing the question and chart URL for each question.
+
+    Example:
+        [
+            {"question": "What is the average age of the customers?", "chart_url": "https://example.com/chart1.png"},
+            {"question": "What is the most common gender?", "chart_url": "https://example.com/chart2.png"}
+        ]
+    """
+    
+    print("LLM using the csv chart function....")
+    print("CSV URL:", csv_url)
+    print("User question:", user_questions)
+
+    # Create an array to accumulate the charts
+    charts = []
+    # Loop through the user questions and generate charts for each
+    for question in user_questions:
+        chart = await csv_chart(csv_url, question, chat_id)
+        charts.append(dict(question=question, image_url=chart))
+    
+    return charts
+
+# Function to create an agent with a specific CSV URL
+def create_agent(csv_url: str, api_key: str, conversation_history: List, chat_id: str) -> Agent:
+    """Create a PydanticAI agent configured for CSV analysis using Cerebras"""
+    csv_metadata = get_csv_basic_info(csv_url)
+    
+    system_prompt = f"""
+# Role: Data Analyst Assistant
+**Specialization:** CSV Analysis & Visualization
+**Powered by:** Cerebras AI
+
+## Key Rules:
+1. **Always provide both:** 
+   - Complete textual answer with explanations
+   - Visualization when applicable
+2. **Output Format:** Markdown compatible (visualizations as `![Image Description](url generated by tool)`)
+3. **Tool Handling:**
+   - Use `generate_csv_answer` for analysis
+   - Use `generate_chart` for visuals
+   - Never disclose tool names
+4. **Visualization Fallback:**
+   - If requested library (plotly, bokeh etc.) isn't available:
+     - Provide closest alternative
+     - Explain the limitation
+
+## Current Context:
+- **Dataset:** {csv_url}
+- **Metadata:** {csv_metadata}
+- **History:** {conversation_history}
+- **Chat ID:** {chat_id}
+
+## Required Output:
+For every question return:
+1. Clear analysis answer
+2. Visualization (when possible, in markdown format)
+3. Follow-up suggestions
+
+**Critical:** 
+- Never return partial responses - always combine both textual answers and visualizations when applicable.
+- Always generate a fresh, tool-assisted response for every query, regardless of its similarity to any prior questions. Never reuse or return a previous answer.
+- Leverage Cerebras's fast inference capabilities for efficient data analysis.
+"""
+
+    return Agent(
+        model=initialize_model(api_key),
+        deps_type=str,
+        tools=[generate_csv_answer, generate_chart],
+        system_prompt=system_prompt
+    )
+
+def csv_orchestrator_chat_cerebras(csv_url: str, user_question: str, conversation_history: List, chat_id: str) -> str:
+    """
+    Main orchestrator function that processes CSV analysis requests using Cerebras AI.
+    
+    Args:
+        csv_url (str): URL of the CSV file to analyze
+        user_question (str): User's question about the CSV data
+        conversation_history (List): Previous conversation context
+        chat_id (str): Unique chat session identifier
+        
+    Returns:
+        str: Analysis response or None if all API keys are exhausted
+    """
+    print("CSV URL:", csv_url)
+    print("User questions:", user_question)
+
+    # Validate API keys
+    if not CEREBRAS_API_KEYS or CEREBRAS_API_KEYS == ['']:
+        print("Error: No Cerebras API keys found. Please set CEREBRAS_API_KEYS environment variable.")
+        return "Configuration error: Cerebras API keys not found."
+
+    # Iterate through all API keys with improved error handling
+    for i, api_key in enumerate(CEREBRAS_API_KEYS):
+        api_key = api_key.strip()  # Remove any whitespace
+        if not api_key:
+            continue
+            
+        try:
+            print(f"Attempting with Cerebras API key #{i+1}")
+            agent = create_agent(csv_url, api_key, conversation_history, chat_id)
+            result = agent.run_sync(user_question)
+            print("Orchestrator Result:", result.data)
+            return result.data
+            
+        except RateLimitError as e:
+            print(f"Rate limit exceeded for API key #{i+1}. Switching to the next key.")
+            continue
+            
+        except APIError as e:
+            print(f"API error with key #{i+1}: {e}")
+            continue
+            
+        except Exception as e:
+            print(f"Unexpected error with API key #{i+1}: {e}")
+            continue
+
+    # If all keys are exhausted or fail
+    error_msg = "All Cerebras API keys have been exhausted or failed. Please check your API keys and quotas."
+    print(error_msg)
+    return error_msg

orchestrator_agent.py

@@ -134,7 +134,7 @@ For every question return:
         system_prompt=system_prompt
     )
 
-def csv_orchestrator_chat(csv_url: str, user_question: str, conversation_history: List, chat_id: str) -> str:
+def csv_orchestrator_chat_gemini(csv_url: str, user_question: str, conversation_history: List, chat_id: str) -> str:
     print("CSV URL:", csv_url)
     print("User questions:", user_question)