Merge pull request #314 from dgageot/simpler-dmr-code

Simpler DMR code

Author: dgageot

pkg/model/provider/dmr/client.go

@@ -212,8 +212,9 @@ func convertMessages(messages []chat.Message) []openai.ChatCompletionMessage {
 		}
 
 		openaiMessage := openai.ChatCompletionMessage{
-			Role: string(msg.Role),
-			Name: msg.Name,
+			Role:       string(msg.Role),
+			Name:       msg.Name,
+			ToolCallID: msg.ToolCallID,
 		}
 
 		if len(msg.MultiContent) == 0 {
@@ -229,22 +230,15 @@ func convertMessages(messages []chat.Message) []openai.ChatCompletionMessage {
 			}
 		}
 
-		if len(msg.ToolCalls) > 0 {
-			openaiMessage.ToolCalls = make([]openai.ToolCall, len(msg.ToolCalls))
-			for j, toolCall := range msg.ToolCalls {
-				openaiMessage.ToolCalls[j] = openai.ToolCall{
-					ID:   toolCall.ID,
-					Type: openai.ToolType(toolCall.Type),
-					Function: openai.FunctionCall{
-						Name:      toolCall.Function.Name,
-						Arguments: toolCall.Function.Arguments,
-					},
-				}
-			}
-		}
-
-		if msg.ToolCallID != "" {
-			openaiMessage.ToolCallID = msg.ToolCallID
+		for _, call := range msg.ToolCalls {
+			openaiMessage.ToolCalls = append(openaiMessage.ToolCalls, openai.ToolCall{
+				ID:   call.ID,
+				Type: openai.ToolType(call.Type),
+				Function: openai.FunctionCall{
+					Name:      call.Function.Name,
+					Arguments: call.Function.Arguments,
+				},
+			})
 		}
 
 		openaiMessages = append(openaiMessages, openaiMessage)
@@ -297,11 +291,7 @@ func convertMessages(messages []chat.Message) []openai.ChatCompletionMessage {
 
 // CreateChatCompletionStream creates a streaming chat completion request
 // It returns a stream that can be iterated over to get completion chunks
-func (c *Client) CreateChatCompletionStream(
-	ctx context.Context,
-	messages []chat.Message,
-	requestTools []tools.Tool,
-) (chat.MessageStream, error) {
+func (c *Client) CreateChatCompletionStream(ctx context.Context, messages []chat.Message, requestTools []tools.Tool) (chat.MessageStream, error) {
 	slog.Debug("Creating DMR chat completion stream",
 		"model", c.config.Model,
 		"message_count", len(messages),
@@ -314,10 +304,7 @@ func (c *Client) CreateChatCompletionStream(
 		return nil, errors.New("at least one message is required")
 	}
 
-	trackUsage := true
-	if c.config.TrackUsage != nil && *c.config.TrackUsage == false {
-		trackUsage = false
-	}
+	trackUsage := c.config.TrackUsage == nil || *c.config.TrackUsage
 
 	request := openai.ChatCompletionRequest{
 		Model:            c.config.Model,
@@ -387,10 +374,7 @@ func (c *Client) CreateChatCompletionStream(
 	return newStreamAdapter(stream, trackUsage), nil
 }
 
-func (c *Client) CreateChatCompletion(
-	ctx context.Context,
-	messages []chat.Message,
-) (string, error) {
+func (c *Client) CreateChatCompletion(ctx context.Context, messages []chat.Message) (string, error) {
 	slog.Debug("Creating DMR chat completion", "model", c.config.Model, "message_count", len(messages), "base_url", c.baseURL)
 
 	request := openai.ChatCompletionRequest{
@@ -412,15 +396,13 @@ func (c *Client) CreateChatCompletion(
 // In particular, it avoids shadowing built-in mapping methods like `keys()` by removing a literal "keys"
 // field from property schemas if present, and guarantees the outer structure is an object with a properties map.
 func sanitizeToolParameters(p tools.FunctionParameters) any {
-	// Start with a safe container
 	out := map[string]any{
 		"type":       "object",
 		"properties": map[string]any{},
 	}
 	if p.Type != "" {
 		out["type"] = p.Type
 	}
-	// Copy required if present
 	if len(p.Required) > 0 {
 		out["required"] = p.Required
 	}

pkg/model/provider/dmr/client_test.go

@@ -1,12 +1,12 @@
 package dmr
 
 import (
-	"reflect"
 	"testing"
 
-	latest "github.com/docker/cagent/pkg/config/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	latest "github.com/docker/cagent/pkg/config/v2"
 )
 
 func TestNewClientWithExplicitBaseURL(t *testing.T) {
@@ -33,27 +33,19 @@ func TestNewClientWithWrongType(t *testing.T) {
 
 func TestBuildDockerConfigureArgs(t *testing.T) {
 	args := buildDockerModelConfigureArgs("ai/qwen3:14B-Q6_K", 8192, []string{"--temp", "0.7", "--top-p", "0.9"})
-	expected := []string{"model", "configure", "--context-size=8192", "ai/qwen3:14B-Q6_K", "--", "--temp", "0.7", "--top-p", "0.9"}
-	if !reflect.DeepEqual(args, expected) {
-		t.Fatalf("unexpected args.\nexpected: %#v\nactual:   %#v", expected, args)
-	}
+
+	assert.Equal(t, []string{"model", "configure", "--context-size=8192", "ai/qwen3:14B-Q6_K", "--", "--temp", "0.7", "--top-p", "0.9"}, args)
 }
 
 func TestBuildRuntimeFlagsFromModelConfig_LlamaCpp(t *testing.T) {
-	cfg := &latest.ModelConfig{
+	flags := buildRuntimeFlagsFromModelConfig("llama.cpp", &latest.ModelConfig{
 		Temperature:      0.6,
 		TopP:             0.95,
 		FrequencyPenalty: 0.2,
 		PresencePenalty:  0.1,
-	}
+	})
 
-	flags := buildRuntimeFlagsFromModelConfig("llama.cpp", cfg)
-
-	// Order matters based on implementation
-	expected := []string{"--temp", "0.6", "--top-p", "0.95", "--frequency-penalty", "0.2", "--presence-penalty", "0.1"}
-	if !reflect.DeepEqual(flags, expected) {
-		t.Fatalf("unexpected runtime flags.\nexpected: %#v\nactual:   %#v", expected, flags)
-	}
+	assert.Equal(t, []string{"--temp", "0.6", "--top-p", "0.95", "--frequency-penalty", "0.2", "--presence-penalty", "0.1"}, flags)
 }
 
 func TestIntegrateFlagsWithProviderOptsOrder(t *testing.T) {
@@ -71,10 +63,7 @@ func TestIntegrateFlagsWithProviderOptsOrder(t *testing.T) {
 	merged := append(derived, []string{"--threads", "6"}...)
 
 	args := buildDockerModelConfigureArgs("ai/qwen3:14B-Q6_K", cfg.MaxTokens, merged)
-	expected := []string{"model", "configure", "--context-size=4096", "ai/qwen3:14B-Q6_K", "--", "--temp", "0.6", "--top-p", "0.9", "--threads", "6"}
-	if !reflect.DeepEqual(args, expected) {
-		t.Fatalf("unexpected configure args.\nexpected: %#v\nactual:   %#v", expected, args)
-	}
+	assert.Equal(t, []string{"model", "configure", "--context-size=4096", "ai/qwen3:14B-Q6_K", "--", "--temp", "0.6", "--top-p", "0.9", "--threads", "6"}, args)
 }
 
 func TestMergeRuntimeFlagsPreferUser_WarnsAndPrefersUser(t *testing.T) {
@@ -85,12 +74,8 @@ func TestMergeRuntimeFlagsPreferUser_WarnsAndPrefersUser(t *testing.T) {
 	merged, warnings := mergeRuntimeFlagsPreferUser(derived, user)
 
 	// Expect 1 warnings for --temp overriding
-	if len(warnings) != 1 {
-		t.Fatalf("expected 1 warning1, got %d: %#v", len(warnings), warnings)
-	}
+	require.Len(t, warnings, 1)
+
 	// Derived conflicting flags should be dropped, user ones kept and appended
-	expected := []string{"--top-p", "0.8", "--temp", "0.7", "--threads", "8"}
-	if !reflect.DeepEqual(merged, expected) {
-		t.Fatalf("unexpected merged flags.\nexpected: %#v\nactual:   %#v", expected, merged)
-	}
+	assert.Equal(t, []string{"--top-p", "0.8", "--temp", "0.7", "--threads", "8"}, merged)
 }