265: pushing wip. Need clarification on role response logic

EinStack · Jun 17, 2024 · 201ef76 · 201ef76
1 parent 516457b
commit 201ef76
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 58 deletions.
diff --git a/pkg/providers/cohere/chat.go b/pkg/providers/cohere/chat.go
@@ -127,7 +127,7 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 				"responseId":   cohereCompletion.ResponseID,
 			},
 			Message: schemas.ChatMessage{
-				Role:    payload.ChatHistory[len(payload.ChatHistory)-1].Role,
+				Role:    payload.Role,
 				Content: cohereCompletion.Text,
 			},
 			TokenUsage: schemas.TokenUsage{

diff --git a/pkg/providers/cohere/schemas.go b/pkg/providers/cohere/schemas.go
@@ -92,6 +92,7 @@ type FinalResponse struct {
 type ChatRequest struct {
 	Model             string                `json:"model"`
 	Message           string                `json:"message"`
+	Role              schemas.Role          `json:"role"`
 	ChatHistory       []schemas.ChatMessage `json:"chat_history"`
 	Temperature       float64               `json:"temperature,omitempty"`
 	Preamble          string                `json:"preamble,omitempty"`
@@ -131,6 +132,7 @@ func (r *ChatRequest) ApplyParams(params *schemas.ChatParams) {
 
 	message.Role = schemas.Role(mapRole(message.Role))
 
+	r.Role = message.Role
 	r.Message = message.Content
 	r.ChatHistory = messageHistory
 }

diff --git a/pkg/providers/cohere/schemas_test.go b/pkg/providers/cohere/schemas_test.go