Coverage for src/local_deep_research/web_search_engines/engines/search_engine_elasticsearch.py: 97%
147 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1from ...security.secure_logging import logger
2from typing import Any, Dict, List, Optional
4from elasticsearch import Elasticsearch
5from langchain_core.language_models import BaseLLM
7from ...constants import SNIPPET_LENGTH_SHORT
8from ..search_engine_base import BaseSearchEngine, Exposure, Sensitivity
9from ...constants import DEFAULT_SEARCH_TOOL
12class ElasticsearchSearchEngine(BaseSearchEngine):
13 """Elasticsearch search engine implementation with two-phase approach"""
15 is_local = True
16 is_lexical = True
17 needs_llm_relevance_filter = True
18 # Egress (ADR-0007): a local document store — sensitive, contained. The
19 # url_setting fail-up reclassifies exposure to EXPOSING when hosts resolve
20 # public (quadrant 4: usable only by itself, contained inference).
21 egress_sensitivity = Sensitivity.SENSITIVE
22 egress_exposure = Exposure.CONTAINED
23 # secrets to redact from error messages (see BaseSearchEngine._scrub_error)
24 _secret_attrs = ("_api_key", "_password")
25 # url_setting feeds the PDP's fail-up URL override: when the configured
26 # hosts resolve to a PUBLIC endpoint (e.g. Elastic Cloud), the engine is
27 # reclassified public so PRIVATE_ONLY denies it at selection time —
28 # queries would leave the box even though the DATA is "local" in nature.
29 # A localhost ES keeps the static is_local classification above.
30 # cloud_id (which is NOT a host the PDP can classify) is handled
31 # separately in __init__: it is rejected when the effective scope forbids
32 # public egress.
33 url_setting = "search.engine.web.elasticsearch.default_params.hosts"
35 @staticmethod
36 def _cloud_id_forbidden_by_scope(
37 settings_snapshot: Optional[Dict[str, Any]],
38 ) -> bool:
39 """True when the effective egress scope forbids the public Elastic
40 Cloud endpoint a ``cloud_id`` targets.
42 Resolves the scope (including ADAPTIVE) via ``context_from_snapshot``
43 and returns True for PRIVATE_ONLY / STRICT. Fails CLOSED (forbidden)
44 if the policy cannot be evaluated, so a snapshot/policy error cannot
45 open a cloud egress under a private posture. A missing/empty snapshot
46 resolves to the permissive default (BOTH) and is allowed.
47 """
48 try:
49 from ...security.egress.policy import (
50 EgressScope,
51 context_from_snapshot,
52 )
53 from ...config.thread_settings import get_setting_from_snapshot
55 snapshot = settings_snapshot or {}
56 primary = (
57 get_setting_from_snapshot(
58 "search.tool",
59 default=DEFAULT_SEARCH_TOOL,
60 settings_snapshot=snapshot,
61 )
62 or DEFAULT_SEARCH_TOOL
63 )
64 ctx = context_from_snapshot(snapshot, primary)
65 return ctx.scope in (
66 EgressScope.PRIVATE_ONLY,
67 EgressScope.STRICT,
68 )
69 except Exception:
70 logger.bind(policy_audit=True).warning(
71 "elasticsearch cloud_id egress check failed; failing closed",
72 exc_info=True,
73 )
74 return True
76 def __init__(
77 self,
78 hosts: Optional[List[str]] = None,
79 index_name: str = "documents",
80 username: Optional[str] = None,
81 password: Optional[str] = None,
82 api_key: Optional[str] = None,
83 cloud_id: Optional[str] = None,
84 max_results: int = 10,
85 highlight_fields: List[str] = ["content", "title"],
86 search_fields: List[str] = ["content", "title"],
87 filter_query: Optional[Dict[str, Any]] = None,
88 llm: Optional[BaseLLM] = None,
89 max_filtered_results: Optional[int] = None,
90 settings_snapshot: Optional[Dict[str, Any]] = None,
91 ):
92 """
93 Initialize the Elasticsearch search engine.
95 Args:
96 hosts: List of Elasticsearch hosts
97 index_name: Name of the index to search
98 username: Optional username for authentication
99 password: Optional password for authentication
100 api_key: Optional API key for authentication
101 cloud_id: Optional Elastic Cloud ID
102 max_results: Maximum number of search results
103 highlight_fields: Fields to highlight in search results
104 search_fields: Fields to search in
105 filter_query: Optional filter query in Elasticsearch DSL format
106 llm: Language model for relevance filtering
107 max_filtered_results: Maximum number of results to keep after filtering
108 """
109 # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
110 super().__init__(
111 llm=llm,
112 max_filtered_results=max_filtered_results,
113 max_results=max_results,
114 settings_snapshot=settings_snapshot,
115 )
117 self.index_name = index_name
118 self.highlight_fields = self._ensure_list(
119 highlight_fields, default=["content", "title"]
120 )
121 self.search_fields = self._ensure_list(
122 search_fields, default=["content", "title"]
123 )
124 self.filter_query = filter_query or {}
126 # Store credentials for error-message redaction
127 self._api_key = api_key
128 self._password = password
130 # Normalize hosts – may arrive as a JSON-encoded string from settings
131 hosts = self._ensure_list(hosts, default=["http://localhost:9200"])
133 # Initialize the Elasticsearch client
134 es_args: Dict[str, Any] = {}
136 # Basic authentication
137 if username and password:
138 es_args["basic_auth"] = (username, password)
140 # API key authentication
141 if api_key:
142 es_args["api_key"] = api_key
144 # Cloud ID for Elastic Cloud
145 if cloud_id:
146 # Egress policy: a cloud_id always targets a public Elastic Cloud
147 # endpoint (*.cloud.es.io), but the url_setting reclassification
148 # only inspects `hosts`. A cloud_id-only config would otherwise
149 # keep the engine's static is_local=True and slip past
150 # evaluate_engine, then connect at self.client.info() below. Reject
151 # it when the effective scope forbids public egress (fail closed).
152 if self._cloud_id_forbidden_by_scope(settings_snapshot):
153 from ...security.egress.policy import (
154 Decision,
155 PolicyDeniedError,
156 )
158 logger.bind(policy_audit=True).warning(
159 "refusing Elasticsearch cloud_id under private egress scope"
160 )
161 raise PolicyDeniedError(
162 Decision(False, "elasticsearch_cloud_id_public_egress"),
163 target="search_engine:elasticsearch",
164 )
165 es_args["cloud_id"] = cloud_id
167 # Connect to Elasticsearch
168 self.client = Elasticsearch(hosts, **es_args)
170 # Verify connection
171 try:
172 info = self.client.info()
173 logger.info(
174 f"Connected to Elasticsearch cluster: {info.get('cluster_name')}"
175 )
176 logger.info(
177 f"Elasticsearch version: {info.get('version', {}).get('number')}"
178 )
179 except Exception as e:
180 safe_msg = self._scrub_error(e)
181 logger.warning(f"Failed to connect to Elasticsearch: {safe_msg}")
182 raise ConnectionError(
183 f"Could not connect to Elasticsearch: {safe_msg}"
184 ) from None
186 def close(self) -> None:
187 """Close the Elasticsearch client and its connection pool."""
188 from ...utilities.resource_utils import safe_close
190 safe_close(self.client, "Elasticsearch client")
191 super().close()
193 def _get_previews(self, query: str) -> List[Dict[str, Any]]:
194 """
195 Get preview information for Elasticsearch documents.
197 Args:
198 query: The search query
200 Returns:
201 List of preview dictionaries
202 """
203 logger.info(
204 f"Getting document previews from Elasticsearch with query: {query}"
205 )
207 try:
208 # Build the search query
209 search_query = {
210 "query": {
211 "multi_match": {
212 "query": query,
213 "fields": self.search_fields,
214 "type": "best_fields",
215 "tie_breaker": 0.3,
216 }
217 },
218 "highlight": {
219 "fields": {field: {} for field in self.highlight_fields},
220 "pre_tags": ["<em>"],
221 "post_tags": ["</em>"],
222 },
223 "size": self.max_results,
224 }
226 # Add filter if provided
227 if self.filter_query:
228 search_query["query"] = {
229 "bool": {
230 "must": search_query["query"],
231 "filter": self.filter_query,
232 }
233 }
235 # Execute the search
236 response = self.client.search(
237 index=self.index_name,
238 body=search_query,
239 )
241 # Process the search results
242 hits = response.get("hits", {}).get("hits", [])
244 # Format results as previews with basic information
245 previews = []
246 for hit in hits:
247 source = hit.get("_source", {})
248 highlight = hit.get("highlight", {})
250 # Extract highlighted snippets or fall back to original content
251 snippet = ""
252 for field in self.highlight_fields:
253 if highlight.get(field):
254 # Join all highlights for this field
255 field_snippets = " ... ".join(highlight[field])
256 snippet += field_snippets + " "
258 # If no highlights, use a portion of the content
259 if not snippet and "content" in source:
260 content = source.get("content", "")
261 snippet = (
262 content[:SNIPPET_LENGTH_SHORT] + "..."
263 if len(content) > SNIPPET_LENGTH_SHORT
264 else content
265 )
267 # Create preview object
268 preview = {
269 "id": hit.get("_id", ""),
270 "title": source.get("title", "Untitled Document"),
271 "link": source.get("url", "")
272 or f"elasticsearch://{self.index_name}/{hit.get('_id', '')}",
273 "snippet": snippet.strip(),
274 "score": hit.get("_score", 0),
275 "_index": hit.get("_index", self.index_name),
276 }
278 previews.append(preview)
280 logger.info(
281 f"Found {len(previews)} preview results from Elasticsearch"
282 )
283 return previews
285 except Exception as e:
286 safe_msg = self._scrub_error(e)
287 logger.warning(f"Error getting Elasticsearch previews: {safe_msg}")
288 return []
290 def _get_full_content(
291 self, relevant_items: List[Dict[str, Any]]
292 ) -> List[Dict[str, Any]]:
293 """
294 Get full content for the relevant Elasticsearch documents.
296 Args:
297 relevant_items: List of relevant preview dictionaries
299 Returns:
300 List of result dictionaries with full content
301 """
302 logger.info("Getting full content for relevant Elasticsearch documents")
304 results = []
305 for item in relevant_items:
306 # Start with the preview data
307 result = item.copy()
309 # Get the document ID
310 doc_id = item.get("id")
311 if not doc_id:
312 # Skip items without ID
313 logger.warning(f"Skipping item without ID: {item}")
314 results.append(result)
315 continue
317 try:
318 # Fetch the full document
319 doc_response = self.client.get(
320 index=self.index_name,
321 id=doc_id,
322 )
324 # Get the source document
325 source = doc_response.get("_source", {})
327 # Add full content to the result
328 result["content"] = source.get(
329 "content", result.get("snippet", "")
330 )
331 result["full_content"] = source.get("content", "")
333 # Add metadata from source
334 for key, value in source.items():
335 if key not in result and key not in ["content"]:
336 result[key] = value
338 except Exception as e:
339 safe_msg = self._scrub_error(e)
340 logger.warning(
341 f"Error fetching full content for document {doc_id}: {safe_msg}"
342 )
343 # Keep the preview data if we can't get the full content
345 results.append(result)
347 return results
349 def search_by_query_string(self, query_string: str) -> List[Dict[str, Any]]:
350 """
351 Perform a search using Elasticsearch Query String syntax.
353 Args:
354 query_string: The query in Elasticsearch Query String syntax
356 Returns:
357 List of search results
358 """
359 try:
360 # Build the search query
361 search_query = {
362 "query": {
363 "query_string": {
364 "query": query_string,
365 "fields": self.search_fields,
366 }
367 },
368 "highlight": {
369 "fields": {field: {} for field in self.highlight_fields},
370 "pre_tags": ["<em>"],
371 "post_tags": ["</em>"],
372 },
373 "size": self.max_results,
374 }
376 # Execute the search
377 response = self.client.search(
378 index=self.index_name,
379 body=search_query,
380 )
382 # Process and return the results
383 previews = self._process_es_response(response)
384 return self._get_full_content(previews)
386 except Exception as e:
387 safe_msg = self._scrub_error(e)
388 logger.warning(f"Error in query_string search: {safe_msg}")
389 return []
391 def search_by_dsl(self, query_dsl: Dict[str, Any]) -> List[Dict[str, Any]]:
392 """
393 Perform a search using Elasticsearch DSL (Query Domain Specific Language).
395 Args:
396 query_dsl: The query in Elasticsearch DSL format
398 Returns:
399 List of search results
400 """
401 try:
402 # Execute the search with the provided DSL
403 response = self.client.search(
404 index=self.index_name,
405 body=query_dsl,
406 )
408 # Process and return the results
409 previews = self._process_es_response(response)
410 return self._get_full_content(previews)
412 except Exception as e:
413 safe_msg = self._scrub_error(e)
414 logger.warning(f"Error in DSL search: {safe_msg}")
415 return []
417 def _process_es_response(self, response: Any) -> List[Dict[str, Any]]:
418 """
419 Process Elasticsearch response into preview dictionaries.
421 Args:
422 response: Elasticsearch response dictionary
424 Returns:
425 List of preview dictionaries
426 """
427 hits = response.get("hits", {}).get("hits", [])
429 # Format results as previews
430 previews = []
431 for hit in hits:
432 source = hit.get("_source", {})
433 highlight = hit.get("highlight", {})
435 # Extract highlighted snippets or fall back to original content
436 snippet = ""
437 for field in self.highlight_fields:
438 if highlight.get(field):
439 field_snippets = " ... ".join(highlight[field])
440 snippet += field_snippets + " "
442 # If no highlights, use a portion of the content
443 if not snippet and "content" in source:
444 content = source.get("content", "")
445 snippet = (
446 content[:SNIPPET_LENGTH_SHORT] + "..."
447 if len(content) > SNIPPET_LENGTH_SHORT
448 else content
449 )
451 # Create preview object
452 preview = {
453 "id": hit.get("_id", ""),
454 "title": source.get("title", "Untitled Document"),
455 "link": source.get("url", "")
456 or f"elasticsearch://{self.index_name}/{hit.get('_id', '')}",
457 "snippet": snippet.strip(),
458 "score": hit.get("_score", 0),
459 "_index": hit.get("_index", self.index_name),
460 }
462 previews.append(preview)
464 return previews