PoliSage / src /utils /groq_client.py
yasserrmd's picture
Upload 80 files
0a40ab8 verified
# 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"