"""
This file contains utility functions used in repo scripts

For API calling, we support:
    - parallelized model calls on the same prompt 
    - batched model calls on different prompts

NOTE: any function beginning with 'a', e.g. 'afetch_together_api_response', is an asynchronous function
"""

import concurrent.futures
from typing import List, Mapping, Dict, Union, Optional, Literal, Any
import asyncio
import litellm
import pydantic
import pprint
import os
from dotenv import load_dotenv 

from judgeval.clients import async_together_client, together_client
from judgeval.constants import *
from judgeval.common.logger import debug, error

LITELLM_SUPPORTED_MODELS = set(litellm.model_list)

class CustomModelParameters(pydantic.BaseModel):
    model_name: str
    secret_key: str
    litellm_base_url: str
    
    @pydantic.field_validator('model_name')
    def validate_model_name(cls, v):
        if not v:
            raise ValueError("Model name cannot be empty")
        return v
    
    @pydantic.field_validator('secret_key')
    def validate_secret_key(cls, v):
        if not v:
            raise ValueError("Secret key cannot be empty")
        return v
    
    @pydantic.field_validator('litellm_base_url')
    def validate_litellm_base_url(cls, v):
        if not v:
            raise ValueError("Litellm base URL cannot be empty")
        return v

class ChatCompletionRequest(pydantic.BaseModel):
    model: str
    messages: List[Dict[str, str]]
    response_format: Optional[Union[pydantic.BaseModel, Dict[str, Any]]] = None
    
    @pydantic.field_validator('messages')
    def validate_messages(cls, messages):
        if not messages:
            raise ValueError("Messages cannot be empty")
            
        for msg in messages:
            if not isinstance(msg, dict):
                raise TypeError("Message must be a dictionary")
            if 'role' not in msg:
                raise ValueError("Message missing required 'role' field")
            if 'content' not in msg:
                raise ValueError("Message missing required 'content' field")
            if msg['role'] not in ['system', 'user', 'assistant']:
                raise ValueError(f"Invalid role '{msg['role']}'. Must be 'system', 'user', or 'assistant'")
        
        return messages

    @pydantic.field_validator('model')
    def validate_model(cls, model):
        if not model:
            raise ValueError("Model cannot be empty")
        if model not in TOGETHER_SUPPORTED_MODELS and model not in LITELLM_SUPPORTED_MODELS:
            raise ValueError(f"Model {model} is not in the list of supported models.")
        return model
    
    @pydantic.field_validator('response_format', mode='before')
    def validate_response_format(cls, response_format):
        if response_format is not None:
            if not isinstance(response_format, (dict, pydantic.BaseModel)):
                raise TypeError("Response format must be a dictionary or pydantic model")
            # Optional: Add additional validation for required fields if needed
            # For example, checking for 'type': 'json' in OpenAI's format
        return response_format

os.environ['LITELLM_LOG'] = 'DEBUG'

load_dotenv()

def read_file(file_path: str) -> str:
    with open(file_path, "r", encoding='utf-8') as file:
        return file.read()


def fetch_together_api_response(model: str, messages: List[Mapping], response_format: pydantic.BaseModel = None) -> str:
    """
    Fetches a single response from the Together API for a given model and messages.
    """
    # Validate request
    if messages is None or messages == []:
        raise ValueError("Messages cannot be empty")
    
    request = ChatCompletionRequest(
            model=model,
            messages=messages,
            response_format=response_format
    )
    
    debug(f"Calling Together API with model: {request.model}")
    debug(f"Messages: {request.messages}")
    
    if request.response_format is not None:
        debug(f"Using response format: {request.response_format}")
        response = together_client.chat.completions.create(
            model=TOGETHER_SUPPORTED_MODELS.get(request.model),
            messages=request.messages,
            response_format=request.response_format
        )
    else:
        response = together_client.chat.completions.create(
            model=TOGETHER_SUPPORTED_MODELS.get(request.model),
            messages=request.messages,
        )
    
    debug(f"Received response: {response.choices[0].message.content[:100]}...")
    return response.choices[0].message.content


async def afetch_together_api_response(model: str, messages: List[Mapping], response_format: pydantic.BaseModel = None) -> str:
    """
    ASYNCHRONOUSLY Fetches a single response from the Together API for a given model and messages.
    """
    request = ChatCompletionRequest(
        model=model,
        messages=messages,
        response_format=response_format
    )
    
    debug(f"Calling Together API with model: {request.model}")
    debug(f"Messages: {request.messages}")
    
    if request.response_format is not None:
        debug(f"Using response format: {request.response_format}")
        response = await async_together_client.chat.completions.create(
            model=TOGETHER_SUPPORTED_MODELS.get(request.model),
            messages=request.messages,
            response_format=request.response_format
        )
    else:
        response = await async_together_client.chat.completions.create(
            model=TOGETHER_SUPPORTED_MODELS.get(request.model),
            messages=request.messages,
        )
    return response.choices[0].message.content


def query_together_api_multiple_calls(models: List[str], messages: List[List[Mapping]], response_formats: List[pydantic.BaseModel] = None) -> List[str]:
    """
    Queries the Together API for multiple calls in parallel

    Args:
        models (List[str]): List of models to query
        messages (List[List[Mapping]]): List of messages to query. Each inner object corresponds to a single prompt.
        response_formats (List[pydantic.BaseModel], optional): A list of the format of the response if JSON forcing. Defaults to None.

    Returns:
        List[str]: TogetherAI responses for each model and message pair in order. Any exceptions in the thread call result in a None.
    """
    # Check for empty models list
    if not models:
        raise ValueError("Models list cannot be empty")

    # Validate all models are supported
    for model in models:
        if model not in TOGETHER_SUPPORTED_MODELS:
            raise ValueError(f"Model {model} is not in the list of supported TogetherAI models: {TOGETHER_SUPPORTED_MODELS}.")

    # Validate input lengths match
    if response_formats is None:
        response_formats = [None] * len(models)
    if not (len(models) == len(messages) == len(response_formats)):
        raise ValueError("Number of models, messages, and response formats must be the same")

    # Validate message format
    validate_batched_chat_messages(messages)

    num_workers = int(os.getenv('NUM_WORKER_THREADS', MAX_WORKER_THREADS))  
    # Initialize results to maintain ordered outputs
    out = [None] * len(messages)
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
        # Submit all queries to together API with index, gets back the response content
        futures = {executor.submit(fetch_together_api_response, model, message, response_format): idx \
                   for idx, (model, message, response_format) in enumerate(zip(models, messages, response_formats))}
        
        # Collect results as they complete -- result is response content
        for future in concurrent.futures.as_completed(futures):
            idx = futures[future]
            try:
                out[idx] = future.result()
            except Exception as e:
                error(f"Error in parallel call {idx}: {str(e)}")
                out[idx] = None 
    return out


async def aquery_together_api_multiple_calls(models: List[str], messages: List[List[Mapping]], response_formats: List[pydantic.BaseModel] = None) -> List[str]:
    """
    Queries the Together API for multiple calls in parallel

    Args:
        models (List[str]): List of models to query
        messages (List[List[Mapping]]): List of messages to query. Each inner object corresponds to a single prompt.
        response_formats (List[pydantic.BaseModel], optional): A list of the format of the response if JSON forcing. Defaults to None.

    Returns:
        List[str]: TogetherAI responses for each model and message pair in order. Any exceptions in the thread call result in a None.
    """
    # Check for empty models list
    if not models:
        raise ValueError("Models list cannot be empty")

    # Validate all models are supported
    for model in models:
        if model not in TOGETHER_SUPPORTED_MODELS:
            raise ValueError(f"Model {model} is not in the list of supported TogetherAI models: {TOGETHER_SUPPORTED_MODELS}.")

    # Validate input lengths match
    if response_formats is None:
        response_formats = [None] * len(models)
    if not (len(models) == len(messages) == len(response_formats)):
        raise ValueError("Number of models, messages, and response formats must be the same")

    # Validate message format
    validate_batched_chat_messages(messages)

    debug(f"Starting parallel Together API calls for {len(messages)} messages")
    out = [None] * len(messages)
    
    async def fetch_and_store(idx, model, message, response_format):
        try:
            debug(f"Processing call {idx} with model {model}")
            out[idx] = await afetch_together_api_response(model, message, response_format)
        except Exception as e:
            error(f"Error in parallel call {idx}: {str(e)}")
            out[idx] = None

    tasks = [
        fetch_and_store(idx, model, message, response_format)
        for idx, (model, message, response_format) in enumerate(zip(models, messages, response_formats))
    ]
    
    await asyncio.gather(*tasks)
    debug(f"Completed {len(messages)} parallel calls")
    return out


def fetch_litellm_api_response(model: str, messages: List[Mapping], response_format: pydantic.BaseModel = None) -> str:
    """
    Fetches a single response from the Litellm API for a given model and messages.
    """
    request = ChatCompletionRequest(
        model=model,
        messages=messages,
        response_format=response_format
    )
    
    debug(f"Calling LiteLLM API with model: {request.model}")
    debug(f"Messages: {request.messages}")
    
    if request.response_format is not None:
        debug(f"Using response format: {request.response_format}")
        response = litellm.completion(
            model=request.model,
            messages=request.messages,
            response_format=request.response_format
        )
    else:
        response = litellm.completion(
            model=request.model,
            messages=request.messages,
        )
    return response.choices[0].message.content


def fetch_custom_litellm_api_response(custom_model_parameters: CustomModelParameters, messages: List[Mapping], response_format: pydantic.BaseModel = None) -> str:
    if messages is None or messages == []:
        raise ValueError("Messages cannot be empty")
    
    if custom_model_parameters is None:
        raise ValueError("Custom model parameters cannot be empty")
    
    if not isinstance(custom_model_parameters, CustomModelParameters):
        raise ValueError("Custom model parameters must be a CustomModelParameters object")
        
    if response_format is not None:
        response = litellm.completion(
            model=custom_model_parameters.model_name,
            messages=messages,
            api_key=custom_model_parameters.secret_key,
            base_url=custom_model_parameters.litellm_base_url,
            response_format=response_format
        )
    else:
        response = litellm.completion(
            model=custom_model_parameters.model_name,
            messages=messages,
            api_key=custom_model_parameters.secret_key,
            base_url=custom_model_parameters.litellm_base_url,
        )
    return response.choices[0].message.content


async def afetch_litellm_api_response(model: str, messages: List[Mapping], response_format: pydantic.BaseModel = None) -> str:
    """
    ASYNCHRONOUSLY Fetches a single response from the Litellm API for a given model and messages.
    """
    if messages is None or messages == []:
        raise ValueError("Messages cannot be empty")
    
    # Add validation
    validate_chat_messages(messages)
    
    if model not in LITELLM_SUPPORTED_MODELS:
        raise ValueError(f"Model {model} is not in the list of supported Litellm models: {LITELLM_SUPPORTED_MODELS}.")
    
    if response_format is not None:
        response = await litellm.acompletion(
            model=model,
            messages=messages,
            response_format=response_format
        )
    else:
        response = await litellm.acompletion(
            model=model,
            messages=messages,  
        )
    return response.choices[0].message.content


async def afetch_custom_litellm_api_response(custom_model_parameters: CustomModelParameters, messages: List[Mapping], response_format: pydantic.BaseModel = None) -> str:
    """
    ASYNCHRONOUSLY Fetches a single response from the Litellm API for a given model and messages.
    """
    if messages is None or messages == []:
        raise ValueError("Messages cannot be empty")
    
    if custom_model_parameters is None:
        raise ValueError("Custom model parameters cannot be empty")
    
    if not isinstance(custom_model_parameters, CustomModelParameters):
        raise ValueError("Custom model parameters must be a CustomModelParameters object")
        
    if response_format is not None:
        response = await litellm.acompletion(
            model=custom_model_parameters.model_name,
            messages=messages,
            api_key=custom_model_parameters.secret_key,
            base_url=custom_model_parameters.litellm_base_url,
            response_format=response_format
        )
    else:
        response = await litellm.acompletion(
            model=custom_model_parameters.model_name,
            messages=messages,
            api_key=custom_model_parameters.secret_key,
            base_url=custom_model_parameters.litellm_base_url,
        )
    return response.choices[0].message.content


def query_litellm_api_multiple_calls(models: List[str], messages: List[Mapping], response_formats: List[pydantic.BaseModel] = None) -> List[str]:
    """
    Queries the Litellm API for multiple calls in parallel

    Args:
        models (List[str]): List of models to query
        messages (List[Mapping]): List of messages to query
        response_formats (List[pydantic.BaseModel], optional): A list of the format of the response if JSON forcing. Defaults to None.

    Returns:
        List[str]: Litellm responses for each model and message pair in order. Any exceptions in the thread call result in a None.
    """
    num_workers = int(os.getenv('NUM_WORKER_THREADS', MAX_WORKER_THREADS))  
    # Initialize results to maintain ordered outputs
    out = [None] * len(messages)
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
        # Submit all queries to Litellm API with index, gets back the response content
        futures = {executor.submit(fetch_litellm_api_response, model, message, response_format): idx \
                    for idx, (model, message, response_format) in enumerate(zip(models, messages, response_formats))}
        
        # Collect results as they complete -- result is response content
        for future in concurrent.futures.as_completed(futures):
            idx = futures[future]
            try:
                out[idx] = future.result()
            except Exception as e:
                error(f"Error in parallel call {idx}: {str(e)}")
                out[idx] = None 
    return out


async def aquery_litellm_api_multiple_calls(models: List[str], messages: List[Mapping], response_formats: List[pydantic.BaseModel] = None) -> List[str]:
    """
    Queries the Litellm API for multiple calls in parallel

    Args:
        models (List[str]): List of models to query
        messages (List[Mapping]): List of messages to query
        response_formats (List[pydantic.BaseModel], optional): A list of the format of the response if JSON forcing. Defaults to None.

    Returns:
        List[str]: Litellm responses for each model and message pair in order. Any exceptions in the thread call result in a None.
    """
    # Initialize results to maintain ordered outputs
    out = [None] * len(messages)
    
    async def fetch_and_store(idx, model, message, response_format):
        try:
            out[idx] = await afetch_litellm_api_response(model, message, response_format)
        except Exception as e:
            error(f"Error in parallel call {idx}: {str(e)}")
            out[idx] = None

    tasks = [
        fetch_and_store(idx, model, message, response_format)
        for idx, (model, message, response_format) in enumerate(zip(models, messages, response_formats))
    ]
    
    await asyncio.gather(*tasks)
    return out


def validate_chat_messages(messages, batched: bool = False):
    """Validate chat message format before API call"""
    if not isinstance(messages, list):
        raise TypeError("Messages must be a list")
    
    for msg in messages:
        if not isinstance(msg, dict):
            if batched and not isinstance(msg, list):
                raise TypeError("Each message must be a list")
            elif not batched:
                raise TypeError("Message must be a dictionary")
        if 'role' not in msg:
            raise ValueError("Message missing required 'role' field")
        if 'content' not in msg:
            raise ValueError("Message missing required 'content' field")
        if msg['role'] not in ['system', 'user', 'assistant']:
            raise ValueError(f"Invalid role '{msg['role']}'. Must be 'system', 'user', or 'assistant'")

def validate_batched_chat_messages(messages: List[List[Mapping]]):
    """
    Validate format of batched chat messages before API call
    
    Args:
        messages (List[List[Mapping]]): List of message lists, where each inner list contains
            message dictionaries with 'role' and 'content' fields
    
    Raises:
        TypeError: If messages format is invalid
        ValueError: If message content is invalid
    """
    if not isinstance(messages, list):
        raise TypeError("Batched messages must be a list")
    
    if not messages:
        raise ValueError("Batched messages cannot be empty")
        
    for message_list in messages:
        if not isinstance(message_list, list):
            raise TypeError("Each batch item must be a list of messages")
            
        # Validate individual messages using existing function
        validate_chat_messages(message_list)

def get_chat_completion(model_type: str, 
                        messages : Union[List[Mapping], List[List[Mapping]]], 
                        response_format: pydantic.BaseModel = None, 
                        batched: bool = False
                        ) -> Union[str, List[str]]:
    """
    Generates chat completions using a single model and potentially several messages. Supports closed-source and OSS models.

    Parameters:
        - model_type (str): The type of model to use for generating completions.
        - messages (Union[List[Mapping], List[List[Mapping]]]): The messages to be used for generating completions. 
            If batched is True, this should be a list of lists of mappings.
        - response_format (pydantic.BaseModel, optional): The format of the response. Defaults to None.
        - batched (bool, optional): Whether to process messages in batch mode. Defaults to False.
    Returns:
        - str: The generated chat completion(s). If batched is True, returns a list of strings.
    Raises:
        - ValueError: If requested model is not supported by Litellm or TogetherAI.
    """
    
    # Check for empty messages list
    if not messages or messages == []:
        raise ValueError("Messages cannot be empty")
    
    # Add validation
    if batched:
        validate_batched_chat_messages(messages)
    else:
        validate_chat_messages(messages)
    
    if batched and model_type in TOGETHER_SUPPORTED_MODELS:
        return query_together_api_multiple_calls(models=[model_type] * len(messages), 
                                                 messages=messages, 
                                                 response_formats=[response_format] * len(messages))
    elif batched and model_type in LITELLM_SUPPORTED_MODELS:
        return query_litellm_api_multiple_calls(models=[model_type] * len(messages), 
                                                messages=messages, 
                                                response_format=response_format)
    elif not batched and model_type in TOGETHER_SUPPORTED_MODELS:
        return fetch_together_api_response(model=model_type, 
                                           messages=messages, 
                                           response_format=response_format)
    elif not batched and model_type in LITELLM_SUPPORTED_MODELS:
        return fetch_litellm_api_response(model=model_type, 
                                          messages=messages, 
                                          response_format=response_format)
        
    
    
    raise ValueError(f"Model {model_type} is not supported by Litellm or TogetherAI for chat completions. Please check the model name and try again.")


async def aget_chat_completion(model_type: str, 
                               messages : Union[List[Mapping], List[List[Mapping]]], 
                               response_format: pydantic.BaseModel = None, 
                               batched: bool = False
                               ) -> Union[str, List[str]]:
    """
    ASYNCHRONOUSLY generates chat completions using a single model and potentially several messages. Supports closed-source and OSS models.

    Parameters:
        - model_type (str): The type of model to use for generating completions.
        - messages (Union[List[Mapping], List[List[Mapping]]]): The messages to be used for generating completions. 
            If batched is True, this should be a list of lists of mappings.
        - response_format (pydantic.BaseModel, optional): The format of the response. Defaults to None.
        - batched (bool, optional): Whether to process messages in batch mode. Defaults to False.
    Returns:
        - str: The generated chat completion(s). If batched is True, returns a list of strings.
    Raises:
        - ValueError: If requested model is not supported by Litellm or TogetherAI.
    """
    debug(f"Starting chat completion for model {model_type}, batched={batched}")
    
    if batched:
        validate_batched_chat_messages(messages)
    else:
        validate_chat_messages(messages)
    
    if batched and model_type in TOGETHER_SUPPORTED_MODELS:
        debug("Using batched Together API call")
        return await aquery_together_api_multiple_calls(models=[model_type] * len(messages), 
                                                        messages=messages, 
                                                        response_formats=[response_format] * len(messages))
    elif batched and model_type in LITELLM_SUPPORTED_MODELS:
        debug("Using batched LiteLLM API call")
        return await aquery_litellm_api_multiple_calls(models=[model_type] * len(messages), 
                                                       messages=messages, 
                                                       response_formats=[response_format] * len(messages))
    elif not batched and model_type in TOGETHER_SUPPORTED_MODELS:
        debug("Using single Together API call")
        return await afetch_together_api_response(model=model_type, 
                                                  messages=messages, 
                                                  response_format=response_format)
    elif not batched and model_type in LITELLM_SUPPORTED_MODELS:
        debug("Using single LiteLLM API call")
        return await afetch_litellm_api_response(model=model_type, 
                                                 messages=messages, 
                                                 response_format=response_format)
    
    error(f"Model {model_type} not supported by either API")
    raise ValueError(f"Model {model_type} is not supported by Litellm or TogetherAI for chat completions. Please check the model name and try again.")


def get_completion_multiple_models(models: List[str], messages: List[List[Mapping]], response_formats: List[pydantic.BaseModel] = None) -> List[str]:
    """
    Retrieves completions for a single prompt from multiple models in parallel. Supports closed-source and OSS models.

    Args:
        models (List[str]): List of models to query
        messages (List[List[Mapping]]): List of messages to query. Each inner object corresponds to a single prompt.
        response_formats (List[pydantic.BaseModel], optional): A list of the format of the response if JSON forcing. Defaults to None.
    
    Returns:
        List[str]: List of completions from the models in the order of the input models
    Raises:
        ValueError: If a model is not supported by Litellm or Together
    """
    debug(f"Starting multiple model completion for {len(models)} models")
    
    if models is None or models == []:
        raise ValueError("Models list cannot be empty")
    
    validate_batched_chat_messages(messages)
    
    if len(models) != len(messages):
        error(f"Model/message count mismatch: {len(models)} vs {len(messages)}")
        raise ValueError(f"Number of models and messages must be the same: {len(models)} != {len(messages)}")
    if response_formats is None:
        response_formats = [None] * len(models)
    # Partition the model requests into TogetherAI and Litellm models, but keep the ordering saved
    together_calls, litellm_calls = {}, {}  # index -> model, message, response_format
    together_responses, litellm_responses = [], []
    for idx, (model, message, r_format) in enumerate(zip(models, messages, response_formats)):
        if model in TOGETHER_SUPPORTED_MODELS:
            debug(f"Model {model} routed to Together API")
            together_calls[idx] = (model, message, r_format)
        elif model in LITELLM_SUPPORTED_MODELS:
            debug(f"Model {model} routed to LiteLLM API")
            litellm_calls[idx] = (model, message, r_format)
        else:
            error(f"Model {model} not supported by either API")
            raise ValueError(f"Model {model} is not supported by Litellm or TogetherAI for chat completions. Please check the model name and try again.")
    
    # Add validation before processing
    for msg_list in messages:
        validate_chat_messages(msg_list)
    
    # Get the responses from the TogetherAI models
    # List of responses from the TogetherAI models in order of the together_calls dict
    if together_calls:
        debug(f"Executing {len(together_calls)} Together API calls")
        together_responses = query_together_api_multiple_calls(models=[model for model, _, _ in together_calls.values()], 
                                                           messages=[message for _, message, _ in together_calls.values()], 
                                                           response_formats=[format for _, _, format in together_calls.values()])  
    
    # Get the responses from the Litellm models
    if litellm_calls:
        debug(f"Executing {len(litellm_calls)} LiteLLM API calls")
        litellm_responses = query_litellm_api_multiple_calls(models=[model for model, _, _ in litellm_calls.values()],
                                                        messages=[message for _, message, _ in litellm_calls.values()],
                                                        response_formats=[format for _, _, format in litellm_calls.values()])

    # Merge the responses in the order of the original models
    debug("Merging responses")
    out = [None] * len(models)
    for idx, (model, message, r_format) in together_calls.items():
        out[idx] = together_responses.pop(0)
    for idx, (model, message, r_format) in litellm_calls.items():
        out[idx] = litellm_responses.pop(0)
    debug("Multiple model completion finished")
    return out 


async def aget_completion_multiple_models(models: List[str], messages: List[List[Mapping]], response_formats: List[pydantic.BaseModel] = None) -> List[str]:
    """
    ASYNCHRONOUSLY retrieves completions for a single prompt from multiple models in parallel. Supports closed-source and OSS models.

    Args:
        models (List[str]): List of models to query
        messages (List[List[Mapping]]): List of messages to query. Each inner object corresponds to a single prompt.
        response_formats (List[pydantic.BaseModel], optional): A list of the format of the response if JSON forcing. Defaults to None.
    
    Returns:
        List[str]: List of completions from the models in the order of the input models
    Raises:
        ValueError: If a model is not supported by Litellm or Together
    """
    if models is None or models == []:
        raise ValueError("Models list cannot be empty")
    
    if len(models) != len(messages):
        raise ValueError(f"Number of models and messages must be the same: {len(models)} != {len(messages)}")
    if response_formats is None:
        response_formats = [None] * len(models)

    validate_batched_chat_messages(messages)
    
    # Partition the model requests into TogetherAI and Litellm models, but keep the ordering saved
    together_calls, litellm_calls = {}, {}  # index -> model, message, response_format
    together_responses, litellm_responses = [], []
    for idx, (model, message, r_format) in enumerate(zip(models, messages, response_formats)):
        if model in TOGETHER_SUPPORTED_MODELS:
            together_calls[idx] = (model, message, r_format)
        elif model in LITELLM_SUPPORTED_MODELS:
            litellm_calls[idx] = (model, message, r_format)
        else:
            raise ValueError(f"Model {model} is not supported by Litellm or TogetherAI for chat completions. Please check the model name and try again.")
    
    # Add validation before processing
    for msg_list in messages:
        validate_chat_messages(msg_list)
    
    # Get the responses from the TogetherAI models
    # List of responses from the TogetherAI models in order of the together_calls dict
    if together_calls:
        together_responses = await aquery_together_api_multiple_calls(
            models=[model for model, _, _ in together_calls.values()], 
            messages=[message for _, message, _ in together_calls.values()], 
            response_formats=[format for _, _, format in together_calls.values()]
        )
    
    # Get the responses from the Litellm models
    if litellm_calls:
        litellm_responses = await aquery_litellm_api_multiple_calls(
            models=[model for model, _, _ in litellm_calls.values()],
            messages=[message for _, message, _ in litellm_calls.values()],
            response_formats=[format for _, _, format in litellm_calls.values()]
        )

    # Merge the responses in the order of the original models
    out = [None] * len(models)
    for idx, (model, message, r_format) in together_calls.items():
        out[idx] = together_responses.pop(0)
    for idx, (model, message, r_format) in litellm_calls.items():
        out[idx] = litellm_responses.pop(0)
    return out


if __name__ == "__main__":
    
    # Batched
    pprint.pprint(get_chat_completion(
        model_type="LLAMA3_405B_INSTRUCT_TURBO",
        messages=[
            [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "What is the capital of France?"},
            ],
            [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "What is the capital of Japan?"},
            ]
        ],
        batched=True
    ))

    # Non batched
    pprint.pprint(get_chat_completion(
        model_type="LLAMA3_8B_INSTRUCT_TURBO",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "What is the capital of France?"},
        ],
        batched=False
    ))

    # Batched single completion to multiple models
    pprint.pprint(get_completion_multiple_models(
        models=[
            "LLAMA3_70B_INSTRUCT_TURBO", "LLAMA3_405B_INSTRUCT_TURBO", "gpt-4o-mini"
        ],
        messages=[
            [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "What is the capital of China?"},
            ],
            [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "What is the capital of France?"},
            ],
            [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "What is the capital of Japan?"},
            ]
        ]
    ))
