feat: add system prompt support

JingYiJun · Feb 9, 2024 · e203139 · e203139
1 parent 372267c
commit e203139
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/apis/record/infer.go b/apis/record/infer.go
@@ -69,15 +69,19 @@ func InferOpenAI(
 	openaiConfig.BaseURL = model.Url
 	client := openai.NewClientWithConfig(openaiConfig)
 
+	var messages = make([]openai.ChatCompletionMessage, 0, len(postRecord)+2)
+	messages = append(messages, openai.ChatCompletionMessage{
+		Role:    "system",
+		Content: model.OpenAISystemPrompt,
+	})
+	messages = append(messages, postRecord.ToOpenAIMessages()...)
+	messages = append(messages, openai.ChatCompletionMessage{
+		Role:    "user",
+		Content: record.Request,
+	})
 	request := openai.ChatCompletionRequest{
-		Model: model.OpenAIModelName,
-		Messages: append(
-			postRecord.ToOpenAIMessages(),
-			openai.ChatCompletionMessage{
-				Role:    "user",
-				Content: record.Request,
-			},
-		),
+		Model:    model.OpenAIModelName,
+		Messages: messages,
 	}
 
 	if ctx == nil {
@@ -137,6 +141,7 @@ func InferOpenAI(
 
 			resultBuilder.WriteString(response.Choices[0].Delta.Content)
 			nowOutput = resultBuilder.String()
+
 			before, _, found := CutLastAny(nowOutput, ",.?!\n，。？！")
 			if !found || before == detectedOutput {
 				continue

diff --git a/models/config.go b/models/config.go
@@ -24,6 +24,7 @@ type ModelConfig struct {
 	CallbackUrl              string          `json:"callback_url"`
 	APIType                  APIType         `json:"api_type"`
 	OpenAIModelName          string          `json:"openai_model_name"`
+	OpenAISystemPrompt       string          `json:"openai_system_prompt"`
 }
 
 type ModelConfigs = []*ModelConfig