Refactor Elasticsearch configuration handling

- Updated Elasticsearch integration to use ElasticSearchConfig from the types module, improving consistency across the codebase.
- Removed the local ElasticsearchConfig dataclass in favor of the new model, streamlining configuration management.
- Adjusted related classes and tests to utilize the updated configuration structure, enhancing maintainability and readability.

Author: Dylan Huang

eval_protocol/logging/elasticsearch_client.py

@@ -10,28 +10,13 @@
 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"
+from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig
 
 
 class ElasticsearchClient:
     """Centralized client for all Elasticsearch operations."""
 
-    def __init__(self, config: ElasticsearchConfig):
+    def __init__(self, config: ElasticSearchConfig):
         """Initialize the Elasticsearch client.
 
         Args:

eval_protocol/logging/elasticsearch_direct_http_handler.py

@@ -8,13 +8,13 @@
 from datetime import datetime
 
 from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig
-from .elasticsearch_client import ElasticsearchClient, ElasticsearchConfig as ESConfig
+from .elasticsearch_client import ElasticsearchClient
 
 
 class ElasticsearchDirectHttpHandler(logging.Handler):
     def __init__(self, elasticsearch_config: ElasticSearchConfig) -> None:
         super().__init__()
-        self.config = ESConfig(
+        self.config = ElasticSearchConfig(
             url=elasticsearch_config.url,
             api_key=elasticsearch_config.api_key,
             index_name=elasticsearch_config.index_name,

eval_protocol/logging/elasticsearch_index_manager.py

@@ -1,5 +1,6 @@
 from typing import Dict, Any, Optional
-from .elasticsearch_client import ElasticsearchClient, ElasticsearchConfig
+from .elasticsearch_client import ElasticsearchClient
+from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig
 
 
 class ElasticsearchIndexManager:
@@ -13,7 +14,7 @@ def __init__(self, base_url: str, index_name: str, api_key: str) -> None:
             index_name: Name of the index to manage
             api_key: API key for authentication
         """
-        self.config = ElasticsearchConfig(url=base_url, api_key=api_key, index_name=index_name)
+        self.config = ElasticSearchConfig(url=base_url, api_key=api_key, index_name=index_name)
         self.client = ElasticsearchClient(self.config)
         self._mapping_created: bool = False
 

eval_protocol/pytest/remote_rollout_processor.py

@@ -4,6 +4,7 @@
 
 import requests
 
+from eval_protocol.logging.elasticsearch_client import ElasticsearchClient
 from eval_protocol.models import EvaluationRow, Status
 from eval_protocol.data_loader.dynamic_data_loader import DynamicDataLoader
 from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig, InitRequest, RolloutMetadata
@@ -182,6 +183,10 @@ def _get_status() -> Dict[str, Any]:
                 r.raise_for_status()
                 return r.json()
 
+            elasticsearch_client = (
+                ElasticsearchClient(self._elastic_search_config) if self._elastic_search_config else None
+            )
+
             while time.time() < deadline:
                 try:
                     status = await asyncio.to_thread(_get_status)

eval_protocol/types/remote_rollout_processor.py

@@ -4,6 +4,7 @@
 
 from typing import Any, Dict, List, Optional
 from pydantic import BaseModel, Field
+from urllib.parse import urlparse
 from eval_protocol.models import Message, Status
 
 
@@ -16,6 +17,12 @@ class ElasticSearchConfig(BaseModel):
     api_key: str
     index_name: str
 
+    @property
+    def verify_ssl(self) -> bool:
+        """Infer verify_ssl from URL scheme."""
+        parsed_url = urlparse(self.url)
+        return parsed_url.scheme == "https"
+
 
 class RolloutMetadata(BaseModel):
     """Metadata for rollout execution."""

tests/logging/test_elasticsearch_direct_http_handler.py

@@ -4,7 +4,7 @@
 import pytest
 
 from eval_protocol.logging.elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
-from eval_protocol.logging.elasticsearch_client import ElasticsearchClient, ElasticsearchConfig as ESConfig
+from eval_protocol.logging.elasticsearch_client import ElasticsearchClient
 from eval_protocol.pytest.elasticsearch_setup import ElasticsearchSetup
 from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig
 
@@ -54,10 +54,8 @@ def elasticsearch_handler(elasticsearch_config: ElasticSearchConfig, rollout_id:
 @pytest.fixture
 def elasticsearch_client(elasticsearch_config: ElasticSearchConfig):
     """Create an Elasticsearch client for testing."""
-    config = ESConfig(
-        url=elasticsearch_config.url, api_key=elasticsearch_config.api_key, index_name=elasticsearch_config.index_name
-    )
-    return ElasticsearchClient(config)
+    # Create a new config instance for the client
+    return ElasticsearchClient(elasticsearch_config)
 
 
 @pytest.fixture