fix sqlite logger tests

Author: Dylan Huang

eval_protocol/dataset_logger/sqlite_dataset_logger_adapter.py

@@ -22,9 +22,8 @@ def __init__(self, db_path: Optional[str] = None, store: Optional[SqliteEvaluati
             self._store = SqliteEvaluationRowStore(self.db_path)
 
     def log(self, row: "EvaluationRow") -> None:
-        rollout_id = row.rollout_id
         data = row.model_dump(exclude_none=True, mode="json")
-        self._store.upsert_row(rollout_id=rollout_id, data=data)
+        self._store.upsert_row(data=data)
         try:
             event_bus.emit(LOG_EVENT_TYPE, EvaluationRow(**data))
         except Exception as e:

eval_protocol/dataset_logger/sqlite_evaluation_row_store.py

@@ -36,7 +36,10 @@ class EvaluationRow(BaseModel):  # type: ignore
     def db_path(self) -> str:
         return self._db_path
 
-    def upsert_row(self, rollout_id: str, data: dict) -> None:
+    def upsert_row(self, data: dict) -> None:
+        rollout_id = data["rollout_id"]
+        if "rollout_id" not in data:
+            raise ValueError("rollout_id is required to upsert a row")
         if self._EvaluationRow.select().where(self._EvaluationRow.rollout_id == rollout_id).exists():
             self._EvaluationRow.update(data=data).where(self._EvaluationRow.rollout_id == rollout_id).execute()
         else:

tests/dataset_logger/test_sqlite_dataset_logger_adapter.py

@@ -18,13 +18,13 @@ def test_update_log_and_read():
     messages = [Message(role="user", content="Hello")]
     input_metadata = InputMetadata(row_id="1")
     row = EvaluationRow(input_metadata=input_metadata, messages=messages)
-    store.upsert_row(row_id="1", data=row.model_dump(exclude_none=True, mode="json"))
+    store.upsert_row(data=row.model_dump(exclude_none=True, mode="json"))
 
     row.messages.append(Message(role="assistant", content="Hello"))
 
-    logger = SqliteDatasetLoggerAdapter()
+    logger = SqliteDatasetLoggerAdapter(store=store)
     logger.log(row)
-    saved = logger.read(rollout_id="1")[0]
+    saved = logger.read(row.rollout_id)[0]
     assert row.messages == saved.messages
     assert row.input_metadata == saved.input_metadata
 
@@ -42,7 +42,7 @@ def test_create_log_and_read():
     row = EvaluationRow(input_metadata=input_metadata, messages=messages)
 
     logger.log(row)
-    saved = logger.read(rollout_id="1")[0]
+    saved = logger.read(rollout_id=row.rollout_id)[0]
     assert row.messages == saved.messages
     assert row.input_metadata == saved.input_metadata