chore: update context max length validating

yiwen-ai · Nov 26, 2023 · deb3d2a · deb3d2a
1 parent 1597f32
commit deb3d2a
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 8 deletions.
diff --git a/src/bll/writing_collection.go b/src/bll/writing_collection.go
@@ -32,6 +32,12 @@ func (i *CreateCollectionInput) Validate() error {
 		return gear.ErrBadRequest.From(err)
 	}
 
+	if i.Context != "" {
+		if tk := util.Tiktokens(i.Context); tk > 2048 {
+			return gear.ErrBadRequest.WithMsgf("context is too long, max tokens is 2048, got %d", tk)
+		}
+	}
+
 	return nil
 }
 
@@ -139,6 +145,12 @@ func (i *UpdateCollectionInput) Validate() error {
 		return gear.ErrBadRequest.From(err)
 	}
 
+	if i.Context != nil {
+		if tk := util.Tiktokens(*i.Context); tk > 2048 {
+			return gear.ErrBadRequest.WithMsgf("context is too long, max tokens is 2048, got %d", tk)
+		}
+	}
+
 	return nil
 }
 

diff --git a/src/bll/writing_message.go b/src/bll/writing_message.go
@@ -183,6 +183,12 @@ func (i *CreateMessageInput) Validate() error {
 		return gear.ErrBadRequest.WithMsg("message is too large")
 	}
 
+	if i.Context != "" {
+		if tk := util.Tiktokens(i.Context); tk > 2048 {
+			return gear.ErrBadRequest.WithMsgf("context is too long, max tokens is 2048, got %d", tk)
+		}
+	}
+
 	return nil
 }
 
@@ -229,6 +235,11 @@ func (i *UpdateMessageInput) Validate() error {
 	if i.Message != nil && i.Language == nil {
 		return gear.ErrBadRequest.WithMsg("language is required with message")
 	}
+	if i.Context != nil {
+		if tk := util.Tiktokens(*i.Context); tk > 2048 {
+			return gear.ErrBadRequest.WithMsgf("context is too long, max tokens is 2048, got %d", tk)
+		}
+	}
 	if i.NewlyAdd == nil {
 		i.NewlyAdd = util.Ptr(true)
 	}

diff --git a/src/bll/writing_publication.go b/src/bll/writing_publication.go
@@ -28,6 +28,11 @@ func (i *CreatePublicationInput) Validate() error {
 	if err := util.Validator.Struct(i); err != nil {
 		return gear.ErrBadRequest.From(err)
 	}
+	if i.Context != nil {
+		if tk := util.Tiktokens(*i.Context); tk > 2048 {
+			return gear.ErrBadRequest.WithMsgf("context is too long, max tokens is 2048, got %d", tk)
+		}
+	}
 
 	return nil
 }

diff --git a/src/util/common_test.go b/src/util/common_test.go
@@ -18,7 +18,7 @@ func TestRemoveDuplicates(t *testing.T) {
 	assert.Equal(t, RemoveDuplicates([]string{"a", "b", "a"}), []string{"a", "b"})
 
 	id := NewID()
-	id2 := mustParseID(id.String())
+	id2 := MustParseID(id.String())
 	assert.Equal(t, RemoveDuplicates([]ID{id, id2}), []ID{id})
 }
 

diff --git a/src/util/http.go b/src/util/http.go
@@ -188,8 +188,8 @@ func RequestCBOR(ctx context.Context, cli *http.Client, method, api string, inpu
 		if e != nil {
 			str = string(data)
 		}
-		er := gear.Err.WithCode(resp.StatusCode).WithMsgf("RequestCBOR failed, url: %q, rid: %s, code: %d, error: %v",
-			api, rid, resp.StatusCode, err)
+		er := gear.Err.WithCode(resp.StatusCode).WithMsgf("RequestCBOR failed, url: %q, rid: %s, code: %d, error: %v, body: %s",
+			api, rid, resp.StatusCode, err, str)
 		er.Data = str
 		return er
 	}
@@ -199,8 +199,8 @@ func RequestCBOR(ctx context.Context, cli *http.Client, method, api string, inpu
 		if e != nil {
 			str = string(data)
 		}
-		er := gear.ErrInternalServerError.WithMsgf("RequestCBOR failed, url: %q, rid: %s, code: %d, error: %v",
-			api, rid, resp.StatusCode, err)
+		er := gear.ErrInternalServerError.WithMsgf("RequestCBOR failed, url: %q, rid: %s, code: %d, error: %v, body: %s",
+			api, rid, resp.StatusCode, err, str)
 		er.Data = str
 		return er
 	}

diff --git a/src/util/id.go b/src/util/id.go
@@ -12,8 +12,8 @@ import (
 )
 
 var ZeroID ID
-var JARVIS ID = mustParseID("0000000000000jarvis0") // system user
-var ANON ID = mustParseID("000000000000000anon0")   // anonymous user
+var JARVIS ID = MustParseID("0000000000000jarvis0") // system user
+var ANON ID = MustParseID("000000000000000anon0")   // anonymous user
 var MinID ID = ID(xid.ID([12]byte{0, 0, 0, 0, 255, 255, 255, 255, 255, 255, 255, 255}))
 
 func NewID() ID {
@@ -28,7 +28,7 @@ func ParseID(s string) (ID, error) {
 	return ID(id), nil
 }
 
-func mustParseID(s string) ID {
+func MustParseID(s string) ID {
 	id, err := xid.FromString(s)
 	if err != nil {
 		panic(err)