src/utils/groq_client.py

# src/utils/groq_client.py
import os
from groq import Groq, RateLimitError, APIError
import time

# Initialize Groq client using environment variable
# Ensure GROQ_API_KEY is set in the environment where the Flask app runs
client = Groq(
    api_key=os.environ.get("GROQ_API_KEY"),
)

# Define a default model
DEFAULT_MODEL = "llama3-8b-8192" # Or choose another suitable model like llama3-70b-8192, mixtral-8x7b-32768

def call_groq(prompt: str, system_prompt: str = "You are a helpful assistant specializing in legislative analysis and drafting.", model: str = DEFAULT_MODEL, max_retries: int = 3, initial_delay: int = 1) -> str | None:
    """Calls the Groq API with a given prompt and handles basic errors/retries."""
    
    if not client.api_key:
        print("Error: GROQ_API_KEY environment variable not set.")
        return "Error: Groq API key not configured."
        
    retries = 0
    delay = initial_delay
    while retries < max_retries:
        try:
            chat_completion = client.chat.completions.create(
                messages=[
                    {
                        "role": "system",
                        "content": system_prompt,
                    },
                    {
                        "role": "user",
                        "content": prompt,
                    }
                ],
                model=model,
                # Optional parameters:
                # temperature=0.7,
                # max_tokens=1024,
                # top_p=1,
                # stop=None,
                # stream=False,
            )
            return chat_completion.choices[0].message.content
        except RateLimitError as e:
            retries += 1
            print(f"Rate limit exceeded. Retrying in {delay} seconds... (Attempt {retries}/{max_retries})")
            time.sleep(delay)
            delay *= 2 # Exponential backoff
        except APIError as e:
            print(f"Groq API Error: {e.status_code} - {e.message}")
            # Depending on the error, you might want to retry or just return None/error message
            return f"Error: Groq API request failed ({e.status_code})."
        except Exception as e:
            print(f"An unexpected error occurred calling Groq: {e}")
            return f"Error: An unexpected error occurred while contacting the AI model."

    print(f"Failed to call Groq API after {max_retries} retries.")
    return "Error: AI model request failed after multiple retries due to rate limits."

# --- Specific Task Functions ---

def generate_draft_content(policy_intent: str) -> str | None:
    """Uses Groq to generate initial draft content based on policy intent."""
    system_prompt = "You are an AI assistant specialized in legislative drafting. Based on the provided policy intent, generate a concise initial draft for a new regulation or law. Focus on clear, actionable language suitable for legislation."
    prompt = f"Policy Intent: {policy_intent}\n\nGenerate an initial legislative draft based on this intent:"
    return call_groq(prompt, system_prompt)

def perform_impact_analysis(document_content: str, analysis_type: str) -> tuple[str | None, float | None, str | None]:
    """Uses Groq to perform impact analysis on a given document.
    Returns: (predicted_impact, confidence_score, rationale)
    """
    system_prompt = f"You are an AI legislative analyst. Analyze the provided legislative text and predict its potential {analysis_type.lower()} impact. Provide a summary of the predicted impact, a confidence score (0.0 to 1.0), and a brief rationale for your prediction."
    prompt = f"Legislative Text:\n```\n{document_content}\n```\n\nAnalyze the {analysis_type.lower()} impact of this text. Format your response strictly as follows:\nPredicted Impact: [Your summary here]\nConfidence Score: [A number between 0.0 and 1.0]\nRationale: [Your explanation here]"
    
    response = call_groq(prompt, system_prompt)
    
    if response and response.startswith("Predicted Impact:"):
        try:
            lines = response.split("\n")
            impact = lines[0].replace("Predicted Impact:", "").strip()
            score_line = next((line for line in lines if line.startswith("Confidence Score:")), None)
            rationale_line = next((line for line in lines if line.startswith("Rationale:")), None)
            
            score = None
            if score_line:
                try:
                    score_str = score_line.replace("Confidence Score:", "").strip()
                    score = float(score_str)
                    if not (0.0 <= score <= 1.0):
                        score = None # Invalid score
                except ValueError:
                    score = None # Could not parse score
            
            rationale = None
            if rationale_line:
                 # Find the start index of Rationale and take everything after it
                rationale_start_index = response.find("Rationale:")
                if rationale_start_index != -1:
                    rationale = response[rationale_start_index + len("Rationale:"):].strip()
                else: # Fallback if split didn't work as expected
                     rationale = rationale_line.replace("Rationale:", "").strip()

            return impact, score, rationale
        except Exception as e:
            print(f"Error parsing Groq impact analysis response: {e}")
            return f"Error parsing AI response: {response}", None, None
    else:
        # Return the raw response or error message if parsing failed
        return response or "Error: No response from AI model.", None, None

def generate_recommendation(criteria: str) -> tuple[str | None, str | None, str | None, str | None]:
    """Uses Groq to generate a legislative recommendation based on criteria.
    Returns: (recommendation_type, details, rationale, priority)
    """
    system_prompt = "You are an AI legislative analyst. Based on the provided criteria, suggest a legislative recommendation (e.g., Update, Consolidation, Removal). Provide details, rationale, and a priority level (Low, Medium, High)."
    prompt = f"Criteria for Recommendation: {criteria}\n\nGenerate a legislative recommendation based on these criteria. Format your response strictly as follows:\nType: [Update/Consolidation/Removal]\nDetails: [Specific recommendation details]\nRationale: [Justification for the recommendation]\nPriority: [Low/Medium/High]"
    
    response = call_groq(prompt, system_prompt)
    
    if response and response.startswith("Type:"):
        try:
            lines = response.split("\n")
            rec_type = lines[0].replace("Type:", "").strip()
            details_line = next((line for line in lines if line.startswith("Details:")), None)
            rationale_line = next((line for line in lines if line.startswith("Rationale:")), None)
            priority_line = next((line for line in lines if line.startswith("Priority:")), None)
            
            details = details_line.replace("Details:", "").strip() if details_line else None
            rationale = rationale_line.replace("Rationale:", "").strip() if rationale_line else None
            priority = priority_line.replace("Priority:", "").strip() if priority_line else None
            
            # Basic validation for priority
            if priority not in ["Low", "Medium", "High"]:
                priority = "Medium" # Default if invalid
                
            return rec_type, details, rationale, priority
        except Exception as e:
            print(f"Error parsing Groq recommendation response: {e}")
            return "Error", f"Error parsing AI response: {response}", None, "Medium"
    else:
        return "Error", response or "Error: No response from AI model.", None, "Medium"