diff --git a/dspy/clients/cache.py b/dspy/clients/cache.py index 187043538e..cca435ae36 100644 --- a/dspy/clients/cache.py +++ b/dspy/clients/cache.py @@ -168,6 +168,35 @@ def load_memory_cache(self, filepath: str) -> None: with open(filepath, "rb") as f: self.memory_cache = cloudpickle.load(f) + def close(self) -> None: + """Close the cache and release all resources.""" + if self.enable_disk_cache and hasattr(self.disk_cache, "close"): + try: + self.disk_cache.close() + except Exception as e: + logger.debug(f"Failed to close disk cache: {e}") + + if self.enable_memory_cache and hasattr(self.memory_cache, "clear"): + with self._lock: + self.memory_cache.clear() + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit - ensures cleanup.""" + self.close() + return False + + def __del__(self): + """Destructor to ensure cleanup on garbage collection.""" + try: + self.close() + except Exception: + # Ignore errors during garbage collection + pass + def request_cache( cache_arg_name: str | None = None, diff --git a/tests/clients/test_cache.py b/tests/clients/test_cache.py index 42a1213957..a3a09505ba 100644 --- a/tests/clients/test_cache.py +++ b/tests/clients/test_cache.py @@ -30,44 +30,47 @@ def cache_config(tmp_path): @pytest.fixture def cache(cache_config): """Create a cache instance with the default configuration.""" - return Cache(**cache_config) + cache_instance = Cache(**cache_config) + yield cache_instance + # Cleanup + cache_instance.close() def test_initialization(tmp_path): """Test different cache initialization configurations.""" # Test memory-only cache - memory_cache = Cache( + with Cache( enable_disk_cache=False, enable_memory_cache=True, disk_cache_dir="", disk_size_limit_bytes=0, memory_max_entries=50, - ) - assert isinstance(memory_cache.memory_cache, LRUCache) - assert memory_cache.memory_cache.maxsize == 50 - assert memory_cache.disk_cache == {} + ) as memory_cache: + assert isinstance(memory_cache.memory_cache, LRUCache) + assert memory_cache.memory_cache.maxsize == 50 + assert memory_cache.disk_cache == {} # Test disk-only cache - disk_cache = Cache( + with Cache( enable_disk_cache=True, enable_memory_cache=False, disk_cache_dir=str(tmp_path), disk_size_limit_bytes=1024, memory_max_entries=0, - ) - assert isinstance(disk_cache.disk_cache, FanoutCache) - assert disk_cache.memory_cache == {} + ) as disk_cache: + assert isinstance(disk_cache.disk_cache, FanoutCache) + assert disk_cache.memory_cache == {} # Test disabled cache - disabled_cache = Cache( + with Cache( enable_disk_cache=False, enable_memory_cache=False, disk_cache_dir="", disk_size_limit_bytes=0, memory_max_entries=0, - ) - assert disabled_cache.memory_cache == {} - assert disabled_cache.disk_cache == {} + ) as disabled_cache: + assert disabled_cache.memory_cache == {} + assert disabled_cache.disk_cache == {} def test_cache_key_generation(cache): @@ -180,22 +183,21 @@ def test_save_and_load_memory_cache(cache, tmp_path): cache.save_memory_cache(str(temp_cache_file)) # Create a new cache instance with disk cache disabled - new_cache = Cache( + with Cache( enable_memory_cache=True, enable_disk_cache=False, disk_cache_dir=tmp_path / "disk_cache", disk_size_limit_bytes=0, memory_max_entries=100, - ) - - # Load the memory cache - new_cache.load_memory_cache(str(temp_cache_file)) - - # Verify items are in the new memory cache - for req in requests: - result = new_cache.get(req) - assert result is not None - assert result == f"Response {requests.index(req)}" + ) as new_cache: + # Load the memory cache + new_cache.load_memory_cache(str(temp_cache_file)) + + # Verify items are in the new memory cache + for req in requests: + result = new_cache.get(req) + assert result is not None + assert result == f"Response {requests.index(req)}" def test_request_cache_decorator(cache): diff --git a/tests/clients/test_lm.py b/tests/clients/test_lm.py index 491c3c5940..9d7b586dbf 100644 --- a/tests/clients/test_lm.py +++ b/tests/clients/test_lm.py @@ -24,7 +24,7 @@ def make_response(output_blocks): model="openai/dspy-test-model", object="response", output=output_blocks, - metadata = {}, + metadata={}, parallel_tool_calls=False, temperature=1.0, tool_choice="auto", @@ -92,6 +92,8 @@ def test_dspy_cache(litellm_test_server, tmp_path): assert len(usage_tracker.usage_data) == 0 + # Cleanup + cache.close() dspy.cache = original_cache @@ -213,6 +215,7 @@ def test_reasoning_model_token_parameter(): assert "max_tokens" in lm.kwargs assert lm.kwargs["max_tokens"] == 1000 + @pytest.mark.parametrize("model_name", ["openai/o1", "openai/gpt-5-nano"]) def test_reasoning_model_requirements(model_name): # Should raise assertion error if temperature or max_tokens requirements not met @@ -369,6 +372,8 @@ async def test_async_lm_call_with_cache(tmp_path): assert len(cache.memory_cache) == 2 assert mock_alitellm_completion.call_count == 2 + # Cleanup + cache.close() dspy.cache = original_cache @@ -407,6 +412,7 @@ def test_disable_history(): model="openai/gpt-4o-mini", ) + def test_responses_api(litellm_test_server): api_base, _ = litellm_test_server expected_text = "This is a test answer from responses API." @@ -418,9 +424,7 @@ def test_responses_api(litellm_test_server): "type": "message", "role": "assistant", "status": "completed", - "content": [ - {"type": "output_text", "text": expected_text, "annotations": []} - ], + "content": [{"type": "output_text", "text": expected_text, "annotations": []}], } ] )