llama-cpp-K2-FAST2 / tools /ui /src /lib /utils /parse-partial-json-args.ts

llama.cpp k2-fast2: MBZUAI model/K2Horizon + FAST2 sliding-window override

c59fd93 verified about 1 month ago

3.73 kB

// JSON delimiters used while scanning partial streamed JSON. Single-char
// tokens so they only need eq-comparison, but naming them keeps the
// scanner readable and keeps the literal source-of-truth in one place.
const JSON_QUOTE = '"';
const JSON_BACKSLASH = '\\';
const JSON_OBJECT_OPEN = '{';
const JSON_OBJECT_CLOSE = '}';
const JSON_ARRAY_OPEN = '[';
const JSON_ARRAY_CLOSE = ']';
// Trailing punctuation to strip before re-closing a partial object/array.
// Matches an optional trailing comma plus any trailing whitespace; lets
// us re-emit a syntactically-valid JSON document without an orphaned
// comma when the model cut off mid-key.
const TRAILING_JSON_PUNCTUATION_REGEX = /,?\s*$/;
/** Bounded cache for parsePartialJsonArgs results. */
const PARTIAL_JSON_CACHE_MAX_SIZE = 32;
const partialJsonCache = new Map<string, Record<string, unknown> | null>();

function cacheResult(input: string, result: Record<string, unknown> | null): void {
	if (partialJsonCache.size >= PARTIAL_JSON_CACHE_MAX_SIZE) {
		partialJsonCache.delete(partialJsonCache.keys().next().value!);
	}

	partialJsonCache.set(input, result);
}

// Parse partial tool-arg JSON streamed token-by-token. Closes any
// unterminated string and dangling open containers (in reverse order),
// so parsers can still surface keys already received while the call
// is still in flight. Memoized: the char-by-char scanner runs on every
// render during streaming even when toolArgs hasn't changed.
export function parsePartialJsonArgs(toolArgsString: string): Record<string, unknown> | null {
	const cached = partialJsonCache.get(toolArgsString);

	if (cached !== undefined) return cached;

	let result: Record<string, unknown> | null;

	try {
		const parsed: unknown = JSON.parse(toolArgsString);

		result =
			parsed && typeof parsed === 'object' && !Array.isArray(parsed)
				? (parsed as Record<string, unknown>)
				: null;
	} catch {
		result = scanPartialJson(toolArgsString);
	}

	cacheResult(toolArgsString, result);

	return result;
}

/** Char-by-char scanner for unterminated partial JSON. */
function scanPartialJson(toolArgsString: string): Record<string, unknown> | null {
	let inString = false;
	let escape = false;

	const stack: ('{' | '[')[] = [];

	for (let i = 0; i < toolArgsString.length; i++) {
		const ch = toolArgsString[i];

		if (escape) {
			escape = false;

			continue;
		}

		if (ch === JSON_BACKSLASH && inString) {
			escape = true;

			continue;
		}

		if (ch === JSON_QUOTE) {
			inString = !inString;

			continue;
		}

		if (inString) continue;

		if (ch === JSON_OBJECT_OPEN) stack.push(JSON_OBJECT_OPEN);
		else if (ch === JSON_OBJECT_CLOSE) {
			if (stack.length === 0 || stack[stack.length - 1] !== JSON_OBJECT_OPEN) return null;

			stack.pop();
		} else if (ch === JSON_ARRAY_OPEN) stack.push(JSON_ARRAY_OPEN);
		else if (ch === JSON_ARRAY_CLOSE) {
			if (stack.length === 0 || stack[stack.length - 1] !== JSON_ARRAY_OPEN) return null;

			stack.pop();
		}
	}

	let completed = toolArgsString;

	if (escape) {
		// Dangling escape at end of partial JSON: escape the trailing
		// backslash as a literal so we can close the string cleanly.
		completed += JSON_BACKSLASH;
	}

	if (inString) completed += JSON_QUOTE;

	if (!inString) completed = completed.replace(TRAILING_JSON_PUNCTUATION_REGEX, '');

	// Close in reverse nesting order: innermost container first.
	for (let i = stack.length - 1; i >= 0; i--) {
		completed += stack[i] === JSON_OBJECT_OPEN ? JSON_OBJECT_CLOSE : JSON_ARRAY_CLOSE;
	}

	try {
		const parsed: unknown = JSON.parse(completed);

		return parsed && typeof parsed === 'object' && !Array.isArray(parsed)
			? (parsed as Record<string, unknown>)
			: null;
	} catch {
		return null;
	}
}

	// JSON delimiters used while scanning partial streamed JSON. Single-char
	// tokens so they only need eq-comparison, but naming them keeps the
	// scanner readable and keeps the literal source-of-truth in one place.
	const JSON_QUOTE = '"';
	const JSON_BACKSLASH = '\\';
	const JSON_OBJECT_OPEN = '{';
	const JSON_OBJECT_CLOSE = '}';
	const JSON_ARRAY_OPEN = '[';
	const JSON_ARRAY_CLOSE = ']';
	// Trailing punctuation to strip before re-closing a partial object/array.
	// Matches an optional trailing comma plus any trailing whitespace; lets
	// us re-emit a syntactically-valid JSON document without an orphaned
	// comma when the model cut off mid-key.
	const TRAILING_JSON_PUNCTUATION_REGEX = /,?\s*$/;
	/** Bounded cache for parsePartialJsonArgs results. */
	const PARTIAL_JSON_CACHE_MAX_SIZE = 32;
	const partialJsonCache = new Map<string, Record<string, unknown> \| null>();

	function cacheResult(input: string, result: Record<string, unknown> \| null): void {
	if (partialJsonCache.size >= PARTIAL_JSON_CACHE_MAX_SIZE) {
	partialJsonCache.delete(partialJsonCache.keys().next().value!);
	}

	partialJsonCache.set(input, result);
	}

	// Parse partial tool-arg JSON streamed token-by-token. Closes any
	// unterminated string and dangling open containers (in reverse order),
	// so parsers can still surface keys already received while the call
	// is still in flight. Memoized: the char-by-char scanner runs on every
	// render during streaming even when toolArgs hasn't changed.
	export function parsePartialJsonArgs(toolArgsString: string): Record<string, unknown> \| null {
	const cached = partialJsonCache.get(toolArgsString);

	if (cached !== undefined) return cached;

	let result: Record<string, unknown> \| null;

	try {
	const parsed: unknown = JSON.parse(toolArgsString);

	result =
	parsed && typeof parsed === 'object' && !Array.isArray(parsed)
	? (parsed as Record<string, unknown>)
	: null;
	} catch {
	result = scanPartialJson(toolArgsString);
	}

	cacheResult(toolArgsString, result);

	return result;
	}

	/** Char-by-char scanner for unterminated partial JSON. */
	function scanPartialJson(toolArgsString: string): Record<string, unknown> \| null {
	let inString = false;
	let escape = false;

	const stack: ('{' \| '[')[] = [];

	for (let i = 0; i < toolArgsString.length; i++) {
	const ch = toolArgsString[i];

	if (escape) {
	escape = false;

	continue;
	}

	if (ch === JSON_BACKSLASH && inString) {
	escape = true;

	continue;
	}

	if (ch === JSON_QUOTE) {
	inString = !inString;

	continue;
	}

	if (inString) continue;

	if (ch === JSON_OBJECT_OPEN) stack.push(JSON_OBJECT_OPEN);
	else if (ch === JSON_OBJECT_CLOSE) {
	if (stack.length === 0 \|\| stack[stack.length - 1] !== JSON_OBJECT_OPEN) return null;

	stack.pop();
	} else if (ch === JSON_ARRAY_OPEN) stack.push(JSON_ARRAY_OPEN);
	else if (ch === JSON_ARRAY_CLOSE) {
	if (stack.length === 0 \|\| stack[stack.length - 1] !== JSON_ARRAY_OPEN) return null;

	stack.pop();
	}
	}

	let completed = toolArgsString;

	if (escape) {
	// Dangling escape at end of partial JSON: escape the trailing
	// backslash as a literal so we can close the string cleanly.
	completed += JSON_BACKSLASH;
	}

	if (inString) completed += JSON_QUOTE;

	if (!inString) completed = completed.replace(TRAILING_JSON_PUNCTUATION_REGEX, '');

	// Close in reverse nesting order: innermost container first.
	for (let i = stack.length - 1; i >= 0; i--) {
	completed += stack[i] === JSON_OBJECT_OPEN ? JSON_OBJECT_CLOSE : JSON_ARRAY_CLOSE;
	}

	try {
	const parsed: unknown = JSON.parse(completed);

	return parsed && typeof parsed === 'object' && !Array.isArray(parsed)
	? (parsed as Record<string, unknown>)
	: null;
	} catch {
	return null;
	}
	}