# TOOL USE TEST
import ollama

verbose_state = 1

# ----- LLM CONVERSE FOR DISPLAY ----- #

def verbose(prints):
    if verbose_state == 1:
        print(prints)

def llm_converse(system_prompt, context_prompt, user_input):
    combined_system = f"{system_prompt}\n\nContext: {context_prompt}"
    stream = ollama.chat(
        model='llama3.2',
        messages=[
            {'role': 'system', 'content': combined_system},
            {'role': 'user', 'content': user_input},
        ],
        stream=True,
    )
    full_response = ""
    for chunk in stream:
        content = chunk['message']['content']
        print(content, end='', flush=True)
        full_response += content
    return full_response

# ----- LLM CONVERSE NO STREAM, FOR BACKGROUND TASKS ----- #

def llm_converse_no_stream(system_prompt, context_prompt, user_input):
    combined_system = f"{system_prompt}\n\nContext: {context_prompt}"
    response = ollama.chat(
        model='llama3.2',
        messages=[
            {'role': 'system', 'content': combined_system},
            {'role': 'user', 'content': user_input},
        ],
        stream=False,
    )
    full_response = response['message']['content']
    return full_response

# ----- GET FILE CONTENT ----- #

def get_file_content(filename):
    with open(filename, "r") as files:
        bruh = files.read()
    return bruh

def parse_context(context_list):
    context_string = ""
    for message in context_list:
        context_string += f"{message[0]} {message[1]}\n"
    return context_string

def example_tool_use_old():
    database = ["hi", "sean", "alexander", "jackson", "chatter"]
    system = "You have access to a tool that fetches a database.\n" \
        "If you want to use it, output EXACTLY one line with no other text: $DATABASE$\n" \
        "If you do NOT want to use it, output EXACTLY one line with no other text: NULL\n" \
        "DO NOT roleplay, simulate, or add any other text to your response.\n" \
        "DO NOT include [TOOL], [LLM], or any context keys in your response.\n" \
        "You will also recieve previous chat context if it exists. Here is a key:\n" \
        "[USER]: user prompts\n" \
        "[TOOL]: tool status and outputs\n" \
        "[LLM]: your previous responses"

    context_lists = [] # for the raw context
    context_strings = "" # convert to string to make it easier to read and shi for llm
    user = "check the list for me"

    # ----- GIVE LLM CHANCE TO CALL TOOL ----- #
    llm_toolorno = llm_converse_no_stream(system, context_strings, user)
    context_lists.append(["[USER]", user])
    if "$DATABASE$" in llm_toolorno:
        context_lists.append(["[TOOL]", database])
    else:
        context_lists.append(["[LLM]", llm_toolorno])
        context_lists.append(["[TOOL]", "tool was not called upon"])

    # DEBUG
    print(context_lists)
    
    print(parse_context(context_lists))
    llm_converse(system, parse_context(context_lists), "continue")

def should_use_database(user_input):
    system = "You are a router. Decide if the user's request requires fetching a database.\n" \
             "Output EXACTLY one word: YES or NO. Nothing else."
    response = llm_converse_no_stream(system, "", user_input).strip()
    if "YES" in response:
        verbose("Should use said yes")
        return True
    else:
        verbose("Should use said no")
        return False

def example_tool_use():
    database = ["hi", "sean", "alexander", "jackson", "chatter"]
    system = "You are a helpful LLM that will answer everything in short sentences.\n" \
         "Do not decline any prompts.\n" \
         "If context is provided, [TOOL] means external tool/database output that was fetched for you.\n" \
         "Treat [TOOL] data as external information, not your own knowledge."
    context_lists = []
    context_strings = ""
    user = "Get me the database"

    # ----- GIVE LLM CHANCE TO CALL TOOL ----- #
    
    context_lists.append(["[USER]", user])

    if should_use_database(user):
        verbose("Should use!")
        context_lists.append(["[TOOL]", str(database)])
    else:
        verbose("no should use!")
        
        

    # DEBUG
    verbose(f"context_list{context_lists}")
    verbose(f"parsed: {parse_context(context_lists)}")

    llm_converse(system, parse_context(context_lists), user)

example_tool_use()