Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 12 additions & 12 deletions haystack/hooks/compaction/hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
"""
Expand Down
Original file line number Diff line number Diff line change
@@ -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.
15 changes: 11 additions & 4 deletions test/hooks/compaction/test_hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import logging
from typing import Annotated, Any
from unittest.mock import Mock

import pytest

Expand Down Expand Up @@ -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"]


Expand Down Expand Up @@ -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
Expand Down
Loading