from dotenv import load_dotenv
import os
import google.generativeai as genai
from helpers import constants

load_dotenv()

# name of the AI model used in this call.
TEXT_MODEL_NAME_LATEST = os.getenv("TEXT_MODEL_NAME_LATEST")
PDF_PAGE_IMAGES_PATH = os.getenv("PDF_PAGE_IMAGES_PATH")
PROMPT_PATH = constants.PROMPT_PATH


class GeminiService:
    generation_config_json: str = {
        'temperature': 0.9,
        'top_p': 1,
        'top_k': 40,
        'max_output_tokens': 2048,
        'stop_sequences': [],
        "response_mime_type": "application/json",
    }

    generation_config: str = {
        'temperature': 0.9,
        'top_p': 1,
        'top_k': 40,
        'max_output_tokens': 2048,
        'stop_sequences': [],
    }

    safety_settings: list[str] = [{"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"}]

    def __init__(self, prompt_path, as_json: bool = True):
        genai.configure(api_key=os.getenv("API_KEY"))
        self.model = genai.GenerativeModel(model_name=os.getenv("TEXT_MODEL_NAME_LATEST"),
                                           system_instruction=GeminiService.read_prompt(prompt_path),
                                           generation_config=GeminiService.generation_config_json if as_json else GeminiService.generation_config,
                                           safety_settings=self.safety_settings)

    @staticmethod
    def read_prompt(prompt_file):
        with open(prompt_file, 'r', encoding="utf-8") as file:
            # Load the JSON data from the file into a Python dictionary
            return file.read()

    def single_completion(self, request) -> str:
        return self.model.generate_content(request).text




if __name__ == '__main__':
    service = GeminiService( os.path.join("..", constants.PROMPT_PATH, "prompt.txt"))
    prompt = "What is the meaning of life"
    response = service.single_completion(prompt)
    print(response)

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1721829902.342524    9100 config.cc:230] gRPC experiments enabled: call_status_override_on_cancellation, event_engine_client, event_engine_dns, event_engine_listener, http2_stats_fix, monitoring_experiment, pick_first_new, trace_record_callops, work_serializer_clears_time_cache