eval-protocol 0.2.35.dev1__tar.gz → 0.2.36__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: eval-protocol
-Version: 0.2.35.dev1
+Version: 0.2.36
 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.35.dev1 → eval_protocol-0.2.36}/eval_protocol/__init__.py

@@ -24,7 +24,7 @@ from .mcp_env import (
 )
 from .data_loader import DynamicDataLoader, InlineDataLoader
 from . import mcp, rewards
-from .models import EvaluateResult, Message, MetricResult, EvaluationRow, InputMetadata
+from .models import EvaluateResult, Message, MetricResult, EvaluationRow, InputMetadata, Status
 from .playback_policy import PlaybackPolicyBase
 from .resources import create_llm_resource
 from .reward_function import RewardFunction
@@ -63,6 +63,7 @@ except ImportError:
 warnings.filterwarnings("default", category=DeprecationWarning, module="eval_protocol")
 
 __all__ = [
+    "Status",
     "RemoteRolloutProcessor",
     "InputMetadata",
     "EvaluationRow",

{eval_protocol-0.2.35.dev1 → eval_protocol-0.2.36}/eval_protocol/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-10-02T01:54:12-0700",
+ "date": "2025-10-02T09:18:41-0700",
  "dirty": false,
  "error": null,
- "full-revisionid": "c4108ba1d87ba6fb76100c63e6ee16f48bc06598",
- "version": "0.2.35-dev1"
+ "full-revisionid": "1283ce226147f55336cd4fe56c744640a0a9ba9b",
+ "version": "0.2.36"
 }
 '''  # END VERSION_JSON
 

eval_protocol-0.2.36/eval_protocol/logging/elasticsearch_client.py

@@ -0,0 +1,286 @@
+"""
+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 eval_protocol.types.remote_rollout_processor import ElasticsearchConfig
+
+
+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-0.2.35.dev1 → eval_protocol-0.2.36}/eval_protocol/logging/elasticsearch_direct_http_handler.py

@@ -1,50 +1,92 @@
 import json
 import logging
 import asyncio
+import os
 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
+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.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 = ElasticsearchConfig(
+            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:
             # Create proper ISO 8601 timestamp
             timestamp = datetime.fromtimestamp(record.created).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
 
+            rollout_id = self._get_rollout_id(record)
+            status_info = self._get_status_info(record)
+
             data: Dict[str, Any] = {
                 "@timestamp": timestamp,
                 "level": record.levelname,
                 "message": record.getMessage(),
                 "logger_name": record.name,
-                # Add other relevant record attributes if needed
+                "rollout_id": rollout_id,
             }
 
+            # Add status information if present
+            if status_info:
+                data.update(status_info)
+
             # 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 _get_rollout_id(self, record: logging.LogRecord) -> str:
+        """Get the rollout ID from environment variables."""
+        rollout_id = os.getenv("EP_ROLLOUT_ID")
+        if rollout_id is None:
+            raise ValueError(
+                "EP_ROLLOUT_ID environment variable is not set but needed for ElasticsearchDirectHttpHandler"
+            )
+        return rollout_id
+
+    def _get_status_info(self, record: logging.LogRecord) -> Optional[Dict[str, Any]]:
+        """Extract status information from the log record's extra data."""
+        # Check if 'status' is in the extra data (passed via extra parameter)
+        if hasattr(record, "status") and record.status is not None:  # type: ignore
+            status = record.status  # type: ignore
+
+            # Handle Status class instances (Pydantic BaseModel)
+            if hasattr(status, "code") and hasattr(status, "message"):
+                # Status object - extract code and message
+                status_code = status.code
+                # Handle both enum values and direct integer values
+                if hasattr(status_code, "value"):
+                    status_code = status_code.value
+
+                return {
+                    "status_code": status_code,
+                    "status_message": status.message,
+                    "status_details": getattr(status, "details", []),
+                }
+            elif isinstance(status, dict):
+                # Dictionary representation of status
+                return {
+                    "status_code": status.get("code"),
+                    "status_message": status.get("message"),
+                    "status_details": status.get("details", []),
+                }
+        return None
+
     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:
@@ -59,13 +101,9 @@ class ElasticsearchDirectHttpHandler(logging.Handler):
     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

{eval_protocol-0.2.35.dev1 → eval_protocol-0.2.36}/eval_protocol/logging/elasticsearch_index_manager.py

@@ -1,6 +1,6 @@
-import requests
 from typing import Dict, Any, Optional
-from urllib.parse import urlparse
+from .elasticsearch_client import ElasticsearchClient
+from eval_protocol.types.remote_rollout_processor import ElasticsearchConfig
 
 
 class ElasticsearchIndexManager:
@@ -14,16 +14,10 @@ class ElasticsearchIndexManager:
             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.config = ElasticsearchConfig(url=base_url, api_key=api_key, index_name=index_name)
+        self.client = ElasticsearchClient(self.config)
         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.
 
@@ -41,25 +35,22 @@ class ElasticsearchIndexManager:
 
             # 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...")
+                print(
+                    f"Warning: Index {self.config.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}")
+                    print(f"Warning: Failed to delete existing index {self.config.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]:
+            success = self.client.create_index(mapping)
+
+            if success:
                 self._mapping_created = True
                 return True
             else:
-                print(f"Warning: Failed to create index mapping: {response.status_code} - {response.text}")
+                print("Warning: Failed to create index mapping")
                 return False
 
         except Exception as e:
@@ -74,46 +65,50 @@ class ElasticsearchIndexManager:
         """
         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:
+            if not self.client.index_exists():
                 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:
+            mapping_data = self.client.get_mapping()
+            if mapping_data is None:
                 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.
+        """Check if the mapping has @timestamp as a date field, rollout_id as a keyword field, and status fields.
 
         Args:
             mapping_data: Elasticsearch mapping response data
 
         Returns:
-            bool: True if @timestamp is correctly mapped as date field
+            bool: True if all required fields are correctly mapped
         """
         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"
-            )
+            if not (
+                self.config.index_name in mapping_data
+                and "mappings" in mapping_data[self.config.index_name]
+                and "properties" in mapping_data[self.config.index_name]["mappings"]
+            ):
+                return False
+
+            properties = mapping_data[self.config.index_name]["mappings"]["properties"]
+
+            # Check @timestamp is mapped as date
+            timestamp_ok = "@timestamp" in properties and properties["@timestamp"].get("type") == "date"
+
+            # Check rollout_id is mapped as keyword
+            rollout_id_ok = "rollout_id" in properties and properties["rollout_id"].get("type") == "keyword"
+
+            # Check status fields are mapped correctly
+            status_code_ok = "status_code" in properties and properties["status_code"].get("type") == "integer"
+            status_message_ok = "status_message" in properties and properties["status_message"].get("type") == "text"
+            status_details_ok = "status_details" in properties and properties["status_details"].get("type") == "object"
+
+            return timestamp_ok and rollout_id_ok and status_code_ok and status_message_ok and status_details_ok
         except (KeyError, TypeError):
             return False
 
@@ -130,6 +125,10 @@ class ElasticsearchIndexManager:
                     "level": {"type": "keyword"},
                     "message": {"type": "text"},
                     "logger_name": {"type": "keyword"},
+                    "rollout_id": {"type": "keyword"},
+                    "status_code": {"type": "integer"},
+                    "status_message": {"type": "text"},
+                    "status_details": {"type": "object"},
                 }
             }
         }
@@ -141,14 +140,12 @@ class ElasticsearchIndexManager:
             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
+            success = self.client.delete_index()
+            if success:
                 self._mapping_created = False
                 return True
             else:
-                print(f"Warning: Failed to delete index: {response.status_code} - {response.text}")
+                print("Warning: Failed to delete index")
                 return False
         except Exception as e:
             print(f"Warning: Failed to delete index: {e}")
@@ -160,13 +157,7 @@ class ElasticsearchIndexManager:
         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
+        return self.client.index_exists()
 
     def get_index_stats(self) -> Optional[Dict[str, Any]]:
         """Get statistics about the index.
@@ -174,14 +165,4 @@ class ElasticsearchIndexManager:
         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
+        return self.client.get_index_stats()