diff --git a/pkg/metrics/openai_recorder.go b/pkg/metrics/openai_recorder.go index f7ac8c3ce..a7669be22 100644 --- a/pkg/metrics/openai_recorder.go +++ b/pkg/metrics/openai_recorder.go @@ -18,6 +18,9 @@ import ( // per model. const maximumRecordsPerModel = 10 +// subscriberChannelBuffer is the buffer size for subscriber channels. +const subscriberChannelBuffer = 100 + type responseRecorder struct { http.ResponseWriter body *bytes.Buffer @@ -69,6 +72,10 @@ type OpenAIRecorder struct { records map[string]*ModelData // key is model ID modelManager *models.Manager // for resolving model tags to IDs m sync.RWMutex + + // streaming + subscribers map[string]chan []ModelRecordsResponse + subMutex sync.RWMutex } func NewOpenAIRecorder(log logging.Logger, modelManager *models.Manager) *OpenAIRecorder { @@ -76,6 +83,7 @@ func NewOpenAIRecorder(log logging.Logger, modelManager *models.Manager) *OpenAI log: log, modelManager: modelManager, records: make(map[string]*ModelData), + subscribers: make(map[string]chan []ModelRecordsResponse), } } @@ -188,6 +196,18 @@ func (r *OpenAIRecorder) RecordResponse(id, model string, rw http.ResponseWriter record.Response = response record.Error = "" // Ensure Error is empty for successful responses } + // Create ModelRecordsResponse with this single updated record to match + // what the non-streaming endpoint returns - []ModelRecordsResponse. + // See getAllRecords and getRecordsByModel. + modelResponse := []ModelRecordsResponse{{ + Count: 1, + Model: model, + ModelData: ModelData{ + Config: modelData.Config, + Records: []*RequestResponsePair{record}, + }, + }} + go r.broadcastToSubscribers(modelResponse) return } } @@ -274,36 +294,124 @@ func (r *OpenAIRecorder) convertStreamingResponse(streamingBody string) string { func (r *OpenAIRecorder) GetRecordsHandler() http.HandlerFunc { return func(w http.ResponseWriter, req *http.Request) { - w.Header().Set("Content-Type", "application/json") + acceptHeader := req.Header.Get("Accept") + + // Check if client wants Server-Sent Events + if acceptHeader == "text/event-stream" { + r.handleStreamingRequests(w, req) + return + } + + // Default to JSON response + r.handleJSONRequests(w, req) + } +} + +func (r *OpenAIRecorder) handleJSONRequests(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + + model := req.URL.Query().Get("model") + + if model == "" { + // Retrieve all records for all models. + allRecords := r.getAllRecords() + if allRecords == nil { + allRecords = []ModelRecordsResponse{} + } + if err := json.NewEncoder(w).Encode(allRecords); err != nil { + http.Error(w, fmt.Sprintf("Failed to encode all records: %v", err), + http.StatusInternalServerError) + return + } + } else { + // Retrieve records for the specified model. + records := r.getRecordsByModel(model) + if records == nil { + records = []ModelRecordsResponse{} + } + if err := json.NewEncoder(w).Encode(records); err != nil { + http.Error(w, fmt.Sprintf("Failed to encode records for model '%s': %v", model, err), + http.StatusInternalServerError) + return + } + } +} - model := req.URL.Query().Get("model") +func (r *OpenAIRecorder) handleStreamingRequests(w http.ResponseWriter, req *http.Request) { + // Set SSE headers. + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + + // Create subscriber channel. + subscriberID := fmt.Sprintf("sub_%d", time.Now().UnixNano()) + ch := make(chan []ModelRecordsResponse, subscriberChannelBuffer) + + // Register subscriber. + r.subMutex.Lock() + r.subscribers[subscriberID] = ch + r.subMutex.Unlock() + + // Clean up on disconnect. + defer func() { + r.subMutex.Lock() + delete(r.subscribers, subscriberID) + close(ch) + r.subMutex.Unlock() + }() + + // Optional: Send existing records first. + model := req.URL.Query().Get("model") + if includeExisting := req.URL.Query().Get("include_existing"); includeExisting == "true" { + r.sendExistingRecords(w, model) + } - if model == "" { - // Retrieve all records for all models. - allRecords := r.getAllRecords() - if allRecords == nil { - // No records found. - http.Error(w, "No records found", http.StatusNotFound) + flusher, ok := w.(http.Flusher) + if !ok { + http.Error(w, "Streaming not supported", http.StatusInternalServerError) + return + } + + // Send heartbeat to establish connection. + if _, err := fmt.Fprintf(w, "event: connected\ndata: {\"status\": \"connected\"}\n\n"); err != nil { + r.log.Errorf("Failed to write connected event to response: %v", err) + } + flusher.Flush() + + for { + select { + case modelRecords, ok := <-ch: + if !ok { return } - if err := json.NewEncoder(w).Encode(allRecords); err != nil { - http.Error(w, fmt.Sprintf("Failed to encode all records: %v", err), - http.StatusInternalServerError) - return + + // Filter by model if specified. + // modelRecords is assumed to have size 1 because that's how we call broadcastToSubscribers. + // We do this so we don't need to query a 2nd time for the model config. + if model != "" && len(modelRecords) > 0 && modelRecords[0].Model != model { + continue } - } else { - // Retrieve records for the specified model. - records := r.getRecordsByModel(model) - if records == nil { - // No records found for the specified model. - http.Error(w, fmt.Sprintf("No records found for model '%s'", model), http.StatusNotFound) - return + + // Send as SSE event. + jsonData, err := json.Marshal(modelRecords) + if err != nil { + r.log.Errorf("Failed to marshal record for streaming: %v", err) + errorMsg := fmt.Sprintf(`{"error": "Failed to marshal record: %v"}`, err) + if _, writeErr := fmt.Fprintf(w, "event: error\ndata: %s\n\n", errorMsg); writeErr != nil { + r.log.Errorf("Failed to write error event to response: %v", writeErr) + } + flusher.Flush() + continue } - if err := json.NewEncoder(w).Encode(records); err != nil { - http.Error(w, fmt.Sprintf("Failed to encode records for model '%s': %v", model, err), - http.StatusInternalServerError) - return + + if _, err := fmt.Fprintf(w, "event: new_request\ndata: %s\n\n", jsonData); err != nil { + r.log.Errorf("Failed to write new_request event to response: %v", err) } + flusher.Flush() + + case <-req.Context().Done(): + // Client disconnected. + return } } } @@ -352,6 +460,60 @@ func (r *OpenAIRecorder) getRecordsByModel(model string) []ModelRecordsResponse return nil } +func (r *OpenAIRecorder) broadcastToSubscribers(modelResponses []ModelRecordsResponse) { + r.subMutex.RLock() + defer r.subMutex.RUnlock() + + for _, ch := range r.subscribers { + select { + case ch <- modelResponses: + default: + // The channel is full, skip this subscriber. + } + } +} + +func (r *OpenAIRecorder) sendExistingRecords(w http.ResponseWriter, model string) { + var records []ModelRecordsResponse + + if model == "" { + records = r.getAllRecords() + } else { + records = r.getRecordsByModel(model) + } + + if records != nil { + // Send each individual request-response pair as a separate event. + for _, modelRecord := range records { + for _, requestRecord := range modelRecord.Records { + // Create a ModelRecordsResponse with a single record to match + // what the non-streaming endpoint returns - []ModelRecordsResponse. + // See getAllRecords and getRecordsByModel. + singleRecord := []ModelRecordsResponse{{ + Count: 1, + Model: modelRecord.Model, + ModelData: ModelData{ + Config: modelRecord.Config, + Records: []*RequestResponsePair{requestRecord}, + }, + }} + jsonData, err := json.Marshal(singleRecord) + if err != nil { + r.log.Errorf("Failed to marshal existing record for streaming: %v", err) + errorMsg := fmt.Sprintf(`{"error": "Failed to marshal existing record: %v"}`, err) + if _, writeErr := fmt.Fprintf(w, "event: error\ndata: %s\n\n", errorMsg); writeErr != nil { + r.log.Errorf("Failed to write error event to response: %v", writeErr) + } + } else { + if _, writeErr := fmt.Fprintf(w, "event: existing_request\ndata: %s\n\n", jsonData); writeErr != nil { + r.log.Errorf("Failed to write existing_request event to response: %v", writeErr) + } + } + } + } + } +} + func (r *OpenAIRecorder) RemoveModel(model string) { modelID := r.modelManager.ResolveModelID(model)