Copilot commented on code in PR #706:
URL: https://github.com/apache/dubbo-go-pixiu/pull/706#discussion_r2228315574
##########
pkg/filter/llm/tokenizer/tokenizer.go:
##########
@@ -77,57 +79,106 @@ func (factory *FilterFactory) Apply() error {
return nil
}
-func (factory *FilterFactory) PrepareFilterChain(ctx *http.HttpContext, chain
filter.FilterChain) error {
+func (factory *FilterFactory) PrepareFilterChain(ctx *contexthttp.HttpContext,
chain filter.FilterChain) error {
f := &Filter{
cfg: factory.cfg,
}
chain.AppendEncodeFilters(f)
return nil
}
-func (f *Filter) Encode(hc *http.HttpContext) filter.FilterStatus {
+func (f *Filter) Encode(hc *contexthttp.HttpContext) filter.FilterStatus {
+ encoding := hc.Writer.Header().Get(constant.HeaderKeyContentEncoding)
+
switch res := hc.TargetResp.(type) {
case *client.StreamResponse:
pr, pw := io.Pipe()
res.Stream = newTeeReadCloser(res.Stream, pw)
- go f.processStreamResponse(pr)
+ go f.processStreamResponse(pr, encoding)
case *client.UnaryResponse:
- f.processUsageData(res.Data)
+ f.processUsageData(res.Data, encoding) // Unary response is not
a stream
default:
logger.Warnf(LoggerFmt+"Response type not suitable for token
calc: %T", res)
}
return filter.Continue
}
-func (f *Filter) processStreamResponse(stream io.Reader) {
- scanner := bufio.NewScanner(stream)
- currentLine := make([]byte, 0, 1024)
- // read the stream by line
- // and process the data lines
- // the data line is prefixed with "data:"
- // the data line is a json string
- // the for loop is to read the streamline by line and concat the
separate "data:" lines
- for scanner.Scan() {
- line := scanner.Text()
- line = strings.TrimSpace(line)
- if strings.HasPrefix(line, "data:") {
- f.processUsageData(currentLine)
- currentLine = make([]byte, 0, 1024)
- line = strings.TrimPrefix(line, "data:")
+// getDecompressedReader returns an io.ReadCloser that decompresses the body
based on the encoding.
+func getDecompressedReader(body io.Reader, encoding string) (io.ReadCloser,
error) {
+ switch encoding {
+ case constant.HeaderValueGzip:
+ return gzip.NewReader(body)
+ case constant.HeaderValueDeflate:
+ return flate.NewReader(body), nil
+ case "":
+ return io.NopCloser(body), nil
+ default:
+ return nil, fmt.Errorf("unsupported content encoding: %s",
encoding)
+ }
+}
+
+func (f *Filter) processStreamResponse(body io.Reader, encoding string) {
+ // For streams, we decompress the entire stream first, then process its
content.
+ // The content itself (with "data:" prefixes) is passed to
processUsageData.
+ decompressedReader, err := getDecompressedReader(body, encoding)
+ if err != nil {
+ logger.Errorf(LoggerFmt+"%v", err)
+ return
+ }
+ defer decompressedReader.Close()
+
+ decompressedData, err := io.ReadAll(decompressedReader)
+ if err != nil {
+ logger.Errorf(LoggerFmt+"Error reading decompressed stream:
%v", err)
+ return
+ }
+
+ decompressedDataTrim := strings.Trim(string(decompressedData), "data:")
Review Comment:
Using strings.Trim() with "data:" will remove individual characters 'd',
'a', 't', ':', from both ends of the string, not the prefix "data:". This could
incorrectly modify valid data. Use strings.TrimPrefix() instead.
```suggestion
decompressedDataTrim := strings.TrimPrefix(string(decompressedData),
"data:")
```
##########
pkg/filter/llm/tokenizer/tokenizer_test.go:
##########
@@ -32,55 +35,174 @@ import (
import (
"github.com/apache/dubbo-go-pixiu/pkg/client"
+ "github.com/apache/dubbo-go-pixiu/pkg/common/constant"
"github.com/apache/dubbo-go-pixiu/pkg/context/mock"
)
-func TestUnaryResponse(t *testing.T) {
- filter := &Filter{}
-
- request, err := http.NewRequest("POST",
"http://www.dubbogopixiu.com/mock/test?name=tc";,
bytes.NewReader([]byte("{\"id\":\"12345\"}")))
- assert.NoError(t, err)
- c := mock.GetMockHTTPContext(request)
- c.TargetResp = &client.UnaryResponse{
- Data: []byte(`{
- "usage": {
- "prompt_tokens": 7,
- "completion_tokens": 32,
- "total_tokens": 39,
- "prompt_tokens_details": {
- "cached_tokens": 0
+// TestUnaryResponseWithEncodings is a table-driven test for unary
(non-streaming) responses.
+// It covers multiple content encodings like gzip and deflate.
+func TestUnaryResponseWithEncodings(t *testing.T) {
+ // This is the payload we expect to process after decompression.
+ const payload = `{
+ "usage": {
+ "prompt_tokens": 7
+ }
+ }`
+
+ // Helper function to compress data with gzip for our test case.
+ compressGzipBytes := func(data string) []byte {
+ var buf bytes.Buffer
+ writer := gzip.NewWriter(&buf)
+ _, err := writer.Write([]byte(data))
+ assert.NoError(t, err)
+ err = writer.Close()
+ assert.NoError(t, err)
+ return buf.Bytes()
+ }
+
+ // Helper function to compress data with flate/deflate for our test
case.
+ compressFlateBytes := func(data string) []byte {
+ var buf bytes.Buffer
+ writer, err := flate.NewWriter(&buf, -1)
+ assert.NoError(t, err)
+ _, err = writer.Write([]byte(data))
+ assert.NoError(t, err)
+ err = writer.Close()
+ assert.NoError(t, err)
+ return buf.Bytes()
+ }
+
+ // Define all test cases in a table.
+ testCases := []struct {
+ name string
+ encoding string
+ getData func(string) []byte
+ }{
+ {
+ name: "No Encoding",
+ encoding: "",
+ getData: func(s string) []byte {
+ return []byte(s)
},
- "prompt_cache_hit_tokens": 0,
- "prompt_cache_miss_tokens": 7
- }
- }`)}
- filter.Encode(c)
+ },
+ {
+ name: "Gzip Encoding",
+ encoding: "gzip",
+ getData: compressGzipBytes,
+ },
+ {
+ name: "Flate Encoding",
+ encoding: "deflate",
+ getData: compressFlateBytes,
+ },
+ }
+
+ // Run the tests for each case.
+ for _, tc := range testCases {
+ t.Run(tc.name, func(t *testing.T) {
+ filter := &Filter{}
+
+ request, err := http.NewRequest("POST",
"http://www.dubbogopixiu.com/mock/test?name=tc";,
bytes.NewReader([]byte("{\"id\":\"12345\"}")))
+ assert.NoError(t, err)
+ c := mock.GetMockHTTPContext(request)
+
+ // Prepare the (potentially) compressed data
+ compressedData := tc.getData(payload)
+
+ c.TargetResp = &client.UnaryResponse{
+ Data: compressedData,
+ }
+ c.AddHeader(constant.HeaderKeyContentEncoding,
tc.encoding)
+
+ // Call the filter's Encode method
+ filter.Encode(c)
+ })
+ }
}
-func TestStreamResponse(t *testing.T) {
- filter := &Filter{}
-
- request, err := http.NewRequest("POST",
"http://www.dubbogopixiu.com/mock/test?name=tc";,
bytes.NewReader([]byte("{\"id\":\"12345\"}")))
- assert.NoError(t, err)
- c := mock.GetMockHTTPContext(request)
- s := io.NopCloser(strings.NewReader(`data: {
- "usage": {
- "prompt_tokens": 7,
- "completion_tokens": 32,
- "total_tokens": 39,
- "prompt_tokens_details": {
- "cached_tokens": 0
+// TestStreamResponseWithEncodings is a table-driven test for streaming
responses.
+// It replaces the old TestStreamResponse.
+func TestStreamResponseWithEncodings(t *testing.T) {
+ // This is the payload we expect to process after decompression.
+ const payload = `data: {
+ "usage": {
+ "prompt_tokens": 7
+ }
+ }`
+
+ // Helper function to compress data with gzip for our test case.
+ compressGzip := func(data string) io.Reader {
+ var buf bytes.Buffer
+ writer := gzip.NewWriter(&buf)
+ _, err := writer.Write([]byte(data))
+ assert.NoError(t, err)
+ err = writer.Close() // IMPORTANT: Close flushes the writer.
+ assert.NoError(t, err)
+ return &buf
+ }
+
+ compressFlate := func(data string) io.Reader {
+ var buf bytes.Buffer
+ writer, _ := flate.NewWriter(&buf, -1)
+ _, err := writer.Write([]byte(data))
+ assert.NoError(t, err)
Review Comment:
[nitpick] Error from flate.NewWriter() is being ignored with blank
identifier. While this specific call is unlikely to fail, it's better practice
to handle the error consistently with the rest of the codebase.
```suggestion
writer, err := flate.NewWriter(&buf, -1)
assert.NoError(t, err)
_, err = writer.Write([]byte(data))
assert.NoError(t, err)
```
##########
pkg/filter/llm/tokenizer/tokenizer.go:
##########
@@ -77,57 +79,106 @@ func (factory *FilterFactory) Apply() error {
return nil
}
-func (factory *FilterFactory) PrepareFilterChain(ctx *http.HttpContext, chain
filter.FilterChain) error {
+func (factory *FilterFactory) PrepareFilterChain(ctx *contexthttp.HttpContext,
chain filter.FilterChain) error {
f := &Filter{
cfg: factory.cfg,
}
chain.AppendEncodeFilters(f)
return nil
}
-func (f *Filter) Encode(hc *http.HttpContext) filter.FilterStatus {
+func (f *Filter) Encode(hc *contexthttp.HttpContext) filter.FilterStatus {
+ encoding := hc.Writer.Header().Get(constant.HeaderKeyContentEncoding)
+
switch res := hc.TargetResp.(type) {
case *client.StreamResponse:
pr, pw := io.Pipe()
res.Stream = newTeeReadCloser(res.Stream, pw)
- go f.processStreamResponse(pr)
+ go f.processStreamResponse(pr, encoding)
case *client.UnaryResponse:
- f.processUsageData(res.Data)
+ f.processUsageData(res.Data, encoding) // Unary response is not
a stream
default:
logger.Warnf(LoggerFmt+"Response type not suitable for token
calc: %T", res)
}
return filter.Continue
}
-func (f *Filter) processStreamResponse(stream io.Reader) {
- scanner := bufio.NewScanner(stream)
- currentLine := make([]byte, 0, 1024)
- // read the stream by line
- // and process the data lines
- // the data line is prefixed with "data:"
- // the data line is a json string
- // the for loop is to read the streamline by line and concat the
separate "data:" lines
- for scanner.Scan() {
- line := scanner.Text()
- line = strings.TrimSpace(line)
- if strings.HasPrefix(line, "data:") {
- f.processUsageData(currentLine)
- currentLine = make([]byte, 0, 1024)
- line = strings.TrimPrefix(line, "data:")
+// getDecompressedReader returns an io.ReadCloser that decompresses the body
based on the encoding.
+func getDecompressedReader(body io.Reader, encoding string) (io.ReadCloser,
error) {
+ switch encoding {
+ case constant.HeaderValueGzip:
+ return gzip.NewReader(body)
+ case constant.HeaderValueDeflate:
+ return flate.NewReader(body), nil
+ case "":
+ return io.NopCloser(body), nil
+ default:
+ return nil, fmt.Errorf("unsupported content encoding: %s",
encoding)
+ }
+}
+
+func (f *Filter) processStreamResponse(body io.Reader, encoding string) {
+ // For streams, we decompress the entire stream first, then process its
content.
+ // The content itself (with "data:" prefixes) is passed to
processUsageData.
+ decompressedReader, err := getDecompressedReader(body, encoding)
+ if err != nil {
+ logger.Errorf(LoggerFmt+"%v", err)
+ return
+ }
+ defer decompressedReader.Close()
+
+ decompressedData, err := io.ReadAll(decompressedReader)
+ if err != nil {
+ logger.Errorf(LoggerFmt+"Error reading decompressed stream:
%v", err)
+ return
+ }
+
+ decompressedDataTrim := strings.Trim(string(decompressedData), "data:")
+
+ // Now process the fully decompressed stream data
+ f.processUsageData([]byte(decompressedDataTrim), "")
+}
+
+func (f *Filter) processUsageData(data []byte, encoding string) {
+ var processedData []byte
+ // Decompress data if an encoding is specified (primarily for unary
responses)
+ if encoding != "" {
+ bodyReader := bytes.NewReader(data)
+ decompressedReader, err := getDecompressedReader(bodyReader,
encoding)
+ if err != nil {
+ logger.Errorf(LoggerFmt+"Failed to create decompressor:
%v", err)
+ return // Cannot proceed if decompression fails
+ }
+ defer decompressedReader.Close()
+
+ decompressedData, err := io.ReadAll(decompressedReader)
+ if err != nil {
+ logger.Errorf(LoggerFmt+"Failed to read decompressed
data: %v", err)
+ return // Cannot proceed if read fails
}
Review Comment:
The logic for handling encoding is duplicated between processStreamResponse
and processUsageData. Consider extracting the decompression logic into a
separate helper function to reduce code duplication.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
