Spaces:
Build error
Build error
| """Tests for Agno hooks integration. | |
| Tests cover: | |
| 1. HeadroomPreHook - Pre-hook for tracking before LLM calls | |
| 2. HeadroomPostHook - Post-hook for tracking after LLM calls | |
| 3. create_headroom_hooks - Convenience function | |
| """ | |
| from datetime import datetime | |
| from unittest.mock import MagicMock | |
| import pytest | |
| # Check if Agno is available | |
| try: | |
| import agno # noqa: F401 | |
| AGNO_AVAILABLE = True | |
| except ImportError: | |
| AGNO_AVAILABLE = False | |
| from headroom import HeadroomConfig, HeadroomMode | |
| # Skip all tests if Agno not installed | |
| pytestmark = pytest.mark.skipif(not AGNO_AVAILABLE, reason="Agno not installed") | |
| class TestHeadroomPreHook: | |
| """Tests for HeadroomPreHook.""" | |
| def test_init_defaults(self): | |
| """Initialize with default settings.""" | |
| from headroom.integrations.agno import HeadroomPreHook | |
| hook = HeadroomPreHook() | |
| assert hook.mode == HeadroomMode.OPTIMIZE | |
| assert hook.model == "gpt-4o" | |
| assert hook.total_tokens_saved == 0 | |
| assert hook.metrics_history == [] | |
| def test_init_with_custom_config(self): | |
| """Initialize with custom config.""" | |
| from headroom.integrations.agno import HeadroomPreHook | |
| config = HeadroomConfig(default_mode=HeadroomMode.AUDIT) | |
| hook = HeadroomPreHook( | |
| config=config, | |
| mode=HeadroomMode.SIMULATE, | |
| model="claude-3-5-sonnet-20241022", | |
| ) | |
| assert hook.config is config | |
| assert hook.mode == HeadroomMode.SIMULATE | |
| assert hook.model == "claude-3-5-sonnet-20241022" | |
| def test_call_returns_input_unchanged(self): | |
| """Hook returns input unchanged (optimization at model level).""" | |
| from headroom.integrations.agno import HeadroomPreHook | |
| hook = HeadroomPreHook() | |
| run_input = "Hello, how are you?" | |
| result = hook(run_input) | |
| assert result == run_input | |
| def test_call_tracks_metrics(self): | |
| """Hook tracks metrics on each call.""" | |
| from headroom.integrations.agno import HeadroomPreHook | |
| hook = HeadroomPreHook() | |
| hook("First input") | |
| hook("Second input") | |
| assert len(hook.metrics_history) == 2 | |
| assert all(m.request_id for m in hook.metrics_history) | |
| assert all(isinstance(m.timestamp, datetime) for m in hook.metrics_history) | |
| def test_metrics_history_limited(self): | |
| """Metrics history is limited to 100 entries.""" | |
| from headroom.integrations.agno import HeadroomPreHook | |
| hook = HeadroomPreHook() | |
| # Call 150 times | |
| for i in range(150): | |
| hook(f"Input {i}") | |
| assert len(hook.metrics_history) == 100 | |
| def test_get_savings_summary_empty(self): | |
| """get_savings_summary with no history.""" | |
| from headroom.integrations.agno import HeadroomPreHook | |
| hook = HeadroomPreHook() | |
| summary = hook.get_savings_summary() | |
| assert summary["total_requests"] == 0 | |
| assert summary["total_tokens_saved"] == 0 | |
| assert summary["average_savings_percent"] == 0 | |
| def test_get_savings_summary_with_data(self): | |
| """get_savings_summary with metrics.""" | |
| from headroom.integrations.agno import HeadroomPreHook | |
| hook = HeadroomPreHook() | |
| # Make some calls | |
| hook("Input 1") | |
| hook("Input 2") | |
| summary = hook.get_savings_summary() | |
| assert summary["total_requests"] == 2 | |
| # Pre-hook doesn't do actual optimization, so tokens_saved is 0 | |
| assert summary["total_tokens_saved"] == 0 | |
| class TestHeadroomPostHook: | |
| """Tests for HeadroomPostHook.""" | |
| def test_init_defaults(self): | |
| """Initialize with default settings.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook() | |
| assert hook.log_level == "INFO" | |
| assert hook.token_alert_threshold is None | |
| assert hook.total_requests == 0 | |
| assert hook.alerts == [] | |
| def test_init_with_threshold(self): | |
| """Initialize with alert threshold.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook( | |
| log_level="DEBUG", | |
| token_alert_threshold=10000, | |
| ) | |
| assert hook.log_level == "DEBUG" | |
| assert hook.token_alert_threshold == 10000 | |
| def test_call_returns_output_unchanged(self): | |
| """Hook returns output unchanged.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook() | |
| output = MagicMock() | |
| output.content = "Hello!" | |
| result = hook(output) | |
| assert result is output | |
| def test_call_tracks_requests(self): | |
| """Hook tracks requests on each call.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook() | |
| output1 = MagicMock() | |
| output1.content = "First response" | |
| output2 = MagicMock() | |
| output2.content = "Second response" | |
| hook(output1) | |
| hook(output2) | |
| assert hook.total_requests == 2 | |
| def test_call_extracts_token_metrics(self): | |
| """Hook extracts token metrics from response.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook() | |
| output = MagicMock() | |
| output.content = "Response" | |
| output.metrics = MagicMock() | |
| output.metrics.input_tokens = 50 | |
| output.metrics.output_tokens = 20 | |
| output.metrics.total_tokens = 70 | |
| hook(output) | |
| assert hook._requests[0]["input_tokens"] == 50 | |
| assert hook._requests[0]["output_tokens"] == 20 | |
| assert hook._requests[0]["total_tokens"] == 70 | |
| def test_call_triggers_alert(self): | |
| """Hook triggers alert when threshold exceeded.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook(token_alert_threshold=50) | |
| output = MagicMock() | |
| output.content = "Response" | |
| output.metrics = MagicMock() | |
| output.metrics.total_tokens = 100 # Exceeds threshold | |
| hook(output) | |
| assert len(hook.alerts) == 1 | |
| assert "Token alert" in hook.alerts[0] | |
| assert "100" in hook.alerts[0] | |
| def test_call_no_alert_below_threshold(self): | |
| """No alert when tokens below threshold.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook(token_alert_threshold=100) | |
| output = MagicMock() | |
| output.content = "Response" | |
| output.metrics = MagicMock() | |
| output.metrics.total_tokens = 50 # Below threshold | |
| hook(output) | |
| assert len(hook.alerts) == 0 | |
| def test_requests_limited(self): | |
| """Request history is limited to 1000 entries.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook() | |
| # Call 1500 times | |
| for i in range(1500): | |
| output = MagicMock() | |
| output.content = f"Response {i}" | |
| hook(output) | |
| assert len(hook._requests) == 1000 | |
| def test_get_summary_empty(self): | |
| """get_summary with no requests.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook() | |
| summary = hook.get_summary() | |
| assert summary["total_requests"] == 0 | |
| assert summary["total_tokens"] == 0 | |
| assert summary["alerts"] == 0 | |
| def test_get_summary_with_data(self): | |
| """get_summary with requests.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook() | |
| # Add some requests directly | |
| hook._requests = [ | |
| {"total_tokens": 100}, | |
| {"total_tokens": 200}, | |
| {"total_tokens": 50}, | |
| ] | |
| summary = hook.get_summary() | |
| assert summary["total_requests"] == 3 | |
| assert summary["total_tokens"] == 350 | |
| assert summary["average_tokens"] == 350 / 3 | |
| def test_reset(self): | |
| """reset() clears all state.""" | |
| from headroom.integrations.agno import HeadroomPostHook | |
| hook = HeadroomPostHook() | |
| # Add some state | |
| hook._requests = [{"test": 1}] | |
| hook._alerts = ["alert"] | |
| hook.reset() | |
| assert hook._requests == [] | |
| assert hook._alerts == [] | |
| class TestCreateHeadroomHooks: | |
| """Tests for create_headroom_hooks convenience function.""" | |
| def test_returns_tuple(self): | |
| """Returns tuple of (pre_hook, post_hook).""" | |
| from headroom.integrations.agno import ( | |
| HeadroomPostHook, | |
| HeadroomPreHook, | |
| create_headroom_hooks, | |
| ) | |
| pre_hook, post_hook = create_headroom_hooks() | |
| assert isinstance(pre_hook, HeadroomPreHook) | |
| assert isinstance(post_hook, HeadroomPostHook) | |
| def test_passes_config_to_pre_hook(self): | |
| """Passes config to pre_hook.""" | |
| from headroom.integrations.agno import create_headroom_hooks | |
| config = HeadroomConfig(default_mode=HeadroomMode.AUDIT) | |
| pre_hook, _ = create_headroom_hooks(config=config) | |
| assert pre_hook.config is config | |
| def test_passes_mode_to_pre_hook(self): | |
| """Passes mode to pre_hook.""" | |
| from headroom.integrations.agno import create_headroom_hooks | |
| pre_hook, _ = create_headroom_hooks(mode=HeadroomMode.SIMULATE) | |
| assert pre_hook.mode == HeadroomMode.SIMULATE | |
| def test_passes_model_to_pre_hook(self): | |
| """Passes model to pre_hook.""" | |
| from headroom.integrations.agno import create_headroom_hooks | |
| pre_hook, _ = create_headroom_hooks(model="claude-3-5-sonnet-20241022") | |
| assert pre_hook.model == "claude-3-5-sonnet-20241022" | |
| def test_passes_log_level_to_post_hook(self): | |
| """Passes log_level to post_hook.""" | |
| from headroom.integrations.agno import create_headroom_hooks | |
| _, post_hook = create_headroom_hooks(log_level="DEBUG") | |
| assert post_hook.log_level == "DEBUG" | |
| def test_passes_threshold_to_post_hook(self): | |
| """Passes token_alert_threshold to post_hook.""" | |
| from headroom.integrations.agno import create_headroom_hooks | |
| _, post_hook = create_headroom_hooks(token_alert_threshold=5000) | |
| assert post_hook.token_alert_threshold == 5000 | |
| def test_all_parameters(self): | |
| """Test with all parameters.""" | |
| from headroom.integrations.agno import create_headroom_hooks | |
| config = HeadroomConfig() | |
| pre_hook, post_hook = create_headroom_hooks( | |
| config=config, | |
| mode=HeadroomMode.AUDIT, | |
| model="gpt-4-turbo", | |
| log_level="WARNING", | |
| token_alert_threshold=8000, | |
| ) | |
| assert pre_hook.config is config | |
| assert pre_hook.mode == HeadroomMode.AUDIT | |
| assert pre_hook.model == "gpt-4-turbo" | |
| assert post_hook.log_level == "WARNING" | |
| assert post_hook.token_alert_threshold == 8000 | |