models/gemini.py

import os
from typing import Optional
import sys

import vertexai
from dotenv import load_dotenv

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

load_dotenv()
VERTEXAI_PROJECT = os.environ["VERTEXAI_PROJECT"]


def get_gemini_response(prompt_text, model, parameters: Optional = None, location: Optional[str] = "us-central1") -> str:
    initialize_vertexai_params()

    if model is None or parameters is None:
        model = "gemini-1.0-pro"
    model = GenerativeModel(model)

    model_response = model.generate_content(prompt_text)

    return model_response.text