Coverage for src/local_deep_research/config/llm_config.py: 78%
560 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1from typing import Any
3from langchain_core.callbacks import BaseCallbackManager, CallbackManager
4from langchain_core.language_models import BaseChatModel, BaseLanguageModel
5from langchain_core.messages import AIMessage
6from langchain_core.runnables import RunnableBinding, ensure_config
7from langchain_core.runnables.config import merge_configs
8from loguru import logger
10from ..llm import get_llm_from_registry, is_llm_registered
11from ..security.log_sanitizer import sanitize_for_log, scrub_error
12from ..security.secure_logging import log_exception_type
13from ..settings.exceptions import SettingsValueError
14from ..utilities.search_utilities import remove_think_tags
16# Import providers module to trigger auto-discovery. get_llm() has no
17# fallback construction path: if this import fails (e.g. a broken
18# langchain install), the module must fail loudly here rather than start
19# with an empty registry and confusing per-call errors.
20from ..llm.providers import discover_providers # noqa: F401
21from ..llm.providers.auto_discovery import is_builtin_provider_factory
22from ..llm.providers.base import normalize_provider
23from .thread_settings import get_setting_from_snapshot
26class LLMNotConfiguredError(ValueError):
27 """``get_llm`` could not build an LLM because none is CONFIGURED.
29 Raised for the "the user has not set this up yet" failures — no
30 provider, an unknown provider, no ``llm.model``, a provider that
31 auto-discovery did not register. It is deliberately a ``ValueError``
32 subclass so every existing ``except ValueError`` around ``get_llm``
33 keeps behaving exactly as before.
35 It exists so an HTTP boundary can tell a configuration failure apart
36 from a request-validation failure *by type* — the alternative is
37 matching on the exception text, which is what the fixed-message
38 boundary rule forbids. Without it, a route that maps ``ValueError`` to
39 "your request was invalid" tells a fresh install with no ``llm.model``
40 that its perfectly valid request was malformed.
41 """
44# The fixed text a route may show for the above. A constant, never
45# ``str(exc)``: the exception messages are setup guidance written for a log
46# reader, not a boundary contract, and a future raise site could put an
47# endpoint URL or a settings path in one.
48LLM_NOT_CONFIGURED_CLIENT_MESSAGE = (
49 "LLM is not configured. Open Settings and choose a provider and model."
50)
53def get_selected_llm_provider(settings_snapshot=None):
54 return normalize_provider(
55 get_setting_from_snapshot(
56 "llm.provider", "ollama", settings_snapshot=settings_snapshot
57 )
58 )
61def _get_context_window_for_provider(provider_type, settings_snapshot=None):
62 """Resolve the context window size for a provider.
64 Thin wrapper around the canonical
65 ``llm/providers/_helpers.get_context_window_for_provider`` so the
66 provider ``create_llm`` path and this path share a single source of
67 truth (the two were previously identical copies). Kept as a named
68 function because ``get_llm``/``wrap_llm_without_think_tags`` and tests
69 reference it by name and signature.
71 NOTE: the helper reads settings through
72 ``thread_settings.get_setting_from_snapshot`` (a function-local import),
73 so tests exercising context-window resolution must patch
74 ``...config.thread_settings.get_setting_from_snapshot`` rather than
75 ``...config.llm_config.get_setting_from_snapshot``.
77 Returns:
78 int or None: The context window size, or None for unrestricted cloud providers.
79 """
80 from ..llm.providers._helpers import get_context_window_for_provider
82 return get_context_window_for_provider(
83 provider_type, settings_snapshot=settings_snapshot
84 )
87def get_llm(
88 model_name=None,
89 temperature=None,
90 provider=None,
91 openai_endpoint_url=None,
92 research_id=None,
93 research_context=None,
94 settings_snapshot=None,
95 username=None,
96):
97 """
98 Get LLM instance based on model name and provider.
100 Args:
101 model_name: Name of the model to use (if None, uses database setting)
102 temperature: Model temperature (if None, uses database setting)
103 provider: Provider to use (if None, uses database setting)
104 openai_endpoint_url: Per-call URL override for the
105 ``openai_endpoint`` provider. When omitted, the provider uses
106 the ``llm.openai_endpoint.url`` setting.
107 research_id: Optional research ID for token tracking
108 research_context: Optional research context for enhanced token tracking
109 username: Requesting user, used to resolve per-user registered LLMs.
110 When omitted, falls back to the ``_username`` in the settings
111 snapshot; ``None`` resolves shared/built-in providers only.
113 Returns:
114 A LangChain LLM instance with automatic think-tag removal
115 """
117 # Resolve the effective username for per-user registry lookups. An
118 # explicit argument wins; otherwise fall back to the ``_username`` the
119 # settings snapshot carries. None resolves the shared namespace (which
120 # holds the auto-discovered built-in providers) only.
121 if username is None:
122 from ..search_system import username_from_snapshot
124 username = username_from_snapshot(settings_snapshot)
126 # Use database values for parameters if not provided
127 if model_name is None:
128 model_name = get_setting_from_snapshot(
129 "llm.model", "", settings_snapshot=settings_snapshot
130 )
131 if temperature is None:
132 temperature = get_setting_from_snapshot(
133 "llm.temperature", 0.7, settings_snapshot=settings_snapshot
134 )
135 if provider is None:
136 provider = get_setting_from_snapshot(
137 "llm.provider", "ollama", settings_snapshot=settings_snapshot
138 )
140 # Clean model name: remove quotes and extra whitespace
141 if model_name:
142 model_name = model_name.strip().strip("\"'").strip()
144 # Clean provider: remove quotes and extra whitespace
145 if provider:
146 provider = provider.strip().strip("\"'").strip()
148 # Normalize provider: convert to lowercase canonical form
149 provider = normalize_provider(provider)
151 had_settings_snapshot = settings_snapshot is not None
153 # The endpoint override must reach both the egress-policy gate and the
154 # provider factory. Overlay it onto a per-call copy so those paths share
155 # one effective value without mutating the caller's settings snapshot.
156 if provider == "openai_endpoint" and openai_endpoint_url is not None:
157 if (
158 not isinstance(openai_endpoint_url, str)
159 or not openai_endpoint_url.strip()
160 ):
161 raise ValueError("openai_endpoint_url must be a non-empty string")
162 settings_snapshot = dict(settings_snapshot or {})
163 settings_snapshot["llm.openai_endpoint.url"] = (
164 openai_endpoint_url.strip()
165 )
167 # Egress policy PEP for LLM endpoints. Fires here (before the registered-
168 # LLM dispatch) because all built-in providers are auto-registered via
169 # discover_providers(), so the registered branch handles every real LLM
170 # call.
171 #
172 # Two paths: snapshot-present runs the full PEP; snapshot-absent runs
173 # an allow-list check so background helpers / scaffolding paths cannot
174 # silently instantiate a cloud LLM. The allow-list (known-local) is
175 # deliberately tight — any provider not in it (including ambiguous
176 # ones like ``openai_endpoint`` and future cloud providers) fails
177 # closed instead of bypassing the PEP.
178 if provider:
179 from ..security.egress.policy import (
180 Decision,
181 EgressContext,
182 EgressScope,
183 PolicyDeniedError,
184 _LOCAL_DEFAULT_LLM_PROVIDERS,
185 _is_user_registered_llm,
186 context_from_snapshot,
187 evaluate_llm_endpoint,
188 resolve_run_primary_engine,
189 )
191 if not had_settings_snapshot:
192 if provider == "openai_endpoint" and settings_snapshot is not None:
193 # A per-call endpoint is enough to classify this otherwise
194 # snapshot-less provider. Preserve the fail-closed contract:
195 # only a local endpoint may proceed without a policy snapshot.
196 ctx = EgressContext(
197 scope=EgressScope.PRIVATE_ONLY,
198 primary_engine="snapshotless-explicit-endpoint",
199 require_local_llm=True,
200 require_local_embeddings=True,
201 )
202 decision = evaluate_llm_endpoint(
203 provider,
204 ctx,
205 settings_snapshot=settings_snapshot,
206 username=username,
207 )
208 if not decision.allowed:
209 logger.bind(policy_audit=True).warning(
210 "Snapshot-less explicit LLM endpoint denied",
211 provider=provider,
212 reason=decision.reason,
213 )
214 raise PolicyDeniedError(decision, target=provider)
215 # User-registered in-process LLMs are exempt here for the same
216 # reason evaluate_llm_endpoint allows them: no endpoint to
217 # classify, operator-injected, audit-hook backstopped.
218 elif provider not in _LOCAL_DEFAULT_LLM_PROVIDERS and not (
219 _is_user_registered_llm(provider, username=username)
220 ):
221 logger.bind(policy_audit=True).warning(
222 "LLM constructed without policy snapshot; refusing "
223 "non-local provider",
224 provider=provider,
225 )
226 raise PolicyDeniedError(
227 Decision(False, "no_snapshot_for_provider"),
228 target=provider,
229 )
230 else:
231 try:
232 # Derive the run's primary the SAME way the search-engine
233 # factory does (single source of truth), instead of the old
234 # ``search.tool`` + searxng fallback. That fallback was a
235 # fail-OPEN: a missing/blank primary defaulted to searxng ->
236 # ADAPTIVE -> PUBLIC_ONLY -> require_local_llm stayed False ->
237 # the endpoint check below was skipped, admitting a CLOUD LLM
238 # for a run whose actual posture was private. resolve_run_
239 # primary_engine raises on a missing/invalid primary, which we
240 # treat as a hard stop here.
241 primary_engine = resolve_run_primary_engine(settings_snapshot)
242 # Pass username so a per-user private retriever set as the
243 # run's primary is classified against the caller's own
244 # namespace. Without it ADAPTIVE can't see the retriever
245 # (shared namespace only), resolves to the permissive BOTH,
246 # leaves require_local_llm False, and admits a cloud LLM for a
247 # run over the user's private corpus (fail-open).
248 ctx = context_from_snapshot(
249 settings_snapshot, primary_engine, username=username
250 )
251 except ValueError as exc:
252 # No configured primary, or an invalid policy config. Fail
253 # closed: previously a missing primary silently fell back to
254 # searxng/PUBLIC_ONLY and skipped the LLM endpoint check
255 # entirely, opening a cloud-LLM bypass under the very
256 # configuration the user asked to be strict about.
257 logger.bind(policy_audit=True).warning(
258 "no/invalid egress policy primary; refusing LLM",
259 provider=provider,
260 reason=str(exc),
261 )
262 raise PolicyDeniedError(
263 Decision(False, "invalid_policy_config"),
264 target=provider,
265 ) from exc
267 if ctx is not None and ctx.require_local_llm:
268 decision = evaluate_llm_endpoint(
269 provider,
270 ctx,
271 settings_snapshot=settings_snapshot,
272 username=username,
273 )
274 if not decision.allowed:
275 logger.bind(policy_audit=True).warning(
276 "LLM endpoint denied by egress policy",
277 provider=provider,
278 reason=decision.reason,
279 )
280 raise PolicyDeniedError(decision, target=provider)
282 # Check if this is a registered custom LLM first. Resolve against the
283 # caller's own namespace before the shared/built-in providers so one
284 # user's registration can't shadow a built-in for anyone else.
285 if provider and is_llm_registered(provider, username=username):
286 logger.info(f"Using registered custom LLM: {provider}")
287 custom_llm = get_llm_from_registry(provider, username=username)
289 # Check if it's a callable (factory function) or a BaseChatModel instance
290 if callable(custom_llm) and not isinstance(custom_llm, BaseChatModel):
291 # It's a callable (factory function), call it with parameters
292 try:
293 llm_instance = custom_llm(
294 model_name=model_name,
295 temperature=temperature,
296 settings_snapshot=settings_snapshot,
297 )
298 except TypeError as e:
299 # Re-raise TypeError with better message
300 raise TypeError(
301 f"Registered LLM factory '{provider}' has invalid signature. "
302 f"Factory functions must accept 'model_name', 'temperature', and 'settings_snapshot' parameters. "
303 f"Error: {e}"
304 )
306 # Validate the result is a BaseChatModel
307 if not isinstance(llm_instance, BaseChatModel):
308 raise ValueError(
309 f"Factory function for {provider} must return a BaseChatModel instance, "
310 f"got {type(llm_instance).__name__}"
311 )
312 elif isinstance(custom_llm, BaseChatModel):
313 # It's already a proper LLM instance, use it directly
314 llm_instance = custom_llm
315 else:
316 raise ValueError(
317 f"Registered LLM {provider} must be either a BaseChatModel instance "
318 f"or a callable factory function. Got: {type(custom_llm).__name__}"
319 )
321 return wrap_llm_without_think_tags(
322 llm_instance,
323 research_id=research_id,
324 provider=provider,
325 research_context=research_context,
326 settings_snapshot=settings_snapshot,
327 # LDR closes only a model it knows it built: a built-in
328 # provider's create_llm builds a new one on every call. A
329 # registered instance, or whatever a registered factory
330 # returns, can be serving other researches (a factory can
331 # return one cached or singleton model), so it stays with the
332 # code that registered it.
333 owns_llm=is_builtin_provider_factory(custom_llm),
334 )
336 # Validate the provider against the auto-discovered set — NOT a hardcoded
337 # list. Auto-discovery (discover_providers, run at module import) registers
338 # every llm/providers/implementations/*.py class, and the
339 # is_llm_registered() check above already serves them, so the registry /
340 # discovery IS the single source of truth for "valid provider".
341 #
342 # We deliberately do not keep a separate VALID_PROVIDERS constant: it was a
343 # third copy of the provider list (besides the implementations directory
344 # and the registry) and it silently drifted from auto-discovery (xai,
345 # ionos and deepseek were valid+registered yet missing from it). Deriving
346 # the set from discovery here means it can never drift. ('none' is the
347 # explicit "unset" sentinel, handled by the guard further down.)
348 from ..llm.providers import get_discovered_provider_options
350 valid_providers = {
351 normalize_provider(option["value"])
352 for option in get_discovered_provider_options()
353 } | {"none"}
354 if provider not in valid_providers:
355 # ``provider`` is a user-settable value (``llm.provider``) reaching a
356 # log sink. The process-wide loguru patcher (``_sanitize_record`` in
357 # utilities/log_utils.py, via ``redact_log_record``) already strips
358 # control characters (\x00-\x1f, \x7f) from every record's message
359 # and redacts credentials, so ``sanitize_for_log`` isn't newline/ANSI
360 # defence here -- it adds the length cap plus the ``None`` guard
361 # below. The exception text is a different audience (it never
362 # reaches a client -- boundaries answer with
363 # LLM_NOT_CONFIGURED_CLIENT_MESSAGE) and is left intact so
364 # `match="Invalid provider: <name>"` assertions and the operator's
365 # "which value did I actually set?" question survive.
366 # ``str()`` first: ``normalize_provider`` returns None for an
367 # empty/whitespace-only ``llm.provider``, and that None reaches here
368 # (the ``if provider:`` gates above all skipped). ``sanitize_for_log``
369 # takes a str, so the bare call would turn this configuration error
370 # into a TypeError. ``str(None)`` renders "None", exactly what the
371 # previous f-string logged.
372 logger.error(
373 f"Invalid provider in settings: {sanitize_for_log(str(provider))}"
374 )
375 raise LLMNotConfiguredError(
376 f"Invalid provider: {provider}. "
377 f"Must be one of: {sorted(valid_providers)}"
378 )
380 # Require an explicit model for built-in providers. Mirrors the
381 # API-key-not-configured pattern in openai_base.py and the URL-not-
382 # configured pattern in providers/implementations/ollama.py: no silent
383 # substitution to a hardcoded default model.
384 if not model_name or not model_name.strip():
385 logger.error("llm.model is not configured (empty/None after lookup)")
386 raise LLMNotConfiguredError(
387 "LLM model not configured. Please open Settings, choose an LLM "
388 "provider, and select a model name (e.g. 'gpt-4o-mini' for "
389 "OpenAI, 'claude-3-5-sonnet-20241022' for Anthropic, "
390 "'llama3.1:8b' for Ollama). The 'llm.model' setting is required."
391 )
392 logger.info(
393 f"Getting LLM with model: {model_name}, temperature: {temperature}, provider: {provider}"
394 )
396 # Set context_limit on research_context for overflow detection. The
397 # actual max_tokens calculation lives in each provider's create_llm
398 # (via providers/_helpers.compute_max_tokens) so the cap is consistent
399 # across the live registered-LLM path.
400 context_window_size = _get_context_window_for_provider(
401 provider, settings_snapshot
402 )
403 if research_context and context_window_size: 403 ↛ 404line 403 didn't jump to line 404 because the condition on line 403 was never true
404 research_context["context_limit"] = context_window_size
405 logger.info(
406 f"Set context_limit={context_window_size} in research_context"
407 )
408 else:
409 logger.debug(
410 f"Context limit not set: research_context={bool(research_context)}, "
411 f"context_window_size={context_window_size}"
412 )
414 # Reaching here means the registered-LLM branch above did NOT fire,
415 # which is unusual — auto-discovery normally registers all 6+ built-in
416 # providers (anthropic, openai, openai_endpoint, ollama, lmstudio,
417 # llamacpp + xai/ionos/openrouter via OpenAICompatibleProvider) at
418 # import time. Two specific guards preserve the user-facing error
419 # messages for known-bad cases.
420 if provider == "none":
421 raise LLMNotConfiguredError(
422 "No LLM provider configured. Please set llm.provider in settings "
423 "to a valid provider (e.g., 'ollama', 'openai', 'anthropic')."
424 )
425 raise LLMNotConfiguredError(
426 f"Provider '{provider}' was not registered by auto-discovery. "
427 f"This usually indicates an import error during startup — check the "
428 f"logs for 'Error loading provider from <module>' messages. "
429 f"Note that clear_llm_registry()/unregister_llm() also remove "
430 f"built-in providers; discover_providers(force_refresh=True) "
431 f"restores them."
432 )
435def _log_llm_error(error: Exception) -> None:
436 """Log an LLM call failure with credential redaction."""
437 safe_msg = scrub_error(error)
438 logger.warning(f"LLM Request - Failed with error: {safe_msg}")
441def _merge_callbacks(existing, added, *, inherit: bool = True):
442 """Return a callback manager holding *existing* plus *added*.
444 *existing* is never changed. It holds a call's callbacks: the ones its
445 caller passed, or the ones it inherits from an enclosing run such as a
446 LangGraph node. A caller can pass one handler list or
447 ``BaseCallbackManager`` to many calls, and an enclosing run's manager
448 serves all of its child runs, so adding to either in place could make
449 those other calls fire this research's token callback too, writing
450 their usage under this research's user and research id.
452 *existing* stays inheritable. With ``inherit=False`` the added handlers
453 see only the run the result is passed to, the way a model's own
454 ``callbacks`` field behaves, and child runs do not inherit them.
456 The result is a manager even when *existing* is a list. A
457 ``RunnableBinding`` merges a call's callbacks with the ones it binds,
458 which the call's callbacks can already hold (see
459 ``_with_research_callbacks``). LangChain's ``merge_configs`` adds each
460 handler once when either side is a manager, but joins two lists with
461 every handler from both.
462 """
463 added = list(added)
464 if isinstance(existing, BaseCallbackManager):
465 merged = existing.copy()
466 if merged is existing: 466 ↛ 467line 466 didn't jump to line 467 because the condition on line 466 was never true
467 raise TypeError("A callback manager must support isolated copies")
468 else:
469 inherited = list(existing or [])
470 merged = CallbackManager(
471 handlers=inherited, inheritable_handlers=list(inherited)
472 )
473 for handler in added:
474 merged.add_handler(handler, inherit=inherit)
475 return merged
478def _config_receiver(runnable):
479 """Return the runnable that receives the config passed to *runnable*.
481 LDR's rate-limiting wrapper hands each call's arguments, the config
482 included, unchanged to the runnable it wraps.
483 """
484 seen = set()
485 while (
486 getattr(type(runnable), "_ldr_forwards_call_arguments", False)
487 and id(runnable) not in seen
488 ):
489 seen.add(id(runnable))
490 runnable = runnable.base_llm
491 return runnable
494def _starts_the_model_run(runnable) -> bool:
495 """Whether the config passed to *runnable* starts the model's own run.
497 A model starts its run with the config it is called with, and a
498 ``RunnableBinding``, which ``bind``, ``bind_tools`` and ``with_config``
499 return for a standard model, passes its merged config on to the
500 runnable it binds. Anything else is taken to start a run of its own
501 and run the model as a child of that run, as the chain that
502 ``with_structured_output`` returns does, or a chain like
503 ``prompt | self`` that a custom ``bind_tools`` returns.
504 """
505 while isinstance(runnable, RunnableBinding):
506 runnable = runnable.bound
507 return isinstance(runnable, BaseLanguageModel)
510class ProcessingLLMWrapper:
511 # Token callbacks of the research this wrapper serves, passed with each
512 # call instead of stored on the model; see _with_research_callbacks.
513 # The class default keeps lookups away from __getattr__, which would
514 # otherwise ask the model.
515 _ldr_research_callbacks: tuple = ()
517 def __init__(
518 self,
519 base_llm,
520 *,
521 owns_llm: bool = True,
522 research_id=None,
523 provider=None,
524 settings_snapshot=None,
525 structured_output: bool = False,
526 research_callbacks=(),
527 ):
528 self.base_llm = base_llm
529 self._owns_llm = owns_llm
530 self._ldr_research_callbacks = tuple(research_callbacks)
531 # Cancellation scope: explicit id wins, otherwise the thread's
532 # search context is consulted per call (see _effective_research_id).
533 # Stored under both names — ``research_id`` for readability and
534 # ``_ldr_research_id`` for the shared bind helper in
535 # llm/cancellation.py (which propagates through nested wrappers).
536 self.research_id = research_id
537 self._ldr_research_id = research_id
538 # Token-accounting gate: the stream-under-invoke path silently
539 # converts every scoped invoke into a streaming request. For the
540 # ``openai_endpoint`` provider that loses token/cost metrics unless
541 # ``llm.openai_endpoint.stream_usage`` is enabled (custom gateways
542 # omit usage on streams, or reject stream_options outright).
543 # The provider + snapshot let _stream_cancel_allowed() make that
544 # tradeoff a decision instead of an accident.
545 self._provider = provider
546 self._settings_snapshot = settings_snapshot
547 # Structured-output contract: a wrapper created via
548 # with_structured_output() wraps a Runnable whose invoke/stream
549 # yield dicts/Pydantic models, not AIMessageChunks. Reassembling
550 # those through _combine_stream_chunks would corrupt them into an
551 # AIMessage holding str(dict). Such wrappers must use the direct
552 # invoke path (pre/post termination checks still apply).
553 # A bound wrapper must not outlive its run: the id is resolved per
554 # call so re-binding for the next run takes effect without
555 # reconstructing the LLM, and a stale binding is only ever used
556 # for a fail-fast/post-check, never to route into another run.
557 self._ldr_structured_output = structured_output
559 @staticmethod
560 def _normalize_response(response: Any) -> Any:
561 """Strip <think> tags and normalize the response shape.
563 This is the SINGLE authorized place that strips <think> tags from a
564 fresh LLM response. Every LLM from get_llm() is wrapped here, so
565 downstream code can use ``response.content`` directly — do NOT add
566 per-site ``remove_think_tags`` / ``get_llm_response_text`` calls on
567 fresh ``invoke``/``ainvoke`` results; they are redundant and hide
568 bugs. (Exceptions: ``.stream()``/``.astream()`` pass chunks on
569 unstripped, Runnable transforms other than ``bind_tools``,
570 ``with_structured_output``, ``bind`` and ``with_config`` still
571 bypass this wrapper via ``__getattr__``, and injected/unwrapped LLMs
572 are not wrapped at all — those may still need explicit handling.
573 Those four are NOT exceptions: they re-wrap so the stripper survives
574 — see the ``bind_tools`` override below and #4804.)
576 A message keeps its object identity (only ``.content`` is rewritten,
577 so ``additional_kwargs``/``reasoning_content``/``tool_calls`` survive).
578 A bare-string return (some providers/wrappers) is wrapped into an
579 ``AIMessage`` so callers can always rely on ``.content``. Anything
580 else is passed through unchanged.
582 Only *string* ``.content`` is stripped: ``remove_think_tags`` is
583 text-only, and the ``<think>...</think>`` artifact is only ever
584 emitted as plain text by some local models. Non-string content
585 (e.g. provider content-block lists like Anthropic's, or ``None``)
586 is passed through untouched — running the regex on it would raise
587 ``TypeError`` or corrupt the structured content.
588 """
589 if hasattr(response, "content"):
590 if isinstance(response.content, str):
591 response.content = remove_think_tags(response.content)
592 elif isinstance(response, str):
593 response = AIMessage(content=remove_think_tags(response))
594 return response
596 @staticmethod
597 def _log_llm_error(error: Exception) -> None:
598 _log_llm_error(error)
600 def _with_research_callbacks(self, args, kwargs):
601 """Return the call arguments with this research's callbacks added.
603 The callbacks go into the call's ``RunnableConfig`` (positional
604 argument two or ``config=``), so the model is invoked as it is and
605 never stores them. Any model can serve other researches and users:
606 a registered instance does, and so does a model that a registered
607 factory returns more than once, such as a cached or singleton
608 model.
610 Only the callbacks change: every other config key resolves as it
611 does for the same call without this wrapper. The config passed on
612 holds the caller's own keys rather than everything the call
613 inherits, plus the callbacks the call would use anyway: the
614 caller's, or those of an enclosing run such as a LangGraph node.
615 A ``RunnableBinding``, such as the result of ``with_config``, lets
616 its bound config win over an inherited one only when it is called
617 without a config. For a binding, the config passed on is therefore
618 its bound config already merged with the call's, as the binding
619 merges them, so the binding merging its bound config in again gives
620 the same values.
622 The handlers are not inheritable when that config starts the
623 model's own run (see ``_starts_the_model_run``), which matches a
624 model's own ``callbacks`` field: a run the model starts inside its
625 own run is not recorded as well. When the config starts a chain
626 instead, such as under ``with_structured_output`` or a custom
627 ``bind_tools`` that returns ``prompt | self``, the model call is a
628 child run, so the handlers are inheritable. If the model itself
629 invokes another runnable during that call, the handlers see that
630 inner run too and record its tokens as well.
632 Only the wrapper's own call methods add the callbacks: ``invoke``,
633 ``ainvoke``, ``stream``, ``astream``, ``batch`` and ``abatch``. The
634 wrappers that ``bind_tools``, ``with_structured_output``, ``bind``
635 and ``with_config`` return have the same methods and callbacks. A
636 call that leaves the wrapper through ``__getattr__`` reaches the
637 model without them, so its tokens are not recorded for the
638 research. That covers the other Runnable transforms, such as
639 ``with_retry`` and ``with_fallbacks``, and methods such as
640 ``generate`` and ``astream_events``.
641 """
642 if not self._ldr_research_callbacks:
643 return args, kwargs
644 positional = len(args) > 1
645 explicit = args[1] if positional else kwargs.get("config")
646 receiver = _config_receiver(self.base_llm)
647 if isinstance(receiver, RunnableBinding):
648 config = merge_configs(receiver.config, explicit)
649 else:
650 config = dict(explicit or {})
651 config["callbacks"] = _merge_callbacks(
652 ensure_config(config).get("callbacks"),
653 self._ldr_research_callbacks,
654 inherit=not _starts_the_model_run(receiver),
655 )
656 if positional:
657 return (args[0], config, *args[2:]), kwargs
658 return args, {**kwargs, "config": config}
660 def _effective_research_id(self):
661 """Resolve the cancellation scope for this call.
663 Precedence: explicit ctor id > ``_ldr_research_id`` binding >
664 current thread's search context. Consulted per call so a late
665 ``bind_llm_to_research`` (or context propagation into pool
666 threads) takes effect without reconstructing the LLM.
667 """
668 from ..llm.cancellation import effective_research_id
670 return effective_research_id(self, self.research_id)
672 def _check_terminated(self, research_id=None):
673 """Raise ``ResearchTerminatedException`` if Stop was requested."""
674 from ..llm.cancellation import raise_if_terminated
676 rid = (
677 research_id
678 if research_id is not None
679 else self._effective_research_id()
680 )
681 if rid:
682 raise_if_terminated(rid)
684 def _walk_base_chain(self):
685 """Yield self.base_llm and each nested ``base_llm`` once.
687 Covers ``Processing -> RateLimited -> ChatModel`` without
688 recursing (each level is a plain attribute read). Cycle-safe:
689 a mis-wired test double that points at itself terminates.
690 """
691 seen = set()
692 cur = getattr(self, "base_llm", None)
693 while cur is not None and id(cur) not in seen:
694 seen.add(id(cur))
695 yield cur
696 try:
697 cur = getattr(cur, "base_llm", None)
698 except Exception:
699 break
701 def _underlying_stream_usage(self) -> bool:
702 """True if any wrapped model opts into usage on streams.
704 Covers ``ChatOpenAI(stream_usage=True)`` and equivalents that
705 stash ``stream_options.include_usage`` in ``model_kwargs``.
706 """
707 for base in self._walk_base_chain():
708 try:
709 if getattr(base, "stream_usage", None):
710 return True
711 mk = getattr(base, "model_kwargs", None)
712 if isinstance(mk, dict): 712 ↛ 713line 712 didn't jump to line 713 because the condition on line 712 was never true
713 so = mk.get("stream_options", {})
714 if isinstance(so, dict) and so.get("include_usage"):
715 return True
716 except Exception:
717 continue
718 return False
720 def _underlying_disables_streaming(self) -> bool:
721 """True if any wrapped model forbids streaming."""
722 for base in self._walk_base_chain():
723 try:
724 if getattr(base, "disable_streaming", False):
725 return True
726 except Exception:
727 continue
728 return False
730 def _stream_cancel_allowed(self) -> bool:
731 """Whether invoke may stream under the hood for cancellation.
733 Two gates, both conservative (False keeps the direct invoke path
734 with pre/post termination checks, so Stop still aborts between
735 calls — only mid-generation abort is lost):
737 1. Structured-output wrappers never stream: their chunks are
738 dicts/Pydantic models, and reassembly would corrupt them into
739 ``AIMessage(str(dict))``.
740 2. ``openai_endpoint`` without usage-on-stream never streams:
741 the reassembled message would carry no ``usage_metadata`` and
742 token/cost metrics would silently go to zero. Opt in via the
743 ``llm.openai_endpoint.stream_usage`` setting (or an underlying
744 model built with ``stream_usage=True``). Other providers stream
745 because their streamed chunks preserve usage via
746 ``usage_metadata``/``response_metadata``.
747 """
748 if getattr(self, "_ldr_structured_output", False):
749 return False
750 if self._underlying_disables_streaming():
751 return False
752 provider = getattr(self, "_provider", None)
753 try:
754 from ..llm.providers.base import normalize_provider
756 provider = normalize_provider(provider) if provider else None
757 except Exception:
758 logger.debug(
759 "Provider normalization failed; treating as unknown provider",
760 exc_info=True,
761 )
762 if provider == "openai_endpoint":
763 try:
764 from .thread_settings import get_setting_from_snapshot
766 snapshot = getattr(self, "_settings_snapshot", None)
767 if get_setting_from_snapshot(
768 "llm.openai_endpoint.stream_usage",
769 False,
770 settings_snapshot=snapshot,
771 ):
772 return True
773 except SettingsValueError:
774 raise
775 except Exception:
776 logger.debug(
777 "stream_usage setting lookup failed; "
778 "falling back to model capability",
779 exc_info=True,
780 )
781 # Model-level opt-in covers callers that built ChatOpenAI with
782 # stream_usage=True but have no settings context on this thread.
783 if self._underlying_stream_usage():
784 return True
785 return False
786 # Unknown provider but recognisably a ChatOpenAI-style model with
787 # a custom base_url and no usage-on-stream: same zero-metrics trap
788 # as openai_endpoint (e.g. thread-context fallback where the wrapper
789 # was built without a provider label). Be conservative.
790 for base in self._walk_base_chain():
791 try:
792 if hasattr(base, "stream_usage") and hasattr( 792 ↛ 795line 792 didn't jump to line 795 because the condition on line 792 was never true
793 base, "openai_api_base"
794 ):
795 if getattr(base, "stream_usage", None):
796 return True
797 if getattr(base, "openai_api_base", None):
798 return False
799 # Default OpenAI endpoint without an explicit opt-in:
800 # the validator enables stream_usage for the default
801 # URL, so a falsy value here is unexpected — fall
802 # through to the safe default below.
803 return False
804 except Exception:
805 continue
806 return True
808 @staticmethod
809 def _is_structured_chunk(chunk: Any) -> bool:
810 """True for structured-output chunks (dicts/Pydantic, no .content)."""
811 if isinstance(chunk, dict):
812 return True
813 if isinstance(chunk, str): 813 ↛ 814line 813 didn't jump to line 814 because the condition on line 813 was never true
814 return False
815 if hasattr(chunk, "content"): 815 ↛ 817line 815 didn't jump to line 817 because the condition on line 815 was always true
816 return False
817 try:
818 from pydantic import BaseModel as _PydanticBaseModel
820 if isinstance(chunk, _PydanticBaseModel):
821 return True
822 except Exception:
823 logger.debug(
824 "Pydantic structured-chunk check failed; treating as unstructured",
825 exc_info=True,
826 )
827 return False
829 @staticmethod
830 def _combine_stream_chunks(chunks):
831 """Combine streamed chunks into a single response object.
833 ``AIMessageChunk`` parts merge with ``+`` (preserves ``tool_calls``,
834 ``additional_kwargs`` and content-block lists); the merged chunk is
835 converted to a plain ``AIMessage`` so downstream code sees the same
836 shape as ``invoke()``. Plain-string chunks (some test doubles)
837 are joined into an ``AIMessage``. Returns ``None`` when *chunks*
838 is empty (test doubles only — real models raise rather than yield
839 zero chunks) so the caller can fall back to a direct ``invoke()``.
841 Structured-output chunks (dicts/Pydantic models without ``.content``)
842 are never reassembled here: returning ``None`` lets the caller fall
843 back to a direct ``invoke()`` that preserves the contract instead of
844 corrupting the payload into ``AIMessage(str(dict))``. That fallback
845 is a second upstream request for one logical call (token callbacks
846 fire twice); callers log a warning when it happens after chunks
847 were already consumed.
848 """
849 if not chunks: 849 ↛ 850line 849 didn't jump to line 850 because the condition on line 849 was never true
850 return None
851 for chunk in chunks:
852 if isinstance(chunk, dict): 852 ↛ 853line 852 didn't jump to line 853 because the condition on line 852 was never true
853 return None
854 if isinstance(chunk, str): 854 ↛ 855line 854 didn't jump to line 855 because the condition on line 854 was never true
855 continue
856 if hasattr(chunk, "content"): 856 ↛ 860line 856 didn't jump to line 860 because the condition on line 856 was always true
857 continue
858 # Pydantic models (and other parsed outputs) have no .content:
859 # treat as structured, not as text to stringify.
860 try:
861 from pydantic import BaseModel as _PydanticBaseModel
863 if isinstance(chunk, _PydanticBaseModel):
864 return None
865 except Exception:
866 logger.debug(
867 "Pydantic structured-chunk check failed; continuing",
868 exc_info=True,
869 )
870 # Unknown non-message shape — only safe to stringify when it
871 # is clearly a message-like object; otherwise fall back.
872 if not hasattr(chunk, "tool_calls") and not hasattr(
873 chunk, "additional_kwargs"
874 ):
875 return None
876 try:
877 # Fast path: message chunks support addition.
878 combined = chunks[0]
879 for chunk in chunks[1:]:
880 try:
881 combined = combined + chunk
882 except Exception:
883 # Mixed shapes — fall back to text join below.
884 combined = None
885 break
886 if combined is not None and hasattr(combined, "content"): 886 ↛ 900line 886 didn't jump to line 900 because the condition on line 886 was always true
887 try:
888 from langchain_core.messages.utils import (
889 message_chunk_to_message,
890 )
892 return message_chunk_to_message(combined)
893 except Exception:
894 return combined
895 except Exception:
896 logger.debug(
897 "Stream-chunk combine via '+' failed; using text fallback"
898 )
899 # Text fallback: join content/text of each chunk.
900 from ..utilities.json_utils import _coerce_content_blocks
902 parts = []
903 for chunk in chunks:
904 text = (
905 chunk
906 if isinstance(chunk, str)
907 else getattr(chunk, "content", None)
908 )
909 if text is None:
910 text = str(chunk)
911 if isinstance(text, list):
912 try:
913 text = _coerce_content_blocks(text)
914 except Exception:
915 text = "".join(str(p) for p in text)
916 if text:
917 parts.append(str(text))
918 if not parts:
919 return None
920 return AIMessage(content="".join(parts))
922 def _invoke_via_stream(self, *args: Any, **kwargs: Any):
923 """Implement ``invoke`` as stream+reassemble with per-chunk cancel.
925 Returns the combined response, or ``None`` when streaming is
926 unavailable (no ``stream`` attr), disallowed (structured-output or
927 usage-unsafe endpoint — see :meth:`_stream_cancel_allowed`), or
928 produced zero chunks (only empty test-double generators — real
929 LangChain models raise inside the generator on failure rather
930 than yielding zero chunks). The caller then falls back to a
931 direct ``invoke()``. Raises ``ResearchTerminatedException``
932 promptly when the flag is set mid-generation; the generator is
933 closed so the underlying HTTP response is released instead of
934 leaking to GC.
936 Mid-stream structured fallback (a dict/Pydantic chunk after one
937 or more message chunks were already pulled) also returns ``None``,
938 so the caller issues a *second* upstream request for the same
939 logical call — token callbacks fire for both attempts. No in-app
940 path produces such mixed shapes today (structured-output wrappers
941 are excluded up front via :meth:`_stream_cancel_allowed`), so this
942 is defense-in-depth only; a warning is logged when it triggers so
943 the double request is visible instead of silent.
944 """
945 if getattr(self, "_ldr_structured_output", False): 945 ↛ 946line 945 didn't jump to line 946 because the condition on line 945 was never true
946 return None
947 stream_fn = getattr(self.base_llm, "stream", None)
948 if stream_fn is None: 948 ↛ 949line 948 didn't jump to line 949 because the condition on line 948 was never true
949 return None
950 research_id = self._effective_research_id()
951 try:
952 gen = stream_fn(*args, **kwargs)
953 except Exception:
954 logger.debug(
955 "Sync stream creation failed; falling back to direct invoke",
956 exc_info=True,
957 )
958 return None
959 chunks = []
960 try:
961 for chunk in gen:
962 if research_id: 962 ↛ 967line 962 didn't jump to line 967 because the condition on line 962 was always true
963 self._check_terminated(research_id)
964 # Structured-output runnables yield dicts/models: abort the
965 # streaming attempt and fall back to direct invoke rather
966 # than corrupting the contract.
967 if self._is_structured_chunk(chunk):
968 # NOTE: package logging is disabled by default
969 # (local_deep_research/__init__.py); this surfaces
970 # once the app enables it at startup.
971 logger.warning(
972 "Sync stream yielded structured chunk after "
973 f"{len(chunks)} message chunk(s); falling back "
974 "to direct invoke issues a second upstream "
975 "request for one logical call (token callbacks "
976 "fire twice)"
977 )
978 try:
979 close = getattr(gen, "close", None)
980 if close is not None: 980 ↛ 987line 980 didn't jump to line 987 because the condition on line 980 was always true
981 close()
982 except Exception:
983 logger.debug(
984 "Sync stream generator close failed "
985 "(structured fallback, non-critical)"
986 )
987 return None
988 chunks.append(chunk)
989 except BaseException:
990 # On abort (or any mid-stream failure) release the HTTP
991 # response promptly. ``GeneratorExit`` from close() surfaces
992 # here for some providers — never mask the original raise.
993 try:
994 close = getattr(gen, "close", None)
995 if close is not None: 995 ↛ 1001line 995 didn't jump to line 1001 because the condition on line 995 was always true
996 close()
997 except Exception:
998 logger.debug(
999 "Sync stream generator close failed on abort (non-critical)"
1000 )
1001 raise
1002 else:
1003 try:
1004 close = getattr(gen, "close", None)
1005 if close is not None: 1005 ↛ 1011line 1005 didn't jump to line 1011 because the condition on line 1005 was always true
1006 close()
1007 except Exception:
1008 logger.debug(
1009 "Sync stream generator close failed (non-critical)"
1010 )
1011 if not chunks:
1012 # Only reachable with empty test-double generators; real
1013 # LangChain models raise inside the generator on failure.
1014 return None
1015 combined = self._combine_stream_chunks(chunks)
1016 if combined is None: 1016 ↛ 1020line 1016 didn't jump to line 1020 because the condition on line 1016 was never true
1017 # Chunks were consumed but not reassemblable (structured or
1018 # unknown shapes): the caller's direct-invoke fallback is a
1019 # second upstream request — see the docstring above.
1020 logger.warning(
1021 "Sync stream produced "
1022 f"{len(chunks)} non-reassemblable chunk(s); falling back "
1023 "to direct invoke issues a second upstream request for "
1024 "one logical call (token callbacks fire twice)"
1025 )
1026 return combined
1028 async def _ainvoke_via_stream(self, *args: Any, **kwargs: Any):
1029 """Async twin of :meth:`_invoke_via_stream` (via ``astream``).
1031 Same fallback contract: zero chunks (test doubles only — real
1032 models raise) and mid-stream structured shapes return ``None``;
1033 a structured shape after pulled chunks means the direct-ainvoke
1034 fallback is a second upstream request (warning logged).
1035 """
1036 if getattr(self, "_ldr_structured_output", False): 1036 ↛ 1037line 1036 didn't jump to line 1037 because the condition on line 1036 was never true
1037 return None
1038 astream_fn = getattr(self.base_llm, "astream", None)
1039 if astream_fn is None: 1039 ↛ 1040line 1039 didn't jump to line 1040 because the condition on line 1039 was never true
1040 return None
1041 research_id = self._effective_research_id()
1042 chunks = []
1043 # ``astream`` is an async-generator function: calling it must not
1044 # raise for a missing implementation (that is what the None check
1045 # above covers); iteration errors propagate to the caller.
1046 try:
1047 agen = astream_fn(*args, **kwargs)
1048 except Exception:
1049 logger.debug(
1050 "Async stream creation failed; falling back to direct ainvoke",
1051 exc_info=True,
1052 )
1053 return None
1054 try:
1055 async for chunk in agen:
1056 if research_id: 1056 ↛ 1058line 1056 didn't jump to line 1058 because the condition on line 1056 was always true
1057 self._check_terminated(research_id)
1058 if self._is_structured_chunk(chunk):
1059 logger.warning(
1060 "Async stream yielded structured chunk after "
1061 f"{len(chunks)} message chunk(s); falling back "
1062 "to direct ainvoke issues a second upstream "
1063 "request for one logical call (token callbacks "
1064 "fire twice)"
1065 )
1066 try:
1067 aclose = getattr(agen, "aclose", None)
1068 if aclose is not None: 1068 ↛ 1075line 1068 didn't jump to line 1075 because the condition on line 1068 was always true
1069 await aclose()
1070 except Exception:
1071 logger.debug(
1072 "Async stream generator close failed "
1073 "(structured fallback, non-critical)"
1074 )
1075 return None
1076 chunks.append(chunk)
1077 except BaseException:
1078 try:
1079 aclose = getattr(agen, "aclose", None)
1080 if aclose is not None: 1080 ↛ 1083line 1080 didn't jump to line 1083 because the condition on line 1080 was always true
1081 await aclose()
1082 else:
1083 close = getattr(agen, "close", None)
1084 if close is not None:
1085 close()
1086 except Exception:
1087 logger.debug(
1088 "Async stream generator close failed on abort (non-critical)"
1089 )
1090 raise
1091 else:
1092 try:
1093 aclose = getattr(agen, "aclose", None)
1094 if aclose is not None: 1094 ↛ 1100line 1094 didn't jump to line 1100 because the condition on line 1094 was always true
1095 await aclose()
1096 except Exception:
1097 logger.debug(
1098 "Async stream generator close failed (non-critical)"
1099 )
1100 if not chunks: 1100 ↛ 1103line 1100 didn't jump to line 1103 because the condition on line 1100 was never true
1101 # Only reachable with empty test-double generators; real
1102 # LangChain models raise inside the generator on failure.
1103 return None
1104 combined = self._combine_stream_chunks(chunks)
1105 if combined is None: 1105 ↛ 1106line 1105 didn't jump to line 1106 because the condition on line 1105 was never true
1106 logger.warning(
1107 "Async stream produced "
1108 f"{len(chunks)} non-reassemblable chunk(s); falling back "
1109 "to direct ainvoke issues a second upstream request for "
1110 "one logical call (token callbacks fire twice)"
1111 )
1112 return combined
1114 def invoke(self, *args: Any, **kwargs: Any) -> Any:
1115 # Fail fast if Stop arrived while queued (e.g. between
1116 # check_termination checkpoints) — no new HTTP request.
1117 self._check_terminated()
1118 try:
1119 args, kwargs = self._with_research_callbacks(args, kwargs)
1120 # When a cancellation scope is known, stream under the hood
1121 # so a Stop mid-generation aborts within one chunk instead
1122 # of waiting for the full response (minutes on slow local
1123 # models). Without a scope there is nothing to check
1124 # between chunks, so keep the direct path. The streaming
1125 # path is additionally gated by _stream_cancel_allowed():
1126 # structured-output wrappers and usage-unsafe endpoints
1127 # (openai_endpoint without stream_usage) stay on direct
1128 # invoke so token/cost metrics and the dict/Pydantic contract
1129 # survive — pre/post checks still abort promptly between calls.
1130 if (
1131 self._effective_research_id() is not None
1132 and self._stream_cancel_allowed()
1133 ):
1134 streamed = self._invoke_via_stream(*args, **kwargs)
1135 if streamed is not None:
1136 self._check_terminated()
1137 return self._normalize_response(streamed)
1138 response = self.base_llm.invoke(*args, **kwargs)
1139 except Exception as e:
1140 self._log_llm_error(e)
1141 raise
1142 # Post-check: the flag may have been set while blocked. Raise
1143 # so the worker aborts before the next pipeline step instead of
1144 # continuing with a stale result.
1145 self._check_terminated()
1146 return self._normalize_response(response)
1148 async def ainvoke(self, *args: Any, **kwargs: Any) -> Any:
1149 # Async counterpart of invoke(); without this, ainvoke() would fall
1150 # through __getattr__ to the base LLM and bypass think-tag stripping.
1151 self._check_terminated()
1152 try:
1153 args, kwargs = self._with_research_callbacks(args, kwargs)
1154 if (
1155 self._effective_research_id() is not None
1156 and self._stream_cancel_allowed()
1157 ):
1158 streamed = await self._ainvoke_via_stream(*args, **kwargs)
1159 if streamed is not None:
1160 self._check_terminated()
1161 return self._normalize_response(streamed)
1162 response = await self.base_llm.ainvoke(*args, **kwargs)
1163 except Exception as e:
1164 self._log_llm_error(e)
1165 raise
1166 self._check_terminated()
1167 return self._normalize_response(response)
1169 def stream(self, *args: Any, **kwargs: Any):
1170 """Yield chunks with per-chunk termination checks.
1172 Previously ``stream`` fell through ``__getattr__`` to the raw
1173 model, bypassing both think-tag normalization and cancellation.
1174 Now the generator checks the flag before yielding each chunk and
1175 raises ``ResearchTerminatedException`` promptly on Stop. Chunks
1176 are forwarded untouched (no per-chunk scrubbing — credentials
1177 can straddle boundaries).
1178 """
1179 research_id = self._effective_research_id()
1180 if research_id:
1181 self._check_terminated(research_id)
1182 args, kwargs = self._with_research_callbacks(args, kwargs)
1183 gen = None
1184 try:
1185 gen = self.base_llm.stream(*args, **kwargs)
1186 for chunk in gen:
1187 if research_id:
1188 self._check_terminated(research_id)
1189 yield chunk
1190 except Exception as error:
1191 log_exception_type(
1192 "LLM Request - Failed with error in stream()", error
1193 )
1194 raise
1195 finally:
1196 try:
1197 close = getattr(gen, "close", None)
1198 if close is not None: 1198 ↛ exitline 1198 didn't return from function 'stream' because the condition on line 1198 was always true
1199 close()
1200 except Exception:
1201 logger.debug(
1202 "Sync stream generator close failed in stream() (non-critical)"
1203 )
1205 async def astream(self, *args: Any, **kwargs: Any):
1206 """Async generator twin of :meth:`stream` with the same checks."""
1207 research_id = self._effective_research_id()
1208 if research_id:
1209 self._check_terminated(research_id)
1210 args, kwargs = self._with_research_callbacks(args, kwargs)
1211 agen = None
1212 try:
1213 agen = self.base_llm.astream(*args, **kwargs)
1214 async for chunk in agen:
1215 if research_id:
1216 self._check_terminated(research_id)
1217 yield chunk
1218 except Exception as error:
1219 log_exception_type(
1220 "LLM Request - Failed with error in astream()", error
1221 )
1222 raise
1223 finally:
1224 try:
1225 aclose = getattr(agen, "aclose", None)
1226 if aclose is not None: 1226 ↛ exitline 1226 didn't return from function 'astream' because the condition on line 1226 was always true
1227 await aclose()
1228 except Exception:
1229 logger.debug(
1230 "Async stream generator close failed in astream() (non-critical)"
1231 )
1233 # Pass through any other attributes to the base LLM
1234 def __getattr__(self, name):
1235 return getattr(self.base_llm, name)
1237 def batch(
1238 self, inputs, config=None, *, return_exceptions=False, **kwargs: Any
1239 ):
1240 """Batch via :meth:`invoke` so per-item cancellation applies."""
1241 results = []
1242 # Resolve shared vs per-input configs like LangChain does.
1243 if isinstance(config, list): 1243 ↛ 1244line 1243 didn't jump to line 1244 because the condition on line 1243 was never true
1244 if len(config) != len(inputs):
1245 raise ValueError(
1246 f"configs must match the number of inputs ({len(config)} != {len(inputs)})"
1247 )
1248 configs = list(config)
1249 else:
1250 configs = [config] * len(inputs)
1251 for single_input, single_config in zip(inputs, configs):
1252 try:
1253 if single_config is not None: 1253 ↛ 1254line 1253 didn't jump to line 1254 because the condition on line 1253 was never true
1254 results.append(
1255 self.invoke(single_input, single_config, **kwargs)
1256 )
1257 else:
1258 results.append(self.invoke(single_input, **kwargs))
1259 except Exception as e:
1260 if return_exceptions:
1261 results.append(e)
1262 else:
1263 raise
1264 # ResearchTerminatedException (BaseException) propagates even with
1265 # return_exceptions=True — a Stop must never be swallowed into a
1266 # result list.
1267 return results
1269 async def abatch(
1270 self, inputs, config=None, *, return_exceptions=False, **kwargs: Any
1271 ):
1272 """Async twin of :meth:`batch`."""
1273 results = []
1274 if isinstance(config, list): 1274 ↛ 1275line 1274 didn't jump to line 1275 because the condition on line 1274 was never true
1275 if len(config) != len(inputs):
1276 raise ValueError(
1277 f"configs must match the number of inputs ({len(config)} != {len(inputs)})"
1278 )
1279 configs = list(config)
1280 else:
1281 configs = [config] * len(inputs)
1282 for single_input, single_config in zip(inputs, configs): 1282 ↛ 1297line 1282 didn't jump to line 1297 because the loop on line 1282 didn't complete
1283 try:
1284 if single_config is not None: 1284 ↛ 1285line 1284 didn't jump to line 1285 because the condition on line 1284 was never true
1285 results.append(
1286 await self.ainvoke(
1287 single_input, single_config, **kwargs
1288 )
1289 )
1290 else:
1291 results.append(await self.ainvoke(single_input, **kwargs))
1292 except Exception as e:
1293 if return_exceptions:
1294 results.append(e)
1295 else:
1296 raise
1297 return results
1299 def close(self):
1300 """Close the model's clients if this wrapper owns the model.
1302 ``get_llm()`` gives a wrapper ownership only of a model that a
1303 built-in provider built for it. A registered instance, or a model
1304 that a registered factory returns, stays with the code that
1305 registered it, since other researches can be using it.
1306 """
1307 if not self._owns_llm:
1308 return
1309 try:
1310 from ..utilities.llm_utils import _close_base_llm
1312 _close_base_llm(self.base_llm)
1313 except Exception:
1314 logger.debug(
1315 "best-effort cleanup of HTTP clients on shutdown",
1316 exc_info=True,
1317 )
1319 def bind_tools(self, tools, **kwargs: Any) -> Any:
1320 """Re-wrap the tool-bound model so the <think>-stripper survives.
1322 ``langchain.agents.create_agent()`` and similar helpers call
1323 ``model.bind_tools(...)`` and then use the *return value* as
1324 the model that actually runs inside the agent loop. Without
1325 this override, Python falls through ``__getattr__`` to
1326 ``self.base_llm.bind_tools(...)``, which returns an unwrapped
1327 ``BaseChatModel``; the wrapper's ``_normalize_response`` (the
1328 single authorized site that strips ``<think>…`` from fresh
1329 responses) is then bypassed for every LLM call the agent
1330 makes.
1332 That bypass is the root cause of the regression tracked in
1333 issue #4804: reasoning-mode models (Qwen 3.x, deepseek-r1,
1334 etc.) emit ``<think>…</think>`` as plain text, the wrapper
1335 fails to strip it, the raw artifact is appended to the
1336 agent's ``AIMessage.content``, and on the next turn the
1337 provider's strict parser rejects the request with
1338 ``Failed to parse input at pos 0: <think>…``.
1340 Re-wrapping the bound model keeps the stripper in the loop
1341 for every LLM call the agent makes, not just the ones we
1342 issue directly via ``self.model.invoke(...)``.
1344 Args:
1345 tools: Tool definitions (schemas, callables, or
1346 ``BaseTool`` instances) to bind to the model.
1347 **kwargs: Provider-specific binding options forwarded
1348 verbatim to ``base_llm.bind_tools``.
1350 Returns:
1351 A ``ProcessingLLMWrapper`` wrapping the bound model, so
1352 every subsequent ``invoke``/``ainvoke`` continues to
1353 strip ``<think>`` tags.
1354 """
1355 bound = self.base_llm.bind_tools(tools, **kwargs)
1356 return self._rewrap(bound, structured_output=False)
1358 def bind(self, **kwargs: Any) -> Any:
1359 """Re-wrap like :meth:`bind_tools`.
1361 ``langchain.agents.create_agent()`` calls ``model.bind(...)`` for an
1362 agent without tools and runs the result. Without this override the
1363 call would fall through ``__getattr__`` to the model, and the bound
1364 model would skip the ``<think>`` stripping, the termination checks
1365 and this research's token callbacks. Like this wrapper, the result
1366 cannot be composed with ``|``.
1367 """
1368 return self._rewrap(self.base_llm.bind(**kwargs))
1370 def with_config(self, config=None, **kwargs: Any) -> Any:
1371 """Re-wrap like :meth:`bind`.
1373 The bound config resolves with each call's config as it does in
1374 LangChain: when the call passes no config, its keys win over those
1375 inherited from an enclosing run. This research's token callbacks
1376 run beside any callbacks it binds.
1377 """
1378 return self._rewrap(self.base_llm.with_config(config, **kwargs))
1380 def _rewrap(self, runnable, *, structured_output=None):
1381 """Wrap *runnable*, derived from this wrapper's model, the same way.
1383 The new wrapper keeps this wrapper's ownership, cancellation scope,
1384 provider, settings and research callbacks. It is a structured-output
1385 wrapper if this one is, unless *structured_output* says otherwise.
1386 """
1387 if structured_output is None:
1388 structured_output = self._ldr_structured_output
1389 return ProcessingLLMWrapper(
1390 runnable,
1391 owns_llm=self._owns_llm,
1392 research_id=self._effective_research_id(),
1393 provider=getattr(self, "_provider", None),
1394 settings_snapshot=getattr(self, "_settings_snapshot", None),
1395 structured_output=structured_output,
1396 research_callbacks=self._ldr_research_callbacks,
1397 )
1399 def with_structured_output(self, schema, **kwargs: Any) -> Any:
1400 """Re-wrap like :meth:`bind_tools` so cancellation survives.
1402 Structured-output runnables are used by several strategies for
1403 JSON extraction; without re-wrapping, the bound runnable would
1404 bypass the termination checks in ``invoke``. The re-wrapped
1405 instance is marked ``structured_output`` so ``invoke``/``ainvoke``
1406 use the direct path (pre/post checks) instead of stream+reassemble,
1407 preserving the dict/Pydantic contract.
1408 """
1409 bound = self.base_llm.with_structured_output(schema, **kwargs)
1410 try:
1411 from ..llm.cancellation import bind_llm_to_research
1413 bind_llm_to_research(bound, self._effective_research_id())
1414 except Exception:
1415 logger.debug(
1416 "Structured-output scope propagation failed (non-critical)"
1417 )
1418 return self._rewrap(bound, structured_output=True)
1421def wrap_llm_without_think_tags(
1422 llm,
1423 research_id=None,
1424 provider=None,
1425 research_context=None,
1426 settings_snapshot=None,
1427 *,
1428 owns_llm: bool = True,
1429):
1430 """Wrap response processing and metrics with explicit resource ownership.
1432 Set ``owns_llm=False`` for a model the wrapper must not close.
1433 ``get_llm()`` keeps ``True`` only for a model that a built-in provider
1434 built for this call, and passes ``False`` for a registered instance and
1435 for any model a registered factory returns: a factory can return the
1436 same cached or singleton model on every call, and closing it when one
1437 research ends would break the others. Whoever registered such a model
1438 remains responsible for closing it after all consumers are done.
1440 Ownership only decides who closes the model. Whether it owns the model
1441 or not, this function writes nothing about the research onto the
1442 model: the research's token callback is passed with each call, and the
1443 cancellation scope is bound to LDR's wrappers only.
1444 """
1446 # The model before any LDR wrapper. No research state lands on it or
1447 # on anything it holds.
1448 model = llm
1450 # First apply rate limiting if enabled
1451 from ..web_search_engines.rate_limiting.llm import (
1452 create_rate_limited_llm_wrapper,
1453 )
1455 # Check if LLM rate limiting is enabled (independent of search rate limiting)
1456 # Use the thread-safe get_db_setting defined in this module
1457 if get_setting_from_snapshot(
1458 "rate_limiting.llm_enabled", False, settings_snapshot=settings_snapshot
1459 ):
1460 llm = create_rate_limited_llm_wrapper(model, provider)
1462 # Set context_limit in research_context for overflow detection.
1463 # This is needed for providers that go through the registered provider path
1464 # (which returns before the code in get_llm that sets context_limit).
1465 if research_context is not None and provider is not None:
1466 if "context_limit" not in research_context:
1467 context_limit = _get_context_window_for_provider(
1468 provider, settings_snapshot
1469 )
1470 if context_limit is not None:
1471 research_context["context_limit"] = context_limit
1472 logger.info(
1473 f"Set context_limit={context_limit} in wrap_llm for provider={provider}"
1474 )
1476 # Import token counting functionality if research_id is provided
1477 callbacks = []
1478 if research_id is not None:
1479 from ..metrics import TokenCounter
1481 token_counter = TokenCounter()
1482 token_callback = token_counter.create_callback(
1483 research_id, research_context
1484 )
1485 # Set provider and model info on the callback
1486 if provider:
1487 token_callback.preset_provider = provider
1488 # Try to extract model name from the LLM instance
1489 if hasattr(model, "model_name"):
1490 token_callback.preset_model = model.model_name
1491 elif hasattr(model, "model"):
1492 token_callback.preset_model = model.model
1493 callbacks.append(token_callback)
1495 wrapped = ProcessingLLMWrapper(
1496 llm,
1497 owns_llm=owns_llm,
1498 research_id=research_id,
1499 provider=provider,
1500 settings_snapshot=settings_snapshot,
1501 # The model can serve other researches and users, owned or not, so
1502 # this research's callbacks travel in each call's RunnableConfig.
1503 # The rate-limiting wrapper passes that config on to the model.
1504 research_callbacks=callbacks,
1505 )
1506 if research_id is not None:
1507 # Propagate the scope to inner wrappers (rate-limiting layer)
1508 # so direct uses of the inner instance also abort promptly. The
1509 # walk stops at the model: it can be shared, and any model can
1510 # hold shared objects under ``base_llm``.
1511 try:
1512 from ..llm.cancellation import bind_llm_to_research
1514 bind_llm_to_research(wrapped, research_id, stop_at=model)
1515 except Exception:
1516 logger.debug(
1517 "Inner-wrapper scope propagation failed (non-critical)"
1518 )
1519 return wrapped