models/gemini.py

from typing import Optional
import sys

sys.path.append("../")
from setup.vertex_ai_setup import initialize_vertexai_params
from vertexai.preview.generative_models import GenerativeModel


def get_gemini_response(prompt_text, model, parameters: Optional = None) -> str:
    initialize_vertexai_params()
    if model is None or parameters is None:
        model = "gemini-1.0-pro"
        parameters = {
            "max_output_tokens": 1024,
            "temperature": 0.2
        }
    model = GenerativeModel.from_pretrained(model, **parameters)
    model_response = model.generate_content(prompt_text)
    return model_response