import dataiku

LLM_SUPPORTED_CONNECTIONS = ['OpenAI', 'AzureOpenAI']

def get_llms_from_OpenAI(llm_connection):
    llms = []
    llm_connection_name = llm_connection['name']
    if llm_connection['params']['allowGPT4']:
        llms.append("openai:{}:gpt-4".format(llm_connection_name))
    if llm_connection['params']['allowGPT432K']:
        llms.append("openai:{}:gpt-4-32k".format(llm_connection_name))
    if llm_connection['params']['allowGPT4Turbo']:
        llms.append("openai:{}:gpt-4-1106-preview".format(llm_connection_name))
    if llm_connection['params']['allowGPT35Turbo']:
        llms.append("openai:{}:gpt-3.5-turbo-16k".format(llm_connection_name))
    if llm_connection['params']['allowGPT3Davinci']:
        llms.append("openai:{}:davinci-002".format(llm_connection_name))
    return llms

def get_llms_from_AzureOpenAI(llm_connection):
    llm_connection_name = llm_connection['name']
    llms = ["azureopenai:{}:{}".format(llm_connection_name, llm['name']) for llm in llm_connection['params']['availableDeployments']]
    return llms

def get_llms_from_connection(llm_connection):
    if llm_connection['type']=='OpenAI':
        return get_llms_from_OpenAI(llm_connection)
    elif llm_connection['type']=='AzureOpenAI':
        return get_llms_from_AzureOpenAI(llm_connection)
    else:
        raise ValueError("Unsupported connection type: ", llm_connection['type'])

def list_available_llms():
    client = dataiku.api_client()
    connections = client.list_connections()
    llm_connections = [connections[connection] for connection in connections.keys() if connections[connection]['type'] in LLM_SUPPORTED_CONNECTIONS]
    llms = [llm for llm_connection in llm_connections for llm in get_llms_from_connection(llm_connection)]
    return llms