from mlx_lm import load, generate

def main():
    model, tokenizer = load("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit")
    system_message = "You are a helpful, concise Assistant"
    messages = [{"role": "system", "content": system_message}]
    run(model, tokenizer, messages, get_user_message)


 # run loops getting the user message and displaying the LLM response
def run(model, tokenizer, messages, get_msg_fn):
    while True:
        user_input, ok = get_msg_fn()
        if not ok:
            break

        messages.append({"role": "user", "content": user_input})
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        response = generate(model, tokenizer, prompt=prompt, max_tokens=10000)
        print(f"\nAssistant: {response}\n")
        messages.append({"role": "assistant", "content": response})


def get_user_message():
    try:
        user_input = input("You (provide input or type quit): ").strip()
        if user_input.lower() in {"exit", "quit"}:
            return "", False
        return user_input, True
    except (EOFError, KeyboardInterrupt):
        print()
        return "", False


if __name__ == "__main__":
    main()

Name:        "read_file"

Description: "Read the contents of a given file path. Use this to see everything in a file. Do not use this with directory names.",

InputSchema:
    "type": "object",
    "properties": {
        "file_path": {
            "type": "string",
            "description": "The path of the file to read",
        },
    },
    "required": ["file_path"],

import json
from mlx_lm import load, generate

def main():
    model, tokenizer = load("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit")
    system_message = """You are a helpful, concise Assistant.
    You have access to a read_file tool, but ONLY use them when the user's request REQUIRES one.
    MOST messages need NO tool call. Just reply in plain text.
    DO NOT call a tool for:
    - Greetings, chitchat, or general questions
    - Questions you already know the answer to
    ONLY call a tool when you literally cannot answer without it.
    When you do NOT need a tool, respond in plain text. Never output JSON unless you are calling a tool."""

    messages = [{"role": "system", "content": system_message}]
    run(model, tokenizer, messages, get_user_message)


 # run loops getting the user message and displaying the LLM response
def run(model, tokenizer, messages, get_msg_fn):
    while True:
        user_input, ok = get_msg_fn()
        if not ok:
            break

        messages.append({"role": "user", "content": user_input})
        user_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=TOOLS)

        # dedicated helper function to handle tool calling complexity
        response = generate_response(model, tokenizer, messages, user_prompt)

        print(f"\nAssistant: {response}\n")
        messages.append({"role": "assistant", "content": response})


def get_user_message():
    try:
        user_input = input("You (provide input or type quit): ").strip()
        if user_input.lower() in {"exit", "quit"}:
            return "", False
        return user_input, True
    except (EOFError, KeyboardInterrupt):
        print()
        return "", False


def generate_response(model, tokenizer, messages, user_prompt):
    response = generate(model, tokenizer, prompt=user_prompt, max_tokens=10000)

    # Check if the LLM believes the User query may need a tool call, if the first LLM response is not JSON just ignore
    try:
        response_as_json = json.loads(response)
        if response_as_json.get("name") == "read_file":
            # DEBUG print(f"\nAssistant believes it needs a tool call: {response}\n")
            tool_input = response_as_json["parameters"]
            tool_result = read_file(tool_input["file_path"])
            messages.append({
                "role": "tool",
                "name": "read_file",
                "content": tool_result,
            })

            # create a user friendly response based on the tool call result
            prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            response = generate(model, tokenizer, prompt=prompt, max_tokens=10000)

    except (json.JSONDecodeError, KeyError, TypeError):
        pass
    return response


def read_file(file_path):
    try:
        with open(file_path, "r") as f:
            return f.read()
    except Exception as e:
        return f"Error reading file: {e}"


TOOLS = [
    {
        "type": "function",
        "name": "read_file",
        "function": "read_file",
        "description": "Read the contents of a given file path. Use this to see everything in a file. Do not use this with directory names.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "file_path": {
                    "type": "string",
                    "description": "The path of the file to read",
                },
            },
            "required": ["file_path"],
        },
    },
]


if __name__ == "__main__":
    main()

You (provide input or type quit): tell me a joke                                                                         

Assistant: I'm a large language model, I don't have a collection of jokes stored in a file. I can try to come up with a joke on the spot, though!
Why couldn't the bicycle stand up by itself? Because it was two-tired!

You (provide input or type quit): summarize the file 02-chat-loop-with-read-file-tool.py

Assistant: This script is a chat loop that uses a large language model (LLM) to respond to user input. The LLM is loaded from a model and tokenizer, and the chat loop continues until the user types "quit" or "exit".

The script also includes a "read_file" tool that can be used to read the contents of a file. If the LLM believes that a user query may require a tool call, it will print a message indicating that it needs to call a tool, and then call the "read_file" tool with the specified file path. The result of the tool call is then used to create a user-friendly response.

<-\[o_o]/--c
x--|---|--C
 _/|___|\_