From 106cab9c26723f50701fef85a14ab520278b3209 Mon Sep 17 00:00:00 2001
From: John Horton <john.joseph.horton@gmail.com>
Date: Wed, 18 Sep 2024 07:31:31 -0400
Subject: [PATCH] Buckets defined at service level, not model level

---
 edsl/jobs/buckets/BucketCollection.py | 39 ++++++++++++++++-----------
 tests/jobs/test_BucketCollection.py   | 14 ++++++++++
 2 files changed, 38 insertions(+), 15 deletions(-)
 create mode 100644 tests/jobs/test_BucketCollection.py

diff --git a/edsl/jobs/buckets/BucketCollection.py b/edsl/jobs/buckets/BucketCollection.py
index 7837580d..28c2408e 100644
--- a/edsl/jobs/buckets/BucketCollection.py
+++ b/edsl/jobs/buckets/BucketCollection.py
@@ -13,6 +13,8 @@ class BucketCollection(UserDict):
     def __init__(self, infinity_buckets=False):
         super().__init__()
         self.infinity_buckets = infinity_buckets
+        self.models_to_services = {}
+        self.services_to_buckets = {}
 
     def __repr__(self):
         return f"BucketCollection({self.data})"
@@ -21,6 +23,7 @@ def add_model(self, model: "LanguageModel") -> None:
         """Adds a model to the bucket collection.
 
         This will create the token and request buckets for the model."""
+
         # compute the TPS and RPS from the model
         if not self.infinity_buckets:
             TPS = model.TPM / 60.0
@@ -29,22 +32,28 @@ def add_model(self, model: "LanguageModel") -> None:
             TPS = float("inf")
             RPS = float("inf")
 
-        # create the buckets
-        requests_bucket = TokenBucket(
-            bucket_name=model.model,
-            bucket_type="requests",
-            capacity=RPS,
-            refill_rate=RPS,
-        )
-        tokens_bucket = TokenBucket(
-            bucket_name=model.model, bucket_type="tokens", capacity=TPS, refill_rate=TPS
-        )
-        model_buckets = ModelBuckets(requests_bucket, tokens_bucket)
-        if model in self:
-            # it if already exists, combine the buckets
-            self[model] += model_buckets
+        if model.model not in self.models_to_services:
+            service = model._inference_service_
+            if service not in self.services_to_buckets:
+                requests_bucket = TokenBucket(
+                    bucket_name=service,
+                    bucket_type="requests",
+                    capacity=RPS,
+                    refill_rate=RPS,
+                )
+                tokens_bucket = TokenBucket(
+                    bucket_name=service,
+                    bucket_type="tokens",
+                    capacity=TPS,
+                    refill_rate=TPS,
+                )
+                self.services_to_buckets[service] = ModelBuckets(
+                    requests_bucket, tokens_bucket
+                )
+            self.models_to_services[model.model] = service
+            self[model] = self.services_to_buckets[service]
         else:
-            self[model] = model_buckets
+            self[model] = self.services_to_buckets[self.models_to_services[model.model]]
 
     def visualize(self) -> dict:
         """Visualize the token and request buckets for each model."""
diff --git a/tests/jobs/test_BucketCollection.py b/tests/jobs/test_BucketCollection.py
new file mode 100644
index 00000000..e3cb1794
--- /dev/null
+++ b/tests/jobs/test_BucketCollection.py
@@ -0,0 +1,14 @@
+import pytest
+
+from edsl import Model, QuestionFreeText
+
+
+def test_one_per_service():
+    models = [Model(temperature=1), Model(temperature=2), Model(temperature=0)]
+    q = QuestionFreeText(
+        question_text="What is your favorite color?", question_name="color"
+    )
+    jobs = q.by(models)
+    bc = jobs.bucket_collection
+    assert len(bc) == 3
+    assert len(set(bc.values())) == 1