parameterize logger

Author: Dylan Huang

eval_protocol/pytest/default_agent_rollout_processor.py

@@ -8,7 +8,7 @@
 from openai.types.chat import ChatCompletionContentPartTextParam, ChatCompletionMessage, ChatCompletionToolParam
 from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
 
-from eval_protocol.dataset_logger import default_logger
+from eval_protocol.dataset_logger.dataset_logger import DatasetLogger
 from eval_protocol.mcp.execution.policy import LiteLLMPolicy
 from eval_protocol.mcp.mcp_multi_client import MCPMultiClient
 from eval_protocol.models import EvaluationRow, Message
@@ -20,12 +20,13 @@ class Agent:
     A really simple agent that calls the model until no more tool calls are needed.
     """
 
-    def __init__(self, model: str, row: EvaluationRow, config_path: str):
+    def __init__(self, model: str, row: EvaluationRow, config_path: str, logger: DatasetLogger):
         self.model = model
         self.evaluation_row: EvaluationRow = row
         self._policy = LiteLLMPolicy(model_id=model)
         self.mcp_client = MCPMultiClient(config_path=config_path) if config_path else None
         self.tools: Union[List[ChatCompletionToolParam], NotGiven] = NOT_GIVEN
+        self.logger: DatasetLogger = logger
 
     async def setup(self):
         if self.mcp_client:
@@ -40,9 +41,9 @@ async def _get_tools(self) -> Optional[List[ChatCompletionToolParam]]:
     def messages(self) -> list[Message]:
         return self.evaluation_row.messages
 
-    def append_message_and_log(self, message: Message):
+    def append_message_and_log(self, message: Message, logger: DatasetLogger):
         self.messages.append(message)
-        default_logger.log(self.evaluation_row)
+        logger.log(self.evaluation_row)
 
     async def call_agent(self) -> str:
         """
@@ -77,7 +78,8 @@ async def call_agent(self) -> str:
                             ChatCompletionContentPartTextParam(text=content.text, type="text") for content in content
                         ],
                         tool_call_id=tool_call_id,
-                    )
+                    ),
+                    logger=self.logger,
                 )
             return await self.call_agent()
         return message.content
@@ -116,7 +118,7 @@ async def default_agent_rollout_processor(
 ) -> List[EvaluationRow]:
     dataset: Dataset = []
     for row in rows:
-        agent = Agent(model=config.model, row=row, config_path=config.mcp_config_path)
+        agent = Agent(model=config.model, row=row, config_path=config.mcp_config_path, logger=config.logger)
         await agent.setup()
         await agent.call_agent()
         dataset.append(agent.evaluation_row)

eval_protocol/pytest/default_single_turn_rollout_process.py

@@ -1,11 +1,9 @@
 import asyncio
-from typing import List
-
 import logging
 import os
+from typing import List
 
-from eval_protocol.dataset_logger import default_logger
-from eval_protocol.models import EvaluationRow, Message, ChatCompletionMessageToolCall
+from eval_protocol.models import ChatCompletionMessageToolCall, EvaluationRow, Message
 from eval_protocol.pytest.types import RolloutProcessorConfig
 
 
@@ -49,6 +47,7 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
 
         # Dynamic import to avoid static dependency/lint errors if LiteLLM isn't installed yet
         import importlib
+
         _litellm = importlib.import_module("litellm")
         acompletion = getattr(_litellm, "acompletion")
         response = await acompletion(**request_params)
@@ -79,7 +78,7 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
         ]
 
         row.messages = messages
-        default_logger.log(row)
+        config.logger.log(row)
         return row
 
     # Process rows with bounded concurrency if configured

eval_protocol/pytest/evaluation_test.py

@@ -8,6 +8,7 @@
 import pytest
 
 from eval_protocol.dataset_logger import default_logger
+from eval_protocol.dataset_logger.dataset_logger import DatasetLogger
 from eval_protocol.human_id import generate_id
 from eval_protocol.models import CompletionParams, EvalMetadata, EvaluationRow, InputMetadata, Message
 from eval_protocol.pytest.default_dataset_adapter import default_dataset_adapter
@@ -55,6 +56,7 @@ def evaluation_test(  # noqa: C901
     steps: int = 30,
     mode: EvaluationTestMode = "batch",
     combine_datasets: bool = True,
+    logger: Optional[DatasetLogger] = None,
 ) -> Callable[
     [TestFunction],
     TestFunction,
@@ -117,8 +119,11 @@ def evaluation_test(  # noqa: C901
         mode: Evaluation mode. "batch" (default) expects test function to handle
             full dataset. "pointwise" applies test function to each row. If your evaluation requires
             the full rollout of all rows to compute the score, use
+        logger: DatasetLogger to use for logging. If not provided, a default logger will be used.
     """
 
+    active_logger: DatasetLogger = logger if logger else default_logger
+
     def decorator(
         test_func: TestFunction,
     ):
@@ -287,7 +292,7 @@ def wrapper_body(**kwargs):
                 def _log_eval_error(
                     status: Literal["finished", "error"], rows: Optional[List[EvaluationRow]] | None, passed: bool
                 ) -> None:
-                    log_eval_status_and_rows(eval_metadata, rows, status, passed, default_logger)
+                    log_eval_status_and_rows(eval_metadata, rows, status, passed, active_logger)
 
                 try:
                     # Handle dataset loading
@@ -369,7 +374,7 @@ def _log_eval_error(
                         # has to be done in the pytest main process since it's
                         # used to determine whether this eval has stopped
                         row.pid = os.getpid()
-                        default_logger.log(row)
+                        active_logger.log(row)
 
                     # Prepare rollout processor config once; we will generate fresh outputs per run
                     config = RolloutProcessorConfig(
@@ -379,6 +384,7 @@ def _log_eval_error(
                         max_concurrent_rollouts=max_concurrent_rollouts,
                         server_script_path=server_script_path,
                         steps=steps,
+                        logger=active_logger,
                     )
 
                     for _ in range(num_runs):
@@ -463,7 +469,7 @@ def _log_eval_error(
                         if r.eval_metadata is not None:
                             r.eval_metadata.status = "finished"
                             r.eval_metadata.passed = passed
-                        default_logger.log(r)
+                        active_logger.log(r)
 
                     # Optional: print and/or persist a summary artifact for CI
                     try:

eval_protocol/pytest/types.py

@@ -5,6 +5,9 @@
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Literal, Optional
 
+from eval_protocol.dataset_logger import default_logger
+from eval_protocol.dataset_logger.dataset_logger import DatasetLogger
+
 from ..models import EvaluationRow, Message
 
 ModelParam = str  # gpt-4o, gpt-4o-mini, accounts/fireworks/models/llama-3.1-8b-instruct
@@ -39,10 +42,13 @@
 class RolloutProcessorConfig:
     model: ModelParam
     input_params: RolloutInputParam  # optional input parameters for inference
-    mcp_config_path: str  
-    server_script_path: Optional[str] = None  # TODO: change from server_script_path to mcp_config_path for agent rollout processor
+    mcp_config_path: str
+    server_script_path: Optional[str] = (
+        None  # TODO: change from server_script_path to mcp_config_path for agent rollout processor
+    )
     max_concurrent_rollouts: int = 8  # maximum number of concurrent rollouts
     steps: int = 30  # max number of rollout steps
+    logger: DatasetLogger = default_logger  # logger to use during rollout for mid-rollout logs
 
 
 RolloutProcessor = Callable[[List[EvaluationRow], RolloutProcessorConfig], List[EvaluationRow]]