Coverage for src/local_deep_research/config/llm_config.py: 96%
144 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +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):
63 """
64 Get LLM instance based on model name and provider.
66 Args:
67 model_name: Name of the model to use (if None, uses database setting)
68 temperature: Model temperature (if None, uses database setting)
69 provider: Provider to use (if None, uses database setting)
70 openai_endpoint_url: NON-FUNCTIONAL, kept for API compatibility.
71 The endpoint URL is read exclusively from the
72 ``llm.openai_endpoint.url`` setting by
73 CustomOpenAIEndpointProvider. This parameter was already
74 ignored before the procedural-chain removal (the registry
75 dispatch never consumed it); honoring or removing it is
76 tracked as a follow-up.
77 research_id: Optional research ID for token tracking
78 research_context: Optional research context for enhanced token tracking
80 Returns:
81 A LangChain LLM instance with automatic think-tag removal
82 """
84 # Use database values for parameters if not provided
85 if model_name is None:
86 model_name = get_setting_from_snapshot(
87 "llm.model", "", settings_snapshot=settings_snapshot
88 )
89 if temperature is None:
90 temperature = get_setting_from_snapshot(
91 "llm.temperature", 0.7, settings_snapshot=settings_snapshot
92 )
93 if provider is None:
94 provider = get_setting_from_snapshot(
95 "llm.provider", "ollama", settings_snapshot=settings_snapshot
96 )
98 # Clean model name: remove quotes and extra whitespace
99 if model_name:
100 model_name = model_name.strip().strip("\"'").strip()
102 # Clean provider: remove quotes and extra whitespace
103 if provider: 103 ↛ 107line 103 didn't jump to line 107 because the condition on line 103 was always true
104 provider = provider.strip().strip("\"'").strip()
106 # Normalize provider: convert to lowercase canonical form
107 provider = normalize_provider(provider)
109 # Egress policy PEP for LLM endpoints. Fires here (before the registered-
110 # LLM dispatch) because all built-in providers are auto-registered via
111 # discover_providers(), so the registered branch handles every real LLM
112 # call.
113 #
114 # Two paths: snapshot-present runs the full PEP; snapshot-absent runs
115 # an allow-list check so background helpers / scaffolding paths cannot
116 # silently instantiate a cloud LLM. The allow-list (known-local) is
117 # deliberately tight — any provider not in it (including ambiguous
118 # ones like ``openai_endpoint`` and future cloud providers) fails
119 # closed instead of bypassing the PEP.
120 if provider: 120 ↛ 189line 120 didn't jump to line 189 because the condition on line 120 was always true
121 from ..security.egress.policy import (
122 Decision,
123 PolicyDeniedError,
124 _LOCAL_DEFAULT_LLM_PROVIDERS,
125 _is_user_registered_llm,
126 context_from_snapshot,
127 evaluate_llm_endpoint,
128 resolve_run_primary_engine,
129 )
131 if settings_snapshot is None:
132 # User-registered in-process LLMs are exempt here for the same
133 # reason evaluate_llm_endpoint allows them: no endpoint to
134 # classify, operator-injected, audit-hook backstopped.
135 if provider not in _LOCAL_DEFAULT_LLM_PROVIDERS and not (
136 _is_user_registered_llm(provider)
137 ):
138 logger.bind(policy_audit=True).warning(
139 "LLM constructed without policy snapshot; refusing "
140 "non-local provider",
141 provider=provider,
142 )
143 raise PolicyDeniedError(
144 Decision(False, "no_snapshot_for_provider"),
145 target=provider,
146 )
147 else:
148 try:
149 # Derive the run's primary the SAME way the search-engine
150 # factory does (single source of truth), instead of the old
151 # ``search.tool`` + searxng fallback. That fallback was a
152 # fail-OPEN: a missing/blank primary defaulted to searxng ->
153 # ADAPTIVE -> PUBLIC_ONLY -> require_local_llm stayed False ->
154 # the endpoint check below was skipped, admitting a CLOUD LLM
155 # for a run whose actual posture was private. resolve_run_
156 # primary_engine raises on a missing/invalid primary, which we
157 # treat as a hard stop here.
158 primary_engine = resolve_run_primary_engine(settings_snapshot)
159 ctx = context_from_snapshot(settings_snapshot, primary_engine)
160 except ValueError as exc:
161 # No configured primary, or an invalid policy config. Fail
162 # closed: previously a missing primary silently fell back to
163 # searxng/PUBLIC_ONLY and skipped the LLM endpoint check
164 # entirely, opening a cloud-LLM bypass under the very
165 # configuration the user asked to be strict about.
166 logger.bind(policy_audit=True).warning(
167 "no/invalid egress policy primary; refusing LLM",
168 provider=provider,
169 reason=str(exc),
170 )
171 raise PolicyDeniedError(
172 Decision(False, "invalid_policy_config"),
173 target=provider,
174 ) from exc
176 if ctx is not None and ctx.require_local_llm:
177 decision = evaluate_llm_endpoint(
178 provider, ctx, settings_snapshot=settings_snapshot
179 )
180 if not decision.allowed:
181 logger.bind(policy_audit=True).warning(
182 "LLM endpoint denied by egress policy",
183 provider=provider,
184 reason=decision.reason,
185 )
186 raise PolicyDeniedError(decision, target=provider)
188 # Check if this is a registered custom LLM first
189 if provider and is_llm_registered(provider):
190 logger.info(f"Using registered custom LLM: {provider}")
191 custom_llm = get_llm_from_registry(provider)
193 # Check if it's a callable (factory function) or a BaseChatModel instance
194 if callable(custom_llm) and not isinstance(custom_llm, BaseChatModel):
195 # It's a callable (factory function), call it with parameters
196 try:
197 llm_instance = custom_llm(
198 model_name=model_name,
199 temperature=temperature,
200 settings_snapshot=settings_snapshot,
201 )
202 except TypeError as e:
203 # Re-raise TypeError with better message
204 raise TypeError(
205 f"Registered LLM factory '{provider}' has invalid signature. "
206 f"Factory functions must accept 'model_name', 'temperature', and 'settings_snapshot' parameters. "
207 f"Error: {e}"
208 )
210 # Validate the result is a BaseChatModel
211 if not isinstance(llm_instance, BaseChatModel):
212 raise ValueError(
213 f"Factory function for {provider} must return a BaseChatModel instance, "
214 f"got {type(llm_instance).__name__}"
215 )
216 elif isinstance(custom_llm, BaseChatModel):
217 # It's already a proper LLM instance, use it directly
218 llm_instance = custom_llm
219 else:
220 raise ValueError(
221 f"Registered LLM {provider} must be either a BaseChatModel instance "
222 f"or a callable factory function. Got: {type(custom_llm).__name__}"
223 )
225 return wrap_llm_without_think_tags(
226 llm_instance,
227 research_id=research_id,
228 provider=provider,
229 research_context=research_context,
230 settings_snapshot=settings_snapshot,
231 )
233 # Validate the provider against the auto-discovered set — NOT a hardcoded
234 # list. Auto-discovery (discover_providers, run at module import) registers
235 # every llm/providers/implementations/*.py class, and the
236 # is_llm_registered() check above already serves them, so the registry /
237 # discovery IS the single source of truth for "valid provider".
238 #
239 # We deliberately do not keep a separate VALID_PROVIDERS constant: it was a
240 # third copy of the provider list (besides the implementations directory
241 # and the registry) and it silently drifted from auto-discovery (xai,
242 # ionos and deepseek were valid+registered yet missing from it). Deriving
243 # the set from discovery here means it can never drift. ('none' is the
244 # explicit "unset" sentinel, handled by the guard further down.)
245 from ..llm.providers import get_discovered_provider_options
247 valid_providers = {
248 normalize_provider(option["value"])
249 for option in get_discovered_provider_options()
250 } | {"none"}
251 if provider not in valid_providers:
252 logger.error(f"Invalid provider in settings: {provider}")
253 raise ValueError(
254 f"Invalid provider: {provider}. "
255 f"Must be one of: {sorted(valid_providers)}"
256 )
258 # Require an explicit model for built-in providers. Mirrors the
259 # API-key-not-configured pattern in openai_base.py and the URL-not-
260 # configured pattern in providers/implementations/ollama.py: no silent
261 # substitution to a hardcoded default model.
262 if not model_name or not model_name.strip():
263 logger.error("llm.model is not configured (empty/None after lookup)")
264 raise ValueError(
265 "LLM model not configured. Please open Settings, choose an LLM "
266 "provider, and select a model name (e.g. 'gpt-4o-mini' for "
267 "OpenAI, 'claude-3-5-sonnet-20241022' for Anthropic, "
268 "'llama3.1:8b' for Ollama). The 'llm.model' setting is required."
269 )
270 logger.info(
271 f"Getting LLM with model: {model_name}, temperature: {temperature}, provider: {provider}"
272 )
274 # Set context_limit on research_context for overflow detection. The
275 # actual max_tokens calculation lives in each provider's create_llm
276 # (via providers/_helpers.compute_max_tokens) so the cap is consistent
277 # across the live registered-LLM path.
278 context_window_size = _get_context_window_for_provider(
279 provider, settings_snapshot
280 )
281 if research_context and context_window_size: 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true
282 research_context["context_limit"] = context_window_size
283 logger.info(
284 f"Set context_limit={context_window_size} in research_context"
285 )
286 else:
287 logger.debug(
288 f"Context limit not set: research_context={bool(research_context)}, "
289 f"context_window_size={context_window_size}"
290 )
292 # Reaching here means the registered-LLM branch above did NOT fire,
293 # which is unusual — auto-discovery normally registers all 6+ built-in
294 # providers (anthropic, openai, openai_endpoint, ollama, lmstudio,
295 # llamacpp + xai/ionos/openrouter via OpenAICompatibleProvider) at
296 # import time. Two specific guards preserve the user-facing error
297 # messages for known-bad cases.
298 if provider == "none":
299 raise ValueError(
300 "No LLM provider configured. Please set llm.provider in settings "
301 "to a valid provider (e.g., 'ollama', 'openai', 'anthropic')."
302 )
303 raise ValueError(
304 f"Provider '{provider}' was not registered by auto-discovery. "
305 f"This usually indicates an import error during startup — check the "
306 f"logs for 'Error loading provider from <module>' messages. "
307 f"Note that clear_llm_registry()/unregister_llm() also remove "
308 f"built-in providers; discover_providers(force_refresh=True) "
309 f"restores them."
310 )
313def _log_llm_error(error: Exception) -> None:
314 """Log an LLM call failure with credential redaction."""
315 safe_msg = scrub_error(error)
316 logger.warning(f"LLM Request - Failed with error: {safe_msg}")
319class ProcessingLLMWrapper:
320 def __init__(self, base_llm):
321 self.base_llm = base_llm
323 @staticmethod
324 def _normalize_response(response: Any) -> Any:
325 """Strip <think> tags and normalize the response shape.
327 This is the SINGLE authorized place that strips <think> tags from a
328 fresh LLM response. Every LLM from get_llm() is wrapped here, so
329 downstream code can use ``response.content`` directly — do NOT add
330 per-site ``remove_think_tags`` / ``get_llm_response_text`` calls on
331 fresh ``invoke``/``ainvoke`` results; they are redundant and hide
332 bugs. (Exceptions: ``.stream()``/``.astream()``, ``with_config``,
333 ``bind`` and other Runnable transforms still bypass this wrapper via
334 ``__getattr__``, as do injected/unwrapped LLMs — those may still need
335 explicit handling. ``bind_tools`` is NOT an exception: it re-wraps so
336 the stripper survives — see the ``bind_tools`` override below and #4804.)
338 A message keeps its object identity (only ``.content`` is rewritten,
339 so ``additional_kwargs``/``reasoning_content``/``tool_calls`` survive).
340 A bare-string return (some providers/wrappers) is wrapped into an
341 ``AIMessage`` so callers can always rely on ``.content``. Anything
342 else is passed through unchanged.
344 Only *string* ``.content`` is stripped: ``remove_think_tags`` is
345 text-only, and the ``<think>...</think>`` artifact is only ever
346 emitted as plain text by some local models. Non-string content
347 (e.g. provider content-block lists like Anthropic's, or ``None``)
348 is passed through untouched — running the regex on it would raise
349 ``TypeError`` or corrupt the structured content.
350 """
351 if hasattr(response, "content"):
352 if isinstance(response.content, str):
353 response.content = remove_think_tags(response.content)
354 elif isinstance(response, str):
355 response = AIMessage(content=remove_think_tags(response))
356 return response
358 @staticmethod
359 def _log_llm_error(error: Exception) -> None:
360 _log_llm_error(error)
362 def invoke(self, *args: Any, **kwargs: Any) -> Any:
363 try:
364 response = self.base_llm.invoke(*args, **kwargs)
365 except Exception as e:
366 self._log_llm_error(e)
367 raise
368 return self._normalize_response(response)
370 async def ainvoke(self, *args: Any, **kwargs: Any) -> Any:
371 # Async counterpart of invoke(); without this, ainvoke() would fall
372 # through __getattr__ to the base LLM and bypass think-tag stripping.
373 try:
374 response = await self.base_llm.ainvoke(*args, **kwargs)
375 except Exception as e:
376 self._log_llm_error(e)
377 raise
378 return self._normalize_response(response)
380 # Pass through any other attributes to the base LLM
381 def __getattr__(self, name):
382 return getattr(self.base_llm, name)
384 def close(self):
385 """Close underlying HTTP clients held by this LLM. Idempotent."""
386 try:
387 from ..utilities.llm_utils import _close_base_llm
389 _close_base_llm(self.base_llm)
390 except Exception:
391 logger.debug(
392 "best-effort cleanup of HTTP clients on shutdown",
393 exc_info=True,
394 )
396 def bind_tools(self, tools, **kwargs: Any) -> Any:
397 """Re-wrap the tool-bound model so the <think>-stripper survives.
399 ``langchain.agents.create_agent()`` and similar helpers call
400 ``model.bind_tools(...)`` and then use the *return value* as
401 the model that actually runs inside the agent loop. Without
402 this override, Python falls through ``__getattr__`` to
403 ``self.base_llm.bind_tools(...)``, which returns an unwrapped
404 ``BaseChatModel``; the wrapper's ``_normalize_response`` (the
405 single authorized site that strips ``<think>…`` from fresh
406 responses) is then bypassed for every LLM call the agent
407 makes.
409 That bypass is the root cause of the regression tracked in
410 issue #4804: reasoning-mode models (Qwen 3.x, deepseek-r1,
411 etc.) emit ``<think>…</think>`` as plain text, the wrapper
412 fails to strip it, the raw artifact is appended to the
413 agent's ``AIMessage.content``, and on the next turn the
414 provider's strict parser rejects the request with
415 ``Failed to parse input at pos 0: <think>…``.
417 Re-wrapping the bound model keeps the stripper in the loop
418 for every LLM call the agent makes, not just the ones we
419 issue directly via ``self.model.invoke(...)``.
421 Args:
422 tools: Tool definitions (schemas, callables, or
423 ``BaseTool`` instances) to bind to the model.
424 **kwargs: Provider-specific binding options forwarded
425 verbatim to ``base_llm.bind_tools``.
427 Returns:
428 A ``ProcessingLLMWrapper`` wrapping the bound model, so
429 every subsequent ``invoke``/``ainvoke`` continues to
430 strip ``<think>`` tags.
431 """
432 bound = self.base_llm.bind_tools(tools, **kwargs)
433 return ProcessingLLMWrapper(bound)
436def wrap_llm_without_think_tags(
437 llm,
438 research_id=None,
439 provider=None,
440 research_context=None,
441 settings_snapshot=None,
442):
443 """Create a wrapper class that processes LLM outputs with remove_think_tags and token counting"""
445 # First apply rate limiting if enabled
446 from ..web_search_engines.rate_limiting.llm import (
447 create_rate_limited_llm_wrapper,
448 )
450 # Check if LLM rate limiting is enabled (independent of search rate limiting)
451 # Use the thread-safe get_db_setting defined in this module
452 if get_setting_from_snapshot(
453 "rate_limiting.llm_enabled", False, settings_snapshot=settings_snapshot
454 ):
455 llm = create_rate_limited_llm_wrapper(llm, provider)
457 # Set context_limit in research_context for overflow detection.
458 # This is needed for providers that go through the registered provider path
459 # (which returns before the code in get_llm that sets context_limit).
460 if research_context is not None and provider is not None:
461 if "context_limit" not in research_context:
462 context_limit = _get_context_window_for_provider(
463 provider, settings_snapshot
464 )
465 if context_limit is not None:
466 research_context["context_limit"] = context_limit
467 logger.info(
468 f"Set context_limit={context_limit} in wrap_llm for provider={provider}"
469 )
471 # Import token counting functionality if research_id is provided
472 callbacks = []
473 if research_id is not None:
474 from ..metrics import TokenCounter
476 token_counter = TokenCounter()
477 token_callback = token_counter.create_callback(
478 research_id, research_context
479 )
480 # Set provider and model info on the callback
481 if provider:
482 token_callback.preset_provider = provider
483 # Try to extract model name from the LLM instance
484 if hasattr(llm, "model_name"):
485 token_callback.preset_model = llm.model_name
486 elif hasattr(llm, "model"):
487 token_callback.preset_model = llm.model
488 callbacks.append(token_callback)
490 # Add callbacks to the LLM if it supports them
491 if callbacks and hasattr(llm, "callbacks"):
492 if llm.callbacks is None:
493 llm.callbacks = callbacks
494 else:
495 llm.callbacks.extend(callbacks)
497 return ProcessingLLMWrapper(llm)