Improvements

simonredfern · simonredfern · commit eea9062d6e51 · 2026-05-22T09:06:43.000+02:00
diff --git a/src/agent/components/nodes.py b/src/agent/components/nodes.py
@@ -88,16 +88,20 @@ async def run_summary_chain(state: OpeyGraphState):
 
     messages = state["messages"]
 
-    # Build a copy of messages with oversized ToolMessage content truncated, so the
-    # summarizer can't blow past the model's context window on a single huge tool result.
-    # The original `messages` in state is unchanged; only the copy sent to the summarizer
-    # is truncated. Summarization is compression, so truncated tool output is acceptable.
-    MAX_TOOL_CONTENT_CHARS = 30000
+    # The summarizer is the recovery path for an oversized conversation, so its OWN
+    # input must be hard-bounded — otherwise summarizing a 375k-token thread sends
+    # 375k tokens to the LLM and overflows (the failure this fixes). Truncate EVERY
+    # message's content with a per-message cap that scales down as the message count
+    # grows, bounding the total to ~SUMMARY_TOTAL_CHAR_BUDGET regardless of thread
+    # size. No messages are dropped, so tool_use/tool_result pairing stays valid.
+    # The original `messages` in state is unchanged; only this copy is truncated.
+    SUMMARY_TOTAL_CHAR_BUDGET = 400_000  # ~100k tokens
+    per_message_cap = max(200, SUMMARY_TOTAL_CHAR_BUDGET // max(len(messages), 1))
     messages_for_summary = []
     for msg in messages:
-        if isinstance(msg, ToolMessage) and isinstance(msg.content, str) and len(msg.content) > MAX_TOOL_CONTENT_CHARS:
-            truncated = msg.content[:MAX_TOOL_CONTENT_CHARS] + "\n\n[TRUNCATED TOOL RESPONSE FOR SUMMARIZATION]"
-            messages_for_summary.append(msg.model_copy(update={"content": truncated}))
+        new_content, was_truncated = _truncate_tool_content(msg.content, per_message_cap)
+        if was_truncated:
+            messages_for_summary.append(msg.model_copy(update={"content": new_content}))
         else:
             messages_for_summary.append(msg)
 
@@ -115,6 +119,16 @@ async def run_summary_chain(state: OpeyGraphState):
         include_system=True
     )
 
+    # trim_messages returns an empty list when even the single most-recent message
+    # exceeds max_tokens. That would delete the entire conversation — including the
+    # user's current question — leaving the next Opey call with only the summary
+    # SystemMessage, which Anthropic rejects ("at least one message is required").
+    # Always keep at least the last message; the repair logic below pulls in its
+    # tool_use/tool_result siblings so the kept set stays valid.
+    if not trimmed_messages and messages:
+        logger.warning("trim_messages returned empty; keeping the last message to avoid an empty Opey call")
+        trimmed_messages = messages[-1:]
+
     # Build an index: tool_call_id -> parent AIMessage (from full conversation)
     tool_call_id_to_ai_msg: Dict[str, AIMessage] = {}
     for msg in messages:
diff --git a/src/agent/components/recovery.py b/src/agent/components/recovery.py
@@ -43,19 +43,36 @@
 HARD_RECAP_CHARS = 1000
 
 
+_OVERFLOW_PHRASES = (
+    "prompt is too long",        # Anthropic
+    "context_length_exceeded",   # OpenAI
+    "maximum context length",    # OpenAI (alt phrasing)
+    "context window",            # Generic / Bedrock
+    "too many tokens",           # Generic fallback
+)
+
+
 def _is_context_overflow(exc: BaseException) -> bool:
-    """True if `exc` looks like an LLM context-window-exceeded error.
+    """True if `exc` (or anything it wraps) looks like a context-window error.
 
-    Provider-agnostic by string-matching the canonical error phrasings.
+    Walks the cause/context chain so we catch wrapped errors raised from
+    inside LangGraph / LangChain / Anthropic SDK layers.
     """
-    msg = str(exc).lower()
-    return (
-        "prompt is too long" in msg                # Anthropic
-        or "context_length_exceeded" in msg        # OpenAI
-        or "maximum context length" in msg         # OpenAI (alt phrasing)
-        or "context window" in msg                 # Generic / Bedrock
-        or "too many tokens" in msg                # Generic fallback
-    )
+    seen: set[int] = set()
+    current: BaseException | None = exc
+    while current is not None and id(current) not in seen:
+        seen.add(id(current))
+        msg = str(current).lower()
+        if any(phrase in msg for phrase in _OVERFLOW_PHRASES):
+            return True
+        # Anthropic SDK BadRequestError exposes the API body in .body or .response
+        body = getattr(current, "body", None)
+        if isinstance(body, dict):
+            body_text = str(body).lower()
+            if any(phrase in body_text for phrase in _OVERFLOW_PHRASES):
+                return True
+        current = current.__cause__ or current.__context__
+    return False
 
 
 def _content_length(content: Any) -> int:
@@ -228,23 +245,88 @@ def drop_largest_tool_message(
 # ---------------------------------------------------------------------------
 
 
-GRACEFUL_DEGRADE_TEXT = (
-    "I had to shorten this conversation to keep it within the model's "
-    "context window. Some earlier tool responses were dropped or "
-    "summarized. If your last question depended on data I just dropped, "
-    "please rephrase it more specifically — for example, ask for a single "
-    "record by id rather than a list, or narrow the time range."
-)
+def _content_to_text(content: Any) -> str:
+    """Best-effort flatten of message content to a plain string for inspection."""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts: List[str] = []
+        for item in content:
+            if isinstance(item, dict):
+                parts.append(item.get("text", "") or str(item))
+            else:
+                parts.append(str(item))
+        return "".join(parts)
+    return str(content)
 
 
-def graceful_failure_message() -> AIMessage:
+def _summarize_last_tool_result(messages: List[BaseMessage]) -> str | None:
+    """Inspect the last ToolMessage and produce a one-line human summary.
+
+    Goal: tell the user *what came back* from the bank, so a failed LLM call
+    doesn't look like "Opey did nothing". Best-effort and safe — returns
+    None if we can't extract anything useful.
+    """
+    last_tool: ToolMessage | None = None
+    for msg in reversed(messages):
+        if isinstance(msg, ToolMessage):
+            last_tool = msg
+            break
+    if last_tool is None:
+        return None
+
+    text = _content_to_text(last_tool.content)
+    if not text:
+        return None
+
+    # Heuristic: parse JSON, look for an obvious list of records.
+    try:
+        import json
+        parsed = json.loads(text)
+    except (ValueError, TypeError):
+        return f"the last tool call returned {len(text):,} characters of data"
+
+    response = parsed.get("response", parsed) if isinstance(parsed, dict) else parsed
+    if isinstance(response, dict):
+        for key, value in response.items():
+            if isinstance(value, list):
+                return f"the OBP API returned {len(value)} `{key}` (≈ {len(text):,} chars)"
+        return f"the OBP API returned an object with keys: {', '.join(list(response.keys())[:6])}"
+    if isinstance(response, list):
+        return f"the OBP API returned {len(response)} items (≈ {len(text):,} chars)"
+    return f"the OBP API returned {len(text):,} characters of data"
+
+
+def graceful_failure_message(messages: List[BaseMessage] | None = None) -> AIMessage:
     """Synthetic AIMessage used when every recovery step has failed.
 
     Returning this to state lets the stream end cleanly with a useful
     next-action for the user instead of bubbling a 400 to the portal.
+
+    If `messages` is provided, inspects the most recent ToolMessage and
+    reports what data was actually retrieved — so the user knows the bank
+    call succeeded, only the summarization step failed.
     """
     logger.error(
         "recovery step 4: all recovery steps exhausted — emitting graceful "
         "degrade message to caller"
     )
-    return AIMessage(content=GRACEFUL_DEGRADE_TEXT)
+
+    retrieved_line = ""
+    if messages:
+        summary = _summarize_last_tool_result(messages)
+        if summary:
+            retrieved_line = (
+                f"\n\n**Good news — the bank call succeeded:** "
+                f"{summary}.\n"
+                f"It was too large for me to summarize in one go."
+            )
+
+    text = (
+        "I couldn't fit the result of this turn into my context window."
+        + retrieved_line
+        + "\n\nPlease try a more specific question — for example, ask for "
+        "a single record by id, narrow the time range, or filter by a "
+        "specific bank or account."
+    )
+    return AIMessage(content=text)
diff --git a/src/agent/graph_builder.py b/src/agent/graph_builder.py
@@ -5,7 +5,7 @@
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langchain_core.tools import BaseTool
 from langchain_core.prompts import SystemMessagePromptTemplate, ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import SystemMessage
+from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.runnables import Runnable
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.runnables import RunnableConfig
@@ -63,6 +63,16 @@ async def _invoke_with_recovery(
         return response, {}, messages
     except Exception as exc:
         if not _is_context_overflow(exc):
+            # Log the full exception chain so we can tell whether overflow is
+            # being wrapped in a way the predicate doesn't recognise.
+            chain = []
+            cur: BaseException | None = exc
+            while cur is not None and len(chain) < 6:
+                chain.append(f"{type(cur).__name__}: {str(cur)[:200]}")
+                cur = cur.__cause__ or cur.__context__
+            logger.error(
+                f"opey: ainvoke raised non-overflow exception: {' <- '.join(chain)}"
+            )
             raise
         logger.warning(f"opey: context overflow on first call: {exc!r}; entering recovery cascade")
 
@@ -111,7 +121,9 @@ async def _invoke_with_recovery(
         logger.error(f"opey: step 3 (drop largest) did not fit: {exc!r}; graceful degrade")
 
     # Step 4: graceful degrade — no LLM call, return synthetic message
-    response = graceful_failure_message()
+    # Pass current state messages so the message can report what data
+    # was actually retrieved (the bank call succeeded; only summarization failed).
+    response = graceful_failure_message(messages=state.get("messages", []))
     return response, state_updates, messages
 
 class OpeyAgentGraphBuilder:
@@ -217,6 +229,20 @@ def _build_system_prompt(self) -> str:
     
     def _get_llm(self) -> Runnable:
         """Get the configured LLM"""
+        # DIAGNOSTIC (temporary): log the size of every bound tool schema. The sum is the
+        # fixed per-request overhead sent to the LLM on every call. Remove once resolved.
+        try:
+            from langchain_core.utils.function_calling import convert_to_openai_tool
+            import json as _json
+            _total = 0
+            for _t in self._tools:
+                _sz = len(_json.dumps(convert_to_openai_tool(_t)))
+                _total += _sz
+                logger.warning(f"[TOOL DIAG] {getattr(_t, 'name', _t)}: {_sz} chars")
+            logger.warning(f"[TOOL DIAG] {len(self._tools)} tools, total schema {_total} chars (~{_total // 4} tokens)")
+        except Exception as _diag_err:
+            logger.warning(f"[TOOL DIAG] failed: {_diag_err}")
+
         return get_model(
             self._model_name,
             temperature=self._temperature,
@@ -238,13 +264,44 @@ def _create_opey_node(self):
         
         @cancellable(preserve_state_keys=["total_tokens"])
         async def run_opey(state: OpeyGraphState, config: RunnableConfig):
-            # Check if we have a conversation summary
+            # Ephemeral per-turn system context, prepended only for this LLM call (never
+            # persisted to message history, so it always reflects the latest values).
+            prepend: list[SystemMessage] = []
+
             summary = state.get("conversation_summary", "")
             if summary:
-                summary_system_message = f"Summary of earlier conversation: {summary}"
-                messages = [SystemMessage(content=summary_system_message)] + state["messages"]
-            else:
-                messages = state["messages"]
+                prepend.append(SystemMessage(content=f"Summary of earlier conversation: {summary}"))
+
+            current_bank_id = (config.get("configurable", {}) or {}).get("current_bank_id")
+            if current_bank_id:
+                prepend.append(SystemMessage(content=(
+                    f"The user currently has OBP bank '{current_bank_id}' selected in the UI. "
+                    f"When a request does not name a bank, use this bank_id."
+                )))
+
+            messages = prepend + state["messages"]
+
+            # DIAGNOSTIC (temporary): locate context bloat. Logs message count, total
+            # content size, and the 8 biggest messages. Remove once context issue is fixed.
+            try:
+                _sizes = sorted(
+                    ((type(m).__name__, len(str(getattr(m, "content", "") or ""))) for m in messages),
+                    key=lambda x: x[1], reverse=True,
+                )
+                logger.warning(
+                    f"[CTX DIAG] {len(messages)} messages, "
+                    f"~{sum(s for _, s in _sizes)} content chars (~{sum(s for _, s in _sizes) // 4} tokens); "
+                    f"biggest 8: {_sizes[:8]}"
+                )
+            except Exception as _diag_err:
+                logger.warning(f"[CTX DIAG] failed: {_diag_err}")
+
+            # Anthropic rejects a request whose messages are all system content
+            # ("at least one message is required"). If summarisation/trimming left
+            # nothing to respond to, add a minimal user turn so the call is valid.
+            if not any(not isinstance(m, SystemMessage) for m in messages):
+                logger.warning("run_opey: only system messages present — appending a continuation HumanMessage")
+                messages = messages + [HumanMessage(content="Please continue, based on the summary above.")]
 
             response, extra_updates, messages = await _invoke_with_recovery(
                 opey_agent, messages, state, config
diff --git a/src/schema/schema.py b/src/schema/schema.py
@@ -265,6 +265,12 @@ class UserInput(BaseModel):
         description="Tool call approval object (with consent_jwt, approval, or batch_decisions) or False if not an approval.",
         default=False,
     )
+    current_bank_id: str | None = Field(
+        description="Bank the user currently has selected in the UI. Used as the default bank "
+                    "for requests that don't name one.",
+        default=None,
+        examples=["gh.29.uk"],
+    )
 
 
 class StreamInput(UserInput):
diff --git a/src/service/streaming/events.py b/src/service/streaming/events.py
@@ -50,6 +50,10 @@ class AssistantCompleteEvent(BaseStreamEvent):
     run_id: str = Field(description="Unique identifier for this run")
     content: str = Field(description="The complete response content")
     tool_calls: Optional[list] = Field(default=[], description="Any tool calls made by the assistant")
+    usage: Optional[dict] = Field(
+        default=None,
+        description="Token usage for this LLM call (input_tokens, output_tokens, total_tokens)."
+    )
 
     def to_sse_data(self) -> str:
         return f"data: {self.model_dump_json()}\n\n"
@@ -323,16 +327,17 @@ def assistant_token(content: str, message_id: str) -> AssistantTokenEvent:
         return event
 
     @staticmethod
-    def assistant_complete(content: str, message_id: str, run_id: str, tool_calls: Optional[list] = None) -> AssistantCompleteEvent:
-        event = AssistantCompleteEvent(content=content, message_id=message_id, run_id=run_id, tool_calls=tool_calls or [])
+    def assistant_complete(content: str, message_id: str, run_id: str, tool_calls: Optional[list] = None, usage: Optional[dict] = None) -> AssistantCompleteEvent:
+        event = AssistantCompleteEvent(content=content, message_id=message_id, run_id=run_id, tool_calls=tool_calls or [], usage=usage)
         StreamEventFactory._log_event(
-            event, 
-            "ASSISTANT_COMPLETE", 
+            event,
+            "ASSISTANT_COMPLETE",
             {
                 "message_id": message_id,
                 "run_id": run_id,
                 "content_length": len(content),
-                "tool_calls": len(tool_calls or [])
+                "tool_calls": len(tool_calls or []),
+                "usage": usage
             }
         )
         return event
diff --git a/src/service/streaming/processors.py b/src/service/streaming/processors.py
@@ -124,11 +124,17 @@ async def process(self, event: LangGraphStreamEvent) -> AsyncGenerator[StreamEve
                         # This ensures each assistant response gets a fresh state
                         self._reset_streaming_state()
 
+                        # usage_metadata carries Anthropic's real token counts for this
+                        # call — input_tokens is the actual prompt size (tools + system +
+                        # history), the number that hits the 200k limit.
+                        usage = getattr(message, 'usage_metadata', None)
+
                         yield StreamEventFactory.assistant_complete(
                             content=content,
                             message_id=message_id,
                             run_id=run_id,
-                            tool_calls=tool_calls
+                            tool_calls=tool_calls,
+                            usage=usage
                         )
                     except Exception as e:
                         error_msg = f"Error processing assistant message completion: {str(e)}"
diff --git a/src/service/streaming/stream_manager.py b/src/service/streaming/stream_manager.py
@@ -40,6 +40,12 @@ async def stream_response(
         """
 
         thread_id = config.get("configurable", {}).get("thread_id")
+
+        # Carry the UI-selected bank into the graph config so the agent node can inject it
+        # as ephemeral context each turn. Kept out of message history so it never goes stale.
+        if stream_input.current_bank_id:
+            config.setdefault("configurable", {})["current_bank_id"] = stream_input.current_bank_id
+
         logger.info("Starting stream response", extra={
             "event_type": "stream_start",
             "thread_id": thread_id,
