helper.py

# helper.py
import streamlit as st
import os
import google.generativeai as genai

# Configure the API key
try:
    genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
except AttributeError:
    st.error("Please set the GOOGLE_API_KEY environment variable.")
    st.stop()

# Set up safety settings to be less restrictive
# This can help prevent the API from blocking reasonable prompts
safety_settings = [
    {
        "category": "HARM_CATEGORY_HARASSMENT",
        "threshold": "BLOCK_NONE"
    },
    {
        "category": "HARM_CATEGORY_HATE_SPEECH",
        "threshold": "BLOCK_NONE"
    },
    {
        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
        "threshold": "BLOCK_NONE"
    },
    {
        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
        "threshold": "BLOCK_NONE"
    },
]

# Set up the model configuration
generation_config = {
  "temperature": 0.7,
  "top_p": 1,
  "top_k": 1,
  "max_output_tokens": 2048, # You can adjust this as needed
}

# Initialize the Gemini model
model = genai.GenerativeModel(
    model_name="gemini-2.0-flash-lite",
    generation_config=generation_config,
    safety_settings=safety_settings
)

def call_gemini(prompt: str) -> str:
    """
        Send a prompt to the Gemini model and return the response.
        Args:
            prompt (str): The input prompt to send to the Gemini model.
        Returns:
            str: The response from the Gemini model.
    """
    try:
        # The Gemini API uses model.generate_content
        response = model.generate_content(prompt)
        
        # Return the text part of the response
        return response.text
    
    except Exception as e:
        # Handle potential errors (e.g., content blocked by safety)
        st.error(f"Error calling Gemini: {e}")
        # Check if the response was blocked
        if "response.prompt_feedback" in locals() or "response.prompt_feedback" in globals():
             st.error(f"Prompt Feedback: {response.prompt_feedback}")
        return "An error occurred, and I couldn't generate a response."