# filename: language_models.py

"""
Module for defining various language model classes and their invocation methods.
"""

import os
import logging
from abc import ABC, abstractmethod
from time import time
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_google_genai import ChatGoogleGenerativeAI
from log_config import get_logger

# Configure logging
logger = get_logger('LanguageModels')

def get_api_key(service_name: str) -> str:
    """
    Retrieve the API key from environment variables.

    Args:
        service_name (str): The name of the service.

    Returns:
        str: The API key for the service.

    Raises:
        ValueError: If the API key is not set in environment variables.
    """
    api_key = os.getenv(service_name)
    if not api_key:
        logger.error(f"{service_name} environment variable not set.")
        raise ValueError(f"API key for {service_name} is required")
    return api_key

class LanguageModel(ABC):
    """
    Abstract base class for language models.
    Defines a common interface for different AI language models.
    """

    @abstractmethod
    def invoke(self, prompt: str) -> str:
        """
        Send a prompt to the language model and return the response.

        Args:
            prompt (str): The text prompt to send to the language model.

        Returns:
            str: The response from the language model.

        Raises:
            Exception: Propagates any exceptions that occur during invocation.
        """
        pass

class GPT4Model(LanguageModel):
    """
    A language model class for OpenAI's GPT-4.
    """
    def __init__(self):
        api_key = get_api_key("OPENAI_API_KEY")
        self.model = ChatOpenAI(
            model=os.getenv('MODEL_NAME_GPT4', 'gpt-4o'),
            temperature=0,
            openai_api_key=api_key
        )

    def invoke(self, prompt: str) -> str:
        if not prompt:
            logger.error("Empty prompt provided.")
            raise ValueError("Prompt cannot be empty.")
        start_time = time()
        try:
            response = self.model.invoke([prompt])
            logger.info(f"Invocation successful in {time() - start_time:.2f} seconds.")
            return response
        except Exception as e:
            logger.error(f"Failed to invoke GPT-4 with prompt: '{prompt}', error: {str(e)}")
            raise

class ClaudeModel(LanguageModel):
    """
    A language model class for Anthropic's Claude.
    """
    def __init__(self):
        api_key = get_api_key("ANTHROPIC_API_KEY")
        self.model = ChatAnthropic(
            model=os.getenv('MODEL_NAME_CLAUDE_OPUS', 'claude-3-opus-20240229'),
            temperature=0,
            anthropic_api_key=api_key
        )

    def invoke(self, prompt: str) -> str:
        if not prompt:
            logger.error("Empty prompt provided.")
            raise ValueError("Prompt cannot be empty.")
        start_time = time()
        try:
            response = self.model.invoke([prompt])
            logger.info(f"Invocation successful in {time() - start_time:.2f} seconds.")
            return response
        except Exception as e:
            logger.error(f"Failed to invoke Claude with prompt: '{prompt}', error: {str(e)}")
            raise

class GeminiModel(LanguageModel):
    """
    A language model class for Google's Generative AI, Gemini.
    """
    def __init__(self):
        api_key = get_api_key("GOOGLE_API_KEY")
        self.model = ChatGoogleGenerativeAI(
            model=os.getenv('MODEL_NAME_GEMINI', 'gemini-1.5-pro-latest'),
            api_key=api_key
        )

    def invoke(self, prompt: str) -> str:
        if not prompt:
            logger.error("Empty prompt provided.")
            raise ValueError("Prompt cannot be empty.")
        start_time = time()
        try:
            response = self.model.invoke([prompt])
            logger.info(f"Invocation successful in {time() - start_time:.2f} seconds.")
            return response
        except Exception as e:
            logger.error(f"Failed to invoke Gemini with prompt: '{prompt}', error: {str(e)}")
            raise

# file: language_models.py (end)