odidere/internal/llm/llm_test.go

package llm

import (
	"context"
	"encoding/json/v2"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestContentPartMarshalJSON(t *testing.T) {
	tests := []struct {
		name string
		part ContentPart
		want string
	}{
		{
			name: "text",
			part: ContentPart{
				Type: ContentTypeText,
				Text: "hello",
			},
			want: `{"text":"hello","type":"text"}`,
		},
		{
			name: "refusal",
			part: ContentPart{
				Type:    ContentTypeRefusal,
				Refusal: "cannot comply",
			},
			want: `{"refusal":"cannot comply","type":"refusal"}`,
		},
		{
			// The payload field is required, also when it is empty.
			name: "empty text",
			part: ContentPart{Type: ContentTypeText},
			want: `{"text":"","type":"text"}`,
		},
		{
			name: "empty refusal",
			part: ContentPart{Type: ContentTypeRefusal},
			want: `{"refusal":"","type":"refusal"}`,
		},
		{
			name: "image URL",
			part: ContentPart{
				Type: ContentTypeImageURL,
				ImageURL: &ImageURL{
					URL: "https://example.com/image.png",
				},
			},
			want: `{"image_url":{"url":"https://example.com/image.png"},"type":"image_url"}`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := json.Marshal(tt.part)
			if err != nil {
				t.Fatalf("json.Marshal() error = %v", err)
			}
			if string(got) != tt.want {
				t.Errorf("json.Marshal() = %s, want %s", got, tt.want)
			}
		})
	}
}

func TestClientReturnsAPIError(t *testing.T) {
	tests := []struct {
		name    string
		call    func(*Client) error
		wantErr string
	}{
		{
			name: "completions",
			call: func(client *Client) error {
				_, err := client.Completions(t.Context(), CompletionsParams{})
				return err
			},
			wantErr: "completions: model_not_found: model unavailable",
		},
		{
			name: "completions stream",
			call: func(client *Client) error {
				_, err := client.CompletionsStream(
					t.Context(), CompletionsParams{},
				)
				return err
			},
			wantErr: "completions stream: model_not_found: model unavailable",
		},
		{
			name: "list models",
			call: func(client *Client) error {
				_, err := client.ListModels(context.Background())
				return err
			},
			wantErr: "list models: model_not_found: model unavailable",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			server := httptest.NewServer(http.HandlerFunc(func(
				w http.ResponseWriter, _ *http.Request,
			) {
				w.WriteHeader(http.StatusNotFound)
				fmt.Fprint(w, `{"error":{"code":"model_not_found","message":"model unavailable","type":"invalid_request_error"}}`)
			}))
			defer server.Close()

			log := slog.New(slog.NewTextHandler(io.Discard, nil))
			client, err := NewClient(Config{URL: server.URL}, log)
			if err != nil {
				t.Fatalf("NewClient() error = %v", err)
			}

			err = tt.call(client)
			if err == nil {
				t.Fatal("call error = nil, want API error")
			}
			if err.Error() != tt.wantErr {
				t.Errorf("call error = %q, want %q", err, tt.wantErr)
			}

			var apiErr *APIError
			if !errors.As(err, &apiErr) {
				t.Fatalf("errors.As(%T) = false, want *APIError", err)
			}
			if apiErr.Code != "model_not_found" {
				t.Errorf("APIError.Code = %q, want %q", apiErr.Code, "model_not_found")
			}
			if apiErr.Message != "model unavailable" {
				t.Errorf("APIError.Message = %q, want %q", apiErr.Message, "model unavailable")
			}
			if apiErr.Type != "invalid_request_error" {
				t.Errorf(
					"APIError.Type = %q, want %q",
					apiErr.Type,
					"invalid_request_error",
				)
			}
		})
	}
}

func TestMessageUnmarshalJSONReplacesArrayContent(t *testing.T) {
	message := Message{
		ContentParts: []ContentPart{{
			Type: ContentTypeText,
			Text: "old",
		}},
	}

	err := json.Unmarshal(
		[]byte(`{"content":[{"type":"text","text":"new"}]}`),
		&message,
	)
	if err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	if got := message.Text(); got != "new" {
		t.Errorf("Message.Text() = %q, want %q", got, "new")
	}
}

// TestMessageEmptyTextRoundTrip is a regression test. ContentPart.Text used
// the omitzero option, so a text part with empty text marshaled as
// {"type":"text"}, without a text key. The message was stored in that form
// and sent to the model server on the next inference, and the server
// rejected the request: the OpenAI API requires the text field, and
// llama.cpp fails to parse the part. MarshalJSON now always writes the
// payload field. Rows stored in the old form are corrected when they are
// decoded and encoded again.
func TestMessageEmptyTextRoundTrip(t *testing.T) {
	tests := []struct {
		name  string
		input string
		want  string
	}{
		{
			// Model replies that contain only tool calls can send
			// an empty content string.
			name:  "empty content string",
			input: `{"role":"assistant","content":""}`,
			want:  `{"content":[{"text":"","type":"text"}],"role":"assistant"}`,
		},
		{
			// The UI sends this when the user edits a message's
			// text to empty.
			name:  "empty text part",
			input: `{"role":"user","content":[{"type":"text","text":""}]}`,
			want:  `{"content":[{"text":"","type":"text"}],"role":"user"}`,
		},
		{
			// Rows stored before the fix have this form.
			name:  "text part without text key",
			input: `{"role":"user","content":[{"type":"text"}]}`,
			want:  `{"content":[{"text":"","type":"text"}],"role":"user"}`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var message Message
			if err := json.Unmarshal([]byte(tt.input), &message); err != nil {
				t.Fatalf("json.Unmarshal() error = %v", err)
			}
			got, err := json.Marshal(message)
			if err != nil {
				t.Fatalf("json.Marshal() error = %v", err)
			}
			if string(got) != tt.want {
				t.Errorf("json.Marshal() = %s, want %s", got, tt.want)
			}
		})
	}
}

// TestResponseFormatMarshalJSON records the schema-omission policy of
// ResponseFormat. The omitzero option drops only a nil schema. An empty
// JSON schema object is valid and permits any output, so it must be
// sent, not dropped. The omitempty option would drop it.
func TestResponseFormatMarshalJSON(t *testing.T) {
	tests := []struct {
		name   string
		format ResponseFormat
		want   string
	}{
		{
			name:   "nil schema",
			format: ResponseFormat{Type: "json_object"},
			want:   `{"type":"json_object"}`,
		},
		{
			name: "schema object",
			format: ResponseFormat{
				JSONSchema: map[string]any{"name": "calendar"},
				Type:       "json_schema",
			},
			want: `{"json_schema":{"name":"calendar"},"type":"json_schema"}`,
		},
		{
			name: "empty schema object",
			format: ResponseFormat{
				JSONSchema: map[string]any{},
				Type:       "json_schema",
			},
			want: `{"json_schema":{},"type":"json_schema"}`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := json.Marshal(tt.format)
			if err != nil {
				t.Fatalf("json.Marshal() error = %v", err)
			}
			if string(got) != tt.want {
				t.Errorf("json.Marshal() = %s, want %s", got, tt.want)
			}
		})
	}
}