324 lines
8.7 KiB
Go
324 lines
8.7 KiB
Go
package jsonparser
|
|
|
|
import (
|
|
"encoding/json"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/maximhq/bifrost/core/schemas"
|
|
)
|
|
|
|
// getRequestID extracts a unique identifier for the request to maintain state
|
|
func (p *JsonParserPlugin) getRequestID(ctx *schemas.BifrostContext, result *schemas.BifrostResponse) string {
|
|
|
|
// Try to get from result
|
|
if result != nil && result.ChatResponse != nil && result.ChatResponse.ID != "" {
|
|
return result.ChatResponse.ID
|
|
}
|
|
|
|
// Try to get from context if not available in result
|
|
if ctx != nil {
|
|
if requestID, ok := ctx.Value(schemas.BifrostContextKeyRequestID).(string); ok && requestID != "" {
|
|
return requestID
|
|
}
|
|
}
|
|
|
|
return ""
|
|
}
|
|
|
|
// shouldRun determines if the plugin should process the request based on usage type
|
|
func (p *JsonParserPlugin) shouldRun(ctx *schemas.BifrostContext, requestType schemas.RequestType) bool {
|
|
// Run only for chat completion stream requests
|
|
if requestType != schemas.ChatCompletionStreamRequest {
|
|
return false
|
|
}
|
|
|
|
switch p.usage {
|
|
case AllRequests:
|
|
return true
|
|
case PerRequest:
|
|
// Check if the context contains the plugin-specific key
|
|
if ctx != nil {
|
|
if value, ok := ctx.Value(EnableStreamingJSONParser).(bool); ok {
|
|
return value
|
|
}
|
|
}
|
|
return false
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
// accumulateContent adds new content to the accumulated content for a specific request
|
|
func (p *JsonParserPlugin) accumulateContent(requestID, newContent string) string {
|
|
p.mutex.Lock()
|
|
defer p.mutex.Unlock()
|
|
|
|
// Get existing accumulated content
|
|
existing := p.accumulatedContent[requestID]
|
|
|
|
if existing != nil {
|
|
// Append to existing builder
|
|
existing.Content.WriteString(newContent)
|
|
return existing.Content.String()
|
|
} else {
|
|
// Create new builder
|
|
builder := &strings.Builder{}
|
|
builder.WriteString(newContent)
|
|
p.accumulatedContent[requestID] = &AccumulatedContent{
|
|
Content: builder,
|
|
Timestamp: time.Now(),
|
|
}
|
|
return builder.String()
|
|
}
|
|
}
|
|
|
|
// parsePartialJSON parses a JSON string that may be missing closing braces
|
|
func (p *JsonParserPlugin) parsePartialJSON(s string) string {
|
|
// Trim whitespace
|
|
s = strings.TrimSpace(s)
|
|
if s == "" {
|
|
return "{}"
|
|
}
|
|
|
|
// Quick check: if it starts with { or [, it might be JSON
|
|
if s[0] != '{' && s[0] != '[' {
|
|
return s
|
|
}
|
|
|
|
// First, try to parse the string as-is (fast path)
|
|
if p.isValidJSON(s) {
|
|
return s
|
|
}
|
|
|
|
// Use a more efficient approach: build the completion directly
|
|
return p.completeJSON(s)
|
|
}
|
|
|
|
// completeJSON completes partial JSON with O(n) time complexity
|
|
func (p *JsonParserPlugin) completeJSON(s string) string {
|
|
// Pre-allocate buffer with estimated capacity
|
|
capacity := len(s) + 10 // Estimate max 10 closing characters needed
|
|
result := make([]byte, 0, capacity)
|
|
|
|
var stack []byte
|
|
inString := false
|
|
escaped := false
|
|
|
|
// Process the string once
|
|
for i := 0; i < len(s); i++ {
|
|
char := s[i]
|
|
result = append(result, char)
|
|
|
|
if escaped {
|
|
escaped = false
|
|
continue
|
|
}
|
|
|
|
if char == '\\' {
|
|
escaped = true
|
|
continue
|
|
}
|
|
|
|
if char == '"' {
|
|
inString = !inString
|
|
continue
|
|
}
|
|
|
|
if inString {
|
|
continue
|
|
}
|
|
|
|
switch char {
|
|
case '{', '[':
|
|
if char == '{' {
|
|
stack = append(stack, '}')
|
|
} else {
|
|
stack = append(stack, ']')
|
|
}
|
|
case '}', ']':
|
|
if len(stack) > 0 && stack[len(stack)-1] == char {
|
|
stack = stack[:len(stack)-1]
|
|
}
|
|
}
|
|
}
|
|
|
|
// Close any unclosed strings
|
|
if inString {
|
|
if escaped {
|
|
// Remove the trailing backslash
|
|
if len(result) > 0 {
|
|
result = result[:len(result)-1]
|
|
}
|
|
}
|
|
result = append(result, '"')
|
|
}
|
|
|
|
// Add closing characters in reverse order
|
|
for i := len(stack) - 1; i >= 0; i-- {
|
|
result = append(result, stack[i])
|
|
}
|
|
|
|
// Validate the result
|
|
if p.isValidJSON(string(result)) {
|
|
return string(result)
|
|
}
|
|
|
|
// If still invalid, try progressive truncation (but more efficiently)
|
|
return p.progressiveTruncation(s, result)
|
|
}
|
|
|
|
// progressiveTruncation efficiently tries different truncation points
|
|
func (p *JsonParserPlugin) progressiveTruncation(original string, completed []byte) string {
|
|
// Try removing characters from the end until we get valid JSON
|
|
// Use binary search for better performance
|
|
left, right := 0, len(completed)
|
|
|
|
for left < right {
|
|
mid := (left + right) / 2
|
|
candidate := completed[:mid]
|
|
|
|
if p.isValidJSON(string(candidate)) {
|
|
left = mid + 1
|
|
} else {
|
|
right = mid
|
|
}
|
|
}
|
|
|
|
// Try the best candidate
|
|
if left > 0 && p.isValidJSON(string(completed[:left-1])) {
|
|
return string(completed[:left-1])
|
|
}
|
|
|
|
// Fallback to original
|
|
return original
|
|
}
|
|
|
|
// isValidJSON checks if a string is valid JSON
|
|
func (p *JsonParserPlugin) isValidJSON(s string) bool {
|
|
// Trim whitespace
|
|
s = strings.TrimSpace(s)
|
|
|
|
// Empty string after trimming is not valid JSON
|
|
if s == "" {
|
|
return false
|
|
}
|
|
|
|
return json.Valid([]byte(s))
|
|
}
|
|
|
|
// DEEP COPY METHODS
|
|
|
|
// deepCopyBifrostResponse creates a deep copy of BifrostResponse to avoid modifying the original
|
|
func (p *JsonParserPlugin) deepCopyBifrostResponse(original *schemas.BifrostResponse) *schemas.BifrostResponse {
|
|
if original == nil {
|
|
return nil
|
|
}
|
|
|
|
// Create a new BifrostResponse
|
|
result := &schemas.BifrostResponse{}
|
|
|
|
// Copy ChatResponse if it exists (this is what we're interested in for the JSON parser)
|
|
if original.ChatResponse != nil {
|
|
result.ChatResponse = p.deepCopyBifrostChatResponse(original.ChatResponse)
|
|
}
|
|
|
|
// Copy other response types if they exist (shallow copy since we don't modify them)
|
|
result.TextCompletionResponse = original.TextCompletionResponse
|
|
result.ResponsesResponse = original.ResponsesResponse
|
|
result.ResponsesStreamResponse = original.ResponsesStreamResponse
|
|
result.EmbeddingResponse = original.EmbeddingResponse
|
|
result.SpeechResponse = original.SpeechResponse
|
|
result.SpeechStreamResponse = original.SpeechStreamResponse
|
|
result.TranscriptionResponse = original.TranscriptionResponse
|
|
result.TranscriptionStreamResponse = original.TranscriptionStreamResponse
|
|
|
|
return result
|
|
}
|
|
|
|
// deepCopyBifrostChatResponse creates a deep copy of BifrostChatResponse
|
|
func (p *JsonParserPlugin) deepCopyBifrostChatResponse(original *schemas.BifrostChatResponse) *schemas.BifrostChatResponse {
|
|
if original == nil {
|
|
return nil
|
|
}
|
|
|
|
result := &schemas.BifrostChatResponse{
|
|
ID: original.ID,
|
|
Created: original.Created,
|
|
Model: original.Model,
|
|
Object: original.Object,
|
|
ServiceTier: original.ServiceTier,
|
|
SystemFingerprint: original.SystemFingerprint,
|
|
Usage: original.Usage, // Shallow copy - usage shouldn't be modified
|
|
ExtraFields: original.ExtraFields, // Shallow copy
|
|
}
|
|
|
|
// Deep copy Choices slice
|
|
if original.Choices != nil {
|
|
result.Choices = make([]schemas.BifrostResponseChoice, len(original.Choices))
|
|
for i, choice := range original.Choices {
|
|
result.Choices[i] = p.deepCopyBifrostResponseChoice(choice)
|
|
}
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
// deepCopyBifrostResponseChoice creates a deep copy of BifrostResponseChoice
|
|
func (p *JsonParserPlugin) deepCopyBifrostResponseChoice(original schemas.BifrostResponseChoice) schemas.BifrostResponseChoice {
|
|
result := schemas.BifrostResponseChoice{
|
|
Index: original.Index,
|
|
FinishReason: original.FinishReason,
|
|
LogProbs: original.LogProbs,
|
|
}
|
|
|
|
// Deep copy ChatStreamResponseChoice if it exists (this is what we modify)
|
|
if original.ChatStreamResponseChoice != nil {
|
|
result.ChatStreamResponseChoice = p.deepCopyChatStreamResponseChoice(original.ChatStreamResponseChoice)
|
|
}
|
|
|
|
// Shallow copy other choice types since we don't modify them
|
|
result.ChatNonStreamResponseChoice = original.ChatNonStreamResponseChoice
|
|
result.TextCompletionResponseChoice = original.TextCompletionResponseChoice
|
|
|
|
return result
|
|
}
|
|
|
|
// deepCopyChatStreamResponseChoice creates a deep copy of ChatStreamResponseChoice
|
|
func (p *JsonParserPlugin) deepCopyChatStreamResponseChoice(original *schemas.ChatStreamResponseChoice) *schemas.ChatStreamResponseChoice {
|
|
if original == nil {
|
|
return nil
|
|
}
|
|
|
|
result := &schemas.ChatStreamResponseChoice{}
|
|
|
|
// Deep copy Delta pointer if it exists
|
|
if original.Delta != nil {
|
|
result.Delta = p.deepCopyChatStreamResponseChoiceDelta(original.Delta)
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
// deepCopyChatStreamResponseChoiceDelta creates a deep copy of ChatStreamResponseChoiceDelta
|
|
func (p *JsonParserPlugin) deepCopyChatStreamResponseChoiceDelta(original *schemas.ChatStreamResponseChoiceDelta) *schemas.ChatStreamResponseChoiceDelta {
|
|
if original == nil {
|
|
return nil
|
|
}
|
|
|
|
result := &schemas.ChatStreamResponseChoiceDelta{
|
|
Role: original.Role,
|
|
Reasoning: original.Reasoning, // Shallow copy
|
|
Refusal: original.Refusal, // Shallow copy
|
|
ToolCalls: original.ToolCalls, // Shallow copy - we don't modify tool calls
|
|
}
|
|
|
|
// Deep copy Content pointer if it exists (this is what we modify)
|
|
if original.Content != nil {
|
|
contentCopy := *original.Content
|
|
result.Content = &contentCopy
|
|
}
|
|
|
|
return result
|
|
}
|