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