ChatGPT Tools

Count GPT tokens and manage context with export functions

// ==UserScript==
// @name			ChatGPT Tools
// @namespace		https://muyi.tw/labs/tokens-tools
// @version			0.5
// @description		Count GPT tokens and manage context with export functions
// @author			MuYi + ChatGPT + Claude
// @match			*://chatgpt.com/*
// @grant			none
// @license			AGPL-3.0-or-later
// ==/UserScript==

(function () {
	'use strict';

	// ====== CONFIG ======
	const CONFIG = {
		selectors: {
			inputText: 'div#prompt-textarea',
			userText: 'article div[data-message-author-role="user"]',
			botText: 'article div[data-message-author-role="assistant"], article div.cm-content'
		},
		validPathPatterns: [
			/\/c\/[0-9a-fA-F\-]{36}/,
		],
		updateInterval: 1000,
		chunkSize: 6000,
		tokenWarningThreshold: 100000,
		summaryText: '請總結上方對話為技術說明。',
		uiStyle: {
			position: 'fixed',
			bottom: '10%',
			right: '2.5em',
			zIndex: '9999',
			padding: '.5em',
			backgroundColor: 'rgba(0,0,0,0.5)',
			color: 'white',
			fontSize: '100%',
			borderRadius: '.5em',
			fontFamily: 'monospace',
			display: 'none',
		}
	};

	// ====== LOCALE ======
	const localeMap = {
		'zh-hant': {
			calculating: 'Token Counter 計算中⋯⋯',
			total: '本窗內容預估:{0} tokens',
			breakdown: '(輸入輸出:{0}/{1})',
			inputBox: '輸入欄內容預估:{0} tokens',
			init: 'Token Counter 初始化中⋯⋯',
			navigation: '偵測到頁面導航',
			errorRegex: '正則取代時發生錯誤:{0}',
			errorText: '讀取文字失敗:{0}',
			regexInfo: '[{0}] 匹配 {1} 次,權重 {2}',
			prefixUser: '使用者說:',
			prefixBot: 'GPT說:',
			exportText: '輸出文字',
			exportJSON: '輸出 JSON'
		},
		'en': {
			calculating: 'Token Counter Calculating...',
			total: 'Total tokens in view: {0}',
			breakdown: '(Input / Output: {0}/{1})',
			inputBox: 'Input box tokens: {0}',
			init: 'Token Counter initializing...',
			navigation: 'Navigation detected',
			errorRegex: 'Token counting regex replacement error: {0}',
			errorText: 'Error getting text: {0}',
			regexInfo: '[{0}] matched {1} times, weight {2}',
			prefixUser: 'You said:',
			prefixBot: 'GPT said:',
			exportText: 'Export Text',
			exportJSON: 'Export JSON'
		}
	};

	function resolveLocale() {
		const lang = navigator.language.toLowerCase();
		if (localeMap[lang]) return lang;
		if (lang.startsWith('zh-')) {
			const fallback = Object.keys(localeMap).find(k => k.startsWith('zh-'));
			if (fallback) return fallback;
		}
		return 'en';
	}
	const locale = localeMap[resolveLocale()];

	// ====== UTILS ======
	const DEBUG = true;
	const format = (s, ...a) => s.replace(/\{(\d+)\}/g, (_, i) => a[i] ?? '');
	const safeIdle = cb => window.requestIdleCallback?.(cb) || setTimeout(() => cb({ didTimeout: false, timeRemaining: () => 0 }), 1);
	const cancelIdle = h => window.cancelIdleCallback?.(h) || clearTimeout(h);
	const debugLog = (...args) => DEBUG && console.log('[TokenCounter]', ...args);

	// ====== ESTIMATE RULES ======
	const gptWeightMap = [
		{ regex: /[\p{Script=Han}]/gu, weight: 0.99 },
		{ regex: /[\p{Script=Hangul}]/gu, weight: 0.79 },
		{ regex: /[\p{Script=Hiragana}\p{Script=Katakana}]/gu, weight: 0.73 },
		{ regex: /[\p{Script=Latin}]+/gu, weight: 1.36 },
		{ regex: /[\p{Script=Greek}]+/gu, weight: 3.14 },
		{ regex: /[\p{Script=Cyrillic}]+/gu, weight: 2.58 },
		{ regex: /[\p{Script=Arabic}]+/gu, weight: 1.78 },
		{ regex: /[\p{Script=Hebrew}]+/gu, weight: 1.9 },
		{ regex: /[\p{Script=Devanagari}]+/gu, weight: 1.28 },
		{ regex: /[\p{Script=Bengali}]+/gu, weight: 1.77 },
		{ regex: /[\p{Script=Thai}]/gu, weight: 0.45 },
		{ regex: /[\p{Script=Myanmar}]/gu, weight: 0.56 },
		{ regex: /[\p{Script=Tibetan}]/gu, weight: 1.58 },
		{ regex: /\p{Number}{1,3}/gu, weight: 1.0 },
		{ regex: /[\u2190-\u2BFF\u1F000-\u1FAFF]/gu, weight: 1.0 },
		{ regex: /[\p{P}]/gu, weight: 0.95 },
		{ regex: /[\S]+/gu, weight: 3.0 }
	];

	// ====== STATE ======
	const state = {
		idleHandle: null,
		intervalId: null,
		uiBox: null,
		operationId: 0,  // 用於標記當前操作的ID
	};

	// ====== CORE ======

	let updateDirty = false;

	function createUI() {
		if (state.uiBox) return;

		const box = document.createElement('div');
		const content = document.createElement('div');
		const actions = document.createElement('div');

		Object.assign(box.style, CONFIG.uiStyle);
		Object.assign(actions.style, { marginTop: '0.5em' });

		box.appendChild(content);
		box.appendChild(actions);
		document.body.appendChild(box);

		state.uiBox = box;
		state.uiContent = content;
		state.uiActions = actions;
		addUIButton('📄', exportAsText, state.uiActions, locale.exportText);
		addUIButton('🧾', exportAsJSON, state.uiActions, locale.exportJSON);
	}

	function extractDialogTurns() {
		return Array.from(document.querySelectorAll('article[data-testid^="conversation-turn-"]'))
			.map(el => {
				const turn = parseInt(el.dataset.testid.replace('conversation-turn-', ''), 10);
				const user = el.querySelector('[data-message-author-role="user"]');
				const bot = el.querySelector('[data-message-author-role="assistant"]');
				return {
					id: turn,
					role: user ? 'user' : bot ? 'bot' : 'unknown',
					text: (user || bot)?.innerText.trim() || ''
				};
			})
			.filter(item => item.role !== 'unknown')
			.sort((a, b) => a.id - b.id);
	}

	function exportAsText() {
		const data = extractDialogTurns();
		const output = data.map(({ role, text }) => {
			const prefix = role === 'user' ? locale.prefixUser : locale.prefixBot;
			const unescape = text
				.replace(/\\r/g, '\r')
				.replace(/\\n/g, '\n')
				.replace(/\\t/g, '\t');
			return `${prefix}\n${unescape}`;
		}).join('\n\n');
		navigator.clipboard.writeText(output).then(() => debugLog('Text copied to clipboard.'));
	}


	function exportAsJSON() {
		const data = extractDialogTurns();
		navigator.clipboard.writeText(JSON.stringify(data, null, 2)).then(() => debugLog('JSON copied to clipboard.'));
	}

	function addUIButton(label, onclick, container = state.uiBox, title = '') {
		const btn = document.createElement('button');
		btn.textContent = label;
		btn.title = title;
		btn.style.marginLeft = '0.5em';
		btn.onclick = onclick;
		container.appendChild(btn);
	}

	let lastPathname = location.pathname;

	function isValidWindow() {
		const now = location.pathname;
		const changed = now !== lastPathname;
		if (changed) lastPathname = now;

		const matched = CONFIG.validPathPatterns.some(re => re.test(now));
		const status = `${matched ? 'valid' : 'invalid'}-${changed ? 'changed' : 'unchanged'}`;

		debugLog('isValidWindow check:', {
			pathname: now,
			status
		});

		return status;
	}

	function getCombinedText(selector) {
		try {
			return Array.from(document.querySelectorAll(selector))
				.map(el => el?.innerText || '')
				.filter(Boolean)
				.join('\n');
		} catch (e) {
			console.error(format(locale.errorText, e));
			return '';
		}
	}

	function estimateTokensAsync(text, callback) {
		if (!text) return callback(0);
		let total = 0, i = 0, remaining = text;
		function process() {
			if (i >= gptWeightMap.length) return callback(Math.round(total));
			const { regex, weight } = gptWeightMap[i++];
			safeIdle(() => {
				const matches = remaining.match(regex) || [];
				total += matches.length * weight;
				try {
					if (matches.length) remaining = remaining.replace(regex, ' ');
				} catch (e) {
					console.error(format(locale.errorRegex, e));
				}
				process();
			});
		}
		process();
	}

	function estimateTokensChunked(text, callback) {
		if (!text) return callback(0);
		const chunks = text.match(new RegExp(`.{1,${CONFIG.chunkSize}}`, 'gs')) || [];
		let total = 0, i = 0;
		function next() {
			if (i >= chunks.length) return callback(total);
			estimateTokensAsync(chunks[i++], count => {
				total += count;
				next();
			});
		}
		next();
	}

	function updateDisplay(user, bot, input) {
		const both = user + bot
		const total = both + input;
		if (total > CONFIG.tokenWarningThreshold) {
			state.uiBox.style.backgroundColor = 'rgba(255,50,50,0.7)';
		} else {
			state.uiBox.style.backgroundColor = CONFIG.uiStyle.backgroundColor;
		}
		state.uiContent.innerHTML = (both === 0)
			? locale.calculating
			: format(locale.total, both) + '<br>' +
			format(locale.breakdown, user, bot) + '<br>' +
			format(locale.inputBox, input);
	}

	function updateCounter() {
		const currentOperation = ++state.operationId;  // 遞增操作ID
		const userText = getCombinedText(CONFIG.selectors.userText);
		const botText = getCombinedText(CONFIG.selectors.botText);
		const inputEl = document.querySelector(CONFIG.selectors.inputText);
		const inputText = inputEl ? inputEl.innerText : '';

		let pending = 3;
		let user = 0, bot = 0, input = 0;

		function tryDisplay() {
			if (currentOperation !== state.operationId) return;  // 檢查是否為最新操作
			if (--pending === 0) updateDisplay(user, bot, input);
		}

		estimateTokensChunked(userText, count => {
			if (currentOperation !== state.operationId) return;  // 檢查是否為最新操作
			user = count;
			tryDisplay();
		});
		estimateTokensChunked(botText, count => {
			if (currentOperation !== state.operationId) return;  // 檢查是否為最新操作
			bot = count;
			tryDisplay();
		});
		estimateTokensChunked(inputText, count => {
			if (currentOperation !== state.operationId) return;  // 檢查是否為最新操作
			input = count;
			tryDisplay();
		});
	}

	function resetAll() {
		if (state.idleHandle) {
			cancelIdle(state.idleHandle);
			state.idleHandle = null;
		}
		if (state.intervalId) {
			clearInterval(state.intervalId);
			state.intervalId = null;
		}
		state.operationId++;  // 遞增操作ID使舊的操作失效
		updateDisplay(0, 0, 0);
		updateDirty = false;
	}

	// DO NOT DELETE: 不要覺得這樣用MutationObserver很沒效率,這是故意的。

	function setupMutationObserver() {
		const observer = new MutationObserver(() => {
			switch (isValidWindow()) {
				case 'valid-changed':
					debugLog(locale.navigation);
					resetAll();
					updateDirty = true;
					initialize();
					state.uiBox.style.display = 'block';
					break;
				case 'valid-unchanged':
					updateDirty = true;
					initialize();
					state.uiBox.style.display = 'block';
					break;
				case 'invalid-changed':
					resetAll();
					if (state.uiBox) state.uiBox.style.display = 'none';
					break;
				case 'invalid-unchanged':
				default:
					// Do nothing
					break;
			}
		});

		observer.observe(document.body, {
			childList: true,
			subtree: true,
			characterData: true
		});
	}

	function initialize() {
		if (state.intervalId) return;
		debugLog(locale.init);
		if (!state.uiBox) createUI();
		state.intervalId = setInterval(() => {
			debugLog('Scheduled update running. UpdateDirty:', updateDirty);
			if (!updateDirty) return;
			updateDirty = false;

			if (state.idleHandle) cancelIdle(state.idleHandle);
			state.idleHandle = safeIdle(() => {
				state.idleHandle = null;
				updateCounter();
			});
		}, CONFIG.updateInterval);
		updateCounter();
	}


	if (document.readyState === 'complete') {
		initialize();
		setupMutationObserver();
	} else {
		window.addEventListener('load', () => {
			initialize();
			setupMutationObserver();
		});
	}

})();