go

cerebras_instance_provider.py

@@ -1,85 +1,50 @@
-# instance_provider.py
 import os
 import logging
-from typing import List, Optional, Tuple
-from openai import OpenAI
+from typing import List, Optional, Dict
 from dotenv import load_dotenv
 
 load_dotenv()
 
-# Setup basic logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class InstanceProvider:
-    """Manages multiple Cerebras/OpenAI clients with simple rotation"""
+    """Manages multiple Cerebras API keys with simple rotation"""
     
     def __init__(self):
-        self.clients: List[OpenAI] = []
+        self.api_keys: List[str] = []
         self.current_index = 0
-        self.model_name = os.getenv("CEREBRAS_MODEL") or "llama3.1-70b"
+        self.base_url = os.getenv("CEREBRAS_BASE_URL")
+        self.model_name = os.getenv("CEREBRAS_MODEL", "llama3.1-70b")
         self._initialize_instances()
     
     def _initialize_instances(self):
-        """Load all API keys and create OpenAI clients"""
-        # Split keys by comma
-        api_keys = os.getenv("CEREBRAS_API_KEYS", "").split(",")
-        base_url = os.getenv("CEREBRAS_BASE_URL")
+        """Load all API keys into a list"""
+        keys_str = os.getenv("CEREBRAS_API_KEYS", "")
+        self.api_keys = [k.strip() for k in keys_str.split(",") if k.strip()]
         
-        for key in api_keys:
-            key = key.strip()
-            if key:
-                try:
-                    # Create a standard OpenAI client for this key
-                    client = OpenAI(
-                        base_url=base_url,
-                        api_key=key
-                    )
-                    self.clients.append(client)
-                except Exception as e:
-                    logger.error(f"Failed to initialize key {key[:4]}...: {e}")
+        if not self.api_keys:
+            logger.error("No API keys found in CEREBRAS_API_KEYS")
 
-    def get_next_instance(self) -> Optional[Tuple[OpenAI, str]]:
+    def get_next_instance(self) -> Optional[Dict[str, str]]:
         """
-        Get next client in rotation.
-        Returns: Tuple (OpenAI_Client, Model_Name)
+        Returns a dictionary with the credentials for the next instance.
+        Returns: {'api_key': str, 'base_url': str, 'model': str}
         """
-        if not self.clients:
+        if not self.api_keys:
             return None
         
-        # Get current client
-        client = self.clients[self.current_index]
+        # Get current key
+        key = self.api_keys[self.current_index]
         
-        # Rotate index for the next call (Round Robin)
-        self.current_index = (self.current_index + 1) % len(self.clients)
+        # Rotate index for the next call
+        self.current_index = (self.current_index + 1) % len(self.api_keys)
         
-        return client, self.model_name
+        return {
+            "api_key": key,
+            "base_url": self.base_url,
+            "model": self.model_name
+        }
     
     def get_total_instances(self) -> int:
-        """Return total number of active clients available"""
-        return len(self.clients)
-
-    def chat_completion_with_retry(self, messages: list, **kwargs):
-        """
-        Helper function that automatically retries across all instances
-        if one fails.
-        """
-        total_attempts = self.get_total_instances()
-        
-        for attempt in range(total_attempts):
-            client, model = self.get_next_instance()
-            
-            try:
-                # Execute the API call
-                response = client.chat.completions.create(
-                    model=model,
-                    messages=messages,
-                    **kwargs
-                )
-                return response
-            except Exception as e:
-                logger.warning(f"Instance failed (Attempt {attempt+1}/{total_attempts}): {e}")
-                # Loop continues to next instance automatically
-                continue
-                
-        raise RuntimeError(f"All {total_attempts} instances failed.")
+        return len(self.api_keys)

orc_agent_main_cerebras.py

@@ -1,20 +1,20 @@
+import os
 import logging
-from typing import List, Any
-from dotenv import load_dotenv
+from typing import List, Any, Dict
 
 # Pydantic AI imports
 from pydantic_ai import Agent
-from pydantic_ai.models.openai import OpenAIModel  # <--- Essential fix
+from pydantic_ai.models.openai import OpenAIModel
 
-# OpenAI imports
+# Error handling
 from openai import RateLimitError, APIError  
 
 # Local application imports
 from csv_service import get_csv_basic_info
 from orchestrator_functions import csv_chart, csv_chat
 from cerebras_instance_provider import InstanceProvider
+from dotenv import load_dotenv
 
-# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -28,13 +28,8 @@ instance_provider = InstanceProvider()
 # ------------------------------------------------------------------
 
 async def generate_csv_answer(csv_url: str, user_questions: List[str]) -> Any:
-    """
-    Generates answers for user questions using the CSV URL.
-    """
-    logger.info("LLM using the csv chat function....")
-    logger.info(f"CSV URL: {csv_url}")
-    logger.info(f"User question: {user_questions}")
-
+    """Generates answers for user questions using the CSV URL."""
+    logger.info(f"Tool: generate_csv_answer | Questions: {user_questions}")
     answers = []
     for question in user_questions:
         answer = await csv_chat(csv_url, question)
@@ -42,140 +37,95 @@ async def generate_csv_answer(csv_url: str, user_questions: List[str]) -> Any:
     return answers
 
 async def generate_chart(csv_url: str, user_questions: List[str], chat_id: str) -> Any:
-    """
-    Generates charts for user questions using the CSV URL.
-    """
-    logger.info("LLM using the csv chart function....")
-    logger.info(f"CSV URL: {csv_url}")
-    logger.info(f"User question: {user_questions}")
-
+    """Generates charts for user questions using the CSV URL."""
+    logger.info(f"Tool: generate_chart | Questions: {user_questions}")
     charts = []
     for question in user_questions:
-        chart = await csv_chart(csv_url, question, chat_id)
-        charts.append(dict(question=question, image_url=chart))
-    
+        chart_url = await csv_chart(csv_url, question, chat_id)
+        charts.append(dict(question=question, image_url=chart_url))
     return charts
 
 # ------------------------------------------------------------------
-# 2. AGENT CREATION (FIXED)
+# 2. AGENT CREATION
 # ------------------------------------------------------------------
 
 def create_orchestrator_agent(csv_url: 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)
+    """Create a PydanticAI agent with a specific API Key instance"""
     
-    system_prompt = f"""
-# Role: Data Analyst Assistant
-**Specialization:** CSV Analysis & Visualization
-
-## Critical Rules:
-
-### 1. Tool Usage - MANDATORY
-- You MUST use `generate_csv_answer` tool for ALL data analysis questions
-- You MUST use `generate_chart` tool ONLY when explicitly asked for visualization, graph, chart, or plot
-- NEVER generate image markdown syntax (![...](url)) unless you have called `generate_chart` tool and received a real URL
-- NEVER fabricate or create placeholder image URLs
-
-### 2. When to Generate Visualizations
-**ONLY create visualizations when the user explicitly requests:**
-- "show me a chart/graph/plot"
-- "visualize this data"
-- "create a visualization"
-- "plot the data"
-- Any similar explicit visualization request
-
-**DO NOT create visualizations for:**
-- Simple data retrieval questions (e.g., "how many rows?", "what is the average?")
-- Questions that can be answered with text alone
-- Questions that don't explicitly ask for visual representation
-
-### 3. Response Format
-- For questions WITHOUT visualization request: Provide only the textual answer from `generate_csv_answer`
-- For questions WITH visualization request: Provide both textual answer AND call `generate_chart`, then include the image using the URL returned by the tool
-
-### 4. Output Guidelines
-- Use markdown formatting for text responses
-- Only include image syntax `![Description](url)` if you actually called `generate_chart` and got a real URL back
-- Provide clear, concise answers with explanations
-- Never mention tool names to the user
-
-## Current Context:
-- **Dataset:** {csv_url}
-- **Metadata:** {csv_metadata}
-- **History:** {conversation_history}
-- **Chat ID:** {chat_id}
-"""
+    # 1. Get credentials dictionary from provider
+    instance_config = instance_provider.get_next_instance()
     
-    # ---------------------------------------------------------
-    # FIX: Unpack tuple and use OpenAIModel wrapper
-    # ---------------------------------------------------------
-    instance_data = instance_provider.get_next_instance()
+    if instance_config is None:
+        raise RuntimeError("No available API instances (Check CEREBRAS_API_KEYS)")
     
-    if instance_data is None:
-        raise RuntimeError("No available API instances")
+    # 2. Create the Model using standard arguments.
+    # We pass api_key and base_url directly. PydanticAI will handle the client creation.
+    model = OpenAIModel(
+        instance_config['model'],
+        base_url=instance_config['base_url'],
+        api_key=instance_config['api_key'],
+    )
     
-    # Unpack the tuple (client, model_name)
-    client, model_name = instance_data
+    csv_metadata = get_csv_basic_info(csv_url)
     
-    # Create the Pydantic AI Model using the specific client for this key
-    model = OpenAIModel(model_name, openai_client=client)
+    system_prompt = f"""
+# Role: Data Analyst Assistant
+**Context:** Analyzing CSV: {csv_url}
+**Metadata:** {csv_metadata}
+**Chat ID:** {chat_id}
+
+## Rules:
+1. Use `generate_csv_answer` for text questions.
+2. Use `generate_chart` ONLY if explicitly asked for visual/plot/graph.
+3. Output format: Markdown. If chart generated, use ![Desc](url).
+"""
     
     return Agent(
         model=model,
         deps_type=str,
         tools=[generate_csv_answer, generate_chart],
         system_prompt=system_prompt,
-        retries=0
+        retries=0 # We handle retries manually in the loop below
     )
 
 # ------------------------------------------------------------------
-# 3. ORCHESTRATOR LOGIC (RETRY/FAILOVER)
+# 3. ORCHESTRATOR LOGIC (RETRY LOOP)
 # ------------------------------------------------------------------
 
 def csv_orchestrator_chat_cerebras(csv_url: str, user_question: str, conversation_history: List, chat_id: str) -> str:
     """
-    CSV orchestrator with automatic failover on 429/API errors using InstanceProvider.
+    Orchestrator that rebuilds the agent with a new key upon failure.
     """
-    logger.info(f"CSV URL: {csv_url}")
-    logger.info(f"User questions: {user_question}")
+    logger.info(f"Starting Orchestrator | Query: {user_question}")
 
     max_attempts = instance_provider.get_total_instances()
-    
-    # Try with different instances until one works
+    if max_attempts == 0:
+         return "System Error: No API keys configured."
+
+    # Loop through available keys
     for attempt in range(max_attempts):
         try:
-            logger.info(f"Attempt {attempt + 1}/{max_attempts}")
-            
-            # Create agent (this internally rotates to the next key)
+            # 1. Create a NEW agent (this fetches the NEXT key automatically)
             agent = create_orchestrator_agent(csv_url, conversation_history, chat_id)
             
-            # Run the agent
+            logger.info(f"Attempt {attempt + 1}/{max_attempts} using key ending in ...{agent.model.client.api_key[-4:] if hasattr(agent.model, 'client') else '****'}")
+
+            # 2. Run the agent
             result = agent.run_sync(user_question)
             
-            logger.info(f"✓ Success with instance {attempt + 1}")
-            logger.info(f"Orchestrator Result: {result.data}")
-            
+            logger.info(f"✓ Success on attempt {attempt + 1}")
             return result.data
             
-        except RateLimitError:
-            logger.error(f"✗ Rate limit (429) hit for instance {attempt + 1}/{max_attempts}")
-            if attempt == max_attempts - 1:
-                raise RuntimeError(f"All {max_attempts} instances failed with rate limits")
-            logger.info("Trying next instance...")
-            continue
-            
-        except APIError as e:
-            logger.error(f"✗ API error with instance {attempt + 1}: {str(e)}")
-            if attempt == max_attempts - 1:
-                raise RuntimeError(f"All {max_attempts} instances failed. Last error: {str(e)}")
-            logger.info("Trying next instance...")
+        except (RateLimitError, APIError) as e:
+            logger.warning(f"⚠ API Error on attempt {attempt + 1}: {e}")
+            logger.info("Rotating to next instance...")
+            # The loop continues, calling create_orchestrator_agent() again, getting the next key.
             continue
             
         except Exception as e:
-            logger.error(f"✗ Unexpected error with instance {attempt + 1}: {str(e)}")
-            if attempt == max_attempts - 1:
-                raise RuntimeError(f"All {max_attempts} instances failed. Last error: {str(e)}")
-            logger.info("Trying next instance...")
+            # Catch unexpected Pydantic/Python errors
+            logger.error(f"✗ Unexpected Error on attempt {attempt + 1}: {e}")
+            logger.info("Rotating to next instance...")
             continue
     
-    raise RuntimeError(f"Failed after {max_attempts} attempts")
+    raise RuntimeError(f"Failed to generate response after {max_attempts} attempts.")