#196: Introduced finish_reason on the error schema

EinStack · Apr 15, 2024 · 017b181 · 017b181
1 parent 81a9169
commit 017b181
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/pkg/api/schemas/chat_stream.go b/pkg/api/schemas/chat_stream.go
@@ -70,8 +70,9 @@ type ChatStreamChunk struct {
 }
 
 type ChatStreamError struct {
-	ErrCode ErrorCode `json:"errCode"`
-	Message string    `json:"message"`
+	ErrCode      ErrorCode     `json:"errCode"`
+	Message      string        `json:"message"`
+	FinishReason *FinishReason `json:"finishReason,omitempty"`
 }
 
 func NewChatStreamChunk(
@@ -95,15 +96,17 @@ func NewChatStreamError(
 	errCode ErrorCode,
 	errMsg string,
 	reqMetadata *Metadata,
+	finishReason *FinishReason,
 ) *ChatStreamMessage {
 	return &ChatStreamMessage{
 		ID:        reqID,
 		RouterID:  routerID,
 		CreatedAt: int(time.Now().UTC().Unix()),
 		Metadata:  reqMetadata,
 		Error: &ChatStreamError{
-			ErrCode: errCode,
-			Message: errMsg,
+			ErrCode:      errCode,
+			Message:      errMsg,
+			FinishReason: finishReason,
 		},
 	}
 }
diff --git a/pkg/routers/router.go b/pkg/routers/router.go
@@ -137,6 +137,7 @@ func (r *LangRouter) ChatStream(
 			schemas.NoModelConfigured,
 			ErrNoModels.Error(),
 			req.Metadata,
+			&schemas.ErrorReason,
 		)
 
 		return
@@ -188,6 +189,7 @@ func (r *LangRouter) ChatStream(
 						schemas.ModelUnavailable,
 						err.Error(),
 						req.Metadata,
+						nil,
 					)
 
 					continue NextModel
@@ -219,6 +221,7 @@ func (r *LangRouter) ChatStream(
 				schemas.UnknownError,
 				err.Error(),
 				req.Metadata,
+				nil,
 			)
 
 			return
@@ -237,5 +240,6 @@ func (r *LangRouter) ChatStream(
 		schemas.AllModelsUnavailable,
 		ErrNoModelAvailable.Error(),
 		req.Metadata,
+		&schemas.ErrorReason,
 	)
 }