Refactor Elasticsearch integration to use dedicated client

- Replaced direct HTTP requests with an ElasticsearchClient for handling interactions with Elasticsearch in ElasticSearchDirectHttpHandler and ElasticsearchIndexManager.
- Updated configuration handling to encapsulate Elasticsearch settings within a single config object.
- Simplified index management and document indexing processes, improving code maintainability and readability.
- Adjusted tests to utilize the new client for searching and verifying log entries, enhancing test reliability.

Author: Dylan Huang

eval_protocol/logging/elasticsearch_client.py

@@ -0,0 +1,301 @@
+"""
+Centralized Elasticsearch client for all Elasticsearch API operations.
+
+This module provides a unified interface for all Elasticsearch operations
+used throughout the codebase, including index management, document operations,
+and search functionality.
+"""
+
+import json
+import requests
+from typing import Any, Dict, List, Optional, Union
+from urllib.parse import urlparse
+from dataclasses import dataclass
+
+
+@dataclass
+class ElasticsearchConfig:
+    """Configuration for Elasticsearch client."""
+
+    url: str
+    api_key: str
+    index_name: str
+    verify_ssl: bool = True
+
+    def __post_init__(self):
+        """Parse URL to determine SSL verification setting."""
+        parsed_url = urlparse(self.url)
+        self.verify_ssl = parsed_url.scheme == "https"
+
+
+class ElasticsearchClient:
+    """Centralized client for all Elasticsearch operations."""
+
+    def __init__(self, config: ElasticsearchConfig):
+        """Initialize the Elasticsearch client.
+
+        Args:
+            config: Elasticsearch configuration
+        """
+        self.config = config
+        self.base_url = config.url.rstrip("/")
+        self.index_url = f"{self.base_url}/{config.index_name}"
+        self._headers = {"Content-Type": "application/json", "Authorization": f"ApiKey {config.api_key}"}
+
+    def _make_request(
+        self,
+        method: str,
+        url: str,
+        json_data: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        timeout: int = 30,
+    ) -> requests.Response:
+        """Make an HTTP request to Elasticsearch.
+
+        Args:
+            method: HTTP method (GET, POST, PUT, DELETE, HEAD)
+            url: Full URL for the request
+            json_data: JSON data to send in request body
+            params: Query parameters
+            timeout: Request timeout in seconds
+
+        Returns:
+            requests.Response object
+
+        Raises:
+            requests.RequestException: If the request fails
+        """
+        return requests.request(
+            method=method,
+            url=url,
+            headers=self._headers,
+            json=json_data,
+            params=params,
+            verify=self.config.verify_ssl,
+            timeout=timeout,
+        )
+
+    # Index Management Operations
+
+    def create_index(self, mapping: Dict[str, Any]) -> bool:
+        """Create an index with the specified mapping.
+
+        Args:
+            mapping: Index mapping configuration
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        try:
+            response = self._make_request("PUT", self.index_url, json_data=mapping)
+            return response.status_code in [200, 201]
+        except Exception:
+            return False
+
+    def index_exists(self) -> bool:
+        """Check if the index exists.
+
+        Returns:
+            bool: True if index exists, False otherwise
+        """
+        try:
+            response = self._make_request("HEAD", self.index_url)
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    def delete_index(self) -> bool:
+        """Delete the index.
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        try:
+            response = self._make_request("DELETE", self.index_url)
+            return response.status_code in [200, 404]  # 404 means index doesn't exist
+        except Exception:
+            return False
+
+    def get_mapping(self) -> Optional[Dict[str, Any]]:
+        """Get the index mapping.
+
+        Returns:
+            Dict containing mapping data, or None if failed
+        """
+        try:
+            response = self._make_request("GET", f"{self.index_url}/_mapping")
+            if response.status_code == 200:
+                return response.json()
+            return None
+        except Exception:
+            return None
+
+    def get_index_stats(self) -> Optional[Dict[str, Any]]:
+        """Get index statistics.
+
+        Returns:
+            Dict containing index statistics, or None if failed
+        """
+        try:
+            response = self._make_request("GET", f"{self.index_url}/_stats")
+            if response.status_code == 200:
+                return response.json()
+            return None
+        except Exception:
+            return None
+
+    # Document Operations
+
+    def index_document(self, document: Dict[str, Any], doc_id: Optional[str] = None) -> bool:
+        """Index a document.
+
+        Args:
+            document: Document to index
+            doc_id: Optional document ID
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        try:
+            if doc_id:
+                url = f"{self.index_url}/_doc/{doc_id}"
+            else:
+                url = f"{self.index_url}/_doc"
+
+            response = self._make_request("POST", url, json_data=document)
+            return response.status_code in [200, 201]
+        except Exception:
+            return False
+
+    def bulk_index_documents(self, documents: List[Dict[str, Any]]) -> bool:
+        """Bulk index multiple documents.
+
+        Args:
+            documents: List of documents to index
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        try:
+            # Prepare bulk request body
+            bulk_body = []
+            for doc in documents:
+                bulk_body.append({"index": {}})
+                bulk_body.append(doc)
+
+            response = self._make_request("POST", f"{self.index_url}/_bulk", json_data=bulk_body)
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    # Search Operations
+
+    def search(
+        self, query: Dict[str, Any], size: int = 10, from_: int = 0, sort: Optional[List[Dict[str, Any]]] = None
+    ) -> Optional[Dict[str, Any]]:
+        """Search documents in the index.
+
+        Args:
+            query: Elasticsearch query
+            size: Number of results to return
+            from_: Starting offset
+            sort: Sort specification
+
+        Returns:
+            Dict containing search results, or None if failed
+        """
+        try:
+            search_body = {"query": query, "size": size, "from": from_}
+
+            if sort:
+                search_body["sort"] = sort
+
+            response = self._make_request("POST", f"{self.index_url}/_search", json_data=search_body)
+
+            if response.status_code == 200:
+                return response.json()
+            return None
+        except Exception:
+            return None
+
+    def search_by_term(self, field: str, value: Any, size: int = 10) -> Optional[Dict[str, Any]]:
+        """Search documents by exact term match.
+
+        Args:
+            field: Field name to search
+            value: Value to match
+            size: Number of results to return
+
+        Returns:
+            Dict containing search results, or None if failed
+        """
+        query = {"term": {field: value}}
+        return self.search(query, size=size)
+
+    def search_by_match(self, field: str, value: str, size: int = 10) -> Optional[Dict[str, Any]]:
+        """Search documents by text match.
+
+        Args:
+            field: Field name to search
+            value: Text to match
+            size: Number of results to return
+
+        Returns:
+            Dict containing search results, or None if failed
+        """
+        query = {"match": {field: value}}
+        return self.search(query, size=size)
+
+    def search_by_match_phrase_prefix(self, field: str, value: str, size: int = 10) -> Optional[Dict[str, Any]]:
+        """Search documents by phrase prefix match.
+
+        Args:
+            field: Field name to search
+            value: Phrase prefix to match
+            size: Number of results to return
+
+        Returns:
+            Dict containing search results, or None if failed
+        """
+        query = {"match_phrase_prefix": {field: value}}
+        return self.search(query, size=size)
+
+    def search_all(self, size: int = 10) -> Optional[Dict[str, Any]]:
+        """Search all documents in the index.
+
+        Args:
+            size: Number of results to return
+
+        Returns:
+            Dict containing search results, or None if failed
+        """
+        query = {"match_all": {}}
+        return self.search(query, size=size)
+
+    # Health and Status Operations
+
+    def health_check(self) -> bool:
+        """Check if Elasticsearch is healthy.
+
+        Returns:
+            bool: True if healthy, False otherwise
+        """
+        try:
+            response = self._make_request("GET", f"{self.base_url}/_cluster/health")
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    def get_cluster_info(self) -> Optional[Dict[str, Any]]:
+        """Get cluster information.
+
+        Returns:
+            Dict containing cluster info, or None if failed
+        """
+        try:
+            response = self._make_request("GET", f"{self.base_url}/_cluster/health")
+            if response.status_code == 200:
+                return response.json()
+            return None
+        except Exception:
+            return None

eval_protocol/logging/elasticsearch_direct_http_handler.py

@@ -6,26 +6,23 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import Optional, Tuple, Any, Dict
 from datetime import datetime
-from urllib.parse import urlparse
-import requests
 
 from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig
+from .elasticsearch_client import ElasticsearchClient, ElasticsearchConfig as ESConfig
 
 
 class ElasticSearchDirectHttpHandler(logging.Handler):
     def __init__(self, elasticsearch_config: ElasticSearchConfig) -> None:
         super().__init__()
-        self.base_url: str = elasticsearch_config.url.rstrip("/")
-        self.index_name: str = elasticsearch_config.index_name
-        self.api_key: str = elasticsearch_config.api_key
-        self.url: str = f"{self.base_url}/{self.index_name}/_doc"
+        self.config = ESConfig(
+            url=elasticsearch_config.url,
+            api_key=elasticsearch_config.api_key,
+            index_name=elasticsearch_config.index_name,
+        )
+        self.client = ElasticsearchClient(self.config)
         self.formatter: logging.Formatter = logging.Formatter()
         self._executor = None
 
-        # Parse URL to determine if we should verify SSL
-        parsed_url = urlparse(elasticsearch_config.url)
-        self.verify_ssl = parsed_url.scheme == "https"
-
     def emit(self, record: logging.LogRecord) -> None:
         """Emit a log record by scheduling it for async transmission."""
         try:
@@ -104,13 +101,9 @@ def _schedule_async_send(self, data: Dict[str, Any], record: logging.LogRecord)
     def _send_to_elasticsearch(self, data: Dict[str, Any], record: logging.LogRecord) -> None:
         """Send data to Elasticsearch (runs in thread pool)."""
         try:
-            response: requests.Response = requests.post(
-                self.url,
-                headers={"Content-Type": "application/json", "Authorization": f"ApiKey {self.api_key}"},
-                data=json.dumps(data),
-                verify=self.verify_ssl,  # If using HTTPS, verify SSL certificate
-            )
-            response.raise_for_status()  # Raise an exception for HTTP errors
+            success = self.client.index_document(data)
+            if not success:
+                raise Exception("Failed to index document to Elasticsearch")
         except Exception as e:
             # Re-raise to be handled by the callback
             raise e