from langchain_openai import ChatOpenAI
from sonika_langchain_bot.langchain_class import ILanguageModel
from langchain_core.messages import HumanMessage
from typing import Generator


class OpenAILanguageModel(ILanguageModel):
    """
    Clase que implementa la interfaz ILanguageModel para interactuar con los modelos de lenguaje de OpenAI.
    Proporciona funcionalidades para generar respuestas y contar tokens.
    """

    def __init__(self, api_key: str, model_name: str = "gpt-4o-mini", temperature: float = 0.7):
        """
        Inicializa el modelo de lenguaje de OpenAI.
        
        Args:
            api_key (str): Clave API de OpenAI
            model_name (str): Nombre del modelo a utilizar
            temperature (float): Temperatura para la generación de respuestas
        """
        self.model = ChatOpenAI(temperature=temperature, model_name=model_name, api_key=api_key)

    def predict(self, prompt: str) -> str:
        """
        Genera una respuesta basada en el prompt proporcionado.
        
        Args:
            prompt (str): Texto de entrada para generar la respuesta
            
        Returns:
            str: Respuesta generada por el modelo
        """
        return self.model.predict(prompt)
    
    def invoke(self, prompt: str) -> str:
        """
        Invokes the language model with a given prompt and returns the generated response.

        Args:
            prompt (str): The input text to be processed by the language model.

        Returns:
            str: The response generated by the language model based on the provided prompt.
        """
        message = HumanMessage(content=prompt)
        response = self.model.invoke([message])
        return response
    
    def stream_response(self, prompt: str) -> Generator[str, None, None]:
        """
        Genera una respuesta en streaming basada en el prompt proporcionado.

        Args:
            prompt (str): Texto de entrada para generar la respuesta
        
        Yields:
            str: Fragmentos de la respuesta generada por el modelo en tiempo real
        """
        message = HumanMessage(content=prompt)
        for chunk in self.model.stream([message]):
            yield chunk.content