Tests for keylook up changes being manifested in token buckets

edsl/jobs/Jobs.py

@@ -408,7 +408,13 @@ def create_bucket_collection(self) -> BucketCollection:
         >>> bc
         BucketCollection(...)
         """
-        return BucketCollection.from_models(self.models)
+        bc = BucketCollection.from_models(self.models)
+
+        if self.run_config.environment.key_lookup is not None:
+            bc.update_from_key_lookup(
+                self.run_config.environment.key_lookup
+            )
+        return bc
 
     def html(self):
         """Return the HTML representations for each scenario"""

tests/jobs/test_KeyLookup_Modify_BucketCollection.py

@@ -0,0 +1,36 @@
+import pytest
+from edsl import Model 
+from edsl import QuestionFreeText
+from edsl.language_models.key_management.KeyLookup import KeyLookup
+
+def test_default_rpm_bucket_collection():
+    # Create model with default RPM
+    m = Model('test')
+    jobs = QuestionFreeText.example().by(m)
+    bc = jobs.create_bucket_collection()
+    rr = bc[Model(model_name='test', temperature=0.5)].requests_bucket.refill_rate
+
+    # Check if actual refill rate is within 10 requests/min of target RPM
+    actual_rpm = rr * 60
+    assert abs(actual_rpm - m.rpm) < 10, \
+        f"Actual RPM ({actual_rpm}) differs from target RPM ({m.rpm}) by more than 10"
+
+def test_custom_rpm_bucket_collection():
+    # Setup custom RPM via KeyLookup
+    kl = KeyLookup.example()
+    target_rpm = 1
+    kl['test'].rpm = target_rpm
+
+    # Create jobs with custom KeyLookup
+    m = Model('test')
+    jobs = QuestionFreeText.example().by(m)
+    jobs2 = jobs.using(kl)
+
+    # Create bucket collection and get refill rate
+    bc = jobs2.create_bucket_collection()
+    rr = bc[Model(model_name='test', temperature=0.5)].requests_bucket.refill_rate
+
+    # Check if actual refill rate is within 1 request/min of target RPM
+    actual_rpm = rr * 60
+    assert abs(actual_rpm - target_rpm) < 1, \
+        f"Actual RPM ({actual_rpm}) differs from target RPM ({target_rpm}) by more than 1"