fix: recover inferOpenAI function and adjust inferMoss2 code logic

Author: SoftEgLi

apis/record/infer.go

@@ -144,7 +144,7 @@ func InferMoss2(
 			if len(response.Choices) == 0 {
 				return unknownError
 			}
-
+			
 			resultBuilder.WriteString(response.Choices[0].Delta.Content)
 			nowOutput = resultBuilder.String()
 
@@ -156,16 +156,17 @@ func InferMoss2(
 
 			if slices.Contains(model.EndDelimiter, FuncCallEnd) && strings.Contains(nowOutput, FuncCallEnd) {
 				// if FuncCallEnd is found, call tool apis
-				var result string
-				err = GetFuncCallResult(nowOutput, &result)
+				var funcCallResult string
+				
+				err = GetFuncCallResult(nowOutput, &funcCallResult)
 				if err != nil {
 					return err
 				}
 
-				// TODO: Do we need this? or ignore <im_start>func_ret and simply send the json content
+				// TODO: Do we need this? or ignore <im_start>func_ret and simply send the funcCallResult
 				message := openai.ChatCompletionMessage{
 					Role: "func_ret",
-					Content: result,
+					Content: funcCallResult,
 				}
 
 				request = openai.ChatCompletionRequest{
@@ -182,9 +183,10 @@ func InferMoss2(
 				if err != nil {
 					return err
 				}
-				// erase the content of funcall
+				// erase the content of fun_call
 				nowOutput = strings.Split(nowOutput, FuncCallStart)[0]
-				continue
+				resultBuilder.Reset()
+				resultBuilder.WriteString(nowOutput)
 			}
 
 			before, _, found := CutLastAny(nowOutput, ",.?!\n，。？！")
@@ -351,6 +353,152 @@ func GetFuncCallResult(
 
 
 
+func InferOpenAI(
+	record *Record,
+	postRecord RecordModels,
+	model *ModelConfig,
+	user *User,
+	ctx *InferWsContext,
+) (
+	err error,
+) {
+	defer func() {
+		if v := recover(); v != nil {
+			Logger.Error("infer openai panicked", zap.Any("error", v))
+			err = unknownError
+		}
+	}()
+
+	openaiConfig := openai.DefaultConfig("")
+	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: messages,
+		Stop:     model.EndDelimiter,
+	}
+
+	if ctx == nil {
+		// openai client may panic when status code is 400
+		response, err := client.CreateChatCompletion(
+			context.Background(),
+			request,
+		)
+		if err != nil {
+			return err
+		}
+
+		if len(response.Choices) == 0 {
+			return unknownError
+		}
+
+		record.Response = response.Choices[0].Message.Content
+	} else {
+		// streaming
+		if config.Config.Debug {
+			Logger.Info("openai streaming",
+				zap.String("model", model.OpenAIModelName),
+				zap.String("url", model.Url),
+			)
+		}
+
+		stream, err := client.CreateChatCompletionStream(
+			context.Background(),
+			request,
+		)
+		if err != nil {
+			return err
+		}
+		defer stream.Close()
+
+		startTime := time.Now()
+
+		var resultBuilder strings.Builder
+		var nowOutput string
+		var detectedOutput string
+
+		for {
+			if ctx.connectionClosed.Load() {
+				return interruptError
+			}
+			response, err := stream.Recv()
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			if err != nil {
+				return err
+			}
+
+			if len(response.Choices) == 0 {
+				return unknownError
+			}
+
+			resultBuilder.WriteString(response.Choices[0].Delta.Content)
+			nowOutput = resultBuilder.String()
+
+			if slices.Contains(model.EndDelimiter, MossEnd) && strings.Contains(nowOutput, MossEnd) {
+				// if MossEnd is found, break the loop
+				nowOutput = strings.Split(nowOutput, MossEnd)[0]
+				break
+			}
+
+			before, _, found := CutLastAny(nowOutput, ",.?!\n，。？！")
+			if !found || before == detectedOutput {
+				continue
+			}
+			detectedOutput = before
+			if model.EnableSensitiveCheck {
+				err = sensitiveCheck(ctx.c, record, detectedOutput, startTime, user)
+				if err != nil {
+					return err
+				}
+			}
+
+			_ = ctx.c.WriteJSON(InferResponseModel{
+				Status: 1,
+				Output: detectedOutput,
+				Stage:  "MOSS",
+			})
+		}
+		if nowOutput != detectedOutput {
+			if model.EnableSensitiveCheck {
+				err = sensitiveCheck(ctx.c, record, nowOutput, startTime, user)
+				if err != nil {
+					return err
+				}
+			}
+
+			_ = ctx.c.WriteJSON(InferResponseModel{
+				Status: 1,
+				Output: nowOutput,
+				Stage:  "MOSS",
+			})
+		}
+
+		record.Response = nowOutput
+		record.Duration = float64(time.Since(startTime)) / 1000_000_000
+		_ = ctx.c.WriteJSON(InferResponseModel{
+			Status: 0,
+			Output: nowOutput,
+			Stage:  "MOSS",
+		})
+	}
+
+	return nil
+}
+
+
 func InferCommon(
 	record *Record,
 	prefix string,
@@ -381,7 +529,9 @@ func InferCommon(
 
 	// dispatch
 	if model.APIType == APITypeMOSS2 {
-		return InferMoss2(record, postRecords, model, user, ctx);
+		return InferMoss2(record, postRecords, model, user, ctx)
+	} else if model.APIType == APITypeMOSS{
+		return InferOpenAI(record, postRecords, model, user, ctx)
 	} else {
 		return errors.New("unknown API type")
 	}

models/config.go

@@ -13,6 +13,7 @@ type APIType string
 
 
 const (
+	APITypeMOSS APIType = "moss"
 	APITypeMOSS2 APIType = "moss2"
 )