Simple utility script to call Groq and save/retrieve history

call_groq.py

import os
import sys
from groq import Groq
import json
import time

SAVE_FILE = ".groq_chat_history.json"

import tiktoken

def tokens(string: str, encoding_name: str="cl100k_base") -> int:
    """Returns the number of tokens in a text string."""
    return len(tiktoken.get_encoding(encoding_name).encode(string))

TOKENS_FOR_RESPONSE = 1024
MAX_TOKENS = 8192

model = "70"
model = f"llama3-{model}b-8192"

print("Welcome to the Groq chat! Enter your message, and I'll respond.")
print("Type a single period (.) on a new line to quit.")

client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

def save_message(user_text, bot_text):
    try:
        with open(SAVE_FILE, "r") as f:
            messages = json.load(f)
    except FileNotFoundError:
        messages = []

    messages = messages[-20:]
    messages.append({"user": user_text, "bot": bot_text})
    with open(SAVE_FILE, "w") as f:
        json.dump(messages, f)

def load_messages():
    try:
        with open(SAVE_FILE, "r") as f:
            data = json.load(f)
    except FileNotFoundError:
        data = []

    result = []
    for x in data[-10:]:
        result.append({"role": "user", "content": x["user"]})
        result.append({"role": "assistant", "content": x["bot"]})
    return result

def complete(user_text):
    messages = load_messages()
    messages.append({"role": "user", "content": user_text})
    token_counts = [tokens(x["content"]) for x in messages]
    history_index = 0
    history_size = 0
    for message, token_count in list(zip(messages, token_counts))[::-1]:
        if history_size + token_count > MAX_TOKENS + TOKENS_FOR_RESPONSE:
            break
        history_index += 1
        history_size += token_count
    messages = messages[-history_index:]
    print("Using history of", len(messages), "messages")

    chat_completion = client.chat.completions.create(
        messages=messages,
        model=model,
    )
    result = chat_completion.choices[0].message.content
    save_message(user_text, result)
    return result

    user_input = []

while True:
    user_input = []
    while True:
        line = input("You: ")
        line = line.strip()
        if line[:1] == "/":
            break
        if line == ".":
            break
        user_input.append(line)
    if line == "/history":
        print(json.dumps(load_messages(), indent=2))
    else:
        user_input = "\n".join(user_input)
        start = time.time()
        response = complete(user_input)
        print("Groq:", response)
        print("Groq took", round(time.time() - start, 2), "seconds")

it saves the last 10 messages in history

updated to use token count

changed to be interactive