Fix MongoKeyValueStore.remove() (#2582)

stanford-crfm · Jun 21, 2024 · 68cbaed · 68cbaed
1 parent 3d1a817
commit 68cbaed
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/helm/common/mongo_key_value_store.py b/src/helm/common/mongo_key_value_store.py
@@ -85,4 +85,5 @@ def multi_put(self, pairs: Iterable[Tuple[Dict, Dict]]) -> None:
         self._collection.bulk_write(operations)
 
     def remove(self, key: Dict) -> None:
-        self._collection.delete_one(key)
+        query = {self._REQUEST_KEY: self._canonicalize_key(key)}
+        self._collection.delete_one(query)