diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index fbaa4400ed..360e081f0f 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -427,26 +427,32 @@ function extractTextParts(m: ThreadMessage | undefined): string { async function generateTitleWithModel(payload: { userText: string; + assistantText?: string; }): Promise { const params = useChatRuntimeStore.getState().params; if (!params.checkpoint) return null; const user = clip(payload.userText, 256); - const parts: string[] = [user]; + const assistant = clip(payload.assistantText ?? "", 384); + const parts: string[] = [`User: ${user}`]; + if (assistant) { + parts.push(`Assistant: ${assistant}`); + } function normalizeTitle(raw: string): string | null { let title = raw.split(/\r?\n/, 1)[0] ?? ""; title = title.replace(/^\s*title\s*:\s*/i, ""); title = title.replace(/[^\x20-\x7E]+/g, " "); title = title.replace(/["'`]+/g, ""); - title = title.replace(/[.!?:;,]+/g, " "); - title = title.replace(/\s+/g, " ").trim(); - // Model echo fail-safe. - if (/\b(user|base|lora|assistant)\s*:/i.test(title)) { + // Echo fail-safe: reject leading role labels before punctuation strips the ":". + if (/^\s*(user|assistant|base|lora)\s*:/i.test(title)) { return null; } + title = title.replace(/[.!?:;,]+/g, " "); + title = title.replace(/\s+/g, " ").trim(); + const words = title.split(" ").filter(Boolean).slice(0, 6); const joined = words.join(" ").trim(); if (!joined) return null; @@ -468,7 +474,7 @@ async function generateTitleWithModel(payload: { { role: "system", content: - "Write 1 concise chat title for the user's message. Rules: 2-6 words, no quotes, no punctuation, ASCII only, do not echo input. Output title only.", + "Write 1 concise chat title summarizing the conversation topic, not the user's exact wording. Use the assistant reply as context when provided. Rules: 2-6 words, no quotes, no punctuation, ASCII only, do not echo input. Output title only.", }, { role: "user", content: parts.join("\n") }, ], @@ -730,8 +736,15 @@ function createStudioDbAdapter( return streamTitle(thread.title); } - const firstUser = messages.find((m) => m.role === "user"); + const firstUserIndex = messages.findIndex((m) => m.role === "user"); + const firstUser = + firstUserIndex === -1 ? undefined : messages[firstUserIndex]; + const firstAssistant = + firstUserIndex === -1 + ? undefined + : messages.find((m, i) => m.role === "assistant" && i > firstUserIndex); const userText = extractTextParts(firstUser) || defaultTitle; + const assistantText = extractTextParts(firstAssistant); if (!autoTitle) { const title = fallbackTitleFromUserText(userText); @@ -769,6 +782,7 @@ function createStudioDbAdapter( const title = (await generateTitleWithModel({ userText, + assistantText, })) || fallbackTitleFromUserText(userText); await persistTitle(title); diff --git a/tests/studio/test_chat_title_generation.py b/tests/studio/test_chat_title_generation.py new file mode 100644 index 0000000000..4f82b4ec03 --- /dev/null +++ b/tests/studio/test_chat_title_generation.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression checks for Studio chat title generation context.""" + +from __future__ import annotations + +from pathlib import Path + + +REPO = Path(__file__).resolve().parents[2] +RUNTIME_TSX = REPO / "studio/frontend/src/features/chat/runtime-provider.tsx" + + +def _source_until(src: str, anchor: str, end_anchor: str) -> str: + start = src.find(anchor) + assert start != -1, f"anchor {anchor!r} not found" + end = src.find(end_anchor, start) + assert end != -1, f"end anchor {end_anchor!r} not found" + return src[start:end] + + +def _balanced_block(src: str, anchor: str) -> str: + # Brace-counting only; assumes no unbalanced braces in strings, regexes, or comments. + start = src.find(anchor) + assert start != -1, f"anchor {anchor!r} not found" + body_start = src.find("{", start) + assert body_start != -1, f"body opener after {anchor!r} not found" + + depth = 0 + for index in range(body_start, len(src)): + char = src[index] + if char == "{": + depth += 1 + elif char == "}": + depth -= 1 + if depth == 0: + return src[start : index + 1] + raise AssertionError(f"unbalanced block after {anchor!r}") + + +def test_title_model_prompt_targets_conversation_topic(): + block = _source_until( + RUNTIME_TSX.read_text(), + "async function generateTitleWithModel", + "\nconst inflightTitleByKey", + ) + + assert "conversation topic" in block + assert "not the user's exact wording" in block + assert "Use the assistant reply as context when provided" in block + assert "Rules: 2-6 words" in block + + +def test_title_model_payload_includes_optional_assistant_reply(): + block = _source_until( + RUNTIME_TSX.read_text(), + "async function generateTitleWithModel", + "\nconst inflightTitleByKey", + ) + + assert "assistantText?: string;" in block + assert 'const assistant = clip(payload.assistantText ?? "", 384);' in block + assert "const parts: string[] = [`User: ${user}`];" in block + assert "if (assistant)" in block + assert "parts.push(`Assistant: ${assistant}`);" in block + assert 'parts.join("\\n")' in block + + +def test_generate_title_passes_first_assistant_reply_after_first_user(): + block = _balanced_block( + RUNTIME_TSX.read_text(), + "async generateTitle(remoteId", + ) + + assert 'const firstUserIndex = messages.findIndex((m) => m.role === "user");' in block + assert '.find((m, i) => m.role === "assistant" && i > firstUserIndex)' in block + assert "const assistantText = extractTextParts(firstAssistant);" in block + assert "generateTitleWithModel({" in block + assert "userText," in block + assert "assistantText," in block + + +def test_auto_title_disabled_uses_deterministic_user_text_fallback(): + block = _balanced_block( + RUNTIME_TSX.read_text(), + "async generateTitle(remoteId", + ) + auto_title_off = _balanced_block(block, "if (!autoTitle)") + + assert "fallbackTitleFromUserText(userText)" in auto_title_off + assert "generateTitleWithModel" not in auto_title_off + + +def test_model_failure_still_falls_back_to_user_text(): + block = _balanced_block( + RUNTIME_TSX.read_text(), + "async generateTitle(remoteId", + ) + + assert "})) || fallbackTitleFromUserText(userText);" in block + + +def test_title_normalizer_still_enforces_output_constraints(): + block = _source_until( + RUNTIME_TSX.read_text(), + "async function generateTitleWithModel", + "\nconst inflightTitleByKey", + ) + + assert r'replace(/[^\x20-\x7E]+/g, " ")' in block + assert 'replace(/["\'`]+/g, "")' in block + assert 'replace(/[.!?:;,]+/g, " ")' in block + assert 'title.split(" ").filter(Boolean).slice(0, 6)' in block + assert "joined.length > 60" in block