ComposableConsult/Tokenusagecalculator

1import gradio as gr2import tiktoken3import sqlite34import matplotlib.pyplot as plt5import io6from datetime import datetime7 8# Load GPT-4 tokenizer9enc = tiktoken.encoding_for_model("gpt-4")10 11# Connect to SQLite DB12conn = sqlite3.connect("usage.db", check_same_thread=False)13cursor = conn.cursor()14 15# Create table if it doesn't exist16cursor.execute('''17    CREATE TABLE IF NOT EXISTS token_logs (18        id INTEGER PRIMARY KEY AUTOINCREMENT,19        user_id TEXT,20        text TEXT,21        token_count INTEGER,22        timestamp TEXT23    )24''')25conn.commit()26 27# Function to count and log tokens28def count_tokens(user_id, text):29    token_count = len(enc.encode(text))30    timestamp = datetime.utcnow().isoformat()31 32    cursor.execute(33        "INSERT INTO token_logs (user_id, text, token_count, timestamp) VALUES (?, ?, ?, ?)",34        (user_id, text, token_count, timestamp)35    )36    conn.commit()37 38    return f"Token Count: {token_count}"39 40# Function to generate usage summary41def generate_chart():42    cursor.execute("SELECT user_id, SUM(token_count) FROM token_logs GROUP BY user_id")43    data = cursor.fetchall()44 45    if not data:46        return "No data yet", None47 48    user_ids, counts = zip(*data)49 50    fig, ax = plt.subplots()51    ax.bar(user_ids, counts)52    ax.set_title("Total Token Usage per User")53    ax.set_ylabel("Tokens")54 55    buf = io.BytesIO()56    plt.savefig(buf, format="png")57    buf.seek(0)58    return "Usage chart:", buf59 60# UI Layout61with gr.Blocks() as demo:62    gr.Markdown("## 📊 Token Usage Dashboard with SQLite + Matplotlib")63 64    with gr.Row():65        user_id = gr.Textbox(label="User ID")66        text = gr.Textbox(label="Enter Prompt", lines=5)67    count_btn = gr.Button("Submit Prompt")68    output = gr.Textbox(label="Token Count")69 70    chart_btn = gr.Button("Show Usage Chart")71    chart_text = gr.Textbox(label="Status")72    chart_img = gr.Image(type="pil")73 74    count_btn.click(fn=count_tokens, inputs=[user_id, text], outputs=output)75    chart_btn.click(fn=generate_chart, outputs=[chart_text, chart_img])76 77demo.launch()78