feat: cached token by key/ip

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)

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
 	}
 

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)
 }
 

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_`
 

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)

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 = `<empty>`
-	}
-	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 ...

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 = `<empty>`
+	}
+	return key
+}

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)
+		}
+	}
+}