diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f259668..22c89777 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to - ✨(back) add model fallback mechanism - ✨(back) add celery for running background tasks - 🧱(helm) add celery worker and beat deployments +- ✨(conversation) summarize messages ### Changed diff --git a/docs/attachments.md b/docs/attachments.md index ee92fd8b..e01e5e98 100644 --- a/docs/attachments.md +++ b/docs/attachments.md @@ -517,11 +517,12 @@ Notes: The decision of which documents are inlined as `full-context` vs left as `tool_call_only` is made by `chat/document_context_builder.py:build_documents_listing` on each turn (called via `_build_document_context_instruction` in `chat/clients/pydantic_ai.py`): -1. Compute the `document_budget` in tokens: +1. Compute budgets in tokens (`chat/tokens.py` splits `max_token_context` by ratio, then subtracts the security buffer from **each** share independently): ```text - document_budget = max(int(model.max_token_context * DOCUMENT_CONTEXT_BUDGET_RATIO) - - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS, 0) + document_budget = max(int(max_token_context * DOCUMENT_CONTEXT_BUDGET_RATIO) - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS, 0) + message_token_budget = max(int(max_token_context * (1 - DOCUMENT_CONTEXT_BUDGET_RATIO)) - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS, 0) ``` + The document share (`document_budget`, via `compute_document_budget`) inlines conversation documents; the history share (`message_token_budget`, via `compute_conversation_budget`) is the conversation-summarization trigger. `build_documents_listing` receives the raw `max_token_context`, `budget_ratio`, and `security_buffer_tokens` and applies the formula above. 2. Load all text attachments from object storage **in parallel** (`asyncio.gather`). Attachments that fail to load are marked `tool_call_only` with their failure logged; other documents are not affected. 3. Iterate documents oldest-first (`order_by("created_at", "id")`). For each document: - If its token count exceeds the whole budget alone → keep `tool_call_only`. @@ -536,6 +537,18 @@ Token estimation uses `tiktoken` with the `cl100k_base` encoding (GPT-4 tokenize The assembled listing is **cached** per turn (in `_build_documents_listing`, `pydantic_ai.py`) keyed on: `conversation_id`, `user_id`, `model_hrid`, `model.max_token_context`, `DOCUMENT_CONTEXT_BUDGET_RATIO`, `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS`, and a fingerprint of `(attachment.id, attachment.updated_at)` for every text attachment — **conversation and project text attachments both contribute to the fingerprint**. Any attachment add / remove / edit (including project files), or any settings change, invalidates the cache. TTL is 30 minutes (`CACHE_TIMEOUT`). +#### Conversation history summarization + +When `message_token_budget` is exceeded, `chat/agents/history_processors.py` calls a separate summarization model (`LLM_SUMMARIZATION_MODEL_HRID`) and stores the result on the conversation (`history_summary`, `history_summary_checkpoint`). This runs at the **start of a new user turn**, before `agent.iter`, against **stored** `pydantic_messages` from previous turns only (the current user prompt is extracted separately and, by design, is not counted in the budget check — see ADR 0002). That stored history usually ends on an assistant `ModelResponse`. + +1. **Trigger**: estimated tokens in the active history slice exceed `message_token_budget` (see formulas above) and there are new messages after the last checkpoint. +2. **After a summary**: the model receives the stored summary text (dynamic instruction) plus the last `CONVERSATION_SUMMARY_CONTEXT_MESSAGES` `ModelMessage` entries before the checkpoint. Use an **even** value so the retained window starts on a user `ModelRequest` in a plain user/assistant alternation (tool messages can break parity). +3. **Summary length**: capped by `CONVERSATION_SUMMARY_MAX_TOKENS` on the summarization LLM call. +4. **Disable summarization**: zero the `message_token_budget`, by setting `max_token_context=0` on the chat model in `LLM_CONFIGURATIONS`, `DOCUMENT_CONTEXT_BUDGET_RATIO=1` (reallocates the whole context to the document share), or a `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS` large enough to consume the history share. Note: *omitting* `max_token_context` does **not** disable summarization — unlike document inlining, `resolve_conversation_budget` falls back to `DEFAULT_MAX_TOKEN_CONTEXT`, so the history budget stays non-zero. +5. **Visible blocking generation**: at the start of a turn whose active history (previous turns only; the new user message is not counted — see ADR 0002) exceeds the history budget, the turn enqueues a Celery task (`chat/tasks.py:summarize_conversation_history`) and waits for it on the open SSE stream, emitting `keep_alive` data parts; the client shows an in-thread progress indicator anchored on the `summarize` tool events. Concurrency is controlled by a time-bounded claim (`ChatConversation.history_summary_claimed_at`, TTL = task hard time limit 120s + 60s): if the task never claims within a 10s grace window or its claim dies and no summary has landed, the turn fails with a `summarization_failed` stream error — there is no inline fallback; generation stays Celery-only (3 task retries with exponential backoff). A turn is never answered with truncated or over-budget context — see `docs/adr/0002-summarization-visible-blocking-phase.md` and ADR 0001. + +The security buffer is **not** a dedicated reserve for system prompts, tool schemas, or completion tokens; those are added on top of the planned document/history split. Size the buffer (and/or plan `max_token_context` below the model nominal window) to leave headroom for that overhead. + #### Targeted document operations (`document_id`) Three tools accept an optional `document_id` argument, each with its own IDOR boundary: @@ -614,8 +627,11 @@ A `READY` attachment whose `rag_document_id` is null (e.g. parse succeeded but t | `RAG_DOCUMENT_SEARCH_BACKEND` | `AlbertRagBackend` | Import path of the vector-search backend used for indexing and search (Albert or Find) | | `PROJECT_FILES_MAX_COUNT` | `10` | Max non-image attachments per project (excludes hidden markdown companions). Enforced at upload-time in `ChatProjectAttachmentViewSet`. Bounds per-turn system-prompt token cost (every entry contributes to `project_documents` on every conversation turn). | | `PROJECT_IMAGES_MAX_COUNT` | `3` | Max image attachments per project. Enforced at upload-time. Bounds per-turn vision token cost - every project image is pinned to every turn alongside conversation-message images, and provider request-level image caps (Anthropic ~20/request) clip the trailing entries first. | -| `DOCUMENT_CONTEXT_BUDGET_RATIO` | `0.5` | Fraction of `model.max_token_context` reserved for inlined documents (0 disables full-context inlining; everything stays `tool_call_only`) | -| `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS` | `1000` | Tokens subtracted from the inlining budget to absorb tokenizer drift on non-OpenAI models | +| `DOCUMENT_CONTEXT_BUDGET_RATIO` | `0.5` | Fraction of `model.max_token_context` reserved for inlined documents; `(1 - ratio)` goes to the conversation history budget (summarization trigger). The security buffer is subtracted from each share afterwards. `0` disables full-context inlining (everything stays `tool_call_only`). | +| `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS` | `1000` | Tokens subtracted from **each** of the document and history budgets (independently), to absorb tokenizer drift on non-OpenAI models and leave headroom beyond the planned split | +| `CONVERSATION_SUMMARY_CONTEXT_MESSAGES` | `10` | Number of pydantic-ai `ModelMessage` entries kept in runtime history after a conversation summary (in addition to the stored summary text). Use an even value (see [Conversation history summarization](#conversation-history-summarization)) | +| `CONVERSATION_SUMMARY_MAX_TOKENS` | `2048` | `max_tokens` for the summarization LLM call when generating or updating `history_summary` | +| `LLM_SUMMARIZATION_MODEL_HRID` | `default-summarization-model` | HRID in `LLM_CONFIGURATIONS` for the conversation summarization agent (see [LLM Configuration](llm-configuration.md)) | | `RAG_COLLECTION_INACTIVITY_DAYS` | `30` | Conversations inactive for this many days have their RAG collection de-indexed by `deindex_inactive_collections`. Resets on re-index. | #### RAG_FILES_ACCEPTED_FORMATS @@ -664,7 +680,7 @@ RAG_FILES_ACCEPTED_FORMATS = [ ### Per-model setting: `max_token_context` -Each entry in `LLM_CONFIGURATIONS` accepts a `max_token_context` integer field declaring the model's context window size. When set, it drives the inlining budget for the documents listing (`document_budget = max_token_context * DOCUMENT_CONTEXT_BUDGET_RATIO - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS`). +Each entry in `LLM_CONFIGURATIONS` accepts a `max_token_context` integer field declaring the model's context window size. When set, it drives document inlining and conversation summarization budgets: `max_token_context` is split by `DOCUMENT_CONTEXT_BUDGET_RATIO`, and `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS` is subtracted from **each** share independently (see the formulas under [Inlining policy and FIFO eviction](#inlining-policy-and-fifo-eviction)). If a model has no `max_token_context`, all of its documents are kept `tool_call_only` regardless of size and a warning is logged on every chat turn. Setting the field accurately matters: too low and small documents get pushed to RAG-only when they could be inlined; too high and the LLM may exceed its real window. diff --git a/docs/env.md b/docs/env.md index bf378bb5..eca14522 100644 --- a/docs/env.md +++ b/docs/env.md @@ -82,6 +82,10 @@ These are the environment variables you can set for the `conversations-backend` | LLM_CONFIGURATION_FILE_PATH | Path to the LLM configuration JSON file. See [LLM Configuration](llm-configuration.md) for details | /conversations/configuration/llm/default.json | | LLM_DEFAULT_MODEL_HRID | HRID of the model used for conversations | default-model | | LLM_SUMMARIZATION_MODEL_HRID | HRID of the model used for summarization | default-summarization-model | +| DOCUMENT_CONTEXT_BUDGET_RATIO | Fraction of `max_token_context` for inlined conversation documents; `(1 - ratio)` is the conversation history share (security buffer subtracted from each afterwards). See [attachments.md](attachments.md) | `0.5` | +| DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS | Tokens subtracted from each of the document and history budgets independently | `1000` | +| CONVERSATION_SUMMARY_CONTEXT_MESSAGES | `ModelMessage` count kept after a conversation summary (use an even value). See [attachments.md](attachments.md) | `10` | +| CONVERSATION_SUMMARY_MAX_TOKENS | Max tokens for the conversation summarization LLM output | `2048` | | AI_API_KEY | AI API key to be used for the default provider (used in default LLM configuration, not for production use) | | | AI_BASE_URL | OpenAI compatible AI base URL (used in default LLM configuration, not for production use) | | | AI_MODEL | AI Model name to use (used in default LLM configuration, not for production use) | | diff --git a/docs/history-processor.md b/docs/history-processor.md index 569974fe..584669b1 100644 --- a/docs/history-processor.md +++ b/docs/history-processor.md @@ -1,28 +1,41 @@ -# History Processor (Sliding Window) +# History Processing (Conversation Summarization) -When a conversation grows long enough to exceed the model's context window, the backend automatically trims the oldest turns before each agent call. The full history is always preserved in the database — only the slice sent to the model is reduced. +When a conversation grows long enough to exceed its token budget, the backend summarizes the oldest turns instead of sending them verbatim to the model. The full history is always preserved in the database — only the slice sent to the model is reduced, with a running summary standing in for the summarized prefix. ## How it works -Before each agent call, `apply_sliding_window` checks whether the estimated token count of the history exceeds the conversation budget. If it does, it removes the oldest complete turns one by one until the history fits, then passes the trimmed slice to the model. +Summarization is a visible blocking phase at the start of the triggering turn. A single check runs at the start of every turn, over the stored history from previous turns (the just-received user message is intentionally not counted — see ADR 0002): when the active history exceeds the conversation budget, the turn enqueues a Celery task to generate the summary, waits on the open stream (emitting keep-alive, with an in-thread progress bar on the client), and only then sends the input to the agent. Generation is **Celery-only** — the web process never runs the summarization LLM — and is guarded by a time-bounded claim on the conversation; checkpoint writes only ever advance, so duplicate completions are no-ops. A turn is never answered with degraded context: if no summary lands while the turn is still over budget — the worker is down, or generation kept failing — the turn errors with `summarization_failed`, and the enqueued task keeps running to persist the summary for the retry (see ADR 0001). -A **turn** is the unit of trimming: a user message plus all the model responses and tool call/return pairs that follow it, up to the next user message. Turns are never split — either the whole turn is kept or the whole turn is dropped. +Two functions in `chat/agents/history_processors.py` split the work. The Celery task calls `generate_history_summary`, which runs a dedicated summarization agent (`LLM_SUMMARIZATION_MODEL_HRID`) to fold the un-summarized messages into the previous summary; the result is persisted on the conversation. The turn itself only trims — `build_active_history` selects the runtime window, never calling the LLM. Two fields carry the summary across turns: -The last turn is always kept, even if it alone exceeds the budget. This guarantees the current user message is never lost. +- `history_summary` — the running summary text, injected into the model's dynamic instructions. +- `history_summary_checkpoint` — the message index up to which the summary is valid. + +The model then receives the summary plus the last `CONVERSATION_SUMMARY_CONTEXT_MESSAGES` `ModelMessage` entries before the checkpoint, so recent detail is kept verbatim. Use an **even** value so the retained window starts on a user message. ```text -Full history: [turn 1] [turn 2] [turn 3] [turn 4] ← exceeds budget -After trim: [turn 3] [turn 4] ← fits within budget -Database: [turn 1] [turn 2] [turn 3] [turn 4] ← untouched +Full history: [msg 1 … msg 20] [msg 21 … msg 30] ← exceeds budget +Sent to model: [summary of 1–20] [msg 21 … msg 30] ← fits within budget +Database: [msg 1 … msg 30] + summary ← history untouched ``` -When trimming occurs, the backend emits a `context_trimmed` SSE event and the frontend displays a notice (session-only — resets on page reload): +Summarization is **incremental**: later turns only summarize messages added since the last checkpoint, folding them into the existing summary. + +While the summarization call runs, the backend emits a `summarize` tool-call event so the frontend can show a "Summarizing conversation..." notice; the matching tool-result event reports `done` (checkpoint advanced) or `error` (summary kept as-is, retried on a later turn). + +Independently of summarization, old tool returns are compacted on every turn (`clean_tool_history`): only the latest tool cycle keeps its full tool responses, older ones are replaced with a `` placeholder. + +## Failure behavior + +Every step degrades gracefully. If the summarization LLM call fails, the previous summary and checkpoint are kept and the turn fails with a `summarization_failed` error; the background task retries with exponential backoff. + +## Operational notes -> *Some older messages are no longer in the model's context.* +An over-budget turn blocks its own SSE request while the worker generates, re-reading the conversation from the DB every `HISTORY_SUMMARY_POLL_INTERVAL_SECONDS` (default 2s) until the summary lands or the claim goes stale — up to the claim TTL (`HISTORY_SUMMARY_CLAIM_TTL_SECONDS`, ~180s). Summarization throughput is therefore bounded by Celery worker headroom, not the web tier: if workers are backlogged past the enqueue grace window (`SUMMARIZATION_ENQUEUE_CLAIM_GRACE_SECONDS`), the turn stalls and then fails with `summarization_failed`. Size the worker pool so summarization jobs are picked up within that grace window under expected concurrent-conversation load; each waiting turn also adds one `arefresh_from_db` per poll tick, so the poll interval trades DB load against how quickly a landed summary is noticed. ## Configuration -Trimming is enabled by setting `max_token_context` on a model in the LLM configuration file: +Summarization is enabled by default: *omitting* `max_token_context` does **not** disable it — `resolve_conversation_budget` falls back to `DEFAULT_MAX_TOKEN_CONTEXT`, leaving a non-zero history budget. Only zeroing `message_token_budget` disables it (`max_token_context=0`, `DOCUMENT_CONTEXT_BUDGET_RATIO=1`, or a large enough `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS` — see [attachments.md](attachments.md#conversation-history-summarization)). The budget is configured by setting `max_token_context` on a model in the LLM configuration file: ```json { @@ -32,54 +45,20 @@ Trimming is enabled by setting `max_token_context` on a model in the LLM configu } ``` -If `max_token_context` is absent or `null` on the model configuration, it falls back to the `DEFAULT_MAX_TOKEN_CONTEXT` setting (default 8192). Set `DEFAULT_MAX_TOKEN_CONTEXT=0` in settings to disable trimming for models without an explicit context limit. - -### Misconfiguration warning - -If the computed conversation budget is 0 (e.g. `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS` exceeds the available tokens, or `DOCUMENT_CONTEXT_BUDGET_RATIO` is 1.0), trimming is disabled and the backend logs a warning: - -``` -Sliding window disabled: conversation budget is 0 -(max_token_context=N, security_buffer=M, budget_ratio=R). -``` - -### Budget formula - -The conversation budget is derived from two Django settings: - -| Setting | Description | Default | -|---|---|---| -| `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS` | Tokens reserved as a safety margin per pool | `1000` | -| `DOCUMENT_CONTEXT_BUDGET_RATIO` | Fraction of the usable window allocated to documents | `0.5` | +The conversation budget is derived from the same settings that drive the document inlining budget: ```text -conversation_budget = int(max_token_context × (1 - DOCUMENT_CONTEXT_BUDGET_RATIO)) - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS +message_token_budget = max(int(max_token_context × (1 - DOCUMENT_CONTEXT_BUDGET_RATIO)) - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS, 0) ``` -The security buffer is subtracted from **both** the conversation pool and the document pool independently — each pool's token count is approximated, so each pool needs its own safety margin. - -For example, with `max_token_context=128000`, `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS=1000`, and `DOCUMENT_CONTEXT_BUDGET_RATIO=0.3`: - -``` -conversation_budget = int(128000 × 0.7) - 1000 = 88600 tokens -``` +See [attachments.md](attachments.md#conversation-history-summarization) for the full budget formulas, the settings reference (`DOCUMENT_CONTEXT_BUDGET_RATIO`, `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS`, `CONVERSATION_SUMMARY_CONTEXT_MESSAGES`, `CONVERSATION_SUMMARY_MAX_TOKENS`, `LLM_SUMMARIZATION_MODEL_HRID`) and how to disable summarization without changing document budgets. ## Token estimation Tokens are estimated without calling the model's tokenizer, using tiktoken (`cl100k_base`). This is intentionally rough — the goal is to stay well within the context limit, not to maximise usage. Each message also carries a small fixed overhead (4 tokens per part + 4 per message) to account for role markers and formatting. -**Images** are counted using a flat constant of 1 500 tokens per image part (`ImageUrl` or `BinaryContent` with an `image/*` MIME type). Precise estimation is model-specific (OpenAI uses tile math, Anthropic uses a different formula), so a conservative constant keeps the estimator model-agnostic. - -**System prompt** tokens are subtracted from the conversation budget at calculation time using the static `configuration.system_prompt` string. Dynamic content injected at runtime (e.g. inlined document context) is not counted here — the security buffer absorbs that drift. - -**Agent instructions** (tool schemas, pydantic-ai framework overhead) are not counted. The security buffer is the backstop. +**System prompt, tool schemas and pydantic-ai framework overhead** are not counted; the security buffer is the backstop. ## Local testing -To trigger trimming locally, set a small `max_token_context` on the default model: - -```json -"max_token_context": 500 -``` - -After ~10 short messages the banner should appear and the chat should continue to work normally. +To trigger summarization locally, set a small `max_token_context` on the default model (large enough to survive `DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS`, or lower that setting too). After a few messages the "Summarizing conversation..." notice should appear and the chat should continue to work normally, with `history_summary` populated on the conversation row. diff --git a/src/backend/chat/agents/history_processors.py b/src/backend/chat/agents/history_processors.py index 3b2d4941..9c35372f 100644 --- a/src/backend/chat/agents/history_processors.py +++ b/src/backend/chat/agents/history_processors.py @@ -1,29 +1,50 @@ -"""Sliding window history processor for LLM conversation context management.""" +"""History processors for model message cleanup.""" +import dataclasses import logging from django.conf import settings from pydantic_ai import ImageUrl -from pydantic_ai.messages import BinaryContent, ModelMessage, ModelRequest, UserPromptPart +from pydantic_ai.messages import ( + BinaryContent, + ModelMessage, + ModelRequest, + ModelResponse, + TextPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) from chat.constants import IMAGE_MIME_PREFIX -from chat.tokens import ( - compute_conversation_budget as _compute_conversation_budget, -) -from chat.tokens import ( - count_approx_tokens, +from chat.tokens import compute_conversation_budget as _compute_conversation_budget +from chat.tokens import count_approx_tokens +from chat.tools.descriptions import CONVERSATION_SUMMARY_TOOL_DESCRIPTION + +from .summarize import SummarizationAgent + +logger = logging.getLogger(__name__) + +SUMMARY_SYSTEM_PREFIX = ( + "[Conversation summary from previous turns] (context only, not a user request):\n" ) -# Conservative flat estimate per image part. Precise estimation is model-specific -# so a constant keeps the estimator model-agnostic -# while ensuring we never silently under-count images. +_TOKENS_PER_PART_OVERHEAD = 4 # framing/delimiters per message part +_TOKENS_PER_MESSAGE_OVERHEAD = 4 # per-message envelope + +# Conservative flat estimate per image part. Precise estimation is model-specific, +# so a constant keeps the estimator model-agnostic while ensuring we never silently +# under-count images. _IMAGE_TOKEN_ESTIMATE = 1500 -logger = logging.getLogger(__name__) +class SummarizationRequiredError(Exception): + """A summary was required to fit the token budget but could not be generated.""" -def _stringify_message_content(content) -> str: + +def _stringify_message_content(content: object) -> str: + """Convert part content to a plain text representation.""" if content is None: return "" if isinstance(content, str): @@ -37,57 +58,248 @@ def _stringify_message_content(content) -> str: return "" +def _format_exchanges_for_summary(messages: list[ModelMessage]) -> str: + """Render user/assistant turns as `Role: text` lines for the summary prompt. + + Only `UserPromptPart` and assistant `TextPart` content is carried into the + running summary. `ToolCallPart`/`ToolReturnPart` content is intentionally + excluded: tool history (web search, RAG lookups) is transient context, not + durable conversation substance, and is stripped separately by + `safe_clean_tool_history`. Note the known asymmetry that + `_estimate_part_tokens` still counts tool-part tokens toward the budget, so + tool-heavy turns can trigger summarization without contributing summary text. + """ + lines = [] + + for msg in messages: + if isinstance(msg, ModelRequest): + role = "User" + parts = (p for p in msg.parts if isinstance(p, UserPromptPart)) + elif isinstance(msg, ModelResponse): + role = "Assistant" + parts = (p for p in msg.parts if isinstance(p, TextPart)) + else: + continue + + for part in parts: + content = part.content + text = _stringify_message_content(content).strip() + if text: + lines.append(f"{role}: {text}") + return "\n".join(lines) + + +async def summarize_conversation( + messages: list[ModelMessage], *, max_tokens: int, previous_summary: str | None = None +) -> str | None: + """Generate an updated running summary, folding `previous_summary` into `messages`. + + Returns the new summary text, or None when the model produces no output. + Provider API errors are left to propagate so the Celery task can retry + transient failures (see `RETRYABLE_SUMMARIZATION_ERRORS`). + """ + summarization_agent = SummarizationAgent() + latest_summary = previous_summary or "" + exchanges = _format_exchanges_for_summary(messages) + prompt = (CONVERSATION_SUMMARY_TOOL_DESCRIPTION + "\n\n").format( + exchanges=exchanges, latest_summary=latest_summary + ) + resp = await summarization_agent.run( + prompt, + model_settings={"max_tokens": max_tokens}, + ) + updated_summary = (resp.output or "").strip() + return updated_summary or None + + +def _latest_tool_call_ids(messages: list[ModelMessage]) -> set[str]: + """Return the tool_call_ids of the most recent assistant tool-call turn.""" + for message in reversed(messages): + if not isinstance(message, ModelResponse): + continue + response_tool_calls = [ + part.tool_call_id + for part in message.parts + if isinstance(part, ToolCallPart) and isinstance(part.tool_call_id, str) + ] + if response_tool_calls: + return set(response_tool_calls) + return set() + + +def _clean_request_parts(parts: list, latest_tool_call_ids: set[str]) -> list: + """Replace stale tool returns with a compact placeholder, keeping the latest cycle.""" + kept_parts = [] + for part in parts: + if not isinstance(part, ToolReturnPart): + kept_parts.append(part) + continue + if part.tool_call_id in latest_tool_call_ids: + kept_parts.append(part) + continue + tool_name = getattr(part, "tool_name", None) or "unknown_tool" + kept_parts.append( + ToolReturnPart( + tool_call_id=part.tool_call_id, + tool_name=tool_name, + content=f"<{tool_name} response compacted>", + ) + ) + return kept_parts + + +def clean_tool_history(messages: list[ModelMessage]) -> list[ModelMessage]: + """Compact old tool returns while preserving the latest tool cycle.""" + latest_tool_call_ids = _latest_tool_call_ids(messages) + cleaned_history: list[ModelMessage] = [] + + for message in messages: + if isinstance(message, ModelRequest): + kept_parts = _clean_request_parts(message.parts, latest_tool_call_ids) + if kept_parts: + cleaned_history.append(dataclasses.replace(message, parts=kept_parts)) + continue + + cleaned_history.append(message) + + return cleaned_history + + +def safe_clean_tool_history(messages: list[ModelMessage]) -> list[ModelMessage]: + """Compact tool history, falling back to the input on unexpected errors.""" + try: + return clean_tool_history(messages) + except Exception as exc: # pylint: disable=broad-except # noqa: BLE001 + logger.warning( + "Tool history cleanup failed, using raw history: %s", + exc, + exc_info=True, + ) + return messages + + +def _estimate_text_tokens(value: object) -> int: + """Count tokens for a stringifiable value, ignoring blank content.""" + text = _stringify_message_content(value).strip() + return count_approx_tokens(text) if text else 0 + + +def _estimate_item_tokens(item: object) -> int: + """Count tokens for one content item, flat-rating images.""" + if isinstance(item, ImageUrl): + return _IMAGE_TOKEN_ESTIMATE + if isinstance(item, BinaryContent) and item.media_type.startswith(IMAGE_MIME_PREFIX): + return _IMAGE_TOKEN_ESTIMATE + return _estimate_text_tokens(item) + + +def _estimate_content_tokens(content: object) -> int: + """Count tokens for a part's content, whether a scalar or a sequence.""" + if isinstance(content, (list, tuple)): + return sum(_estimate_item_tokens(item) for item in content) + return _estimate_text_tokens(content) + + +def _estimate_args_tokens(args: object) -> int: + """Count tokens for a tool-call part's args.""" + if isinstance(args, dict): + return count_approx_tokens(str(args)) + if isinstance(args, str) and args.strip(): + return count_approx_tokens(args) + return 0 + + +def _estimate_part_tokens(part: object) -> int: + """Count tokens for one message part: its content plus any tool-call args.""" + return _estimate_content_tokens(getattr(part, "content", "")) + _estimate_args_tokens( + getattr(part, "args", "") + ) + + def _estimate_message_tokens(message: ModelMessage) -> int: + """Estimate the token weight of one model message.""" parts = getattr(message, "parts", []) or [] if not parts: return 0 - part_tokens = 0 - for part in parts: - content = getattr(part, "content", "") - if isinstance(content, (list, tuple)): - for item in content: - if isinstance(item, ImageUrl): - part_tokens += _IMAGE_TOKEN_ESTIMATE - elif isinstance(item, BinaryContent) and item.media_type.startswith( - IMAGE_MIME_PREFIX - ): - part_tokens += _IMAGE_TOKEN_ESTIMATE - else: - text = _stringify_message_content(item).strip() - if text: - part_tokens += count_approx_tokens(text) - else: - text = _stringify_message_content(content).strip() - if text: - part_tokens += count_approx_tokens(text) - args = getattr(part, "args", "") - if isinstance(args, dict): - part_tokens += count_approx_tokens(str(args)) - elif isinstance(args, str) and args.strip(): - part_tokens += count_approx_tokens(args) - # Estimated overhead: ~4 tokens per part + ~4 tokens for message wrapper - return part_tokens + (4 * len(parts)) + 4 - - -def estimate_history_tokens(history: list[ModelMessage]) -> int: - """Return the approximate token count for a list of messages.""" - return sum(_estimate_message_tokens(m) for m in history) - - -def _group_into_turns(history: list[ModelMessage]) -> list[list[ModelMessage]]: - turns: list[list[ModelMessage]] = [] - current_turn: list[ModelMessage] = [] - for message in history: - is_user_turn_start = isinstance(message, ModelRequest) and any( - isinstance(p, UserPromptPart) for p in message.parts + + part_tokens = sum(_estimate_part_tokens(part) for part in parts) + return part_tokens + (_TOKENS_PER_PART_OVERHEAD * len(parts)) + _TOKENS_PER_MESSAGE_OVERHEAD + + +def _estimate_history_tokens(messages: list[ModelMessage]) -> int: + """Estimate token weight for a message list.""" + return sum(_estimate_message_tokens(message) for message in messages) + + +def _safe_checkpoint(messages: list[ModelMessage], summary_checkpoint: int) -> int: + """Clamp the stored summary checkpoint to the current history size.""" + return max(0, min(summary_checkpoint, len(messages))) + + +def build_active_history( + messages: list[ModelMessage], summary_checkpoint: int, context_messages: int +) -> list[ModelMessage]: + """Trim history to the runtime window: the last `context_messages` entries + before the checkpoint, plus everything after it. + + Pure list-slicing, no LLM. This is the history the model actually sees for + the current turn; the summarized prefix (everything before the window) is + represented by the persisted summary instead. Never returns empty while + `messages` is non-empty. + """ + checkpoint = _safe_checkpoint(messages, summary_checkpoint) + active_start = max(0, checkpoint - max(context_messages, 1)) + active = messages[active_start:] + if active: + return active + return messages[-1:] if messages else [] + + +def _active_window( + messages: list[ModelMessage], summary_checkpoint: int, context_messages: int +) -> tuple[int, list[ModelMessage], int]: + """Return (clamped checkpoint, active-history slice, its token estimate).""" + checkpoint = _safe_checkpoint(messages, summary_checkpoint) + active_history = build_active_history(messages, checkpoint, context_messages) + return checkpoint, active_history, _estimate_history_tokens(active_history) + + +async def generate_history_summary( + messages: list[ModelMessage], + *, + previous_summary: str | None = None, + summary_checkpoint: int = 0, + summary_max_tokens: int = 2048, +) -> tuple[str, int]: + """Summarize everything after the checkpoint; return (summary, new_checkpoint). + + Called by the Celery task once `should_generate_conversation_summary` has + confirmed the history is over budget. Summarizes `messages[checkpoint:]`, + folding in `previous_summary`, and returns the text plus the advanced + checkpoint (the full message count). Raises `SummarizationRequiredError` + when the model yields nothing, so the task retries rather than persisting an + empty summary. + """ + checkpoint = _safe_checkpoint(messages, summary_checkpoint) + next_checkpoint = len(messages) + summary_input = messages[checkpoint:next_checkpoint] + logger.debug( + "generate_history_summary summarizing messages %s:%s (%s message(s)).", + checkpoint, + next_checkpoint, + len(summary_input), + ) + summary = await summarize_conversation( + summary_input, + max_tokens=summary_max_tokens, + previous_summary=previous_summary, + ) + if not summary: + raise SummarizationRequiredError( + "Summarization produced no output for the over-budget history " + f"({len(summary_input)} message(s) after checkpoint {checkpoint})." ) - if is_user_turn_start and current_turn: - turns.append(current_turn) - current_turn = [] - current_turn.append(message) - if current_turn: - turns.append(current_turn) - return turns + return summary, next_checkpoint def resolve_conversation_budget(configuration) -> int: @@ -106,7 +318,7 @@ def resolve_conversation_budget(configuration) -> int: ) if conversation_budget == 0 and max_token_context > 0: logger.warning( - "Sliding window disabled: conversation budget is 0 " + "Summarization is disabled: conversation budget is 0 " "(max_token_context=%d, security_buffer=%d, budget_ratio=%.2f).", max_token_context, security_buffer, @@ -115,25 +327,19 @@ def resolve_conversation_budget(configuration) -> int: return conversation_budget -def apply_sliding_window( - history: list[ModelMessage], - conversation_budget: int, -) -> tuple[list[ModelMessage], bool]: - """Drop oldest turns until history fits within conversation_budget tokens. +def should_generate_conversation_summary( + messages: list[ModelMessage], + *, + summary_checkpoint: int = 0, + message_token_budget: int = 0, + context_messages: int = 10, +) -> bool: + """Return True when active history exceeds budget and new messages can be summarized.""" + if message_token_budget <= 0: + return False - Returns the (possibly trimmed) history and a flag indicating whether trimming occurred. - """ - if conversation_budget == 0 or not history: - return history, False - if estimate_history_tokens(history) <= conversation_budget: - return history, False - turns = _group_into_turns(history) - was_trimmed = False - # len(turns) > 1 guarantees we never drop the last turn (the current exchange) - while ( - len(turns) > 1 - and estimate_history_tokens([m for t in turns for m in t]) > conversation_budget - ): - turns.pop(0) - was_trimmed = True - return [m for turn in turns for m in turn], was_trimmed + cleaned_history = safe_clean_tool_history(messages) + checkpoint, _active, active_tokens = _active_window( + cleaned_history, summary_checkpoint, context_messages + ) + return active_tokens > message_token_budget and len(cleaned_history) > checkpoint diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index c222ed59..514f26b5 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -90,7 +90,6 @@ from django.conf import settings from django.contrib.auth import get_user_model from django.core.cache import cache -from django.core.exceptions import ImproperlyConfigured from django.core.files.storage import default_storage from django.utils.module_loading import import_string from django.utils.translation import gettext_lazy as _ @@ -99,7 +98,7 @@ from langfuse import get_client, propagate_attributes from mistralai.client.errors import HTTPValidationError, SDKError from pydantic_ai import Agent, InstrumentationSettings, RunContext, RunUsage -from pydantic_ai.capabilities import Instrumentation +from pydantic_ai.capabilities import Instrumentation, ProcessHistory from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError from pydantic_ai.messages import ( BinaryContent, @@ -130,7 +129,14 @@ from chat import models from chat.agents.conversation import ConversationAgent, TitleGenerationAgent -from chat.agents.history_processors import apply_sliding_window, resolve_conversation_budget +from chat.agents.history_processors import ( + SUMMARY_SYSTEM_PREFIX, + SummarizationRequiredError, + build_active_history, + resolve_conversation_budget, + safe_clean_tool_history, + should_generate_conversation_summary, +) from chat.agents.local_media_url_processors import ( build_project_image_urls, project_has_image_attachments, @@ -158,12 +164,15 @@ ContextDeps, DocumentParsingResult, ImagePostRunActions, + PreparedHistory, StreamingState, ) from chat.constants import ( ACCESS_FULL_CONTEXT, + HISTORY_SUMMARY_POLL_INTERVAL_SECONDS, IMAGE_MIME_PREFIX, MARKDOWN_MIME_TYPE, + SUMMARIZATION_ENQUEUE_CLAIM_GRACE_SECONDS, TEXT_MIME_PREFIX, ) from chat.document_context_builder import ( @@ -172,8 +181,10 @@ render_listing, ) from chat.enums import CollectionIndexState +from chat.llm_configuration import get_model_configuration from chat.mcp_servers import get_mcp_servers from chat.rate_limiting import record_and_compute_cooldown +from chat.tasks import summarize_conversation_history from chat.tools.descriptions import ( DOCUMENT_SUMMARIZE_PROJECT_TOOL_DESCRIPTION, DOCUMENT_SUMMARIZE_SYSTEM_PROMPT, @@ -207,14 +218,6 @@ IMAGE_SKIP_REASON_TEXT_ONLY = "model_text_only" -def get_model_configuration(model_hrid: str): - """Get the model configuration from settings.""" - try: - return settings.LLM_CONFIGURATIONS[model_hrid] - except KeyError as exc: - raise ImproperlyConfigured(f"LLM model configuration '{model_hrid}' not found.") from exc - - def _strip_thinking_parts(history: list[ModelMessage]) -> list[ModelMessage]: """Remove ThinkingPart from ModelResponse history for models that don't support it. @@ -298,6 +301,7 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument self.conversation = conversation self.user = user # authenticated user only self.model_hrid = model_hrid or settings.LLM_DEFAULT_MODEL_HRID # HRID of the model to use + self.model_configuration = get_model_configuration(self.model_hrid) self.language = language # might be None self._last_stop_check = 0 # Events queued during _prepare_agent_run for _run_agent to yield before @@ -311,7 +315,7 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument self.event_encoder = EventEncoder(CURRENT_EVENT_ENCODER_VERSION) # We use v4 for now self._support_streaming = True - if (streaming := get_model_configuration(self.model_hrid).supports_streaming) is not None: + if (streaming := self.model_configuration.supports_streaming) is not None: self._support_streaming = streaming # Feature flags @@ -328,19 +332,19 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument ) self._web_search_tool_registered = False self._self_documentation_tool_registered = False + self._history_summary = conversation.history_summary.strip() or None + self._history_summary_checkpoint = max(conversation.history_summary_checkpoint, 0) - _capabilities = ( - [ + _capabilities: list = [ProcessHistory(safe_clean_tool_history)] + if self._langfuse_available: + _capabilities.append( Instrumentation( settings=InstrumentationSettings( include_binary_content=self._store_analytics, include_content=self._store_analytics, ) ) - ] - if self._langfuse_available - else [] - ) + ) self.conversation_agent = ConversationAgent( model_hrid=self.model_hrid, language=self.language, @@ -348,6 +352,9 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument deps_type=ContextDeps, ) add_document_rag_search_tool_from_setting(self.conversation_agent, self.user) + self._conversation_message_token_budget = resolve_conversation_budget( + self.model_configuration + ) # Inject project-level custom instructions if the conversation belongs to a project llm_user_instructions = ( @@ -359,10 +366,19 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument @self.conversation_agent.instructions def project_instructions() -> str: + """Inject the project's custom user instructions.""" return llm_user_instructions + @self.conversation_agent.instructions + def history_summary_instructions() -> str: + """Inject the stored conversation summary as context-only instructions.""" + if not self._history_summary: + return "" + return f"{SUMMARY_SYSTEM_PREFIX}{self._history_summary.strip()}" + @property def _stop_cache_key(self): + """Cache key holding the stop signal for this conversation's stream.""" return f"streaming:stop:{self.conversation.pk}" # --------------------------------------------------------------------- # @@ -446,13 +462,14 @@ def _add_unreadable_images_instruction(self, subject: Optional[str]) -> None: @self.conversation_agent.instructions def acknowledge_unreadable_images() -> str: + """Instruct the model to tell the user it cannot read attached images.""" return ( f"This conversation includes {subject or 'some images'}, but you cannot read or " "process images. If a request depends on it, tell the user you " "can't read images." ) - async def _stream_content( + async def _stream_content( # noqa: PLR0912 # pylint: disable=too-many-branches self, messages: List[UIMessage], force_web_search: bool = False, encoder_fn: Callable = None ): """Common streaming logic with configurable encoder.""" @@ -511,6 +528,19 @@ async def _stream_content( yield encoded else: raise + except SummarizationRequiredError as exc: + logger.error( + "Summarization required but failed for conversation %s: %s", + self.conversation.pk, + exc, + ) + if messages: + await self._persist_user_message_on_error(messages[-1]) + error_event = events_v4.ErrorPart(error="summarization_failed") + if encoded := encoder_fn(error_event): + yield encoded + else: + raise async def stream_text_async(self, messages: List[UIMessage], force_web_search: bool = False): """Return only the assistant text deltas (legacy text mode).""" @@ -692,6 +722,107 @@ async def _prepare_agent_run( conversation_has_own_documents, ) + def _build_model_history(self, history: list[ModelMessage]) -> list[ModelMessage]: + """Compact tool cycles and trim to the runtime window for this turn. + + Pure transformation of the message list sent to the model: no LLM, no + DB writes. The summarized prefix is represented by the persisted summary + (injected via dynamic instructions), not by these messages. + """ + cleaned_history = safe_clean_tool_history(history) + return build_active_history( + cleaned_history, + self._history_summary_checkpoint, + settings.CONVERSATION_SUMMARY_CONTEXT_MESSAGES, + ) + + async def _wait_for_history_summary( + self, + claim_deadline: float | None = None, + ) -> AsyncGenerator[events_v4.DataPart, None]: + """Wait for the summarization task while its claim is live. + + Yields keep-alive data parts so proxies don't kill the idle SSE + stream. Returns when the summary lands (checkpoint advanced past + ours) or when the claim is dead — immediately, or, when + ``claim_deadline`` (monotonic clock) is given, once the deadline has + passed without a worker claiming. Bounded by the claim TTL plus the + grace window by construction. + """ + while True: + await self.conversation.arefresh_from_db( + fields=[ + "history_summary", + "history_summary_checkpoint", + "history_summary_claimed_at", + ] + ) + if self.conversation.history_summary_checkpoint > self._history_summary_checkpoint: + return + if not self.conversation.history_summarization_claim_is_live and ( + claim_deadline is None or time.monotonic() >= claim_deadline + ): + return + yield events_v4.DataPart(data=[{"type": "keep_alive"}]) + await asyncio.sleep(HISTORY_SUMMARY_POLL_INTERVAL_SECONDS) + + async def _run_history_summary_phase( + self, history: list[ModelMessage] + ) -> AsyncGenerator[events_v4.Event | PreparedHistory, None]: + """Run the conversation-summary phase, then trim the history. + + This runs before agent.iter so dynamic instructions can read the + updated summary in the same model request. When the message budget is + exceeded it emits a visible `summarize` tool call (running -> done) and + blocks the turn while a Celery worker generates the summary: enqueue the + task, wait on its claim, then adopt whatever landed. Generation is + Celery-only — if no summary lands (worker down) and the turn is still + over budget, it raises `SummarizationRequiredError` rather than answering + with degraded context. Always trims the history and yields the result as + the final item (the `DocumentParsingResult` sentinel pattern). + """ + should_emit_summary_event = should_generate_conversation_summary( + history, + summary_checkpoint=self._history_summary_checkpoint, + message_token_budget=self._conversation_message_token_budget, + context_messages=settings.CONVERSATION_SUMMARY_CONTEXT_MESSAGES, + ) + if should_emit_summary_event: + tool_call_id = str(uuid.uuid4()) + yield events_v4.ToolCallPart( + tool_call_id=tool_call_id, + tool_name="summarize", + args={"state": "running", "summary_scope": "conversation"}, + ) + claim_deadline = None + if not self.conversation.history_summarization_claim_is_live: + # Nobody is generating: enqueue the task (the summarized prefix + # predates this turn, so DB state is current), then give the + # worker a grace window to pick it up and claim. + await sync_to_async(summarize_conversation_history.delay)(str(self.conversation.pk)) + claim_deadline = time.monotonic() + SUMMARIZATION_ENQUEUE_CLAIM_GRACE_SECONDS + async for keep_alive in self._wait_for_history_summary(claim_deadline=claim_deadline): + yield keep_alive + # Adopt whatever landed while we waited. + self._history_summary = self.conversation.history_summary.strip() or None + self._history_summary_checkpoint = max(self.conversation.history_summary_checkpoint, 0) + if should_generate_conversation_summary( + history, + summary_checkpoint=self._history_summary_checkpoint, + message_token_budget=self._conversation_message_token_budget, + context_messages=settings.CONVERSATION_SUMMARY_CONTEXT_MESSAGES, + ): + # No summary landed and still over budget: no degraded turn. + raise SummarizationRequiredError( + "Conversation summary did not land while the history exceeds the budget." + ) + yield events_v4.ToolResultPart( + tool_call_id=tool_call_id, + result={"state": "done"}, + ) + + yield PreparedHistory(history=self._build_model_history(history)) + def _setup_web_search(self, force_web_search: bool) -> bool: """Configure web search if forced. Returns whether web search is actually forced.""" @@ -816,6 +947,7 @@ async def _fetch_document_data( key = document.url[len(DOCUMENT_URL_PREFIX) :] def _read(): + """Read the document's raw bytes from object storage.""" with default_storage.open(key, "rb") as file: return file.read() @@ -964,6 +1096,7 @@ def _prepare_prompt( # noqa: PLR0912 # pylint: disable=too-many-branches return user_prompt[0], attachment_images, attachment_documents async def _build_documents_listing(self) -> DocumentsListing | None: + """Assemble the per-turn document listing (inlined vs tool-call-only).""" budget_ratio = settings.DOCUMENT_CONTEXT_BUDGET_RATIO security_buffer_tokens = settings.DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS # Restrict to READY uploads on both arms: PENDING / ANALYZING / @@ -985,7 +1118,7 @@ async def _build_documents_listing(self) -> DocumentsListing | None: upload_state=models.AttachmentStatus.READY, ).order_by("created_at", "id") ) - max_token_context = self.conversation_agent.configuration.max_token_context + max_token_context = self.model_configuration.max_token_context # Cache the DocumentsListing: building it requires reading every text # attachment from object storage and tokenizing it. The fingerprint includes @@ -1023,6 +1156,7 @@ async def _build_documents_listing(self) -> DocumentsListing | None: return listing async def _build_document_context_instruction(self) -> str: + """Render the document listing into a system-prompt instruction string.""" listing = await self._build_documents_listing() return render_listing(listing) if listing is not None else "" @@ -1032,11 +1166,13 @@ def _setup_rag_tools(self, document_context_instruction: str = "") -> None: @self.conversation_agent.instructions def summarization_system_prompt() -> str: + """Inject the document-summarization system prompt.""" return DOCUMENT_SUMMARIZE_SYSTEM_PROMPT # Inform the model (system-level) that documents are attached and available @self.conversation_agent.instructions def attached_documents_note() -> str: + """Inform the model that user documents are attached and available.""" base_note = ( "[Internal context] User documents are attached to this conversation. " "Do not request re-upload of documents; consider them already available " @@ -1105,6 +1241,7 @@ def _setup_self_documentation_tool(self) -> None: @self.conversation_agent.instructions def self_documentation_instruction() -> str: + """Inject the self-documentation system prompt.""" return SELF_DOCUMENTATION_SYSTEM_PROMPT @self.conversation_agent.tool( @@ -1198,6 +1335,7 @@ async def _handle_input_documents( yield DocumentParsingResult(success=True, has_documents=True) async def _handle_non_streaming_response(self, node, run_ctx): + """Run a model node without streaming, emitting its parts as events.""" result = await node.run(run_ctx) logger.debug("node.run result: %s", result) for part in result.model_response.parts: @@ -1225,6 +1363,7 @@ async def _handle_non_streaming_response(self, node, run_ctx): logger.warning("Unknown part type: %s %s", type(part), dataclasses.asdict(part)) async def _handle_streaming_response(self, node, run_ctx, state: StreamingState): + """Stream a model node, emitting text/tool/reasoning deltas as events.""" async with node.stream(run_ctx) as request_stream: async for event in request_stream: await self._agent_stop_streaming() @@ -1265,6 +1404,7 @@ async def _handle_model_request_node( async def _handle_call_tools_node( self, node, run_ctx, state: StreamingState ) -> AsyncGenerator[events_v4.Event, None]: + """Stream a tool-call node, emitting tool result/error events.""" async with node.stream(run_ctx) as handle_stream: async for event in handle_stream: await self._agent_stop_streaming() @@ -1442,10 +1582,6 @@ async def _run_agent( # pylint: disable=too-many-locals,too-many-branches,too-m yield events_v4.DataPart(data=[pre_event]) self._pre_stream_events = [] - history, was_trimmed = apply_sliding_window( - history, resolve_conversation_budget(self.conversation_agent.configuration) - ) - # Re-index (or report busy) when the conversation has READY attachments and # its index is not current. INDEXING is included: reindex_conversation handles # the concurrent-claim case by yielding a busy error when it cannot claim the row. @@ -1487,11 +1623,17 @@ async def _run_agent( # pylint: disable=too-many-locals,too-many-branches,too-m if doc_result is None or not doc_result.success: return - if was_trimmed: - yield events_v4.DataPart(data=[{"type": "context_trimmed"}]) - conversation_has_own_documents = doc_result.has_documents + # The budget check runs on `history` (stored previous turns) only; the incoming + # user message is not counted, by design — a turn tipped over by it alone is caught + # next turn, and the security buffer absorbs the overflow meanwhile (see ADR 0002). + async for item in self._run_history_summary_phase(history): + if isinstance(item, PreparedHistory): + history = item.history + else: + yield item + await self._agent_stop_streaming(force_cache_check=True) self._setup_self_documentation_tool() self._setup_web_search_tool() @@ -1509,10 +1651,12 @@ async def _run_agent( # pylint: disable=too-many-locals,too-many-branches,too-m if history and history[-1].kind == "request": if history[-1].parts and history[-1].parts[-1].part_kind == "tool-return": history.append(ModelResponse(parts=[TextPart(content="ok")], kind="response")) + message_history = history if history else None async with self.conversation_agent.iter( [user_prompt] + input_images, - message_history=history, # history will pass through agent's history_processors + # History passes through the ProcessHistory capability set on the agent. + message_history=message_history, deps=self._context_deps, toolsets=mcp_servers, ) as run: diff --git a/src/backend/chat/clients/schema.py b/src/backend/chat/clients/schema.py index 8b1d4642..eda38114 100644 --- a/src/backend/chat/clients/schema.py +++ b/src/backend/chat/clients/schema.py @@ -21,6 +21,13 @@ class DocumentParsingResult: has_documents: bool +@dataclasses.dataclass +class PreparedHistory: + """Result marker carrying the prepared (trimmed) history out of the summary phase.""" + + history: List + + @dataclasses.dataclass class StreamingState: """ diff --git a/src/backend/chat/constants.py b/src/backend/chat/constants.py index 3b8e4a9e..bb1f8c60 100644 --- a/src/backend/chat/constants.py +++ b/src/backend/chat/constants.py @@ -19,3 +19,24 @@ # can't reference module-level constants). ACCESS_FULL_CONTEXT = "full-context" ACCESS_TOOL_CALL_ONLY = "tool_call_only" + +# Conversation summarization task limits. The claim TTL is the hard time +# limit plus a margin: past it, the claiming worker is provably dead +# (SIGKILLed at time_limit, OOM-killed, or crashed) and the claim stops +# blocking. Keep the three values consistent — the liveness math depends +# on TTL > TIME_LIMIT. +SUMMARIZATION_TASK_SOFT_TIME_LIMIT = 110 # seconds, raises SoftTimeLimitExceeded +SUMMARIZATION_TASK_TIME_LIMIT = 120 # seconds, worker is SIGKILLed +HISTORY_SUMMARY_CLAIM_TTL_SECONDS = SUMMARIZATION_TASK_TIME_LIMIT + 60 + +# After the triggering turn enqueues the summarization task, how long the +# wait loop tolerates "no live claim yet" before giving up and failing the +# turn (covers broker latency and a short worker backlog). Generation stays +# Celery-only; there is no inline fallback (see ADR 0002). +SUMMARIZATION_ENQUEUE_CLAIM_GRACE_SECONDS = 10 + +# How often the waiting turn re-reads the conversation from the DB while a +# worker generates the summary. Each tick is an `arefresh_from_db`, so this +# bounds the DB load per concurrent over-budget turn; summaries take seconds, +# so a couple-second cadence is responsive enough. +HISTORY_SUMMARY_POLL_INTERVAL_SECONDS = 2 diff --git a/src/backend/chat/llm_configuration.py b/src/backend/chat/llm_configuration.py index 9da6fc0b..94191a0b 100644 --- a/src/backend/chat/llm_configuration.py +++ b/src/backend/chat/llm_configuration.py @@ -181,20 +181,20 @@ def validate_web_search(cls, value: str | None) -> str | None: @field_validator("max_token_context", mode="before") @classmethod - def validate_max_token_context(cls, value: int | str | None) -> int | None: - """Accept integer-like values from JSON for model context size.""" + def validate_max_token_context(cls, value: Any) -> int | None: + """Parse max_token_context from literal, setting, or env value.""" if value is None or value == "": return None if isinstance(value, bool): # bool is an int subclass in Python, reject explicitly so True/False # don't become 1/0. raise ValueError("max_token_context must be an integer value.") - if isinstance(value, int): - parsed_value = value - elif isinstance(value, str): + if isinstance(value, str): + value = _get_setting_or_env_or_value(value) + try: parsed_value = int(value) - else: - raise ValueError("max_token_context must be an integer value.") + except (TypeError, ValueError) as exc: + raise ValueError("max_token_context must be an integer value.") from exc if parsed_value <= 0: raise ValueError("max_token_context must be a positive integer") @@ -263,3 +263,15 @@ def load_llm_configuration(llm_configuration_file_path) -> dict[str, LLModel]: def cached_load_llm_configuration(llm_configuration_file_path) -> dict[str, LLModel]: """Load the LLM configuration with caching to avoid redundant loading.""" return load_llm_configuration(llm_configuration_file_path) + + +def get_model_configuration(model_hrid: str): + """Get the model configuration from settings.""" + # pylint: disable=import-outside-toplevel + from django.conf import settings # noqa: PLC0415 + from django.core.exceptions import ImproperlyConfigured # noqa: PLC0415 + + try: + return settings.LLM_CONFIGURATIONS[model_hrid] + except KeyError as exc: + raise ImproperlyConfigured(f"LLM model configuration '{model_hrid}' not found.") from exc diff --git a/src/backend/chat/migrations/0013_chatconversation_history_summary_and_more.py b/src/backend/chat/migrations/0013_chatconversation_history_summary_and_more.py new file mode 100644 index 00000000..6d8ef5f5 --- /dev/null +++ b/src/backend/chat/migrations/0013_chatconversation_history_summary_and_more.py @@ -0,0 +1,41 @@ +# Generated by Django 5.2.14 on 2026-07-02 13:17 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("chat", "0012_chatconversationattachment_index_state_and_more"), + ] + + operations = [ + migrations.AddField( + model_name="chatconversation", + name="history_summary", + field=models.TextField( + blank=True, + default="", + help_text="Latest generated conversation summary used as system context", + ), + ), + migrations.AddField( + model_name="chatconversation", + name="history_summary_checkpoint", + field=models.PositiveIntegerField( + default=0, + help_text="Number of pydantic history messages already compacted into history_summary", + ), + ), + migrations.AddField( + model_name="chatconversation", + name="history_summary_claimed_at", + field=models.DateTimeField( + blank=True, + help_text=( + "When a worker claimed summary generation; claims expire after " + "HISTORY_SUMMARY_CLAIM_TTL_SECONDS (dead-worker liveness bound)" + ), + null=True, + ), + ), + ] diff --git a/src/backend/chat/models.py b/src/backend/chat/models.py index 335a6643..70dada93 100644 --- a/src/backend/chat/models.py +++ b/src/backend/chat/models.py @@ -1,9 +1,11 @@ """Models for chat conversations.""" +from datetime import timedelta from typing import Sequence from django.contrib.auth import get_user_model from django.db import models +from django.utils import timezone from django_pydantic_field import SchemaField @@ -11,6 +13,7 @@ from core.models import BaseModel from chat.ai_sdk_types import UIMessage +from chat.constants import HISTORY_SUMMARY_CLAIM_TTL_SECONDS from chat.enums import AttachmentIndexState, CollectionIndexState User = get_user_model() @@ -137,6 +140,15 @@ class ChatConversation(BaseModel): blank=True, help_text="Pydantic messages for the chat conversation, used for history", ) + history_summary = models.TextField( + blank=True, + default="", + help_text="Latest generated conversation summary used as system context", + ) + history_summary_checkpoint = models.PositiveIntegerField( + default=0, + help_text="Number of pydantic history messages already compacted into history_summary", + ) messages: Sequence[UIMessage] = SchemaField( schema=list[UIMessage], default=list, @@ -182,6 +194,13 @@ class ChatConversation(BaseModel): ), ) + history_summary_claimed_at = models.DateTimeField( + null=True, + blank=True, + help_text="When a worker claimed summary generation; claims expire after " + "HISTORY_SUMMARY_CLAIM_TTL_SECONDS (dead-worker liveness bound)", + ) + class Meta: # pylint: disable=missing-class-docstring indexes = [ models.Index(fields=["owner", "-created_at"]), @@ -191,6 +210,68 @@ class Meta: # pylint: disable=missing-class-docstring def __str__(self): return self.title or str(self.pk) + def claim_history_summarization(self) -> bool: + """Atomically claim summary generation; True when we now hold the claim. + + Succeeds when unclaimed or when the previous claim exceeded the TTL + (its worker is provably dead). + """ + now = timezone.now() + expiry = now - timedelta(seconds=HISTORY_SUMMARY_CLAIM_TTL_SECONDS) + updated = ( + type(self) + .objects.filter(pk=self.pk) + .filter( + models.Q(history_summary_claimed_at__isnull=True) + | models.Q(history_summary_claimed_at__lt=expiry) + ) + .update(history_summary_claimed_at=now) + ) + if updated: + self.history_summary_claimed_at = now + return bool(updated) + + def release_history_summarization_claim(self) -> None: + """Release the claim, but only while we still hold it. + + Guards against a stale worker wiping a newer worker's live claim: if + our claim already expired and another worker reclaimed, the timestamp + no longer matches ours and the update is a no-op. Mirrors the guarded + checkpoint write in `persist_history_summary`. + """ + type(self).objects.filter( + pk=self.pk, history_summary_claimed_at=self.history_summary_claimed_at + ).update(history_summary_claimed_at=None) + self.history_summary_claimed_at = None + + @property + def history_summarization_claim_is_live(self) -> bool: + """True while a claim exists and its worker may still be alive.""" + claimed_at = self.history_summary_claimed_at + if claimed_at is None: + return False + return claimed_at > timezone.now() - timedelta(seconds=HISTORY_SUMMARY_CLAIM_TTL_SECONDS) + + def persist_history_summary(self, summary: str, checkpoint: int) -> bool: + """Persist a generated summary only if it advances the checkpoint. + + Messages are append-only, so a larger checkpoint is always the newer + result; late or duplicate completions are rejected. + """ + updated = ( + type(self) + .objects.filter(pk=self.pk, history_summary_checkpoint__lt=checkpoint) + .update( + history_summary=summary, + history_summary_checkpoint=checkpoint, + updated_at=timezone.now(), + ) + ) + if updated: + self.history_summary = summary + self.history_summary_checkpoint = checkpoint + return bool(updated) + class ChatConversationAttachment(BaseModel): """ diff --git a/src/backend/chat/tasks.py b/src/backend/chat/tasks.py index d631df0c..95fb80a2 100644 --- a/src/backend/chat/tasks.py +++ b/src/backend/chat/tasks.py @@ -1,13 +1,39 @@ -"""Celery tasks for the chat app.""" +"""Celery tasks for the chat application.""" import logging +from django.conf import settings +from django.core.exceptions import ImproperlyConfigured + +from asgiref.sync import async_to_sync +from pydantic_ai.exceptions import ModelAPIError +from pydantic_ai.messages import ModelMessagesTypeAdapter + from chat.agent_rag.indexing import index_project_attachment -from chat.models import ChatConversationAttachment +from chat.agents.history_processors import ( + generate_history_summary, + resolve_conversation_budget, + safe_clean_tool_history, + should_generate_conversation_summary, +) +from chat.constants import ( + SUMMARIZATION_TASK_SOFT_TIME_LIMIT, + SUMMARIZATION_TASK_TIME_LIMIT, +) +from chat.llm_configuration import get_model_configuration +from chat.models import ChatConversation, ChatConversationAttachment from conversations.celery_app import app logger = logging.getLogger(__name__) +# Transient provider failures worth retrying: every model-provider API error +# (HTTP 4xx/5xx, connection and timeout errors) raised by the summarization +# agent. Deterministic failures (empty output -> SummarizationRequiredError, +# application bugs) fall through and fail fast; the claim is released either way +# and the budget re-check makes stale retries no-ops. Note 4xx is not singled +# out here, so a bad-request/auth error would still retry (a known trade-off). +RETRYABLE_SUMMARIZATION_ERRORS = (ModelAPIError,) + @app.task def index_project_attachment_task(attachment_id): @@ -33,3 +59,65 @@ def index_project_attachment_task(attachment_id): return index_project_attachment(attachment) + + +@app.task( + autoretry_for=RETRYABLE_SUMMARIZATION_ERRORS, + retry_backoff=True, + max_retries=3, + soft_time_limit=SUMMARIZATION_TASK_SOFT_TIME_LIMIT, + time_limit=SUMMARIZATION_TASK_TIME_LIMIT, +) +def summarize_conversation_history(conversation_id: str) -> None: + """Generate the running summary for a conversation's history (see ADR 0002). + + Enqueued by the over-budget turn, which waits on the claim. Claims the + conversation, re-checks the budget against current DB state, generates, + and persists behind the checkpoint guard. The claim is always released: + a retry re-claims, and the re-check makes stale retries no-ops. + """ + try: + conversation = ChatConversation.objects.get(pk=conversation_id) + except ChatConversation.DoesNotExist: + logger.info("Conversation %s gone before summarization ran.", conversation_id) + return + + if not conversation.claim_history_summarization(): + logger.debug("Conversation %s already claimed, skipping.", conversation_id) + return + + try: + model_hrid = conversation.model_hrid or settings.LLM_DEFAULT_MODEL_HRID + try: + model_configuration = get_model_configuration(model_hrid) + except ImproperlyConfigured: + logger.warning( + "Model %s no longer configured, skipping summarization for %s.", + model_hrid, + conversation_id, + ) + return + + budget = resolve_conversation_budget(model_configuration) + if budget <= 0: + return + + messages = ModelMessagesTypeAdapter.validate_python(conversation.pydantic_messages) + if not should_generate_conversation_summary( + messages, + summary_checkpoint=conversation.history_summary_checkpoint, + message_token_budget=budget, + context_messages=settings.CONVERSATION_SUMMARY_CONTEXT_MESSAGES, + ): + return + + cleaned = safe_clean_tool_history(messages) + summary, checkpoint = async_to_sync(generate_history_summary)( + cleaned, + previous_summary=conversation.history_summary.strip() or None, + summary_checkpoint=conversation.history_summary_checkpoint, + summary_max_tokens=settings.CONVERSATION_SUMMARY_MAX_TOKENS, + ) + conversation.persist_history_summary(summary, checkpoint) + finally: + conversation.release_history_summarization_claim() diff --git a/src/backend/chat/tests/agents/test_history_processors.py b/src/backend/chat/tests/agents/test_history_processors.py index 4fd97886..aff27e9c 100644 --- a/src/backend/chat/tests/agents/test_history_processors.py +++ b/src/backend/chat/tests/agents/test_history_processors.py @@ -1,263 +1,421 @@ -"""Tests for chat.agents.history_processors.""" +"""Tests for history processors.""" import logging +from types import SimpleNamespace -from django.test import override_settings - -from pydantic_ai import ImageUrl -from pydantic_ai.messages import ( - BinaryContent, +import pytest +from pydantic_ai import ( + Agent, + ModelMessage, ModelRequest, ModelResponse, TextPart, + UserPromptPart, +) +from pydantic_ai.capabilities import ProcessHistory +from pydantic_ai.exceptions import ModelHTTPError +from pydantic_ai.messages import ( ToolCallPart, ToolReturnPart, - UserPromptPart, ) +from pydantic_ai.models.function import AgentInfo, FunctionModel -from chat.agents.history_processors import ( - _IMAGE_TOKEN_ESTIMATE, - _estimate_message_tokens, - _group_into_turns, - _stringify_message_content, - apply_sliding_window, - estimate_history_tokens, - resolve_conversation_budget, +from chat.agents import history_processors +from chat.tokens import ( + compute_conversation_budget, + compute_document_budget, + count_approx_tokens, ) -from chat.tokens import count_approx_tokens -# ── helpers ────────────────────────────────────────────────────────────────── +@pytest.fixture +def _received_messages_fixture() -> list[ModelMessage]: + """Fixture to capture messages received by the function model.""" + return [] -def user_msg(text: str) -> ModelRequest: - """Build a minimal user ModelRequest.""" - return ModelRequest(parts=[UserPromptPart(content=text)], kind="request") +@pytest.fixture(name="received_messages") +def received_messages_fixture_alias( + _received_messages_fixture: list[ModelMessage], +) -> list[ModelMessage]: + """Expose received messages fixture with a stable pytest name.""" + return _received_messages_fixture -def assistant_msg(text: str) -> ModelResponse: - """Build a minimal assistant ModelResponse.""" - return ModelResponse(parts=[TextPart(content=text)], kind="response") +@pytest.fixture(name="function_model") +def function_model_fixture(received_messages: list[ModelMessage]) -> FunctionModel: + """Fixture to capture model function.""" -def tool_call_msg(tool_call_id: str = "tc1") -> ModelResponse: - """Build a ModelResponse containing a single tool call.""" - return ModelResponse( - parts=[ToolCallPart(tool_name="web_search", tool_call_id=tool_call_id, args="{}")], - kind="response", - ) + def capture_model_function(messages: list[ModelMessage], _info: AgentInfo) -> ModelResponse: + # Capture exactly what reaches the provider after history processors. + received_messages.clear() + received_messages.extend(messages) + return ModelResponse(parts=[TextPart(content="Provider response")]) + return FunctionModel(capture_model_function) -def tool_return_msg(tool_call_id: str = "tc1") -> ModelRequest: - """Build a ModelRequest containing a single tool return.""" - return ModelRequest( - parts=[ToolReturnPart(tool_name="web_search", tool_call_id=tool_call_id, content="result")], - kind="request", - ) + +def _build_turns(turn_count: int) -> list: + """Build a list of turns for testing.""" + messages = [] + for turn in range(1, turn_count + 1): + messages.append(ModelRequest(parts=[UserPromptPart(content=[f"user-{turn}"])])) + messages.append(ModelResponse(parts=[TextPart(content=f"assistant-{turn}")])) + return messages -# ── _estimate_message_tokens ───────────────────────────────────────────────── +def test_history_processors_are_applied_before_provider_call( + function_model: FunctionModel, received_messages: list[ModelMessage] +): + """History processors should run before provider invocation.""" + def keep_only_requests(messages: list[ModelMessage]) -> list[ModelMessage]: + return [msg for msg in messages if isinstance(msg, ModelRequest)] -def test_estimate_message_tokens_empty_parts(): - """Messages with no parts return 0 tokens.""" - msg = ModelRequest(parts=[], kind="request") - assert _estimate_message_tokens(msg) == 0 + agent = Agent(function_model, capabilities=[ProcessHistory(keep_only_requests)]) + message_history = [ + ModelRequest(parts=[UserPromptPart(content="Question 1")]), + ModelResponse(parts=[TextPart(content="Answer 1")]), + ] + agent.run_sync("Question 2", message_history=message_history) + assert len(received_messages) == 1 + assert isinstance(received_messages[0], ModelRequest) + user_prompt_contents = [ + part.content for part in received_messages[0].parts if isinstance(part, UserPromptPart) + ] + assert user_prompt_contents == ["Question 1", "Question 2"] -def test_estimate_message_tokens_tool_call(): - """ToolCallPart with args='{}' accounts for arg tokens plus per-message overhead.""" - # args="{}" → tiktoken("{}")=1 token; overhead: 4*1+4=8 → total 9 - msg = tool_call_msg() - assert _estimate_message_tokens(msg) == 9 +def test_build_active_history_keeps_full_history_within_context_window(): + """The whole history is kept when it fits inside the context window.""" + messages = _build_turns(2) -def test_estimate_message_tokens_image_url_adds_constant(): - """An ImageUrl part adds the flat IMAGE_TOKEN_ESTIMATE constant.""" - msg = ModelRequest( - parts=[UserPromptPart(content=[ImageUrl(url="https://example.com/photo.jpg")])], - kind="request", + result = history_processors.build_active_history( + messages, summary_checkpoint=0, context_messages=10 ) - # overhead: 4*1 parts + 4 message = 8 - assert _estimate_message_tokens(msg) == _IMAGE_TOKEN_ESTIMATE + 8 + assert result == messages + + +def test_clean_tool_history_redacts_old_tool_returns_but_keeps_latest_tool_result(): + """Old tool results are compacted, latest tool result stays intact.""" + messages = [ + ModelResponse(parts=[ToolCallPart(tool_call_id="old", tool_name="search", args="{}")]), + ModelRequest( + parts=[ + ToolReturnPart( + tool_call_id="old", + tool_name="search", + content="large old result", + ) + ] + ), + ModelResponse(parts=[ToolCallPart(tool_call_id="latest", tool_name="search", args="{}")]), + ModelRequest( + parts=[ + ToolReturnPart( + tool_call_id="latest", + tool_name="search", + content="fresh result", + ) + ] + ), + ] + + result = history_processors.clean_tool_history(messages) + + old_return = result[1].parts[0] + latest_return = result[3].parts[0] + assert isinstance(old_return, ToolReturnPart) + assert isinstance(latest_return, ToolReturnPart) + assert old_return.content == "" + assert latest_return.content == "fresh result" -def test_estimate_message_tokens_binary_image_adds_constant(): - """A BinaryContent image part adds the flat IMAGE_TOKEN_ESTIMATE constant.""" - msg = ModelRequest( - parts=[UserPromptPart(content=[BinaryContent(data=b"\x89PNG", media_type="image/png")])], - kind="request", + +@pytest.mark.asyncio +async def test_generate_history_summary_returns_summary_and_advances_checkpoint(monkeypatch): + """Generation summarizes everything after the checkpoint and advances it to the end.""" + summarized_messages = [] + + async def fake_summary(messages, *, max_tokens=300, previous_summary=None): + summarized_messages.extend(messages) + _ = max_tokens + _ = previous_summary + return "summary-v1" + + monkeypatch.setattr(history_processors, "summarize_conversation", fake_summary) + messages = _build_turns(3) + + summary, checkpoint = await history_processors.generate_history_summary(messages) + + assert summary == "summary-v1" + assert checkpoint == len(messages) + assert summarized_messages == messages + + +@pytest.mark.asyncio +async def test_generate_history_summary_summarizes_only_after_checkpoint(monkeypatch): + """With an existing checkpoint, only messages after it are sent to the model.""" + summarized_messages = [] + + async def fake_summary(messages, *, max_tokens=300, previous_summary=None): + summarized_messages.extend(messages) + _ = max_tokens + _ = previous_summary + return "summary-v2" + + monkeypatch.setattr(history_processors, "summarize_conversation", fake_summary) + messages = _build_turns(5) + + summary, checkpoint = await history_processors.generate_history_summary( + messages, summary_checkpoint=6 ) - assert _estimate_message_tokens(msg) == _IMAGE_TOKEN_ESTIMATE + 8 + + assert summary == "summary-v2" + assert checkpoint == len(messages) + assert summarized_messages == messages[6:] -# ── estimate_history_tokens ─────────────────────────────────────────────────── +@pytest.mark.asyncio +async def test_summarize_conversation_propagates_provider_errors(monkeypatch): + """Provider API errors propagate so the Celery task can retry the transient case. + Swallowing them (returning None) would turn every transient blip into a + non-retryable SummarizationRequiredError — the regression this guards against. + """ -def test_estimate_history_tokens_empty(): - """Empty history returns 0 tokens.""" - assert estimate_history_tokens([]) == 0 + class _FailingAgent: + async def run(self, *_args, **_kwargs): + """Simulate a transient provider failure.""" + raise ModelHTTPError(status_code=503, model_name="summ") + monkeypatch.setattr(history_processors, "SummarizationAgent", _FailingAgent) -def test_estimate_history_tokens_two_messages(): - """Token count for a user+assistant exchange matches the sum of individual estimates.""" - history = [user_msg("hello world"), assistant_msg("hello world")] - assert estimate_history_tokens(history) == 24 + messages = _build_turns(1) + with pytest.raises(ModelHTTPError): + await history_processors.summarize_conversation(messages, max_tokens=100) -# ── _group_into_turns ───────────────────────────────────────────────────────── +@pytest.mark.asyncio +async def test_generate_history_summary_raises_when_output_empty(monkeypatch): + """Genuinely empty model output raises SummarizationRequiredError (non-retryable).""" + class _EmptyAgent: + async def run(self, *_args, **_kwargs): + """Simulate the model returning empty output.""" + return SimpleNamespace(output="") -def test_group_into_turns_single_turn(): - """A single user+assistant exchange forms one turn.""" - history = [user_msg("hi"), assistant_msg("hello")] - turns = _group_into_turns(history) - assert len(turns) == 1 - assert turns[0] == history + monkeypatch.setattr(history_processors, "SummarizationAgent", _EmptyAgent) + messages = _build_turns(2) + with pytest.raises(history_processors.SummarizationRequiredError): + await history_processors.generate_history_summary(messages) -def test_group_into_turns_two_turns(): - """Two user+assistant exchanges form two turns.""" - h = [user_msg("q1"), assistant_msg("a1"), user_msg("q2"), assistant_msg("a2")] - turns = _group_into_turns(h) - assert len(turns) == 2 - assert turns[0] == [h[0], h[1]] - assert turns[1] == [h[2], h[3]] +def test_build_active_history_starts_at_checkpoint_minus_context(): + """The runtime window starts `context_messages` before the checkpoint.""" + messages = _build_turns(5) + + result = history_processors.build_active_history( + messages, summary_checkpoint=6, context_messages=1 + ) + + assert result == messages[5:] -def test_group_into_turns_tool_calls_stay_in_same_turn(): - """Tool call/return messages stay grouped with the user turn that triggered them.""" - h = [ - user_msg("search this"), - tool_call_msg(), - tool_return_msg(), - assistant_msg("found it"), - user_msg("thanks"), - assistant_msg("welcome"), + +def test_build_active_history_drops_the_summarized_prefix(): + """A large already-summarized prefix is dropped from the runtime window.""" + messages = [ + ModelRequest(parts=[UserPromptPart(content=["old user " + ("x " * 500)])]), + ModelResponse(parts=[TextPart(content="old assistant " + ("x " * 500))]), + ModelRequest(parts=[UserPromptPart(content=["context user"])]), + ModelResponse(parts=[TextPart(content="context assistant")]), + ModelRequest(parts=[UserPromptPart(content=["new user"])]), + ModelResponse(parts=[TextPart(content="new assistant")]), ] - turns = _group_into_turns(h) - assert len(turns) == 2 - assert turns[0] == h[:4] - assert turns[1] == h[4:] + result = history_processors.build_active_history( + messages, summary_checkpoint=4, context_messages=1 + ) -def test_group_into_turns_empty(): - """Empty history produces no turns.""" - assert not _group_into_turns([]) + assert result == messages[3:] -# ── apply_sliding_window ────────────────────────────────────────────────────── +def test_build_active_history_never_empty_when_checkpoint_at_end(): + """pydantic-ai rejects empty processed history, so the last message is kept.""" + messages = _build_turns(2) + result = history_processors.build_active_history( + messages, summary_checkpoint=len(messages), context_messages=1 + ) -def test_apply_sliding_window_no_trim_needed(): - """History within budget is returned unchanged with trimmed=False.""" - history = [user_msg("hi"), assistant_msg("hello")] - result, trimmed = apply_sliding_window(history, 10_000) - assert result == history - assert trimmed is False + assert result == messages[3:] -def test_apply_sliding_window_budget_zero_disables(): - """Budget of 0 disables trimming regardless of history size.""" - history = [user_msg("a" * 4000), assistant_msg("b" * 4000)] - result, trimmed = apply_sliding_window(history, 0) - assert result == history - assert trimmed is False +def test_clean_tool_history_has_no_summary_checkpoint_behavior(): + """The pydantic-ai history processor path only cleans tools.""" + messages = _build_turns(2) + result = history_processors.clean_tool_history(messages) -def test_apply_sliding_window_trims_oldest_turn(): - """Oldest turn is evicted when total history exceeds budget.""" - # old_turn: user("a"*400)+assistant("b"*400) ≈ 216 tokens - # new_turn: user("c"*400)+assistant("d"*400) ≈ 216 tokens - # total ≈ 432 tokens; budget=250 → old_turn evicted - old_turn = [user_msg("a" * 400), assistant_msg("b" * 400)] - new_turn = [user_msg("c" * 400), assistant_msg("d" * 400)] - result, trimmed = apply_sliding_window(old_turn + new_turn, 250) - assert trimmed is True - assert result == new_turn + assert result == messages -def test_apply_sliding_window_keeps_last_turn_even_if_too_big(): - """A single oversized turn is kept as-is; trimmed remains False.""" - big_turn = [user_msg("a" * 4000), assistant_msg("b" * 4000)] - result, trimmed = apply_sliding_window(big_turn, 100) - assert result == big_turn - assert trimmed is False +@pytest.mark.asyncio +async def test_generate_history_summary_raises_when_model_returns_nothing(monkeypatch): + """No degraded turn: an empty summary raises so the task retries.""" + async def fake_summary(_messages, *, max_tokens=300, previous_summary=None): + _ = max_tokens, previous_summary -# ── resolve_conversation_budget ─────────────────────────────────────────────── + monkeypatch.setattr(history_processors, "summarize_conversation", fake_summary) + messages = _build_turns(20) + with pytest.raises(history_processors.SummarizationRequiredError): + await history_processors.generate_history_summary(messages) -@override_settings( - DEFAULT_MAX_TOKEN_CONTEXT=8192, - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS=1000, - DOCUMENT_CONTEXT_BUDGET_RATIO=0.5, -) -def test_resolve_conversation_budget_no_max_context_uses_default(): + +def test_should_generate_conversation_summary_when_budget_exceeded(): + """Frontend summary event should trigger only when over budget.""" + messages = _build_turns(3) + assert history_processors.should_generate_conversation_summary(messages, message_token_budget=1) + assert not history_processors.should_generate_conversation_summary( + messages, message_token_budget=10_000 + ) + assert not history_processors.should_generate_conversation_summary( + messages, + summary_checkpoint=len(messages), + message_token_budget=1, + context_messages=1, + ) + assert not history_processors.should_generate_conversation_summary( + messages, + summary_checkpoint=28, + message_token_budget=1, + context_messages=1, + ) + messages_with_large_summarized_prefix = [ + ModelRequest(parts=[UserPromptPart(content=["old user " + ("x " * 500)])]), + ModelResponse(parts=[TextPart(content="old assistant " + ("x " * 500))]), + ModelRequest(parts=[UserPromptPart(content=["context user"])]), + ModelResponse(parts=[TextPart(content="context assistant")]), + ModelRequest(parts=[UserPromptPart(content=["new user"])]), + ModelResponse(parts=[TextPart(content="new assistant")]), + ] + assert not history_processors.should_generate_conversation_summary( + messages_with_large_summarized_prefix, + summary_checkpoint=4, + message_token_budget=100, + context_messages=1, + ) + + +def test_resolve_conversation_budget_no_max_context_uses_default(settings): """Models without max_token_context fall back to DEFAULT_MAX_TOKEN_CONTEXT.""" + settings.DEFAULT_MAX_TOKEN_CONTEXT = 8192 + settings.DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS = 1000 + settings.DOCUMENT_CONTEXT_BUDGET_RATIO = 0.5 class Cfg: # pylint: disable=missing-class-docstring max_token_context = None # int(8192*0.5)-1000=3096 - assert resolve_conversation_budget(Cfg()) == 3096 + assert history_processors.resolve_conversation_budget(Cfg()) == 3096 -def test_resolve_conversation_budget_zero_max_context(): +def test_resolve_conversation_budget_zero_max_context(settings): """max_token_context=0 returns 0 (trimming disabled).""" + settings.DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS = 1000 + settings.DOCUMENT_CONTEXT_BUDGET_RATIO = 0.5 class Cfg: # pylint: disable=missing-class-docstring max_token_context = 0 - assert resolve_conversation_budget(Cfg()) == 0 + assert history_processors.resolve_conversation_budget(Cfg()) == 0 -@override_settings( - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS=1000, - DOCUMENT_CONTEXT_BUDGET_RATIO=0.5, -) -def test_resolve_conversation_budget_formula(): +def test_resolve_conversation_budget_formula(settings): """Budget is computed as int(max*(1-ratio))-buffer.""" + settings.DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS = 1000 + settings.DOCUMENT_CONTEXT_BUDGET_RATIO = 0.5 class Cfg: # pylint: disable=missing-class-docstring max_token_context = 10_000 # int(10000*0.5)-1000=4000 - assert resolve_conversation_budget(Cfg()) == 4000 + assert history_processors.resolve_conversation_budget(Cfg()) == 4000 -@override_settings( - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS=0, - DOCUMENT_CONTEXT_BUDGET_RATIO=0.0, -) -def test_resolve_conversation_budget_zero_ratio(): +def test_resolve_conversation_budget_zero_ratio(settings): """budget_ratio=0 allocates the full context window to conversation.""" + settings.DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS = 0 + settings.DOCUMENT_CONTEXT_BUDGET_RATIO = 0.0 class Cfg: # pylint: disable=missing-class-docstring max_token_context = 10_000 - assert resolve_conversation_budget(Cfg()) == 10_000 + assert history_processors.resolve_conversation_budget(Cfg()) == 10_000 -@override_settings(DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS=10_000) -def test_resolve_conversation_budget_logs_warning_when_buffer_exceeds_context(caplog): +@pytest.mark.parametrize("budget_ratio", [0.0, 0.3, 0.5, 0.7, 1.0]) +@pytest.mark.parametrize("security_buffer", [0, 1000]) +@pytest.mark.parametrize("max_token_context", [4000, 10_000, 128_000]) +def test_document_and_conversation_budgets_never_jointly_exceed_context( + max_token_context, budget_ratio, security_buffer +): + """Documents and history partition the window; together they never overflow it. + + The two budgets split `max_token_context` by `budget_ratio` / `(1 - ratio)`, + each minus the security buffer, so summarization bounding the history share + is enough to keep the whole request within the model's context window. + """ + document_budget = compute_document_budget(max_token_context, budget_ratio, security_buffer) + conversation_budget = compute_conversation_budget( + max_token_context, budget_ratio, security_buffer + ) + + assert document_budget + conversation_budget <= max_token_context + + +def test_budget_partition_leaves_only_the_security_buffer_as_slack(): + """At a large-doc ratio the buffer is the sole headroom between the two budgets.""" + max_token_context = 10_000 + budget_ratio = 0.7 + security_buffer = 1000 + + document_budget = compute_document_budget(max_token_context, budget_ratio, security_buffer) + conversation_budget = compute_conversation_budget( + max_token_context, budget_ratio, security_buffer + ) + + # 7000-1000=6000 docs + 3000-1000=2000 history = 8000; the missing 2000 is + # exactly the two security buffers. A ratio/buffer change that breaks the + # partition would push this sum past max_token_context and fail loudly. + assert document_budget == 6000 + assert conversation_budget == 2000 + assert max_token_context - (document_budget + conversation_budget) == 2 * security_buffer + + +def test_resolve_conversation_budget_logs_warning_when_buffer_exceeds_context(settings, caplog): """A warning is logged when the security buffer leaves no tokens for conversation.""" + settings.DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS = 10_000 + settings.DOCUMENT_CONTEXT_BUDGET_RATIO = 0.5 class Cfg: # pylint: disable=missing-class-docstring max_token_context = 5_000 with caplog.at_level(logging.WARNING, logger="chat.agents.history_processors"): - result = resolve_conversation_budget(Cfg()) + result = history_processors.resolve_conversation_budget(Cfg()) assert result == 0 - assert "Sliding window disabled" in caplog.text + assert "Summarization is disabled" in caplog.text -@override_settings( - DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS=0, - DOCUMENT_CONTEXT_BUDGET_RATIO=0.0, -) -def test_resolve_conversation_budget_deducts_system_prompt(): +def test_resolve_conversation_budget_deducts_system_prompt(settings): """System prompt tokens are subtracted from the conversation budget.""" + settings.DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS = 0 + settings.DOCUMENT_CONTEXT_BUDGET_RATIO = 0.0 system_prompt = "You are a helpful assistant." class CfgNoPrompt: # pylint: disable=missing-class-docstring @@ -268,24 +426,20 @@ class CfgWithPrompt: # pylint: disable=missing-class-docstring max_token_context = 10_000 system_prompt = "You are a helpful assistant." - budget_without = resolve_conversation_budget(CfgNoPrompt()) - budget_with = resolve_conversation_budget(CfgWithPrompt()) + budget_without = history_processors.resolve_conversation_budget(CfgNoPrompt()) + budget_with = history_processors.resolve_conversation_budget(CfgWithPrompt()) assert budget_with == budget_without - count_approx_tokens(system_prompt) -# ── _stringify_message_content ─────────────────────────────────────────────── - - -def test_stringify_none_returns_empty(): - """None content stringifies to empty string.""" - assert _stringify_message_content(None) == "" +def test_safe_clean_tool_history_falls_back_to_raw_history(monkeypatch): + """Unexpected cleanup errors should not break the conversation flow.""" + messages = _build_turns(2) + def raise_cleanup(_messages): + raise RuntimeError("cleanup failed") -def test_stringify_str_returns_same(): - """String content is returned unchanged.""" - assert _stringify_message_content("hello") == "hello" + monkeypatch.setattr(history_processors, "clean_tool_history", raise_cleanup) + result = history_processors.safe_clean_tool_history(messages) -def test_stringify_list_joins(): - """List content items are joined with a space.""" - assert _stringify_message_content(["a", "b"]) == "a b" + assert result == messages diff --git a/src/backend/chat/tests/clients/pydantic_ai/test_document_context_window.py b/src/backend/chat/tests/clients/pydantic_ai/test_document_context_window.py index eaf9265f..379a1f34 100644 --- a/src/backend/chat/tests/clients/pydantic_ai/test_document_context_window.py +++ b/src/backend/chat/tests/clients/pydantic_ai/test_document_context_window.py @@ -38,6 +38,9 @@ def _parse_listing(instruction: str) -> dict: def _llm_config_with_context(settings): """Configure a model with max_token_context for context window tests.""" settings.DOCUMENT_CONTEXT_BUDGET_RATIO = 0.5 + # Pin the buffer so the budget math stays deterministic regardless of the + # environment default: int(4000 * 0.5) - 1000 = 1000. + settings.DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS = 1000 settings.LLM_CONFIGURATIONS = { "default-model": LLModel( hrid="default-model", diff --git a/src/backend/chat/tests/clients/pydantic_ai/test_run_agent_summary.py b/src/backend/chat/tests/clients/pydantic_ai/test_run_agent_summary.py new file mode 100644 index 00000000..a56a39de --- /dev/null +++ b/src/backend/chat/tests/clients/pydantic_ai/test_run_agent_summary.py @@ -0,0 +1,455 @@ +"""Tests for _run_agent summary-event wiring and history-summary persistence.""" +# pylint: disable=protected-access # tests intentionally access private members +# pylint: disable=using-constant-test,unreachable # if False: generator stubs + +from contextlib import ExitStack, asynccontextmanager, contextmanager +from datetime import timedelta +from unittest.mock import AsyncMock, MagicMock, patch + +from django.utils import timezone + +import pytest +from asgiref.sync import sync_to_async + +from chat.agents.history_processors import SummarizationRequiredError +from chat.ai_sdk_types import UIMessage +from chat.clients.pydantic_ai import AIAgentService, DocumentParsingResult +from chat.constants import HISTORY_SUMMARY_POLL_INTERVAL_SECONDS +from chat.factories import ChatConversationFactory +from chat.llm_configuration import LLModel, LLMProvider +from chat.vercel_ai_sdk.core import events_v4 + +pytestmark = pytest.mark.django_db() + + +@pytest.fixture(autouse=True) +def _llm_config(settings): + """Provide a valid model configuration so AIAgentService can be constructed.""" + settings.LLM_CONFIGURATIONS = { + "default-model": LLModel( + hrid="default-model", + model_name="amazing-llm", + human_readable_name="Amazing LLM", + is_active=True, + icon=None, + system_prompt="You are an amazing assistant.", + tools=[], + max_token_context=4000, + provider=LLMProvider( + hrid="unused", + base_url="https://example.com", + api_key="key", + ), + ), + } + + +@pytest.fixture(name="ui_messages") +def ui_messages_fixture(): + """A single user message, enough to drive _run_agent.""" + return [UIMessage(id="msg-1", role="user", content="Hello", parts=[])] + + +@pytest.mark.asyncio +async def test_build_model_history_trims_without_generating_or_claiming(): + """The turn only trims history: no LLM, no summary write, no claim taken.""" + conversation = await sync_to_async(ChatConversationFactory)() + service = AIAgentService(conversation, user=conversation.owner) + service._history_summary_checkpoint = 3 + + history = ["m0", "m1", "m2", "m3", "m4"] + with patch( + "chat.clients.pydantic_ai.build_active_history", + return_value=history[2:], + ) as trim: + result = service._build_model_history(history) + + assert result == history[2:] + trim.assert_called_once() + # No summary generated, nothing persisted, no claim taken. + assert service._history_summary is None + await service.conversation.arefresh_from_db() + assert service.conversation.history_summary_claimed_at is None + assert service.conversation.history_summary == "" + + +class _FakeRun: + """Minimal stand-in for a pydantic-ai agent run result.""" + + def __init__(self): + """Build a fake run exposing empty results and zero usage.""" + self.result = MagicMock() + self.result.new_messages.return_value = [] + self.result.output = "" + self.usage = MagicMock(input_tokens=0, output_tokens=0) + + +async def _empty_async_gen(*_args, **_kwargs): + """An async generator that yields nothing.""" + if False: # pragma: no cover - never yields, only makes this a generator + yield + + +@pytest.mark.asyncio +async def test_run_agent_emits_summary_events_when_summarization_triggers(ui_messages): + """Regression test: _run_agent must define should_emit_summary_event, tool_call_id + and conversation_has_own_documents. + + A dropped block during history rewriting previously left these undefined, crashing + every agent run with UnboundLocalError/NameError. All existing tests mock _run_agent + itself, so nothing exercised the real method. This drives the real _run_agent with the + summarization path enabled and asserts the running/done event pair brackets the run. + """ + conversation = await sync_to_async(ChatConversationFactory)() + service = AIAgentService(conversation, user=conversation.owner) + + async def fake_prepare(_messages): + """Return a stub agent-run preparation tuple.""" + usage = {"promptTokens": 0, "completionTokens": 0, "co2_impact": 0} + return ("prompt", [], [], [], usage, [], False) + + async def fake_handle_docs(*_args, **_kwargs): + """Yield a stub document-parsing result with no documents.""" + yield DocumentParsingResult(success=True, has_documents=False) + + @asynccontextmanager + async def fake_iter(*_args, **_kwargs): + """Yield a fake agent run as an async context manager.""" + yield _FakeRun() + + service.conversation_agent = MagicMock() + service.conversation_agent.iter = fake_iter + + with ( + patch.object(service, "_prepare_agent_run", side_effect=fake_prepare), + patch.object(service, "_handle_input_documents", side_effect=fake_handle_docs), + patch.object(service, "_build_model_history", return_value=[]), + patch.object(service, "_agent_stop_streaming", AsyncMock()), + patch.object(service, "_setup_self_documentation_tool"), + patch.object(service, "_setup_web_search_tool"), + patch.object(service, "_setup_web_search"), + patch.object(service, "_check_should_enable_rag", AsyncMock(return_value=False)), + patch.object(service, "_process_agent_nodes", side_effect=_empty_async_gen), + patch.object(service, "_finalize_conversation", side_effect=_empty_async_gen), + # True at start (phase triggers), False on the re-check (summary landed). + patch( + "chat.clients.pydantic_ai.should_generate_conversation_summary", + side_effect=[True, False], + ), + patch("chat.clients.pydantic_ai.get_mcp_servers", return_value=[]), + patch("chat.clients.pydantic_ai._extract_co2_from_usage", return_value=0), + patch.object(service, "_wait_for_history_summary", side_effect=_empty_async_gen), + ): + events = [event async for event in service._run_agent(ui_messages)] + + tool_calls = [e for e in events if isinstance(e, events_v4.ToolCallPart)] + tool_results = [e for e in events if isinstance(e, events_v4.ToolResultPart)] + + assert len(tool_calls) == 1 + assert tool_calls[0].tool_name == "summarize" + assert tool_calls[0].args["state"] == "running" + + assert len(tool_results) == 1 + assert tool_results[0].result == {"state": "done"} + assert tool_results[0].tool_call_id == tool_calls[0].tool_call_id + + +@pytest.mark.asyncio +async def test_stream_content_emits_summarization_failed_error(ui_messages): + """SummarizationRequiredError becomes an ErrorPart and persists the user message.""" + conversation = await sync_to_async(ChatConversationFactory)() + service = AIAgentService(conversation, user=conversation.owner) + + async def failing_run_agent(_messages, _force_web_search=False): + """Fake agent run that raises SummarizationRequiredError.""" + raise SummarizationRequiredError("model down") + yield # pragma: no cover - makes this an async generator + + with ( + patch.object(service, "_run_agent", side_effect=failing_run_agent), + patch.object(service, "_persist_user_message_on_error", AsyncMock()) as persist, + ): + chunks = [chunk async for chunk in service.stream_data_async(ui_messages)] + + assert any("summarization_failed" in chunk for chunk in chunks) + persist.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_run_agent_skips_summary_events_when_summarization_not_triggered(ui_messages): + """No summary events are emitted when summarization is not needed (and still no crash).""" + conversation = await sync_to_async(ChatConversationFactory)() + service = AIAgentService(conversation, user=conversation.owner) + + async def fake_prepare(_messages): + """Return a stub agent-run preparation tuple.""" + usage = {"promptTokens": 0, "completionTokens": 0, "co2_impact": 0} + return ("prompt", [], [], [], usage, [], False) + + async def fake_handle_docs(*_args, **_kwargs): + """Yield a stub document-parsing result with no documents.""" + yield DocumentParsingResult(success=True, has_documents=False) + + @asynccontextmanager + async def fake_iter(*_args, **_kwargs): + """Yield a fake agent run as an async context manager.""" + yield _FakeRun() + + service.conversation_agent = MagicMock() + service.conversation_agent.iter = fake_iter + + with ( + patch.object(service, "_prepare_agent_run", side_effect=fake_prepare), + patch.object(service, "_handle_input_documents", side_effect=fake_handle_docs), + patch.object(service, "_build_model_history", return_value=[]), + patch.object(service, "_agent_stop_streaming", AsyncMock()), + patch.object(service, "_setup_self_documentation_tool"), + patch.object(service, "_setup_web_search_tool"), + patch.object(service, "_setup_web_search"), + patch.object(service, "_check_should_enable_rag", AsyncMock(return_value=False)), + patch.object(service, "_process_agent_nodes", side_effect=_empty_async_gen), + patch.object(service, "_finalize_conversation", side_effect=_empty_async_gen), + patch("chat.clients.pydantic_ai.should_generate_conversation_summary", return_value=False), + patch("chat.clients.pydantic_ai.get_mcp_servers", return_value=[]), + patch("chat.clients.pydantic_ai._extract_co2_from_usage", return_value=0), + patch.object(service, "_wait_for_history_summary", side_effect=_empty_async_gen), + ): + events = [event async for event in service._run_agent(ui_messages)] + + assert not [e for e in events if isinstance(e, events_v4.ToolCallPart)] + assert not [e for e in events if isinstance(e, events_v4.ToolResultPart)] + + +@pytest.mark.asyncio +async def test_wait_for_history_summary_returns_when_checkpoint_advances(): + """The wait ends as soon as another worker's summary lands.""" + conversation = await sync_to_async(ChatConversationFactory)( + history_summary_claimed_at=timezone.now(), + ) + service = AIAgentService(conversation, user=conversation.owner) + ticks = 0 + + async def fake_sleep(_seconds): + """Patched sleep that lands the summary on the second poll tick.""" + nonlocal ticks + ticks += 1 + if ticks == 2: + # Simulate the task completing on another worker. + await sync_to_async(type(conversation).objects.filter(pk=conversation.pk).update)( + history_summary="done elsewhere", history_summary_checkpoint=8 + ) + + with patch("chat.clients.pydantic_ai.asyncio.sleep", side_effect=fake_sleep): + events = [event async for event in service._wait_for_history_summary()] + + assert ticks >= 2 + assert all(e.data == [{"type": "keep_alive"}] for e in events) + assert service.conversation.history_summary_checkpoint == 8 + + +@pytest.mark.asyncio +async def test_wait_for_history_summary_returns_immediately_on_dead_claim(): + """A claim past the TTL belongs to a dead worker: no waiting.""" + conversation = await sync_to_async(ChatConversationFactory)( + history_summary_claimed_at=timezone.now() - timedelta(seconds=181), + ) + service = AIAgentService(conversation, user=conversation.owner) + + with patch("chat.clients.pydantic_ai.asyncio.sleep") as sleep: + events = [event async for event in service._wait_for_history_summary()] + + sleep.assert_not_awaited() + assert events == [] + + +@pytest.mark.asyncio +async def test_wait_for_history_summary_waits_for_claim_grace_deadline(): + """No live claim but a grace deadline: the wait holds until the deadline passes.""" + conversation = await sync_to_async(ChatConversationFactory)() + service = AIAgentService(conversation, user=conversation.owner) + fake_now = 0.0 + + async def fake_sleep(seconds): + """Patched sleep that advances the fake monotonic clock.""" + nonlocal fake_now + fake_now += seconds + + # Deadline three poll intervals out: keep-alive fire at t=0, τ and 2τ, + # then the loop returns once t reaches the 3τ deadline. + claim_deadline = 3 * HISTORY_SUMMARY_POLL_INTERVAL_SECONDS + with ( + patch("chat.clients.pydantic_ai.asyncio.sleep", side_effect=fake_sleep), + patch("chat.clients.pydantic_ai.time.monotonic", side_effect=lambda: fake_now), + ): + events = [ + event + async for event in service._wait_for_history_summary(claim_deadline=claim_deadline) + ] + + assert len(events) == 3 + assert all(e.data == [{"type": "keep_alive"}] for e in events) + + +@pytest.mark.asyncio +async def test_wait_for_history_summary_grace_ends_early_when_summary_lands(): + """A checkpoint advance ends the grace wait immediately.""" + conversation = await sync_to_async(ChatConversationFactory)() + service = AIAgentService(conversation, user=conversation.owner) + fake_now = 0.0 + + async def fake_sleep(seconds): + """Patched sleep that lands the summary on the first poll tick.""" + nonlocal fake_now + first_tick = fake_now == 0.0 + fake_now += seconds + if first_tick: + await sync_to_async(type(conversation).objects.filter(pk=conversation.pk).update)( + history_summary="landed", history_summary_checkpoint=6 + ) + + with ( + patch("chat.clients.pydantic_ai.asyncio.sleep", side_effect=fake_sleep), + patch("chat.clients.pydantic_ai.time.monotonic", side_effect=lambda: fake_now), + ): + events = [event async for event in service._wait_for_history_summary(claim_deadline=30.0)] + + assert len(events) == 1 + assert service.conversation.history_summary_checkpoint == 6 + + +@contextmanager +def _run_agent_patches(service): + """The boilerplate patch stack for driving the real _run_agent.""" + + async def fake_prepare(_messages): + """Return a stub agent-run preparation tuple.""" + usage = {"promptTokens": 0, "completionTokens": 0, "co2_impact": 0} + return ("prompt", [], [], [], usage, [], False) + + async def fake_handle_docs(*_args, **_kwargs): + """Yield a stub document-parsing result with no documents.""" + yield DocumentParsingResult(success=True, has_documents=False) + + @asynccontextmanager + async def fake_iter(*_args, **_kwargs): + """Yield a fake agent run as an async context manager.""" + yield _FakeRun() + + service.conversation_agent = MagicMock() + service.conversation_agent.iter = fake_iter + + patches = ( + patch.object(service, "_prepare_agent_run", side_effect=fake_prepare), + patch.object(service, "_handle_input_documents", side_effect=fake_handle_docs), + patch.object(service, "_build_model_history", return_value=[]), + patch.object(service, "_agent_stop_streaming", AsyncMock()), + patch.object(service, "_setup_self_documentation_tool"), + patch.object(service, "_setup_web_search_tool"), + patch.object(service, "_setup_web_search"), + patch.object(service, "_check_should_enable_rag", AsyncMock(return_value=False)), + patch.object(service, "_process_agent_nodes", side_effect=_empty_async_gen), + patch.object(service, "_finalize_conversation", side_effect=_empty_async_gen), + patch("chat.clients.pydantic_ai.get_mcp_servers", return_value=[]), + patch("chat.clients.pydantic_ai._extract_co2_from_usage", return_value=0), + ) + with ExitStack() as stack: + for item in patches: + stack.enter_context(item) + yield + + +@pytest.mark.asyncio +async def test_run_agent_enqueues_task_and_adopts_its_summary(ui_messages): + """The over-budget turn enqueues the Celery task, waits, and adopts the result.""" + conversation = await sync_to_async(ChatConversationFactory)() + service = AIAgentService(conversation, user=conversation.owner) + + async def fake_wait(claim_deadline=None): # pylint: disable=unused-argument + """Fake wait that lands a summary on the conversation while the turn waits.""" + # Simulate the task completing while the turn waited. + await sync_to_async(type(conversation).objects.filter(pk=conversation.pk).update)( + history_summary="task summary", history_summary_checkpoint=6 + ) + await service.conversation.arefresh_from_db() + if False: # pragma: no cover - never yields, only makes this a generator + yield + + # True at turn start (phase triggers), False after adopting the summary. + should_generate = MagicMock(side_effect=[True, False]) + + with ( + _run_agent_patches(service), + patch( + "chat.clients.pydantic_ai.should_generate_conversation_summary", + should_generate, + ), + patch("chat.clients.pydantic_ai.summarize_conversation_history") as task, + patch.object(service, "_wait_for_history_summary", side_effect=fake_wait), + ): + events = [event async for event in service._run_agent(ui_messages)] + + task.delay.assert_called_once_with(str(conversation.pk)) + assert service._history_summary == "task summary" + assert service._history_summary_checkpoint == 6 + tool_results = [e for e in events if isinstance(e, events_v4.ToolResultPart)] + assert [r.result for r in tool_results] == [{"state": "done"}] + + +@pytest.mark.asyncio +async def test_run_agent_does_not_enqueue_when_claim_is_live(ui_messages): + """Another actor is already summarizing: wait for it, never double-enqueue.""" + conversation = await sync_to_async(ChatConversationFactory)( + history_summary_claimed_at=timezone.now(), + ) + service = AIAgentService(conversation, user=conversation.owner) + + async def fake_wait(claim_deadline=None): # pylint: disable=unused-argument + """Fake wait that lands a summary on the conversation while the turn waits.""" + await sync_to_async(type(conversation).objects.filter(pk=conversation.pk).update)( + history_summary="from the other actor", + history_summary_checkpoint=6, + history_summary_claimed_at=None, + ) + await service.conversation.arefresh_from_db() + if False: # pragma: no cover + yield + + should_generate = MagicMock(side_effect=[True, False]) + + with ( + _run_agent_patches(service), + patch( + "chat.clients.pydantic_ai.should_generate_conversation_summary", + should_generate, + ), + patch("chat.clients.pydantic_ai.summarize_conversation_history") as task, + patch.object(service, "_wait_for_history_summary", side_effect=fake_wait), + ): + _ = [event async for event in service._run_agent(ui_messages)] + + task.delay.assert_not_called() + assert service._history_summary == "from the other actor" + + +@pytest.mark.asyncio +async def test_run_agent_fails_when_summary_never_lands(ui_messages): + """Grace expires with nothing landed and still over budget: the turn fails, no degraded turn.""" + conversation = await sync_to_async(ChatConversationFactory)() + service = AIAgentService(conversation, user=conversation.owner) + + # Still over budget at turn start AND after the fruitless wait. + should_generate = MagicMock(side_effect=[True, True]) + + with ( + _run_agent_patches(service), + patch( + "chat.clients.pydantic_ai.should_generate_conversation_summary", + should_generate, + ), + patch("chat.clients.pydantic_ai.summarize_conversation_history") as task, + patch.object(service, "_wait_for_history_summary", side_effect=_empty_async_gen), + ): + with pytest.raises(SummarizationRequiredError): + _ = [event async for event in service._run_agent(ui_messages)] + + task.delay.assert_called_once_with(str(conversation.pk)) diff --git a/src/backend/chat/tests/test_assistant_health.py b/src/backend/chat/tests/test_assistant_health.py index 9765cbae..b41c6eb0 100644 --- a/src/backend/chat/tests/test_assistant_health.py +++ b/src/backend/chat/tests/test_assistant_health.py @@ -233,6 +233,7 @@ def test_main_red_fb1_unknown_hrid_unavailable(settings): assert result["blocked"] is True +@pytest.mark.django_db def test_main_red_fb1_empty_fb2_red_unavailable(settings, llm_configs): # fb1="" (not configured → down) + fb2=red (down) → all_down=True → unavailable. settings.LLM_FALLBACK_MODEL_HRID_2 = "fallback-2" diff --git a/src/backend/chat/tests/test_models_history_summary_claim.py b/src/backend/chat/tests/test_models_history_summary_claim.py new file mode 100644 index 00000000..c23f0855 --- /dev/null +++ b/src/backend/chat/tests/test_models_history_summary_claim.py @@ -0,0 +1,82 @@ +"""Tests for the history-summarization claim on ChatConversation.""" + +from datetime import timedelta + +from django.utils import timezone + +import pytest + +from chat.factories import ChatConversationFactory + +pytestmark = pytest.mark.django_db() + + +def test_claim_succeeds_when_unclaimed(): + """Claim is granted and recorded when no prior claim exists.""" + conversation = ChatConversationFactory() + assert conversation.claim_history_summarization() is True + conversation.refresh_from_db() + assert conversation.history_summary_claimed_at is not None + assert conversation.history_summarization_claim_is_live + + +def test_claim_fails_while_another_claim_is_live(): + """A second concurrent claim is rejected while the first is still active.""" + conversation = ChatConversationFactory() + assert conversation.claim_history_summarization() is True + other = type(conversation).objects.get(pk=conversation.pk) + assert other.claim_history_summarization() is False + + +def test_claim_succeeds_over_an_expired_claim(): + """A claim older than the TTL belongs to a provably dead worker.""" + conversation = ChatConversationFactory( + history_summary_claimed_at=timezone.now() - timedelta(seconds=181) + ) + assert not conversation.history_summarization_claim_is_live + assert conversation.claim_history_summarization() is True + + +def test_release_clears_the_claim(): + """Releasing a claim sets history_summary_claimed_at back to None.""" + conversation = ChatConversationFactory() + conversation.claim_history_summarization() + conversation.release_history_summarization_claim() + conversation.refresh_from_db() + assert conversation.history_summary_claimed_at is None + + +def test_release_is_a_noop_when_another_worker_reclaimed(): + """A stale worker's release must not wipe a newer worker's live claim.""" + conversation = ChatConversationFactory() + conversation.claim_history_summarization() + + # A second worker reclaims (e.g. after the first worker's TTL elapsed). + reclaimed_at = timezone.now() + other = type(conversation).objects.get(pk=conversation.pk) + type(other).objects.filter(pk=other.pk).update(history_summary_claimed_at=reclaimed_at) + + # The first worker releases: it no longer owns the claim, so this is a no-op. + conversation.release_history_summarization_claim() + + conversation.refresh_from_db() + assert conversation.history_summary_claimed_at == reclaimed_at + + +def test_persist_history_summary_advances_checkpoint(): + """Persisting a summary with a newer checkpoint succeeds and saves both fields.""" + conversation = ChatConversationFactory() + assert conversation.persist_history_summary("a summary", 4) is True + conversation.refresh_from_db() + assert conversation.history_summary == "a summary" + assert conversation.history_summary_checkpoint == 4 + + +def test_persist_history_summary_rejects_stale_checkpoint(): + """Late/duplicate task completions must be no-ops (write-back guard).""" + conversation = ChatConversationFactory(history_summary="newer", history_summary_checkpoint=6) + assert conversation.persist_history_summary("older", 6) is False + assert conversation.persist_history_summary("older", 4) is False + conversation.refresh_from_db() + assert conversation.history_summary == "newer" + assert conversation.history_summary_checkpoint == 6 diff --git a/src/backend/chat/tests/test_tasks.py b/src/backend/chat/tests/test_tasks.py index b9b50eac..018b96d8 100644 --- a/src/backend/chat/tests/test_tasks.py +++ b/src/backend/chat/tests/test_tasks.py @@ -4,18 +4,29 @@ `.delay()` executes the task synchronously in-process. """ +from unittest.mock import patch + from django.core.files.base import ContentFile from django.core.files.storage import default_storage +from django.utils import timezone import pytest import responses +from pydantic_ai.exceptions import ModelHTTPError from rest_framework import status from core.file_upload.enums import AttachmentStatus from chat import factories +from chat.agents.history_processors import SummarizationRequiredError from chat.enums import AttachmentIndexState -from chat.tasks import index_project_attachment_task +from chat.factories import ChatConversationFactory +from chat.llm_configuration import LLModel, LLMProvider +from chat.tasks import ( + RETRYABLE_SUMMARIZATION_ERRORS, + index_project_attachment_task, + summarize_conversation_history, +) pytestmark = pytest.mark.django_db @@ -66,3 +77,109 @@ def test_index_project_attachment_task_indexes_attachment(): def test_index_project_attachment_task_ignores_missing_attachment(): """A deleted attachment id is logged and skipped, never raised.""" index_project_attachment_task.delay(999999) + + +@pytest.fixture(name="_llm_config", autouse=True) +def llm_config_fixture(settings): + """Configure a single active model for the summarization task tests.""" + settings.LLM_CONFIGURATIONS = { + "default-model": LLModel( + hrid="default-model", + model_name="amazing-llm", + human_readable_name="Amazing LLM", + is_active=True, + icon=None, + system_prompt="You are an amazing assistant.", + tools=[], + max_token_context=4000, + provider=LLMProvider(hrid="p", base_url="https://example.com", api_key="k"), + ), + } + settings.LLM_DEFAULT_MODEL_HRID = "default-model" + settings.DOCUMENT_CONTEXT_SECURITY_BUFFER_TOKENS = 1000 + settings.DOCUMENT_CONTEXT_BUDGET_RATIO = 0.5 + + +def _long_history(turns: int = 30) -> list: + """Raw pydantic message dicts long enough to exceed any small budget.""" + messages = [] + for i in range(turns): + messages.append( + { + "kind": "request", + "parts": [{"part_kind": "user-prompt", "content": f"user {i} " + "x " * 200}], + } + ) + messages.append( + { + "kind": "response", + "parts": [{"part_kind": "text", "content": f"assistant {i} " + "y " * 200}], + } + ) + return messages + + +def test_task_generates_summary_and_advances_checkpoint(): + """Happy path: summary is generated and persisted, claim released.""" + conversation = ChatConversationFactory(messages=[], pydantic_messages=_long_history()) + + with patch( + "chat.tasks.generate_history_summary", return_value=("task summary", 60) + ) as summarize: + summarize_conversation_history(str(conversation.pk)) + + summarize.assert_called_once() + conversation.refresh_from_db() + assert conversation.history_summary == "task summary" + assert conversation.history_summary_checkpoint == 60 + assert conversation.history_summary_claimed_at is None # released + + +def test_task_exits_when_claim_is_held(): + """Task is a no-op when another worker holds an active claim.""" + conversation = ChatConversationFactory( + pydantic_messages=_long_history(), + history_summary_claimed_at=timezone.now(), + ) + with patch("chat.tasks.generate_history_summary") as summarize: + summarize_conversation_history(str(conversation.pk)) + summarize.assert_not_called() + + +def test_task_noop_when_under_budget(): + """Task is a no-op when history fits within the token budget.""" + conversation = ChatConversationFactory(pydantic_messages=_long_history(turns=1)) + with patch("chat.tasks.generate_history_summary") as summarize: + summarize_conversation_history(str(conversation.pk)) + summarize.assert_not_called() + conversation.refresh_from_db() + assert conversation.history_summary_claimed_at is None # released + + +def test_task_releases_claim_and_raises_on_failure(): + """A non-retryable failure propagates and the claim is still released.""" + conversation = ChatConversationFactory(pydantic_messages=_long_history()) + + with ( + patch( + "chat.tasks.generate_history_summary", + side_effect=RuntimeError("provider down"), + ), + pytest.raises(RuntimeError), + ): + summarize_conversation_history(str(conversation.pk)) + + conversation.refresh_from_db() + assert conversation.history_summary_claimed_at is None + + +def test_retryable_errors_exclude_deterministic_failures(): + """Only transient provider errors are retried; deterministic ones fail fast.""" + assert issubclass(ModelHTTPError, RETRYABLE_SUMMARIZATION_ERRORS) + assert not issubclass(SummarizationRequiredError, RETRYABLE_SUMMARIZATION_ERRORS) + assert not issubclass(RuntimeError, RETRYABLE_SUMMARIZATION_ERRORS) + + +def test_task_ignores_deleted_conversation(): + """Task exits silently for a non-existent conversation PK.""" + summarize_conversation_history("00000000-0000-0000-0000-000000000000") diff --git a/src/backend/chat/tools/descriptions.py b/src/backend/chat/tools/descriptions.py index 1402f076..d198c441 100644 --- a/src/backend/chat/tools/descriptions.py +++ b/src/backend/chat/tools/descriptions.py @@ -142,3 +142,25 @@ - "Explain how a for loop works" - "What does this report say about Q3 revenue?" """ + + +CONVERSATION_SUMMARY_TOOL_DESCRIPTION = ( + "You are a conversation summarization assistant. Your role is to maintain\n" + "a concise and accurate running summary of a conversation, " + "omitting small talk and unrelated topics.\n\n" + "Given the previous summary (if any) and the new exchanges provided,\n" + "generate an updated summary that:\n\n" + "- **Preserves** every key information, decisions, and important facts\n" + "- **Integrates** the new exchanges in a coherent way\n" + "- **Removes** redundant or non-essential details\n" + "- **Maintains** the context needed for the conversation to continue\n" + "- Is written in a neutral, factual, third-person style\n" + "- Stays **concise** (5-10 lines maximum)\n\n" + "## Previous Summary:\n" + "{latest_summary}\n\n" + "## New Exchanges:\n" + "{exchanges}\n\n" + "Only answer with the updated summary, including the new exchanges " + "information and the previous summary.\n\n" + "## Updated Summary:\n" +) diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index 6fa0eb7d..1966a3b9 100755 --- a/src/backend/conversations/settings.py +++ b/src/backend/conversations/settings.py @@ -708,6 +708,25 @@ class Base(BraveSettings, Configuration): environ_name="DEFAULT_ALLOW_SMART_WEB_SEARCH", environ_prefix=None, ) + # Conversation summary: at the start of a new user turn (before agent.iter), when the + # active slice of stored pydantic_messages (previous turns only) exceeds the message + # token budget (int(max_token_context * (1 - DOCUMENT_CONTEXT_BUDGET_RATIO)) - security_buffer). + # The incoming user message is intentionally NOT counted in this check: it is not + # persisted yet, so the check stays on stored history. A turn tipped over budget by the + # new message alone is caught on the next turn; the security buffer absorbs the overflow + # meanwhile (see ADR 0002). That history usually ends + # on an assistant ModelResponse. After a summary, keep the last N ModelMessage + # entries before the checkpoint. Use an even N so the window starts on a user message. + CONVERSATION_SUMMARY_CONTEXT_MESSAGES = values.PositiveIntegerValue( + default=10, + environ_name="CONVERSATION_SUMMARY_CONTEXT_MESSAGES", + environ_prefix=None, + ) + CONVERSATION_SUMMARY_MAX_TOKENS = values.PositiveIntegerValue( + default=2048, + environ_name="CONVERSATION_SUMMARY_MAX_TOKENS", + environ_prefix=None, + ) # These settings are default values used in the default LLM_CONFIGURATIONS # They allow a deployment with only one model without a specific configuration file diff --git a/src/frontend/apps/conversations/src/features/chat/api/__tests__/useChat.test.tsx b/src/frontend/apps/conversations/src/features/chat/api/__tests__/useChat.test.tsx index a4618c59..64a75c25 100644 --- a/src/frontend/apps/conversations/src/features/chat/api/__tests__/useChat.test.tsx +++ b/src/frontend/apps/conversations/src/features/chat/api/__tests__/useChat.test.tsx @@ -27,7 +27,7 @@ describe('isImagesSkippedEvent', () => { }); it('rejects events of other types', () => { - expect(isImagesSkippedEvent({ type: 'context_trimmed' })).toBe(false); + expect(isImagesSkippedEvent({ type: 'some_other_event' })).toBe(false); }); it('rejects events with unknown kind', () => { diff --git a/src/frontend/apps/conversations/src/features/chat/api/useChat.tsx b/src/frontend/apps/conversations/src/features/chat/api/useChat.tsx index d3d9e392..b847f223 100644 --- a/src/frontend/apps/conversations/src/features/chat/api/useChat.tsx +++ b/src/frontend/apps/conversations/src/features/chat/api/useChat.tsx @@ -87,21 +87,6 @@ async function fetchChatCooldown(): Promise<{ cooldown_seconds: number }> { return response.json() as Promise<{ cooldown_seconds: number }>; } -interface ContextTrimmedEvent { - type: 'context_trimmed'; -} - -export function isContextTrimmedEvent( - item: unknown, -): item is ContextTrimmedEvent { - return ( - typeof item === 'object' && - item !== null && - 'type' in item && - item.type === 'context_trimmed' - ); -} - // Stream-protocol contract with the backend. Mirrored in // ``pydantic_ai.py`` (``IMAGES_SKIPPED_EVENT_TYPE`` / // ``IMAGE_SKIP_REASON_TEXT_ONLY``). Keep both sides in sync when adding new diff --git a/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx b/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx index 6583f313..06c53504 100644 --- a/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx +++ b/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx @@ -20,7 +20,6 @@ import { useProjectAttachments } from '@/features/attachments/api/useProjectAtta import { useReindexProjectAttachment } from '@/features/attachments/api/useReindexProjectAttachment'; import { useUploadFile } from '@/features/attachments/hooks/useUploadFile'; import { - isContextTrimmedEvent, isImagesSkippedEvent, stampImagesSkippedOnLatestUserMessage, useChat, @@ -59,6 +58,7 @@ const PROVIDER_ERROR_CODES = new Set([ 'model_not_found', 'model_wrong_type', 'model_busy', + 'summarization_failed', ]); const IMAGES_BANNER_STORAGE_PREFIX = 'conversations:images-banner-dismissed:'; @@ -213,7 +213,6 @@ export const Chat = ({ const { mutate: createChatConversation } = useCreateChatConversation(); const queryClient = useQueryClient(); const [isReadingInstructions, setIsReadingInstructions] = useState(false); - const [contextTrimmed, setContextTrimmed] = useState(false); const readingInstructionsStartRef = useRef(0); const aprilFools = useAprilFools(); @@ -668,15 +667,8 @@ export const Chat = ({ // eslint-disable-next-line react-hooks/exhaustive-deps }, [shouldRetry, input, files]); - useEffect(() => { - setContextTrimmed(false); - }, [conversationId]); - useEffect(() => { if (!data || !Array.isArray(data)) return; - if (data.some(isContextTrimmedEvent)) { - setContextTrimmed(true); - } if ( data.some( (item) => isImagesSkippedEvent(item) && item.kind === 'chat_notice', @@ -979,6 +971,8 @@ export const Chat = ({ isFirstConversationMessage={isFirstConversationMessage} streamingMessageHeight={streamingMessageHeight} status={status} + chatErrorType={chatErrorType} + onRetry={handleRetry} conversationId={conversationId} isSourceOpen={isSourceOpen} isMobile={isMobile} @@ -1013,21 +1007,6 @@ export const Chat = ({ })} )} - {contextTrimmed && ( - - - {t("Some older messages are no longer in the model's context.")} - - - )} {!aprilFools.isActive && ((status !== 'ready' && status !== 'streaming' && status !== 'error') || isUploadingFiles) ? ( @@ -1054,13 +1033,15 @@ export const Chat = ({ ) : null} - {status === 'error' && !isUploadingFiles && ( - - )} + {status === 'error' && + !isUploadingFiles && + chatErrorType !== 'summarization_failed' && ( + + )} - {!isProviderError && ( - - )} + {!isProviderError && } {!isProviderError && !hasLastSubmission && ( + ); +}; diff --git a/src/frontend/apps/conversations/src/features/chat/components/SummarizationError.tsx b/src/frontend/apps/conversations/src/features/chat/components/SummarizationError.tsx new file mode 100644 index 00000000..60bdf397 --- /dev/null +++ b/src/frontend/apps/conversations/src/features/chat/components/SummarizationError.tsx @@ -0,0 +1,28 @@ +import { useTranslation } from 'react-i18next'; + +import { Box, Text } from '@/components'; + +import { RetryButton } from './RetryButton'; +import { getChatErrorMessage } from './chatErrorMessages'; + +interface SummarizationErrorProps { + onRetry: () => void; +} + +// Rendered in the same slot as SummarizationProgress: when the summary can't +// be prepared, the progress bar is replaced in place by this failure notice +// plus a Retry button that resends the same message. +export const SummarizationError = ({ onRetry }: SummarizationErrorProps) => { + const { t } = useTranslation(); + + return ( + + + {getChatErrorMessage(t, 'summarization_failed')} + + + + + + ); +}; diff --git a/src/frontend/apps/conversations/src/features/chat/components/SummarizationProgress.tsx b/src/frontend/apps/conversations/src/features/chat/components/SummarizationProgress.tsx new file mode 100644 index 00000000..43ad9c9f --- /dev/null +++ b/src/frontend/apps/conversations/src/features/chat/components/SummarizationProgress.tsx @@ -0,0 +1,80 @@ +import React from 'react'; +import { useTranslation } from 'react-i18next'; + +import { Box, Text } from '@/components'; + +// Fake, purely time-based progress (the backend reports no real progress): +// 1 - e^(-t/τ), capped below 100% until the summarization actually completes. +// τ is sized to the worst-case summarization latency (near the ~110s soft time +// limit) so the bar keeps creeping instead of parking at the cap on long runs. +const TIME_CONSTANT_MS = 30000; +const MAX_RATIO = 0.95; +const TICK_MS = 100; +const HIDE_DELAY_MS = 400; + +interface SummarizationProgressProps { + done: boolean; +} + +export const SummarizationProgress = ({ done }: SummarizationProgressProps) => { + const { t } = useTranslation(); + const [progress, setProgress] = React.useState(0); + const [hidden, setHidden] = React.useState(false); + + React.useEffect(() => { + if (done) { + return; + } + const startedAt = Date.now(); + const interval = setInterval(() => { + const elapsed = Date.now() - startedAt; + setProgress(MAX_RATIO * (1 - Math.exp(-elapsed / TIME_CONSTANT_MS))); + }, TICK_MS); + return () => clearInterval(interval); + }, [done]); + + React.useEffect(() => { + if (!done) { + return; + } + setProgress(1); + const timeout = setTimeout(() => setHidden(true), HIDE_DELAY_MS); + return () => clearTimeout(timeout); + }, [done]); + + if (hidden) { + return null; + } + + return ( + + + {t('Summarizing conversation...')} + + + + + + ); +}; diff --git a/src/frontend/apps/conversations/src/features/chat/components/__tests__/MessageItem.test.tsx b/src/frontend/apps/conversations/src/features/chat/components/__tests__/MessageItem.test.tsx index 959493de..10f957c4 100644 --- a/src/frontend/apps/conversations/src/features/chat/components/__tests__/MessageItem.test.tsx +++ b/src/frontend/apps/conversations/src/features/chat/components/__tests__/MessageItem.test.tsx @@ -539,4 +539,119 @@ describe('MessageItem', () => { expect(screen.queryByTestId('feedback-buttons')).not.toBeInTheDocument(); }); }); + + describe('summarization progress', () => { + const conversationSummarizeMessage = { + id: 'msg-sum', + role: 'assistant' as const, + content: '', + parts: [ + { + type: 'tool-invocation' as const, + toolInvocation: { + toolCallId: 'call-1', + toolName: 'summarize', + state: 'call' as const, + args: { state: 'running', summary_scope: 'conversation' }, + }, + }, + ], + }; + + it('renders the progress bar while conversation summarization runs', async () => { + await act(async () => { + renderWithProviders( + , + ); + }); + + expect(screen.getByTestId('summarization-progress')).toBeInTheDocument(); + // The generic spinner branch must not double-render for this tool + // (exact-string queries: 'Summarizing...' does not match the progress + // bar's 'Summarizing conversation...' label). + expect(screen.queryByText('Summarizing...')).not.toBeInTheDocument(); + expect(screen.queryByText('Search...')).not.toBeInTheDocument(); + }); + + it('keeps the generic spinner for document-scope summarize', async () => { + const documentSummarizeMessage = { + ...conversationSummarizeMessage, + parts: [ + { + type: 'tool-invocation' as const, + toolInvocation: { + toolCallId: 'call-2', + toolName: 'summarize', + state: 'call' as const, + args: { state: 'running', summary_scope: 'document' }, + }, + }, + ], + }; + + await act(async () => { + renderWithProviders( + , + ); + }); + + expect(screen.getByText('Summarizing...')).toBeInTheDocument(); + expect( + screen.queryByTestId('summarization-progress'), + ).not.toBeInTheDocument(); + }); + + it('replaces the progress bar with the error and retry when summarization fails', async () => { + const onRetry = jest.fn(); + await act(async () => { + renderWithProviders( + , + ); + }); + + expect(screen.getByTestId('summarization-error')).toBeInTheDocument(); + expect( + screen.queryByTestId('summarization-progress'), + ).not.toBeInTheDocument(); + + await userEvent.click(screen.getByRole('button', { name: /retry/i })); + expect(onRetry).toHaveBeenCalledTimes(1); + }); + + it('does not render the summarization error for other error types', async () => { + await act(async () => { + renderWithProviders( + , + ); + }); + + expect( + screen.queryByTestId('summarization-error'), + ).not.toBeInTheDocument(); + }); + }); }); diff --git a/src/frontend/apps/conversations/src/features/chat/components/__tests__/SummarizationError.test.tsx b/src/frontend/apps/conversations/src/features/chat/components/__tests__/SummarizationError.test.tsx new file mode 100644 index 00000000..7fea9b30 --- /dev/null +++ b/src/frontend/apps/conversations/src/features/chat/components/__tests__/SummarizationError.test.tsx @@ -0,0 +1,25 @@ +import { fireEvent, render, screen } from '@testing-library/react'; + +import '@/i18n/initI18n'; +import { AppWrapper } from '@/tests/utils'; + +import { SummarizationError } from '../SummarizationError'; + +describe('SummarizationError', () => { + it('renders the failure message and a retry button', () => { + render(, { wrapper: AppWrapper }); + + expect( + screen.getByText('Summarization failed. Please try again later.'), + ).toBeInTheDocument(); + expect(screen.getByRole('button', { name: /retry/i })).toBeInTheDocument(); + }); + + it('calls onRetry when the retry button is clicked', () => { + const onRetry = jest.fn(); + render(, { wrapper: AppWrapper }); + + fireEvent.click(screen.getByRole('button', { name: /retry/i })); + expect(onRetry).toHaveBeenCalledTimes(1); + }); +}); diff --git a/src/frontend/apps/conversations/src/features/chat/components/__tests__/SummarizationProgress.test.tsx b/src/frontend/apps/conversations/src/features/chat/components/__tests__/SummarizationProgress.test.tsx new file mode 100644 index 00000000..b354a57e --- /dev/null +++ b/src/frontend/apps/conversations/src/features/chat/components/__tests__/SummarizationProgress.test.tsx @@ -0,0 +1,63 @@ +import { act, render, screen } from '@testing-library/react'; + +import { SummarizationProgress } from '../SummarizationProgress'; + +jest.mock('react-i18next', () => ({ + useTranslation: () => ({ + t: (key: string) => key, + }), +})); + +describe('SummarizationProgress', () => { + beforeEach(() => { + jest.useFakeTimers(); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('renders the label and starts near zero', () => { + render(); + + expect(screen.getByText('Summarizing conversation...')).toBeInTheDocument(); + const fill = screen.getByTestId('summarization-progress-fill'); + expect(parseInt(fill.style.width, 10)).toBeLessThanOrEqual(5); + }); + + it('advances along the logarithmic curve without reaching 100%', () => { + render(); + + act(() => { + jest.advanceTimersByTime(30000); // one time constant + }); + const fill = screen.getByTestId('summarization-progress-fill'); + const afterOneTau = parseInt(fill.style.width, 10); + expect(afterOneTau).toBeGreaterThan(50); // ~60% at t = τ + expect(afterOneTau).toBeLessThan(95); + + act(() => { + jest.advanceTimersByTime(120000); + }); + expect(parseInt(fill.style.width, 10)).toBeLessThanOrEqual(95); + }); + + it('snaps to 100% on done and hides shortly after', () => { + const { rerender } = render(); + act(() => { + jest.advanceTimersByTime(3000); + }); + + rerender(); + expect(screen.getByTestId('summarization-progress-fill').style.width).toBe( + '100%', + ); + + act(() => { + jest.advanceTimersByTime(400); + }); + expect( + screen.queryByTestId('summarization-progress'), + ).not.toBeInTheDocument(); + }); +}); diff --git a/src/frontend/apps/conversations/src/features/chat/components/chatErrorMessages.ts b/src/frontend/apps/conversations/src/features/chat/components/chatErrorMessages.ts index ec30d792..072b52f1 100644 --- a/src/frontend/apps/conversations/src/features/chat/components/chatErrorMessages.ts +++ b/src/frontend/apps/conversations/src/features/chat/components/chatErrorMessages.ts @@ -26,6 +26,7 @@ export const getChatErrorMessage = ( model_busy: t( 'The AI inference provider is too busy. Please try again later.', ), + summarization_failed: t('Summarization failed. Please try again later.'), }; return messages[errorType]; };