397 lines
17 KiB
Python
397 lines
17 KiB
Python
"""Middleware for automatic thread title generation."""
|
|
|
|
import copy
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, NotRequired, override
|
|
|
|
from langchain.agents import AgentState
|
|
from langchain.agents.middleware import AgentMiddleware
|
|
from langgraph.runtime import Runtime
|
|
|
|
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:
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.config.title_config import TitleConfig
|
|
|
|
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):
|
|
"""Compatible with the `ThreadState` schema."""
|
|
|
|
title: NotRequired[str | None]
|
|
title_provisional: NotRequired[bool]
|
|
|
|
|
|
class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
|
|
"""Automatically generate a title for the thread after the first user message."""
|
|
|
|
state_schema = TitleMiddlewareState
|
|
|
|
def __init__(self, *, app_config: "AppConfig | None" = None, title_config: "TitleConfig | None" = None):
|
|
super().__init__()
|
|
self._app_config = app_config
|
|
self._title_config = title_config
|
|
|
|
def _get_title_config(self):
|
|
if self._title_config is not None:
|
|
return self._title_config
|
|
if self._app_config is not None:
|
|
return self._app_config.title
|
|
return get_title_config()
|
|
|
|
def _normalize_content(self, content: object) -> str:
|
|
if isinstance(content, str):
|
|
return content
|
|
|
|
if isinstance(content, list):
|
|
parts = [self._normalize_content(item) for item in content]
|
|
return "\n".join(part for part in parts if part)
|
|
|
|
if isinstance(content, dict):
|
|
text_value = content.get("text")
|
|
if isinstance(text_value, str):
|
|
return text_value
|
|
|
|
nested_content = content.get("content")
|
|
if nested_content is not None:
|
|
return self._normalize_content(nested_content)
|
|
|
|
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:
|
|
"""Check if we should set an instant, question-derived placeholder title.
|
|
|
|
Runs in ``before_model`` so the thread shows a sensible title the moment
|
|
the user sends their first message — eliminating the "untitled" window
|
|
while the (possibly long) first answer streams. No LLM call, no latency.
|
|
"""
|
|
config = self._get_title_config()
|
|
if not config.enabled:
|
|
return False
|
|
|
|
# Already has a title (provisional or final) — nothing to do.
|
|
if state.get("title"):
|
|
return False
|
|
|
|
# Only on the very first turn, as soon as the user message arrives.
|
|
messages = state.get("messages", [])
|
|
user_messages = [m for m in messages if self._message_type(m) in {"human", "user"}]
|
|
return len(user_messages) == 1
|
|
|
|
def _should_generate_title(self, state: TitleMiddlewareState) -> bool:
|
|
"""Check if we should generate the final LLM title for this thread."""
|
|
config = self._get_title_config()
|
|
if not config.enabled:
|
|
return False
|
|
|
|
# Skip only if we already have a *finalized* title. A provisional
|
|
# placeholder set in before_model should still be upgraded by the LLM.
|
|
if state.get("title") and not state.get("title_provisional"):
|
|
return False
|
|
|
|
# Check if this is the first turn (has at least one user message and one assistant response)
|
|
messages = state.get("messages", [])
|
|
if len(messages) < 2:
|
|
return False
|
|
|
|
# Count user and assistant messages
|
|
user_messages = [m for m in messages if self._message_type(m) in {"human", "user"}]
|
|
assistant_messages = [m for m in messages if self._message_type(m) in {"ai", "assistant"}]
|
|
|
|
# Generate title after first complete exchange
|
|
return len(user_messages) == 1 and len(assistant_messages) >= 1
|
|
|
|
def _build_title_prompt(self, state: TitleMiddlewareState) -> tuple[str, str]:
|
|
"""Extract user/assistant messages and build the title prompt.
|
|
|
|
Returns (prompt_string, user_msg) so callers can use user_msg as fallback.
|
|
"""
|
|
config = self._get_title_config()
|
|
messages = state.get("messages", [])
|
|
|
|
user_msg_content = next((self._message_content(m) for m in messages if self._message_type(m) in {"human", "user"}), "")
|
|
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)
|
|
assistant_msg = self._strip_think_tags(self._normalize_content(assistant_msg_content))
|
|
|
|
prompt = config.prompt_template.format(
|
|
max_words=config.max_words,
|
|
user_msg=user_msg[:500],
|
|
assistant_msg=assistant_msg[:500],
|
|
)
|
|
return prompt, user_msg
|
|
|
|
def _strip_think_tags(self, text: str) -> str:
|
|
"""Remove <think>...</think> blocks emitted by reasoning models (e.g. minimax, GLM, DeepSeek-R1).
|
|
|
|
Also handles malformed cases: an unclosed `<think>` (response truncated mid-reasoning)
|
|
or a stray `</think>` (open tag eaten by upstream parsing) — both can otherwise leak
|
|
raw reasoning into the title.
|
|
"""
|
|
text = re.sub(r"<think>[\s\S]*?</think>", "", text, flags=re.IGNORECASE)
|
|
text = re.sub(r"<think>[\s\S]*$", "", text, flags=re.IGNORECASE)
|
|
text = re.sub(r"^[\s\S]*?</think>", "", text, flags=re.IGNORECASE)
|
|
return text.strip()
|
|
|
|
def _parse_title(self, content: object) -> str:
|
|
"""Normalize model output into a clean title string (strips <think> blocks)."""
|
|
config = self._get_title_config()
|
|
title_content = self._strip_think_tags(self._normalize_content(content))
|
|
title = title_content.strip().strip('"').strip("'")
|
|
return title[: config.max_chars] if len(title) > config.max_chars else title
|
|
|
|
def _isolated_config(self) -> dict[str, Any]:
|
|
"""Build a RunnableConfig that does NOT inherit the parent's streaming callbacks.
|
|
|
|
The leak that pushed title-generation reasoning into the live UI stream came
|
|
from inheriting the parent RunnableConfig (and with it the SSE/streaming
|
|
callbacks) for the title model call. Here we deliberately start from an empty
|
|
config with ``callbacks: []`` so the title model's tokens — including any
|
|
``<think>`` reasoning from a reasoning-only deployment — never reach the
|
|
frontend stream or the run journal as conversation output. The stored title
|
|
is additionally cleaned by ``_parse_title``/``_strip_think_tags``.
|
|
"""
|
|
return {
|
|
"run_name": "title_agent",
|
|
"tags": ["middleware:title"],
|
|
"callbacks": [],
|
|
}
|
|
|
|
def _fallback_title(self, user_msg: str) -> str:
|
|
config = self._get_title_config()
|
|
fallback_chars = min(config.max_chars, 50)
|
|
if len(user_msg) > fallback_chars:
|
|
return user_msg[:fallback_chars].rstrip() + "..."
|
|
return user_msg if user_msg else "New Conversation"
|
|
|
|
def _provisional_title_result(self, state: TitleMiddlewareState) -> dict | None:
|
|
"""Set an instant placeholder title from the first user message (no LLM)."""
|
|
if not self._should_set_provisional_title(state):
|
|
return None
|
|
|
|
messages = state.get("messages", [])
|
|
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)
|
|
if not user_msg.strip():
|
|
return None
|
|
title = self._fallback_title(user_msg)
|
|
logger.debug("TitleMiddleware: set provisional title %r before first model call", title)
|
|
return {"title": title, "title_provisional": True}
|
|
|
|
def _generate_title_result(self, state: TitleMiddlewareState) -> dict | None:
|
|
"""Generate a local fallback title without blocking on an LLM call."""
|
|
if not self._should_generate_title(state):
|
|
return None
|
|
|
|
_, user_msg = self._build_title_prompt(state)
|
|
return {"title": self._fallback_title(user_msg), "title_provisional": False}
|
|
|
|
def _force_disable_thinking(self, model: object) -> None:
|
|
"""Force-disable thinking on the title model via `extra_body` overrides.
|
|
|
|
Title generation never needs reasoning. The model factory only emits
|
|
disable params when the model config declares `when_thinking_disabled`
|
|
or a matching `when_thinking_enabled` shape — many user-added reasoning
|
|
models (e.g. GLM-4.5/5) only set `supports_thinking: true` and would
|
|
otherwise keep thinking on, which is slow and leaks <think> tags into
|
|
the title prompt response. Injecting both common OpenAI-compatible
|
|
disable shapes covers Zhipu GLM (`thinking.type`) and Qwen/vLLM
|
|
(`chat_template_kwargs.enable_thinking`); other providers ignore
|
|
unknown body fields.
|
|
"""
|
|
if not hasattr(model, "extra_body"):
|
|
return
|
|
current = getattr(model, "extra_body", None) or {}
|
|
merged_chat_template_kwargs = {
|
|
**(current.get("chat_template_kwargs") or {}),
|
|
"enable_thinking": False,
|
|
}
|
|
merged = {
|
|
**current,
|
|
"thinking": {**(current.get("thinking") or {}), "type": "disabled"},
|
|
"chat_template_kwargs": merged_chat_template_kwargs,
|
|
}
|
|
try:
|
|
model.extra_body = merged
|
|
except Exception:
|
|
logger.debug("Could not override extra_body to disable thinking on title model", exc_info=True)
|
|
|
|
async def _agenerate_title_result(self, state: TitleMiddlewareState) -> dict | None:
|
|
"""Polish the title with an LLM, without leaking its reasoning to the UI.
|
|
|
|
The title model is invoked with an *isolated* config (``_isolated_config``)
|
|
that does not inherit the parent's streaming callbacks, so even a
|
|
reasoning-only deployment's ``<think>`` tokens never reach the live UI
|
|
stream or the run journal. Thinking is also force-disabled on the model,
|
|
and the final stored title is stripped of any residual ``<think>`` blocks.
|
|
On any failure we fall back to the local question-derived title.
|
|
"""
|
|
if not self._should_generate_title(state):
|
|
return None
|
|
|
|
config = self._get_title_config()
|
|
prompt, user_msg = self._build_title_prompt(state)
|
|
|
|
try:
|
|
model_kwargs: dict[str, Any] = {"thinking_enabled": False, "force_disable_thinking": True}
|
|
if self._app_config is not None:
|
|
model_kwargs["app_config"] = self._app_config
|
|
if config.model_name:
|
|
model = create_chat_model(name=config.model_name, **model_kwargs)
|
|
else:
|
|
model = create_chat_model(**model_kwargs)
|
|
self._force_disable_thinking(model)
|
|
response = await model.ainvoke(prompt, config=self._isolated_config())
|
|
title = self._parse_title(response.content)
|
|
if title:
|
|
return {"title": title, "title_provisional": False}
|
|
except Exception:
|
|
logger.debug("Failed to generate async title via LLM; falling back to local title", exc_info=True)
|
|
return {"title": self._fallback_title(user_msg), "title_provisional": False}
|
|
|
|
@override
|
|
def before_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
|
|
return self._provisional_title_result(state)
|
|
|
|
@override
|
|
async def abefore_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
|
|
return self._provisional_title_result(state)
|
|
|
|
@override
|
|
def after_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
|
|
# Keep the provisional title in place during the main run. The gateway
|
|
# worker finalizes it in a background task after the answer is complete.
|
|
return None
|
|
|
|
@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
|