fix(rate-limiter): prevent deepcopy from creating duplicate instances (#89)

* fix(rate-limiter): prevent deepcopy from creating duplicate instances

Implement singleton pattern for RateLimiter to ensure shared instances
are preserved during deepcopy operations. This fixes a bug where
rate limiters were being incorrectly duplicated when generators
were copied or serialized.

- Make RateLimiter frozen to prevent mutation
- Add model_validator to enforce singleton pattern
- Override __deepcopy__ to return same instance
- Add tests to verify rate limiter preservation in with_params and serialization

* chore: code formatting

* chore: improve comments

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* fix(test): unique rate limiter name

* chore: code formatting

* fix(test): update expectation

* fix(test): update expectation

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: kevinmessiaen

src/giskard/agents/rate_limiter.py

@@ -3,15 +3,15 @@
 import uuid
 from contextlib import asynccontextmanager
 
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, Field, PrivateAttr, model_validator
 
 
 class RateLimiterStrategy(BaseModel):
     min_interval: float
     max_concurrent: int = Field(default=5)
 
 
-class RateLimiter(BaseModel):
+class RateLimiter(BaseModel, frozen=True):
     rate_limiter_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
     strategy: RateLimiterStrategy
 
@@ -80,6 +80,28 @@ async def acquire(self) -> None:
     def release(self) -> None:
         self._semaphore.release()
 
+    @model_validator(mode="wrap")
+    @classmethod
+    def _validate_singleton(cls, v, handler, _info) -> "RateLimiter":
+        # Determine the rate limiter ID from the input
+        rate_limiter_id = None
+        if isinstance(v, dict):
+            rate_limiter_id = v.get("rate_limiter_id")
+        elif isinstance(v, RateLimiter):
+            rate_limiter_id = v.rate_limiter_id
+
+        # If ID exists in our global registry, return the singleton
+        if rate_limiter_id and rate_limiter_id in _rate_limiters:
+            return _rate_limiters[rate_limiter_id]
+
+        # 3. Otherwise, proceed with standard validation/creation
+        instance = handler(v)
+        return instance
+
+    def __deepcopy__(self, memo) -> "RateLimiter":
+        # RateLimiter is a shared resource, so we can just return the same instance.
+        return self
+
 
 _rate_limiters: dict[str, RateLimiter] = {}
 

tests/test_generator.py

@@ -5,7 +5,7 @@
 from giskard.agents.chat import Chat, Message
 from giskard.agents.generators.base import GenerationParams, Response
 from giskard.agents.generators.litellm_generator import LiteLLMGenerator
-from giskard.agents.rate_limiter import RateLimiter
+from giskard.agents.rate_limiter import RateLimiter, _rate_limiters
 from giskard.agents.templates import MessageTemplate
 from giskard.agents.workflow import ChatWorkflow
 from litellm import ModelResponse
@@ -119,8 +119,12 @@ async def test_generator_without_rate_limiter(mock_response):
 
 
 async def test_generator_rate_limiter_context():
-    rate_limiter = RateLimiter.from_rpm(rpm=100, rate_limiter_id="test")
-    generator = LiteLLMGenerator(model="test-model", rate_limiter="test")
+    rate_limiter = RateLimiter.from_rpm(
+        rpm=100, rate_limiter_id="test_generator_rate_limiter_context"
+    )
+    generator = LiteLLMGenerator(
+        model="test-model", rate_limiter="test_generator_rate_limiter_context"
+    )
     assert generator.rate_limiter is rate_limiter
 
 
@@ -139,6 +143,62 @@ def test_generator_with_params():
     assert generator.params.response_format is None
 
 
+def test_generator_with_params_and_rate_limiter():
+    """Test that with_params works correctly with a rate limiter."""
+    rate_limiter = RateLimiter.from_rpm(rpm=100, max_concurrent=5)
+    generator = LiteLLMGenerator(model="test-model", rate_limiter=rate_limiter)
+
+    # Verify initial state
+    assert generator.rate_limiter is rate_limiter
+
+    # Call with_params and verify rate limiter is preserved
+    generator_with_params = generator.with_params(temperature=0.5, max_tokens=100)
+    assert generator_with_params.params.temperature == 0.5
+    assert generator_with_params.params.max_tokens == 100
+    # Verify rate limiter is preserved and the same instance
+    assert generator_with_params.rate_limiter is rate_limiter
+
+    # Verify original generator is unchanged
+    assert generator.params.temperature == 1.0  # default value
+    assert generator.params.max_tokens is None
+    assert generator.rate_limiter is rate_limiter
+
+
+def test_generator_serialization_keep_rate_limiter_instance():
+    """Test that serializing and deserializing a generator preserves the rate limiter instance."""
+    rate_limiter = RateLimiter.from_rpm(rpm=100, max_concurrent=5)
+    generator = LiteLLMGenerator(model="test-model", rate_limiter=rate_limiter)
+
+    json_str = generator.model_dump_json()
+    deserialized_generator = LiteLLMGenerator.model_validate_json(json_str)
+
+    assert deserialized_generator.rate_limiter is rate_limiter
+
+
+def test_generator_serialization_recreate_rate_limiter_instance_if_not_in_registry():
+    """Test that deserializing a generator recreates the rate limiter if it's not in the registry."""
+    rate_limiter = RateLimiter.from_rpm(rpm=100, max_concurrent=5)
+    generator = LiteLLMGenerator(model="test-model", rate_limiter=rate_limiter)
+
+    json_str = generator.model_dump_json()
+    del _rate_limiters[rate_limiter.rate_limiter_id]
+    deserialized_generator = LiteLLMGenerator.model_validate_json(json_str)
+
+    assert deserialized_generator.rate_limiter is not rate_limiter
+    assert (
+        deserialized_generator.rate_limiter.rate_limiter_id
+        == rate_limiter.rate_limiter_id
+    )
+    assert (
+        deserialized_generator.rate_limiter.strategy.min_interval
+        == rate_limiter.strategy.min_interval
+    )
+    assert (
+        deserialized_generator.rate_limiter.strategy.max_concurrent
+        == rate_limiter.strategy.max_concurrent
+    )
+
+
 async def test_generator_with_params_overwrite(mock_response):
     # ARRANGE: Create a generator with base parameters.
     generator = LiteLLMGenerator(model="test-model").with_params(

tests/test_rate_limiter.py

@@ -121,9 +121,9 @@ async def test_rate_limiter_max_concurrent():
 
 def test_can_serialize_rate_limiter():
     rate_limiter = RateLimiter.from_rpm(
-        rpm=600, max_concurrent=10, rate_limiter_id="test"
+        rpm=600, max_concurrent=10, rate_limiter_id="test_can_serialize_rate_limiter"
     )
     assert rate_limiter.model_dump() == {
-        "rate_limiter_id": "test",
+        "rate_limiter_id": "test_can_serialize_rate_limiter",
         "strategy": {"min_interval": 0.1, "max_concurrent": 10},
     }