first commit

2026-04-26 21:52:23 +03:00
commit 880f412e2c
2662 changed files with 866266 additions and 0 deletions
--- a/plugins/maxim/main.go
+++ b/plugins/maxim/main.go
@@ -0,0 +1,688 @@
+// Package maxim provides integration for Maxim's SDK as a Bifrost plugin.
+// This file contains the main plugin implementation.
+package maxim
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/bytedance/sonic"
+	"github.com/google/uuid"
+	bifrost "github.com/maximhq/bifrost/core"
+	"github.com/maximhq/bifrost/core/schemas"
+	"github.com/maximhq/bifrost/framework/streaming"
+
+	"github.com/maximhq/maxim-go"
+	"github.com/maximhq/maxim-go/logging"
+	maximSchemas "github.com/maximhq/maxim-go/schemas"
+)
+
+// PluginName is the canonical name for the maxim plugin.
+const (
+	PluginName         string = "maxim"
+	PluginLoggerPrefix string = "[Maxim Plugin]"
+)
+
+// Config is the configuration for the maxim plugin.
+//   - APIKey: API key for Maxim SDK authentication
+//   - LogRepoID: Optional default ID for the Maxim logger instance
+type Config struct {
+	LogRepoID string `json:"log_repo_id,omitempty"` // Optional - can be empty
+	APIKey    string `json:"api_key"`
+}
+
+// Plugin implements the schemas.LLMPlugin interface for Maxim's logger.
+// It provides request and response tracing functionality using Maxim logger,
+// allowing detailed tracking of requests and responses across different log repositories.
+//
+// Fields:
+//   - mx: The Maxim SDK instance for creating new loggers
+//   - defaultLogRepoId: Default log repository ID from config (optional)
+//   - loggers: Map of log repo ID to logger instances
+//   - loggerMutex: RW mutex for thread-safe access to loggers map
+type Plugin struct {
+	mx               *maxim.Maxim
+	defaultLogRepoID string
+	loggers          map[string]*logging.Logger
+	loggerMutex      *sync.RWMutex
+	logger           schemas.Logger
+}
+
+// Init initializes and returns a Plugin instance for Maxim's logger.
+//
+// Parameters:
+//   - config: Configuration for the maxim plugin
+//
+// Returns:
+//   - schemas.LLMPlugin: A configured plugin instance for request/response tracing
+//   - error: Any error that occurred during plugin initialization
+func Init(config *Config, logger schemas.Logger) (schemas.LLMPlugin, error) {
+	if config == nil {
+		return nil, fmt.Errorf("config is required")
+	}
+	// check if Maxim Logger variables are set
+	if config.APIKey == "" {
+		return nil, fmt.Errorf("apiKey is not set")
+	}
+
+	mx := maxim.Init(&maxim.MaximSDKConfig{ApiKey: config.APIKey})
+
+	plugin := &Plugin{
+		mx:               mx,
+		defaultLogRepoID: config.LogRepoID,
+		loggers:          make(map[string]*logging.Logger),
+		loggerMutex:      &sync.RWMutex{},
+		logger:           logger,
+	}
+
+	// Initialize default logger if LogRepoId is provided
+	if config.LogRepoID != "" {
+		logger, err := mx.GetLogger(&logging.LoggerConfig{Id: config.LogRepoID})
+		if err != nil {
+			return nil, fmt.Errorf("failed to initialize default logger: %w", err)
+		}
+		plugin.loggers[config.LogRepoID] = logger
+	}
+
+	return plugin, nil
+}
+
+// TraceIDKey is the context key used to store and retrieve trace IDs.
+// This constant provides a consistent key for tracking request traces
+// throughout the request/response lifecycle.
+const (
+	SessionIDKey      schemas.BifrostContextKey = "session-id"
+	TraceIDKey        schemas.BifrostContextKey = "trace-id"
+	TraceNameKey      schemas.BifrostContextKey = "trace-name"
+	GenerationIDKey   schemas.BifrostContextKey = "generation-id"
+	GenerationNameKey schemas.BifrostContextKey = "generation-name"
+	TagsKey           schemas.BifrostContextKey = "maxim-tags"
+	LogRepoIDKey      schemas.BifrostContextKey = "log-repo-id"
+)
+
+// convertAccResultToProcessedStreamResponse converts StreamAccumulatorResult to ProcessedStreamResponse
+func convertAccResultToProcessedStreamResponse(accResult *schemas.StreamAccumulatorResult) *streaming.ProcessedStreamResponse {
+	if accResult == nil {
+		return nil
+	}
+	// Determine StreamType based on the response content
+	streamType := streaming.StreamTypeChat
+	if accResult.AudioOutput != nil {
+		streamType = streaming.StreamTypeAudio
+	} else if accResult.TranscriptionOutput != nil {
+		streamType = streaming.StreamTypeTranscription
+	} else if len(accResult.OutputMessages) > 0 {
+		streamType = streaming.StreamTypeResponses
+	} else if accResult.ImageGenerationOutput != nil {
+		streamType = streaming.StreamTypeImage
+	}
+	return &streaming.ProcessedStreamResponse{
+		RequestID:      accResult.RequestID,
+		StreamType:     streamType,
+		RequestedModel: accResult.RequestedModel,
+		ResolvedModel:  accResult.ResolvedModel,
+		Provider:       accResult.Provider,
+		Data: &streaming.AccumulatedData{
+			Status:              accResult.Status,
+			Latency:             accResult.Latency,
+			TimeToFirstToken:    accResult.TimeToFirstToken,
+			OutputMessage:       accResult.OutputMessage,
+			OutputMessages:      accResult.OutputMessages,
+			TokenUsage:          accResult.TokenUsage,
+			Cost:                accResult.Cost,
+			ErrorDetails:        accResult.ErrorDetails,
+			AudioOutput:         accResult.AudioOutput,
+			TranscriptionOutput: accResult.TranscriptionOutput,
+			FinishReason:        accResult.FinishReason,
+			RawResponse:         accResult.RawResponse,
+		},
+		RawRequest: &accResult.RawRequest,
+	}
+}
+
+// The plugin provides request/response tracing functionality by integrating with Maxim's logging system.
+// It supports both chat completion and text completion requests, tracking the entire lifecycle of each request
+// including inputs, parameters, and responses.
+//
+// Key Features:
+// - Automatic trace and generation ID management
+// - Support for both chat and text completion requests
+// - Contextual tracking across request lifecycle
+// - Graceful handling of existing trace/generation IDs
+//
+// The plugin uses context values to maintain trace and generation IDs throughout the request lifecycle.
+// These IDs can be propagated from external systems through HTTP headers (x-bf-maxim-trace-id and x-bf-maxim-generation-id).
+
+// GetName returns the name of the plugin.
+func (plugin *Plugin) GetName() string {
+	return PluginName
+}
+
+// HTTPTransportPreHook is not used for this plugin
+func (plugin *Plugin) HTTPTransportPreHook(ctx *schemas.BifrostContext, req *schemas.HTTPRequest) (*schemas.HTTPResponse, error) {
+	return nil, nil
+}
+
+// HTTPTransportPostHook is not used for this plugin
+func (plugin *Plugin) HTTPTransportPostHook(ctx *schemas.BifrostContext, req *schemas.HTTPRequest, resp *schemas.HTTPResponse) error {
+	return nil
+}
+
+// HTTPTransportStreamChunkHook passes through streaming chunks unchanged
+func (plugin *Plugin) HTTPTransportStreamChunkHook(ctx *schemas.BifrostContext, req *schemas.HTTPRequest, chunk *schemas.BifrostStreamChunk) (*schemas.BifrostStreamChunk, error) {
+	return chunk, nil
+}
+
+// getEffectiveLogRepoID determines which single log repo ID to use based on priority:
+// 1. Header log repo ID (if provided)
+// 2. Default log repo ID from config (if configured)
+// 3. Empty string (skip logging)
+func (plugin *Plugin) getEffectiveLogRepoID(ctx *schemas.BifrostContext) string {
+	// Check for header log repo ID first (highest priority)
+	if ctx != nil {
+		if headerRepoID, ok := ctx.Value(LogRepoIDKey).(string); ok && headerRepoID != "" {
+			return headerRepoID
+		}
+	}
+
+	// Fall back to default log repo ID from config
+	if plugin.defaultLogRepoID != "" {
+		return plugin.defaultLogRepoID
+	}
+
+	// Return empty string if neither header nor default is available
+	return ""
+}
+
+// getOrCreateLogger gets an existing logger or creates a new one for the given log repo ID
+func (plugin *Plugin) getOrCreateLogger(logRepoID string) (*logging.Logger, error) {
+	// First, try to get existing logger (read lock)
+	plugin.loggerMutex.RLock()
+	if logger, exists := plugin.loggers[logRepoID]; exists {
+		plugin.loggerMutex.RUnlock()
+		return logger, nil
+	}
+	plugin.loggerMutex.RUnlock()
+
+	// Logger doesn't exist, create it (write lock)
+	plugin.loggerMutex.Lock()
+	defer plugin.loggerMutex.Unlock()
+
+	// Double-check in case another goroutine created it while we were waiting
+	if logger, exists := plugin.loggers[logRepoID]; exists {
+		return logger, nil
+	}
+
+	// Create new logger
+	logger, err := plugin.mx.GetLogger(&logging.LoggerConfig{Id: logRepoID})
+	if err != nil {
+		return nil, fmt.Errorf("failed to create logger for repo ID %s: %w", logRepoID, err)
+	}
+
+	plugin.loggers[logRepoID] = logger
+	return logger, nil
+}
+
+// PreLLMHook is called before a request is processed by Bifrost.
+// It manages trace and generation tracking for incoming requests by either:
+// - Creating a new trace if none exists
+// - Reusing an existing trace ID from the context
+// - Creating a new generation within an existing trace
+// - Skipping trace/generation creation if they already exist
+//
+// The function handles both chat completion and text completion requests,
+// capturing relevant metadata such as:
+// - Request type (chat/text completion)
+// - Model information
+// - Message content and role
+// - Model parameters
+//
+// Parameters:
+//   - ctx: Pointer to the schemas.BifrostContext that may contain existing trace/generation IDs
+//   - req: The incoming Bifrost request to be traced
+//
+// Returns:
+//   - *schemas.BifrostRequest: The original request, unmodified
+//   - error: Any error that occurred during trace/generation creation
+func (plugin *Plugin) PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) {
+	if req != nil && req.RequestType == schemas.RealtimeRequest {
+		return req, nil, nil
+	}
+
+	var traceID string
+	var traceName string
+	var sessionID string
+	var generationName string
+
+	// Get effective log repo ID (header > default > skip)
+	effectiveLogRepoID := plugin.getEffectiveLogRepoID(ctx)
+
+	// If no log repo ID available, skip logging
+	if effectiveLogRepoID == "" {
+		return req, nil, nil
+	}
+
+	// Check if context already has traceID and generationID
+	if ctx != nil {
+		if existingGenerationID, ok := ctx.Value(GenerationIDKey).(string); ok && existingGenerationID != "" {
+			// If generationID exists, return early
+			return req, nil, nil
+		}
+
+		if existingTraceID, ok := ctx.Value(TraceIDKey).(string); ok && existingTraceID != "" {
+			// If traceID exists, and no generationID, create a new generation on the trace
+			traceID = existingTraceID
+		}
+
+		if existingSessionID, ok := ctx.Value(SessionIDKey).(string); ok && existingSessionID != "" {
+			sessionID = existingSessionID
+		}
+
+		if existingTraceName, ok := ctx.Value(TraceNameKey).(string); ok && existingTraceName != "" {
+			traceName = existingTraceName
+		}
+
+		if existingGenerationName, ok := ctx.Value(GenerationNameKey).(string); ok && existingGenerationName != "" {
+			generationName = existingGenerationName
+		}
+	}
+
+	provider, model, _ := req.GetRequestFields()
+
+	// Determine request type and set appropriate tags
+	var messages []maximSchemas.CompletionRequest
+	var latestMessage string
+
+	modelParams := make(map[string]interface{})
+
+	switch req.RequestType {
+	case schemas.TextCompletionRequest, schemas.TextCompletionStreamRequest:
+		messages = append(messages, maximSchemas.CompletionRequest{
+			Role:    string(schemas.ChatMessageRoleUser),
+			Content: req.TextCompletionRequest.Input,
+		})
+		if req.TextCompletionRequest.Input.PromptStr != nil {
+			latestMessage = *req.TextCompletionRequest.Input.PromptStr
+		} else {
+			var stringBuilder strings.Builder
+			for _, prompt := range req.TextCompletionRequest.Input.PromptArray {
+				stringBuilder.WriteString(prompt)
+			}
+			latestMessage = stringBuilder.String()
+		}
+
+		if req.TextCompletionRequest.Params != nil {
+			// Convert the struct to a map using reflection or JSON marshaling
+			jsonData, err := sonic.Marshal(req.TextCompletionRequest.Params)
+			if err == nil {
+				sonic.Unmarshal(jsonData, &modelParams)
+			}
+		}
+	case schemas.ChatCompletionRequest, schemas.ChatCompletionStreamRequest:
+		for _, message := range req.ChatRequest.Input {
+			messages = append(messages, maximSchemas.CompletionRequest{
+				Role:    string(message.Role),
+				Content: message.Content,
+			})
+		}
+		if len(req.ChatRequest.Input) > 0 {
+			lastMsg := req.ChatRequest.Input[len(req.ChatRequest.Input)-1]
+			if lastMsg.Content.ContentStr != nil {
+				latestMessage = *lastMsg.Content.ContentStr
+			} else if lastMsg.Content.ContentBlocks != nil {
+				// Find the last text content block
+				for i := len(lastMsg.Content.ContentBlocks) - 1; i >= 0; i-- {
+					block := (lastMsg.Content.ContentBlocks)[i]
+					if block.Type == schemas.ChatContentBlockTypeText && block.Text != nil {
+						latestMessage = *block.Text
+						break
+					}
+				}
+				// If no text block found, use placeholder
+				if latestMessage == "" {
+					latestMessage = "-"
+				}
+			}
+		}
+
+		if req.ChatRequest.Params != nil {
+			// Convert the struct to a map using reflection or JSON marshaling
+			jsonData, err := sonic.Marshal(req.ChatRequest.Params)
+			if err == nil {
+				sonic.Unmarshal(jsonData, &modelParams)
+			}
+		}
+	case schemas.ResponsesRequest, schemas.ResponsesStreamRequest, schemas.WebSocketResponsesRequest:
+		for _, message := range req.ResponsesRequest.Input {
+			if message.Content != nil {
+				role := schemas.ChatMessageRoleUser
+				if message.Role != nil {
+					role = schemas.ChatMessageRole(*message.Role)
+				}
+				messages = append(messages, maximSchemas.CompletionRequest{
+					Role:    string(role),
+					Content: message.Content,
+				})
+			}
+		}
+		if len(req.ResponsesRequest.Input) > 0 {
+			lastMsg := req.ResponsesRequest.Input[len(req.ResponsesRequest.Input)-1]
+			// Initialize to placeholder in case content is missing or empty
+			latestMessage = "-"
+
+			// Check if Content is nil before accessing its fields
+			if lastMsg.Content != nil {
+				if lastMsg.Content.ContentStr != nil {
+					latestMessage = *lastMsg.Content.ContentStr
+				} else if lastMsg.Content.ContentBlocks != nil {
+					// Find the last text content block
+					for i := len(lastMsg.Content.ContentBlocks) - 1; i >= 0; i-- {
+						block := (lastMsg.Content.ContentBlocks)[i]
+						if block.Text != nil {
+							latestMessage = *block.Text
+							break
+						}
+					}
+					// If no text block found, keep the placeholder
+				}
+			}
+		}
+
+		if req.ResponsesRequest.Params != nil {
+			// Convert the struct to a map using reflection or JSON marshaling
+			jsonData, err := sonic.Marshal(req.ResponsesRequest.Params)
+			if err == nil {
+				sonic.Unmarshal(jsonData, &modelParams)
+			}
+		}
+	case schemas.ImageGenerationRequest, schemas.ImageGenerationStreamRequest:
+		if req.ImageGenerationRequest == nil || req.ImageGenerationRequest.Input == nil {
+			break
+		}
+		messages = append(messages, maximSchemas.CompletionRequest{
+			Role:    string(schemas.ChatMessageRoleUser),
+			Content: req.ImageGenerationRequest.Input.Prompt,
+		})
+		latestMessage = req.ImageGenerationRequest.Input.Prompt
+		if req.ImageGenerationRequest.Params != nil {
+			jsonData, err := sonic.Marshal(req.ImageGenerationRequest.Params)
+			if err == nil {
+				sonic.Unmarshal(jsonData, &modelParams)
+			}
+		}
+	case schemas.ImageEditRequest, schemas.ImageEditStreamRequest:
+		if req.ImageEditRequest == nil || req.ImageEditRequest.Input == nil {
+			break
+		}
+		messages = append(messages, maximSchemas.CompletionRequest{
+			Role:    string(schemas.ChatMessageRoleUser),
+			Content: req.ImageEditRequest.Input.Prompt,
+		})
+		latestMessage = req.ImageEditRequest.Input.Prompt
+		if req.ImageEditRequest.Params != nil {
+			jsonData, err := sonic.Marshal(req.ImageEditRequest.Params)
+			if err == nil {
+				sonic.Unmarshal(jsonData, &modelParams)
+			}
+		}
+	}
+
+	if traceID == "" {
+		// If traceID is not set, create a new trace
+		traceID = uuid.New().String()
+	}
+
+	name := fmt.Sprintf("bifrost_%s", string(req.RequestType))
+	if traceName != "" {
+		name = traceName
+	}
+
+	traceConfig := logging.TraceConfig{
+		Id:   traceID,
+		Name: maxim.StrPtr(name),
+	}
+
+	if sessionID != "" {
+		traceConfig.SessionId = &sessionID
+	}
+
+	// Create trace in the effective log repository
+	logger, err := plugin.getOrCreateLogger(effectiveLogRepoID)
+	if err != nil {
+		return req, nil, fmt.Errorf("failed to create trace: %w", err)
+	}
+
+	trace := logger.Trace(&traceConfig)
+	trace.SetInput(latestMessage)
+	generationID := uuid.New().String()
+
+	generationConfig := logging.GenerationConfig{
+		Id:              generationID,
+		Model:           model,
+		Provider:        string(provider),
+		Messages:        messages,
+		ModelParameters: modelParams,
+	}
+
+	if generationName != "" {
+		generationConfig.Name = &generationName
+	}
+
+	// Add generation to the effective log repository
+	logger.AddGenerationToTrace(traceID, &generationConfig)
+
+	// Extract and log attachments from message content
+	for _, att := range ExtractAttachmentsFromRequest(req) {
+		if att != nil {
+			logger.GenerationAddAttachment(generationID, att)
+		}
+	}
+
+	if ctx != nil {
+		if _, ok := ctx.Value(TraceIDKey).(string); !ok {
+			ctx.SetValue(TraceIDKey, traceID)
+		}
+		ctx.SetValue(GenerationIDKey, generationID)
+
+		// Extract request ID from context, if not present, create a new one
+		requestID, ok := ctx.Value(schemas.BifrostContextKeyRequestID).(string)
+		if !ok || requestID == "" {
+			// This should never happen since core/bifrost.go guarantees it's set before PreHooks
+			requestID = uuid.New().String()
+			plugin.logger.Warn("%s request ID missing in PreLLMHook, using fallback: %s", PluginLoggerPrefix, requestID)
+		}
+
+		// If streaming, create accumulator via central tracer using traceID
+		if bifrost.IsStreamRequestType(req.RequestType) {
+			tracer, bifrostTraceID, err := bifrost.GetTracerFromContext(ctx)
+			if err == nil && tracer != nil && bifrostTraceID != "" {
+				tracer.CreateStreamAccumulator(bifrostTraceID, time.Now())
+			}
+		}
+	}
+
+	return req, nil, nil
+}
+
+// PostLLMHook is called after a request has been processed by Bifrost.
+// It completes the request trace by:
+// - Adding response data to the generation if a generation ID exists
+// - Logging error details if bifrostErr is provided
+// - Ending the generation if it exists
+// - Ending the trace if a trace ID exists
+// - Flushing all pending log data
+//
+// The function gracefully handles cases where trace or generation IDs may be missing,
+// ensuring that partial logging is still performed when possible.
+//
+// Parameters:
+//   - ctx: Pointer to the schemas.BifrostContext containing trace/generation IDs
+//   - result: The Bifrost response to be traced
+//   - bifrostErr: The BifrostError returned by the request, if any
+//
+// Returns:
+//   - *schemas.BifrostResponse: The original response, unmodified
+//   - *schemas.BifrostError: The original error, unmodified
+//   - error: Never returns an error as it handles missing IDs gracefully
+func (plugin *Plugin) PostLLMHook(ctx *schemas.BifrostContext, result *schemas.BifrostResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) {
+	requestType, _, _, _ := bifrost.GetResponseFields(result, bifrostErr)
+	if requestType == schemas.RealtimeRequest {
+		return result, bifrostErr, nil
+	}
+
+	// Get effective log repo ID for this request
+	effectiveLogRepoID := plugin.getEffectiveLogRepoID(ctx)
+	if effectiveLogRepoID == "" {
+		return result, bifrostErr, nil
+	}
+	if ctx == nil {
+		return result, bifrostErr, nil
+	}
+
+	requestID, ok := ctx.Value(schemas.BifrostContextKeyRequestID).(string)
+	if !ok || requestID == "" {
+		return result, bifrostErr, nil
+	}
+
+	// Capture context values BEFORE goroutine to avoid race conditions
+	// when the same context is reused across multiple requests
+	generationID, hasGenerationID := ctx.Value(GenerationIDKey).(string)
+	traceID, hasTraceID := ctx.Value(TraceIDKey).(string)
+	tags, hasTags := ctx.Value(TagsKey).(map[string]string)
+
+	isFinalChunk := bifrost.IsFinalChunk(ctx)
+
+	go func() {
+		requestType, _, originalModel, resolvedModel := bifrost.GetResponseFields(result, bifrostErr)
+		modelTag := resolvedModel
+		if modelTag == "" {
+			modelTag = originalModel
+		}
+
+		var streamResponse *streaming.ProcessedStreamResponse
+		if bifrost.IsStreamRequestType(requestType) {
+			// Use central tracer's accumulator
+			tracer, bifrostTraceID, err := bifrost.GetTracerFromContext(ctx)
+			if err == nil && tracer != nil && bifrostTraceID != "" {
+				accResult := tracer.ProcessStreamingChunk(bifrostTraceID, isFinalChunk, result, bifrostErr)
+				if accResult != nil {
+					streamResponse = convertAccResultToProcessedStreamResponse(accResult)
+				}
+			}
+
+			// For streaming: only process on final chunk. Skip intermediate chunks.
+			// When there's an error, streamResponse may be nil but we must still log bifrostErr.
+			if !isFinalChunk {
+				return
+			}
+		}
+
+		logger, err := plugin.getOrCreateLogger(effectiveLogRepoID)
+		if err != nil {
+			return
+		}
+		if hasGenerationID {
+			if bifrostErr != nil {
+				// Safely extract message from nested error
+				message := ""
+				code := ""
+				errorType := ""
+				if bifrostErr.Error != nil {
+					message = bifrostErr.Error.Message
+					if bifrostErr.Error.Code != nil {
+						code = *bifrostErr.Error.Code
+					}
+					if bifrostErr.Error.Type != nil {
+						errorType = *bifrostErr.Error.Type
+					}
+				}
+				genErr := maximSchemas.GenerationError{
+					Message: message,
+					Code:    &code,
+					Type:    &errorType,
+				}
+				logger.SetGenerationError(generationID, &genErr)
+
+				if bifrost.IsStreamRequestType(requestType) {
+					// Cleanup via central tracer
+					tracer, bifrostTraceID, err := bifrost.GetTracerFromContext(ctx)
+					if err == nil && tracer != nil && bifrostTraceID != "" {
+						tracer.CleanupStreamAccumulator(bifrostTraceID)
+					}
+				}
+			} else if result != nil {
+				switch requestType {
+				case schemas.TextCompletionRequest, schemas.TextCompletionStreamRequest:
+					if streamResponse != nil {
+						logger.AddResultToGeneration(generationID, streamResponse.ToBifrostResponse().TextCompletionResponse)
+					} else {
+						logger.AddResultToGeneration(generationID, result.TextCompletionResponse)
+					}
+				case schemas.ChatCompletionRequest, schemas.ChatCompletionStreamRequest:
+					if streamResponse != nil {
+						logger.AddResultToGeneration(generationID, streamResponse.ToBifrostResponse().ChatResponse)
+					} else {
+						logger.AddResultToGeneration(generationID, result.ChatResponse)
+					}
+				case schemas.ResponsesRequest, schemas.ResponsesStreamRequest, schemas.WebSocketResponsesRequest:
+					if streamResponse != nil {
+						logger.AddResultToGeneration(generationID, streamResponse.ToBifrostResponse().ResponsesResponse)
+					} else {
+						logger.AddResultToGeneration(generationID, result.ResponsesResponse)
+					}
+				case schemas.ImageGenerationRequest, schemas.ImageGenerationStreamRequest,
+					schemas.ImageEditRequest, schemas.ImageEditStreamRequest:
+					if streamResponse != nil {
+						logger.AddResultToGeneration(generationID, streamResponse.ToBifrostResponse().ImageGenerationResponse)
+					} else if result != nil {
+						logger.AddResultToGeneration(generationID, result.ImageGenerationResponse)
+					}
+				}
+				if streamResponse != nil && isFinalChunk {
+					// Cleanup via central tracer
+					tracer, bifrostTraceID, err := bifrost.GetTracerFromContext(ctx)
+					if err == nil && tracer != nil && bifrostTraceID != "" {
+						tracer.CleanupStreamAccumulator(bifrostTraceID)
+					}
+				}
+			}
+		}
+		if hasTraceID {
+			logger.EndTrace(traceID)
+		}
+
+		// add tags to the generation and trace
+		if hasTags {
+			for key, value := range tags {
+				if generationID != "" {
+					logger.AddTagToGeneration(generationID, key, value)
+				}
+				if traceID != "" {
+					logger.AddTagToTrace(traceID, key, value)
+				}
+			}
+		}
+		if hasGenerationID && generationID != "" && modelTag != "" {
+			logger.AddTagToGeneration(generationID, "model", string(modelTag))
+		}
+		if hasTraceID && traceID != "" && modelTag != "" {
+			logger.AddTagToTrace(traceID, "model", string(modelTag))
+		}
+		// Flush only the effective logger that was used for this request
+		logger.Flush()
+	}()
+	return result, bifrostErr, nil
+}
+
+func (plugin *Plugin) Cleanup() error {
+	// Flush all loggers
+	plugin.loggerMutex.RLock()
+	for _, logger := range plugin.loggers {
+		logger.Flush()
+	}
+	plugin.loggerMutex.RUnlock()
+
+	return nil
+}