Fix async thread title generation flow

This commit is contained in:
Administrator committed 2026-06-28 18:09:48 +08:00
1 parent d522498cc9
commit 2d3033434b
6 files changed
+561 -105

No files matched your search

+164 -44
View File
@@ -43,9 +43,13 @@ export type ThreadStreamOptions = {
onToolEnd?: (event: ToolEndEvent) => void; onToolEnd?: (event: ToolEndEvent) => void;
}; };
type SendMessageOptions = { type SendMessageOptions = {
additionalKwargs?: Record<string, unknown>; additionalKwargs?: Record<string, unknown>;
}; };
type ThreadTitleRefreshState = Pick<AgentThreadState, "title"> & {
title_provisional?: boolean;
};
function normalizeStoredRunId(runId: string | null): string | null { function normalizeStoredRunId(runId: string | null): string | null {
if (!runId) { if (!runId) {
@@ -155,7 +159,7 @@ export async function stopThreadRun(
} }
} }
function getStreamErrorMessage(error: unknown): string { function getStreamErrorMessage(error: unknown): string {
if (typeof error === "string" && error.trim()) { if (typeof error === "string" && error.trim()) {
return error; return error;
} }
@@ -175,8 +179,31 @@ function getStreamErrorMessage(error: unknown): string {
return nestedError; return nestedError;
} }
} }
return "Request failed."; return "Request failed.";
} }
async function fetchThreadTitleRefreshState(
threadId: string,
signal?: AbortSignal,
): Promise<ThreadTitleRefreshState | null> {
const response = await apiFetch(
`${getBackendBaseURL()}/api/threads/${encodeURIComponent(threadId)}/state`,
{ signal },
);
if (!response.ok) {
return null;
}
const json = (await response.json()) as {
values?: { title?: string; title_provisional?: boolean };
};
if (!json.values) {
return null;
}
return {
title: json.values.title ?? "",
title_provisional: json.values.title_provisional,
};
}
/** /**
* Detect "thread not found" / forbidden errors from the LangGraph SDK. * Detect "thread not found" / forbidden errors from the LangGraph SDK.
@@ -596,9 +623,12 @@ export function useThreadStream({
}); });
// Optimistic messages shown before the server stream responds // Optimistic messages shown before the server stream responds
const [optimisticMessages, setOptimisticMessages] = useState<Message[]>([]); const [optimisticMessages, setOptimisticMessages] = useState<Message[]>([]);
const [isUploading, setIsUploading] = useState(false); const [isUploading, setIsUploading] = useState(false);
const sendInFlightRef = useRef(false); const [threadValueOverrides, setThreadValueOverrides] = useState<
Partial<ThreadTitleRefreshState>
>({});
const sendInFlightRef = useRef(false);
// Carries the ask time of an in-flight send until the server-assigned human // Carries the ask time of an in-flight send until the server-assigned human
// message id shows up, so we can bind the timestamp to the real id (see the // message id shows up, so we can bind the timestamp to the real id (see the
// reconciliation effect below). ``known`` is the set of human ids that // reconciliation effect below). ``known`` is the set of human ids that
@@ -626,12 +656,13 @@ export function useThreadStream({
// clearing optimistic messages in that case is exactly what made submitted // clearing optimistic messages in that case is exactly what made submitted
// questions visually disappear — the live run that just started would be // questions visually disappear — the live run that just started would be
// aborted by the very effect intended to reset stale state. // aborted by the very effect intended to reset stale state.
useEffect(() => { useEffect(() => {
if (threadId != null) { if (threadId != null) {
// Forward direction: starting or continuing with a real thread. // Forward direction: starting or continuing with a real thread.
// No reset needed; the upstream useStream will load history for this id. // No reset needed; the upstream useStream will load history for this id.
return; setThreadValueOverrides({});
} return;
}
const prevThread = threadStreamRef.current; const prevThread = threadStreamRef.current;
if (prevThread?.isLoading) { if (prevThread?.isLoading) {
try { try {
@@ -643,9 +674,88 @@ export function useThreadStream({
startedRef.current = false; startedRef.current = false;
sendInFlightRef.current = false; sendInFlightRef.current = false;
prevMsgCountRef.current = 0; prevMsgCountRef.current = 0;
setOptimisticMessages([]); setOptimisticMessages([]);
setIsUploading(false); setIsUploading(false);
}, [threadId]); setThreadValueOverrides({});
}, [threadId]);
useEffect(() => {
const activeThreadId = threadIdRef.current ?? onStreamThreadId ?? null;
const effectiveTitleProvisional =
(threadValueOverrides.title_provisional ??
thread.values?.title_provisional) === true;
if (!activeThreadId || thread.isLoading || !effectiveTitleProvisional) {
return;
}
let cancelled = false;
let timeoutId: ReturnType<typeof setTimeout> | null = null;
const controller = new AbortController();
const poll = async () => {
while (!cancelled) {
try {
const nextState = await fetchThreadTitleRefreshState(
activeThreadId,
controller.signal,
);
if (cancelled || !nextState) {
return;
}
setThreadValueOverrides({
title: nextState.title,
title_provisional: nextState.title_provisional,
});
if (nextState.title) {
updateCachedThreadLists(queryClient, (threads) =>
threads.map((t) =>
t.thread_id === activeThreadId
? {
...t,
values: {
...t.values,
title: nextState.title,
title_provisional: nextState.title_provisional,
},
}
: t,
),
);
}
if (nextState.title_provisional !== true) {
void queryClient.invalidateQueries({ queryKey: ["threads", "search"] });
return;
}
} catch {
return;
}
await new Promise<void>((resolve) => {
timeoutId = setTimeout(resolve, 2000);
});
}
};
void poll();
return () => {
cancelled = true;
controller.abort();
if (timeoutId) {
clearTimeout(timeoutId);
}
};
}, [
onStreamThreadId,
queryClient,
thread.isLoading,
thread.values?.title_provisional,
threadValueOverrides.title_provisional,
]);
// Clear optimistic when server messages arrive (count increases) // Clear optimistic when server messages arrive (count increases)
useEffect(() => { useEffect(() => {
@@ -910,13 +1020,22 @@ export function useThreadStream({
); );
// Merge thread with optimistic messages for display // Merge thread with optimistic messages for display
const mergedThread = const mergedValues =
optimisticMessages.length > 0 Object.keys(threadValueOverrides).length > 0
? ({ ? {
...thread, ...thread.values,
messages: [...thread.messages, ...optimisticMessages], ...threadValueOverrides,
} as typeof thread) }
: thread; : thread.values;
const mergedThread =
optimisticMessages.length > 0 || mergedValues !== thread.values
? ({
...thread,
values: mergedValues,
messages: [...thread.messages, ...optimisticMessages],
} as typeof thread)
: thread;
// 第 4 个返回值 `onStreamThreadId` = 底层 useStream 当前实际绑定的线程 id(新建/未保存时为 null)。 // 第 4 个返回值 `onStreamThreadId` = 底层 useStream 当前实际绑定的线程 id(新建/未保存时为 null)。
// 调用方据此判断「流是否已切到全新线程」,避免在切换会话的那一帧把消息发进上一个会话 // 调用方据此判断「流是否已切到全新线程」,避免在切换会话的那一帧把消息发进上一个会话
@@ -1284,28 +1403,29 @@ export function useRenameThread() {
mutationFn: async ({ mutationFn: async ({
threadId, threadId,
title, title,
}: { }: {
threadId: string; threadId: string;
title: string; title: string;
}) => { }) => {
await apiClient.threads.updateState(threadId, { await apiClient.threads.updateState(threadId, {
values: { title }, values: { title, title_provisional: false },
}); });
}, },
onSuccess(_, { threadId, title }) { onSuccess(_, { threadId, title }) {
updateCachedThreadLists(queryClient, (threads) => updateCachedThreadLists(queryClient, (threads) =>
threads.map((t) => threads.map((t) =>
t.thread_id === threadId t.thread_id === threadId
? { ? {
...t, ...t,
values: { values: {
...t.values, ...t.values,
title, title,
}, title_provisional: false,
} },
: t, }
), : t,
); ),
);
}, },
}); });
} }
+7 -6
View File
@@ -2,12 +2,13 @@ import type { Message, Thread } from "@langchain/langgraph-sdk";
import type { Todo } from "../todos"; import type { Todo } from "../todos";
export interface AgentThreadState extends Record<string, unknown> { export interface AgentThreadState extends Record<string, unknown> {
title: string; title: string;
messages: Message[]; title_provisional?: boolean;
artifacts: string[]; messages: Message[];
todos?: Todo[]; artifacts: string[];
} todos?: Todo[];
}
export interface AgentThreadContext extends Record<string, unknown> { export interface AgentThreadContext extends Record<string, unknown> {
thread_id: string; thread_id: string;
@@ -1,21 +1,49 @@
"""Middleware for automatic thread title generation.""" """Middleware for automatic thread title generation."""
import logging import copy
import re import logging
from typing import TYPE_CHECKING, Any, NotRequired, override import re
from pathlib import Path
from langchain.agents import AgentState from typing import TYPE_CHECKING, Any, NotRequired, override
from langchain.agents.middleware import AgentMiddleware
from langgraph.runtime import Runtime from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from deerflow.config.title_config import get_title_config from langgraph.runtime import Runtime
from deerflow.models import create_chat_model
from deerflow.config.title_config import get_title_config
from deerflow.models import create_chat_model
from deerflow.utils.time import now_iso
if TYPE_CHECKING: if TYPE_CHECKING:
from deerflow.config.app_config import AppConfig from deerflow.config.app_config import AppConfig
from deerflow.config.title_config import TitleConfig from deerflow.config.title_config import TitleConfig
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_TITLE_DEBUG_LOGGER_NAME = "title_generation_debug"
def _get_title_debug_logger() -> logging.Logger:
debug_logger = logging.getLogger(_TITLE_DEBUG_LOGGER_NAME)
if getattr(debug_logger, "_title_file_configured", False):
return debug_logger
log_path = Path("logs/title_generation_debug.log")
if not log_path.is_absolute():
log_path = Path.cwd() / log_path
log_path.parent.mkdir(parents=True, exist_ok=True)
handler = logging.FileHandler(log_path, encoding="utf-8")
handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
debug_logger.addHandler(handler)
debug_logger.setLevel(logging.INFO)
debug_logger.propagate = False
setattr(debug_logger, "_title_file_configured", True)
return debug_logger
def log_title_debug(message: str, *args: Any) -> None:
logger.warning(message, *args)
_get_title_debug_logger().info(message, *args)
class TitleMiddlewareState(AgentState): class TitleMiddlewareState(AgentState):
@@ -25,7 +53,7 @@ class TitleMiddlewareState(AgentState):
title_provisional: NotRequired[bool] title_provisional: NotRequired[bool]
class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]): class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
"""Automatically generate a title for the thread after the first user message.""" """Automatically generate a title for the thread after the first user message."""
state_schema = TitleMiddlewareState state_schema = TitleMiddlewareState
@@ -42,9 +70,9 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
return self._app_config.title return self._app_config.title
return get_title_config() return get_title_config()
def _normalize_content(self, content: object) -> str: def _normalize_content(self, content: object) -> str:
if isinstance(content, str): if isinstance(content, str):
return content return content
if isinstance(content, list): if isinstance(content, list):
parts = [self._normalize_content(item) for item in content] parts = [self._normalize_content(item) for item in content]
@@ -58,8 +86,18 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
nested_content = content.get("content") nested_content = content.get("content")
if nested_content is not None: if nested_content is not None:
return self._normalize_content(nested_content) return self._normalize_content(nested_content)
return "" return ""
def _message_type(self, message: object) -> str:
if isinstance(message, dict):
return str(message.get("type") or message.get("role") or "").lower()
return str(getattr(message, "type", "") or getattr(message, "role", "") or "").lower()
def _message_content(self, message: object) -> object:
if isinstance(message, dict):
return message.get("content", "")
return getattr(message, "content", "")
def _should_set_provisional_title(self, state: TitleMiddlewareState) -> bool: def _should_set_provisional_title(self, state: TitleMiddlewareState) -> bool:
"""Check if we should set an instant, question-derived placeholder title. """Check if we should set an instant, question-derived placeholder title.
@@ -78,8 +116,8 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
# Only on the very first turn, as soon as the user message arrives. # Only on the very first turn, as soon as the user message arrives.
messages = state.get("messages", []) messages = state.get("messages", [])
user_messages = [m for m in messages if m.type == "human"] user_messages = [m for m in messages if self._message_type(m) in {"human", "user"}]
return len(user_messages) == 1 return len(user_messages) == 1
def _should_generate_title(self, state: TitleMiddlewareState) -> bool: def _should_generate_title(self, state: TitleMiddlewareState) -> bool:
"""Check if we should generate the final LLM title for this thread.""" """Check if we should generate the final LLM title for this thread."""
@@ -98,8 +136,8 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
return False return False
# Count user and assistant messages # Count user and assistant messages
user_messages = [m for m in messages if m.type == "human"] user_messages = [m for m in messages if self._message_type(m) in {"human", "user"}]
assistant_messages = [m for m in messages if m.type == "ai"] assistant_messages = [m for m in messages if self._message_type(m) in {"ai", "assistant"}]
# Generate title after first complete exchange # Generate title after first complete exchange
return len(user_messages) == 1 and len(assistant_messages) >= 1 return len(user_messages) == 1 and len(assistant_messages) >= 1
@@ -112,8 +150,8 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
config = self._get_title_config() config = self._get_title_config()
messages = state.get("messages", []) messages = state.get("messages", [])
user_msg_content = next((m.content for m in messages if m.type == "human"), "") user_msg_content = next((self._message_content(m) for m in messages if self._message_type(m) in {"human", "user"}), "")
assistant_msg_content = next((m.content for m in messages if m.type == "ai"), "") assistant_msg_content = next((self._message_content(m) for m in messages if self._message_type(m) in {"ai", "assistant"}), "")
user_msg = self._normalize_content(user_msg_content) user_msg = self._normalize_content(user_msg_content)
assistant_msg = self._strip_think_tags(self._normalize_content(assistant_msg_content)) assistant_msg = self._strip_think_tags(self._normalize_content(assistant_msg_content))
@@ -174,7 +212,7 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
return None return None
messages = state.get("messages", []) messages = state.get("messages", [])
user_msg_content = next((m.content for m in messages if m.type == "human"), "") user_msg_content = next((self._message_content(m) for m in messages if self._message_type(m) in {"human", "user"}), "")
user_msg = self._normalize_content(user_msg_content) user_msg = self._normalize_content(user_msg_content)
if not user_msg.strip(): if not user_msg.strip():
return None return None
@@ -261,10 +299,98 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
async def abefore_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None: async def abefore_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
return self._provisional_title_result(state) return self._provisional_title_result(state)
@override @override
def after_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None: def after_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
return self._generate_title_result(state) # Keep the provisional title in place during the main run. The gateway
# worker finalizes it in a background task after the answer is complete.
@override return None
async def aafter_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
return await self._agenerate_title_result(state) @override
async def aafter_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
# Keep the provisional title in place during the main async run so the
# worker's post-run background finalizer can replace it with the real
# LLM-generated title without blocking the input box.
return None
async def finalize_thread_title(
*,
checkpointer: Any,
thread_store: Any | None,
thread_id: str,
source_state: dict[str, Any],
app_config: "AppConfig | None" = None,
) -> bool:
"""Finalize a provisional thread title in the background.
``source_state`` must be the checkpoint snapshot captured when the first run
completed. That avoids races where the user asks a second question before
the background title task starts, which would otherwise make the current
checkpoint fail the "first exchange only" predicate and leave the title
provisional forever.
"""
log_title_debug("Title finalize: evaluating background title generation for thread %s", thread_id)
middleware = TitleMiddleware(app_config=app_config)
source_snapshot = copy.deepcopy(source_state)
if not middleware._should_generate_title(source_snapshot):
log_title_debug(
"Title finalize: skipped for thread %s because source snapshot is not eligible",
thread_id,
)
return False
log_title_debug("Title finalize: starting background title generation for thread %s", thread_id)
result = await middleware._agenerate_title_result(source_snapshot)
if not result:
log_title_debug("Title finalize: no title result produced for thread %s", thread_id)
return False
read_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
checkpoint_tuple = await checkpointer.aget_tuple(read_config)
if checkpoint_tuple is None:
log_title_debug("Title finalize: skipped for thread %s because checkpoint was not found", thread_id)
return False
checkpoint: dict[str, Any] = dict(getattr(checkpoint_tuple, "checkpoint", {}) or {})
metadata: dict[str, Any] = dict(getattr(checkpoint_tuple, "metadata", {}) or {})
channel_values: dict[str, Any] = dict(checkpoint.get("channel_values", {}) or {})
# If another path already finalized or manually renamed the title, keep it.
if channel_values.get("title_provisional") is not True:
log_title_debug(
"Title finalize: skipped for thread %s because title is no longer provisional",
thread_id,
)
return False
channel_values.update(result)
checkpoint["channel_values"] = channel_values
metadata["updated_at"] = now_iso()
write_config = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": "",
}
}
new_versions: dict[str, Any] = {}
get_next_version = getattr(checkpointer, "get_next_version", None)
if callable(get_next_version):
versions = dict(checkpoint.get("channel_versions", {}) or {})
for key in ("title", "title_provisional"):
next_version = get_next_version(versions.get(key), None)
versions[key] = next_version
new_versions[key] = next_version
checkpoint["channel_versions"] = versions
await checkpointer.aput(write_config, checkpoint, metadata, new_versions)
title = result.get("title")
if title and thread_store is not None:
await thread_store.update_display_name(thread_id, title, user_id=None)
log_title_debug("Title finalize: completed for thread %s with title %r", thread_id, title)
return True
@@ -15,10 +15,10 @@ internal checkpoint callbacks that are not exposed in the Python public API.
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import copy import copy
import inspect import inspect
import logging import logging
from dataclasses import dataclass, field from dataclasses import dataclass, field
from functools import lru_cache from functools import lru_cache
from typing import TYPE_CHECKING, Any, Literal, cast from typing import TYPE_CHECKING, Any, Literal, cast
@@ -28,17 +28,31 @@ from langgraph.checkpoint.base import empty_checkpoint
if TYPE_CHECKING: if TYPE_CHECKING:
from langchain_core.messages import HumanMessage from langchain_core.messages import HumanMessage
from deerflow.config.app_config import AppConfig from deerflow.agents.middlewares.title_middleware import (
from deerflow.runtime.serialization import serialize finalize_thread_title,
from deerflow.runtime.stream_bridge import StreamBridge log_title_debug,
)
from deerflow.config.app_config import AppConfig
from deerflow.runtime.serialization import serialize
from deerflow.runtime.stream_bridge import StreamBridge
from .manager import RunManager, RunRecord from .manager import RunManager, RunRecord
from .schemas import RunStatus from .schemas import RunStatus
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Valid stream_mode values for LangGraph's graph.astream() # Valid stream_mode values for LangGraph's graph.astream()
_VALID_LG_MODES = {"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"} _VALID_LG_MODES = {"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"}
def _log_background_title_finalize_failure(thread_id: str):
def _done(task: asyncio.Task[bool]) -> None:
try:
task.result()
except Exception:
logger.debug("Failed to finalize background title for thread %s", thread_id, exc_info=True)
return _done
def _build_runtime_context( def _build_runtime_context(
@@ -400,18 +414,47 @@ async def run_agent(
except Exception: except Exception:
logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True) logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True)
# Sync title from checkpoint to threads_meta.display_name # Sync title from checkpoint to threads_meta.display_name
if checkpointer is not None and thread_store is not None: if checkpointer is not None and thread_store is not None:
try: try:
ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
ckpt_tuple = await checkpointer.aget_tuple(ckpt_config) ckpt_tuple = await checkpointer.aget_tuple(ckpt_config)
if ckpt_tuple is not None: if ckpt_tuple is not None:
ckpt = getattr(ckpt_tuple, "checkpoint", {}) or {} ckpt = getattr(ckpt_tuple, "checkpoint", {}) or {}
title = ckpt.get("channel_values", {}).get("title") channel_values = ckpt.get("channel_values", {}) or {}
if title: title = channel_values.get("title")
await thread_store.update_display_name(thread_id, title) if title:
except Exception: await thread_store.update_display_name(thread_id, title)
logger.debug("Failed to sync title for thread %s (non-fatal)", thread_id) if (
record.status == RunStatus.success
and channel_values.get("title_provisional") is True
):
log_title_debug(
"Run %s: scheduling background title generation for thread %s",
run_id,
thread_id,
)
title_task = asyncio.create_task(
finalize_thread_title(
checkpointer=checkpointer,
thread_store=thread_store,
thread_id=thread_id,
source_state=copy.deepcopy(channel_values),
app_config=ctx.app_config,
),
name=f"title-finalize:{thread_id}",
)
title_task.add_done_callback(_log_background_title_finalize_failure(thread_id))
elif record.status == RunStatus.success:
log_title_debug(
"Run %s: background title generation not scheduled for thread %s because title_provisional=%r title=%r",
run_id,
thread_id,
channel_values.get("title_provisional"),
channel_values.get("title"),
)
except Exception:
logger.debug("Failed to sync title for thread %s (non-fatal)", thread_id)
# Update threads_meta status based on run outcome # Update threads_meta status based on run outcome
if thread_store is not None: if thread_store is not None:
@@ -22,10 +22,10 @@ from langgraph.checkpoint.base import empty_checkpoint
from deerflow.config.agents_config import AGENT_ID_PATTERN, load_agent_config, load_agent_soul from deerflow.config.agents_config import AGENT_ID_PATTERN, load_agent_config, load_agent_soul
from deerflow.config.paths import get_paths from deerflow.config.paths import get_paths
from deerflow.persistence.scheduled_tasks import ScheduledTaskStore from deerflow.persistence.scheduled_tasks import ScheduledTaskStore
from deerflow.runtime import serialize_channel_values
from deerflow.runtime.scheduler import html_page as html_page_utils from deerflow.runtime.scheduler import html_page as html_page_utils
from deerflow.runtime.scheduler import template_fill from deerflow.runtime.scheduler import template_fill
from deerflow.runtime.scheduler.element import ElementDeliveryClient from deerflow.runtime.scheduler.element import ElementDeliveryClient
from deerflow.runtime.serialization import serialize_channel_values
from deerflow.runtime.user_context import reset_current_user, set_current_user from deerflow.runtime.user_context import reset_current_user, set_current_user
from deerflow.utils.time import now_iso from deerflow.utils.time import now_iso
@@ -0,0 +1,166 @@
from __future__ import annotations
from types import SimpleNamespace
import pytest
from deerflow.agents.middlewares.title_middleware import (
TitleMiddleware,
finalize_thread_title,
)
class _FakeCheckpointer:
def __init__(self, checkpoint: dict, metadata: dict | None = None):
self._tuple = SimpleNamespace(
checkpoint=checkpoint,
metadata=metadata or {},
config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
)
self.aput_calls: list[tuple[dict, dict, dict, dict]] = []
self._version = 0
async def aget_tuple(self, _config):
return self._tuple
def get_next_version(self, _current, _channel):
self._version += 1
return self._version
async def aput(self, config, checkpoint, metadata, versions):
self.aput_calls.append((config, checkpoint, metadata, versions))
self._tuple = SimpleNamespace(
checkpoint=checkpoint,
metadata=metadata,
config={
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"checkpoint_ns": config["configurable"]["checkpoint_ns"],
"checkpoint_id": "new-ckpt",
}
},
)
return self._tuple.config
class _FakeThreadStore:
def __init__(self):
self.updated: list[tuple[str, str, str | None]] = []
async def update_display_name(self, thread_id: str, title: str, *, user_id=None):
self.updated.append((thread_id, title, user_id))
def _source_state() -> dict:
return {
"messages": [
SimpleNamespace(type="human", content="please analyze the first topic"),
SimpleNamespace(type="ai", content="this is the first answer"),
],
"title": "please analyze the first topic",
"title_provisional": True,
}
def _dict_source_state() -> dict:
return {
"messages": [
{"type": "human", "content": "please analyze the first topic"},
{"type": "ai", "content": "this is the first answer"},
],
"title": "please analyze the first topic",
"title_provisional": True,
}
@pytest.mark.asyncio
async def test_finalize_thread_title_updates_checkpoint_and_display_name(monkeypatch) -> None:
async def _fake_generate(self, state):
assert state["title_provisional"] is True
return {"title": "first topic analysis", "title_provisional": False}
monkeypatch.setattr(TitleMiddleware, "_agenerate_title_result", _fake_generate)
checkpointer = _FakeCheckpointer({"channel_values": _source_state()})
thread_store = _FakeThreadStore()
updated = await finalize_thread_title(
checkpointer=checkpointer,
thread_store=thread_store,
thread_id="thread-1",
source_state=_source_state(),
)
assert updated is True
assert len(checkpointer.aput_calls) == 1
_, written_checkpoint, _, new_versions = checkpointer.aput_calls[0]
assert written_checkpoint["channel_values"]["title"] == "first topic analysis"
assert written_checkpoint["channel_values"]["title_provisional"] is False
assert set(new_versions) == {"title", "title_provisional"}
assert thread_store.updated == [("thread-1", "first topic analysis", None)]
@pytest.mark.asyncio
async def test_finalize_thread_title_skips_if_current_thread_already_finalized(monkeypatch) -> None:
async def _fake_generate(self, state):
return {"title": "first topic analysis", "title_provisional": False}
monkeypatch.setattr(TitleMiddleware, "_agenerate_title_result", _fake_generate)
checkpointer = _FakeCheckpointer(
{
"channel_values": {
**_source_state(),
"title": "manually renamed title",
"title_provisional": False,
}
}
)
thread_store = _FakeThreadStore()
updated = await finalize_thread_title(
checkpointer=checkpointer,
thread_store=thread_store,
thread_id="thread-1",
source_state=_source_state(),
)
assert updated is False
assert checkpointer.aput_calls == []
assert thread_store.updated == []
def test_sync_after_model_keeps_provisional_title_for_background_finalize() -> None:
middleware = TitleMiddleware()
result = middleware.after_model(_source_state(), runtime=SimpleNamespace())
assert result is None
@pytest.mark.asyncio
async def test_async_after_model_keeps_provisional_title_for_background_finalize() -> None:
middleware = TitleMiddleware()
result = await middleware.aafter_model(_source_state(), runtime=SimpleNamespace())
assert result is None
@pytest.mark.asyncio
async def test_finalize_thread_title_accepts_serialized_dict_messages(monkeypatch) -> None:
async def _fake_generate(self, state):
return {"title": "first topic analysis", "title_provisional": False}
monkeypatch.setattr(TitleMiddleware, "_agenerate_title_result", _fake_generate)
checkpointer = _FakeCheckpointer({"channel_values": _dict_source_state()})
updated = await finalize_thread_title(
checkpointer=checkpointer,
thread_store=None,
thread_id="thread-1",
source_state=_dict_source_state(),
)
assert updated is True