research_agent/llm_utils.py

import json
import google.generativeai as genai
from google.generativeai.types import GenerationConfig, HarmCategory, HarmBlockThreshold

def run_gemini_json_completion(model, prompt: str, temperature: float):
    """Runs a Gemini call expecting a JSON response."""
    try:
        response = model.generate_content(prompt, generation_config=GenerationConfig(response_mime_type="application/json", temperature=temperature))
        return json.loads(response.text)
    except Exception as e:
        print(f"Warning: Failed to parse JSON from Gemini. Error: {e}")
        return {}

def run_gemini_text_completion(model, prompt: str, temperature: float):
    """Runs a standard Gemini text completion call."""
    generation_config = GenerationConfig(temperature=temperature)
    safety_settings = [{"category": HarmCategory.HARM_CATEGORY_HARASSMENT, "threshold": HarmBlockThreshold.BLOCK_ONLY_HIGH}]
    try:
        response = model.generate_content(prompt, generation_config=generation_config, safety_settings=safety_settings)
        return response.text
    except Exception as e:
        return f"[Error: Could not generate response. Details: {e}]"