eval-protocol 0.2.34__tar.gz → 0.2.35__tar.gz

{eval_protocol-0.2.34/eval_protocol.egg-info → eval_protocol-0.2.35}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: eval-protocol
-Version: 0.2.34
+Version: 0.2.35
 Summary: The official Python SDK for Eval Protocol (EP.) EP is an open protocol that standardizes how developers author evals for large language model (LLM) applications.
 Author-email: Fireworks AI <info@fireworks.ai>
 License-Expression: MIT

{eval_protocol-0.2.34 → eval_protocol-0.2.35}/eval_protocol/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-09-30T15:39:15-0700",
+ "date": "2025-10-01T13:28:59-0700",
  "dirty": false,
  "error": null,
- "full-revisionid": "c09755b30386c03c95bd79d7b142ed614419c7c4",
- "version": "0.2.34"
+ "full-revisionid": "43ea8eaa8329931e6a9e61aa23a7aeca359f1d1c",
+ "version": "0.2.35"
 }
 '''  # END VERSION_JSON
 

eval_protocol-0.2.35/eval_protocol/logging/elasticsearch_direct_http_handler.py

@@ -0,0 +1,91 @@
+import json
+import logging
+import asyncio
+import threading
+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
+
+
+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.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:
+            # Create proper ISO 8601 timestamp
+            timestamp = datetime.fromtimestamp(record.created).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+            data: Dict[str, Any] = {
+                "@timestamp": timestamp,
+                "level": record.levelname,
+                "message": record.getMessage(),
+                "logger_name": record.name,
+                # Add other relevant record attributes if needed
+            }
+
+            # Schedule the HTTP request to run asynchronously
+            self._schedule_async_send(data, record)
+        except Exception as e:
+            self.handleError(record)
+            print(f"Error preparing log for Elasticsearch: {e}")
+
+    def _schedule_async_send(self, data: Dict[str, Any], record: logging.LogRecord) -> None:
+        """Schedule an async task to send the log data to Elasticsearch."""
+        if self._executor is None:
+            self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="elasticsearch-logger")
+
+        # Submit the HTTP request to the thread pool
+        future = self._executor.submit(self._send_to_elasticsearch, data, record)
+
+        # Add error handling callback
+        future.add_done_callback(lambda f: self._handle_async_result(f, record))
+
+    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
+        except Exception as e:
+            # Re-raise to be handled by the callback
+            raise e
+
+    def _handle_async_result(self, future, record: logging.LogRecord) -> None:
+        """Handle the result of the async send operation."""
+        try:
+            future.result()  # This will raise any exception that occurred
+        except Exception as e:
+            self.handleError(record)
+            # You might want to log this error to a file or console
+            # to prevent a logging loop.
+            if hasattr(e, "response") and getattr(e, "response", None) is not None:
+                print(f"Error sending log to Elasticsearch: {e}")
+                print(f"Response content: {getattr(e, 'response').text}")
+            else:
+                print(f"Error sending log to Elasticsearch: {e}")
+
+    def close(self) -> None:
+        """Clean up resources when the handler is closed."""
+        super().close()
+        if self._executor:
+            self._executor.shutdown(wait=True)

eval_protocol-0.2.35/eval_protocol/logging/elasticsearch_index_manager.py

@@ -0,0 +1,187 @@
+import requests
+from typing import Dict, Any, Optional
+from urllib.parse import urlparse
+
+
+class ElasticsearchIndexManager:
+    """Manages Elasticsearch index creation and mapping configuration."""
+
+    def __init__(self, base_url: str, index_name: str, api_key: str) -> None:
+        """Initialize the Elasticsearch index manager.
+
+        Args:
+            base_url: Elasticsearch base URL (e.g., "https://localhost:9200")
+            index_name: Name of the index to manage
+            api_key: API key for authentication
+        """
+        self.base_url: str = base_url.rstrip("/")
+        self.index_name: str = index_name
+        self.api_key: str = api_key
+        self.index_url: str = f"{self.base_url}/{self.index_name}"
+        self._mapping_created: bool = False
+
+        # Parse URL to determine if we should verify SSL
+        parsed_url = urlparse(base_url)
+        self.verify_ssl = parsed_url.scheme == "https"
+
+    def create_logging_index_mapping(self) -> bool:
+        """Create index with proper mapping for logging data.
+
+        Returns:
+            bool: True if mapping was created successfully, False otherwise.
+        """
+        if self._mapping_created:
+            return True
+
+        try:
+            # Check if index exists and has correct mapping
+            if self._index_exists_with_correct_mapping():
+                self._mapping_created = True
+                return True
+
+            # If index exists but has wrong mapping, delete and recreate it
+            if self.index_exists():
+                print(f"Warning: Index {self.index_name} exists with incorrect mapping. Deleting and recreating...")
+                if not self.delete_index():
+                    print(f"Warning: Failed to delete existing index {self.index_name}")
+                    return False
+
+            # Create index with proper mapping
+            mapping = self._get_logging_mapping()
+            response = requests.put(
+                self.index_url,
+                headers={"Content-Type": "application/json", "Authorization": f"ApiKey {self.api_key}"},
+                json=mapping,
+                verify=self.verify_ssl,
+            )
+
+            if response.status_code in [200, 201]:
+                self._mapping_created = True
+                return True
+            else:
+                print(f"Warning: Failed to create index mapping: {response.status_code} - {response.text}")
+                return False
+
+        except Exception as e:
+            print(f"Warning: Failed to create index mapping: {e}")
+            return False
+
+    def _index_exists_with_correct_mapping(self) -> bool:
+        """Check if index exists and has the correct @timestamp mapping.
+
+        Returns:
+            bool: True if index exists with correct mapping, False otherwise.
+        """
+        try:
+            # Check if index exists
+            response = requests.head(
+                self.index_url, headers={"Authorization": f"ApiKey {self.api_key}"}, verify=self.verify_ssl
+            )
+
+            if response.status_code != 200:
+                return False
+
+            # Check if mapping is correct
+            mapping_response = requests.get(
+                f"{self.index_url}/_mapping",
+                headers={"Authorization": f"ApiKey {self.api_key}"},
+                verify=self.verify_ssl,
+            )
+
+            if mapping_response.status_code != 200:
+                return False
+
+            mapping_data = mapping_response.json()
+            return self._has_correct_timestamp_mapping(mapping_data)
+
+        except Exception:
+            return False
+
+    def _has_correct_timestamp_mapping(self, mapping_data: Dict[str, Any]) -> bool:
+        """Check if the mapping has @timestamp as a date field.
+
+        Args:
+            mapping_data: Elasticsearch mapping response data
+
+        Returns:
+            bool: True if @timestamp is correctly mapped as date field
+        """
+        try:
+            return (
+                self.index_name in mapping_data
+                and "mappings" in mapping_data[self.index_name]
+                and "properties" in mapping_data[self.index_name]["mappings"]
+                and "@timestamp" in mapping_data[self.index_name]["mappings"]["properties"]
+                and mapping_data[self.index_name]["mappings"]["properties"]["@timestamp"].get("type") == "date"
+            )
+        except (KeyError, TypeError):
+            return False
+
+    def _get_logging_mapping(self) -> Dict[str, Any]:
+        """Get the standard mapping for logging data.
+
+        Returns:
+            Dict containing the index mapping configuration
+        """
+        return {
+            "mappings": {
+                "properties": {
+                    "@timestamp": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+                    "level": {"type": "keyword"},
+                    "message": {"type": "text"},
+                    "logger_name": {"type": "keyword"},
+                }
+            }
+        }
+
+    def delete_index(self) -> bool:
+        """Delete the managed index.
+
+        Returns:
+            bool: True if index was deleted successfully, False otherwise.
+        """
+        try:
+            response = requests.delete(
+                self.index_url, headers={"Authorization": f"ApiKey {self.api_key}"}, verify=self.verify_ssl
+            )
+            if response.status_code in [200, 404]:  # 404 means index doesn't exist, which is fine
+                self._mapping_created = False
+                return True
+            else:
+                print(f"Warning: Failed to delete index: {response.status_code} - {response.text}")
+                return False
+        except Exception as e:
+            print(f"Warning: Failed to delete index: {e}")
+            return False
+
+    def index_exists(self) -> bool:
+        """Check if the index exists.
+
+        Returns:
+            bool: True if index exists, False otherwise.
+        """
+        try:
+            response = requests.head(
+                self.index_url, headers={"Authorization": f"ApiKey {self.api_key}"}, verify=self.verify_ssl
+            )
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    def get_index_stats(self) -> Optional[Dict[str, Any]]:
+        """Get statistics about the index.
+
+        Returns:
+            Dict containing index statistics, or None if failed
+        """
+        try:
+            response = requests.get(
+                f"{self.index_url}/_stats",
+                headers={"Authorization": f"ApiKey {self.api_key}"},
+                verify=self.verify_ssl,
+            )
+            if response.status_code == 200:
+                return response.json()
+            return None
+        except Exception:
+            return None

eval_protocol-0.2.35/eval_protocol/pytest/elasticsearch_setup.py

@@ -0,0 +1,167 @@
+import os
+import subprocess
+import tempfile
+import logging
+from typing import Optional
+
+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.logging.elasticsearch_index_manager import ElasticsearchIndexManager
+
+logger = logging.getLogger(__name__)
+
+
+class ElasticsearchSetupError(Exception):
+    """Exception raised when Elasticsearch setup fails."""
+
+    pass
+
+
+class ElasticsearchSetup:
+    """Handles Elasticsearch setup with retry logic for existing containers."""
+
+    def __init__(self):
+        self.eval_protocol_dir = find_eval_protocol_dir()
+
+    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.
+        """
+        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")
+
+        # If elastic-start-local directory exists, use existing Docker script
+        if os.path.exists(elastic_start_local_dir):
+            config = self._setup_existing_docker_elasticsearch(elastic_start_local_dir, env_file_path)
+        else:
+            # Otherwise, initialize Docker setup from scratch
+            config = self._setup_initialized_docker_elasticsearch(env_file_path)
+
+        # Create the logging index with proper mapping
+        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)
+
+    def _setup_existing_docker_elasticsearch(
+        self, elastic_start_local_dir: str, env_file_path: str
+    ) -> ElasticSearchConfig:
+        """Set up Elasticsearch using existing Docker start.sh script."""
+        from eval_protocol.utils.subprocess_utils import run_script_and_wait
+
+        run_script_and_wait(
+            script_name="start.sh",
+            working_directory=elastic_start_local_dir,
+            inherit_stdout=True,
+        )
+        return self._parse_elastic_env_file(env_file_path)
+
+    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):
+            # Use a temporary file to capture output while also showing it in parent stdout
+            with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file:
+                temp_file_path = temp_file.name
+
+            try:
+                # Run the command and tee output to both stdout and temp file
+                # Use set -o pipefail to ensure we get the return code of the first failing command
+                process = subprocess.Popen(
+                    [
+                        "sh",
+                        "-c",
+                        f"set -o pipefail; curl -fsSL https://elastic.co/start-local | sh -s -- --esonly | tee {temp_file_path}",
+                    ],
+                    cwd=self.eval_protocol_dir,
+                )
+                returncode = process.wait()
+
+                # Read the captured output
+                with open(temp_file_path, "r") as f:
+                    stdout = f.read()
+
+                if returncode == 0:
+                    return self._parse_elastic_env_file(env_file_path)
+
+                # Check if container is already running and handle it
+                if self._handle_existing_elasticsearch_container(stdout):
+                    logger.info(f"Retrying Elasticsearch setup (attempt {attempt + 1}/{max_retries})")
+                    continue
+
+                # If we get here, it's a different error
+                raise ElasticsearchSetupError(
+                    f"Failed to start Elasticsearch (attempt {attempt + 1}/{max_retries}): {stdout}"
+                )
+
+            finally:
+                # Clean up the temporary file
+                try:
+                    os.unlink(temp_file_path)
+                except OSError:
+                    pass
+
+        raise ElasticsearchSetupError(f"Failed to start Elasticsearch after {max_retries} attempts")
+
+    def _handle_existing_elasticsearch_container(self, output: str) -> bool:
+        """
+        Check if the curl command output indicates that the Elasticsearch container is already running.
+        If so, stop the existing container and return True to indicate a retry is needed.
+        """
+        if "docker stop es-local-dev" in output:
+            logger.info("Elasticsearch container 'es-local-dev' is already running. Stopping it...")
+            try:
+                subprocess.run(["docker", "stop", "es-local-dev"], check=True, capture_output=True, text=True)
+                logger.info("Successfully stopped existing Elasticsearch container")
+                return True  # Indicate retry is needed
+            except subprocess.CalledProcessError as e:
+                logger.warning(f"Failed to stop existing container: {e}")
+                return False
+        return False
+
+    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:
+            raise ElasticsearchSetupError("Failed to load .env file")
+
+        api_key = os.getenv("ES_LOCAL_API_KEY")
+        url = os.getenv("ES_LOCAL_URL")
+
+        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")
+
+    def create_logging_index(self, index_name: str) -> bool:
+        """Create an Elasticsearch index with proper mapping for logging data.
+
+        Args:
+            index_name: Name of the index to create
+
+        Returns:
+            bool: True if index was created successfully, False otherwise.
+        """
+        try:
+            # Get the config from the .env file
+            config = self._parse_elastic_env_file(self._get_env_file_path())
+
+            # Create index manager and set up mapping
+            index_manager = ElasticsearchIndexManager(config.url, index_name, config.api_key)
+            return index_manager.create_logging_index_mapping()
+
+        except Exception as e:
+            logger.error(f"Failed to create logging index {index_name}: {e}")
+            return False
+
+    def _get_env_file_path(self) -> str:
+        """Get the path to the .env file."""
+        elastic_start_local_dir = os.path.join(self.eval_protocol_dir, "elastic-start-local")
+        return os.path.join(elastic_start_local_dir, ".env")

{eval_protocol-0.2.34 → eval_protocol-0.2.35}/eval_protocol/pytest/evaluation_test.py

@@ -367,6 +367,8 @@ def evaluation_test(
                         exception_handler_config=exception_handler_config,
                     )
 
+                    rollout_processor.setup()
+
                     async def execute_run(run_idx: int, config: RolloutProcessorConfig):
                         nonlocal all_results
 

{eval_protocol-0.2.34 → eval_protocol-0.2.35}/eval_protocol/pytest/remote_rollout_processor.py

@@ -6,11 +6,16 @@ import requests
 
 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 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
+import logging
+
 import os
 
+logger = logging.getLogger(__name__)
+
 
 class RemoteRolloutProcessor(RolloutProcessor):
     """
@@ -27,6 +32,8 @@ class RemoteRolloutProcessor(RolloutProcessor):
         poll_interval: float = 1.0,
         timeout_seconds: float = 120.0,
         output_data_loader: Callable[[str], DynamicDataLoader],
+        disable_elastic_search: bool = False,
+        elastic_search_config: Optional[ElasticSearchConfig] = None,
     ):
         # Prefer constructor-provided configuration. These can be overridden via
         # config.kwargs at call time for backward compatibility.
@@ -37,6 +44,21 @@ class RemoteRolloutProcessor(RolloutProcessor):
         self._poll_interval = poll_interval
         self._timeout_seconds = timeout_seconds
         self._output_data_loader = output_data_loader
+        self._disable_elastic_search = disable_elastic_search
+        self._elastic_search_config = elastic_search_config
+
+    def setup(self) -> None:
+        if self._disable_elastic_search:
+            logger.info("Elasticsearch is disabled, skipping setup")
+            return
+        logger.info("Setting up Elasticsearch")
+        self._elastic_search_config = self._setup_elastic_search()
+        logger.info("Elasticsearch setup complete")
+
+    def _setup_elastic_search(self) -> ElasticSearchConfig:
+        """Set up Elasticsearch using the dedicated setup module."""
+        setup = ElasticsearchSetup()
+        return setup.setup_elasticsearch()
 
     def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[asyncio.Task[EvaluationRow]]:
         tasks: List[asyncio.Task[EvaluationRow]] = []
@@ -113,12 +135,23 @@ class RemoteRolloutProcessor(RolloutProcessor):
             if row.execution_metadata.rollout_id is None:
                 raise ValueError("Rollout ID is required in RemoteRolloutProcessor")
 
+            final_model_base_url = model_base_url
+            if model_base_url and model_base_url.startswith("https://tracing.fireworks.ai/project_id/"):
+                final_model_base_url = (
+                    f"{model_base_url}/rollout_id/{meta.rollout_id}"
+                    f"/invocation_id/{meta.invocation_id}"
+                    f"/experiment_id/{meta.experiment_id}"
+                    f"/run_id/{meta.run_id}"
+                    f"/row_id/{meta.row_id}"
+                )
+
             init_payload: InitRequest = InitRequest(
                 model=model,
                 messages=clean_messages,
                 tools=row.tools,
                 metadata=meta,
-                model_base_url=model_base_url,
+                model_base_url=final_model_base_url,
+                elastic_search_config=self._elastic_search_config,
             )
 
             # Fire-and-poll
@@ -197,6 +230,11 @@ class RemoteRolloutProcessor(RolloutProcessor):
                     langfuse_row.input_metadata.dataset_info = row.input_metadata.dataset_info
                 langfuse_row.eval_metadata = row.eval_metadata
                 langfuse_row.ground_truth = row.ground_truth
+
+                # this is useful to detect stopped evaluations so we can update
+                # the status in the logs server
+                langfuse_row.pid = row.pid
+
                 return langfuse_row
             else:
                 raise ValueError("RemoteRolloutProcessor's output_data_loader should return exactly one row.")

{eval_protocol-0.2.34 → eval_protocol-0.2.35}/eval_protocol/pytest/rollout_processor.py

@@ -10,6 +10,10 @@ class RolloutProcessor(ABC):
     Abstract base class for all rollout processor strategies.
     """
 
+    def setup(self) -> None:
+        """Setup resources. Override in subclasses if setup is needed. Executed once per invocation."""
+        pass
+
     @abstractmethod
     def __call__(self, rows: list[EvaluationRow], config: RolloutProcessorConfig) -> list[asyncio.Task[EvaluationRow]]:
         """Process evaluation rows and return async tasks. Must be implemented by subclasses."""

{eval_protocol-0.2.34 → eval_protocol-0.2.35}/eval_protocol/types/remote_rollout_processor.py

@@ -7,6 +7,16 @@ from pydantic import BaseModel, Field
 from eval_protocol.models import Message, Status
 
 
+class ElasticSearchConfig(BaseModel):
+    """
+    Configuration for Elasticsearch.
+    """
+
+    url: str
+    api_key: str
+    index_name: str
+
+
 class RolloutMetadata(BaseModel):
     """Metadata for rollout execution."""
 
@@ -21,6 +31,7 @@ class InitRequest(BaseModel):
     """Request model for POST /init endpoint."""
 
     model: str
+    elastic_search_config: Optional[ElasticSearchConfig] = None
     messages: Optional[List[Message]] = None
     tools: Optional[List[Dict[str, Any]]] = None