import os

from openai import OpenAI
from .openai import LLM_GPT_OpenAI


class LLM_Maritalk(LLM_GPT_OpenAI):
    def __init__(self, model_size):
        """Constructor
        Arguments:
            model_size - Maritaca model to use to make LLM calls
        """
        if model_size == "Sabia3 Maritaca":
            self.model_id = "sabia-3.1"
            self.llm_description = (
                "Sabia-3.1 (medium-sized LLM) - directly from Maritaca"
            )
            self.price_per_M_input_tokens = 0.95
            self.price_per_M_output_tokens = 1.9

        try:
            self.openai_client = OpenAI(
                api_key=os.environ.get("MARITACA_API_KEY"),
                base_url="https://chat.maritaca.ai/api",
            )
        except Exception:
            self.openai_client = None

        self.config = {
            # "messages": prompt,
            # "system": sysprompt,
            "max_tokens": 4000,
            "temperature": 0.5,  # 0.5 is default,
            "stream": True,
            # "top_k": 250,
            # "top_p": 1,
            "stop": None,  # the regular is already implemented
            "model": self.model_id,
        }
        # requests and answer word count
        self.word_counts = []
