llm/gemini_client.py

import os

import google.generativeai as genai


class GeminiClient:
    def __init__(self, system_message=None):
        self._system_message = system_message
        self._connect_client()

    def _connect_client(self):
        if not os.getenv("GOOGLE_PALM_KEY"):
            raise Exception("Please set your Google MakerSuite API key")

        api_key = os.getenv("GOOGLE_PALM_KEY")
        genai.configure(api_key=api_key)

        safety_settings = [
            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
            {
                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                "threshold": "BLOCK_ONLY_HIGH",
            },
            {
                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                "threshold": "BLOCK_ONLY_HIGH",
            },
        ]

        defaults = {
            "temperature": 0.7,
            "top_k": 40,
            "top_p": 0.95,
            "max_output_tokens": 1024,
        }

        self._model = genai.GenerativeModel(
            model_name="gemini-pro",
            generation_config=defaults,
            safety_settings=safety_settings,
        )

    def generate_text(self, prompt: str) -> str:
        full_prompt = self._system_message + prompt
        try:
            response = self._model.generate_content(full_prompt).text
        except Exception as e:
            print(f"Error: {e}")
            response = ""
        return response