fix[go]: support retrieval kb in agent chat (#16944)

### Summary As title - [x] fix tool & component `Retrieval KB` - [x] fix agent cannot use `Retrieval KB` in agent chat #### Main cause Model provider do not set up: ```Go if chatModelConfig.Tools != nil { reqBody["tools"] = chatModelConfig.Tools } ``` #### Working Now <img width="3774" height="2128" alt="image" src="https://github.com/user-attachments/assets/400a349d-0211-43e5-a7ec-7a014acf77a6" /> ``` ____________________________________ < This PR takes me all day to do it. > ------------------------------------ \ \ (\__/) (•ㅅ•) / づ ```
2026-07-16 20:57:21 +08:00 · 2026-07-15 21:44:28 +08:00
parent a7da78d0d7
commit ff3d566a4f
15 changed files with 1559 additions and 95 deletions
--- a/internal/entity/models/deepseek.go
+++ b/internal/entity/models/deepseek.go
@@ -24,6 +24,7 @@ import (
 	"io"
 	"net/http"
 	"ragflow/internal/common"
+	"sort"
 	"strconv"
 	"strings"
 )
@@ -70,10 +71,17 @@ func (d *DeepSeekModel) ChatWithMessages(modelName string, messages []Message, a
 	// Convert messages to the format expected by API
 	apiMessages := make([]map[string]interface{}, len(messages))
 	for i, msg := range messages {
-		apiMessages[i] = map[string]interface{}{
+		apiMsg := map[string]interface{}{
 			"role":    msg.Role,
 			"content": msg.Content,
 		}
+		if msg.ToolCallID != "" {
+			apiMsg["tool_call_id"] = msg.ToolCallID
+		}
+		if len(msg.ToolCalls) > 0 {
+			apiMsg["tool_calls"] = msg.ToolCalls
+		}
+		apiMessages[i] = apiMsg
 	}

 	// Build request body
@@ -105,6 +113,13 @@ func (d *DeepSeekModel) ChatWithMessages(modelName string, messages []Message, a
 			reqBody["stop"] = *chatModelConfig.Stop
 		}

+		if chatModelConfig.Tools != nil {
+			reqBody["tools"] = chatModelConfig.Tools
+		}
+		if chatModelConfig.ToolChoice != nil {
+			reqBody["tool_choice"] = *chatModelConfig.ToolChoice
+		}
+
 		if chatModelConfig.Thinking != nil {
 			if *chatModelConfig.Thinking {
 				var thinkingFlag string
@@ -198,26 +213,38 @@ func (d *DeepSeekModel) ChatWithMessages(modelName string, messages []Message, a
 		return nil, fmt.Errorf("invalid message format")
 	}

-	content, ok := messageMap["content"].(string)
-	if !ok {
-		return nil, fmt.Errorf("invalid content format")
+	var content string
+	if c, ok := messageMap["content"].(string); ok {
+		content = c
 	}

 	var reasonContent string
-	if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking {
-		reasonContent, ok = messageMap["reasoning_content"].(string)
-		if !ok {
-			return nil, fmt.Errorf("invalid content format")
-		}
+	if rc, ok := messageMap["reasoning_content"].(string); ok {
+		reasonContent = rc
 		// if first char of reasonContent is \n remove the '\n'
 		if reasonContent != "" && reasonContent[0] == '\n' {
 			reasonContent = reasonContent[1:]
 		}
 	}

+	var toolCalls []map[string]interface{}
+	if tcs, ok := messageMap["tool_calls"].([]interface{}); ok {
+		for _, tc := range tcs {
+			if tcMap, ok := tc.(map[string]interface{}); ok {
+				toolCalls = append(toolCalls, tcMap)
+			}
+		}
+	}
+
 	chatResponse := &ChatResponse{
 		Answer:        &content,
 		ReasonContent: &reasonContent,
+		ToolCalls:     toolCalls,
+	}
+	if pt, ct, tt := extractUsageFromMap(result); tt > 0 {
+		chatResponse.Usage = &ChatUsage{
+			PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt,
+		}
 	}

 	return chatResponse, nil
@@ -242,10 +269,17 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 	// Convert messages to API format
 	apiMessages := make([]map[string]interface{}, len(messages))
 	for i, msg := range messages {
-		apiMessages[i] = map[string]interface{}{
+		apiMsg := map[string]interface{}{
 			"role":    msg.Role,
 			"content": msg.Content,
 		}
+		if msg.ToolCallID != "" {
+			apiMsg["tool_call_id"] = msg.ToolCallID
+		}
+		if len(msg.ToolCalls) > 0 {
+			apiMsg["tool_calls"] = msg.ToolCalls
+		}
+		apiMessages[i] = apiMsg
 	}

 	// Build request body with streaming enabled
@@ -280,6 +314,12 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 		if chatModelConfig.Stop != nil {
 			reqBody["stop"] = *chatModelConfig.Stop
 		}
+		if chatModelConfig.Tools != nil {
+			reqBody["tools"] = chatModelConfig.Tools
+		}
+		if chatModelConfig.ToolChoice != nil {
+			reqBody["tool_choice"] = *chatModelConfig.ToolChoice
+		}

 		if chatModelConfig.Thinking != nil {
 			if *chatModelConfig.Thinking {
@@ -351,8 +391,8 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 		return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
 	}

-	// SSE parsing: read line by line
 	sawTerminal := false
+	accumulatedToolCalls := make(map[int]map[string]interface{})
 	done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
 		common.Info(fmt.Sprintf("%v", event))

@@ -371,6 +411,37 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 			return nil
 		}

+		if tcs, ok := delta["tool_calls"].([]interface{}); ok {
+			for _, tc := range tcs {
+				tcMap, ok := tc.(map[string]interface{})
+				if !ok {
+					continue
+				}
+				idxF, ok := tcMap["index"].(float64)
+				if !ok {
+					continue
+				}
+				idx := int(idxF)
+				existing, hasExisting := accumulatedToolCalls[idx]
+				if hasExisting {
+					if fn, ok := tcMap["function"].(map[string]interface{}); ok {
+						if args, ok := fn["arguments"].(string); ok {
+							if ef, ok := existing["function"].(map[string]interface{}); ok {
+								if ea, ok := ef["arguments"].(string); ok {
+									ef["arguments"] = ea + args
+								} else {
+									ef["arguments"] = args
+								}
+							}
+						}
+					}
+				} else {
+					accumulatedToolCalls[idx] = cloneMap(tcMap)
+				}
+			}
+			return nil
+		}
+
 		content, ok := delta["content"].(string)
 		if ok && content != "" {
 			if err := sender(&content, nil); err != nil {
@@ -398,6 +469,19 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 		return fmt.Errorf("deepseek: stream ended before [DONE] or finish_reason")
 	}

+	if len(accumulatedToolCalls) > 0 && chatModelConfig != nil {
+		indices := make([]int, 0, len(accumulatedToolCalls))
+		for idx := range accumulatedToolCalls {
+			indices = append(indices, idx)
+		}
+		sort.Ints(indices)
+		tcs := make([]map[string]interface{}, 0, len(accumulatedToolCalls))
+		for _, idx := range indices {
+			tcs = append(tcs, accumulatedToolCalls[idx])
+		}
+		chatModelConfig.ToolCallsResult = &tcs
+	}
+
 	// Send [DONE] marker for OpenAI compatibility
 	endOfStream := "[DONE]"
 	return sender(&endOfStream, nil)
--- a/internal/entity/models/deepseek_test.go
+++ b/internal/entity/models/deepseek_test.go
@@ -0,0 +1,161 @@
+package models
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func newDeepSeekForTest(baseURL string) *DeepSeekModel {
+	return NewDeepSeekModel(
+		map[string]string{"default": baseURL},
+		URLSuffix{
+			Chat:   "chat/completions",
+			Models: "models",
+		},
+	)
+}
+
+func TestDeepSeekChatWithMessagesSupportsToolCalls(t *testing.T) {
+	var requestBody map[string]interface{}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/chat/completions" {
+			t.Errorf("path=%s, want /chat/completions", r.URL.Path)
+		}
+		if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
+			t.Errorf("Authorization=%q, want Bearer test-key", got)
+		}
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Errorf("decode request: %v", err)
+			http.Error(w, "bad request", http.StatusBadRequest)
+			return
+		}
+		_ = json.NewEncoder(w).Encode(map[string]interface{}{
+			"choices": []map[string]interface{}{
+				{
+					"message": map[string]interface{}{
+						"role":    "assistant",
+						"content": nil,
+						"tool_calls": []map[string]interface{}{
+							{
+								"id":   "call-1",
+								"type": "function",
+								"function": map[string]interface{}{
+									"name":      "search_my_dateset",
+									"arguments": `{"query":"marigold"}`,
+								},
+							},
+						},
+					},
+				},
+			},
+			"usage": map[string]interface{}{
+				"prompt_tokens":     1,
+				"completion_tokens": 2,
+				"total_tokens":      3,
+			},
+		})
+	}))
+	defer srv.Close()
+
+	apiKey := "test-key"
+	toolChoice := "auto"
+	resp, err := newDeepSeekForTest(srv.URL).ChatWithMessages(
+		"deepseek-chat",
+		[]Message{
+			{Role: "user", Content: "what is marigold"},
+		},
+		&APIConfig{ApiKey: &apiKey},
+		&ChatConfig{
+			Tools: []map[string]interface{}{
+				{
+					"type": "function",
+					"function": map[string]interface{}{
+						"name":        "search_my_dateset",
+						"description": "Search datasets.",
+					},
+				},
+			},
+			ToolChoice: &toolChoice,
+		},
+	)
+	if err != nil {
+		t.Fatalf("ChatWithMessages: %v", err)
+	}
+	if requestBody["tool_choice"] != "auto" {
+		t.Fatalf("tool_choice=%#v, want auto", requestBody["tool_choice"])
+	}
+	if _, ok := requestBody["tools"].([]interface{}); !ok {
+		t.Fatalf("tools missing or wrong type: %#v", requestBody["tools"])
+	}
+	if resp.Answer == nil || *resp.Answer != "" {
+		t.Fatalf("Answer=%#v, want empty string for tool-call response", resp.Answer)
+	}
+	if len(resp.ToolCalls) != 1 {
+		t.Fatalf("ToolCalls len=%d, want 1", len(resp.ToolCalls))
+	}
+	fn, _ := resp.ToolCalls[0]["function"].(map[string]interface{})
+	if fn["name"] != "search_my_dateset" {
+		t.Fatalf("tool call function=%#v, want search_my_dateset", fn)
+	}
+	if resp.Usage == nil || resp.Usage.TotalTokens != 3 {
+		t.Fatalf("Usage=%#v, want total tokens 3", resp.Usage)
+	}
+}
+
+func TestDeepSeekChatWithMessagesForwardsToolHistory(t *testing.T) {
+	var requestBody map[string]interface{}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Errorf("decode request: %v", err)
+			http.Error(w, "bad request", http.StatusBadRequest)
+			return
+		}
+		_ = json.NewEncoder(w).Encode(map[string]interface{}{
+			"choices": []map[string]interface{}{
+				{"message": map[string]interface{}{"role": "assistant", "content": "answer"}},
+			},
+		})
+	}))
+	defer srv.Close()
+
+	apiKey := "test-key"
+	_, err := newDeepSeekForTest(srv.URL).ChatWithMessages(
+		"deepseek-chat",
+		[]Message{
+			{
+				Role:    "assistant",
+				Content: nil,
+				ToolCalls: []map[string]interface{}{
+					{
+						"id":   "call-1",
+						"type": "function",
+						"function": map[string]interface{}{
+							"name":      "search_my_dateset",
+							"arguments": `{"query":"marigold"}`,
+						},
+					},
+				},
+			},
+			{Role: "tool", ToolCallID: "call-1", Content: `{"formalized_content":"marigold"}`},
+		},
+		&APIConfig{ApiKey: &apiKey},
+		nil,
+	)
+	if err != nil {
+		t.Fatalf("ChatWithMessages: %v", err)
+	}
+	messages, ok := requestBody["messages"].([]interface{})
+	if !ok || len(messages) != 2 {
+		t.Fatalf("messages=%#v, want 2 messages", requestBody["messages"])
+	}
+	assistant, _ := messages[0].(map[string]interface{})
+	if _, ok := assistant["tool_calls"].([]interface{}); !ok {
+		t.Fatalf("assistant tool_calls missing: %#v", assistant)
+	}
+	toolMsg, _ := messages[1].(map[string]interface{})
+	if toolMsg["tool_call_id"] != "call-1" {
+		t.Fatalf("tool_call_id=%#v, want call-1", toolMsg["tool_call_id"])
+	}
+}
--- a/internal/entity/models/llm.go
+++ b/internal/entity/models/llm.go
@@ -32,6 +32,7 @@ package models

 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"sync"

@@ -89,7 +90,14 @@ func toInternalMessages(msgs []*schema.Message) []Message {
 		if role == "" {
 			role = "user"
 		}
-		out = append(out, Message{Role: role, Content: mm.Content})
+		msg := Message{Role: role, Content: mm.Content}
+		if len(mm.ToolCalls) > 0 {
+			msg.ToolCalls = toolCallsToInternal(mm.ToolCalls)
+		}
+		if mm.ToolCallID != "" {
+			msg.ToolCallID = mm.ToolCallID
+		}
+		out = append(out, msg)
 	}
 	return out
 }
@@ -105,7 +113,14 @@ func fromInternalResponse(resp *ChatResponse) *schema.Message {
 	if resp.Answer != nil {
 		content = *resp.Answer
 	}
-	return &schema.Message{Role: schema.Assistant, Content: content}
+	msg := &schema.Message{Role: schema.Assistant, Content: content}
+	if resp.ReasonContent != nil {
+		msg.ReasoningContent = *resp.ReasonContent
+	}
+	if len(resp.ToolCalls) > 0 {
+		msg.ToolCalls = toolCallsFromInternal(resp.ToolCalls)
+	}
+	return msg
 }

 // Generate blocks until the model returns a complete response. Mirrors
@@ -128,7 +143,11 @@ func (m *EinoChatModel) Generate(ctx context.Context, msgs []*schema.Message, op
 	// without a usage block doesn't leak the previous call's data.
 	// Mirrors Python's LLMBundle._reset_last_usage().
 	m.inner.LastUsage = nil
-	resp, err := m.inner.ModelDriver.ChatWithMessages(*m.inner.ModelName, internal, m.inner.APIConfig, m.chatCfg)
+	chatCfg, err := m.chatConfigForGenerate()
+	if err != nil {
+		return nil, err
+	}
+	resp, err := m.inner.ModelDriver.ChatWithMessages(*m.inner.ModelName, internal, m.inner.APIConfig, chatCfg)
 	if err != nil {
 		return nil, fmt.Errorf("models: EinoChatModel.Generate(%s): %w", *m.inner.ModelName, err)
 	}
@@ -144,6 +163,106 @@ func (m *EinoChatModel) Generate(ctx context.Context, msgs []*schema.Message, op
 	return fromInternalResponse(resp), nil
 }

+func (m *EinoChatModel) chatConfigForGenerate() (*ChatConfig, error) {
+	if len(m.tools) == 0 {
+		return m.chatCfg, nil
+	}
+	cfg := &ChatConfig{}
+	if m.chatCfg != nil {
+		cp := *m.chatCfg
+		cfg = &cp
+	}
+	tools, err := openAIToolsFromEino(m.tools)
+	if err != nil {
+		return nil, err
+	}
+	cfg.Tools = tools
+	choice := "auto"
+	cfg.ToolChoice = &choice
+	return cfg, nil
+}
+
+func openAIToolsFromEino(infos []*schema.ToolInfo) ([]map[string]any, error) {
+	tools := make([]map[string]any, 0, len(infos))
+	for _, info := range infos {
+		if info == nil {
+			continue
+		}
+		fn := map[string]any{
+			"name":        info.Name,
+			"description": info.Desc,
+		}
+		if info.ParamsOneOf != nil {
+			params, err := info.ParamsOneOf.ToJSONSchema()
+			if err != nil {
+				return nil, fmt.Errorf("models: convert tool %q schema: %w", info.Name, err)
+			}
+			fn["parameters"] = params
+		} else {
+			fn["parameters"] = map[string]any{"type": "object", "properties": map[string]any{}}
+		}
+		tools = append(tools, map[string]any{
+			"type":     "function",
+			"function": fn,
+		})
+	}
+	return tools, nil
+}
+
+func toolCallsToInternal(calls []schema.ToolCall) []map[string]interface{} {
+	out := make([]map[string]interface{}, 0, len(calls))
+	for _, call := range calls {
+		fn := map[string]interface{}{
+			"name":      call.Function.Name,
+			"arguments": call.Function.Arguments,
+		}
+		out = append(out, map[string]interface{}{
+			"id":       call.ID,
+			"type":     call.Type,
+			"function": fn,
+		})
+	}
+	return out
+}
+
+func toolCallsFromInternal(calls []map[string]interface{}) []schema.ToolCall {
+	out := make([]schema.ToolCall, 0, len(calls))
+	for i, call := range calls {
+		id, _ := call["id"].(string)
+		if id == "" {
+			id = fmt.Sprintf("call_%d", i)
+		}
+		callType, _ := call["type"].(string)
+		if callType == "" {
+			callType = "function"
+		}
+		var fnName, fnArgs string
+		if fn, ok := call["function"].(map[string]interface{}); ok {
+			fnName, _ = fn["name"].(string)
+			switch args := fn["arguments"].(type) {
+			case string:
+				fnArgs = args
+			case nil:
+				fnArgs = "{}"
+			default:
+				b, err := json.Marshal(args)
+				if err == nil {
+					fnArgs = string(b)
+				}
+			}
+		}
+		out = append(out, schema.ToolCall{
+			ID:   id,
+			Type: callType,
+			Function: schema.FunctionCall{
+				Name:      fnName,
+				Arguments: fnArgs,
+			},
+		})
+	}
+	return out
+}
+
 // Stream returns a schema.StreamReader that yields message chunks
 // incrementally. Uses the existing ChatStreamlyWithSender pathway; the
 // sender callback pushes the streamed delta into the StreamReader.
@@ -157,6 +276,13 @@ func (m *EinoChatModel) Stream(ctx context.Context, msgs []*schema.Message, opts
 	if m.inner.ModelName == nil {
 		return nil, fmt.Errorf("models: EinoChatModel: nil model name")
 	}
+	if len(m.tools) > 0 {
+		msg, err := m.Generate(ctx, msgs, opts...)
+		if err != nil {
+			return nil, err
+		}
+		return schema.StreamReaderFromArray([]*schema.Message{msg}), nil
+	}
 	internal := toInternalMessages(msgs)

 	sr, sw := schema.Pipe[*schema.Message](1)
--- a/internal/entity/models/llm_test.go
+++ b/internal/entity/models/llm_test.go
@@ -0,0 +1,196 @@
+package models
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cloudwego/eino/schema"
+)
+
+func TestEinoChatModelGenerateSendsBoundTools(t *testing.T) {
+	apiKey := "key"
+	modelName := "chat"
+	driver := &captureToolDriver{
+		resp: &ChatResponse{
+			ToolCalls: []map[string]interface{}{
+				{
+					"id":   "call-1",
+					"type": "function",
+					"function": map[string]interface{}{
+						"name":      "search_my_dateset",
+						"arguments": `{"query":"hello"}`,
+					},
+				},
+			},
+		},
+	}
+	base := NewChatModel(driver, &modelName, &APIConfig{ApiKey: &apiKey})
+	model := NewEinoChatModel(base, nil)
+	bound, err := model.WithTools([]*schema.ToolInfo{
+		{
+			Name: "search_my_dateset",
+			Desc: "Search datasets.",
+			ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
+				"query": &schema.ParameterInfo{Type: schema.String, Required: true},
+			}),
+		},
+	})
+	if err != nil {
+		t.Fatalf("WithTools: %v", err)
+	}
+
+	msg, err := bound.Generate(context.Background(), []*schema.Message{schema.UserMessage("hello")})
+	if err != nil {
+		t.Fatalf("Generate: %v", err)
+	}
+	if driver.lastConfig == nil || driver.lastConfig.Tools == nil {
+		t.Fatal("Generate did not send tools to driver")
+	}
+	tools, ok := driver.lastConfig.Tools.([]map[string]any)
+	if !ok || len(tools) != 1 {
+		t.Fatalf("driver tools = %#v, want one OpenAI-style tool", driver.lastConfig.Tools)
+	}
+	fn, _ := tools[0]["function"].(map[string]any)
+	if fn["name"] != "search_my_dateset" {
+		t.Fatalf("tool function name = %#v, want search_my_dateset", fn["name"])
+	}
+	if driver.lastConfig.ToolChoice == nil || *driver.lastConfig.ToolChoice != "auto" {
+		t.Fatalf("ToolChoice = %#v, want auto", driver.lastConfig.ToolChoice)
+	}
+	if len(msg.ToolCalls) != 1 {
+		t.Fatalf("msg.ToolCalls len = %d, want 1", len(msg.ToolCalls))
+	}
+	if msg.ToolCalls[0].Function.Name != "search_my_dateset" || msg.ToolCalls[0].Function.Arguments != `{"query":"hello"}` {
+		t.Fatalf("tool call = %#v, want search_my_dateset query call", msg.ToolCalls[0])
+	}
+}
+
+func TestEinoChatModelStreamWithToolsYieldsToolCalls(t *testing.T) {
+	apiKey := "key"
+	modelName := "chat"
+	driver := &captureToolDriver{
+		resp: &ChatResponse{
+			ToolCalls: []map[string]interface{}{
+				{
+					"id":   "call-1",
+					"type": "function",
+					"function": map[string]interface{}{
+						"name":      "search_my_dateset",
+						"arguments": `{"query":"hello"}`,
+					},
+				},
+			},
+		},
+	}
+	base := NewChatModel(driver, &modelName, &APIConfig{ApiKey: &apiKey})
+	model := NewEinoChatModel(base, nil)
+	bound, err := model.WithTools([]*schema.ToolInfo{
+		{
+			Name: "search_my_dateset",
+			Desc: "Search datasets.",
+			ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
+				"query": &schema.ParameterInfo{Type: schema.String, Required: true},
+			}),
+		},
+	})
+	if err != nil {
+		t.Fatalf("WithTools: %v", err)
+	}
+
+	stream, err := bound.Stream(context.Background(), []*schema.Message{schema.UserMessage("hello")})
+	if err != nil {
+		t.Fatalf("Stream: %v", err)
+	}
+	msg, err := stream.Recv()
+	if err != nil {
+		t.Fatalf("stream.Recv: %v", err)
+	}
+	if msg == nil {
+		t.Fatal("stream ended before yielding message")
+	}
+	if len(msg.ToolCalls) != 1 || msg.ToolCalls[0].Function.Name != "search_my_dateset" {
+		t.Fatalf("stream message tool calls = %#v, want search_my_dateset", msg.ToolCalls)
+	}
+}
+
+func TestToInternalMessagesPreservesToolMessages(t *testing.T) {
+	internal := toInternalMessages([]*schema.Message{
+		{
+			Role: schema.Assistant,
+			ToolCalls: []schema.ToolCall{{
+				ID:   "call-1",
+				Type: "function",
+				Function: schema.FunctionCall{
+					Name:      "search_my_dateset",
+					Arguments: `{"query":"hello"}`,
+				},
+			}},
+		},
+		{
+			Role:       schema.Tool,
+			Content:    `{"formalized_content":"answer"}`,
+			ToolCallID: "call-1",
+		},
+	})
+	if len(internal) != 2 {
+		t.Fatalf("len(internal) = %d, want 2", len(internal))
+	}
+	if len(internal[0].ToolCalls) != 1 {
+		t.Fatalf("assistant ToolCalls = %#v, want one tool call", internal[0].ToolCalls)
+	}
+	if internal[1].ToolCallID != "call-1" || internal[1].Role != "tool" {
+		t.Fatalf("tool message = %#v, want tool role with call id", internal[1])
+	}
+}
+
+type captureToolDriver struct {
+	resp       *ChatResponse
+	lastConfig *ChatConfig
+}
+
+func (d *captureToolDriver) NewInstance(baseURL map[string]string) ModelDriver { return d }
+func (d *captureToolDriver) Name() string                                      { return "capture" }
+func (d *captureToolDriver) ChatWithMessages(_ string, _ []Message, _ *APIConfig, cfg *ChatConfig) (*ChatResponse, error) {
+	d.lastConfig = cfg
+	return d.resp, nil
+}
+func (d *captureToolDriver) ChatStreamlyWithSender(_ string, _ []Message, _ *APIConfig, _ *ChatConfig, _ func(*string, *string) error) error {
+	return nil
+}
+func (d *captureToolDriver) Embed(_ *string, _ []string, _ *APIConfig, _ *EmbeddingConfig) ([]EmbeddingData, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) Rerank(_ *string, _ string, _ []string, _ *APIConfig, _ *RerankConfig) (*RerankResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) TranscribeAudio(_ *string, _ *string, _ *APIConfig, _ *ASRConfig) (*ASRResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) TranscribeAudioWithSender(_ *string, _ *string, _ *APIConfig, _ *ASRConfig, _ func(*string, *string) error) error {
+	return nil
+}
+func (d *captureToolDriver) AudioSpeech(_ *string, _ *string, _ *APIConfig, _ *TTSConfig) (*TTSResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) AudioSpeechWithSender(_ *string, _ *string, _ *APIConfig, _ *TTSConfig, _ func(*string, *string) error) error {
+	return nil
+}
+func (d *captureToolDriver) OCRFile(_ *string, _ []byte, _ *string, _ *APIConfig, _ *OCRConfig) (*OCRFileResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) ParseFile(_ *string, _ []byte, _ *string, _ *APIConfig, _ *ParseFileConfig) (*ParseFileResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) ListModels(_ *APIConfig) ([]ListModelResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) Balance(_ *APIConfig) (map[string]interface{}, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) CheckConnection(_ *APIConfig) error { return nil }
+func (d *captureToolDriver) ListTasks(_ *APIConfig) ([]ListTaskStatus, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) ShowTask(_ string, _ *APIConfig) (*TaskResponse, error) {
+	return nil, nil
+}