go

orc_agent_main_cerebras.py

@@ -1,48 +1,41 @@
-import os
+import logging
 from typing import List, Any
+from dotenv import load_dotenv
+
+# Pydantic AI imports
 from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIModel  # <--- Essential fix
+
+# OpenAI imports
 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
-import logging
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 load_dotenv()
 
 # Initialize the instance provider
 instance_provider = InstanceProvider()
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# ------------------------------------------------------------------
+# 1. DEFINE TOOLS
+# ------------------------------------------------------------------
 
-# 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."}
-        ]
+    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}")
 
-    # 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))
@@ -50,38 +43,22 @@ async def generate_csv_answer(csv_url: str, user_questions: List[str]) -> Any:
 
 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"}
-        ]
+    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}")
 
-    # 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
 
+# ------------------------------------------------------------------
+# 2. AGENT CREATION (FIXED)
+# ------------------------------------------------------------------
 
 def create_orchestrator_agent(csv_url: str, conversation_history: List, chat_id: str) -> Agent:
     """Create a PydanticAI agent configured for CSV analysis using Cerebras"""
@@ -127,29 +104,22 @@ def create_orchestrator_agent(csv_url: str, conversation_history: List, chat_id:
 - **Metadata:** {csv_metadata}
 - **History:** {conversation_history}
 - **Chat ID:** {chat_id}
-
-## Example Behavior:
-
-**Question: "How many rows are in the dataset?"**
-✅ Correct Response: "The dataset contains 1,000 rows."
-❌ Wrong Response: "The dataset contains 1,000 rows. ![Number of Rows](url)"
-
-**Question: "Show me a chart of the distribution of ages"**
-✅ Correct Response: "The age distribution shows... [call generate_chart tool] ![Age Distribution](actual_url_from_tool)"
-❌ Wrong Response: "The age distribution shows... ![Age Distribution](https://example.com/chart.png)"
-
-**Remember:** 
-- Always generate fresh, tool-assisted responses
-- Never reuse previous answers
-- Never create fake image URLs
-- Only use visualization when explicitly requested
 """
     
-    # Get next available model instance
-    model = instance_provider.get_next_instance()
-    if model is None:
+    # ---------------------------------------------------------
+    # FIX: Unpack tuple and use OpenAIModel wrapper
+    # ---------------------------------------------------------
+    instance_data = instance_provider.get_next_instance()
+    
+    if instance_data is None:
         raise RuntimeError("No available API instances")
     
+    # Unpack the tuple (client, model_name)
+    client, model_name = instance_data
+    
+    # Create the Pydantic AI Model using the specific client for this key
+    model = OpenAIModel(model_name, openai_client=client)
+    
     return Agent(
         model=model,
         deps_type=str,
@@ -158,11 +128,13 @@ def create_orchestrator_agent(csv_url: str, conversation_history: List, chat_id:
         retries=0
     )
 
+# ------------------------------------------------------------------
+# 3. ORCHESTRATOR LOGIC (RETRY/FAILOVER)
+# ------------------------------------------------------------------
 
 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 errors using InstanceProvider.
-    Follows the same pattern as query_csv_agent_cerebras.
+    CSV orchestrator with automatic failover on 429/API errors using InstanceProvider.
     """
     logger.info(f"CSV URL: {csv_url}")
     logger.info(f"User questions: {user_question}")
@@ -174,10 +146,10 @@ def csv_orchestrator_chat_cerebras(csv_url: str, user_question: str, conversatio
         try:
             logger.info(f"Attempt {attempt + 1}/{max_attempts}")
             
-            # Create agent with next instance
+            # Create agent (this internally rotates to the next key)
             agent = create_orchestrator_agent(csv_url, conversation_history, chat_id)
             
-            # Run the agent - this is where rate limits typically occur
+            # Run the agent
             result = agent.run_sync(user_question)
             
             logger.info(f"✓ Success with instance {attempt + 1}")
@@ -185,30 +157,24 @@ def csv_orchestrator_chat_cerebras(csv_url: str, user_question: str, conversatio
             
             return result.data
             
-        except RateLimitError as e:
+        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...")
             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...")
             continue