Coverage for src/local_deep_research/utilities/es_utils.py: 97%
125 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""
2Elasticsearch utilities for indexing and managing documents.
3"""
5from loguru import logger
6from pathlib import Path
7from typing import Any, Dict, List, Optional
9from elasticsearch import Elasticsearch
10from elasticsearch.helpers import bulk
13class ElasticsearchManager:
14 """
15 Utility class for managing Elasticsearch indices and documents.
17 This class provides methods for creating indices, indexing documents,
18 and performing other Elasticsearch management tasks.
19 """
21 def __init__(
22 self,
23 hosts: Optional[List[str]] = None,
24 username: Optional[str] = None,
25 password: Optional[str] = None,
26 api_key: Optional[str] = None,
27 cloud_id: Optional[str] = None,
28 ):
29 """
30 Initialize the Elasticsearch manager.
32 Args:
33 hosts: List of Elasticsearch hosts
34 username: Optional username for authentication
35 password: Optional password for authentication
36 api_key: Optional API key for authentication
37 cloud_id: Optional Elastic Cloud ID
38 """
39 hosts = hosts or ["http://localhost:9200"]
40 # Initialize the Elasticsearch client
41 es_args: Dict[str, Any] = {}
43 # Basic authentication
44 if username and password:
45 es_args["basic_auth"] = (username, password)
47 # API key authentication
48 if api_key:
49 es_args["api_key"] = api_key
51 # Cloud ID for Elastic Cloud
52 if cloud_id:
53 es_args["cloud_id"] = cloud_id
55 # Connect to Elasticsearch
56 self.client = Elasticsearch(hosts, **es_args)
58 # Verify connection
59 try:
60 info = self.client.info()
61 logger.info(
62 f"Connected to Elasticsearch cluster: {info.get('cluster_name')}"
63 )
64 logger.info(
65 f"Elasticsearch version: {info.get('version', {}).get('number')}"
66 )
67 except Exception as e:
68 logger.exception("Failed to connect to Elasticsearch")
69 raise ConnectionError(f"Could not connect to Elasticsearch: {e!s}")
71 def close(self) -> None:
72 """Close the Elasticsearch client and its connection pool."""
73 from .resource_utils import safe_close
75 safe_close(self.client, "Elasticsearch client")
77 def create_index(
78 self,
79 index_name: str,
80 mappings: Optional[Dict[str, Any]] = None,
81 settings: Optional[Dict[str, Any]] = None,
82 ) -> bool:
83 """
84 Create an Elasticsearch index with optional mappings and settings.
86 Args:
87 index_name: Name of the index to create
88 mappings: Optional mappings for the index fields
89 settings: Optional settings for the index
91 Returns:
92 bool: True if successful, False otherwise
93 """
94 try:
95 # Check if index already exists
96 if self.client.indices.exists(index=index_name):
97 logger.warning(
98 f"Index '{index_name}' already exists - skipping creation"
99 )
100 return True
102 # Default mappings for better text search if none provided
103 if mappings is None:
104 mappings = {
105 "properties": {
106 "title": {
107 "type": "text",
108 "analyzer": "standard",
109 "fields": {
110 "keyword": {
111 "type": "keyword",
112 "ignore_above": 256,
113 }
114 },
115 },
116 "content": {"type": "text", "analyzer": "standard"},
117 "url": {"type": "keyword"},
118 "source": {"type": "keyword"},
119 "timestamp": {"type": "date"},
120 "metadata": {"type": "object", "enabled": True},
121 }
122 }
124 # Default settings if none provided
125 if settings is None:
126 settings = {
127 "number_of_shards": 1,
128 "number_of_replicas": 0,
129 "analysis": {
130 "analyzer": {"standard": {"type": "standard"}}
131 },
132 }
134 # Create the index with mappings and settings
135 create_response = self.client.indices.create(
136 index=index_name,
137 mappings=mappings,
138 settings=settings,
139 )
141 logger.info(f"Created index '{index_name}': {create_response}")
142 return True
144 except Exception:
145 logger.exception(f"Error creating index '{index_name}'")
146 return False
148 def delete_index(self, index_name: str) -> bool:
149 """
150 Delete an Elasticsearch index.
152 Args:
153 index_name: Name of the index to delete
155 Returns:
156 bool: True if successful, False otherwise
157 """
158 try:
159 # Check if index exists
160 if not self.client.indices.exists(index=index_name):
161 logger.warning(
162 f"Index '{index_name}' does not exist - skipping deletion"
163 )
164 return True
166 # Delete the index
167 delete_response = self.client.indices.delete(index=index_name)
168 logger.info(f"Deleted index '{index_name}': {delete_response}")
169 return True
171 except Exception:
172 logger.exception(f"Error deleting index '{index_name}'")
173 return False
175 def index_document(
176 self,
177 index_name: str,
178 document: Dict[str, Any],
179 document_id: Optional[str] = None,
180 refresh: bool = False,
181 ) -> Optional[str]:
182 """
183 Index a single document in Elasticsearch.
185 Args:
186 index_name: Name of the index to add the document to
187 document: The document to index
188 document_id: Optional document ID (will be generated if not provided)
189 refresh: Whether to refresh the index after indexing
191 Returns:
192 str: Document ID if successful, None otherwise
193 """
194 try:
195 # Index the document
196 response = self.client.index(
197 index=index_name,
198 document=document,
199 id=document_id,
200 refresh=refresh,
201 )
203 logger.info(
204 f"Indexed document in '{index_name}' with ID: {response['_id']}"
205 )
206 return response["_id"]
208 except Exception:
209 logger.exception(f"Error indexing document in '{index_name}'")
210 return None
212 def bulk_index_documents(
213 self,
214 index_name: str,
215 documents: List[Dict[str, Any]],
216 id_field: Optional[str] = None,
217 refresh: bool = False,
218 ) -> int:
219 """
220 Bulk index multiple documents in Elasticsearch.
222 Args:
223 index_name: Name of the index to add the documents to
224 documents: List of documents to index
225 id_field: Optional field in the documents to use as the document ID
226 refresh: Whether to refresh the index after indexing
228 Returns:
229 int: Number of successfully indexed documents
230 """
231 try:
232 # Prepare the bulk actions
233 actions = []
234 for doc in documents:
235 action = {
236 "_index": index_name,
237 "_source": doc,
238 }
240 # Use the specified field as the document ID if provided
241 if id_field and id_field in doc:
242 action["_id"] = doc[id_field]
244 actions.append(action)
246 # Execute the bulk indexing
247 success, failed = bulk(
248 self.client,
249 actions,
250 refresh=refresh,
251 stats_only=True,
252 )
254 logger.info(
255 f"Bulk indexed {success} documents in '{index_name}', failed: {failed}"
256 )
257 return success
259 except Exception:
260 logger.exception(f"Error bulk indexing documents in '{index_name}'")
261 return 0
263 def index_file(
264 self,
265 index_name: str,
266 file_path: str,
267 content_field: str = "content",
268 title_field: Optional[str] = "title",
269 extract_metadata: bool = True,
270 refresh: bool = False,
271 ) -> Optional[str]:
272 """
273 Index a file in Elasticsearch, extracting text content and metadata.
275 Args:
276 index_name: Name of the index to add the document to
277 file_path: Path to the file to index
278 content_field: Field name to store the file content
279 title_field: Field name to store the file title (filename if not specified)
280 extract_metadata: Whether to extract file metadata
281 refresh: Whether to refresh the index after indexing
283 Returns:
284 str: Document ID if successful, None otherwise
285 """
286 try:
287 from unstructured.partition.auto import partition
289 # This is the same local partitioning + single-document join used
290 # by the deprecated LangChain loader, without constructing unused
291 # sync and async Unstructured API clients for every indexed file.
292 elements = partition(filename=file_path)
293 content = "\n\n".join(str(element) for element in elements)
295 # Get the filename for the title
296 filename = Path(file_path).name
297 title = filename
299 # Prepare the document
300 document = {
301 content_field: content,
302 }
304 # Add title if requested
305 if title_field:
306 document[title_field] = title
308 # Add metadata if requested
309 if extract_metadata: 309 ↛ 319line 309 didn't jump to line 319 because the condition on line 309 was always true
310 # Legacy UnstructuredFileLoader(mode="single") exposed only
311 # file-level metadata. Preserve that behavior for empty files
312 # and multi-element documents alike.
313 document["metadata"] = {"source": file_path}
314 document["source"] = file_path
315 document["file_extension"] = Path(filename).suffix.lstrip(".")
316 document["filename"] = filename
318 # Index the document
319 return self.index_document(index_name, document, refresh=refresh)
321 except ImportError:
322 logger.error(
323 "Unstructured partitioner not available. Please install the "
324 "'unstructured' package."
325 )
326 return None
327 except Exception:
328 logger.exception(f"Error indexing file '{file_path}'")
329 return None
331 def index_directory(
332 self,
333 index_name: str,
334 directory_path: str,
335 file_patterns: Optional[List[str]] = None,
336 content_field: str = "content",
337 title_field: str = "title",
338 extract_metadata: bool = True,
339 refresh: bool = False,
340 ) -> int:
341 """
342 Index all matching files in a directory in Elasticsearch.
344 Args:
345 index_name: Name of the index to add the documents to
346 directory_path: Path to the directory containing files to index
347 file_patterns: List of file patterns to match (glob patterns)
348 content_field: Field name to store the file content
349 title_field: Field name to store the file title
350 extract_metadata: Whether to extract file metadata
351 refresh: Whether to refresh the index after indexing
353 Returns:
354 int: Number of successfully indexed files
355 """
356 file_patterns = file_patterns or ["*.txt", "*.pdf", "*.docx", "*.md"]
357 try:
358 # Find all matching files
359 all_files = []
360 directory = Path(directory_path)
361 for pattern in file_patterns:
362 matching_files = list(directory.glob(pattern))
363 all_files.extend([str(f) for f in matching_files])
365 logger.info(
366 f"Found {len(all_files)} files matching patterns {file_patterns} in {directory_path}"
367 )
369 # Index each file
370 successful_count = 0
371 for file_path in all_files:
372 logger.info(f"Indexing file: {file_path}")
373 doc_id = self.index_file(
374 index_name=index_name,
375 file_path=file_path,
376 content_field=content_field,
377 title_field=title_field,
378 extract_metadata=extract_metadata,
379 refresh=refresh,
380 )
382 if doc_id:
383 successful_count += 1
385 logger.info(
386 f"Successfully indexed {successful_count} files out of {len(all_files)}"
387 )
388 return successful_count
390 except Exception:
391 logger.exception(f"Error indexing directory '{directory_path}'")
392 return 0
394 def search(
395 self,
396 index_name: str,
397 query: str,
398 fields: List[str] = ["content", "title"],
399 size: int = 10,
400 highlight: bool = True,
401 ) -> Dict[str, Any]:
402 """
403 Search for documents in Elasticsearch.
405 Args:
406 index_name: Name of the index to search
407 query: Search query
408 fields: Fields to search in
409 size: Maximum number of results to return
410 highlight: Whether to include highlighted excerpts in results
412 Returns:
413 Dict: Elasticsearch search response
414 """
415 try:
416 search_query = {
417 "query": {
418 "multi_match": {
419 "query": query,
420 "fields": fields,
421 "type": "best_fields",
422 "tie_breaker": 0.3,
423 }
424 },
425 "size": size,
426 }
428 # Add highlighting if requested
429 if highlight:
430 search_query["highlight"] = {
431 "fields": {field: {} for field in fields},
432 "pre_tags": ["<em>"],
433 "post_tags": ["</em>"],
434 }
436 # Execute the search
437 return self.client.search(
438 index=index_name,
439 body=search_query,
440 )
442 except Exception:
443 logger.exception(f"Error searching index '{index_name}'")
444 return {"error": "Search failed"}