chore: add finishReason (#96)

henomis · Jul 12, 2023 · eb30a61 · eb30a61
1 parent ee3df70
commit eb30a61
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/llm/openai/openai.go b/llm/openai/openai.go
@@ -70,6 +70,7 @@ type openAI struct {
 	functions              map[string]Function
 	functionsMaxIterations uint
 	calledFunctionName     *string
+	finishReason           string
 }
 
 func New(model Model, temperature float32, maxTokens int, verbose bool) *openAI {
@@ -131,6 +132,10 @@ func (o *openAI) CalledFunctionName() *string {
 	return o.calledFunctionName
 }
 
+func (o *openAI) FinishReason() string {
+	return o.finishReason
+}
+
 func NewCompletion() *openAI {
 	return New(
 		GPT3TextDavinci003,
@@ -275,6 +280,7 @@ func (o *openAI) Chat(ctx context.Context, prompt *chat.Chat) (string, error) {
 
 	content := response.Choices[0].Message.Content
 
+	o.finishReason = string(response.Choices[0].FinishReason)
 	o.calledFunctionName = nil
 	if response.Choices[0].FinishReason == "function_call" && len(o.functions) > 0 {
 		content, err = o.functionCall(response)