odidere/internal/conversation/title_test.go

package conversation

import (
	"fmt"
	"net/http"
	"strings"
	"testing"

	"code.chimeric.al/chimerical/odidere/internal/db"
	"code.chimeric.al/chimerical/odidere/internal/llm"
)

func TestDeriveTitle(t *testing.T) {
	tests := []struct {
		name string
		text string
		want string
	}{
		{
			name: "single line",
			text: "Check the cameras",
			want: "Check the cameras",
		},
		{
			name: "blank lines come before the first line",
			text: "\n   \n\t\nCheck the cameras\nSecond line",
			want: "Check the cameras",
		},
		{
			name: "spaces inside the line collapse to one",
			text: "Can   you\tcheck the cameras?   ",
			want: "Can you check the cameras?",
		},
		{
			name: "leading indentation is trimmed",
			text: "    def main():",
			want: "def main():",
		},
		{
			name: "a line of MaxTitleLength characters stays whole",
			text: strings.Repeat("x", MaxTitleLength),
			want: strings.Repeat("x", MaxTitleLength),
		},
		{
			name: "a longer line cuts to MaxTitleLength characters",
			text: strings.Repeat("x", 300),
			want: strings.Repeat("x", MaxTitleLength),
		},
		{
			name: "the cut counts characters, not bytes",
			text: strings.Repeat("文", 300),
			want: strings.Repeat("文", MaxTitleLength),
		},
		{
			name: "empty text",
			text: "",
			want: "",
		},
		{
			name: "spaces only",
			text: " \n\t  \n",
			want: "",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := deriveTitle(tt.text); got != tt.want {
				t.Fatalf("deriveTitle(%q) = %q, want %q", tt.text, got, tt.want)
			}
		})
	}
}

func title(t *testing.T, c *Conversation) *string {
	t.Helper()
	conv, err := c.dbr.SelectConversationId(t.Context(), c.id)
	if err != nil {
		t.Fatalf("SelectConversationId: %v", err)
	}
	return conv.Title
}

func TestInsertMessageTitlesConversation(t *testing.T) {
	ctx := t.Context()

	// The first user message sets the title from its first non-empty
	// line.
	c := newTestConversation(t, nil)
	if _, err := c.InsertMessage(
		ctx, *userMessage("  Check the cameras\nplease"),
	); err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}
	if got := title(t, c); got == nil || *got != "Check the cameras" {
		t.Fatalf("Title = %v, want Check the cameras", got)
	}

	// A later user message does not overwrite an existing title.
	if _, err := c.InsertMessage(
		ctx, *userMessage("A different subject"),
	); err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}
	if got := title(t, c); got == nil || *got != "Check the cameras" {
		t.Fatalf("Title = %v, want Check the cameras", got)
	}

	// Non-user messages do not set the title.
	c2 := newTestConversation(t, nil)
	assistant := llm.Message{
		Role: llm.RoleAssistant,
		ContentParts: []llm.ContentPart{{
			Type: llm.ContentTypeText,
			Text: "Hello there",
		}},
	}
	if _, err := c2.InsertMessage(ctx, assistant); err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}
	if got := title(t, c2); got != nil {
		t.Fatalf("Title = %v, want nil", got)
	}

	// A title set before the first user message is kept.
	c3 := newTestConversation(t, nil)
	if _, err := c3.dbw.UpdateConversationTitle(
		ctx,
		db.UpdateConversationTitleParams{
			Id:    c3.id,
			Title: "Nightly report",
		},
	); err != nil {
		t.Fatalf("UpdateConversationTitle: %v", err)
	}
	if _, err := c3.InsertMessage(
		ctx, *userMessage("A different subject"),
	); err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}
	if got := title(t, c3); got == nil || *got != "Nightly report" {
		t.Fatalf("Title = %v, want Nightly report", got)
	}
}

func TestInferenceTitlesConversation(t *testing.T) {
	var (
		ctx    = t.Context()
		client = newTestModelClient(
			t,
			func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "application/json")
				fmt.Fprint(w, `{
					"choices":[{
						"message":{"role":"assistant","content":"ok"},
						"finish_reason":"stop"
					}]
				}`)
			},
		)
	)

	// Inference with a new user message sets an unset title.
	c := newTestConversation(t, nil)
	if err := infer(ctx, c, client, InferenceRequest{
		Message: userMessage("Turn off the porch light"),
	}); err != nil {
		t.Fatalf("infer: %v", err)
	}
	if got := title(t, c); got == nil || *got != "Turn off the porch light" {
		t.Fatalf("Title = %v, want Turn off the porch light", got)
	}

	// Later inference does not overwrite the title.
	if err := infer(ctx, c, client, InferenceRequest{
		Message: userMessage("A different subject"),
	}); err != nil {
		t.Fatalf("infer: %v", err)
	}
	if got := title(t, c); got == nil || *got != "Turn off the porch light" {
		t.Fatalf("Title = %v, want Turn off the porch light", got)
	}

	// Inference does not overwrite a title that was already set.
	c2 := newTestConversation(t, nil)
	if _, err := c2.dbw.UpdateConversationTitle(
		ctx,
		db.UpdateConversationTitleParams{
			Id:    c2.id,
			Title: "Nightly report",
		},
	); err != nil {
		t.Fatalf("UpdateConversationTitle: %v", err)
	}
	if err := infer(ctx, c2, client, InferenceRequest{
		Message: userMessage("A different subject"),
	}); err != nil {
		t.Fatalf("infer: %v", err)
	}
	if got := title(t, c2); got == nil || *got != "Nightly report" {
		t.Fatalf("Title = %v, want Nightly report", got)
	}
}