quickprompt.py -

quickprompt.py

import os
import json
import requests
import questionary
import inquirer

def ensure_api_key():
    """
    Checks if the environment variable 'GROQ_API_KEY' is set.

    If the 'GROQ_API_KEY' environment variable is not set, it prints an error message
    and exits the program with a status code of 1.

    Raises:
        SystemExit: If 'GROQ_API_KEY' is not set in the environment variables.
    """
    if not os.getenv("GROQ_API_KEY"):
        print("Error: GROQ_API_KEY is not set. Please export your API key.")
        exit(1)

def select_option(prompt_text, options):
    """
    Display a selection prompt to the user and return the selected option.

    Args:
        prompt_text (str): The text to display as the prompt.
        options (list): A list of options for the user to choose from.

    Returns:
        str: The option selected by the user.
    """
    answer = questionary.select(
        prompt_text,
        choices=options
    ).ask()
    return answer

def prompt_for_files():
    """
    Prompts the user to select files from the current directory and its subdirectories,
    and saves the selected files to 'files.txt'. If 'files.txt' already exists, the user
    is asked whether to add more files to the existing list.
    Returns:
        list: A list of file paths selected by the user.
    """
    files = []
    for root, _, filenames in os.walk('.'):
        for filename in filenames:
            files.append(os.path.join(root, filename))
    
    existing_files = []
    if os.path.isfile("files.txt"):
        with open("files.txt", "r") as f:
            existing_files = f.read().splitlines()
        if existing_files:
            add_more_files = questionary.confirm(
                message="files.txt already contains files. Would you like to add more files?",
                default=True
            ).ask()
            if not add_more_files:
                return existing_files

    questions = [
        inquirer.Checkbox(
            'selected_files',
            message='Select files to pass to the model',
            choices=files
        )
    ]
    answers = inquirer.prompt(questions)
    selected_files = answers['selected_files']
    
    all_files = existing_files + selected_files
    print("Files to be processed:", ', '.join(all_files))
    with open("files.txt", "w") as f:
        f.write('\n'.join(all_files))
    return all_files

def send_message(user_message, api_key):
    """
    Sends a message to the OpenAI API and retrieves the assistant's reply.

    Args:
        user_message (str): The message to send to the assistant.
        api_key (str): The API key for authenticating the request.

    Returns:
        str: The assistant's reply.

    Raises:
        requests.exceptions.HTTPError: If the HTTP request returned an unsuccessful status code.

    The function sends a POST request to the OpenAI API with the user's message and API key.
    It then processes the response to extract the assistant's reply, prints it, and writes it to a file named 'prompt.txt'.
    """
    payload = {
        "model": "llama3-8b-8192",
        "messages": [{"role": "user", "content": user_message}]
    }

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    response = requests.post(
        "https://api.groq.com/openai/v1/chat/completions",
        headers=headers,
        data=json.dumps(payload)
    )

    response.raise_for_status()

    response_json = response.json()
    assistant_reply = response_json['choices'][0]['message']['content']
    print("Assistant:", assistant_reply)
    with open("prompt.txt", "w") as f:
        f.write(assistant_reply)
    return assistant_reply

def main():
    """
    Main function to execute the script.

    This function performs the following steps:
    1. Ensures the API key is available.
    2. Prompts the user to select an action, focus, and subject from predefined lists.
    3. Reads context information from a file named 'context.txt'.
    4. Prompts the user to specify file paths.
    5. Constructs an initial prompt using the selected options and context information.
    6. Writes the initial prompt to a file named 'initial_prompt.md'.
    7. Sends the initial prompt to an assistant and handles the assistant's replies in a loop.
        - If the assistant's reply ends with a question mark, prompts the user for a response.
        - Appends the user's response to the initial prompt and sends it to the assistant again.
        - Exits the loop if the user types 'exit' or if the assistant has no further questions.

    Raises:
         SystemExit: If the context file 'context.txt' does not exist.
    """
    ensure_api_key()

    actions = [
        "Implement", "Debug", "Optimize", "Refactor", "Review",
        "Integrate", "Document", "Test", "Deploy"
    ]
    focuses = [
        "Login System", "Sorting Algorithm", "Database Connection",
        "User Interface", "API Endpoint", "Vue.js Component",
        "Vue.js State Management", "Ansible Playbook", "Ansible Role", "Vite"
    ]
    subjects = [
        "JavaScript", "Python", "Java", "C++", "React", "Angular",
        "Django", "Flask", "TensorFlow", "Vue.js", "Ansible", "Vite"
    ]

    action = select_option("Select an action", actions)
    focus = select_option("Select a focus", focuses)
    subject = select_option("Select a subject", subjects)

    context_file = "context.txt"

    if os.path.isfile(context_file):
        with open(context_file, "r") as f:
            context = f.read()
    else:
        print(f"The context file '{context_file}' does not exist.")
        print("Please create the file with the necessary context information.")
        print("You can use the following command to create the file:")
        print(f"  nano {context_file}")
        print("After creating the file, run this script again.")
        exit(1)

    file_paths = prompt_for_files()

    initial_prompt = (
        f"Action: {action}\n"
        f"Focus: {focus}\n"
        f"Subject: {subject}\n"
        f"Context:\n{context}"
    )

    with open("initial_prompt.md", "w") as f:
        f.write(initial_prompt)

    assistant_reply = send_message(initial_prompt, os.getenv("GROQ_API_KEY"))

    while True:
        if assistant_reply.strip().endswith("?"):
            user_input = input("Your response (or type 'exit' to end): ")
            if user_input.lower() == "exit":
                print("Exiting chat.")
                break

            initial_prompt += f"\nUser: {user_input}"

            with open("initial_prompt.md", "w") as f:
                f.write(initial_prompt)

            assistant_reply = send_message(initial_prompt, os.getenv("GROQ_API_KEY"))
        else:
            print("No further questions from the assistant. Exiting chat.")
            break

if __name__ == "__main__":
    main()