Coverage for src/local_deep_research/llm/cancellation.py: 69%
68 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""Cooperative cancellation for in-flight LLM calls.
3Research, chat and other long-running operations run on worker threads that
4block inside a synchronous LangChain ``llm.invoke()``. Setting the
5termination flag (``web.research_state.set_termination_flag``) only takes
6effect at the *next* ``check_termination`` checkpoint — i.e. after the
7current LLM HTTP request finishes, which can be minutes on a slow local
8model. The worker therefore keeps burning GPU/API budget after the user
9clicks Stop.
11This module is the single place that maps a ``research_id`` to "should this
12LLM call abort now?". Wrappers (``config.llm_config.ProcessingLLMWrapper``
13and the rate-limiting wrapper) resolve the effective research id from:
151. an explicitly bound id (``bind_llm_to_research`` / ``research_id`` ctor
16 arg — survives thread-pool hops where the ContextVar is not copied), or
172. the current thread's search context (``utilities.thread_context``), so
18 LLMs created without an explicit id inside a research worker are still
19 covered.
21When an id is known, wrappers prefer a streaming implementation of
22``invoke``: they consume ``base_llm.stream()`` chunk by chunk, check the
23flag between chunks, and close the generator on abort. Cancellation
24latency drops from "full response time" to "time to next token". Closing
25the generator releases the underlying HTTP response promptly instead of
26leaving it to GC.
27"""
29from typing import Any, Optional
31from loguru import logger
34def get_current_research_id() -> Optional[str]:
35 """Return the research id from the current thread's search context.
37 Returns ``None`` outside a research worker (request threads, CLI,
38 benchmarks without a context). Never raises — a missing/broken
39 context means "no cancellation scope", not an error.
40 """
41 try:
42 from ..utilities.thread_context import get_search_context
44 ctx = get_search_context()
45 if isinstance(ctx, dict):
46 rid = ctx.get("research_id")
47 if rid: 47 ↛ 54line 47 didn't jump to line 54 because the condition on line 47 was always true
48 return str(rid)
49 except Exception:
50 logger.debug(
51 "get_current_research_id: failed to read search context",
52 exc_info=True,
53 )
54 return None
57def resolve_research_id(explicit: Optional[str] = None) -> Optional[str]:
58 """Return the explicit id, falling back to the thread's search context."""
59 if explicit:
60 return str(explicit)
61 return get_current_research_id()
64def is_termination_requested(research_id: Optional[str]) -> bool:
65 """Return True iff Stop was requested for *research_id*.
67 ``None`` (no cancellation scope) is never considered terminated.
68 Import is function-local so this module never creates an import cycle
69 with ``web.research_state``.
70 """
71 if not research_id:
72 return False
73 try:
74 from ..web.research_state import is_termination_requested as _check
76 return bool(_check(research_id))
77 except Exception:
78 logger.debug(
79 "is_termination_requested: flag lookup failed",
80 exc_info=True,
81 )
82 return False
85def raise_if_terminated(research_id: Optional[str]) -> None:
86 """Raise ``ResearchTerminatedException`` if Stop was requested.
88 No-op when *research_id* is falsy or no flag is set. The exception
89 inherits from ``BaseException`` so ``except Exception`` blocks in
90 strategy code let it propagate (same pattern as CancelledError).
91 """
92 if is_termination_requested(research_id):
93 from ..exceptions import ResearchTerminatedException
95 raise ResearchTerminatedException(
96 f"Research {research_id} was terminated by user during LLM call"
97 )
100def _try_bind(target: Any, rid: str) -> None:
101 """Best-effort ``_ldr_research_id`` setattr on one wrapper level.
103 Plain objects always accept it; exotic proxies or frozen test
104 doubles may refuse — that only loses the fast-path scope for that
105 level (the thread-context fallback still applies), so debug-log
106 and move on.
107 """
108 try:
109 setattr(target, "_ldr_research_id", rid)
110 except Exception:
111 logger.debug("bind_llm_to_research: setattr refused", exc_info=True)
114def bind_llm_to_research(
115 llm: Any, research_id: Optional[str], *, stop_at: Any = None
116) -> Any:
117 """Attach *research_id* to an LLM wrapper chain (in place, chain-safe).
119 Sets ``_ldr_research_id`` on *llm* and, when present, on its
120 ``base_llm`` (covers ``ProcessingLLMWrapper -> RateLimitedWrapper ->
121 raw ChatModel`` without unwrapping). Returns *llm* for chaining.
122 ``None`` clears nothing — a falsy id leaves existing bindings intact
123 so rebinding with an empty value cannot silently drop cancellation.
125 *stop_at* names an object the binding must not reach. The walk stops
126 there, so neither that object nor anything below it is written. Pass
127 the model under LDR's wrappers: other researches can share it (a
128 registered instance, or a model a registered factory returns more than
129 once), or it can hold shared objects of its own.
131 Invariant: a bound wrapper must not outlive its run. The binding is
132 resolved per call (explicit id > binding > thread context), so
133 re-binding for the next run takes effect without reconstructing the
134 LLM — but a wrapper bound to run A and reused for run B without
135 re-binding would fail-fast/post-check against A's flag. Callers must
136 re-bind (or construct with the new id) when crossing runs.
137 """
138 if not research_id: 138 ↛ 139line 138 didn't jump to line 139 because the condition on line 138 was never true
139 return llm
140 if stop_at is not None and llm is stop_at: 140 ↛ 141line 140 didn't jump to line 141 because the condition on line 140 was never true
141 return llm
142 rid = str(research_id)
143 _try_bind(llm, rid)
144 # Propagate one level down so inner wrappers that are invoked
145 # directly (e.g. via __getattr__ passthrough) see the same scope,
146 # down to *stop_at*. Depth is bounded: each level sets a plain
147 # attribute, never recursing.
148 try:
149 inner = getattr(llm, "base_llm", None)
150 except Exception:
151 logger.debug(
152 "bind_llm_to_research: base_llm lookup failed", exc_info=True
153 )
154 return llm
155 if (
156 inner is not None
157 and inner is not stop_at
158 and getattr(inner, "_ldr_research_id", None) != rid
159 ):
160 _try_bind(inner, rid)
161 # One more level for the double-wrapped shape
162 # (Processing -> RateLimited -> ChatModel).
163 try:
164 inner_inner = getattr(inner, "base_llm", None)
165 except Exception:
166 logger.debug(
167 "bind_llm_to_research: nested base_llm lookup failed",
168 exc_info=True,
169 )
170 return llm
171 if ( 171 ↛ 176line 171 didn't jump to line 176 because the condition on line 171 was never true
172 inner_inner is not None
173 and inner_inner is not stop_at
174 and getattr(inner_inner, "_ldr_research_id", None) != rid
175 ):
176 _try_bind(inner_inner, rid)
177 return llm
180def effective_research_id(
181 llm: Any, explicit: Optional[str] = None
182) -> Optional[str]:
183 """Resolve the cancellation scope for an LLM call.
185 Precedence: explicit ctor arg > ``_ldr_research_id`` binding >
186 current thread's search context. Used by wrappers on every call so
187 rebinding (or context propagation into pool threads) takes effect
188 without reconstructing the LLM.
189 """
190 if explicit:
191 return str(explicit)
192 try:
193 bound = getattr(llm, "_ldr_research_id", None)
194 except Exception:
195 logger.debug(
196 "effective_research_id: binding lookup failed", exc_info=True
197 )
198 return get_current_research_id()
199 if bound:
200 return str(bound)
201 return get_current_research_id()