-
Notifications
You must be signed in to change notification settings - Fork 129
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
refactor(llm-bridge): improve log and trace (#938)
1. Enhance llm bridge logging, add `service`, `namespace`, `traceId`, `requestId`, `duration` field. 2. The Bridge's http handler supports the injection of tracer. 3. Yomo exports the global trace client. 4. Add new `ResponseWriter` struct, the `ResponseWriter` supports record errors and TTFT timing.
- Loading branch information
Showing
11 changed files
with
288 additions
and
108 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,86 @@ | ||
| package ai | ||
|
|
||
| import ( | ||
| "encoding/json" | ||
| "io" | ||
| "net/http" | ||
| "time" | ||
| ) | ||
|
|
||
| // ResponseWriter is a wrapper for http.ResponseWriter. | ||
| // It is used to add TTFT and Err to the response. | ||
| type ResponseWriter struct { | ||
| IsStream bool | ||
| Err error | ||
| TTFT time.Time | ||
| underlying http.ResponseWriter | ||
| } | ||
|
|
||
| // NewResponseWriter returns a new ResponseWriter. | ||
| func NewResponseWriter(w http.ResponseWriter) *ResponseWriter { | ||
| return &ResponseWriter{ | ||
| underlying: w, | ||
| } | ||
| } | ||
|
|
||
| // Header returns the headers of the underlying ResponseWriter. | ||
| func (w *ResponseWriter) Header() http.Header { | ||
| return w.underlying.Header() | ||
| } | ||
|
|
||
| // Write writes the data to the underlying ResponseWriter. | ||
| func (w *ResponseWriter) Write(b []byte) (int, error) { | ||
| return w.underlying.Write(b) | ||
| } | ||
|
|
||
| // WriteHeader writes the header to the underlying ResponseWriter. | ||
| func (w *ResponseWriter) WriteHeader(code int) { | ||
| w.underlying.WriteHeader(code) | ||
| } | ||
|
|
||
| // WriteStreamEvent writes the event to the underlying ResponseWriter. | ||
| func (w *ResponseWriter) WriteStreamEvent(event any) error { | ||
| if _, err := io.WriteString(w, "data: "); err != nil { | ||
| return err | ||
| } | ||
| if err := json.NewEncoder(w).Encode(event); err != nil { | ||
| return err | ||
| } | ||
| if _, err := io.WriteString(w, "\n"); err != nil { | ||
| return err | ||
| } | ||
| flusher, ok := w.underlying.(http.Flusher) | ||
| if ok { | ||
| flusher.Flush() | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| // WriteStreamDone writes the done event to the underlying ResponseWriter. | ||
| func (w *ResponseWriter) WriteStreamDone() error { | ||
| _, err := io.WriteString(w, "data: [DONE]") | ||
|
|
||
| flusher, ok := w.underlying.(http.Flusher) | ||
| if ok { | ||
| flusher.Flush() | ||
| } | ||
|
|
||
| return err | ||
| } | ||
|
|
||
| // SetStreamHeader sets the stream headers of the underlying ResponseWriter. | ||
| func (w *ResponseWriter) SetStreamHeader() http.Header { | ||
| h := w.Header() | ||
| h.Set("Content-Type", "text/event-stream") | ||
| h.Set("Cache-Control", "no-cache, must-revalidate") | ||
| h.Set("x-content-type-options", "nosniff") | ||
| return h | ||
| } | ||
|
|
||
| // Flush flushes the underlying ResponseWriter. | ||
| func (w *ResponseWriter) Flush() { | ||
| flusher, ok := w.underlying.(http.Flusher) | ||
| if ok { | ||
| flusher.Flush() | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| package ai | ||
|
|
||
| import ( | ||
| "net/http/httptest" | ||
| "testing" | ||
|
|
||
| "github.com/sashabaranov/go-openai" | ||
| "github.com/stretchr/testify/assert" | ||
| ) | ||
|
|
||
| func TestResponseWriter(t *testing.T) { | ||
| recorder := httptest.NewRecorder() | ||
|
|
||
| w := NewResponseWriter(recorder) | ||
|
|
||
| h := w.SetStreamHeader() | ||
|
|
||
| err := w.WriteStreamEvent(openai.ChatCompletionResponse{ | ||
| ID: "chatcmpl-123", | ||
| }) | ||
| assert.NoError(t, err) | ||
|
|
||
| err = w.WriteStreamDone() | ||
| assert.NoError(t, err) | ||
|
|
||
| got := recorder.Body.String() | ||
|
|
||
| want := `data: {"id":"chatcmpl-123","object":"","created":0,"model":"","choices":null,"usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0,"prompt_tokens_details":null,"completion_tokens_details":null},"system_fingerprint":""} | ||
| data: [DONE]` | ||
|
|
||
| assert.Equal(t, want, got) | ||
| assert.Equal(t, recorder.Header(), h) | ||
| } |
Oops, something went wrong.