Refactor Elasticsearch configuration naming for consistency

- Updated all instances of ElasticSearchConfig to ElasticsearchConfig across the codebase to align with official naming conventions.
- Ensured that related classes and methods reflect this change, enhancing code clarity and maintainability.

Author: Dylan Huang

eval_protocol/logging/elasticsearch_client.py

@@ -10,13 +10,13 @@
 import requests
 from typing import Any, Dict, List, Optional, Union
 from urllib.parse import urlparse
-from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig
+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

@@ -7,14 +7,14 @@
 from typing import Optional, Tuple, Any, Dict
 from datetime import datetime
 
-from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig
+from eval_protocol.types.remote_rollout_processor import ElasticsearchConfig
 from .elasticsearch_client import ElasticsearchClient
 
 
 class ElasticsearchDirectHttpHandler(logging.Handler):
-    def __init__(self, elasticsearch_config: ElasticSearchConfig) -> None:
+    def __init__(self, elasticsearch_config: ElasticsearchConfig) -> None:
         super().__init__()
-        self.config = ElasticSearchConfig(
+        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,6 +1,6 @@
 from typing import Dict, Any, Optional
 from .elasticsearch_client import ElasticsearchClient
-from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig
+from eval_protocol.types.remote_rollout_processor import ElasticsearchConfig
 
 
 class ElasticsearchIndexManager:
@@ -14,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/elasticsearch_setup.py

@@ -6,7 +6,7 @@
 
 from dotenv import load_dotenv
 from eval_protocol.directory_utils import find_eval_protocol_dir
-from eval_protocol.types.remote_rollout_processor import ElasticSearchConfig
+from eval_protocol.types.remote_rollout_processor import ElasticsearchConfig
 from eval_protocol.logging.elasticsearch_index_manager import ElasticsearchIndexManager
 
 logger = logging.getLogger(__name__)
@@ -24,15 +24,15 @@ class ElasticsearchSetup:
     def __init__(self):
         self.eval_protocol_dir = find_eval_protocol_dir()
 
-    def setup_elasticsearch(self, index_name: str = "default-logs") -> ElasticSearchConfig:
+    def setup_elasticsearch(self, index_name: str = "default-logs") -> ElasticsearchConfig:
         """
         Set up Elasticsearch, handling both local and remote scenarios.
 
         Args:
             index_name: Name of the Elasticsearch index to use for logging
 
         Returns:
-            ElasticSearchConfig for the running instance with the specified index name.
+            ElasticsearchConfig for the running instance with the specified index name.
         """
         elastic_start_local_dir = os.path.join(self.eval_protocol_dir, "elastic-start-local")
         env_file_path = os.path.join(elastic_start_local_dir, ".env")
@@ -48,11 +48,11 @@ def setup_elasticsearch(self, index_name: str = "default-logs") -> ElasticSearch
         self.create_logging_index(index_name)
 
         # Return config with the specified index name
-        return ElasticSearchConfig(url=config.url, api_key=config.api_key, index_name=index_name)
+        return ElasticsearchConfig(url=config.url, api_key=config.api_key, index_name=index_name)
 
     def _setup_existing_docker_elasticsearch(
         self, elastic_start_local_dir: str, env_file_path: str
-    ) -> ElasticSearchConfig:
+    ) -> ElasticsearchConfig:
         """Set up Elasticsearch using existing Docker start.sh script."""
         from eval_protocol.utils.subprocess_utils import run_script_and_wait
 
@@ -63,7 +63,7 @@ def _setup_existing_docker_elasticsearch(
         )
         return self._parse_elastic_env_file(env_file_path)
 
-    def _setup_initialized_docker_elasticsearch(self, env_file_path: str) -> ElasticSearchConfig:
+    def _setup_initialized_docker_elasticsearch(self, env_file_path: str) -> ElasticsearchConfig:
         """Set up Elasticsearch by initializing Docker setup from scratch with retry logic."""
         max_retries = 2
         for attempt in range(max_retries):
@@ -126,7 +126,7 @@ def _handle_existing_elasticsearch_container(self, output: str) -> bool:
                 return False
         return False
 
-    def _parse_elastic_env_file(self, env_file_path: str) -> ElasticSearchConfig:
+    def _parse_elastic_env_file(self, env_file_path: str) -> ElasticsearchConfig:
         """Parse ES_LOCAL_API_KEY and ES_LOCAL_URL from .env file."""
         loaded = load_dotenv(env_file_path)
         if not loaded:
@@ -138,7 +138,7 @@ def _parse_elastic_env_file(self, env_file_path: str) -> ElasticSearchConfig:
         if not url or not api_key:
             raise ElasticsearchSetupError("Failed to parse ES_LOCAL_API_KEY and ES_LOCAL_URL from .env file")
 
-        return ElasticSearchConfig(url=url, api_key=api_key, index_name="default-logs")
+        return ElasticsearchConfig(url=url, api_key=api_key, index_name="default-logs")
 
     def create_logging_index(self, index_name: str) -> bool:
         """Create an Elasticsearch index with proper mapping for logging data.

eval_protocol/pytest/remote_rollout_processor.py

@@ -7,7 +7,7 @@
 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
+from eval_protocol.types.remote_rollout_processor import ElasticsearchConfig, InitRequest, RolloutMetadata
 from .rollout_processor import RolloutProcessor
 from .types import RolloutProcessorConfig
 from .elasticsearch_setup import ElasticsearchSetup
@@ -34,7 +34,7 @@ def __init__(
         timeout_seconds: float = 120.0,
         output_data_loader: Callable[[str], DynamicDataLoader],
         disable_elastic_search: bool = False,
-        elastic_search_config: Optional[ElasticSearchConfig] = None,
+        elastic_search_config: Optional[ElasticsearchConfig] = None,
     ):
         # Prefer constructor-provided configuration. These can be overridden via
         # config.kwargs at call time for backward compatibility.
@@ -56,7 +56,7 @@ def setup(self) -> None:
         self._elastic_search_config = self._setup_elastic_search()
         logger.info("Elasticsearch setup complete")
 
-    def _setup_elastic_search(self) -> ElasticSearchConfig:
+    def _setup_elastic_search(self) -> ElasticsearchConfig:
         """Set up Elasticsearch using the dedicated setup module."""
         setup = ElasticsearchSetup()
         return setup.setup_elasticsearch()

eval_protocol/types/remote_rollout_processor.py

@@ -8,7 +8,7 @@
 from eval_protocol.models import Message, Status
 
 
-class ElasticSearchConfig(BaseModel):
+class ElasticsearchConfig(BaseModel):
     """
     Configuration for Elasticsearch.
     """
@@ -38,7 +38,7 @@ class InitRequest(BaseModel):
     """Request model for POST /init endpoint."""
 
     model: str
-    elastic_search_config: Optional[ElasticSearchConfig] = None
+    elastic_search_config: Optional[ElasticsearchConfig] = None
     messages: Optional[List[Message]] = None
     tools: Optional[List[Dict[str, Any]]] = None
 

tests/logging/test_elasticsearch_direct_http_handler.py

@@ -6,7 +6,7 @@
 from eval_protocol.logging.elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
 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
+from eval_protocol.types.remote_rollout_processor import ElasticsearchConfig
 
 
 @pytest.fixture
@@ -39,7 +39,7 @@ def elasticsearch_config():
 
 
 @pytest.fixture
-def elasticsearch_handler(elasticsearch_config: ElasticSearchConfig, rollout_id: str):
+def elasticsearch_handler(elasticsearch_config: ElasticsearchConfig, rollout_id: str):
     """Create and configure ElasticsearchDirectHttpHandler."""
     # Use a unique test-specific index name with timestamp
 
@@ -52,7 +52,7 @@ def elasticsearch_handler(elasticsearch_config: ElasticSearchConfig, rollout_id:
 
 
 @pytest.fixture
-def elasticsearch_client(elasticsearch_config: ElasticSearchConfig):
+def elasticsearch_client(elasticsearch_config: ElasticsearchConfig):
     """Create an Elasticsearch client for testing."""
     # Create a new config instance for the client
     return ElasticsearchClient(elasticsearch_config)