fix: handle tool calls in retryStream

Author: Bajahaw

backend/cmd/chat/chat.go

@@ -273,7 +273,6 @@ func retryStream(w http.ResponseWriter, r *http.Request) {
 		ReasoningEffort: provider.ReasoningEffort(reasoningSetting),
 	}
 
-	// Save assistant message after streaming completes
 	responseMessage := Message{
 		ID:        -1,
 		ConvID:    req.ConversationID,
@@ -285,6 +284,9 @@ func retryStream(w http.ResponseWriter, r *http.Request) {
 		Children:  []int{},
 	}
 
+	var toolCalls []tools.ToolCall
+	var isToolsUsed bool
+
 	// Stream assistant content
 	completion, err := providerClient.SendChatCompletionStreamRequest(providerParams, w)
 	if err != nil {
@@ -295,6 +297,8 @@ func retryStream(w http.ResponseWriter, r *http.Request) {
 	} else {
 		responseMessage.Content = completion.Content
 		responseMessage.Reasoning = completion.Reasoning
+		toolCalls = completion.ToolCalls
+		isToolsUsed = len(toolCalls) > 0
 	}
 
 	responseID, saveErr := saveMessage(responseMessage)
@@ -306,6 +310,68 @@ func retryStream(w http.ResponseWriter, r *http.Request) {
 		parent.Children = append(parent.Children, responseID)
 	}
 
+	for len(toolCalls) > 0 {
+
+		toolCall := toolCalls[0]
+
+		providerParams.Messages = append(providerParams.Messages, provider.SimpleMessage{
+			Role:     "assistant",
+			ToolCall: toolCall,
+		})
+
+		toolCall.MessageID = responseMessage.ID
+		toolCall.ConvID = req.ConversationID
+
+		output := tools.ExecuteToolCall(toolCall)
+		toolCall.Output = output
+
+		chunk, _ := json.Marshal(provider.StreamChunk{
+			ToolCall: toolCall,
+		})
+		fmt.Fprintf(w, "data: %s\n\n", chunk)
+		flusher.Flush()
+
+		// Append tool result message to context for continued completion
+		providerParams.Messages = append(providerParams.Messages, provider.SimpleMessage{
+			Role: "tool",
+			ToolCall: tools.ToolCall{
+				ID:          toolCall.ID,
+				ReferenceID: toolCall.ReferenceID,
+				Name:        toolCall.Name,
+				Output:      output,
+			},
+		})
+
+		toolCalls = toolCalls[1:]
+		if len(toolCalls) == 0 {
+			completion, err = providerClient.SendChatCompletionStreamRequest(providerParams, w)
+			if err != nil {
+				log.Error("Error streaming chat completion after tool call", "err", err)
+				fmt.Fprintf(w, "event: error\ndata: {\"error\": \"%s\"}\n\n", err.Error())
+				flusher.Flush()
+				responseMessage.Error = err.Error()
+				break
+			}
+			toolCalls = append(toolCalls, completion.ToolCalls...)
+		}
+
+		// Accumulate reasoning for all tool calls
+		if responseMessage.Reasoning != "" {
+			responseMessage.Reasoning += " \n`using tool:" + toolCall.Name + "`\n " + completion.Reasoning
+		}
+	}
+
+	// Update assistant message with full content after all tool calls
+	if isToolsUsed {
+		if err == nil {
+			responseMessage.Content = completion.Content
+		}
+		_, err = updateMessage(responseMessage.ID, responseMessage)
+		if err != nil {
+			log.Error("Error updating assistant message after tool calls", "err", err)
+		}
+	}
+
 	// Send completion event with the new assistant message id
 	completionData := provider.StreamComplete{
 		UserMessageID:      parent.ID,
@@ -344,157 +410,3 @@ func update(W http.ResponseWriter, R *http.Request) {
 
 	utils.RespondWithJSON(W, &response, http.StatusOK)
 }
-
-// 	// Temporarily disabled
-// func chat(w http.ResponseWriter, r *http.Request) {
-// 	var req Request
-// 	err := utils.ExtractJSONBody(r, &req)
-// 	if err != nil || req.ConversationID == "" || req.Content == "" {
-// 		log.Error("Error unmarshalling request body", "err", err)
-// 		http.Error(w, "Invalid request body", http.StatusBadRequest)
-// 		return
-// 	}
-
-// 	// find or create conversation
-// 	convID := req.ConversationID
-// 	err = repo.touchConversation(req.ConversationID)
-// 	if err != nil {
-// 		conv := newConversation("admin")
-// 		if err = repo.saveConversation(conv); err != nil {
-// 			log.Error("Error creating conversation", "err", err)
-// 			http.Error(w, fmt.Sprintf("Error creating conversation: %v", err), http.StatusInternalServerError)
-// 			return
-// 		}
-// 		convID = conv.ID
-// 	}
-
-// 	userMessage := Message{
-// 		ID:         -1,
-// 		ConvID:     convID,
-// 		Role:       "user",
-// 		Content:    req.Content,
-// 		ParentID:   req.ParentID,
-// 		Children:   []int{},
-// 		Attachment: req.Attachment,
-// 	}
-
-// 	userMessage.ID, err = saveMessage(userMessage)
-// 	if err != nil {
-// 		log.Error("Error saving user message", "err", err)
-// 		http.Error(w, fmt.Sprintf("Error saving user message: %v", err), http.StatusInternalServerError)
-// 		return
-// 	}
-
-// 	ctx := buildContext(convID, userMessage.ID)
-// 	reasoningSetting, _ := getSetting("reasoningEffort")
-
-// 	providerParams := provider.ProviderRequestParams{
-// 		Messages:        ctx,
-// 		Model:           req.Model,
-// 		ReasoningEffort: provider.ReasoningEffort(reasoningSetting),
-// 	}
-
-// 	// send to provider
-// 	completion, err := providerClient.SendChatCompletionRequest(providerParams)
-// 	if err != nil {
-// 		log.Error("Error sending chat completion request", "err", err)
-// 		http.Error(w, fmt.Sprintf("Chat completion error: %v", err), http.StatusInternalServerError)
-// 		return
-// 	}
-
-// 	responseMessage := Message{
-// 		ID:        -1,
-// 		ConvID:    convID,
-// 		Role:      "assistant",
-// 		Model:     req.Model,
-// 		Content:   completion.Choices[0].Message.Content,
-// 		Reasoning: completion.Choices[0].Message.Reasoning,
-// 		ParentID:  userMessage.ID,
-// 		Children:  []int{},
-// 	}
-
-// 	responseMessage.ID, err = saveMessage(responseMessage)
-// 	if err != nil {
-// 		log.Error("Error saving response message", "err", err)
-// 		http.Error(w, fmt.Sprintf("Error saving response message: %v", err), http.StatusInternalServerError)
-// 		return
-// 	}
-
-// 	response := &Response{
-// 		Messages: make(map[int]*Message),
-// 	}
-// 	response.Messages[userMessage.ID] = &userMessage
-// 	response.Messages[responseMessage.ID] = &responseMessage
-
-// 	utils.RespondWithJSON(w, &response, http.StatusOK)
-// }
-
-// // Temporarily disabled
-// func retry(w http.ResponseWriter, r *http.Request) {
-// 	var req Retry
-// 	err := utils.ExtractJSONBody(r, &req)
-// 	if err != nil || req.ConversationID == "" {
-// 		log.Error("Error unmarshalling request body", "err", err)
-// 		http.Error(w, "Invalid request body", http.StatusBadRequest)
-// 		return
-// 	}
-
-// 	err = repo.touchConversation(req.ConversationID)
-// 	if err != nil {
-// 		log.Error("Error retrieving conversation", "err", err)
-// 		http.Error(w, fmt.Sprintf("Error retrieving conversation: %v", err), http.StatusNotFound)
-// 		return
-// 	}
-
-// 	parent, err := getMessage(req.ParentID)
-// 	if err != nil || parent.Role != "user" {
-// 		log.Error("Error retrieving parent message or invalid role", "err", err)
-// 		http.Error(w, "Invalid parent message", http.StatusBadRequest)
-// 		return
-// 	}
-
-// 	ctx := buildContext(req.ConversationID, parent.ID)
-// 	reasoningSetting, _ := getSetting("reasoningEffort")
-
-// 	providerParams := provider.ProviderRequestParams{
-// 		Messages:        ctx,
-// 		Model:           req.Model,
-// 		ReasoningEffort: provider.ReasoningEffort(reasoningSetting),
-// 	}
-
-// 	completion, err := providerClient.SendChatCompletionRequest(providerParams)
-// 	if err != nil {
-// 		log.Error("Error sending chat completion request", "err", err)
-// 		http.Error(w, fmt.Sprintf("Chat completion error: %v", err), http.StatusInternalServerError)
-// 		return
-// 	}
-
-// 	responseMessage := Message{
-// 		ID:        -1,
-// 		ConvID:    req.ConversationID,
-// 		Model:     req.Model,
-// 		Role:      "assistant",
-// 		Content:   completion.Choices[0].Message.Content,
-// 		Reasoning: completion.Choices[0].Message.Reasoning,
-// 		ParentID:  parent.ID,
-// 		Children:  []int{},
-// 	}
-
-// 	responseMessage.ID, err = saveMessage(responseMessage)
-// 	if err != nil {
-// 		log.Error("Error saving message", "err", err)
-// 		http.Error(w, fmt.Sprintf("Error saving message: %v", err), http.StatusInternalServerError)
-// 		return
-// 	}
-
-// 	parent.Children = append(parent.Children, responseMessage.ID)
-
-// 	response := &Response{
-// 		Messages: make(map[int]*Message),
-// 	}
-
-// 	response.Messages[parent.ID] = parent
-// 	response.Messages[responseMessage.ID] = &responseMessage
-
-// 	utils.RespondWithJSON(w, &response, http.StatusOK)
-// }