go mdrfkr

cerebras_openrouter_chart_generator.py

@@ -121,97 +121,75 @@ def extract_code(content: str) -> str:
     return match.group(1).strip() if match else content.replace("```", "").strip()
 
 # ==================================================================================
-#                               CEREBRAS AGENT (MULTI-KEY)
+#                               CEREBRAS AGENT
 # ==================================================================================
 
-def generate_cerebras_chart(csv_url: str, query: str, max_retries: int = 3) -> Optional[str]:
-    """
-    Cerebras-specific orchestrator.
-    Iterates through ALL available Cerebras API keys if failures occur.
-    """
-    logger.info("Starting CEREBRAS chart generation...")
+def get_cerebras_client() -> OpenAI:
+    api_keys = os.getenv("CEREBRAS_API_KEYS", "").split(",")
+    base_url = os.getenv("CEREBRAS_BASE_URL", "https://api.cerebras.ai/v1")
     
-    # Get all keys and clean them
-    raw_keys = os.getenv("CEREBRAS_API_KEYS", "").split(",")
-    api_keys = [k.strip() for k in raw_keys if k.strip()]
+    if not api_keys or not api_keys[0]:
+        raise ValueError("Missing CEREBRAS_API_KEYS")
     
-    if not api_keys:
-        logger.error("No CEREBRAS_API_KEYS found.")
-        return None
+    # Use the first key (can implement rotation if needed)
+    return OpenAI(base_url=base_url, api_key=api_keys[0])
 
-    base_url = os.getenv("CEREBRAS_BASE_URL", "https://api.cerebras.ai/v1")
-    model_name = os.getenv("CEREBRAS_CODING_MODEL", "llama3.1-70b")
+def generate_cerebras_chart(csv_url: str, query: str, max_retries: int = 3) -> Optional[str]:
+    """Cerebras-specific orchestrator."""
+    logger.info("Starting CEREBRAS chart generation...")
+    model_name = os.getenv("CEREBRAS_MODEL", "llama3.1-70b")
     data_context = get_data_context(csv_url)
+    error_history = []
+    
+    client = get_cerebras_client()
 
-    # LOOP THROUGH EACH KEY
-    for key_index, api_key in enumerate(api_keys):
-        logger.info(f"--- Attempting Cerebras with Key Index {key_index + 1}/{len(api_keys)} ---")
-        
+    for attempt in range(1, max_retries + 1):
         try:
-            client = OpenAI(base_url=base_url, api_key=api_key)
-            error_history = []
+            system_prompt, user_prompt = construct_prompt(query, csv_url, data_context, error_history)
             
-            # Inner Loop: Attempt to fix code `max_retries` times with CURRENT key
-            for attempt in range(1, max_retries + 1):
-                try:
-                    system_prompt, user_prompt = construct_prompt(query, csv_url, data_context, error_history)
-                    
-                    response = client.chat.completions.create(
-                        model=model_name,
-                        messages=[
-                            {"role": "system", "content": system_prompt},
-                            {"role": "user", "content": user_prompt}
-                        ],
-                        temperature=0.1
-                    )
-                    
-                    code = extract_code(response.choices[0].message.content)
-                    file_path, error = execute_and_capture(code, csv_url)
-                    
-                    if file_path:
-                        logger.info(f"Cerebras Success (Key Index {key_index}): {file_path}")
-                        return file_path
-                    
-                    # Logic error in code, record and retry loop
-                    error_history.append(f"Attempt {attempt} Error:\n{error}")
-                    logger.warning(f"Cerebras Key {key_index} - Attempt {attempt} Logic Fail: {error}")
-                
-                except Exception as api_e:
-                    # If it's an API error (Auth, Rate Limit, Network), break inner loop to switch keys
-                    logger.error(f"Cerebras API Error with Key {key_index}: {api_e}")
-                    break 
+            response = client.chat.completions.create(
+                model=model_name,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
+                ],
+                temperature=0.1
+            )
+            
+            code = extract_code(response.choices[0].message.content)
+            file_path, error = execute_and_capture(code, csv_url)
+            
+            if file_path:
+                logger.info(f"Cerebras Success: {file_path}")
+                return file_path
+            
+            error_history.append(f"Attempt {attempt} Error:\n{error}")
             
-            logger.warning(f"Cerebras Key {key_index} exhausted retries or failed. Moving to next key...")
-
         except Exception as e:
-            logger.error(f"Fatal error initializing Cerebras Key {key_index}: {e}")
-            continue
-
-    logger.error("All Cerebras keys failed.")
+            logger.error(f"Cerebras Attempt {attempt} failed: {e}")
+            error_history.append(f"System Error: {str(e)}")
+            
     return None
 
 # ==================================================================================
 #                               OPENROUTER AGENT (FALLBACK)
 # ==================================================================================
 
+def get_openrouter_client() -> OpenAI:
+    api_keys = os.getenv("OPENROUTER_API_KEYS", "").split(",")
+    base_url = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
+    if not api_keys or not api_keys[0]:
+        raise ValueError("Missing OPENROUTER_API_KEYS")
+    return OpenAI(base_url=base_url, api_key=api_keys[0])
+
 def generate_openrouter_chart(csv_url: str, query: str, max_retries: int = 3) -> Optional[str]:
     """OpenRouter-specific orchestrator (Fallback)."""
     logger.info("Starting OPENROUTER chart generation (Fallback)...")
-    
-    # We stick to the first key for OpenRouter as per original logic, but can be expanded
-    api_keys = os.getenv("OPENROUTER_API_KEYS", "").split(",")
-    if not api_keys or not api_keys[0]:
-        logger.error("Missing OPENROUTER_API_KEYS")
-        return None
-        
-    client = OpenAI(
-        base_url=os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"),
-        api_key=api_keys[0]
-    )
-    
     model_name = os.getenv("OPENROUTER_MODEL", "openai/gpt-4o")
     data_context = get_data_context(csv_url)
     error_history = []
+    
+    client = get_openrouter_client()
 
     for attempt in range(1, max_retries + 1):
         try:
@@ -240,39 +218,3 @@ def generate_openrouter_chart(csv_url: str, query: str, max_retries: int = 3) ->
             error_history.append(f"System Error: {str(e)}")
             
     return None
-
-# ==================================================================================
-#                               MAIN ORCHESTRATOR
-# ==================================================================================
-
-def generate_cerebras_openrouter_chart(csv_url: str, query: str, max_retries: int = 3) -> Dict[str, str]:
-    """
-    Orchestrator that:
-    1. Tries all Cerebras keys.
-    2. If all fail, tries OpenRouter.
-    3. Returns result dictionary.
-    """
-    
-    # 1. Try Cerebras (Handles key rotation internally)
-    chart_path = generate_cerebras_chart(csv_url, query, max_retries)
-    
-    # 2. Fallback to OpenRouter if Cerebras failed
-    if not chart_path:
-        logger.warning("Cerebras failed (all keys). Switching to OpenRouter...")
-        chart_path = generate_openrouter_chart(csv_url, query, max_retries)
-        
-    # 3. Construct Response
-    if chart_path:
-        return {
-            "image_url": chart_path,
-            "message": "Chart generated successfully."
-        }
-    else:
-        return {
-            "image_url": "",
-            "message": "Failed to generate chart after trying all providers."
-        }
-
-# ==================================================================================
-#                               ENTRY POINT
-# ==================================================================================

orchestrator_functions.py

@@ -12,7 +12,7 @@ from pandasai import SmartDataframe
 from langchain_groq.chat_models import ChatGroq
 from dotenv import load_dotenv
 from pydantic import BaseModel
-from cerebras_openrouter_chart_generator import generate_cerebras_openrouter_chart
+from cerebras_openrouter_chart_generator import generate_cerebras_chart, generate_openrouter_chart
 from csv_service import clean_data, extract_chart_filenames
 from langchain_groq import ChatGroq
 import pandas as pd
@@ -395,21 +395,61 @@ 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.
     Strategy:
-    1. Try Cerebras Agent (All Keys).
+    1. Try Cerebras Agent.
     2. If failed/empty, try OpenRouter Agent.
-    3. Return result.
+    3. Upload and return.
     """
     logger.info(f"Received csv_chart request. Chat ID: {chat_id}")
+    error_messages = []
+    
+    async def upload_and_return(image_path: str) -> Dict[str, str]:
+        """Handle image upload and return public URL"""
+        unique_name = f'{uuid.uuid4()}.png'
+        # Ensure upload_file_to_supabase is imported in your file
+        public_url = await upload_file_to_supabase(image_path, unique_name, chat_id)
+        
+        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. Primary Attempt: Cerebras ---
     try:
-        logger.info("Attempting Chart Generation...")
-        # Run synchronous generation logic in a thread
-        csv_chart_response = await asyncio.to_thread(
-            generate_cerebras_openrouter_chart, csv_url, query, 3
+        logger.info("Attempting with Cerebras...")
+        cerebras_path = await asyncio.to_thread(
+            generate_cerebras_chart, csv_url, query, 3
         )
         
-        logger.info(f"Chart Response: {csv_chart_response}")
-        return csv_chart_response
+        if cerebras_path:
+            return await upload_and_return(cerebras_path)
+        
+        error_messages.append("Cerebras returned no valid image.")
+        logger.warning("Cerebras failed. Switching to fallback.")
+        
     except Exception as e:
-        error_message = f"Chart Generation Error: {str(e)}"
-        logger.error(error_message)
-        return {"image_url": "", "message": error_message}
+        msg = f"Cerebras critical error: {str(e)}"
+        error_messages.append(msg)
+        logger.error(msg)
+
+    # --- 2. Fallback Attempt: OpenRouter ---
+    try:
+        logger.info("Attempting with OpenRouter...")
+        openrouter_path = await asyncio.to_thread(
+            generate_openrouter_chart, csv_url, query, 3
+        )
+        
+        if openrouter_path:
+            return await upload_and_return(openrouter_path)
+            
+        error_messages.append("OpenRouter fallback returned no valid image.")
+        
+    except Exception as e:
+        msg = f"OpenRouter critical error: {str(e)}"
+        error_messages.append(msg)
+        logger.error(msg)
+
+    # --- Final Error Handling ---
+    logger.error(f"All chart generation providers failed. Errors: {'; '.join(error_messages)}")
+    return {"error": "Could not generate chart. Both primary and fallback agents failed."}