Coverage for src/local_deep_research/config/llm_config.py: 97%
159 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1from typing import Any
3from langchain_core.language_models import BaseChatModel
4from langchain_core.messages import AIMessage
5from loguru import logger
7from ..llm import get_llm_from_registry, is_llm_registered
8from ..security.log_sanitizer import scrub_error
9from ..utilities.search_utilities import remove_think_tags
11# Import providers module to trigger auto-discovery. get_llm() has no
12# fallback construction path: if this import fails (e.g. a broken
13# langchain install), the module must fail loudly here rather than start
14# with an empty registry and confusing per-call errors.
15from ..llm.providers import discover_providers # noqa: F401
16from ..llm.providers.base import normalize_provider
17from .thread_settings import get_setting_from_snapshot
20def get_selected_llm_provider(settings_snapshot=None):
21 return normalize_provider(
22 get_setting_from_snapshot(
23 "llm.provider", "ollama", settings_snapshot=settings_snapshot
24 )
25 )
28def _get_context_window_for_provider(provider_type, settings_snapshot=None):
29 """Resolve the context window size for a provider.
31 Thin wrapper around the canonical
32 ``llm/providers/_helpers.get_context_window_for_provider`` so the
33 provider ``create_llm`` path and this path share a single source of
34 truth (the two were previously identical copies). Kept as a named
35 function because ``get_llm``/``wrap_llm_without_think_tags`` and tests
36 reference it by name and signature.
38 NOTE: the helper reads settings through
39 ``thread_settings.get_setting_from_snapshot`` (a function-local import),
40 so tests exercising context-window resolution must patch
41 ``...config.thread_settings.get_setting_from_snapshot`` rather than
42 ``...config.llm_config.get_setting_from_snapshot``.
44 Returns:
45 int or None: The context window size, or None for unrestricted cloud providers.
46 """
47 from ..llm.providers._helpers import get_context_window_for_provider
49 return get_context_window_for_provider(
50 provider_type, settings_snapshot=settings_snapshot
51 )
54def get_llm(
55 model_name=None,
56 temperature=None,
57 provider=None,
58 openai_endpoint_url=None,
59 research_id=None,
60 research_context=None,
61 settings_snapshot=None,
62 username=None,
63):
64 """
65 Get LLM instance based on model name and provider.
67 Args:
68 model_name: Name of the model to use (if None, uses database setting)
69 temperature: Model temperature (if None, uses database setting)
70 provider: Provider to use (if None, uses database setting)
71 openai_endpoint_url: Per-call URL override for the
72 ``openai_endpoint`` provider. When omitted, the provider uses
73 the ``llm.openai_endpoint.url`` setting.
74 research_id: Optional research ID for token tracking
75 research_context: Optional research context for enhanced token tracking
76 username: Requesting user, used to resolve per-user registered LLMs.
77 When omitted, falls back to the ``_username`` in the settings
78 snapshot; ``None`` resolves shared/built-in providers only.
80 Returns:
81 A LangChain LLM instance with automatic think-tag removal
82 """
84 # Resolve the effective username for per-user registry lookups. An
85 # explicit argument wins; otherwise fall back to the ``_username`` the
86 # settings snapshot carries. None resolves the shared namespace (which
87 # holds the auto-discovered built-in providers) only.
88 if username is None:
89 from ..search_system import username_from_snapshot
91 username = username_from_snapshot(settings_snapshot)
93 # Use database values for parameters if not provided
94 if model_name is None:
95 model_name = get_setting_from_snapshot(
96 "llm.model", "", settings_snapshot=settings_snapshot
97 )
98 if temperature is None:
99 temperature = get_setting_from_snapshot(
100 "llm.temperature", 0.7, settings_snapshot=settings_snapshot
101 )
102 if provider is None:
103 provider = get_setting_from_snapshot(
104 "llm.provider", "ollama", settings_snapshot=settings_snapshot
105 )
107 # Clean model name: remove quotes and extra whitespace
108 if model_name:
109 model_name = model_name.strip().strip("\"'").strip()
111 # Clean provider: remove quotes and extra whitespace
112 if provider: 112 ↛ 116line 112 didn't jump to line 116 because the condition on line 112 was always true
113 provider = provider.strip().strip("\"'").strip()
115 # Normalize provider: convert to lowercase canonical form
116 provider = normalize_provider(provider)
118 had_settings_snapshot = settings_snapshot is not None
120 # The endpoint override must reach both the egress-policy gate and the
121 # provider factory. Overlay it onto a per-call copy so those paths share
122 # one effective value without mutating the caller's settings snapshot.
123 if provider == "openai_endpoint" and openai_endpoint_url is not None:
124 if (
125 not isinstance(openai_endpoint_url, str)
126 or not openai_endpoint_url.strip()
127 ):
128 raise ValueError("openai_endpoint_url must be a non-empty string")
129 settings_snapshot = dict(settings_snapshot or {})
130 settings_snapshot["llm.openai_endpoint.url"] = (
131 openai_endpoint_url.strip()
132 )
134 # Egress policy PEP for LLM endpoints. Fires here (before the registered-
135 # LLM dispatch) because all built-in providers are auto-registered via
136 # discover_providers(), so the registered branch handles every real LLM
137 # call.
138 #
139 # Two paths: snapshot-present runs the full PEP; snapshot-absent runs
140 # an allow-list check so background helpers / scaffolding paths cannot
141 # silently instantiate a cloud LLM. The allow-list (known-local) is
142 # deliberately tight — any provider not in it (including ambiguous
143 # ones like ``openai_endpoint`` and future cloud providers) fails
144 # closed instead of bypassing the PEP.
145 if provider: 145 ↛ 252line 145 didn't jump to line 252 because the condition on line 145 was always true
146 from ..security.egress.policy import (
147 Decision,
148 EgressContext,
149 EgressScope,
150 PolicyDeniedError,
151 _LOCAL_DEFAULT_LLM_PROVIDERS,
152 _is_user_registered_llm,
153 context_from_snapshot,
154 evaluate_llm_endpoint,
155 resolve_run_primary_engine,
156 )
158 if not had_settings_snapshot:
159 if provider == "openai_endpoint" and settings_snapshot is not None:
160 # A per-call endpoint is enough to classify this otherwise
161 # snapshot-less provider. Preserve the fail-closed contract:
162 # only a local endpoint may proceed without a policy snapshot.
163 ctx = EgressContext(
164 scope=EgressScope.PRIVATE_ONLY,
165 primary_engine="snapshotless-explicit-endpoint",
166 require_local_llm=True,
167 require_local_embeddings=True,
168 )
169 decision = evaluate_llm_endpoint(
170 provider,
171 ctx,
172 settings_snapshot=settings_snapshot,
173 username=username,
174 )
175 if not decision.allowed:
176 logger.bind(policy_audit=True).warning(
177 "Snapshot-less explicit LLM endpoint denied",
178 provider=provider,
179 reason=decision.reason,
180 )
181 raise PolicyDeniedError(decision, target=provider)
182 # User-registered in-process LLMs are exempt here for the same
183 # reason evaluate_llm_endpoint allows them: no endpoint to
184 # classify, operator-injected, audit-hook backstopped.
185 elif provider not in _LOCAL_DEFAULT_LLM_PROVIDERS and not (
186 _is_user_registered_llm(provider, username=username)
187 ):
188 logger.bind(policy_audit=True).warning(
189 "LLM constructed without policy snapshot; refusing "
190 "non-local provider",
191 provider=provider,
192 )
193 raise PolicyDeniedError(
194 Decision(False, "no_snapshot_for_provider"),
195 target=provider,
196 )
197 else:
198 try:
199 # Derive the run's primary the SAME way the search-engine
200 # factory does (single source of truth), instead of the old
201 # ``search.tool`` + searxng fallback. That fallback was a
202 # fail-OPEN: a missing/blank primary defaulted to searxng ->
203 # ADAPTIVE -> PUBLIC_ONLY -> require_local_llm stayed False ->
204 # the endpoint check below was skipped, admitting a CLOUD LLM
205 # for a run whose actual posture was private. resolve_run_
206 # primary_engine raises on a missing/invalid primary, which we
207 # treat as a hard stop here.
208 primary_engine = resolve_run_primary_engine(settings_snapshot)
209 # Pass username so a per-user private retriever set as the
210 # run's primary is classified against the caller's own
211 # namespace. Without it ADAPTIVE can't see the retriever
212 # (shared namespace only), resolves to the permissive BOTH,
213 # leaves require_local_llm False, and admits a cloud LLM for a
214 # run over the user's private corpus (fail-open).
215 ctx = context_from_snapshot(
216 settings_snapshot, primary_engine, username=username
217 )
218 except ValueError as exc:
219 # No configured primary, or an invalid policy config. Fail
220 # closed: previously a missing primary silently fell back to
221 # searxng/PUBLIC_ONLY and skipped the LLM endpoint check
222 # entirely, opening a cloud-LLM bypass under the very
223 # configuration the user asked to be strict about.
224 logger.bind(policy_audit=True).warning(
225 "no/invalid egress policy primary; refusing LLM",
226 provider=provider,
227 reason=str(exc),
228 )
229 raise PolicyDeniedError(
230 Decision(False, "invalid_policy_config"),
231 target=provider,
232 ) from exc
234 if ctx is not None and ctx.require_local_llm:
235 decision = evaluate_llm_endpoint(
236 provider,
237 ctx,
238 settings_snapshot=settings_snapshot,
239 username=username,
240 )
241 if not decision.allowed:
242 logger.bind(policy_audit=True).warning(
243 "LLM endpoint denied by egress policy",
244 provider=provider,
245 reason=decision.reason,
246 )
247 raise PolicyDeniedError(decision, target=provider)
249 # Check if this is a registered custom LLM first. Resolve against the
250 # caller's own namespace before the shared/built-in providers so one
251 # user's registration can't shadow a built-in for anyone else.
252 if provider and is_llm_registered(provider, username=username):
253 logger.info(f"Using registered custom LLM: {provider}")
254 custom_llm = get_llm_from_registry(provider, username=username)
256 # Check if it's a callable (factory function) or a BaseChatModel instance
257 if callable(custom_llm) and not isinstance(custom_llm, BaseChatModel):
258 # It's a callable (factory function), call it with parameters
259 try:
260 llm_instance = custom_llm(
261 model_name=model_name,
262 temperature=temperature,
263 settings_snapshot=settings_snapshot,
264 )
265 except TypeError as e:
266 # Re-raise TypeError with better message
267 raise TypeError(
268 f"Registered LLM factory '{provider}' has invalid signature. "
269 f"Factory functions must accept 'model_name', 'temperature', and 'settings_snapshot' parameters. "
270 f"Error: {e}"
271 )
273 # Validate the result is a BaseChatModel
274 if not isinstance(llm_instance, BaseChatModel):
275 raise ValueError(
276 f"Factory function for {provider} must return a BaseChatModel instance, "
277 f"got {type(llm_instance).__name__}"
278 )
279 elif isinstance(custom_llm, BaseChatModel):
280 # It's already a proper LLM instance, use it directly
281 llm_instance = custom_llm
282 else:
283 raise ValueError(
284 f"Registered LLM {provider} must be either a BaseChatModel instance "
285 f"or a callable factory function. Got: {type(custom_llm).__name__}"
286 )
288 return wrap_llm_without_think_tags(
289 llm_instance,
290 research_id=research_id,
291 provider=provider,
292 research_context=research_context,
293 settings_snapshot=settings_snapshot,
294 )
296 # Validate the provider against the auto-discovered set — NOT a hardcoded
297 # list. Auto-discovery (discover_providers, run at module import) registers
298 # every llm/providers/implementations/*.py class, and the
299 # is_llm_registered() check above already serves them, so the registry /
300 # discovery IS the single source of truth for "valid provider".
301 #
302 # We deliberately do not keep a separate VALID_PROVIDERS constant: it was a
303 # third copy of the provider list (besides the implementations directory
304 # and the registry) and it silently drifted from auto-discovery (xai,
305 # ionos and deepseek were valid+registered yet missing from it). Deriving
306 # the set from discovery here means it can never drift. ('none' is the
307 # explicit "unset" sentinel, handled by the guard further down.)
308 from ..llm.providers import get_discovered_provider_options
310 valid_providers = {
311 normalize_provider(option["value"])
312 for option in get_discovered_provider_options()
313 } | {"none"}
314 if provider not in valid_providers:
315 logger.error(f"Invalid provider in settings: {provider}")
316 raise ValueError(
317 f"Invalid provider: {provider}. "
318 f"Must be one of: {sorted(valid_providers)}"
319 )
321 # Require an explicit model for built-in providers. Mirrors the
322 # API-key-not-configured pattern in openai_base.py and the URL-not-
323 # configured pattern in providers/implementations/ollama.py: no silent
324 # substitution to a hardcoded default model.
325 if not model_name or not model_name.strip():
326 logger.error("llm.model is not configured (empty/None after lookup)")
327 raise ValueError(
328 "LLM model not configured. Please open Settings, choose an LLM "
329 "provider, and select a model name (e.g. 'gpt-4o-mini' for "
330 "OpenAI, 'claude-3-5-sonnet-20241022' for Anthropic, "
331 "'llama3.1:8b' for Ollama). The 'llm.model' setting is required."
332 )
333 logger.info(
334 f"Getting LLM with model: {model_name}, temperature: {temperature}, provider: {provider}"
335 )
337 # Set context_limit on research_context for overflow detection. The
338 # actual max_tokens calculation lives in each provider's create_llm
339 # (via providers/_helpers.compute_max_tokens) so the cap is consistent
340 # across the live registered-LLM path.
341 context_window_size = _get_context_window_for_provider(
342 provider, settings_snapshot
343 )
344 if research_context and context_window_size: 344 ↛ 345line 344 didn't jump to line 345 because the condition on line 344 was never true
345 research_context["context_limit"] = context_window_size
346 logger.info(
347 f"Set context_limit={context_window_size} in research_context"
348 )
349 else:
350 logger.debug(
351 f"Context limit not set: research_context={bool(research_context)}, "
352 f"context_window_size={context_window_size}"
353 )
355 # Reaching here means the registered-LLM branch above did NOT fire,
356 # which is unusual — auto-discovery normally registers all 6+ built-in
357 # providers (anthropic, openai, openai_endpoint, ollama, lmstudio,
358 # llamacpp + xai/ionos/openrouter via OpenAICompatibleProvider) at
359 # import time. Two specific guards preserve the user-facing error
360 # messages for known-bad cases.
361 if provider == "none":
362 raise ValueError(
363 "No LLM provider configured. Please set llm.provider in settings "
364 "to a valid provider (e.g., 'ollama', 'openai', 'anthropic')."
365 )
366 raise ValueError(
367 f"Provider '{provider}' was not registered by auto-discovery. "
368 f"This usually indicates an import error during startup — check the "
369 f"logs for 'Error loading provider from <module>' messages. "
370 f"Note that clear_llm_registry()/unregister_llm() also remove "
371 f"built-in providers; discover_providers(force_refresh=True) "
372 f"restores them."
373 )
376def _log_llm_error(error: Exception) -> None:
377 """Log an LLM call failure with credential redaction."""
378 safe_msg = scrub_error(error)
379 logger.warning(f"LLM Request - Failed with error: {safe_msg}")
382class ProcessingLLMWrapper:
383 def __init__(self, base_llm):
384 self.base_llm = base_llm
386 @staticmethod
387 def _normalize_response(response: Any) -> Any:
388 """Strip <think> tags and normalize the response shape.
390 This is the SINGLE authorized place that strips <think> tags from a
391 fresh LLM response. Every LLM from get_llm() is wrapped here, so
392 downstream code can use ``response.content`` directly — do NOT add
393 per-site ``remove_think_tags`` / ``get_llm_response_text`` calls on
394 fresh ``invoke``/``ainvoke`` results; they are redundant and hide
395 bugs. (Exceptions: ``.stream()``/``.astream()``, ``with_config``,
396 ``bind`` and other Runnable transforms still bypass this wrapper via
397 ``__getattr__``, as do injected/unwrapped LLMs — those may still need
398 explicit handling. ``bind_tools`` is NOT an exception: it re-wraps so
399 the stripper survives — see the ``bind_tools`` override below and #4804.)
401 A message keeps its object identity (only ``.content`` is rewritten,
402 so ``additional_kwargs``/``reasoning_content``/``tool_calls`` survive).
403 A bare-string return (some providers/wrappers) is wrapped into an
404 ``AIMessage`` so callers can always rely on ``.content``. Anything
405 else is passed through unchanged.
407 Only *string* ``.content`` is stripped: ``remove_think_tags`` is
408 text-only, and the ``<think>...</think>`` artifact is only ever
409 emitted as plain text by some local models. Non-string content
410 (e.g. provider content-block lists like Anthropic's, or ``None``)
411 is passed through untouched — running the regex on it would raise
412 ``TypeError`` or corrupt the structured content.
413 """
414 if hasattr(response, "content"):
415 if isinstance(response.content, str):
416 response.content = remove_think_tags(response.content)
417 elif isinstance(response, str):
418 response = AIMessage(content=remove_think_tags(response))
419 return response
421 @staticmethod
422 def _log_llm_error(error: Exception) -> None:
423 _log_llm_error(error)
425 def invoke(self, *args: Any, **kwargs: Any) -> Any:
426 try:
427 response = self.base_llm.invoke(*args, **kwargs)
428 except Exception as e:
429 self._log_llm_error(e)
430 raise
431 return self._normalize_response(response)
433 async def ainvoke(self, *args: Any, **kwargs: Any) -> Any:
434 # Async counterpart of invoke(); without this, ainvoke() would fall
435 # through __getattr__ to the base LLM and bypass think-tag stripping.
436 try:
437 response = await self.base_llm.ainvoke(*args, **kwargs)
438 except Exception as e:
439 self._log_llm_error(e)
440 raise
441 return self._normalize_response(response)
443 # Pass through any other attributes to the base LLM
444 def __getattr__(self, name):
445 return getattr(self.base_llm, name)
447 def close(self):
448 """Close underlying HTTP clients held by this LLM. Idempotent."""
449 try:
450 from ..utilities.llm_utils import _close_base_llm
452 _close_base_llm(self.base_llm)
453 except Exception:
454 logger.debug(
455 "best-effort cleanup of HTTP clients on shutdown",
456 exc_info=True,
457 )
459 def bind_tools(self, tools, **kwargs: Any) -> Any:
460 """Re-wrap the tool-bound model so the <think>-stripper survives.
462 ``langchain.agents.create_agent()`` and similar helpers call
463 ``model.bind_tools(...)`` and then use the *return value* as
464 the model that actually runs inside the agent loop. Without
465 this override, Python falls through ``__getattr__`` to
466 ``self.base_llm.bind_tools(...)``, which returns an unwrapped
467 ``BaseChatModel``; the wrapper's ``_normalize_response`` (the
468 single authorized site that strips ``<think>…`` from fresh
469 responses) is then bypassed for every LLM call the agent
470 makes.
472 That bypass is the root cause of the regression tracked in
473 issue #4804: reasoning-mode models (Qwen 3.x, deepseek-r1,
474 etc.) emit ``<think>…</think>`` as plain text, the wrapper
475 fails to strip it, the raw artifact is appended to the
476 agent's ``AIMessage.content``, and on the next turn the
477 provider's strict parser rejects the request with
478 ``Failed to parse input at pos 0: <think>…``.
480 Re-wrapping the bound model keeps the stripper in the loop
481 for every LLM call the agent makes, not just the ones we
482 issue directly via ``self.model.invoke(...)``.
484 Args:
485 tools: Tool definitions (schemas, callables, or
486 ``BaseTool`` instances) to bind to the model.
487 **kwargs: Provider-specific binding options forwarded
488 verbatim to ``base_llm.bind_tools``.
490 Returns:
491 A ``ProcessingLLMWrapper`` wrapping the bound model, so
492 every subsequent ``invoke``/``ainvoke`` continues to
493 strip ``<think>`` tags.
494 """
495 bound = self.base_llm.bind_tools(tools, **kwargs)
496 return ProcessingLLMWrapper(bound)
499def wrap_llm_without_think_tags(
500 llm,
501 research_id=None,
502 provider=None,
503 research_context=None,
504 settings_snapshot=None,
505):
506 """Create a wrapper class that processes LLM outputs with remove_think_tags and token counting"""
508 # First apply rate limiting if enabled
509 from ..web_search_engines.rate_limiting.llm import (
510 create_rate_limited_llm_wrapper,
511 )
513 # Check if LLM rate limiting is enabled (independent of search rate limiting)
514 # Use the thread-safe get_db_setting defined in this module
515 if get_setting_from_snapshot(
516 "rate_limiting.llm_enabled", False, settings_snapshot=settings_snapshot
517 ):
518 llm = create_rate_limited_llm_wrapper(llm, provider)
520 # Set context_limit in research_context for overflow detection.
521 # This is needed for providers that go through the registered provider path
522 # (which returns before the code in get_llm that sets context_limit).
523 if research_context is not None and provider is not None:
524 if "context_limit" not in research_context:
525 context_limit = _get_context_window_for_provider(
526 provider, settings_snapshot
527 )
528 if context_limit is not None:
529 research_context["context_limit"] = context_limit
530 logger.info(
531 f"Set context_limit={context_limit} in wrap_llm for provider={provider}"
532 )
534 # Import token counting functionality if research_id is provided
535 callbacks = []
536 if research_id is not None:
537 from ..metrics import TokenCounter
539 token_counter = TokenCounter()
540 token_callback = token_counter.create_callback(
541 research_id, research_context
542 )
543 # Set provider and model info on the callback
544 if provider:
545 token_callback.preset_provider = provider
546 # Try to extract model name from the LLM instance
547 if hasattr(llm, "model_name"):
548 token_callback.preset_model = llm.model_name
549 elif hasattr(llm, "model"):
550 token_callback.preset_model = llm.model
551 callbacks.append(token_callback)
553 # Add callbacks to the LLM if it supports them
554 if callbacks and hasattr(llm, "callbacks"):
555 if llm.callbacks is None:
556 llm.callbacks = callbacks
557 else:
558 llm.callbacks.extend(callbacks)
560 return ProcessingLLMWrapper(llm)