address comments

Author: xzrderek

eval_protocol/adapters/__init__.py

@@ -25,9 +25,9 @@
 except ImportError:
     pass
 
-from .fireworks_tracing import FireworksTracingAdapter, create_fireworks_tracing_adapter
+from .fireworks_tracing import FireworksTracingAdapter
 
-__all__.extend(["FireworksTracingAdapter", "create_fireworks_tracing_adapter"])
+__all__.extend(["FireworksTracingAdapter"])
 
 try:
     from .huggingface import (

eval_protocol/adapters/fireworks_tracing.py

@@ -375,18 +375,3 @@ def get_evaluation_rows(
 
         logger.info("Successfully converted %d traces to evaluation rows", len(eval_rows))
         return eval_rows
-
-
-def create_fireworks_tracing_adapter(
-    project_id: Optional[str] = None, base_url: str = "https://tracing.fireworks.ai"
-) -> FireworksTracingAdapter:
-    """Factory function to create a Fireworks Tracing adapter.
-
-    Args:
-        project_id: Optional project ID. If not provided, uses the default project configured on the server.
-        base_url: The base URL of the tracing proxy
-
-    Returns:
-        FireworksTracingAdapter instance
-    """
-    return FireworksTracingAdapter(project_id=project_id, base_url=base_url)

eval_protocol/pytest/remote_rollout_processor.py

@@ -13,7 +13,7 @@
     InitRequest,
     RolloutMetadata,
 )
-from eval_protocol.adapters.fireworks_tracing import create_fireworks_tracing_adapter
+from eval_protocol.adapters.fireworks_tracing import FireworksTracingAdapter
 from eval_protocol.quickstart.utils import filter_longest_conversation
 from .rollout_processor import RolloutProcessor
 from .types import RolloutProcessorConfig
@@ -25,6 +25,26 @@
 logger = logging.getLogger(__name__)
 
 
+def _build_fireworks_tracing_url(base_url: str, metadata: RolloutMetadata) -> str:
+    """Build a Fireworks tracing URL by appending rollout metadata to the base URL path,
+    allowing the Fireworks tracing proxy to automatically tag traces.
+
+    Format: {base_url}/rollout_id/{id}/invocation_id/{id}/experiment_id/{id}/run_id/{id}/row_id/{id}
+
+    Args:
+        base_url: Fireworks tracing proxy URL (we expect this to be https://tracing.fireworks.ai or
+                  https://tracing.fireworks.ai/project_id/{project_id})
+        metadata: Rollout metadata containing IDs to embed in the URL
+    """
+    return (
+        f"{base_url}/rollout_id/{metadata.rollout_id}"
+        f"/invocation_id/{metadata.invocation_id}"
+        f"/experiment_id/{metadata.experiment_id}"
+        f"/run_id/{metadata.run_id}"
+        f"/row_id/{metadata.row_id}"
+    )
+
+
 def _default_output_data_loader(config: DataLoaderConfig) -> DynamicDataLoader:
     """Default output data loader that fetches traces from Fireworks tracing proxy.
 
@@ -37,7 +57,7 @@ def _default_output_data_loader(config: DataLoaderConfig) -> DynamicDataLoader:
 
     def fetch_traces() -> List[EvaluationRow]:
         base_url = config.model_base_url or "https://tracing.fireworks.ai"
-        adapter = create_fireworks_tracing_adapter(base_url=base_url)
+        adapter = FireworksTracingAdapter(base_url=base_url)
         return adapter.get_evaluation_rows(tags=[f"rollout_id:{config.rollout_id}"], max_retries=5)
 
     return DynamicDataLoader(generators=[fetch_traces], preprocess_fn=filter_longest_conversation)
@@ -169,13 +189,7 @@ async def _process_row(row: EvaluationRow) -> EvaluationRow:
 
             final_model_base_url = model_base_url
             if model_base_url and model_base_url.startswith("https://tracing.fireworks.ai"):
-                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}"
-                )
+                final_model_base_url = _build_fireworks_tracing_url(model_base_url, meta)
 
             init_payload: InitRequest = InitRequest(
                 model=model,

tests/remote_server/test_remote_fireworks.py

@@ -20,7 +20,7 @@
 from eval_protocol.models import EvaluationRow, Message
 from eval_protocol.pytest import evaluation_test
 from eval_protocol.pytest.remote_rollout_processor import RemoteRolloutProcessor
-from eval_protocol.adapters.fireworks_tracing import create_fireworks_tracing_adapter
+from eval_protocol.adapters.fireworks_tracing import FireworksTracingAdapter
 from eval_protocol.quickstart.utils import filter_longest_conversation
 from eval_protocol.types.remote_rollout_processor import DataLoaderConfig
 
@@ -42,7 +42,7 @@ def fetch_fireworks_traces(config: DataLoaderConfig) -> List[EvaluationRow]:
     ROLLOUT_IDS.add(config.rollout_id)
 
     base_url = config.model_base_url or "https://tracing.fireworks.ai"
-    adapter = create_fireworks_tracing_adapter(base_url=base_url)
+    adapter = FireworksTracingAdapter(base_url=base_url)
     return adapter.get_evaluation_rows(tags=[f"rollout_id:{config.rollout_id}"], max_retries=5)