diff --git a/misc/docker/Makefile b/misc/docker/Makefile index 4b6d4b0..89e7577 100644 --- a/misc/docker/Makefile +++ b/misc/docker/Makefile @@ -8,6 +8,9 @@ build: git run: build sudo docker run \ --env "OPENAI_API_KEY=$(OPENAI_API_KEY)" \ + --env "ORCA_ES_ADDR=$(ORCA_ES_ADDR)" \ + --env "ORCA_ES_USER=$(ORCA_ES_USER)" \ + --env "ORCA_ES_PASS=$(ORCA_ES_PASS)" \ --mount type=bind,source=/www/orca/static,target=/tmp \ --mount type=bind,source=/www/orca/server/dist/prod/log,target=/log \ -p 127.0.0.1:21035:80 \ diff --git a/proto/es.proto b/proto/es.proto index bc9a659..e33b17e 100644 --- a/proto/es.proto +++ b/proto/es.proto @@ -12,6 +12,7 @@ message EsMetrics { uint32 reqBytes = 7; uint32 rspBytes = 8; uint64 ts = 9; + string hash = 10; } message EsMetricsToken { diff --git a/server/src/config/config.go b/server/src/config/config.go index 93e119d..fd606f0 100644 --- a/server/src/config/config.go +++ b/server/src/config/config.go @@ -14,6 +14,7 @@ var ( OpenAIKey = `` + ESAddr = `` ESUser = `` ESPass = `` ) diff --git a/server/src/config/init.go b/server/src/config/init.go index cf86ec0..a9e20ac 100644 --- a/server/src/config/init.go +++ b/server/src/config/init.go @@ -17,6 +17,7 @@ func init() { `STATIC_DIR`: &StaticDir, `ORCA_WEB`: &WebAddr, `ORCA_LOG`: &LogDir, + `ORCA_ES_ADDR`: &ESAddr, `ORCA_ES_USER`: &ESUser, `ORCA_ES_PASS`: &ESPass, } diff --git a/server/src/es/init.go b/server/src/es/init.go index acd2ecf..015d823 100644 --- a/server/src/es/init.go +++ b/server/src/es/init.go @@ -12,10 +12,15 @@ var theClient *elasticsearch.Client // Init ... func Init() (err error) { - theClient, err = elasticsearch.NewClient(elasticsearch.Config{ + cfg := elasticsearch.Config{ Username: config.ESUser, Password: config.ESPass, - }) + } + if config.ESAddr != `` { + cfg.Addresses = []string{config.ESAddr} + } + + theClient, err = elasticsearch.NewClient(cfg) if err != nil { return err } diff --git a/server/src/es/insert.go b/server/src/es/insert.go index f63b4ae..8952ce7 100644 --- a/server/src/es/insert.go +++ b/server/src/es/insert.go @@ -2,8 +2,9 @@ package es import ( "bytes" - "encoding/json" "project/pb" + + "google.golang.org/protobuf/encoding/protojson" ) // Insert ... @@ -13,7 +14,10 @@ func Insert(d *pb.EsMetrics) { return } - ab, err := json.Marshal(d) + ab, err := protojson.MarshalOptions{ + EmitUnpopulated: true, + AllowPartial: true, + }.Marshal(d) if err != nil { return } diff --git a/server/src/es/test.go b/server/src/es/test.go index 7f2f320..93e76cb 100644 --- a/server/src/es/test.go +++ b/server/src/es/test.go @@ -1,6 +1,7 @@ package es import ( + "crypto/md5" "fmt" "math/rand" "project/pb" @@ -14,6 +15,8 @@ func Test() { for { + ts := zu.MS() + minute := time.Now().Minute() - 30 if minute < 0 { minute = -minute @@ -27,20 +30,37 @@ func Test() { } model := modelChoose[rand.Intn(len(modelChoose))] + ipChoose := []string{ + `10.0.32.43`, + `10.0.84.231`, + `10.0.84.49`, + `10.2.197.197`, + } + ip := ipChoose[rand.Intn(len(ipChoose))] + + keyChoose := []string{ + `zhengkai.orca`, + `zhengkai.debug`, + } + key := keyChoose[rand.Intn(len(keyChoose))] + + hash := md5.Sum([]byte(fmt.Sprintf(`%s-%s-%s-%d`, model, ip, key, ts))) + d := &pb.EsMetrics{ - ID: fmt.Sprintf(`chatcmpl-%d`, zu.MS()), + ID: fmt.Sprintf(`chatcmpl-%d`, ts), Token: &pb.EsMetricsToken{ Total: uint32(rand.Intn(minute * 10)), Completion: uint32(rand.Intn(minute * 10)), Prompt: uint32(rand.Intn(minute * 10)), }, - Cached: true, - Ip: `127.0.0.1`, + Cached: rand.Intn(2) == 1, + Ip: ip, Model: model, - Key: `zhengkai.orca`, + Key: key, ReqBytes: uint32(rand.Intn(minute * 57)), RspBytes: uint32(rand.Intn(minute * 21)), - Ts: zu.MS(), + Ts: ts, + Hash: fmt.Sprintf(`%x`, hash), } Insert(d) diff --git a/server/src/es/tpl/mapping.json b/server/src/es/tpl/mapping.json index e278209..aaf6b06 100644 --- a/server/src/es/tpl/mapping.json +++ b/server/src/es/tpl/mapping.json @@ -37,6 +37,9 @@ }, "ts": { "type": "date" + }, + "hash": { + "type": "keyword" } } }