diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 6337b638..82dcfabf 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -26,10 +26,10 @@ jobs: run: pip install httpx==0.28.1 pytest pytest-asyncio anyio pyyaml fastapi "pydantic>=2.0,<3.0" uvicorn python-multipart asyncpg langfuse redis aiofiles==25.1.0 - name: Run Unit Tests - run: CONFIG_PATH=router/config.yaml PYTHONPATH=.:router pytest --ignore=tests/test_agy_behavior.py --ignore=tests/test_agy_tiers.py --ignore=tests/test_antigravity.py + run: CONFIG_PATH=router/config.yaml PYTHONPATH=. pytest --ignore=tests/test_agy_behavior.py --ignore=tests/test_agy_tiers.py --ignore=tests/test_antigravity.py - name: Run Breaker Verification run: PYTHONPATH=. python3 scripts/verification/verify_breaker.py - name: Run Integration Verification - run: PYTHONPATH=.:router python3 tests/test_a2_verify.py + run: PYTHONPATH=. python3 tests/test_a2_verify.py diff --git a/router/__init__.py b/router/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/router/agy_proxy.py b/router/agy_proxy.py index 8c1f8631..2ad75b2a 100644 --- a/router/agy_proxy.py +++ b/router/agy_proxy.py @@ -38,7 +38,7 @@ async def save(self) -> None: """Push updated cooldown state to Valkey.""" ... -from circuit_breaker import get_google_breaker, get_vendor_breaker +from router.circuit_breaker import get_google_breaker, get_vendor_breaker logger = logging.getLogger("agy-proxy") diff --git a/router/main.py b/router/main.py index 0c976f4f..02b63c75 100644 --- a/router/main.py +++ b/router/main.py @@ -19,7 +19,7 @@ from fastapi.staticfiles import StaticFiles from pathlib import Path from urllib.parse import urlparse -from circuit_breaker import get_breaker +from router.circuit_breaker import get_breaker from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator, RootModel from typing import Dict, Optional, Union diff --git a/router/tests/test_dashboard_data.py b/router/tests/test_dashboard_data.py index 1e811bed..fa2b12f2 100644 --- a/router/tests/test_dashboard_data.py +++ b/router/tests/test_dashboard_data.py @@ -1,27 +1,20 @@ import pytest from unittest.mock import AsyncMock, patch, MagicMock -import sys -import os @pytest.mark.asyncio async def test_get_dashboard_data_structure(): - # Ensure router directory is in sys.path - router_path = os.path.join(os.getcwd(), "router") - if router_path not in sys.path: - sys.path.insert(0, router_path) - - import main + from router import main # Mocking all I/O and external calls - with patch("main.sync_cooldowns_from_valkey", new_callable=AsyncMock) as mock_sync, \ - patch("main.check_tcp_port", new_callable=AsyncMock) as mock_tcp, \ - patch("main.check_http_endpoint", new_callable=AsyncMock) as mock_http, \ - patch("main.get_gemini_oauth_status", new_callable=AsyncMock) as mock_oauth, \ - patch("main.get_best_free_model", new_callable=AsyncMock) as mock_best_model, \ - patch("main.get_goose_sessions") as mock_goose, \ - patch("main.get_llamacpp_metrics", new_callable=AsyncMock) as mock_llamacpp, \ - patch("main.get_pie_chart_gradient") as mock_gradient, \ - patch("main.stats") as mock_stats: + with patch("router.main.sync_cooldowns_from_valkey", new_callable=AsyncMock) as mock_sync, \ + patch("router.main.check_tcp_port", new_callable=AsyncMock) as mock_tcp, \ + patch("router.main.check_http_endpoint", new_callable=AsyncMock) as mock_http, \ + patch("router.main.get_gemini_oauth_status", new_callable=AsyncMock) as mock_oauth, \ + patch("router.main.get_best_free_model", new_callable=AsyncMock) as mock_best_model, \ + patch("router.main.get_goose_sessions") as mock_goose, \ + patch("router.main.get_llamacpp_metrics", new_callable=AsyncMock) as mock_llamacpp, \ + patch("router.main.get_pie_chart_gradient") as mock_gradient, \ + patch("router.main.stats") as mock_stats: # Setup mock return values mock_sync.return_value = None diff --git a/router/tests/test_detect_active_tool.py b/router/tests/test_detect_active_tool.py index 0bebb391..95774be2 100644 --- a/router/tests/test_detect_active_tool.py +++ b/router/tests/test_detect_active_tool.py @@ -1,5 +1,5 @@ import pytest -from main import detect_active_tool +from router.main import detect_active_tool def test_detect_active_tool_empty(): assert detect_active_tool({}) == "none" diff --git a/router/tests/test_estimate_prompt_tokens.py b/router/tests/test_estimate_prompt_tokens.py index 32da0463..71d17532 100644 --- a/router/tests/test_estimate_prompt_tokens.py +++ b/router/tests/test_estimate_prompt_tokens.py @@ -1,5 +1,5 @@ import pytest -from main import estimate_prompt_tokens +from router.main import estimate_prompt_tokens def test_estimate_prompt_tokens_empty(): assert estimate_prompt_tokens({}) == 50 diff --git a/router/tests/test_get_gemini_oauth_status.py b/router/tests/test_get_gemini_oauth_status.py index 9d506334..016b5356 100644 --- a/router/tests/test_get_gemini_oauth_status.py +++ b/router/tests/test_get_gemini_oauth_status.py @@ -2,7 +2,7 @@ import json from unittest.mock import patch, mock_open -import main +from router import main @pytest.mark.asyncio async def test_get_gemini_oauth_status_missing_file(): diff --git a/router/tests/test_get_goose_sessions.py b/router/tests/test_get_goose_sessions.py index 45de1a05..4987d4e0 100644 --- a/router/tests/test_get_goose_sessions.py +++ b/router/tests/test_get_goose_sessions.py @@ -2,7 +2,7 @@ import sys from unittest.mock import patch, MagicMock -from main import get_goose_sessions +from router.main import get_goose_sessions def test_get_goose_sessions_no_db(): with patch('os.path.exists', return_value=False): diff --git a/router/tests/test_load_persisted_stats.py b/router/tests/test_load_persisted_stats.py index 67458d50..8e32663a 100644 --- a/router/tests/test_load_persisted_stats.py +++ b/router/tests/test_load_persisted_stats.py @@ -2,8 +2,8 @@ import pytest from unittest.mock import patch, mock_open -import main -from main import load_persisted_stats +from router import main +from router.main import load_persisted_stats @pytest.fixture def mock_stats(): @@ -17,7 +17,7 @@ def mock_stats(): yield main.stats def test_load_persisted_stats_file_not_exists(mock_stats): - with patch("main.os.path.exists", return_value=False) as mock_exists: + with patch("router.main.os.path.exists", return_value=False) as mock_exists: load_persisted_stats() mock_exists.assert_called_once_with(main.STATS_JSON_PATH) # Stats should remain unchanged @@ -31,9 +31,9 @@ def test_load_persisted_stats_success(mock_stats): } mock_json = json.dumps(mock_data) - with patch("main.os.path.exists", side_effect=lambda p: p == main.STATS_JSON_PATH): - with patch("main.open", mock_open(read_data=mock_json)): - with patch("main.logger.info") as mock_logger: + with patch("router.main.os.path.exists", side_effect=lambda p: p == main.STATS_JSON_PATH): + with patch("router.main.open", mock_open(read_data=mock_json)): + with patch("router.main.logger.info") as mock_logger: load_persisted_stats() # Assert simple value updated via else block @@ -48,9 +48,9 @@ def test_load_persisted_stats_success(mock_stats): mock_logger.assert_called_once_with("✓ Successfully loaded persisted gateway statistics from disk.") def test_load_persisted_stats_exception(mock_stats): - with patch("main.os.path.exists", return_value=True): - with patch("main.open", side_effect=Exception("Mock read error")): - with patch("main.logger.error") as mock_logger: + with patch("router.main.os.path.exists", return_value=True): + with patch("router.main.open", side_effect=Exception("Mock read error")): + with patch("router.main.logger.error") as mock_logger: load_persisted_stats() # Stats should remain unchanged diff --git a/router/tests/test_resolve_external_urls.py b/router/tests/test_resolve_external_urls.py index b86bf4c5..8cc79f0d 100644 --- a/router/tests/test_resolve_external_urls.py +++ b/router/tests/test_resolve_external_urls.py @@ -1,15 +1,9 @@ import pytest import os -import sys from unittest.mock import MagicMock -# Ensure router directory is in sys.path -router_path = os.path.join(os.getcwd(), "router") -if router_path not in sys.path: - sys.path.insert(0, router_path) - -import main -from main import resolve_external_urls +from router import main +from router.main import resolve_external_urls class MockRequest: diff --git a/scripts/benchmark_classifier.py b/scripts/benchmark_classifier.py index 2878c402..b2b7f1a7 100644 --- a/scripts/benchmark_classifier.py +++ b/scripts/benchmark_classifier.py @@ -1,14 +1,16 @@ """Benchmark gemma4-26a4b-routing classifier against labeled dataset.""" import os -import json, urllib.request, urllib.error, time, sys +import json, urllib.request, urllib.error, time import concurrent.futures import threading from collections import defaultdict, Counter from pathlib import Path # Shared chat response parser (used by verification scripts too) -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) -from scripts.chat_helpers import parse_chat_response +try: + from scripts.chat_helpers import parse_chat_response +except ImportError: + from chat_helpers import parse_chat_response # Load dataset dataset_path = Path(__file__).resolve().parent.parent / "data" / "classified_dataset.json" diff --git a/scripts/benchmark_tokens.py b/scripts/benchmark_tokens.py index 34c9a2f0..0737fa01 100644 --- a/scripts/benchmark_tokens.py +++ b/scripts/benchmark_tokens.py @@ -5,11 +5,12 @@ # Set CONFIG_PATH and ROUTER_API_KEY for import os.environ["CONFIG_PATH"] = str(Path(__file__).resolve().parent.parent / "router" / "config.yaml") os.environ["ROUTER_API_KEY"] = "local-token" -# Add the parent directory and the router directory to the path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) -sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "router")) -from router.main import estimate_prompt_tokens, METADATA_OVERHEAD +try: + from router.main import estimate_prompt_tokens, METADATA_OVERHEAD +except ImportError: + sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + from router.main import estimate_prompt_tokens, METADATA_OVERHEAD def verify_accuracy(): """Benchmarking utility to verify token estimation accuracy across content types.""" diff --git a/scripts/classify_direct.py b/scripts/classify_direct.py index ee699daf..237609c7 100644 --- a/scripts/classify_direct.py +++ b/scripts/classify_direct.py @@ -1,11 +1,13 @@ """Direct classification of Hermes prompts using gemma4-26a4b-routing.""" import os -import json, urllib.request, time, sys +import json, urllib.request, time from pathlib import Path # Shared chat response parser (used by verification scripts too) -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) -from scripts.chat_helpers import parse_chat_response +try: + from scripts.chat_helpers import parse_chat_response +except ImportError: + from chat_helpers import parse_chat_response PROMPT_TEMPLATE = """Classify the coding task complexity. Output ONLY the tier name. diff --git a/scripts/reclassify_all.py b/scripts/reclassify_all.py index 38675878..b74e5653 100644 --- a/scripts/reclassify_all.py +++ b/scripts/reclassify_all.py @@ -8,8 +8,10 @@ from collections import Counter # Shared chat response parser (used by verification scripts too) -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) -from scripts.chat_helpers import parse_chat_response +try: + from scripts.chat_helpers import parse_chat_response +except ImportError: + from chat_helpers import parse_chat_response TIERS = ['agent-simple-core','agent-medium-core','agent-complex-core','agent-reasoning-core','agent-advanced-core'] diff --git a/scripts/retry_errors.py b/scripts/retry_errors.py index ed273721..c4a30754 100644 --- a/scripts/retry_errors.py +++ b/scripts/retry_errors.py @@ -4,8 +4,10 @@ from collections import Counter # Shared chat response parser (used by verification scripts too) -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) -from scripts.chat_helpers import parse_chat_response +try: + from scripts.chat_helpers import parse_chat_response +except ImportError: + from chat_helpers import parse_chat_response PROMPT_TEMPLATE = """Classify the coding task complexity. Output ONLY the tier name. diff --git a/scripts/verification/verification_helpers.py b/scripts/verification/verification_helpers.py index e4500d61..a3a1c839 100644 --- a/scripts/verification/verification_helpers.py +++ b/scripts/verification/verification_helpers.py @@ -1,8 +1,11 @@ # Shared verification helpers for cooldown and routing tests -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) -from scripts.chat_helpers import parse_chat_response +try: + from scripts.chat_helpers import parse_chat_response +except ImportError: + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from scripts.chat_helpers import parse_chat_response import os import uuid import time diff --git a/scripts/verification/verify_breaker.py b/scripts/verification/verify_breaker.py index 6ed9e7c7..eebd6687 100644 --- a/scripts/verification/verify_breaker.py +++ b/scripts/verification/verify_breaker.py @@ -1,11 +1,13 @@ #!/usr/bin/env python3 """Verification test for the agy circuit breaker.""" -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) - -from router.circuit_breaker import get_breaker +try: + from router.circuit_breaker import get_breaker +except ImportError: + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from router.circuit_breaker import get_breaker b = get_breaker() assert b.is_allowed(), 'Tier 0 should be open' diff --git a/scripts/verification/verify_canonical_endpoints.py b/scripts/verification/verify_canonical_endpoints.py index 13c7d9b2..c0c6ce62 100644 --- a/scripts/verification/verify_canonical_endpoints.py +++ b/scripts/verification/verify_canonical_endpoints.py @@ -19,8 +19,11 @@ WORKDIR = Path(__file__).resolve().parent.parent.parent # Import shared chat response parser (also used by classifier scripts) -sys.path.insert(0, str(WORKDIR)) -from scripts.chat_helpers import parse_chat_response +try: + from scripts.chat_helpers import parse_chat_response +except ImportError: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from scripts.chat_helpers import parse_chat_response def load_env(dev: bool = False) -> dict: diff --git a/scripts/verification/verify_ollama_routing.py b/scripts/verification/verify_ollama_routing.py index c844e549..6f666b73 100644 --- a/scripts/verification/verify_ollama_routing.py +++ b/scripts/verification/verify_ollama_routing.py @@ -5,8 +5,11 @@ from pathlib import Path WORKDIR = Path(__file__).resolve().parent.parent.parent -sys.path.insert(0, str(WORKDIR)) -from scripts.chat_helpers import parse_chat_response +try: + from scripts.chat_helpers import parse_chat_response +except ImportError: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from scripts.chat_helpers import parse_chat_response URL = "http://localhost:5000/v1/chat/completions" diff --git a/tests/test_a2_verify.py b/tests/test_a2_verify.py index 95c6be31..f6caf19b 100644 --- a/tests/test_a2_verify.py +++ b/tests/test_a2_verify.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """Verify circuit breaker integration into agy_proxy.py""" -from circuit_breaker import get_breaker -from agy_proxy import try_agy_proxy +from router.circuit_breaker import get_breaker +from router.agy_proxy import try_agy_proxy import asyncio, time b = get_breaker() diff --git a/tests/test_models_proxy.py b/tests/test_models_proxy.py index 2ad68535..2c90ec8f 100644 --- a/tests/test_models_proxy.py +++ b/tests/test_models_proxy.py @@ -4,11 +4,11 @@ from fastapi import Response from fastapi.responses import JSONResponse -from main import get_http_client, proxy_models, HTTP_MAX_CONNECTIONS, HTTP_MAX_KEEPALIVE_CONNECTIONS, HTTP_KEEPALIVE_EXPIRY +from router.main import get_http_client, proxy_models, HTTP_MAX_CONNECTIONS, HTTP_MAX_KEEPALIVE_CONNECTIONS, HTTP_KEEPALIVE_EXPIRY def test_http_client_limits(): # Verify that get_http_client initializes with configured limits using public mocks - import main + from router import main import httpx original_init = httpx.Limits.__init__ @@ -43,7 +43,7 @@ async def test_proxy_models_success(): mock_client = AsyncMock() mock_client.get.return_value = mock_resp - with patch("main.get_http_client", return_value=mock_client): + with patch("router.main.get_http_client", return_value=mock_client): response = await proxy_models() assert isinstance(response, JSONResponse) assert response.status_code == 200 @@ -67,7 +67,7 @@ async def test_proxy_models_error_status(): mock_client = AsyncMock() mock_client.get.return_value = mock_resp - with patch("main.get_http_client", return_value=mock_client): + with patch("router.main.get_http_client", return_value=mock_client): response = await proxy_models() assert isinstance(response, Response) assert response.status_code == 500 @@ -85,7 +85,7 @@ async def test_proxy_models_invalid_json(): mock_client = AsyncMock() mock_client.get.return_value = mock_resp - with patch("main.get_http_client", return_value=mock_client): + with patch("router.main.get_http_client", return_value=mock_client): response = await proxy_models() assert isinstance(response, Response) assert response.status_code == 200