config.py

import os

openai_flag = True

try:
    import openai
except: 
    openai_flag = False 

import gpt4all
import questionary
import yaml

from talk_codebase.consts import MODEL_TYPES




config_path = os.path.join(os.path.expanduser("~"), ".talk_codebase_config.yaml")


def get_config():
    if os.path.exists(config_path):
        with open(config_path, "r") as f:
            config = yaml.safe_load(f)
    else:
        config = {}
    return config


def save_config(config):
    with open(config_path, "w") as f:
        yaml.dump(config, f)


def api_key_is_invalid(api_key):
    if not api_key:
        return True
    try:
        openai.api_key = api_key
        openai.Engine.list()
    except Exception:
        return True
    return False


def get_gpt_models(openai):
    try:
        model_lst = openai.Model.list()
    except Exception:
        print("✘ Failed to retrieve model list")
        return []

    return [i['id'] for i in model_lst['data'] if 'gpt' in i['id']]


def configure_model_name_openai(config):
    api_key = config.get("api_key")

    if config.get("model_type") != MODEL_TYPES["OPENAI"] or config.get("openai_model_name"):
        return

    openai.api_key = api_key
    gpt_models = get_gpt_models(openai)
    choices = [{"name": model, "value": model} for model in gpt_models]

    if not choices:
        print("ℹ No GPT models available")
        return

    model_name = questionary.select("🤖 Select model name:", choices).ask()

    if not model_name:
        print("✘ No model selected")
        return

    config["openai_model_name"] = model_name
    save_config(config)
    print("🤖 Model name saved!")


def remove_model_name_openai():
    config = get_config()
    config["openai_model_name"] = None
    save_config(config)


def configure_model_name_local(config):
    if config.get("model_type") != MODEL_TYPES["LOCAL"] or config.get("local_model_name"):
        return

    list_models = gpt4all.GPT4All.list_models()

    def get_model_info(model):
        return (
            f"{model['name']} "
            f"| {model['filename']} "
            f"| {model['filesize']} "
            f"| {model['parameters']} "
            f"| {model['quant']} "
            f"| {model['type']}"
        )

    choices = [
        {"name": get_model_info(model), "value": model['filename']} for model in list_models
    ]

    model_name = questionary.select("🤖 Select model name:", choices).ask()
    config["local_model_name"] = model_name
    save_config(config)
    print("🤖 Model name saved!")


def remove_model_name_local():
    config = get_config()
    config["local_model_name"] = None
    save_config(config)


def get_and_validate_api_key():
    prompt = "🤖 Enter your OpenAI API key: "
    api_key = input(prompt)
    while api_key_is_invalid(api_key):
        print("✘ Invalid API key")
        api_key = input(prompt)
    return api_key


def configure_api_key(config):
    if config.get("model_type") != MODEL_TYPES["OPENAI"]:
        return

    if api_key_is_invalid(config.get("api_key")):
        api_key = get_and_validate_api_key()
        config["api_key"] = api_key
        save_config(config)


def remove_api_key():
    config = get_config()
    config["api_key"] = None
    save_config(config)


def remove_model_type():
    config = get_config()
    config["model_type"] = None
    save_config(config)


def configure_model_type(config):
    if config.get("model_type"):
        return
    
    choices = [{"name": "Local", "value": MODEL_TYPES["LOCAL"]}]

    if openai_flag: choices.append(
            {"name": "OpenAI", "value": MODEL_TYPES["OPENAI"]})


    model_type = questionary.select(
        "🤖 Select model type:",
        choices=choices 
    ).ask()
    config["model_type"] = model_type
    save_config(config)


CONFIGURE_STEPS = [
    configure_model_type,
    configure_api_key,
    configure_model_name_openai,
    configure_model_name_local,
]