Coverage for src/local_deep_research/utilities/llm_utils.py: 100%

1# utilities/llm_utils.py 

2""" 

3LLM utilities for Local Deep Research. 

4 

5This module provides utility functions for working with language models 

6when the user's llm_config.py is missing or incomplete. 

7""" 

8 

9import asyncio 

10import threading 

11from typing import Any, Optional, Dict 

12 

13from loguru import logger 

14 

15from ..config.constants import DEFAULT_OLLAMA_URL 

16from ..config.thread_settings import get_setting_from_snapshot 

17 

18 

19__all__ = [ 

20 "ainvoke_llm", 

21 "invoke_llm_sync", 

22 "get_model_identifier", 

23 "get_ollama_base_url", 

24 "get_server_url", 

25 "fetch_ollama_models", 

26] 

27 

28 

29def invoke_llm_sync(llm: Any, prompt: Any, **kwargs: Any) -> Any: 

30 """Invoke an LLM on the synchronous client (#5854). 

31 

32 This is deliberately a plain ``llm.invoke(prompt, **kwargs)`` and starts no 

33 event loop: langchain's cached ``httpx.AsyncClient`` is bound to the 

34 loop that first used it, so bridging a sync caller onto ``ainvoke`` 

35 with a throwaway loop double-sends or hard-fails from the second call 

36 onward (see #6293). The helper is the research pipeline's one 

37 synchronous dispatch seam — question generators, filters, explorers, 

38 context managers and evaluators call through it (exception below) — 

39 and it is the canonical production path. If I/O-bound fan-out is wired, it means 

40 callers submit :func:`ainvoke_llm` coroutines to the process's single 

41 shared event loop, not routing this helper through that loop; this 

42 helper itself never opens or uses a per-call loop. 

43 

44 Returns the RAW response object, so each caller's existing parsing 

45 (``get_llm_response_text`` / ``.content`` / ``extract_json``) is 

46 untouched. Per-call model options pass through the existing wrappers 

47 without changing the shared model's defaults or binding around a 

48 rate-limiting wrapper. 

49 """ 

50 return llm.invoke(prompt, **kwargs) 

51 

52 

53async def ainvoke_llm(llm: Any, prompt: Any) -> Any: 

54 """Async core matching :func:`invoke_llm_sync` (#5854). 

55 

56 Additive: nothing in production awaits this yet. Before it is usable, 

57 the process needs a single long-lived event loop for the LLM clients 

58 (#6293), and three ADR-0013 preconditions 

59 (docs/decisions/0013-research-runs-keep-daemon-thread.md) must hold: 

60 ``TokenCountingCallback`` needs to stop writing its call-usage row to 

61 the user's encrypted database from the loop's default executor (rule 

62 5; #6294 tracks the related attribution gap); the egress audit 

63 context is thread-local and must be passed in and re-armed on the 

64 loop side; and the submit helper must run the coroutine in an 

65 explicit minimal context so it does not inherit the daemon thread's 

66 ambient search context, including the password (rule 4). 

67 ``RateLimitedLLMWrapper.ainvoke`` already awaits the wrapped 

68 provider's own async client and was hardened in #6347; that is the 

69 correct shape for a fan-out interface and the pattern this helper 

70 follows. 

71 

72 Every production model is a ``ProcessingLLMWrapper``, which defines 

73 ``ainvoke``; a wrapped object without it is a configuration error and 

74 raises ``AttributeError`` rather than silently falling back to the 

75 sync client. 

76 """ 

77 return await llm.ainvoke(prompt) 

78 

79 

80def get_model_identifier(llm: Any) -> str: 

81 """Return a stable string identifier for an LLM instance. 

82 

83 The identifier is used as a cache key: `Journal.quality_model` records 

84 which LLM scored a cached journal, and the lookup predicate filters on 

85 it so scores from a superseded model don't get served. 

86 

87 Discovery order: 

88 1. Unwrap `ProcessingLLMWrapper` (or any wrapper exposing `.base_llm`) 

89 so we key on the underlying model, not the wrapper identity. 

90 2. Prefer `model_name` (some LangChain classes). Then `model` 

91 (ChatOpenAI, ChatAnthropic, ChatOllama use this). Fallback to the 

92 class name so we never return an ephemeral `repr(object)` that 

93 poisons the cache. 

94 

95 Returns a plain string; never None. Values written by `getattr(llm, 

96 "name", str(llm))` in earlier versions (e.g. `"<ProcessingLLMWrapper 

97 object at 0x…>"`) naturally miss this cache and re-score once. 

98 """ 

99 base = getattr(llm, "base_llm", llm) 

100 for attr in ("model_name", "model"): 

101 val = getattr(base, attr, None) 

102 if val: 

103 return str(val) 

104 return type(base).__name__ 

105 

106 

107def _close_base_llm(llm): 

108 """Close per-instance HTTP clients on a raw LLM or embeddings instance. 

109 Internal use only. 

110 

111 Applies to every ``langchain_ollama`` class — ``ChatOllama``, 

112 ``OllamaLLM``, and ``OllamaEmbeddings`` all carry the same 

113 ``_client`` / ``_async_client`` shape, eagerly constructed at instance 

114 init by a Pydantic ``@model_validator(mode="after")``. ``ChatAnthropic`` 

115 and ``ChatOpenAI`` use ``@lru_cache``'d shared httpx clients that must 

116 NOT be closed; the ``ollama``-module check below short-circuits cleanly 

117 for them. Raw ``HuggingFaceEmbeddings`` / 

118 ``SentenceTransformerEmbeddings`` instances are skipped the same way (no 

119 per-instance httpx client at all). The Sentence Transformers provider 

120 returns a shared-model handle with its own ``close()``, which the 

121 delegation below calls; it closes only that handle, not the shared model. 

122 

123 Each Ollama class owns both ``_client`` (sync ``ollama.Client`` wrapping 

124 ``httpx.Client``) and ``_async_client`` (async ``ollama.AsyncClient`` 

125 wrapping ``httpx.AsyncClient``). Async paths via ``ainvoke()`` 

126 (exercised by browsecomp_entity_strategy, llm_driven_modular_strategy, 

127 modular_strategy) leak the async transport per call if only the sync 

128 side is released — investigated in #3816 where ~72% of leaked FDs were 

129 ``a_inode [eventpoll]`` selectors bound to those async clients. The 

130 same shape reappeared for embeddings after the langchain_community → 

131 langchain_ollama migration (#4352/#4353); see the resource-cleanup doc 

132 for the post-mortem. 

133 

134 Idempotent via an ``_ldr_closed`` sentinel on the inner httpx clients. 

135 

136 The async path always runs ``aclose()`` to completion: when no event 

137 loop is currently running we use ``asyncio.run()`` directly; when a 

138 loop is running in the calling thread (e.g. ``_close_base_llm`` is 

139 invoked inside async code or in a ``finally`` block reached through 

140 LangGraph/LangChain async dispatch) we hand the close off to a brief 

141 daemon thread whose own ``asyncio.run()`` is unaffected by the 

142 caller's loop state. A prior implementation skipped the close in 

143 that case and relied on the "loop owner" to close — but no loop 

144 owner code actually does, so the ``httpx.AsyncClient`` and its 

145 ``epoll_create`` FD were silently leaked. See the regression history 

146 in ``docs/developing/resource-cleanup.md`` (this is the gap left by 

147 #3855 when reaching for the in-async-context close). 

148 """ 

149 # If the llm is another wrapper with its own close(), delegate. 

150 # NOTE: if a future ChatOllama version adds a public close() method, 

151 # this short-circuit fires and the introspection below is skipped — 

152 # that future close() must then handle BOTH sync AND async clients. 

153 if hasattr(type(llm), "close"): 

154 llm.close() 

155 return 

156 

157 # Sync side: ollama.Client._client is an httpx.Client. 

158 # ``_ldr_closed is True`` (not just truthy) so we don't trip on Mock 

159 # objects without a spec, where attribute access auto-generates a child 

160 # Mock that is truthy by default. 

161 sync_ollama = getattr(llm, "_client", None) 

162 if sync_ollama is not None and type(sync_ollama).__module__.startswith( 

163 "ollama" 

164 ): 

165 sync_httpx = getattr(sync_ollama, "_client", None) 

166 if ( 

167 sync_httpx is not None 

168 and getattr(sync_httpx, "_ldr_closed", None) is not True 

169 and hasattr(sync_httpx, "close") 

170 ): 

171 try: 

172 sync_httpx.close() 

173 except Exception: 

174 logger.warning("Failed to close Ollama sync httpx client") 

175 sync_httpx._ldr_closed = True 

176 

177 # Async side: ollama.AsyncClient._client is an httpx.AsyncClient 

178 async_ollama = getattr(llm, "_async_client", None) 

179 if async_ollama is not None and type(async_ollama).__module__.startswith( 

180 "ollama" 

181 ): 

182 async_httpx = getattr(async_ollama, "_client", None) 

183 if ( 

184 async_httpx is not None 

185 and getattr(async_httpx, "_ldr_closed", None) is not True 

186 and hasattr(async_httpx, "aclose") 

187 ): 

188 try: 

189 asyncio.get_running_loop() 

190 except RuntimeError: 

191 # No running loop in this thread: spin a temporary one 

192 # to await aclose() right here. 

193 try: 

194 asyncio.run(async_httpx.aclose()) 

195 except Exception: 

196 logger.warning("Failed to close Ollama async httpx client") 

197 # Mark closed unconditionally — matches the sync-side 

198 # invariant: on a known-broken close we don't want to 

199 # retry endlessly. The WARNING above is the signal. 

200 async_httpx._ldr_closed = True 

201 else: 

202 # A loop is running in this thread. ``asyncio.run`` cannot 

203 # be called here and fire-and-forget tasks scheduled on the 

204 # caller's loop from a finally block are unreliable (the 

205 # loop may exit before the task is awaited). Hand the close 

206 # off to a brief daemon thread whose own loop is 

207 # independent of ours; cap with a bounded ``join`` so a 

208 # stuck server can't hold up shutdown. 

209 

210 def _close_in_thread() -> None: 

211 try: 

212 asyncio.run(async_httpx.aclose()) 

213 except Exception: 

214 logger.warning( 

215 "Failed to close Ollama async httpx client " 

216 "in cleanup thread" 

217 ) 

218 

219 t = threading.Thread( 

220 target=_close_in_thread, 

221 daemon=True, 

222 name="ldr-async-llm-close", 

223 ) 

224 t.start() 

225 t.join(timeout=5) 

226 if t.is_alive(): 

227 # Don't set ``_ldr_closed`` — the cleanup thread is 

228 # still running and the FD is therefore still open. 

229 # Surface at WARNING so operators can correlate 

230 # against rising eventpoll-FD counts. A subsequent 

231 # call to _close_base_llm will retry the close. 

232 logger.warning( 

233 "Async httpx close exceeded 5s; abandoning to GC. " 

234 "If this fires repeatedly, check Ollama server " 

235 "responsiveness and look for rising " 

236 "anon_inode:[eventpoll] FDs on the process." 

237 ) 

238 else: 

239 # Thread completed (with or without an inner 

240 # exception). Mark closed to match the sync-side 

241 # invariant; the inner exception, if any, was 

242 # already logged from inside the thread. 

243 async_httpx._ldr_closed = True 

244 

245 

246def _close_inner_ollama_clients(sync_client, async_client): 

247 """Close just the inner sync/async ``ollama.Client`` pair. 

248 

249 A ``weakref.finalize`` callback that strong-refs the wrapping LLM or 

250 embeddings instance would defeat its own purpose — the registry's 

251 reference would keep the instance alive forever. Callers (the Ollama 

252 provider factory's safety net) pass the inner clients directly 

253 instead. This shim wraps them in a ``_close_base_llm``-shaped proxy 

254 so we reuse the same idempotent sync+async close logic without 

255 duplicating its asyncio/eventpoll handling. 

256 """ 

257 

258 class _Proxy: 

259 pass 

260 

261 proxy = _Proxy() 

262 proxy._client = sync_client 

263 proxy._async_client = async_client 

264 _close_base_llm(proxy) 

265 

266 

267def get_ollama_base_url( 

268 settings_snapshot: Optional[Dict[str, Any]] = None, 

269) -> str: 

270 """ 

271 Get Ollama base URL from settings with normalization. 

272 

273 Checks both embeddings.ollama.url and llm.ollama.url settings, 

274 falling back to http://localhost:11434. 

275 

276 Args: 

277 settings_snapshot: Optional settings snapshot 

278 

279 Returns: 

280 Normalized Ollama base URL 

281 """ 

282 from .url_utils import normalize_url 

283 

284 raw_base_url = get_setting_from_snapshot( 

285 "embeddings.ollama.url", 

286 default=get_setting_from_snapshot( 

287 "llm.ollama.url", # Fall back to LLM setting 

288 default=DEFAULT_OLLAMA_URL, 

289 settings_snapshot=settings_snapshot, 

290 ), 

291 settings_snapshot=settings_snapshot, 

292 ) 

293 return normalize_url(raw_base_url) if raw_base_url else DEFAULT_OLLAMA_URL 

294 

295 

296def get_server_url(settings_snapshot: Optional[Dict[str, Any]] = None) -> str: 

297 """ 

298 Get server URL from settings with fallback logic. 

299 

300 Checks multiple sources in order: 

301 1. Direct server_url in settings snapshot 

302 2. system.server_url in settings 

303 3. Constructs from web.host, web.port, and web.use_https 

304 4. Fallback to http://127.0.0.1:5000/ 

305 

306 Args: 

307 settings_snapshot: Optional settings snapshot 

308 

309 Returns: 

310 Server URL. URLs built in steps 3 and 4 end with a slash; a 

311 server_url found in steps 1 or 2 is returned as given. 

312 """ 

313 

314 server_url = None 

315 

316 if settings_snapshot: 

317 # Try to get server URL from research metadata first (where we added it) 

318 server_url = settings_snapshot.get("server_url") 

319 

320 # If not found, try system settings 

321 if not server_url: 

322 system_settings = settings_snapshot.get("system", {}) 

323 server_url = system_settings.get("server_url") 

324 

325 # If not found, try web.host and web.port settings. Keyword 

326 # arguments matter: the signature is (key, default, username, 

327 # settings_snapshot), so a positional snapshot became the default 

328 # and was never read. 

329 if not server_url: 

330 host = get_setting_from_snapshot( 

331 "web.host", 

332 default="127.0.0.1", 

333 settings_snapshot=settings_snapshot, 

334 ) 

335 port = get_setting_from_snapshot( 

336 "web.port", 

337 default=5000, 

338 settings_snapshot=settings_snapshot, 

339 ) 

340 use_https = get_setting_from_snapshot( 

341 "web.use_https", 

342 default=True, 

343 settings_snapshot=settings_snapshot, 

344 ) 

345 

346 # Use localhost for 0.0.0.0 bindings as that's what users will use 

347 # Replace a bind-all address with localhost for a client URL. 

348 if host == "0.0.0.0": # nosec B104 

349 host = "127.0.0.1" 

350 

351 scheme = "https" if use_https else "http" 

352 server_url = f"{scheme}://{host}:{port}/" 

353 

354 # Fallback to default if still not found 

355 if not server_url: 

356 server_url = "http://127.0.0.1:5000/" 

357 logger.warning("Could not determine server URL, using default") 

358 

359 return server_url 

360 

361 

362# Timeout (seconds) for listing installed Ollama models via /api/tags. A cold 

363# Ollama — process just started, connection not yet warm, host busy right after 

364# a restart — regularly needs more than a couple seconds, and a timeout here is 

365# silently turned into an empty model list (which surfaces as an empty model 

366# dropdown in the UI). Shared by the LLM and embedding Ollama providers so the 

367# two can't drift back apart. 

368OLLAMA_MODEL_LIST_TIMEOUT = 5.0 

369 

370 

371def fetch_ollama_models( 

372 base_url: str, 

373 timeout: float = OLLAMA_MODEL_LIST_TIMEOUT, 

374 auth_headers: Optional[Dict[str, str]] = None, 

375) -> list[Dict[str, str]]: 

376 """ 

377 Fetch available models from Ollama API. 

378 

379 Centralized function to avoid duplication between LLM and embedding providers. 

380 

381 Args: 

382 base_url: Ollama base URL (should be normalized) 

383 timeout: Request timeout in seconds 

384 auth_headers: Optional authentication headers 

385 

386 Returns: 

387 List of model dicts with 'value' (model name) and 'label' (display name) keys. 

388 Returns empty list on error. 

389 """ 

390 from ..security import safe_get 

391 

392 models = [] 

393 

394 try: 

395 response = safe_get( 

396 f"{base_url}/api/tags", 

397 timeout=timeout, 

398 headers=auth_headers or {}, 

399 allow_localhost=True, 

400 allow_private_ips=True, 

401 ) 

402 

403 if response.status_code == 200: 

404 data = response.json() 

405 

406 # Handle both newer and older Ollama API formats 

407 ollama_models = ( 

408 data.get("models", []) if isinstance(data, dict) else data 

409 ) 

410 

411 for model_data in ollama_models: 

412 model_name = model_data.get("name", "") 

413 if model_name: 

414 models.append({"value": model_name, "label": model_name}) 

415 

416 logger.info(f"Found {len(models)} Ollama models") 

417 else: 

418 logger.warning( 

419 f"Failed to fetch Ollama models: HTTP {response.status_code}" 

420 ) 

421 

422 except Exception: 

423 logger.exception("Error fetching Ollama models") 

424 

425 return models