from os import getenv
from pprint import pprint

from dotenv import load_dotenv
from openai import OpenAI

from srs2std.llms.llm import LLM, Prompt

load_dotenv()
client = OpenAI(
    base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
    api_key=getenv("GEMINI_API_KEY"),
)


class Gemini_2_5_Pro(LLM):
    client = client
    model = "gemini-2.5-pro"
    max_completion_tokens = 65_535
    max_retries = 5
    temperature = 0


class Gemini_2_5_Flash(LLM):
    client = client
    model = "gemini-2.5-flash"
    max_completion_tokens = 65_535
    max_retries = 5
    temperature = 0


if __name__ == "__main__":
    gemini = Gemini_2_5_Flash()
    response = gemini.ask(Prompt(system="", user="Explain AI in short."))
    pprint(response)
