@juspay/neurolink
Version:
Universal AI Development Platform with working MCP integration, multi-provider support, and professional CLI. Built-in tools operational, 58+ external MCP servers discoverable. Connect to filesystem, GitHub, database operations, and more. Build, test, and
627 lines (626 loc) • 21.7 kB
JavaScript
/**
* SageMaker Streaming Response Parsers
*
* This module provides protocol-specific parsers for different streaming
* formats used by SageMaker endpoints (HuggingFace, LLaMA, custom models).
*/
import { isNonNullObject } from "../../utils/typeUtils.js";
import { createStructuredOutputParser, isStructuredContent, } from "./structured-parser.js";
import { SageMakerError } from "./errors.js";
import { logger } from "../../utils/logger.js";
import { randomUUID } from "crypto";
/**
* Constants for JSON parsing and validation
*/
const MIN_JSON_OBJECT_LENGTH = 2; // Minimum length for JSON object "{}"
/**
* Process a single character for bracket counting logic
* Shared utility to avoid code duplication between parsers
*/
export function processBracketCharacter(char, state) {
if (state.escapeNext) {
state.escapeNext = false;
return { isValid: true };
}
if (char === "\\") {
state.escapeNext = true;
return { isValid: true };
}
if (char === '"' && !state.escapeNext) {
state.inString = !state.inString;
return { isValid: true };
}
// Only count brackets outside of strings
if (!state.inString) {
switch (char) {
case "{":
state.braceCount++;
break;
case "}":
state.braceCount--;
if (state.braceCount < 0) {
return { isValid: false, reason: "Unmatched closing brace" };
}
break;
case "[":
state.bracketCount++;
break;
case "]":
state.bracketCount--;
if (state.bracketCount < 0) {
return { isValid: false, reason: "Unmatched closing bracket" };
}
break;
}
}
return { isValid: true };
}
/**
* Utility function to validate JSON completeness using efficient bracket counting
* Extracted from parseArgumentsForToolCall for reusability
*/
export function validateJSONCompleteness(jsonString) {
// Basic length check - minimum for empty object "{}"
if (jsonString.length < MIN_JSON_OBJECT_LENGTH) {
return { isComplete: false, reason: "Too short" };
}
// Must start and end with braces for object
if (!jsonString.startsWith("{") || !jsonString.endsWith("}")) {
return { isComplete: false, reason: "Missing object braces" };
}
// Use shared bracket counting logic
const state = {
braceCount: 0,
bracketCount: 0,
inString: false,
escapeNext: false,
};
for (let i = 0; i < jsonString.length; i++) {
const char = jsonString[i];
const result = processBracketCharacter(char, state);
if (!result.isValid) {
return { isComplete: false, reason: result.reason };
}
}
// Check for unterminated string
if (state.inString) {
return { isComplete: false, reason: "Unterminated string" };
}
// Check if all brackets are balanced
if (state.braceCount !== 0) {
return {
isComplete: false,
reason: `Unbalanced braces: ${state.braceCount}`,
};
}
if (state.bracketCount !== 0) {
return {
isComplete: false,
reason: `Unbalanced brackets: ${state.bracketCount}`,
};
}
return { isComplete: true };
}
/**
* Utility function to parse tool call arguments with robust validation
* Extracted from parseArgumentsForToolCall for reusability
*/
export function parseToolCallArguments(args) {
const trimmedArgs = args.trim();
// Handle empty arguments
if (trimmedArgs.length === 0) {
return { arguments: "{}", complete: true };
}
// Use robust bracket counting for JSON validation
const validationResult = validateJSONCompleteness(trimmedArgs);
if (validationResult.isComplete) {
try {
const parsed = JSON.parse(trimmedArgs);
// Additional validation: ensure it's actually an object
if (typeof parsed === "object" &&
parsed !== null &&
!Array.isArray(parsed)) {
return { arguments: trimmedArgs, complete: true };
}
else {
// Not a valid object, treat as delta
return { argumentsDelta: trimmedArgs };
}
}
catch (parseError) {
// Log parsing error for debugging
logger.debug("JSON parsing failed for tool arguments", {
args: trimmedArgs.substring(0, 100),
error: formatErrorMessage(parseError),
});
// Not valid JSON despite looking complete, treat as delta
return { argumentsDelta: trimmedArgs };
}
}
else {
// String doesn't look like complete JSON, treat as delta
return { argumentsDelta: trimmedArgs };
}
}
/**
* Abstract base parser with common functionality
*/
class BaseStreamingParser {
buffer = "";
isCompleted = false;
totalUsage;
structuredParser;
responseSchema;
isComplete(chunk) {
return chunk.done === true || chunk.finishReason !== undefined;
}
extractUsage(finalChunk) {
return finalChunk.usage || this.totalUsage;
}
reset() {
this.buffer = "";
this.isCompleted = false;
this.totalUsage = undefined;
this.structuredParser?.reset();
}
/**
* Enable structured output parsing with optional schema
*/
enableStructuredOutput(schema) {
this.responseSchema = schema;
this.structuredParser = createStructuredOutputParser(schema);
}
/**
* Parse structured content if enabled
*/
parseStructuredContent(content) {
if (!this.structuredParser || !isStructuredContent(content)) {
return undefined;
}
return this.structuredParser.parseChunk(content);
}
decodeChunk(chunk) {
return new TextDecoder().decode(chunk);
}
parseJSON(text) {
try {
return JSON.parse(text);
}
catch (error) {
logger.warn("Failed to parse JSON in streaming response", {
text,
error,
});
return null;
}
}
}
/**
* HuggingFace Transformers streaming parser (Server-Sent Events)
*/
export class HuggingFaceStreamParser extends BaseStreamingParser {
getName() {
return "HuggingFace SSE Parser";
}
parse(chunk) {
const text = this.decodeChunk(chunk);
this.buffer += text;
const chunks = [];
const lines = this.buffer.split("\n");
// Keep the last potentially incomplete line in buffer
this.buffer = lines.pop() || "";
for (const line of lines) {
const trimmed = line.trim();
// Skip empty lines and comments
if (!trimmed || trimmed.startsWith(":")) {
continue;
}
// Parse Server-Sent Events format
if (trimmed.startsWith("data: ")) {
const data = trimmed.substring(6);
// Check for stream end
if (data === "[DONE]") {
chunks.push({
content: "",
done: true,
finishReason: "stop",
});
this.isCompleted = true;
continue;
}
// Parse JSON data
const parsed = this.parseJSON(data);
if (parsed && isNonNullObject(parsed)) {
const chunk = this.parseHuggingFaceChunk(parsed);
if (chunk) {
chunks.push(chunk);
}
}
}
}
return chunks;
}
parseHuggingFaceChunk(data) {
// HuggingFace streaming format
if (data.token) {
const token = data.token;
return {
content: typeof token.text === "string" ? token.text : String(token),
done: false,
};
}
// Alternative format with generated_text
if (data.generated_text !== undefined) {
const details = data.details;
return {
content: String(data.generated_text),
done: details?.finish_reason !== undefined,
finishReason: this.mapFinishReason(details?.finish_reason),
usage: details ? this.parseHuggingFaceUsage(details) : undefined,
};
}
// Error format
if (data.error) {
const errorMessage = extractApiErrorMessage(data.error);
throw new SageMakerError(`HuggingFace streaming error: ${errorMessage}`, "MODEL_ERROR", 500);
}
return null;
}
parseHuggingFaceUsage(details) {
if (!details.tokens) {
return undefined;
}
const tokens = details.tokens;
return {
promptTokens: Number(tokens.input) || 0,
completionTokens: Number(tokens.generated) || 0,
totalTokens: Number(tokens.total) || 0,
};
}
mapFinishReason(reason) {
switch (reason) {
case "stop":
return "stop";
case "length":
return "length";
case "eos_token":
return "stop";
default:
return "unknown";
}
}
}
/**
* LLaMA/OpenAI-compatible streaming parser (JSON Lines)
*/
export class LlamaStreamParser extends BaseStreamingParser {
getName() {
return "LLaMA JSONL Parser";
}
parse(chunk) {
const text = this.decodeChunk(chunk);
this.buffer += text;
const chunks = [];
const lines = this.buffer.split("\n");
// Keep the last potentially incomplete line in buffer
this.buffer = lines.pop() || "";
for (const line of lines) {
const trimmed = line.trim();
if (!trimmed) {
continue;
}
// Parse each line as JSON
const parsed = this.parseJSON(trimmed);
if (parsed && isNonNullObject(parsed)) {
const chunk = this.parseLlamaChunk(parsed);
if (chunk) {
chunks.push(chunk);
}
}
}
return chunks;
}
parseLlamaChunk(data) {
// OpenAI-compatible format
if (Array.isArray(data.choices) && data.choices[0]) {
const choice = data.choices[0];
// Delta format (streaming)
if (choice.delta) {
const delta = choice.delta;
const content = String(delta.content || "");
const finishReason = choice.finish_reason;
const chunk = {
content,
done: finishReason !== null && finishReason !== undefined,
finishReason: this.mapFinishReason(finishReason || null),
usage: data.usage
? this.parseLlamaUsage(data.usage)
: undefined,
};
// Phase 2.3: Handle structured output if enabled
if (content && this.structuredParser) {
chunk.structuredOutput = this.parseStructuredContent(content);
}
// Phase 2.3: Handle streaming tool calls
if (Array.isArray(delta.tool_calls) && delta.tool_calls[0]) {
const toolCall = delta.tool_calls[0];
chunk.toolCall = this.parseStreamingToolCall(toolCall);
// If tool call is complete and we have finish_reason, mark chunk complete
if (finishReason === "function_call" && chunk.toolCall.arguments) {
chunk.toolCall.complete = true;
}
}
return chunk;
}
// Text format (non-streaming fallback)
if (choice.text !== undefined) {
return {
content: String(choice.text),
done: choice.finish_reason !== null,
finishReason: this.mapFinishReason(choice.finish_reason),
usage: data.usage
? this.parseLlamaUsage(data.usage)
: undefined,
};
}
}
// Direct content format
if (data.content !== undefined) {
return {
content: String(data.content),
done: Boolean(data.done),
finishReason: this.mapFinishReason(data.finish_reason),
usage: data.usage
? this.parseLlamaUsage(data.usage)
: undefined,
};
}
// Error format
if (data.error) {
const errorData = data.error;
const errorMessage = extractApiErrorMessage(errorData);
throw new SageMakerError(`LLaMA streaming error: ${errorMessage}`, "MODEL_ERROR", 500);
}
return null;
}
/**
* Parse tool call arguments with robust validation and error handling
*/
parseToolCallArguments(functionData, toolCall) {
if (typeof functionData.arguments === "string") {
const result = parseToolCallArguments(functionData.arguments);
if (result.complete) {
toolCall.arguments = result.arguments;
toolCall.complete = true;
}
else if (result.argumentsDelta !== undefined) {
toolCall.argumentsDelta = result.argumentsDelta;
}
}
else if (functionData.arguments !== undefined) {
// Handle non-string arguments (objects, numbers, etc.)
try {
toolCall.arguments = JSON.stringify(functionData.arguments);
toolCall.complete = true;
}
catch (stringifyError) {
logger.warn("Failed to stringify tool arguments", {
args: functionData.arguments,
error: formatErrorMessage(stringifyError),
});
toolCall.arguments = "{}";
toolCall.complete = true;
}
}
else {
// No arguments provided, default to empty object
toolCall.arguments = "{}";
toolCall.complete = true;
}
}
/**
* Parse streaming tool call from OpenAI-compatible format (Phase 2.3)
*/
parseStreamingToolCall(toolCallData) {
const toolCall = {
id: String(toolCallData.id || `call_${randomUUID()}`),
type: "function",
};
// Handle function name (usually sent in first chunk)
const functionData = toolCallData.function;
if (functionData?.name) {
toolCall.name = String(functionData.name);
}
// Handle streaming arguments
if (functionData?.arguments) {
this.parseToolCallArguments(functionData, toolCall);
}
return toolCall;
}
parseLlamaUsage(usage) {
return {
promptTokens: Number(usage.prompt_tokens) || 0,
completionTokens: Number(usage.completion_tokens) || 0,
totalTokens: Number(usage.total_tokens) || 0,
};
}
mapFinishReason(reason) {
switch (reason) {
case "stop":
return "stop";
case "length":
return "length";
case "function_call":
return "tool-calls";
case "content_filter":
return "content-filter";
default:
return reason ? "unknown" : undefined;
}
}
}
/**
* Custom/Generic streaming parser (Chunked Transfer)
*/
export class CustomStreamParser extends BaseStreamingParser {
expectedFormat = "json";
constructor(format = "json") {
super();
this.expectedFormat = format;
}
getName() {
return `Custom ${this.expectedFormat.toUpperCase()} Parser`;
}
parse(chunk) {
const text = this.decodeChunk(chunk);
this.buffer += text;
if (this.expectedFormat === "json") {
return this.parseJSONFormat();
}
else {
return this.parseTextFormat();
}
}
parseJSONFormat() {
const chunks = [];
// Try to parse complete JSON objects
let startIndex = 0;
while (startIndex < this.buffer.length) {
try {
const remaining = this.buffer.substring(startIndex);
const parsed = JSON.parse(remaining);
// Successfully parsed - this is probably the complete response
const chunk = {
content: parsed.text ||
parsed.generated_text ||
parsed.output ||
String(parsed),
done: true,
finishReason: "stop",
};
chunks.push(chunk);
this.buffer = "";
this.isCompleted = true;
break;
}
catch {
// JSON parsing failed - look for newline-separated JSON
const newlineIndex = this.buffer.indexOf("\n", startIndex);
if (newlineIndex === -1) {
// No complete line found, wait for more data
break;
}
const line = this.buffer.substring(startIndex, newlineIndex);
if (line.trim()) {
const parsed = this.parseJSON(line.trim());
if (parsed && isNonNullObject(parsed)) {
const chunk = this.parseCustomChunk(parsed);
if (chunk) {
chunks.push(chunk);
}
}
}
startIndex = newlineIndex + 1;
}
}
// Clean processed content from buffer
if (startIndex > 0) {
this.buffer = this.buffer.substring(startIndex);
}
return chunks;
}
parseTextFormat() {
// Simple text streaming - treat each chunk as content
if (this.buffer) {
const content = this.buffer;
this.buffer = "";
return [
{
content,
done: false,
},
];
}
return [];
}
parseCustomChunk(data) {
// Generic parsing for various custom formats
const content = data.text ||
data.generated_text ||
data.output ||
data.response ||
data.content ||
(typeof data === "string" ? data : JSON.stringify(data));
return {
content: String(content),
done: Boolean(data.done || data.finished || data.complete),
finishReason: data.finish_reason || data.status === "complete" ? "stop" : undefined,
usage: data.usage || data.tokens ? this.parseCustomUsage(data) : undefined,
};
}
parseCustomUsage(data) {
const usage = (data.usage || data.tokens || {});
return {
promptTokens: Number(usage.prompt_tokens || usage.input_tokens) || 0,
completionTokens: Number(usage.completion_tokens || usage.output_tokens) || 0,
totalTokens: Number(usage.total_tokens) || 0,
};
}
}
/**
* Parser factory to create appropriate parser for detected protocol
*/
export class StreamingParserFactory {
static createParser(protocol, options) {
switch (protocol) {
case "sse":
return new HuggingFaceStreamParser();
case "jsonl":
return new LlamaStreamParser();
case "chunked": {
const format = options?.format;
return new CustomStreamParser(format || "json");
}
case "none":
default:
// Return a no-op parser that just converts complete responses
return new CustomStreamParser("text");
}
}
static getSupportedProtocols() {
return ["sse", "jsonl", "chunked", "text"];
}
}
/**
* Helper function to safely format error messages for logging and error handling
* Consolidates the duplicated error formatting pattern used throughout the parsers
*/
function formatErrorMessage(error) {
if (error instanceof Error) {
return error.message;
}
return String(error);
}
/**
* Helper function to extract error messages from API response error data
* Handles both string and object error formats consistently
*/
function extractApiErrorMessage(errorData) {
if (isNonNullObject(errorData)) {
return errorData.message || String(errorData);
}
return String(errorData);
}
/**
* Utility function to estimate token usage when not provided
*/
export function estimateTokenUsage(prompt, completion) {
// Rough estimation: ~4 characters per token for English text
const promptTokens = Math.ceil(prompt.length / 4);
const completionTokens = Math.ceil(completion.length / 4);
return {
promptTokens,
completionTokens,
totalTokens: promptTokens + completionTokens,
};
}