Merge pull request #849 from trheyi/main

Add support for reasoning content and thinking state in OpenAI streaming

neo/assistant/api.go

@@ -283,6 +283,8 @@ func (ast *Assistant) streamChat(
 
 	errorRaw := ""
 	isFirst := true
+	isFirstThink := true
+	isThinking := false
 	currentMessageID := ""
 	err := ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
 		select {
@@ -290,7 +292,7 @@ func (ast *Assistant) streamChat(
 			return 0 // break
 
 		default:
-			msg := chatMessage.NewOpenAI(data)
+			msg := chatMessage.NewOpenAI(data, isThinking)
 			if msg == nil {
 				return 1 // continue
 			}
@@ -314,6 +316,29 @@ func (ast *Assistant) streamChat(
 				return 0 // break
 			}
 
+			// for api reasoning_content response
+			if msg.Type == "think" {
+				if isFirstThink {
+					msg.Text = "<think>\n" + msg.Text // add the think begin tag
+					isFirstThink = false
+					isThinking = true
+				}
+			}
+
+			// for api reasoning_content response
+			if isThinking && msg.Type != "think" {
+				// add the think close tag
+				end := chatMessage.New().Map(map[string]interface{}{"text": "\n</think>\n", "type": "think", "delta": true})
+				end.Write(c.Writer)
+				end.ID = currentMessageID
+				end.AppendTo(contents)
+				isThinking = false
+
+				// Clear the token and make a new line
+				contents.NewText([]byte{}, currentMessageID)
+				contents.ClearToken()
+			}
+
 			delta := msg.String()
 
 			// Chunk the delta

neo/message/contents.go

@@ -66,7 +66,7 @@ func (c *Contents) ScanTokens(currentID string, cb func(token string, id string,
 			c.UpdateType(c.token, map[string]interface{}{"text": text}, c.id)
 			c.NewText([]byte(tails), c.id) // Create new text with the tails
 			cb(c.token, c.id, false, text, tails)
-			c.token = "" // clear the token
+			c.ClearToken() // clear the token
 			return
 		}
 
@@ -88,6 +88,11 @@ func (c *Contents) ScanTokens(currentID string, cb func(token string, id string,
 	}
 }
 
+// ClearToken clear the token
+func (c *Contents) ClearToken() {
+	c.token = ""
+}
+
 // RemoveLastEmpty remove the last empty data
 func (c *Contents) RemoveLastEmpty() {
 	if c.Current == -1 {

neo/message/message.go

@@ -183,7 +183,7 @@ func NewAny(content interface{}) (*Message, error) {
 }
 
 // NewOpenAI create a new message from OpenAI response
-func NewOpenAI(data []byte) *Message {
+func NewOpenAI(data []byte, isThinking bool) *Message {
 	if data == nil || len(data) == 0 {
 		return nil
 	}
@@ -212,7 +212,7 @@ func NewOpenAI(data []byte) *Message {
 		}
 
 	case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"content":`):
-		var message openai.Message
+		var message openai.MessageWithReasoningContent
 		if err := jsoniter.Unmarshal(data, &message); err != nil {
 			color.Red("JSON parse error: %s", err.Error())
 			color.White(string(data))
@@ -224,7 +224,26 @@ func NewOpenAI(data []byte) *Message {
 
 		msg.Type = "text"
 		if len(message.Choices) > 0 {
-			msg.Text = message.Choices[0].Delta.Content
+			if reasoningContent, ok := message.Choices[0].Delta["reasoning_content"].(string); ok {
+				msg.Text = reasoningContent
+				msg.Type = "think"
+				return msg
+			}
+
+			if content, ok := message.Choices[0].Delta["content"].(string); ok && content != "" {
+				msg.Text = content
+				msg.Type = "text"
+				return msg
+			}
+
+			if isThinking {
+				msg.Type = "think"
+				msg.Text = ""
+				return msg
+			}
+
+			msg.Text = ""
+			return msg
 		}
 
 	case strings.Index(text, `{"code":`) == 0:

neo/neo.go

@@ -124,13 +124,16 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
 		}
 
 		errorRaw := ""
+		isFirstThink := true
+		isThinking := false
+		currentMessageID := ""
 		err := ast.Chat(c.Request.Context(), msgList, neo.Option, func(data []byte) int {
 			select {
 			case <-clientBreak:
 				return 0 // break
 
 			default:
-				msg := message.NewOpenAI(data)
+				msg := message.NewOpenAI(data, isThinking)
 				if msg == nil {
 					return 1 // continue
 				}
@@ -146,14 +149,60 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
 					return 0 // break
 				}
 
+				// for api reasoning_content response
+				if msg.Type == "think" {
+					if isFirstThink {
+						msg.Text = "<think>\n" + msg.Text // add the think begin tag
+						isFirstThink = false
+						isThinking = true
+					}
+				}
+
+				// for api reasoning_content response
+				if isThinking && msg.Type != "think" {
+					// add the think close tag
+					end := message.New().Map(map[string]interface{}{"text": "\n</think>\n", "type": "think", "delta": true})
+					end.Write(c.Writer)
+					end.ID = currentMessageID
+					end.AppendTo(contents)
+					isThinking = false
+
+					// Clear the token and make a new line
+					contents.NewText([]byte{}, currentMessageID)
+					contents.ClearToken()
+				}
+
 				// Append content and send message
 				msg.AppendTo(contents)
+
+				// Scan the tokens
+				contents.ScanTokens(currentMessageID, func(token string, id string, begin bool, text string, tails string) {
+					currentMessageID = id
+					msg.ID = id
+					msg.Type = token
+					msg.Text = ""                                    // clear the text
+					msg.Props = map[string]interface{}{"text": text} // Update props
+
+					// End of the token clear the text
+					if begin {
+						return
+					}
+
+					// New message with the tails
+					newMsg, err := message.NewString(tails, id)
+					if err != nil {
+						return
+					}
+					msgList = append(msgList, *newMsg)
+				})
+
 				if !silent {
 					value := msg.String()
 					if value != "" {
 						message.New().
 							Map(map[string]interface{}{
 								"text":  value,
+								"delta": true,
 								"done":  msg.IsDone,
 							}).
 							Write(c.Writer)

openai/openai.go

@@ -28,6 +28,7 @@ type OpenAI struct {
 	key          string
 	model        string
 	host         string
+	baseURL      string
 	organization string
 	maxToken     int
 }
@@ -71,8 +72,16 @@ func NewOpenAI(setting map[string]interface{}) (*OpenAI, error) {
 	}
 
 	host := "https://api.openai.com"
+	baseURL := "/v1"
 	if v, ok := setting["host"].(string); ok {
+		// Trim trailing slashes
+		v = strings.TrimRight(v, "/")
 		host = v
+		parts := strings.Split(v, "/")
+		if len(parts) > 3 {
+			host = strings.Join(parts[0:3], "/")
+			baseURL = "/" + strings.Join(parts[3:], "/")
+		}
 	}
 
 	organization := ""
@@ -89,6 +98,7 @@ func NewOpenAI(setting map[string]interface{}) (*OpenAI, error) {
 		key:          key,
 		model:        model,
 		host:         host,
+		baseURL:      baseURL,
 		organization: organization,
 		maxToken:     maxToken,
 	}, nil
@@ -142,11 +152,11 @@ func (openai OpenAI) Completions(prompt interface{}, option map[string]interface
 
 	if cb != nil {
 		option["stream"] = true
-		return nil, openai.stream(context.Background(), "/v1/completions", option, cb)
+		return nil, openai.stream(context.Background(), openai.baseURL+"/completions", option, cb)
 	}
 
 	option["stream"] = false
-	return openai.post("/v1/completions", option)
+	return openai.post(openai.baseURL+"/completions", option)
 }
 
 // CompletionsWith Creates a completion for the provided prompt and parameters.
@@ -159,11 +169,11 @@ func (openai OpenAI) CompletionsWith(ctx context.Context, prompt interface{}, op
 
 	if cb != nil {
 		option["stream"] = true
-		return nil, openai.stream(ctx, "/v1/completions", option, cb)
+		return nil, openai.stream(ctx, openai.baseURL+"/completions", option, cb)
 	}
 
 	option["stream"] = false
-	return openai.post("/v1/completions", option)
+	return openai.post(openai.baseURL+"/completions", option)
 }
 
 // ChatCompletions Creates a model response for the given chat conversation.
@@ -176,11 +186,11 @@ func (openai OpenAI) ChatCompletions(messages []map[string]interface{}, option m
 
 	if cb != nil {
 		option["stream"] = true
-		return nil, openai.stream(context.Background(), "/v1/chat/completions", option, cb)
+		return nil, openai.stream(context.Background(), openai.baseURL+"/chat/completions", option, cb)
 	}
 
 	option["stream"] = false
-	return openai.post("/v1/chat/completions", option)
+	return openai.post(openai.baseURL+"/chat/completions", option)
 }
 
 // ChatCompletionsWith Creates a model response for the given chat conversation.
@@ -193,11 +203,11 @@ func (openai OpenAI) ChatCompletionsWith(ctx context.Context, messages []map[str
 
 	if cb != nil {
 		option["stream"] = true
-		return nil, openai.stream(ctx, "/v1/chat/completions", option, cb)
+		return nil, openai.stream(ctx, openai.baseURL+"/chat/completions", option, cb)
 	}
 
 	option["stream"] = false
-	return openai.post("/v1/chat/completions", option)
+	return openai.post(openai.baseURL+"/chat/completions", option)
 }
 
 // Edits Creates a new edit for the provided input, instruction, and parameters.
@@ -207,7 +217,7 @@ func (openai OpenAI) Edits(instruction string, option map[string]interface{}) (i
 		option = map[string]interface{}{}
 	}
 	option["instruction"] = instruction
-	return openai.post("/v1/edits", option)
+	return openai.post(openai.baseURL+"/edits", option)
 }
 
 // Embeddings Creates an embedding vector representing the input text.
@@ -217,7 +227,7 @@ func (openai OpenAI) Embeddings(input interface{}, user string) (interface{}, *e
 	if user != "" {
 		payload["user"] = user
 	}
-	return openai.post("/v1/embeddings", payload)
+	return openai.post(openai.baseURL+"/embeddings", payload)
 }
 
 // AudioTranscriptions Transcribes audio into the input language.
@@ -231,7 +241,7 @@ func (openai OpenAI) AudioTranscriptions(dataBase64 string, option map[string]in
 	if option == nil {
 		option = map[string]interface{}{}
 	}
-	return openai.postFile("/v1/audio/transcriptions", map[string][]byte{"file": data}, option)
+	return openai.postFile(openai.baseURL+"/audio/transcriptions", map[string][]byte{"file": data}, option)
 }
 
 // ImagesGenerations Creates an image given a prompt.
@@ -246,7 +256,7 @@ func (openai OpenAI) ImagesGenerations(prompt string, option map[string]interfac
 	}
 
 	option["prompt"] = prompt
-	return openai.postWithoutModel("/v1/images/generations", option)
+	return openai.postWithoutModel(openai.baseURL+"/images/generations", option)
 }
 
 // ImagesEdits Creates an edited or extended image given an original image and a prompt.
@@ -277,7 +287,7 @@ func (openai OpenAI) ImagesEdits(imageBase64 string, prompt string, option map[s
 	}
 
 	option["prompt"] = prompt
-	return openai.postFileWithoutModel("/v1/images/edits", files, option)
+	return openai.postFileWithoutModel(openai.baseURL+"/images/edits", files, option)
 }
 
 // ImagesVariations Creates a variation of a given image.
@@ -298,7 +308,7 @@ func (openai OpenAI) ImagesVariations(imageBase64 string, option map[string]inte
 		option["response_format"] = "b64_json"
 	}
 
-	return openai.postFileWithoutModel("/v1/images/variations", files, option)
+	return openai.postFileWithoutModel(openai.baseURL+"/images/variations", files, option)
 }
 
 // Tiktoken get number of tokens

openai/types.go

@@ -16,6 +16,19 @@ type Message struct {
 	} `json:"choices,omitempty"`
 }
 
+// MessageWithReasoningContent is the response from OpenAI
+type MessageWithReasoningContent struct {
+	ID      string `json:"id,omitempty"`
+	Object  string `json:"object,omitempty"`
+	Created int64  `json:"created,omitempty"`
+	Model   string `json:"model,omitempty"`
+	Choices []struct {
+		Delta        map[string]interface{} `json:"delta,omitempty"`
+		Index        int                    `json:"index,omitempty"`
+		FinishReason string                 `json:"finish_reason,omitempty"`
+	} `json:"choices,omitempty"`
+}
+
 // ToolCalls is the response from OpenAI
 type ToolCalls struct {
 	ID      string `json:"id,omitempty"`