Merge branch 'main' into completionparams

Author: shreymodi1

.github/workflows/elasticsearch-tests.yml …ub/workflows/fireworks-tracing-tests.yml.github/workflows/elasticsearch-tests.yml renamed to .github/workflows/fireworks-tracing-tests.yml .github/workflows/elasticsearch-tests.yml renamed to .github/workflows/fireworks-tracing-tests.yml

@@ -1,4 +1,4 @@
-name: Elasticsearch Tests
+name: Fireworks Tracing Tests
 
 on:
   push:
@@ -13,8 +13,8 @@ on:
   workflow_dispatch: # Allow manual triggering
 
 jobs:
-  elasticsearch-tests:
-    name: Elasticsearch Integration Tests
+  fireworks-tracing-tests:
+    name: Fireworks Tracing Integration Tests
     runs-on: ubuntu-latest
 
     steps:
@@ -36,14 +36,15 @@ jobs:
       - name: Install the project
         run: uv sync --locked --all-extras --dev
 
-      - name: Run Elasticsearch Tests
+      - name: Run Fireworks Tracing Tests
         env:
           FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }}
           PYTHONWARNINGS: "ignore::DeprecationWarning,ignore::RuntimeWarning"
         run: |
-          # Run Elasticsearch direct HTTP handler tests
-          uv run pytest tests/logging/test_elasticsearch_direct_http_handler.py -v --tb=short
+          # Run RemoteRolloutProcessor End-to-End Test (auto server startup)
+          uv run pytest tests/remote_server/test_remote_fireworks.py::test_remote_rollout_and_fetch_fireworks \
+            -v --tb=short
 
-          # Run RemoteRolloutProcessor Propagate Status Smoke Test (also uses Elasticsearch)
+          # Run RemoteRolloutProcessor Propagate Status Test (auto server startup)
           uv run pytest tests/remote_server/test_remote_fireworks_propagate_status.py::test_remote_rollout_and_fetch_fireworks_propagate_status \
             -v --tb=short

eval_protocol/__init__.py

@@ -36,7 +36,6 @@
     filter_longest_conversation,
 )
 from .pytest import evaluation_test, SingleTurnRolloutProcessor, RemoteRolloutProcessor, GithubActionRolloutProcessor
-from .pytest.remote_rollout_processor import create_elasticsearch_config_from_env
 from .pytest.parameterize import DefaultParameterIdGenerator
 from .log_utils.elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
 from .log_utils.rollout_id_filter import RolloutIdFilter
@@ -90,7 +89,6 @@
 warnings.filterwarnings("default", category=DeprecationWarning, module="eval_protocol")
 
 __all__ = [
-    "create_elasticsearch_config_from_env",
     "ElasticsearchConfig",
     "ElasticsearchDirectHttpHandler",
     "RolloutIdFilter",

eval_protocol/adapters/fireworks_tracing.py

@@ -265,6 +265,55 @@ def __init__(
         self.base_url = base_url.rstrip("/")
         self.timeout = timeout
 
+    def search_logs(self, tags: List[str], limit: int = 100, hours_back: int = 24) -> List[Dict[str, Any]]:
+        """Fetch logs from Fireworks tracing gateway /logs endpoint.
+
+        Returns entries with keys: timestamp, message, severity, tags.
+        """
+        if not tags:
+            raise ValueError("At least one tag is required to fetch logs")
+
+        headers = {"Authorization": f"Bearer {os.environ.get('FIREWORKS_API_KEY')}"}
+        params: Dict[str, Any] = {"tags": tags, "limit": limit, "hours_back": hours_back, "program": "eval_protocol"}
+
+        # Try /logs first, fall back to /v1/logs if not found
+        urls_to_try = [f"{self.base_url}/logs", f"{self.base_url}/v1/logs"]
+        data: Dict[str, Any] = {}
+        last_error: Optional[str] = None
+        for url in urls_to_try:
+            try:
+                response = requests.get(url, params=params, timeout=self.timeout, headers=headers)
+                if response.status_code == 404:
+                    # Try next variant
+                    last_error = f"404 for {url}"
+                    continue
+                response.raise_for_status()
+                data = response.json() or {}
+                break
+            except requests.exceptions.RequestException as e:
+                last_error = str(e)
+                continue
+        else:
+            # All attempts failed
+            if last_error:
+                logger.error("Failed to fetch logs from Fireworks (tried %s): %s", urls_to_try, last_error)
+            return []
+
+        entries: List[Dict[str, Any]] = data.get("entries", []) or []
+        # Normalize minimal shape
+        results: List[Dict[str, Any]] = []
+        for e in entries:
+            results.append(
+                {
+                    "timestamp": e.get("timestamp"),
+                    "message": e.get("message"),
+                    "severity": e.get("severity", "INFO"),
+                    "tags": e.get("tags", []),
+                    "status": e.get("status"),
+                }
+            )
+        return results
+
     def get_evaluation_rows(
         self,
         tags: List[str],

eval_protocol/cli.py

@@ -307,6 +307,16 @@ def parse_args(args=None):
         action="store_true",
         help="Use env vars for Elasticsearch config (requires ELASTICSEARCH_URL, ELASTICSEARCH_API_KEY, ELASTICSEARCH_INDEX_NAME)",
     )
+    logs_parser.add_argument(
+        "--use-fireworks",
+        action="store_true",
+        help="Force Fireworks tracing backend for logs UI (overrides env auto-detection)",
+    )
+    logs_parser.add_argument(
+        "--use-elasticsearch",
+        action="store_true",
+        help="Force Elasticsearch backend for logs UI (overrides env auto-detection)",
+    )
 
     # Upload command
     upload_parser = subparsers.add_parser(

eval_protocol/cli_commands/logs.py

@@ -5,6 +5,7 @@
 import sys
 from pathlib import Path
 
+import os
 from ..utils.logs_server import serve_logs
 
 
@@ -20,53 +21,33 @@ def logs_command(args):
     print("Press Ctrl+C to stop the server")
     print("-" * 50)
 
-    # Setup Elasticsearch based on flags
-    elasticsearch_config = None
-    try:
-        if getattr(args, "use_env_elasticsearch_config", False):
-            # Use environment variables for configuration
-            print("⚙️ Using environment variables for Elasticsearch config")
-            from eval_protocol.pytest.remote_rollout_processor import (
-                create_elasticsearch_config_from_env,
-            )
-
-            elasticsearch_config = create_elasticsearch_config_from_env()
-            # Ensure index exists with correct mapping, mirroring Docker setup path
-            try:
-                from eval_protocol.log_utils.elasticsearch_index_manager import (
-                    ElasticsearchIndexManager,
-                )
-
-                index_manager = ElasticsearchIndexManager(
-                    elasticsearch_config.url,
-                    elasticsearch_config.index_name,
-                    elasticsearch_config.api_key,
-                )
-                created = index_manager.create_logging_index_mapping()
-                if created:
-                    print(
-                        f"🧭 Verified Elasticsearch index '{elasticsearch_config.index_name}' mapping (created or already correct)"
-                    )
-                else:
-                    print(
-                        f"⚠️ Could not verify/create mapping for index '{elasticsearch_config.index_name}'. Searches may behave unexpectedly."
-                    )
-            except Exception as e:
-                print(f"⚠️ Failed to ensure index mapping via IndexManager: {e}")
-        elif not getattr(args, "disable_elasticsearch_setup", False):
-            # Default behavior: start or connect to local Elasticsearch via Docker helper
-            from eval_protocol.pytest.elasticsearch_setup import ElasticsearchSetup
+    # Backend selection: Fireworks first when API key present, unless overridden
+    use_fireworks = False
+    if getattr(args, "use_fireworks", False):
+        use_fireworks = True
+    elif getattr(args, "use_elasticsearch", False):
+        use_fireworks = False
+    else:
+        use_fireworks = bool(os.environ.get("FIREWORKS_API_KEY"))
 
-            print("🧰 Auto-configuring local Elasticsearch (Docker)")
-            elasticsearch_config = ElasticsearchSetup().setup_elasticsearch()
-        else:
-            print("🚫 Elasticsearch setup disabled; running without Elasticsearch integration")
-    except Exception as e:
-        print(f"❌ Failed to configure Elasticsearch: {e}")
-        return 1
+    # Setup backend configs
+    elasticsearch_config = None
+    # Prefer explicit FW_TRACING_GATEWAY_BASE_URL, then GATEWAY_URL from env (remote validation),
+    # finally default to public tracing.fireworks.ai
+    fireworks_base_url = (
+        os.environ.get("FW_TRACING_GATEWAY_BASE_URL")
+        or os.environ.get("GATEWAY_URL")
+        or "https://tracing.fireworks.ai"
+    )
 
     try:
-        serve_logs(port=args.port, elasticsearch_config=elasticsearch_config, debug=args.debug)
+        serve_logs(
+            port=args.port,
+            elasticsearch_config=elasticsearch_config,
+            debug=args.debug,
+            backend="fireworks" if use_fireworks else "elasticsearch",
+            fireworks_base_url=fireworks_base_url if use_fireworks else None,
+        )
         return 0
     except KeyboardInterrupt:
         print("\n🛑 Server stopped by user")

eval_protocol/cli_commands/upload.py

@@ -267,32 +267,29 @@ def _parse_entry(entry: str, cwd: str) -> tuple[str, str]:
 def _resolve_entry_to_qual_and_source(entry: str, cwd: str) -> tuple[str, str]:
     target, func = _parse_entry(entry, cwd)
 
-    # Check if target looks like a file path
+    # Determine the file path to load
     if "/" in target or "\\" in target or os.path.exists(target):
-        # It's a file path - convert to absolute and load as module
+        # It's a file path - convert to absolute
         if not os.path.isabs(target):
             target = os.path.abspath(os.path.join(cwd, target))
-
         if not target.endswith(".py"):
             target = target + ".py"
-
         if not os.path.isfile(target):
             raise ValueError(f"File not found: {target}")
-
-        # Import module from file path
-        spec = importlib.util.spec_from_file_location(Path(target).stem, target)
-        if not spec or not spec.loader:
-            raise ValueError(f"Unable to load module from path: {target}")
-        module = importlib.util.module_from_spec(spec)
-        sys.modules[spec.name] = module
-        spec.loader.exec_module(module)  # type: ignore[attr-defined]
-        module_name = spec.name
         source_file_path = target
     else:
-        # Treat as module path (e.g., "my_package.my_module")
-        module_name = target
-        module = importlib.import_module(module_name)
-        source_file_path = getattr(module, "__file__", "") or ""
+        # Treat dotted name as a file path
+        dotted_as_path = target.replace(".", "/") + ".py"
+        source_file_path = os.path.join(cwd, dotted_as_path)
+
+    # Load the module from the file path
+    spec = importlib.util.spec_from_file_location(Path(source_file_path).stem, source_file_path)
+    if not spec or not spec.loader:
+        raise ValueError(f"Unable to load module from path: {source_file_path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)  # type: ignore[attr-defined]
+    module_name = spec.name
 
     if not hasattr(module, func):
         raise ValueError(f"Function '{func}' not found in module '{module_name}'")
@@ -591,8 +588,7 @@ def upload_command(args: argparse.Namespace) -> int:
 
         print(f"\nUploading evaluator '{evaluator_id}' for {qualname.split('.')[-1]}...")
         try:
-            # Always treat as a single evaluator (single-metric) even if folder has helper modules
-            test_dir = os.path.dirname(source_file_path) if source_file_path else root
+            test_dir = root
             metric_name = os.path.basename(test_dir) or "metric"
             result = create_evaluation(
                 evaluator_id=evaluator_id,

eval_protocol/log_utils/elasticsearch_direct_http_handler.py

@@ -60,6 +60,17 @@ def emit(self, record: logging.LogRecord) -> None:
             if status_info:
                 data.update(status_info)
 
+            # Optional correlation enrichment
+            experiment_id = getattr(record, "experiment_id", None)
+            if experiment_id is not None:
+                data["experiment_id"] = experiment_id
+            run_id = getattr(record, "run_id", None)
+            if run_id is not None:
+                data["run_id"] = run_id
+            rollout_ids = getattr(record, "rollout_ids", None)
+            if rollout_ids is not None:
+                data["rollout_ids"] = rollout_ids
+
             # Schedule the HTTP request to run asynchronously
             self._schedule_async_send(data, record)
         except Exception as e: