Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 42 additions & 17 deletions core/controller/relay-controller.go
Original file line number Diff line number Diff line change
Expand Up @@ -401,23 +401,14 @@ func recordResult(
firstByteAt = result.BodyDetail.FirstByteAt
}

if config.GetSaveAllLogDetail() || meta.ModelConfig.ForceSaveDetail || code != http.StatusOK {
if result.BodyDetail != nil {
requestBodyMaxSize := effectiveDetailBodyMaxSize(
meta.ModelConfig.RequestBodyStorageMaxSize,
config.GetLogDetailRequestBodyMaxSize(),
)
responseBodyMaxSize := effectiveDetailBodyMaxSize(
meta.ModelConfig.ResponseBodyStorageMaxSize,
config.GetLogDetailResponseBodyMaxSize(),
)

detail = &model.RequestDetail{
RequestBody: result.BodyDetail.RequestBody,
ResponseBody: result.BodyDetail.ResponseBody,
}
detail.ApplyBodySizeLimits(requestBodyMaxSize, responseBodyMaxSize)
}
forceSaveDetail := config.GetSaveAllLogDetail() || meta.ModelConfig.ForceSaveDetail
if forceSaveDetail || code != http.StatusOK {
detail = buildRequestDetailForLog(
result.BodyDetail,
meta.ModelConfig,
code,
forceSaveDetail,
)
}

gbc := middleware.GetGroupBalanceConsumerFromContext(c)
Expand Down Expand Up @@ -502,6 +493,40 @@ func effectiveDetailBodyMaxSize(modelLimit, globalLimit int64) int64 {
return globalLimit
}

func buildRequestDetailForLog(
bodyDetail *controller.BodyDetail,
modelConfig model.ModelConfig,
code int,
forceSaveDetail bool,
) *model.RequestDetail {
if bodyDetail == nil {
return nil
}

requestBodyMaxSize := effectiveDetailBodyMaxSize(
modelConfig.RequestBodyStorageMaxSize,
config.GetLogDetailRequestBodyMaxSize(),
)
responseBodyMaxSize := effectiveDetailBodyMaxSize(
modelConfig.ResponseBodyStorageMaxSize,
config.GetLogDetailResponseBodyMaxSize(),
)

detail := &model.RequestDetail{
RequestBody: bodyDetail.RequestBody,
ResponseBody: bodyDetail.ResponseBody,
}
detail.DropInvalidUTF8Bodies()

if controller.ShouldSkipRequestBodyDetailForStatus(code) && !forceSaveDetail {
detail.RequestBody = ""
}

detail.ApplyBodySizeLimits(requestBodyMaxSize, responseBodyMaxSize)

return detail
}

func buildBodyDetailOption(meta *meta.Meta) controller.BodyDetailOption {
requestBodyMaxSize := effectiveDetailBodyMaxSize(
meta.ModelConfig.RequestBodyStorageMaxSize,
Expand Down
83 changes: 83 additions & 0 deletions core/controller/relay-controller_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
package controller

import (
"net/http"
"reflect"
"testing"
"time"
Expand Down Expand Up @@ -209,3 +210,85 @@ func TestSaveAsyncUsageInfoDoesNotStoreInitialUsage(t *testing.T) {
require.Zero(t, captured.Usage.TotalTokens)
require.Equal(t, "priority", captured.UsageContext.ServiceTier)
}

func TestBuildRequestDetailForLogSkipsRequestBodyForUpstreamOnlyStatuses(t *testing.T) {
t.Parallel()

bodyDetail := &relaycontroller.BodyDetail{
RequestBody: `{"prompt":"secret"}`,
ResponseBody: `{"error":"upstream"}`,
}

for _, statusCode := range []int{
http.StatusUnauthorized,
http.StatusPaymentRequired,
http.StatusForbidden,
http.StatusNotFound,
http.StatusMethodNotAllowed,
http.StatusTooManyRequests,
} {
t.Run(http.StatusText(statusCode), func(t *testing.T) {
t.Parallel()

detail := buildRequestDetailForLog(bodyDetail, model.ModelConfig{}, statusCode, false)

require.NotNil(t, detail)
assert.Empty(t, detail.RequestBody)
assert.Equal(t, `{"error":"upstream"}`, detail.ResponseBody)
})
}
}

func TestBuildRequestDetailForLogKeepsRequestBodyWhenForced(t *testing.T) {
t.Parallel()

detail := buildRequestDetailForLog(
&relaycontroller.BodyDetail{
RequestBody: `{"prompt":"secret"}`,
ResponseBody: `{"error":"limited"}`,
},
model.ModelConfig{},
http.StatusTooManyRequests,
true,
)

require.NotNil(t, detail)
assert.Equal(t, `{"prompt":"secret"}`, detail.RequestBody)
assert.Equal(t, `{"error":"limited"}`, detail.ResponseBody)
}

func TestBuildRequestDetailForLogKeepsRequestBodyForClientPayloadErrors(t *testing.T) {
t.Parallel()

detail := buildRequestDetailForLog(
&relaycontroller.BodyDetail{
RequestBody: `{"prompt":"secret"}`,
ResponseBody: `{"error":"bad request"}`,
},
model.ModelConfig{},
http.StatusBadRequest,
false,
)

require.NotNil(t, detail)
assert.Equal(t, `{"prompt":"secret"}`, detail.RequestBody)
assert.Equal(t, `{"error":"bad request"}`, detail.ResponseBody)
}

func TestBuildRequestDetailForLogDropsInvalidUTF8Bodies(t *testing.T) {
t.Parallel()

detail := buildRequestDetailForLog(
&relaycontroller.BodyDetail{
RequestBody: string([]byte{0xff, 0xfe}),
ResponseBody: string([]byte{'o', 'k', 0xff}),
},
model.ModelConfig{},
http.StatusBadRequest,
false,
)

require.NotNil(t, detail)
assert.Empty(t, detail.RequestBody)
assert.Empty(t, detail.ResponseBody)
}
13 changes: 13 additions & 0 deletions core/model/log.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"fmt"
"strings"
"time"
"unicode/utf8"

"github.com/bytedance/sonic"
"github.com/labring/aiproxy/core/common"
Expand Down Expand Up @@ -45,6 +46,18 @@ func (d *RequestDetail) ApplyBodySizeLimits(requestMaxSize, responseMaxSize int6
d.ResponseBody, d.ResponseBodyTruncated = truncateDetailBody(d.ResponseBody, responseMaxSize)
}

func (d *RequestDetail) DropInvalidUTF8Bodies() {
if d.RequestBody != "" && !utf8.ValidString(d.RequestBody) {
d.RequestBody = ""
d.RequestBodyTruncated = false
}

if d.ResponseBody != "" && !utf8.ValidString(d.ResponseBody) {
d.ResponseBody = ""
d.ResponseBodyTruncated = false
}
}

type Log struct {
RequestDetail *RequestDetail `gorm:"foreignKey:LogID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE;" json:"request_detail,omitempty"`
RequestAt time.Time ` json:"request_at"`
Expand Down
27 changes: 27 additions & 0 deletions core/model/log_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,33 @@ func TestRequestDetailApplyBodySizeLimitsZeroKeepsOriginalBody(t *testing.T) {
}
}

func TestRequestDetailDropInvalidUTF8Bodies(t *testing.T) {
detail := &model.RequestDetail{
RequestBody: string([]byte{0xff, 0xfe}),
ResponseBody: "valid",
RequestBodyTruncated: true,
ResponseBodyTruncated: true,
}

detail.DropInvalidUTF8Bodies()

if detail.RequestBody != "" {
t.Fatalf("expected invalid request body to be cleared, got %q", detail.RequestBody)
}

if detail.RequestBodyTruncated {
t.Fatal("expected request body truncated flag to be cleared")
}

if detail.ResponseBody != "valid" {
t.Fatalf("expected valid response body to remain unchanged, got %q", detail.ResponseBody)
}

if !detail.ResponseBodyTruncated {
t.Fatal("expected valid response body truncated flag to remain unchanged")
}
}

func TestRecordConsumeLogPersistsWebSearchCount(t *testing.T) {
db, err := model.OpenSQLite(filepath.Join(t.TempDir(), "logs.db"))
if err != nil {
Expand Down
Loading
Loading