From 43f957dca4b57eb153be17e5f0d580ad1f401631 Mon Sep 17 00:00:00 2001 From: Zheng Kai Date: Wed, 6 Sep 2023 11:01:15 +0800 Subject: [PATCH] feat: cached token by key/ip --- server/src/core/metrics.go | 5 ++++- server/src/core/req.go | 3 +-- server/src/core/row.go | 1 + server/src/metrics/alias.go | 4 +++- server/src/metrics/init.go | 4 +++- server/src/metrics/rsp.go | 31 ++++++++++++++++++++++------- server/src/util/key.go | 23 ++++++++++++++++++++++ server/src/util/key_test.go | 39 +++++++++++++++++++++++++++++++++++++ 8 files changed, 98 insertions(+), 12 deletions(-) create mode 100644 server/src/util/key.go create mode 100644 server/src/util/key_test.go diff --git a/server/src/core/metrics.go b/server/src/core/metrics.go index c7477c1..df8f688 100644 --- a/server/src/core/metrics.go +++ b/server/src/core/metrics.go @@ -42,7 +42,10 @@ func doMetrics(ab []byte, cached bool, r *http.Request, req *pb.Req) { key := strings.TrimPrefix(r.Header.Get(`Authorization`), `Bearer `) metrics.RspToken(u.PromptTokens, u.TotalTokens, cached) - if !cached { + if cached { + metrics.RspTokenCachedByKey(key, u.TotalTokens) + metrics.RspTokenCachedByIP(sip, u.TotalTokens) + } else { metrics.RspTokenByModel(o.Model, u.TotalTokens) metrics.RspTokenByKey(key, u.TotalTokens) metrics.RspTokenByIP(sip, u.TotalTokens) diff --git a/server/src/core/req.go b/server/src/core/req.go index a41aeb3..b5e9801 100644 --- a/server/src/core/req.go +++ b/server/src/core/req.go @@ -5,13 +5,12 @@ import ( "net/http" "project/pb" "project/util" - "strings" ) func (c *Core) getAB(p *pb.Req, r *http.Request) (ab []byte, cached bool, pr *row, err error) { canCache := p.Method != http.MethodGet && p.Method != http.MethodDelete - if strings.Contains(r.Header.Get(`Authorization`), `no-cache`) { + if util.KeyNoCache(r.Header.Get(`Authorization`)) { canCache = false } diff --git a/server/src/core/row.go b/server/src/core/row.go index 2d0a58c..edf929c 100644 --- a/server/src/core/row.go +++ b/server/src/core/row.go @@ -57,6 +57,7 @@ func (pr *row) wait() { func (pr *row) saveFile() { rspFile := rspCacheFile(pr.req) + util.Mkdir(rspFile) util.WriteFile(rspFile, pr.rsp) } diff --git a/server/src/metrics/alias.go b/server/src/metrics/alias.go index e91f14a..c57abbe 100644 --- a/server/src/metrics/alias.go +++ b/server/src/metrics/alias.go @@ -1,6 +1,8 @@ package metrics -import "github.com/prometheus/client_golang/prometheus" +import ( + "github.com/prometheus/client_golang/prometheus" +) var baseName = `orca_` diff --git a/server/src/metrics/init.go b/server/src/metrics/init.go index 3594808..7696f56 100644 --- a/server/src/metrics/init.go +++ b/server/src/metrics/init.go @@ -15,9 +15,11 @@ func init() { prometheus.MustRegister(rspTokenCount) prometheus.MustRegister(rspTokenCachedCount) prometheus.MustRegister(rspJSONFailCount) - prometheus.MustRegister(rspTokenByIP) prometheus.MustRegister(rspTokenByModel) + prometheus.MustRegister(rspTokenByIP) + prometheus.MustRegister(rspTokenCachedByIP) prometheus.MustRegister(rspTokenByKey) + prometheus.MustRegister(rspTokenCachedByKey) prometheus.MustRegister(limitReq) prometheus.MustRegister(limitToken) diff --git a/server/src/metrics/rsp.go b/server/src/metrics/rsp.go index 9ace5e3..5b219a5 100644 --- a/server/src/metrics/rsp.go +++ b/server/src/metrics/rsp.go @@ -1,5 +1,7 @@ package metrics +import "project/util" + var ( rspBytes = newCounter(`rsp_bytes`, `rsp bytes`) rspPromptTokenCount = newCounter(`rsp_prompt_token`, `prompt token`) @@ -13,7 +15,12 @@ var ( ) rspTokenByKey = newCounterVec( `token_by_key`, - `openai key`, + `token by key`, + `key`, + ) + rspTokenCachedByKey = newCounterVec( + `token_cached_by_key`, + `cached token by key`, `key`, ) rspTokenByIP = newCounterVec( @@ -21,6 +28,11 @@ var ( `token by ip`, `ip`, ) + rspTokenCachedByIP = newCounterVec( + `token_cached_by_ip`, + `cached token by ip`, + `ip`, + ) ) // RspToken ... @@ -48,14 +60,19 @@ func RspTokenByIP(ip string, token uint32) { rspTokenByIP.WithLabelValues(ip).Add(float64(token)) } +// RspTokenCachedByIP ... +func RspTokenCachedByIP(ip string, token uint32) { + rspTokenCachedByIP.WithLabelValues(ip).Add(float64(token)) +} + // RspTokenByKey ... func RspTokenByKey(key string, token uint32) { - if len(key) > 30 { - key = key[:30] - } else if key == `` { - key = `` - } - rspTokenByKey.WithLabelValues(key).Add(float64(token)) + rspTokenByKey.WithLabelValues(util.FormatKey(key)).Add(float64(token)) +} + +// RspTokenCachedByKey ... +func RspTokenCachedByKey(key string, token uint32) { + rspTokenCachedByKey.WithLabelValues(util.FormatKey(key)).Add(float64(token)) } // RspTokenByModel ... diff --git a/server/src/util/key.go b/server/src/util/key.go new file mode 100644 index 0000000..b5e9f38 --- /dev/null +++ b/server/src/util/key.go @@ -0,0 +1,23 @@ +package util + +import ( + "regexp" + "strings" +) + +var pattrenNoCache = regexp.MustCompile(`(\s+|^)no-cache(\s+|$)`) + +// KeyNoCache ... +func KeyNoCache(key string) bool { + return pattrenNoCache.MatchString(key) +} + +// FormatKey ... +func FormatKey(key string) string { + key = pattrenNoCache.ReplaceAllString(key, ` `) + key = strings.TrimSpace(key) + if key == `` { + key = `` + } + return key +} diff --git a/server/src/util/key_test.go b/server/src/util/key_test.go new file mode 100644 index 0000000..fe9d288 --- /dev/null +++ b/server/src/util/key_test.go @@ -0,0 +1,39 @@ +package util + +import ( + "testing" +) + +func TestKey(t *testing.T) { + + li := []string{ + `foo no-cache`, + `no-cache foo`, + `foo no-cache`, + `no-cache foo`, + ` no-cache foo`, + ` no-cache foo `, + ` foo no-cache `, + } + + for _, s := range li { + if !KeyNoCache(s) { + t.Error(`KeyNoCache fail`, s) + } + if FormatKey(s) != `foo` { + t.Error(`FormatKey fail`, s) + } + } + + li = []string{ + `foo`, + `foono-cache`, + `no-cachefoo`, + ` nocache foo`, + } + for _, s := range li { + if KeyNoCache(s) { + t.Error(`KeyNoCache fail`, s) + } + } +}