bug: role assistant message maybe have function_call

tmc · Aug 16, 2023 · 7575c90 · 7575c90
1 parent eb0cbd3
commit 7575c90
Showing 1 changed file with 9 additions and 11 deletions.
diff --git a/llms/openai/openaillm_chat.go b/llms/openai/openaillm_chat.go
@@ -58,6 +58,10 @@ func (o *Chat) Generate(ctx context.Context, messageSets [][]schema.ChatMessage,
 				msg.Role = "system"
 			case schema.ChatMessageTypeAI:
 				msg.Role = "assistant"
+				if aiChatMsg, ok := m.(schema.AIChatMessage); ok && aiChatMsg.FunctionCall != nil {
+					msg.FunctionCall = &openaiclient.FunctionCall{Name: aiChatMsg.FunctionCall.Name,
+						Arguments: aiChatMsg.FunctionCall.Arguments}
+				}
 			case schema.ChatMessageTypeHuman:
 				msg.Role = "user"
 			case schema.ChatMessageTypeGeneric:
@@ -71,16 +75,11 @@ func (o *Chat) Generate(ctx context.Context, messageSets [][]schema.ChatMessage,
 			msgs[i] = msg
 		}
 		req := &openaiclient.ChatRequest{
-			Model:            opts.Model,
-			StopWords:        opts.StopWords,
-			Messages:         msgs,
-			StreamingFunc:    opts.StreamingFunc,
-			Temperature:      opts.Temperature,
-			MaxTokens:        opts.MaxTokens,
-			N:                opts.N,
-			FrequencyPenalty: opts.FrequencyPenalty,
-			PresencePenalty:  opts.PresencePenalty,
-
+			Model: opts.Model, StopWords: opts.StopWords,
+			Messages: msgs, StreamingFunc: opts.StreamingFunc,
+			Temperature: opts.Temperature, MaxTokens: opts.MaxTokens,
+			N: opts.N, FrequencyPenalty: opts.FrequencyPenalty,
+			PresencePenalty:      opts.PresencePenalty,
 			FunctionCallBehavior: openaiclient.FunctionCallBehavior(opts.FunctionCallBehavior),
 		}
 		for _, fn := range opts.Functions {
@@ -116,7 +115,6 @@ func (o *Chat) Generate(ctx context.Context, messageSets [][]schema.ChatMessage,
 			GenerationInfo: generationInfo,
 		})
 	}
-
 	return generations, nil
 }