Skip to content

Commit b2862e3

Browse files
sawkasawka
committed
move toolusedata to the usechat pkg and out of the backends
1 parent f11cc8d commit b2862e3

File tree

6 files changed

+42
-47
lines changed

6 files changed

+42
-47
lines changed

‎frontend/app/aipanel/aipanel.tsx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,8 @@ const AIPanelComponentInner = memo(() => {
246246
model.registerUseChatData(sendMessage, setMessages, status, stop);
247247

248248
// console.log("AICHAT messages", messages);
249+
(window as any).aichatmessages = messages;
250+
(window as any).aichatstatus = status;
249251

250252
const handleKeyDown = (waveEvent: WaveKeyboardEvent): boolean => {
251253
if (checkKeyPressed(waveEvent, "Cmd:k")) {

‎pkg/aiusechat/aiutil/aiutil.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -195,8 +195,8 @@ func IsOpenAIReasoningModel(model string) bool {
195195
}
196196

197197
// CreateToolUseData creates a UIMessageDataToolUse from tool call information
198-
func CreateToolUseData(toolCallID, toolName string, arguments string, chatOpts uctypes.WaveChatOpts) *uctypes.UIMessageDataToolUse {
199-
toolUseData := &uctypes.UIMessageDataToolUse{
198+
func CreateToolUseData(toolCallID, toolName string, arguments string, chatOpts uctypes.WaveChatOpts) uctypes.UIMessageDataToolUse {
199+
toolUseData := uctypes.UIMessageDataToolUse{
200200
ToolCallId: toolCallID,
201201
ToolName: toolName,
202202
Status: uctypes.ToolUseStatusPending,

‎pkg/aiusechat/openai/openai-backend.go‎

Lines changed: 9 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import (
99
"errors"
1010
"fmt"
1111
"io"
12-
"log"
1312
"net/http"
1413
"net/url"
1514
"strings"
@@ -396,8 +395,7 @@ type openaiBlockState struct {
396395
}
397396

398397
type openaiStreamingState struct {
399-
blockMap map[string]*openaiBlockState // Use item_id as key for UI streaming
400-
toolUseData map[string]*uctypes.UIMessageDataToolUse // Use toolCallId as key
398+
blockMap map[string]*openaiBlockState // Use item_id as key for UI streaming
401399
msgID string
402400
model string
403401
stepStarted bool
@@ -407,7 +405,7 @@ type openaiStreamingState struct {
407405

408406
// ---------- Public entrypoint ----------
409407

410-
func UpdateToolUseData(chatId string, callId string, newToolUseData *uctypes.UIMessageDataToolUse) error {
408+
func UpdateToolUseData(chatId string, callId string, newToolUseData uctypes.UIMessageDataToolUse) error {
411409
chat := chatstore.DefaultChatStore.Get(chatId)
412410
if chat == nil {
413411
return fmt.Errorf("chat not found: %s", chatId)
@@ -422,7 +420,7 @@ func UpdateToolUseData(chatId string, callId string, newToolUseData *uctypes.UIM
422420
if chatMsg.FunctionCall != nil && chatMsg.FunctionCall.CallId == callId {
423421
updatedMsg := *chatMsg
424422
updatedFunctionCall := *chatMsg.FunctionCall
425-
updatedFunctionCall.ToolUseData = newToolUseData
423+
updatedFunctionCall.ToolUseData = &newToolUseData
426424
updatedMsg.FunctionCall = &updatedFunctionCall
427425

428426
aiOpts := &uctypes.AIOptsType{
@@ -592,9 +590,8 @@ func parseOpenAIHTTPError(resp *http.Response) error {
592590
func handleOpenAIStreamingResp(ctx context.Context, sse *sse.SSEHandlerCh, decoder *eventsource.Decoder, cont *uctypes.WaveContinueResponse, chatOpts uctypes.WaveChatOpts) (*uctypes.WaveStopReason, []*OpenAIChatMessage) {
593591
// Per-response state
594592
state := &openaiStreamingState{
595-
blockMap: map[string]*openaiBlockState{},
596-
toolUseData: map[string]*uctypes.UIMessageDataToolUse{},
597-
chatOpts: chatOpts,
593+
blockMap: map[string]*openaiBlockState{},
594+
chatOpts: chatOpts,
598595
}
599596

600597
var rtnStopReason *uctypes.WaveStopReason
@@ -875,8 +872,6 @@ func handleOpenAIEvent(
875872

876873
// Get the function call info from the block state
877874
if st := state.blockMap[ev.ItemId]; st != nil && st.kind == openaiBlockToolUse {
878-
toolUseData := aiutil.CreateToolUseData(st.toolCallID, st.toolName, ev.Arguments, state.chatOpts)
879-
state.toolUseData[st.toolCallID] = toolUseData
880875
aiutil.SendToolProgress(st.toolCallID, st.toolName, []byte(ev.Arguments), state.chatOpts, sse, false)
881876
}
882877
return nil, nil
@@ -934,7 +929,6 @@ func handleOpenAIEvent(
934929
}
935930
}
936931

937-
938932
// extractMessageAndToolsFromResponse extracts the final OpenAI message and tool calls from the completed response
939933
func extractMessageAndToolsFromResponse(resp openaiResponse, state *openaiStreamingState) ([]*OpenAIChatMessage, []uctypes.WaveToolCall) {
940934
var messageContent []OpenAIMessageContent
@@ -969,32 +963,20 @@ func extractMessageAndToolsFromResponse(resp openaiResponse, state *openaiStream
969963
}
970964
}
971965

972-
// Attach UIToolUseData if available
973-
if data, ok := state.toolUseData[outputItem.CallId]; ok {
974-
toolCall.ToolUseData = data
975-
} else {
976-
log.Printf("AI no data-tooluse for %s (callid: %s)\n", outputItem.Id, outputItem.CallId)
977-
}
978-
979966
toolCalls = append(toolCalls, toolCall)
980967

981968
// Create separate FunctionCall message
982969
var argsStr string
983970
if outputItem.Arguments != "" {
984971
argsStr = outputItem.Arguments
985972
}
986-
var toolUseDataPtr *uctypes.UIMessageDataToolUse
987-
if data, ok := state.toolUseData[outputItem.CallId]; ok {
988-
toolUseDataPtr = data
989-
}
990973
functionCallMsg := &OpenAIChatMessage{
991974
MessageId: uuid.New().String(),
992975
FunctionCall: &OpenAIFunctionCallInput{
993-
Type: "function_call",
994-
CallId: outputItem.CallId,
995-
Name: outputItem.Name,
996-
Arguments: argsStr,
997-
ToolUseData: toolUseDataPtr,
976+
Type: "function_call",
977+
CallId: outputItem.CallId,
978+
Name: outputItem.Name,
979+
Arguments: argsStr,
998980
},
999981
}
1000982
messages = append(messages, functionCallMsg)

‎pkg/aiusechat/openaicomp/openaicomp-convertmessage.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -317,7 +317,7 @@ func GetFunctionCallInputByToolCallId(aiChat uctypes.AIChat, toolCallId string)
317317
}
318318

319319
// UpdateToolUseData updates the ToolUseData for a specific tool call in the chat history
320-
func UpdateToolUseData(chatId string, callId string, newToolUseData *uctypes.UIMessageDataToolUse) error {
320+
func UpdateToolUseData(chatId string, callId string, newToolUseData uctypes.UIMessageDataToolUse) error {
321321
chat := chatstore.DefaultChatStore.Get(chatId)
322322
if chat == nil {
323323
return fmt.Errorf("chat not found: %s", chatId)
@@ -333,7 +333,7 @@ func UpdateToolUseData(chatId string, callId string, newToolUseData *uctypes.UIM
333333
continue
334334
}
335335
updatedMsg := compMsg.Copy()
336-
updatedMsg.Message.ToolCalls[idx].ToolUseData = newToolUseData
336+
updatedMsg.Message.ToolCalls[idx].ToolUseData = &newToolUseData
337337
aiOpts := &uctypes.AIOptsType{
338338
APIType: chat.APIType,
339339
Model: chat.Model,

‎pkg/aiusechat/usechat-backend.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ type UseChatBackend interface {
2929

3030
// UpdateToolUseData updates the tool use data for a specific tool call in the chat.
3131
// This is used to update the UI state for tool execution (approval status, results, etc.)
32-
UpdateToolUseData(chatId string, toolCallId string, toolUseData *uctypes.UIMessageDataToolUse) error
32+
UpdateToolUseData(chatId string, toolCallId string, toolUseData uctypes.UIMessageDataToolUse) error
3333

3434
// ConvertToolResultsToNativeChatMessage converts tool execution results into native chat messages
3535
// that can be sent back to the AI backend. Returns a slice of messages (some backends may
@@ -86,7 +86,7 @@ func (b *openaiResponsesBackend) RunChatStep(
8686
return stopReason, genMsgs, rateLimitInfo, err
8787
}
8888

89-
func (b *openaiResponsesBackend) UpdateToolUseData(chatId string, toolCallId string, toolUseData *uctypes.UIMessageDataToolUse) error {
89+
func (b *openaiResponsesBackend) UpdateToolUseData(chatId string, toolCallId string, toolUseData uctypes.UIMessageDataToolUse) error {
9090
return openai.UpdateToolUseData(chatId, toolCallId, toolUseData)
9191
}
9292

@@ -140,7 +140,7 @@ func (b *openaiCompletionsBackend) RunChatStep(
140140
return stopReason, genMsgs, rateLimitInfo, err
141141
}
142142

143-
func (b *openaiCompletionsBackend) UpdateToolUseData(chatId string, toolCallId string, toolUseData *uctypes.UIMessageDataToolUse) error {
143+
func (b *openaiCompletionsBackend) UpdateToolUseData(chatId string, toolCallId string, toolUseData uctypes.UIMessageDataToolUse) error {
144144
return openaicomp.UpdateToolUseData(chatId, toolCallId, toolUseData)
145145
}
146146

@@ -173,7 +173,7 @@ func (b *anthropicBackend) RunChatStep(
173173
return stopReason, []uctypes.GenAIMessage{msg}, rateLimitInfo, err
174174
}
175175

176-
func (b *anthropicBackend) UpdateToolUseData(chatId string, toolCallId string, toolUseData *uctypes.UIMessageDataToolUse) error {
176+
func (b *anthropicBackend) UpdateToolUseData(chatId string, toolCallId string, toolUseData uctypes.UIMessageDataToolUse) error {
177177
return fmt.Errorf("UpdateToolUseData not implemented for anthropic backend")
178178
}
179179

‎pkg/aiusechat/usechat.go‎

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
"time"
1818

1919
"github.com/google/uuid"
20+
"github.com/wavetermdev/waveterm/pkg/aiusechat/aiutil"
2021
"github.com/wavetermdev/waveterm/pkg/aiusechat/chatstore"
2122
"github.com/wavetermdev/waveterm/pkg/aiusechat/uctypes"
2223
"github.com/wavetermdev/waveterm/pkg/secretstore"
@@ -239,7 +240,7 @@ func GetChatUsage(chat *uctypes.AIChat) uctypes.AIUsage {
239240
return usage
240241
}
241242

242-
func updateToolUseDataInChat(backend UseChatBackend, chatOpts uctypes.WaveChatOpts, toolCallID string, toolUseData *uctypes.UIMessageDataToolUse) {
243+
func updateToolUseDataInChat(backend UseChatBackend, chatOpts uctypes.WaveChatOpts, toolCallID string, toolUseData uctypes.UIMessageDataToolUse) {
243244
if err := backend.UpdateToolUseData(chatOpts.ChatId, toolCallID, toolUseData); err != nil {
244245
log.Printf("failed to update tool use data in chat: %v\n", err)
245246
}
@@ -279,7 +280,7 @@ func processToolCallInternal(backend UseChatBackend, toolCall uctypes.WaveToolCa
279280
}
280281
// ToolVerifyInput can modify the toolusedata. re-send it here.
281282
_ = sseHandler.AiMsgData("data-tooluse", toolCall.ID, *toolCall.ToolUseData)
282-
updateToolUseDataInChat(backend, chatOpts, toolCall.ID, toolCall.ToolUseData)
283+
updateToolUseDataInChat(backend, chatOpts, toolCall.ID, *toolCall.ToolUseData)
283284
}
284285

285286
if toolCall.ToolUseData.Approval == uctypes.ApprovalNeedsApproval {
@@ -308,7 +309,7 @@ func processToolCallInternal(backend UseChatBackend, toolCall uctypes.WaveToolCa
308309

309310
// this still happens here because we need to update the FE to say the tool call was approved
310311
_ = sseHandler.AiMsgData("data-tooluse", toolCall.ID, *toolCall.ToolUseData)
311-
updateToolUseDataInChat(backend, chatOpts, toolCall.ID, toolCall.ToolUseData)
312+
updateToolUseDataInChat(backend, chatOpts, toolCall.ID, *toolCall.ToolUseData)
312313
}
313314

314315
toolCall.ToolUseData.RunTs = time.Now().UnixMilli()
@@ -344,7 +345,7 @@ func processToolCall(backend UseChatBackend, toolCall uctypes.WaveToolCall, chat
344345

345346
if toolCall.ToolUseData != nil {
346347
_ = sseHandler.AiMsgData("data-tooluse", toolCall.ID, *toolCall.ToolUseData)
347-
updateToolUseDataInChat(backend, chatOpts, toolCall.ID, toolCall.ToolUseData)
348+
updateToolUseDataInChat(backend, chatOpts, toolCall.ID, *toolCall.ToolUseData)
348349
}
349350

350351
return result
@@ -356,17 +357,27 @@ func processToolCalls(backend UseChatBackend, stopReason *uctypes.WaveStopReason
356357
defer activeToolMap.Delete(toolCall.ID)
357358
}
358359

359-
// Send all data-tooluse packets at the beginning
360-
for _, toolCall := range stopReason.ToolCalls {
361-
if toolCall.ToolUseData != nil {
362-
log.Printf("AI data-tooluse %s\n", toolCall.ID)
363-
_ = sseHandler.AiMsgData("data-tooluse", toolCall.ID, *toolCall.ToolUseData)
364-
updateToolUseDataInChat(backend, chatOpts, toolCall.ID, toolCall.ToolUseData)
365-
if toolCall.ToolUseData.Approval == uctypes.ApprovalNeedsApproval && chatOpts.RegisterToolApproval != nil {
366-
chatOpts.RegisterToolApproval(toolCall.ID)
360+
// Create and send all data-tooluse packets at the beginning
361+
for i := range stopReason.ToolCalls {
362+
toolCall := &stopReason.ToolCalls[i]
363+
// Create toolUseData from the tool call input
364+
var argsJSON string
365+
if toolCall.Input != nil {
366+
argsBytes, err := json.Marshal(toolCall.Input)
367+
if err == nil {
368+
argsJSON = string(argsBytes)
367369
}
368370
}
371+
toolUseData := aiutil.CreateToolUseData(toolCall.ID, toolCall.Name, argsJSON, chatOpts)
372+
stopReason.ToolCalls[i].ToolUseData = &toolUseData
373+
log.Printf("AI data-tooluse %s\n", toolCall.ID)
374+
_ = sseHandler.AiMsgData("data-tooluse", toolCall.ID, toolUseData)
375+
updateToolUseDataInChat(backend, chatOpts, toolCall.ID, toolUseData)
376+
if toolUseData.Approval == uctypes.ApprovalNeedsApproval && chatOpts.RegisterToolApproval != nil {
377+
chatOpts.RegisterToolApproval(toolCall.ID)
378+
}
369379
}
380+
// At this point, all ToolCalls are guaranteed to have non-nil ToolUseData
370381

371382
var toolResults []uctypes.AIToolResult
372383
for _, toolCall := range stopReason.ToolCalls {

0 commit comments

Comments
 (0)