Coverage for src/local_deep_research/web_search_engines/engines/search_engine_elasticsearch.py: 94%
240 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1import json
2import threading
3import time as _time
4from typing import Any, Dict, List, Optional
6from elasticsearch import Elasticsearch
7from langchain_core.language_models import BaseLLM
9from ...constants import DEFAULT_SEARCH_TOOL, SNIPPET_LENGTH_SHORT
10from ...security.secure_logging import logger
11from ...settings.exceptions import SettingsValueError
12from ..search_engine_base import BaseSearchEngine, Exposure, Sensitivity
15class ElasticsearchSearchEngine(BaseSearchEngine):
16 """Elasticsearch search engine implementation with two-phase approach"""
18 is_local = True
19 is_lexical = True
20 needs_llm_relevance_filter = True
21 # Egress (ADR-0007): a local document store — sensitive, contained. The
22 # url_setting fail-up reclassifies exposure to EXPOSING when hosts resolve
23 # public (quadrant 4: usable only by itself, contained inference).
24 egress_sensitivity = Sensitivity.SENSITIVE
25 egress_exposure = Exposure.CONTAINED
26 # secrets to redact from error messages (see BaseSearchEngine._scrub_error)
27 _secret_attrs = ("_api_key", "_password")
28 # url_setting feeds the PDP's fail-up URL override: when the configured
29 # hosts resolve to a PUBLIC endpoint (e.g. Elastic Cloud), the engine is
30 # reclassified public so PRIVATE_ONLY denies it at selection time —
31 # queries would leave the box even though the DATA is "local" in nature.
32 # A localhost ES keeps the static is_local classification above.
33 # cloud_id (which is NOT a host the PDP can classify) is handled
34 # separately in __init__: it is rejected when the effective scope forbids
35 # public egress.
36 url_setting = "search.engine.web.elasticsearch.default_params.hosts"
38 @staticmethod
39 def _cloud_id_forbidden_by_scope(
40 settings_snapshot: Optional[Dict[str, Any]],
41 ) -> bool:
42 """True when the effective egress scope forbids the public Elastic
43 Cloud endpoint a ``cloud_id`` targets.
45 Resolves the scope (including ADAPTIVE) via ``context_from_snapshot``
46 and returns True for PRIVATE_ONLY / STRICT. Fails CLOSED (forbidden)
47 if the policy cannot be evaluated, so a snapshot/policy error cannot
48 open a cloud egress under a private posture. A missing/empty snapshot
49 resolves to the permissive default (BOTH) and is allowed.
50 """
51 try:
52 from ...security.egress.policy import (
53 EgressScope,
54 context_from_snapshot,
55 resolve_run_primary_engine,
56 )
57 from ...search_system import username_from_snapshot
59 snapshot = settings_snapshot or {}
60 primary = resolve_run_primary_engine(
61 snapshot, default=DEFAULT_SEARCH_TOOL
62 )
63 ctx = context_from_snapshot(
64 snapshot, primary, username=username_from_snapshot(snapshot)
65 )
66 return ctx.scope in (
67 EgressScope.PRIVATE_ONLY,
68 EgressScope.STRICT,
69 )
70 except Exception:
71 logger.bind(policy_audit=True).warning(
72 "elasticsearch cloud_id egress check failed; failing closed",
73 exc_info=True,
74 )
75 return True
77 # TTL cache for ``is_available`` probes. Keyed by JSON string representation of hosts so
78 # two instances pointing at different clusters don't poison each other (and dict hosts don't raise TypeError).
79 # The cache lives on the class so test patches that replace the class don't fight a stale module-level dict.
80 # A negative result (connection refused / timeout) is cached for the same TTL — we don't
81 # want a misconfigured engine to cost a TCP round-trip on every research run.
82 _availability_cache: Dict[str, tuple[float, bool]] = {}
83 _availability_cache_lock = threading.Lock()
84 _AVAILABILITY_TTL_SECONDS = 60.0
86 @classmethod
87 def clear_availability_cache(cls) -> None:
88 """Clear the cached availability probe results."""
89 with cls._availability_cache_lock:
90 cls._availability_cache.clear()
92 @classmethod
93 def _get_cached_availability(
94 cls, cache_key: str, now: float
95 ) -> Optional[bool]:
96 """Prune expired probe entries and return a fresh cached result."""
97 with cls._availability_cache_lock:
98 for key, (timestamp, _available) in list(
99 cls._availability_cache.items()
100 ):
101 if now - timestamp >= cls._AVAILABILITY_TTL_SECONDS:
102 cls._availability_cache.pop(key, None)
104 cached = cls._availability_cache.get(cache_key)
105 return cached[1] if cached is not None else None
107 @classmethod
108 def is_available(
109 cls, settings_snapshot: Optional[Dict[str, Any]] = None
110 ) -> bool:
111 """Probe the configured Elasticsearch host(s) with a cheap TCP connect.
113 ``ElasticsearchSearchEngine.__init__`` calls ``self.client.info()``
114 which raises ``ConnectionError`` when the cluster is unreachable.
115 Without this probe the factory logs ``Failed to create search
116 engine 'elasticsearch' (ConnectionError)`` on every tool call, and
117 worse the langgraph agent still advertises the engine as a tool in
118 its per-step "selecting next action from …" heartbeat.
120 We do a bare TCP connect (not the full ES handshake) so the probe
121 stays cheap enough to call from ``list_eligible_engine_configs``.
122 Negative results are cached for ``_AVAILABILITY_TTL_SECONDS`` so a
123 single down cluster doesn't cost a round-trip per research run.
125 Returns True (fail-open) for ``cloud_id`` configurations — those
126 target Elastic Cloud and a TCP probe isn't meaningful. They'll
127 surface their own error in ``__init__`` if misconfigured.
128 """
129 from ...config.thread_settings import get_setting_from_snapshot
131 try:
132 snapshot = settings_snapshot or {}
134 # cloud_id configs don't have a host we can TCP-probe. Trust them;
135 # __init__ surfaces a useful error if the cluster is unreachable.
136 cloud_id = get_setting_from_snapshot(
137 "search.engine.web.elasticsearch.default_params.cloud_id",
138 "",
139 settings_snapshot=snapshot,
140 )
141 if cloud_id:
142 return True
144 hosts_setting = get_setting_from_snapshot(
145 "search.engine.web.elasticsearch.default_params.hosts",
146 ["http://localhost:9200"],
147 settings_snapshot=snapshot,
148 )
149 # The setting can arrive as a JSON string (ui_element=json) or a
150 # real list — mirror the normalization __init__ does.
151 hosts = cls._ensure_list(
152 hosts_setting, default=["http://localhost:9200"]
153 )
154 try:
155 cache_key = json.dumps(hosts, sort_keys=True, default=str)
156 except Exception:
157 cache_key = str(hosts)
159 now = _time.monotonic()
160 cached = cls._get_cached_availability(cache_key, now)
161 if cached is not None:
162 return cached
164 available = cls._probe_hosts_available(hosts)
165 with cls._availability_cache_lock:
166 cls._availability_cache[cache_key] = (now, available)
167 if not available:
168 logger.info(
169 "Elasticsearch availability probe failed; excluding from "
170 "agent tool list. Will re-probe in {:.0f}s.",
171 cls._AVAILABILITY_TTL_SECONDS,
172 )
173 return available
174 except SettingsValueError:
175 raise
176 except Exception as exc:
177 logger.debug(
178 "Elasticsearch is_available probe raised exception ({}) — failing open",
179 type(exc).__name__,
180 )
181 return True
183 @staticmethod
184 def _probe_hosts_available(hosts: List[Any]) -> bool:
185 """TCP-connect to the first reachable host in ``hosts``.
187 Any single host responding makes the engine available — matches the
188 Elasticsearch client behavior of treating the list as failover.
189 Cheap (no HTTP) and short-timeout (1s per host, capped at a 2.0s
190 connection budget across all hosts). SSRF validation runs before each
191 connection and can perform DNS resolution that cannot be interrupted
192 by this connection-only budget.
193 """
194 import socket
195 from urllib.parse import urlparse
196 from ...security.ssrf_validator import validate_url
198 start_time = _time.monotonic()
199 total_budget = 2.0 # Cap aggregate probe duration across all hosts
201 for host in hosts:
202 elapsed = _time.monotonic() - start_time
203 remaining = total_budget - elapsed
204 if remaining <= 0: 204 ↛ 205line 204 didn't jump to line 205 because the condition on line 204 was never true
205 break
207 if isinstance(host, dict):
208 hostname = host.get("host") or host.get("hostname")
209 if not hostname: 209 ↛ 210line 209 didn't jump to line 210 because the condition on line 209 was never true
210 continue
211 scheme = host.get("scheme", "http")
212 port = host.get("port")
213 if port is None: 213 ↛ 214line 213 didn't jump to line 214 because the condition on line 213 was never true
214 port = 443 if scheme == "https" else 9200
215 elif isinstance(host, str): 215 ↛ 227line 215 didn't jump to line 227 because the condition on line 215 was always true
216 url_str = host if "://" in host else f"http://{host}"
217 try:
218 parsed = urlparse(url_str)
219 hostname = parsed.hostname or host
220 scheme = parsed.scheme or "http"
221 port = parsed.port
222 if port is None:
223 port = 443 if scheme == "https" else 9200
224 except ValueError:
225 continue
226 else:
227 continue
229 connection_host = str(hostname).strip("[]")
230 url_host = (
231 f"[{connection_host}]"
232 if ":" in connection_host
233 else connection_host
234 )
235 probe_url = f"{scheme}://{url_host}:{port}"
236 if not validate_url(
237 probe_url, allow_localhost=True, allow_private_ips=True
238 ):
239 logger.warning(
240 "Elasticsearch host failed SSRF validation: {}", probe_url
241 )
242 continue
244 timeout = min(remaining, 1.0)
245 try:
246 with socket.create_connection(
247 (connection_host, port), timeout=timeout
248 ):
249 return True
250 except (OSError, socket.timeout, ValueError):
251 continue
252 return False
254 def __init__(
255 self,
256 hosts: Optional[List[str]] = None,
257 index_name: str = "documents",
258 username: Optional[str] = None,
259 password: Optional[str] = None,
260 api_key: Optional[str] = None,
261 cloud_id: Optional[str] = None,
262 max_results: int = 10,
263 highlight_fields: List[str] = ["content", "title"],
264 search_fields: List[str] = ["content", "title"],
265 filter_query: Optional[Dict[str, Any]] = None,
266 llm: Optional[BaseLLM] = None,
267 max_filtered_results: Optional[int] = None,
268 settings_snapshot: Optional[Dict[str, Any]] = None,
269 ):
270 """
271 Initialize the Elasticsearch search engine.
273 Args:
274 hosts: List of Elasticsearch hosts
275 index_name: Name of the index to search
276 username: Optional username for authentication
277 password: Optional password for authentication
278 api_key: Optional API key for authentication
279 cloud_id: Optional Elastic Cloud ID
280 max_results: Maximum number of search results
281 highlight_fields: Fields to highlight in search results
282 search_fields: Fields to search in
283 filter_query: Optional filter query in Elasticsearch DSL format
284 llm: Language model for relevance filtering
285 max_filtered_results: Maximum number of results to keep after filtering
286 """
287 # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
288 super().__init__(
289 llm=llm,
290 max_filtered_results=max_filtered_results,
291 max_results=max_results,
292 settings_snapshot=settings_snapshot,
293 )
295 self.index_name = index_name
296 self.highlight_fields = self._ensure_list(
297 highlight_fields, default=["content", "title"]
298 )
299 self.search_fields = self._ensure_list(
300 search_fields, default=["content", "title"]
301 )
302 self.filter_query = filter_query or {}
304 # Store credentials for error-message redaction
305 self._api_key = api_key
306 self._password = password
308 # Normalize hosts – may arrive as a JSON-encoded string from settings
309 hosts = self._ensure_list(hosts, default=["http://localhost:9200"])
311 # Initialize the Elasticsearch client
312 es_args: Dict[str, Any] = {}
314 # Basic authentication
315 if username and password:
316 es_args["basic_auth"] = (username, password)
318 # API key authentication
319 if api_key:
320 es_args["api_key"] = api_key
322 # Cloud ID for Elastic Cloud
323 if cloud_id:
324 # Egress policy: a cloud_id always targets a public Elastic Cloud
325 # endpoint (*.cloud.es.io), but the url_setting reclassification
326 # only inspects `hosts`. A cloud_id-only config would otherwise
327 # keep the engine's static is_local=True and slip past
328 # evaluate_engine, then connect at self.client.info() below. Reject
329 # it when the effective scope forbids public egress (fail closed).
330 if self._cloud_id_forbidden_by_scope(settings_snapshot):
331 from ...security.egress.policy import (
332 Decision,
333 PolicyDeniedError,
334 )
336 logger.bind(policy_audit=True).warning(
337 "refusing Elasticsearch cloud_id under private egress scope"
338 )
339 raise PolicyDeniedError(
340 Decision(False, "elasticsearch_cloud_id_public_egress"),
341 target="search_engine:elasticsearch",
342 )
343 es_args["cloud_id"] = cloud_id
345 # Connect to Elasticsearch
346 self.client = Elasticsearch(hosts, **es_args)
348 # Verify connection
349 try:
350 info = self.client.info()
351 logger.info(
352 f"Connected to Elasticsearch cluster: {info.get('cluster_name')}"
353 )
354 logger.info(
355 f"Elasticsearch version: {info.get('version', {}).get('number')}"
356 )
357 except Exception as e:
358 safe_msg = self._scrub_error(e)
359 logger.warning(f"Failed to connect to Elasticsearch: {safe_msg}")
360 raise ConnectionError(
361 f"Could not connect to Elasticsearch: {safe_msg}"
362 ) from None
364 def close(self) -> None:
365 """Close the Elasticsearch client and its connection pool."""
366 from ...utilities.resource_utils import safe_close
368 safe_close(self.client, "Elasticsearch client")
369 super().close()
371 def _get_previews(self, query: str) -> List[Dict[str, Any]]:
372 """
373 Get preview information for Elasticsearch documents.
375 Args:
376 query: The search query
378 Returns:
379 List of preview dictionaries
380 """
381 logger.info("Getting document previews from Elasticsearch")
383 try:
384 # Build the search query
385 search_query = {
386 "query": {
387 "multi_match": {
388 "query": query,
389 "fields": self.search_fields,
390 "type": "best_fields",
391 "tie_breaker": 0.3,
392 }
393 },
394 "highlight": {
395 "fields": {field: {} for field in self.highlight_fields},
396 "pre_tags": ["<em>"],
397 "post_tags": ["</em>"],
398 },
399 "size": self.max_results,
400 }
402 # Add filter if provided
403 if self.filter_query:
404 search_query["query"] = {
405 "bool": {
406 "must": search_query["query"],
407 "filter": self.filter_query,
408 }
409 }
411 # Execute the search
412 response = self.client.search(
413 index=self.index_name,
414 body=search_query,
415 )
417 # Process the search results
418 hits = response.get("hits", {}).get("hits", [])
420 # Format results as previews with basic information
421 previews = []
422 for hit in hits:
423 source = hit.get("_source", {})
424 highlight = hit.get("highlight", {})
426 # Extract highlighted snippets or fall back to original content
427 snippet = ""
428 for field in self.highlight_fields:
429 if highlight.get(field):
430 # Join all highlights for this field
431 field_snippets = " ... ".join(highlight[field])
432 snippet += field_snippets + " "
434 # If no highlights, use a portion of the content
435 if not snippet and "content" in source:
436 content = source.get("content", "")
437 snippet = (
438 content[:SNIPPET_LENGTH_SHORT] + "..."
439 if len(content) > SNIPPET_LENGTH_SHORT
440 else content
441 )
443 # Create preview object
444 preview = {
445 "id": hit.get("_id", ""),
446 "title": source.get("title", "Untitled Document"),
447 "link": source.get("url", "")
448 or f"elasticsearch://{self.index_name}/{hit.get('_id', '')}",
449 "snippet": snippet.strip(),
450 "score": hit.get("_score", 0),
451 "_index": hit.get("_index", self.index_name),
452 }
454 previews.append(preview)
456 logger.info(
457 f"Found {len(previews)} preview results from Elasticsearch"
458 )
459 return previews
461 except Exception as e:
462 safe_msg = self._scrub_error(e)
463 logger.warning(f"Error getting Elasticsearch previews: {safe_msg}")
464 return []
466 def _get_full_content(
467 self, relevant_items: List[Dict[str, Any]]
468 ) -> List[Dict[str, Any]]:
469 """
470 Get full content for the relevant Elasticsearch documents.
472 Args:
473 relevant_items: List of relevant preview dictionaries
475 Returns:
476 List of result dictionaries with full content
477 """
478 logger.info("Getting full content for relevant Elasticsearch documents")
480 results = []
481 for item in relevant_items:
482 # Start with the preview data
483 result = item.copy()
485 # Get the document ID
486 doc_id = item.get("id")
487 if not doc_id:
488 # Skip items without ID
489 logger.warning(f"Skipping item without ID: {item}")
490 results.append(result)
491 continue
493 try:
494 # Fetch the full document
495 doc_response = self.client.get(
496 index=self.index_name,
497 id=doc_id,
498 )
500 # Get the source document
501 source = doc_response.get("_source", {})
503 # Add full content to the result
504 result["content"] = source.get(
505 "content", result.get("snippet", "")
506 )
507 result["full_content"] = source.get("content", "")
508 result["full_content_is_complete"] = bool(source.get("content"))
510 # Add metadata from source
511 for key, value in source.items():
512 if key not in result and key not in ["content"]:
513 result[key] = value
515 except Exception as e:
516 safe_msg = self._scrub_error(e)
517 logger.warning(
518 f"Error fetching full content for document {doc_id}: {safe_msg}"
519 )
520 # Keep the preview data if we can't get the full content
522 results.append(result)
524 return results
526 def search_by_query_string(self, query_string: str) -> List[Dict[str, Any]]:
527 """
528 Perform a search using Elasticsearch Query String syntax.
530 Args:
531 query_string: The query in Elasticsearch Query String syntax
533 Returns:
534 List of search results
535 """
536 try:
537 # Build the search query
538 search_query = {
539 "query": {
540 "query_string": {
541 "query": query_string,
542 "fields": self.search_fields,
543 }
544 },
545 "highlight": {
546 "fields": {field: {} for field in self.highlight_fields},
547 "pre_tags": ["<em>"],
548 "post_tags": ["</em>"],
549 },
550 "size": self.max_results,
551 }
553 # Execute the search
554 response = self.client.search(
555 index=self.index_name,
556 body=search_query,
557 )
559 # Process and return the results
560 previews = self._process_es_response(response)
561 return self._get_full_content_cached(previews)
563 except Exception as e:
564 safe_msg = self._scrub_error(e)
565 logger.warning(f"Error in query_string search: {safe_msg}")
566 return []
568 def search_by_dsl(self, query_dsl: Dict[str, Any]) -> List[Dict[str, Any]]:
569 """
570 Perform a search using Elasticsearch DSL (Query Domain Specific Language).
572 Args:
573 query_dsl: The query in Elasticsearch DSL format
575 Returns:
576 List of search results
577 """
578 try:
579 # Execute the search with the provided DSL
580 response = self.client.search(
581 index=self.index_name,
582 body=query_dsl,
583 )
585 # Process and return the results
586 previews = self._process_es_response(response)
587 return self._get_full_content_cached(previews)
589 except Exception as e:
590 safe_msg = self._scrub_error(e)
591 logger.warning(f"Error in DSL search: {safe_msg}")
592 return []
594 def _process_es_response(self, response: Any) -> List[Dict[str, Any]]:
595 """
596 Process Elasticsearch response into preview dictionaries.
598 Args:
599 response: Elasticsearch response dictionary
601 Returns:
602 List of preview dictionaries
603 """
604 hits = response.get("hits", {}).get("hits", [])
606 # Format results as previews
607 previews = []
608 for hit in hits:
609 source = hit.get("_source", {})
610 highlight = hit.get("highlight", {})
612 # Extract highlighted snippets or fall back to original content
613 snippet = ""
614 for field in self.highlight_fields:
615 if highlight.get(field):
616 field_snippets = " ... ".join(highlight[field])
617 snippet += field_snippets + " "
619 # If no highlights, use a portion of the content
620 if not snippet and "content" in source:
621 content = source.get("content", "")
622 snippet = (
623 content[:SNIPPET_LENGTH_SHORT] + "..."
624 if len(content) > SNIPPET_LENGTH_SHORT
625 else content
626 )
628 # Create preview object
629 preview = {
630 "id": hit.get("_id", ""),
631 "title": source.get("title", "Untitled Document"),
632 "link": source.get("url", "")
633 or f"elasticsearch://{self.index_name}/{hit.get('_id', '')}",
634 "snippet": snippet.strip(),
635 "score": hit.get("_score", 0),
636 "_index": hit.get("_index", self.index_name),
637 }
639 previews.append(preview)
641 return previews