diff --git a/haystack/hooks/compaction/hooks.py b/haystack/hooks/compaction/hooks.py index cc1cc124a4d..d6aafecd41a 100644 --- a/haystack/hooks/compaction/hooks.py +++ b/haystack/hooks/compaction/hooks.py @@ -290,24 +290,24 @@ async def warm_up_async(self) -> None: """Warm up the token counter and the compactor on the serving event loop.""" if hasattr(self.token_counter, "warm_up"): self.token_counter.warm_up() - warm_up_async = getattr(self.compactor, "warm_up_async", None) - if warm_up_async is not None: - await warm_up_async() + if hasattr(self.compactor, "warm_up_async"): + await self.compactor.warm_up_async() elif hasattr(self.compactor, "warm_up"): self.compactor.warm_up() def close(self) -> None: - """Release the compactor's resources.""" - if hasattr(self.compactor, "close"): - self.compactor.close() + """Release the token counter's and the compactor's resources.""" + for resource in (self.token_counter, self.compactor): + if hasattr(resource, "close"): + resource.close() async def close_async(self) -> None: - """Release the compactor's async resources.""" - close_async = getattr(self.compactor, "close_async", None) - if close_async is not None: - await close_async() - elif hasattr(self.compactor, "close"): - self.compactor.close() + """Release the token counter's and the compactor's async resources.""" + for resource in (self.token_counter, self.compactor): + if hasattr(resource, "close_async"): + await resource.close_async() + elif hasattr(resource, "close"): + resource.close() def to_dict(self) -> dict[str, Any]: """ diff --git a/releasenotes/notes/compaction-hook-close-token-counter-0495043806039056.yaml b/releasenotes/notes/compaction-hook-close-token-counter-0495043806039056.yaml new file mode 100644 index 00000000000..81ba3bc8cc6 --- /dev/null +++ b/releasenotes/notes/compaction-hook-close-token-counter-0495043806039056.yaml @@ -0,0 +1,6 @@ +--- +fixes: + - | + Fixed ``CompactionHook.close()`` and ``close_async()`` to release the token counter's + resources as well as the compactor's. Previously, resources such as + ``OpenAITokenCounter``'s HTTP client remained open after the hook or Agent was closed. diff --git a/test/hooks/compaction/test_hooks.py b/test/hooks/compaction/test_hooks.py index e32f6458068..1bb8f8d5ee2 100644 --- a/test/hooks/compaction/test_hooks.py +++ b/test/hooks/compaction/test_hooks.py @@ -4,6 +4,7 @@ import logging from typing import Annotated, Any +from unittest.mock import Mock import pytest @@ -387,11 +388,14 @@ def test_chains_tool_result_pruning_before_sliding_window(self): ) assert compacted[-2:] == messages[-2:] - def test_lifecycle_delegates_to_the_compactor(self): + def test_lifecycle_delegates_to_the_counter_and_compactor(self): + counter = Mock(spec=["warm_up", "close"]) compactor = _RecordingCompactor() - hook = _hook(compactor) + hook = _hook(compactor=compactor, token_counter=counter) hook.warm_up() hook.close() + counter.warm_up.assert_called_once_with() + counter.close.assert_called_once_with() assert compactor.calls == ["warm_up", "close"] @@ -422,11 +426,14 @@ async def test_run_async_uses_the_async_compaction_path(self): assert compactor.calls == ["compact_async"] @pytest.mark.asyncio - async def test_lifecycle_prefers_the_async_methods(self): + async def test_lifecycle_prefers_async_methods_with_sync_fallback(self): + counter = Mock(spec=["warm_up", "close"]) compactor = _RecordingCompactor() - hook = _hook(compactor) + hook = _hook(compactor=compactor, token_counter=counter) await hook.warm_up_async() await hook.close_async() + counter.warm_up.assert_called_once_with() + counter.close.assert_called_once_with() assert compactor.calls == ["warm_up_async", "close_async"] @pytest.mark.asyncio