From a406b43c9918a12431d3469a3755d0a916d73993 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 30 Jun 2026 19:35:17 +0000 Subject: [PATCH 01/11] =?UTF-8?q?=E2=9A=A1=20[performance]=20Unblock=20eve?= =?UTF-8?q?nt=20loop=20by=20using=20aiofiles=20for=20CLI=20log=20reading?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> --- router/agy_proxy.py | 14 +++++++----- router/tests/test_agy_proxy.py | 41 ++++++++++++++++++++-------------- 2 files changed, 32 insertions(+), 23 deletions(-) diff --git a/router/agy_proxy.py b/router/agy_proxy.py index 5da0c335..844a04f0 100644 --- a/router/agy_proxy.py +++ b/router/agy_proxy.py @@ -22,6 +22,7 @@ import json import logging import os +import aiofiles import time import httpx from typing import Optional, Protocol, runtime_checkable @@ -127,7 +128,7 @@ async def _run_agy_print(client: httpx.AsyncClient, prompt: str, model_override: # Track the last log check time to avoid hammering the file _last_log_check: float = 0 -def _is_quota_exhausted(returncode: int, stdout: str, stderr: str) -> bool: +async def _is_quota_exhausted(returncode: int, stdout: str, stderr: str) -> bool: """ Detect quota exhaustion from agy subprocess results. @@ -152,8 +153,9 @@ def _is_quota_exhausted(returncode: int, stdout: str, stderr: str) -> bool: log_path = os.path.expanduser("~/.gemini/antigravity-cli/cli.log") try: if os.path.exists(log_path): - with open(log_path, "r") as f: - for line in f.readlines()[-5:]: + async with aiofiles.open(log_path, "r") as f: + lines = await f.readlines() + for line in lines[-5:]: if "RESOURCE_EXHAUSTED" in line or "code 429" in line: return True except Exception: @@ -346,8 +348,8 @@ async def try_agy_proxy(prompt: str, messages: list = None, rc = 0 if raw_rc is None else raw_rc raw_stderr = first_data.get("stderr", "") stderr_content = "" if raw_stderr is None else raw_stderr - if _is_quota_exhausted(rc, "", stderr_content) or rc != 0: - if _is_quota_exhausted(rc, "", stderr_content): + if await _is_quota_exhausted(rc, "", stderr_content) or rc != 0: + if await _is_quota_exhausted(rc, "", stderr_content): tier_breaker.record_failure() if cooldown_persistence is not None: try: @@ -420,7 +422,7 @@ async def token_generator(stream_resp, httpx_client, initial_line, current_conv_ last_conv_id = result_conv_id # Check for quota exhaustion - if _is_quota_exhausted(returncode, stdout, stderr): + if await _is_quota_exhausted(returncode, stdout, stderr): tier_breaker.record_failure() if cooldown_persistence is not None: try: diff --git a/router/tests/test_agy_proxy.py b/router/tests/test_agy_proxy.py index 404059eb..f4cd540d 100644 --- a/router/tests/test_agy_proxy.py +++ b/router/tests/test_agy_proxy.py @@ -1,4 +1,5 @@ -from unittest.mock import patch, MagicMock +import pytest +from unittest.mock import patch, MagicMock, AsyncMock from router.agy_proxy import _wrap_response, _is_quota_exhausted def test_wrap_response_basic(): @@ -58,45 +59,51 @@ def test_wrap_response_long_strings(): assert result["usage"]["completion_tokens"] == 250 assert result["usage"]["total_tokens"] == 375 -def test_is_quota_exhausted_stderr_markers(): +@pytest.mark.asyncio +async def test_is_quota_exhausted_stderr_markers(): markers = ["RESOURCE_EXHAUSTED", "code 429", "quota reached", "rate limit"] for marker in markers: - assert _is_quota_exhausted(0, "", marker) is True - assert _is_quota_exhausted(1, "", f"Error: {marker}") is True + assert await _is_quota_exhausted(0, "", marker) is True + assert await _is_quota_exhausted(1, "", f"Error: {marker}") is True -def test_is_quota_exhausted_success(): - assert _is_quota_exhausted(0, "some valid response", "") is False +@pytest.mark.asyncio +async def test_is_quota_exhausted_success(): + assert await _is_quota_exhausted(0, "some valid response", "") is False -def test_is_quota_exhausted_other_error(): - assert _is_quota_exhausted(1, "", "some other random error") is False +@pytest.mark.asyncio +async def test_is_quota_exhausted_other_error(): + assert await _is_quota_exhausted(1, "", "some other random error") is False @patch("router.agy_proxy.os.path.exists") -@patch("builtins.open") +@patch("aiofiles.open") @patch("router.agy_proxy.time.time") -def test_is_quota_exhausted_empty_reads_log(mock_time, mock_open, mock_exists): +@pytest.mark.asyncio +async def test_is_quota_exhausted_empty_reads_log(mock_time, mock_open, mock_exists): mock_time.return_value = 1000.0 mock_exists.return_value = True - mock_file = MagicMock() + mock_file = AsyncMock() mock_file.readlines.return_value = ["line 1\n", "RESOURCE_EXHAUSTED info\n", "line 3\n"] - mock_open.return_value.__enter__.return_value = mock_file + mock_open.return_value.__aenter__.return_value = mock_file with patch("router.agy_proxy._last_log_check", 0): - assert _is_quota_exhausted(0, "", "") is True + assert await _is_quota_exhausted(0, "", "") is True @patch("router.agy_proxy.time.time") -def test_is_quota_exhausted_empty_throttled(mock_time): +@pytest.mark.asyncio +async def test_is_quota_exhausted_empty_throttled(mock_time): # Set time diff to be < 2.0 mock_time.return_value = 1001.0 with patch("router.agy_proxy._last_log_check", 1000.0): # Even without reading log, falls back to True - assert _is_quota_exhausted(0, "", "") is True + assert await _is_quota_exhausted(0, "", "") is True @patch("router.agy_proxy.os.path.exists") @patch("router.agy_proxy.time.time") -def test_is_quota_exhausted_empty_no_log_fallback(mock_time, mock_exists): +@pytest.mark.asyncio +async def test_is_quota_exhausted_empty_no_log_fallback(mock_time, mock_exists): mock_time.return_value = 1000.0 mock_exists.return_value = False with patch("router.agy_proxy._last_log_check", 0): - assert _is_quota_exhausted(0, "", "") is True + assert await _is_quota_exhausted(0, "", "") is True From 176d2242a05aaa05a0c7fd72d036aacdcdd73638 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 30 Jun 2026 19:46:06 +0000 Subject: [PATCH 02/11] =?UTF-8?q?=E2=9A=A1=20[performance]=20Unblock=20eve?= =?UTF-8?q?nt=20loop=20by=20using=20aiofiles=20for=20CLI=20log=20reading?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> --- .github/workflows/test.yml | 2 +- router/Dockerfile | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 0b8e5878..e4c73991 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -23,7 +23,7 @@ jobs: python-version: '3.11' - name: Install dependencies - run: pip install httpx==0.28.1 pytest pytest-asyncio anyio pyyaml fastapi "pydantic>=2.0,<3.0" uvicorn python-multipart asyncpg langfuse redis + 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 - name: Run Unit Tests run: CONFIG_PATH=router/config.yaml PYTHONPATH=.:router pytest --ignore=test_agy_behavior.py --ignore=test_agy_tiers.py diff --git a/router/Dockerfile b/router/Dockerfile index cd0f3653..49d9c6a6 100644 --- a/router/Dockerfile +++ b/router/Dockerfile @@ -3,7 +3,7 @@ FROM python:3.14-slim WORKDIR /app # Install deps in a single layer (no pip cache, no extra files) -RUN pip install --no-cache-dir fastapi "pydantic>=2.0,<3.0" uvicorn httpx pyyaml python-multipart asyncpg langfuse redis +RUN pip install --no-cache-dir fastapi "pydantic>=2.0,<3.0" uvicorn httpx pyyaml python-multipart asyncpg langfuse redis aiofiles # Copy all source in one layer — removes dead config COPY (volume-mounted at runtime) COPY main.py agy_proxy.py circuit_breaker.py aa_scores.json free_models_roster.json /app/ From abd03cb5c3429fcb7ef406a26dcfa55aba744b69 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 09:23:23 +0000 Subject: [PATCH 03/11] =?UTF-8?q?=E2=9A=A1=20[performance]=20Unblock=20eve?= =?UTF-8?q?nt=20loop=20by=20using=20aiofiles=20for=20CLI=20log=20reading?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> --- .github/dependabot.yml | 4 +- .github/workflows/test.yml | 6 +- README.md | 48 +- scripts/get_pr_status.py => get_pr_status.py | 0 .../host_agy_daemon.py => host_agy_daemon.py | 106 ++-- litellm/entrypoint.py | 20 - pod.yaml | 40 +- router/main.py | 44 +- router/test_memory_mcp.py | 131 +++++ router/tests/test_agy_proxy.py | 13 +- router/tests/test_dashboard_data.py | 2 +- router/tests/test_estimate_prompt_tokens.py | 1 + router/tests/test_get_redis.py | 101 ---- router/tests/test_load_persisted_stats.py | 99 ---- router/tests/test_memory_mcp.py | 327 ------------ scripts/README.md | 49 +- start-stack.sh | 85 +-- ...nc_gemini_token.py => sync_gemini_token.py | 0 tests/test_a2_verify.py => test_a2_verify.py | 8 +- ...st_agy_behavior.py => test_agy_behavior.py | 0 tests/test_agy_tiers.py => test_agy_tiers.py | 0 ...test_antigravity.py => test_antigravity.py | 0 ...st_atomic_write.py => test_atomic_write.py | 0 ...endpoint.py => test_check_http_endpoint.py | 0 ...cuit_breaker.py => test_circuit_breaker.py | 7 +- ...accuracy.py => test_classifier_accuracy.py | 0 ...ore.py => test_compute_free_model_score.py | 0 ...ategory.py => test_map_tool_to_category.py | 0 test_memory_mcp.py | 167 ++++++ ...st_models_proxy.py => test_models_proxy.py | 8 +- ..._gradient.py => test_pie_chart_gradient.py | 0 ...test_quota_reset.sh => test_quota_reset.sh | 0 ...tool_usage.py => test_record_tool_usage.py | 2 +- tests/test_src_badge.py => test_src_badge.py | 0 ...tream_latency.py => test_stream_latency.py | 0 ...mini_token.py => test_sync_gemini_token.py | 8 - tests/test_host_agy_daemon.py | 490 ------------------ triage_upgrade_plan.md | 4 +- .../verify_breaker.py => verify_breaker.py | 8 +- scripts/watch_quota.sh => watch_quota.sh | 3 +- 40 files changed, 442 insertions(+), 1339 deletions(-) rename scripts/get_pr_status.py => get_pr_status.py (100%) rename scripts/host_agy_daemon.py => host_agy_daemon.py (72%) create mode 100644 router/test_memory_mcp.py delete mode 100644 router/tests/test_get_redis.py delete mode 100644 router/tests/test_load_persisted_stats.py delete mode 100644 router/tests/test_memory_mcp.py rename scripts/sync_gemini_token.py => sync_gemini_token.py (100%) rename tests/test_a2_verify.py => test_a2_verify.py (72%) rename tests/test_agy_behavior.py => test_agy_behavior.py (100%) rename tests/test_agy_tiers.py => test_agy_tiers.py (100%) rename tests/test_antigravity.py => test_antigravity.py (100%) rename tests/test_atomic_write.py => test_atomic_write.py (100%) rename tests/test_check_http_endpoint.py => test_check_http_endpoint.py (100%) rename tests/test_circuit_breaker.py => test_circuit_breaker.py (98%) rename tests/test_classifier_accuracy.py => test_classifier_accuracy.py (100%) rename tests/test_compute_free_model_score.py => test_compute_free_model_score.py (100%) rename tests/test_map_tool_to_category.py => test_map_tool_to_category.py (100%) create mode 100644 test_memory_mcp.py rename tests/test_models_proxy.py => test_models_proxy.py (93%) rename tests/test_pie_chart_gradient.py => test_pie_chart_gradient.py (100%) rename scripts/test_quota_reset.sh => test_quota_reset.sh (100%) rename tests/test_record_tool_usage.py => test_record_tool_usage.py (98%) rename tests/test_src_badge.py => test_src_badge.py (100%) rename tests/test_stream_latency.py => test_stream_latency.py (100%) rename tests/test_sync_gemini_token.py => test_sync_gemini_token.py (96%) delete mode 100644 tests/test_host_agy_daemon.py rename scripts/verification/verify_breaker.py => verify_breaker.py (76%) rename scripts/watch_quota.sh => watch_quota.sh (92%) diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 102b4c74..2010274f 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -37,9 +37,9 @@ updates: # - dockerhub ignore: # Example: ignore semver-major updates for Postgres and ClickHouse while testing - - dependency-name: "pgvector/pgvector" + - dependency-name: "docker.io/pgvector/pgvector" update-types: ["version-update:semver-major"] - - dependency-name: "clickhouse/clickhouse-server" + - dependency-name: "docker.io/clickhouse/clickhouse-server" update-types: ["version-update:semver-major"] # Dockerfiles (keeps Dockerfile FROM lines updated) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 790f2ff9..e4c73991 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 - 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 + run: CONFIG_PATH=router/config.yaml PYTHONPATH=.:router pytest --ignore=test_agy_behavior.py --ignore=test_agy_tiers.py - name: Run Breaker Verification - run: python3 scripts/verification/verify_breaker.py + run: python3 verify_breaker.py - name: Run Integration Verification - run: python3 tests/test_a2_verify.py + run: python3 test_a2_verify.py diff --git a/README.md b/README.md index bdfac3cf..4a162e8c 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,8 @@ graph TD style QwenLocal fill:#f0f0f0,stroke:#999,stroke-width:1px; ``` +> **Version Pin**: LiteLLM Gateway runs `ghcr.io/berriai/litellm:v1.88.0`. See §3B for pinning policy. + --- ## 1b. Container Health Checks & Auto-Restart @@ -74,12 +76,12 @@ All core containers are configured with **Kubernetes-style liveness and readines | Container | Liveness Probe | Readiness Probe | |:---|---:|---:| -| **valkey-cache** | `tcpSocket` on port 6379 every 10s | Same, every 5s | +| **valkey-cache** (9.1.0-alpine) | `tcpSocket` on port 6379 every 10s | Same, every 5s | | **litellm-gateway** | Python `urllib` GET `/ping` (port 4000) every 15s | Python `urllib` GET `/health/readiness` (port 4000) every 10s | | **llm-triage-router** | Python `urllib` GET `/metrics` (port 5000) every 15s | Same, every 10s | | **postgres-db** | `pg_isready -U postgres` every 10s | Same, every 5s | | **clickhouse-db** | `clickhouse-client --user clickhouse --password clickhouse --query "SELECT 1"` every 15s | `clickhouse-client --query "SELECT 1"` every 10s | -| **valkey-lf** | `tcpSocket` on port 6380 every 10s | Same, every 5s | +| **valkey-lf** (9.1.0-alpine) | `tcpSocket` on port 6380 every 10s | Same, every 5s | | **langfuse-web** | `wget` GET `/api/health` (port 3001) every 15s | Same, every 10s | | **langfuse-worker** | `pgrep node` every 15s | — | | **minio-s3** | `httpGet` `/minio/health/live` (port 9002) every 15s | `httpGet` `/minio/health/ready` (port 9002) every 10s | @@ -210,7 +212,7 @@ The gateway supports multiple routing modes controlled by the `model` field: All configurations, automation scripts, and databases are self-contained within this repository directory: ``` -LLM-Routing/ +/home/gpav/Vrac/LAB/AI/LLM-Routing/ ├── .env # Environment file for API keys, passwords, and generated secrets (ignored by git) ├── .gitignore # Git ignore policy protecting secrets & database files ├── README.md # In-depth system and operational guide @@ -226,26 +228,16 @@ LLM-Routing/ │ ├── agy_proxy.py # 3-tier agy fallback with session continuation │ ├── circuit_breaker.py # Exponential cooldown breaker for agy proxy │ └── memory_mcp.py # MCP bridge server for Goose memory integration -├── scripts/ # Automation, maintenance, and verification scripts -│ ├── backup.sh # Database backup with pg_isready retry logic -│ ├── host_agy_daemon.py # Real-time PTY-based streaming daemon for agy -│ ├── sync_gemini_token.py # Extraction/sync script for keyring credentials -│ ├── get_pr_status.py # PR state helper -│ ├── watch_quota.sh # Quota reset watcher -│ ├── test_quota_reset.sh # Quota reset test simulator -│ └── verification/ # Routing & cooldown verification tests -│ ├── verify_breaker.py # Circuit breaker verification -│ └── ... -├── tests/ # Integration & unit test suite -│ ├── test_agy_tiers.py # agy proxy model tier test suite -│ ├── test_classifier_accuracy.py # Classifier accuracy benchmark -│ └── ... +├── scripts/ +│ └── backup.sh # Database backup with pg_isready retry logic ├── backups/ # Timestamped PostgreSQL dumps + config snapshots ├── valkey-data/ # [Git Ignored] Persistent memory volumes for Valkey Cache ├── postgres-data/ # [Git Ignored] Persistent tables for PostgreSQL ├── clickhouse-data/ # [Git Ignored] Persistent traces for Langfuse v3 ├── valkey-lf-data/ # [Git Ignored] Persistent job queues for Langfuse v3 -└── minio-data/ # [Git Ignored] S3-compatible event storage for Langfuse v3 +├── minio-data/ # [Git Ignored] S3-compatible event storage for Langfuse v3 +├── test_agy_tiers.py # agy proxy model tier test suite +└── test_classifier_accuracy.py # Classifier accuracy benchmark ``` --- @@ -277,8 +269,7 @@ Exposes the entry endpoint (`http://localhost:5000/v1`) and evaluates prompt com > Model capabilities, token limits, and costs are visible in LiteLLM's Model Hub Table at `http://localhost:4000/ui/?page=model-hub-table` (or port 4000 on the gateway host). ### B. LiteLLM Proxy Gateway (`litellm/config.yaml`) -- **Version Pinning**: The tags are explicitly pinned in `pod.yaml` — never use `:latest`. Check available tags with `skopeo list-tags docker://ghcr.io/repo/image` before upgrading. - +- **Version Pinning**: The LiteLLM gateway runs `ghcr.io/berriai/litellm:v1.88.0` (latest stable as of June 2026). The tag is explicitly pinned in `pod.yaml` — never use `:latest`. Check available tags with `skopeo list-tags docker://ghcr.io/berriai/litellm` before upgrading. ClickHouse runs `docker.io/clickhouse/clickhouse-server:26.5.1` (upgraded from 24.8, June 2026). Valkey Cache runs `docker.io/valkey/valkey:9.1.0-alpine` (upgraded from 8, June 2026). Orchestrates routing fallback chains, Redis caching, and telemetry callbacks: - **`drop_params: true`**: Automatically strips unsupported arguments when transitioning to models that don't support them. - **Request Timeouts (`300s`)**: Provides ample padding to prevent connection aborts during dynamic RAM swapping operations on the local GPU `llama-server`. @@ -288,7 +279,7 @@ Orchestrates routing fallback chains, Redis caching, and telemetry callbacks: - `embedding_model: "local-nomic-embed"` — uses the local nomic-embed model (no API costs) - `collection_name: "litellm_semantic_cache"` — stores embeddings for similarity-based cache lookups - **Cascading Fallback Chains** (configured in `litellm_settings.fallbacks`): - Each tier escalates through increasingly capable free models, then paid local/remote Ollama models, and finally falls back to `openrouter-auto` (LiteLLM's internal fallback to OpenRouter `/auto`). + Each tier escalates through increasingly capable free models, then paid local/remote Ollama models, and finally falls back to `openrouter-auto` (LiteLLM's internal fallback to OpenRouter `/auto`). `local-qwen-3.6` (35B) was disabled 2026-06-08 to free 23GB RAM/GTT. ```mermaid graph TD @@ -704,6 +695,15 @@ Additional mounts required in `pod.yaml`: mountPath: /root/.gemini # agy expects config at ~/.gemini ``` +### Model Identifiers (found in agy binary) + +| Model | Env Var Value | Backend | +|-------|---------------|---------| +| Gemini 3.5 Flash | `""` (auto-select) | Cloud Code Assist (default) | +| Claude Opus 4.6 | `claude-opus-4-6@default` | Anthropic premium tier | +| Claude Sonnet 4.5 | `claude-sonnet-4-5@20250929` | Anthropic via Vertex AI | +| Claude Haiku 4.5 | `claude-haiku-4-5@20251001` | Anthropic lightweight | + ### Verification ```bash @@ -719,7 +719,7 @@ agy --print "First message" # creates conversation agy --conversation --print "Follow-up" # continues same session # Run the full tier test suite -python3 tests/test_agy_tiers.py +python3 test_agy_tiers.py ``` ### 9b. Streaming & Concurrency Optimizations @@ -727,7 +727,7 @@ python3 tests/test_agy_tiers.py To support production agentic environments (such as `goose-cli` or similar tools) that require low-latency streaming and high concurrent throughput, the following components were introduced: #### 1. Real-Time PTY-Based Streaming Bridge for `agy` Response -To support low-latency streaming for agent clients (such as `goose-cli`), the host-side `scripts/host_agy_daemon.py` runs `agy --print` inside a pseudo-terminal (PTY) using `pty.openpty()`. +To support low-latency streaming for agent clients (such as `goose-cli`), the host-side `host_agy_daemon.py` runs `agy --print` inside a pseudo-terminal (PTY) using `pty.openpty()`. * Running `agy` inside a PTY disables internal buffering, forcing it to write generated characters/lines progressively. * The host daemon streams these chunks in real-time as `application/x-ndjson` lines to the Triage Router. * The Triage Router immediately transforms these incoming chunks into standard OpenAI Server-Sent Event (SSE) packets and yields them to the client. This results in a true, low-latency stream with minimal Time-To-First-Token (TTFT) and eliminates synthetic buffering. @@ -789,7 +789,7 @@ For auto-routing modes, the Triage Router handles failures by silently falling b | Triage Evaluation Layer | Latency Footprint | Hardware Offload | Efficiency Ratio | | :--- | :---: | :---: | :---: | | **Cold-Run Triage** (First query) | ~15 - 24s | Dynamic HF Download | Includes GGUF fetch & initialization | -| **Warm-Run Triage** (Local inference) | **~449 ms** | 100% GPU | **12x speedup** compared to 35B model | +| **Warm-Run Triage** (Local inference) | **~449 ms** | 100% Vulkan GPU (Ryzen APU) | **12x speedup** compared to 35B model | | **Triage Cache Hit** (Repeat query) | **0.0 ms** | RAM In-Memory TTL | Infinite speedup, zero backend requests | | **Valkey Gateway Cache Hit** | **< 10 ms** | Redis RAM Cache | Zero provider cost, immediate response | diff --git a/scripts/get_pr_status.py b/get_pr_status.py similarity index 100% rename from scripts/get_pr_status.py rename to get_pr_status.py diff --git a/scripts/host_agy_daemon.py b/host_agy_daemon.py similarity index 72% rename from scripts/host_agy_daemon.py rename to host_agy_daemon.py index 9f26ff3d..09162856 100755 --- a/scripts/host_agy_daemon.py +++ b/host_agy_daemon.py @@ -65,91 +65,67 @@ async def run_stream(): cmd.extend(["--print", prompt]) master_fd, slave_fd = pty.openpty() - proc = None try: proc = await asyncio.create_subprocess_exec( *cmd, env=env, stdout=slave_fd, stderr=slave_fd, ) + os.close(slave_fd) except Exception as e: - try: - os.close(slave_fd) - except OSError: - pass - try: - os.close(master_fd) - except OSError: - pass + os.close(slave_fd) + os.close(master_fd) # Write failure details as status - try: - err_msg = json.dumps({"type": "status", "returncode": -1, "stderr": str(e)}) + "\n" - self.wfile.write(err_msg.encode('utf-8')) - self.wfile.flush() - except Exception: - pass + err_msg = json.dumps({"type": "status", "returncode": -1, "stderr": str(e)}) + "\n" + self.wfile.write(err_msg.encode('utf-8')) + self.wfile.flush() return - finally: - # Always close the slave end in the parent process + + loop_ref = asyncio.get_running_loop() + + def read_bytes(): try: - os.close(slave_fd) + return os.read(master_fd, 1024) except OSError: - pass + return b"" + + while True: + data = await loop_ref.run_in_executor(None, read_bytes) + if not data: + break + text = data.decode('utf-8', errors='replace') + # PTY text can have \r\n, normalize to \n + text_norm = text.replace('\r\n', '\n') + # Yield token JSON line + chunk_json = json.dumps({"type": "token", "content": text_norm}) + "\n" + self.wfile.write(chunk_json.encode('utf-8')) + self.wfile.flush() - returncode = -1 try: - loop_ref = asyncio.get_running_loop() - - def read_bytes(): - try: - return os.read(master_fd, 1024) - except OSError: - return b"" - - while True: - data = await loop_ref.run_in_executor(None, read_bytes) - if not data: - break - text = data.decode('utf-8', errors='replace') - text_norm = text.replace('\r\n', '\n') - chunk_json = json.dumps({"type": "token", "content": text_norm}) + "\n" - self.wfile.write(chunk_json.encode('utf-8')) - self.wfile.flush() - - # Wait for subprocess await asyncio.wait_for(proc.wait(), timeout=timeout) returncode = proc.returncode or 0 except asyncio.TimeoutError: + try: + proc.kill() + except Exception: + pass returncode = -1 except Exception: returncode = -1 - finally: - # Ensure process is killed and cleaned up - if proc and proc.returncode is None: - try: - proc.kill() - await proc.wait() - except Exception: - pass - - # Ensure master FD is closed - try: - os.close(master_fd) - except OSError: - pass - # Retrieve last conversation ID and write closing status - try: - result_conv_id = get_last_conversation_id() - meta_json = json.dumps({ - "type": "status", - "returncode": returncode, - "conversation_id": result_conv_id - }) + "\n" - self.wfile.write(meta_json.encode('utf-8')) - self.wfile.flush() - except Exception: - pass + os.close(master_fd) + + # Retrieve last conversation ID + result_conv_id = get_last_conversation_id() + + # Write closing metadata + meta_json = json.dumps({ + "type": "status", + "returncode": returncode, + "conversation_id": result_conv_id + }) + "\n" + self.wfile.write(meta_json.encode('utf-8')) + self.wfile.flush() loop.run_until_complete(run_stream()) loop.close() diff --git a/litellm/entrypoint.py b/litellm/entrypoint.py index 20d4baf0..1b987365 100644 --- a/litellm/entrypoint.py +++ b/litellm/entrypoint.py @@ -98,26 +98,6 @@ def strptime(cls, date_str: str, fmt: str) -> original_datetime: datetime.datetime = RobustDatetime sys.stdout.flush() -# Register both RobustDatetime AND the original datetime with Prisma's -# singledispatch serializer. When entrypoint.py replaces datetime.datetime -# with RobustDatetime before Prisma loads, Prisma's own -# @serializer.register(datetime.datetime) ends up registering RobustDatetime. -# But database drivers (psycopg2) return the *original* C-level datetime -# instances, which no longer match. We must register both classes. -try: - from prisma.builder import serializer - def _serialize_dt(dt): - """Serialize datetime to ISO8601 with timezone (UTC if naive).""" - if dt.tzinfo is None: - dt = dt.replace(tzinfo=timezone.utc) - return dt.isoformat() - serializer.register(original_datetime, _serialize_dt) - serializer.register(RobustDatetime, _serialize_dt) - print("🩹 Registered original_datetime + RobustDatetime with Prisma serializer") -except Exception: - pass -sys.stdout.flush() - # Start LiteLLM Proxy import litellm from litellm.proxy.proxy_cli import run_server diff --git a/pod.yaml b/pod.yaml index 65a013aa..9d3af258 100644 --- a/pod.yaml +++ b/pod.yaml @@ -4,14 +4,14 @@ metadata: name: agent-router-pod spec: containers: - # Valkey Cache — LiteLLM response cache (exact-match & semantic) + # Valkey Cache 9.1.0-alpine — LiteLLM response cache (exact-match & semantic) - command: - valkey-server - --protected-mode - 'no' - --loglevel - warning - image: valkey/valkey:9.1.0-alpine + image: docker.io/valkey/valkey:9.1.0-alpine livenessProbe: tcpSocket: port: 6379 @@ -34,7 +34,7 @@ spec: - /app/entrypoint.py env: - name: DATABASE_URL - value: postgresql://postgres:POSTGRES_PASSWORD_ENCODED_PLACEHOLDER@127.0.0.1:5432/postgres + value: postgresql://postgres:***@127.0.0.1:5432/postgres - name: STORE_MODEL_IN_DB value: 'True' - name: LITELLM_LOG @@ -42,10 +42,10 @@ spec: - name: LANGFUSE_HOST value: http://127.0.0.1:3001 - name: LITELLM_MASTER_KEY - value: LITELLM_MASTER_KEY_PLACEHOLDER + value: sk-lit...33bf - name: OLLAMA_API_KEY - value: OLLAMA_API_KEY_PLACEHOLDER - image: berriai/litellm:v1.90.2 + value: 3cd542f8fca14fd08bedfd4f2ab36f9f.6G7Iukbvu1Keyi9x8eKckNEO + image: ghcr.io/berriai/litellm:v1.89.4 livenessProbe: exec: command: @@ -89,19 +89,19 @@ spec: - name: LITELLM_CONFIG_PATH value: /config/litellm_dir/config.yaml - name: DATABASE_URL - value: postgresql://postgres:POSTGRES_PASSWORD_ENCODED_PLACEHOLDER@127.0.0.1:5432/postgres + value: postgresql://postgres:***@127.0.0.1:5432/postgres - name: DBUS_SESSION_BUS_ADDRESS value: unix:path=/run/user/1000/bus - name: LITELLM_MASTER_KEY - value: LITELLM_MASTER_KEY_PLACEHOLDER + value: sk-lit...33bf - name: LANGFUSE_PUBLIC_KEY - value: LANGFUSE_PUBLIC_KEY_PLACEHOLDER + value: pk-lf-200bbf48-1d23-447d-a79a-8614106497e6 - name: LANGFUSE_SECRET_KEY - value: LANGFUSE_SECRET_KEY_PLACEHOLDER + value: sk-lf-...3b49 - name: LANGFUSE_HOST value: http://127.0.0.1:3001 - name: OLLAMA_API_KEY - value: OLLAMA_API_KEY_PLACEHOLDER + value: 3cd542f8fca14fd08bedfd4f2ab36f9f.6G7Iukbvu1Keyi9x8eKckNEO - name: LOG_LEVEL value: info image: localhost/llm-triage-router:latest @@ -154,10 +154,10 @@ spec: - name: POSTGRES_USER value: postgres - name: POSTGRES_PASSWORD - value: POSTGRES_PASSWORD_RAW_PLACEHOLDER + value: postgres-password-*** - name: POSTGRES_DB value: langfuse - image: pgvector/pgvector:0.8.4-pg18 + image: docker.io/pgvector/pgvector:pg18 livenessProbe: exec: command: @@ -188,7 +188,7 @@ spec: value: clickhouse - name: CLICKHOUSE_PASSWORD value: clickhouse - image: clickhouse/clickhouse-server:26.6.1.1193 + image: docker.io/clickhouse/clickhouse-server:26.5.1 livenessProbe: exec: command: @@ -227,7 +227,7 @@ spec: - noeviction - --loglevel - warning - image: valkey/valkey:9.1.0-alpine + image: docker.io/valkey/valkey:9.1.0-alpine livenessProbe: tcpSocket: port: 6380 @@ -253,7 +253,7 @@ spec: # Langfuse Web — observability dashboard - env: - name: DATABASE_URL - value: postgresql://postgres:POSTGRES_PASSWORD_ENCODED_PLACEHOLDER@127.0.0.1:5432/langfuse + value: postgresql://postgres:***@127.0.0.1:5432/langfuse - name: NEXTAUTH_SECRET value: NEXTAUTH_SECRET_PLACEHOLDER - name: NEXTAUTH_URL @@ -310,7 +310,7 @@ spec: value: admin-local-pw-2026 - name: LANGFUSE_LOG_LEVEL value: warn - image: langfuse/langfuse:3.202.1 + image: docker.io/langfuse/langfuse:3 livenessProbe: exec: command: @@ -337,7 +337,7 @@ spec: # Langfuse Worker — background job processor - env: - name: DATABASE_URL - value: postgresql://postgres:POSTGRES_PASSWORD_ENCODED_PLACEHOLDER@127.0.0.1:5432/langfuse + value: postgresql://postgres:***@127.0.0.1:5432/langfuse - name: CLICKHOUSE_URL value: http://127.0.0.1:8123 - name: CLICKHOUSE_USER @@ -372,7 +372,7 @@ spec: value: minioadmin - name: LANGFUSE_LOG_LEVEL value: warn - image: langfuse/langfuse-worker:3.202.1 + image: docker.io/langfuse/langfuse-worker:3 livenessProbe: exec: command: @@ -396,7 +396,7 @@ spec: value: minioadmin - name: MINIO_ROOT_PASSWORD value: minioadmin - image: minio/minio:RELEASE.2025-09-07T16-13-09Z + image: docker.io/minio/minio:latest livenessProbe: httpGet: path: /minio/health/live diff --git a/router/main.py b/router/main.py index b51581ae..7dbb750b 100644 --- a/router/main.py +++ b/router/main.py @@ -1312,18 +1312,17 @@ def get_goose_sessions() -> list: import sqlite3 conn = sqlite3.connect(db_path, timeout=1.0) - try: - conn.row_factory = sqlite3.Row - cursor = conn.cursor() - cursor.execute(""" - SELECT id, name, description, created_at, updated_at, accumulated_total_tokens, goose_mode - FROM sessions - ORDER BY updated_at DESC - LIMIT 5 - """) - sessions_list = [dict(row) for row in cursor] - finally: - conn.close() + conn.row_factory = sqlite3.Row + cursor = conn.cursor() + cursor.execute(""" + SELECT id, name, description, created_at, updated_at, accumulated_total_tokens, goose_mode + FROM sessions + ORDER BY updated_at DESC + LIMIT 5 + """) + for row in cursor.fetchall(): + sessions_list.append(dict(row)) + conn.close() except Exception as e: logger.error(f"Failed to query goose sessions SQLite DB: {e}") return sessions_list @@ -1672,8 +1671,8 @@ async def proxy_models(): "context_length": 524288, }, ] - data["data"] = routing_models + data["data"] - + for entry in reversed(routing_models): + data["data"].insert(0, entry) return JSONResponse(content=data, status_code=200) except Exception as parse_err: logger.warning( @@ -3887,23 +3886,10 @@ def validate_payload(self) -> "AnnotationPayload": # the atomic file-replace mechanism, which is acceptable for this dashboard feature. annotations_lock = asyncio.Lock() -_annotations_cache = {} def _read_annotations_sync(path) -> dict: - import copy - - # Do not swallow OSError if file doesn't exist to preserve original behavior. - # The caller (save_annotations) handles the exception when reading existing annotations. - current_mtime = os.path.getmtime(path) - - cache_entry = _annotations_cache.get(path) - - if cache_entry is None or current_mtime != cache_entry["mtime"]: - with open(path, "r", encoding="utf-8") as f: - data = json.load(f) - _annotations_cache[path] = {"mtime": current_mtime, "data": data} - - return copy.deepcopy(_annotations_cache[path]["data"]) + with open(path, "r", encoding="utf-8") as f: + return json.load(f) @app.post("/dashboard/save-annotations") diff --git a/router/test_memory_mcp.py b/router/test_memory_mcp.py new file mode 100644 index 00000000..24d23a34 --- /dev/null +++ b/router/test_memory_mcp.py @@ -0,0 +1,131 @@ +import time +import re +import sys +import json +from pathlib import Path +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from memory_mcp import _make_key, SCOPE_GLOBAL, SCOPE_LOCAL, PREFIX, _memory_value, _parse_memory_value + +def test_make_key_global(): + """Test generating a key for global scope.""" + category = "test_cat" + data = "test_data" + + before_ts = int(time.time() * 1000) + key = _make_key(category, True, data) + after_ts = int(time.time() * 1000) + + # Expected format: f"{PREFIX}:{scope}:{category}::{ts}:{h}" + parts = key.split(":") + + assert key.startswith(f"{PREFIX}:{SCOPE_GLOBAL}:{category}::") + + # Extract timestamp and hash part + # Format is memory:global:test_cat::1717612345:a1b2c3d4... + match = re.match(rf"^{PREFIX}:{SCOPE_GLOBAL}:{category}::(\d+):([a-f0-9]+)$", key) + assert match is not None, f"Key {key} does not match expected format" + + ts = int(match.group(1)) + h = match.group(2) + + assert before_ts <= ts <= after_ts + assert len(h) == 20 + +def test_make_key_local(): + """Test generating a key for local scope.""" + category = "another_cat" + data = "more_data" + + before_ts = int(time.time() * 1000) + key = _make_key(category, False, data) + after_ts = int(time.time() * 1000) + + assert key.startswith(f"{PREFIX}:{SCOPE_LOCAL}:{category}::") + + match = re.match(rf"^{PREFIX}:{SCOPE_LOCAL}:{category}::(\d+):([a-f0-9]+)$", key) + assert match is not None, f"Key {key} does not match expected format" + + ts = int(match.group(1)) + h = match.group(2) + + assert before_ts <= ts <= after_ts + assert len(h) == 20 + + +def test_make_key_formatting_details(monkeypatch): + """Test the exact output formatting of _make_key using deterministic BLAKE2b.""" + # Mock time.time to return a predictable float so ts = 1620000000123 + monkeypatch.setattr(time, "time", lambda: 1620000000.123) + + # data="data", ts=1620000000123 -> blake2b("data1620000000123", digest_size=10) -> 5e5dad075ca7764bc51f + key1 = _make_key("cat1", True, "data") + assert key1 == f"{PREFIX}:{SCOPE_GLOBAL}:cat1::1620000000123:5e5dad075ca7764bc51f" + + key2 = _make_key("cat2", False, "data") + assert key2 == f"{PREFIX}:{SCOPE_LOCAL}:cat2::1620000000123:5e5dad075ca7764bc51f" + + +def test_make_key_determinism_and_uniqueness(): + """Test determinism for same inputs within same timestamp, and uniqueness across timestamps/data.""" + category = "test_cat" + data1 = "data1" + data2 = "data2" + + key1 = _make_key(category, True, data1) + time.sleep(0.002) + key2 = _make_key(category, True, data1) + key3 = _make_key(category, True, data2) + + # Uniqueness across data + assert key1 != key3 + + # Check determinism: if the timestamp parts are the same, the keys should be identical + ts1 = key1.split("::")[1].split(":")[0] + ts2 = key2.split("::")[1].split(":")[0] + if ts1 == ts2: + assert key1 == key2 + else: + # If timestamp is different, keys should be different + assert key1 != key2 + +def test_memory_value_happy_path(): + """Test _memory_value with standard data and tags.""" + result = _memory_value("some data", ["tag1", "tag2"]) + parsed = json.loads(result) + assert parsed == {"data": "some data", "tags": ["tag1", "tag2"]} + +def test_memory_value_missing_tags(): + """Test _memory_value when tags is None.""" + result = _memory_value("some data", None) + parsed = json.loads(result) + assert parsed == {"data": "some data", "tags": []} + +def test_memory_value_unicode(): + """Test _memory_value properly handles unicode and ensure_ascii=False.""" + result = _memory_value("こんにちは", ["世界"]) + # If ensure_ascii=False, the unicode characters shouldn't be escaped (no \uXXXX) + assert "こんにちは" in result + assert "世界" in result + parsed = json.loads(result) + assert parsed == {"data": "こんにちは", "tags": ["世界"]} + +def test_parse_memory_value_success(): + """Test _parse_memory_value successfully decodes valid JSON.""" + raw = '{"data": "info", "tags": ["a"]}' + result = _parse_memory_value(raw) + assert result == {"data": "info", "tags": ["a"]} + +def test_parse_memory_value_invalid_json(): + """Test _parse_memory_value with invalid JSON.""" + result = _parse_memory_value("{invalid_json:") + assert result == {"data": "{invalid_json:", "tags": []} + +def test_parse_memory_value_type_error(): + """Test _parse_memory_value with TypeError (e.g. passing None).""" + result = _parse_memory_value(None) + assert result == {"data": None, "tags": []} + +def test_parse_memory_value_invalid_json_string(): + """Test _parse_memory_value with invalid JSON string.""" + result = _parse_memory_value("this is not a valid json string") + assert result == {"data": "this is not a valid json string", "tags": []} diff --git a/router/tests/test_agy_proxy.py b/router/tests/test_agy_proxy.py index 0474f080..f4cd540d 100644 --- a/router/tests/test_agy_proxy.py +++ b/router/tests/test_agy_proxy.py @@ -1,14 +1,5 @@ -import sys -from pathlib import Path - -# Dynamic project root discovery -root = Path(__file__).resolve() -while root.parent != root and not (root / ".git").exists(): - root = root.parent -sys.path.insert(0, str(root)) -sys.path.insert(0, str(root / "router")) - -from unittest.mock import patch, MagicMock +import pytest +from unittest.mock import patch, MagicMock, AsyncMock from router.agy_proxy import _wrap_response, _is_quota_exhausted def test_wrap_response_basic(): diff --git a/router/tests/test_dashboard_data.py b/router/tests/test_dashboard_data.py index 6c960d4d..643425f8 100644 --- a/router/tests/test_dashboard_data.py +++ b/router/tests/test_dashboard_data.py @@ -1,6 +1,6 @@ import pytest import asyncio -from unittest.mock import AsyncMock, patch # Removed unused MagicMock import +from unittest.mock import AsyncMock, patch, MagicMock import sys import os diff --git a/router/tests/test_estimate_prompt_tokens.py b/router/tests/test_estimate_prompt_tokens.py index 2a12f97c..e93390f9 100644 --- a/router/tests/test_estimate_prompt_tokens.py +++ b/router/tests/test_estimate_prompt_tokens.py @@ -1,3 +1,4 @@ +import pytest import sys import os from pathlib import Path diff --git a/router/tests/test_get_redis.py b/router/tests/test_get_redis.py deleted file mode 100644 index b92f4ae7..00000000 --- a/router/tests/test_get_redis.py +++ /dev/null @@ -1,101 +0,0 @@ -import os -import time -from unittest.mock import patch, MagicMock - -import pytest - -import router.main as main - -@pytest.fixture(autouse=True) -def reset_redis_globals(): - """Reset the global variables before and after each test.""" - original_client = main._redis_client - original_last_attempt = main._redis_last_init_attempt - - main._redis_client = None - main._redis_last_init_attempt = 0.0 - - yield - - main._redis_client = original_client - main._redis_last_init_attempt = original_last_attempt - -def test_get_redis_already_initialized(): - """If the client is already initialized, it should return the client immediately.""" - mock_client = MagicMock() - main._redis_client = mock_client - - assert main.get_redis() is mock_client - -@patch("router.main.time.monotonic") -def test_get_redis_cooldown(mock_monotonic): - """If init failed recently, it should return None without attempting to initialize.""" - main._redis_client = None - main._redis_last_init_attempt = 100.0 - - # Time elapsed is less than 5.0 seconds - mock_monotonic.return_value = 103.0 - - assert main.get_redis() is None - -@patch("router.main.time.monotonic") -@patch("router.main.aioredis.Redis") -@patch.dict(os.environ, {"VALKEY_HOST": "my-host", "VALKEY_PORT": "1234"}) -def test_get_redis_initialization_success(mock_redis, mock_monotonic): - """If sufficient time has passed, it should initialize and return the client.""" - main._redis_client = None - main._redis_last_init_attempt = 100.0 - - # Time elapsed is 10.0 seconds (greater than 5.0) - mock_monotonic.return_value = 110.0 - - mock_redis_instance = MagicMock() - mock_redis.return_value = mock_redis_instance - - client = main.get_redis() - - assert client is mock_redis_instance - assert main._redis_client is mock_redis_instance - assert main._redis_last_init_attempt == 110.0 - mock_redis.assert_called_once_with(host="my-host", port=1234, decode_responses=True, socket_timeout=1.0) - -@patch("router.main.time.monotonic") -@patch("router.main.logger.warning") -@patch.dict(os.environ, {"VALKEY_HOST": "my-host", "VALKEY_PORT": "invalid"}) -def test_get_redis_initialization_failure(mock_logger_warning, mock_monotonic): - """If initialization fails, it should catch the exception, log a warning, and return None.""" - main._redis_client = None - main._redis_last_init_attempt = 100.0 - - # Time elapsed is 10.0 seconds - mock_monotonic.return_value = 110.0 - - # The int() conversion of VALKEY_PORT will raise ValueError - client = main.get_redis() - - assert client is None - assert main._redis_client is None - assert main._redis_last_init_attempt == 110.0 - mock_logger_warning.assert_called_once() - assert "Failed to initialize Valkey client" in mock_logger_warning.call_args[0][0] - -@patch("router.main.time.monotonic") -@patch("router.main.aioredis.Redis") -@patch("router.main.logger.warning") -@patch.dict(os.environ, {"VALKEY_HOST": "my-host", "VALKEY_PORT": "1234"}) -def test_get_redis_initialization_exception(mock_logger_warning, mock_redis, mock_monotonic): - """If aioredis.Redis throws an exception, it should catch it and return None.""" - main._redis_client = None - main._redis_last_init_attempt = 100.0 - - mock_monotonic.return_value = 110.0 - - mock_redis.side_effect = Exception("Test Exception") - - client = main.get_redis() - - assert client is None - assert main._redis_client is None - assert main._redis_last_init_attempt == 110.0 - mock_logger_warning.assert_called_once() - assert "Test Exception" in mock_logger_warning.call_args[0][0] diff --git a/router/tests/test_load_persisted_stats.py b/router/tests/test_load_persisted_stats.py deleted file mode 100644 index 43cf52ff..00000000 --- a/router/tests/test_load_persisted_stats.py +++ /dev/null @@ -1,99 +0,0 @@ -import pytest -import os -import json -import sys -from unittest.mock import patch, mock_open - -# Ensure router directory is in sys.path based on __file__ instead of getcwd -router_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) -if router_path not in sys.path: - sys.path.insert(0, router_path) - -import main - -def test_load_persisted_stats_success(): - mock_stats = { - "total_requests": 100, - "some_dict": {"a": 1, "b": 2}, - "new_key": "new_value" - } - - mock_timeline = [ - {"time": "12:00", "total_requests": 10, "avg_latency": 50.0} - ] - - # We patch the stats dict directly to avoid replacing the reference - # and to ensure automatic teardown even if an assertion fails. - initial_stats = {"some_dict": {"c": 3}, "timeline": []} - - def mock_exists(path): - if path == main.STATS_JSON_PATH: - return True - if path.endswith("router_timeline.json"): - return True - return False - - # We only intercept specific files, otherwise call the real open - real_open = open - def mock_open_file(file, mode="r", *args, **kwargs): - if file == main.STATS_JSON_PATH: - return mock_open(read_data=json.dumps(mock_stats))() - if type(file) is str and file.endswith("router_timeline.json"): - return mock_open(read_data=json.dumps(mock_timeline))() - return real_open(file, mode, *args, **kwargs) - - with patch.dict('main.stats', initial_stats, clear=True), \ - patch('os.path.exists', side_effect=mock_exists), \ - patch('builtins.open', side_effect=mock_open_file): - - main.load_persisted_stats() - - # Verify stats updated correctly - assert main.stats["total_requests"] == 100 - # Check dictionary merging - assert "some_dict" in main.stats - assert main.stats["some_dict"]["a"] == 1 - assert main.stats["some_dict"]["c"] == 3 - # Check new key added - assert main.stats["new_key"] == "new_value" - # Check timeline loaded - assert main.stats["timeline"] == mock_timeline - -def test_load_persisted_stats_files_missing(): - initial_stats = {"total_requests": 50} - - def mock_exists(path): - return False - - with patch.dict('main.stats', initial_stats, clear=True), \ - patch('os.path.exists', side_effect=mock_exists): - - main.load_persisted_stats() - - # Verify stats didn't change - assert main.stats == initial_stats - -def test_load_persisted_stats_invalid_json(): - initial_stats = {"total_requests": 50} - - def mock_exists(path): - if path == main.STATS_JSON_PATH: - return True - return False - - real_open = open - def mock_open_file(file, mode="r", *args, **kwargs): - if file == main.STATS_JSON_PATH: - return mock_open(read_data="invalid json")() - return real_open(file, mode, *args, **kwargs) - - with patch.dict('main.stats', initial_stats, clear=True), \ - patch('os.path.exists', side_effect=mock_exists), \ - patch('builtins.open', side_effect=mock_open_file), \ - patch('main.logger.error') as mock_logger: - - main.load_persisted_stats() - - assert main.stats == initial_stats - mock_logger.assert_called_once() - assert "Failed to load persisted stats" in mock_logger.call_args[0][0] diff --git a/router/tests/test_memory_mcp.py b/router/tests/test_memory_mcp.py deleted file mode 100644 index 92c5bc96..00000000 --- a/router/tests/test_memory_mcp.py +++ /dev/null @@ -1,327 +0,0 @@ -import json -import re -import sys -import time -from pathlib import Path - -# Dynamic project root discovery -root = Path(__file__).resolve() -while root.parent != root and not (root / ".git").exists(): - root = root.parent -sys.path.insert(0, str(root)) -sys.path.insert(0, str(root / "router")) - -import pytest -from memory_mcp import ( - PREFIX, - SCOPE_GLOBAL, - SCOPE_LOCAL, - _make_key, - _memory_entry, - _memory_value, - _parse_key, - _parse_memory_value, -) - - -# ===================================================================== -# Tests from router/test_memory_mcp.py -# ===================================================================== - -def test_make_key_global(): - """Test generating a key for global scope.""" - category = "test_cat" - data = "test_data" - - before_ts = int(time.time() * 1000) - key = _make_key(category, True, data) - after_ts = int(time.time() * 1000) - - # Expected format: f"{PREFIX}:{scope}:{category}::{ts}:{h}" - assert key.startswith(f"{PREFIX}:{SCOPE_GLOBAL}:{category}::") - - # Extract timestamp and hash part - match = re.match(rf"^{PREFIX}:{SCOPE_GLOBAL}:{category}::(\d+):([a-f0-9]+)$", key) - assert match is not None, f"Key {key} does not match expected format" - - ts = int(match.group(1)) - h = match.group(2) - - assert before_ts <= ts <= after_ts - assert len(h) == 20 - - -def test_make_key_local(): - """Test generating a key for local scope.""" - category = "another_cat" - data = "more_data" - - before_ts = int(time.time() * 1000) - key = _make_key(category, False, data) - after_ts = int(time.time() * 1000) - - assert key.startswith(f"{PREFIX}:{SCOPE_LOCAL}:{category}::") - - match = re.match(rf"^{PREFIX}:{SCOPE_LOCAL}:{category}::(\d+):([a-f0-9]+)$", key) - assert match is not None, f"Key {key} does not match expected format" - - ts = int(match.group(1)) - h = match.group(2) - - assert before_ts <= ts <= after_ts - assert len(h) == 20 - - -def test_make_key_formatting_details(monkeypatch): - """Test the exact output formatting of _make_key using deterministic BLAKE2b.""" - # Mock time.time to return a predictable float so ts = 1620000000123 - monkeypatch.setattr(time, "time", lambda: 1620000000.123) - - # data="data", ts=1620000000123 -> blake2b("data1620000000123", digest_size=10) -> 5e5dad075ca7764bc51f - key1 = _make_key("cat1", True, "data") - assert key1 == f"{PREFIX}:{SCOPE_GLOBAL}:cat1::1620000000123:5e5dad075ca7764bc51f" - - key2 = _make_key("cat2", False, "data") - assert key2 == f"{PREFIX}:{SCOPE_LOCAL}:cat2::1620000000123:5e5dad075ca7764bc51f" - - -def test_make_key_determinism_and_uniqueness(): - """Test determinism for same inputs within same timestamp, and uniqueness across timestamps/data.""" - category = "test_cat" - data1 = "data1" - data2 = "data2" - - key1 = _make_key(category, True, data1) - time.sleep(0.002) - key2 = _make_key(category, True, data1) - key3 = _make_key(category, True, data2) - - # Uniqueness across data - assert key1 != key3 - - # Check determinism: if the timestamp parts are the same, the keys should be identical - ts1 = key1.split("::")[1].split(":")[0] - ts2 = key2.split("::")[1].split(":")[0] - if ts1 == ts2: - assert key1 == key2 - else: - # If timestamp is different, keys should be different - assert key1 != key2 - - -def test_memory_value_happy_path(): - """Test _memory_value with standard data and tags.""" - result = _memory_value("some data", ["tag1", "tag2"]) - parsed = json.loads(result) - assert parsed == {"data": "some data", "tags": ["tag1", "tag2"]} - - -def test_memory_value_missing_tags(): - """Test _memory_value when tags is None.""" - result = _memory_value("some data", None) - parsed = json.loads(result) - assert parsed == {"data": "some data", "tags": []} - - -def test_memory_value_unicode(): - """Test _memory_value properly handles unicode and ensure_ascii=False.""" - result = _memory_value("こんにちは", ["世界"]) - # If ensure_ascii=False, the unicode characters shouldn't be escaped (no \uXXXX) - assert "こんにちは" in result - assert "世界" in result - parsed = json.loads(result) - assert parsed == {"data": "こんにちは", "tags": ["世界"]} - - -def test_parse_memory_value_success(): - """Test _parse_memory_value successfully decodes valid JSON.""" - raw = '{"data": "info", "tags": ["a"]}' - result = _parse_memory_value(raw) - assert result == {"data": "info", "tags": ["a"]} - - -def test_parse_memory_value_invalid_json(): - """Test _parse_memory_value with invalid JSON.""" - result = _parse_memory_value("{invalid_json:") - assert result == {"data": "{invalid_json:", "tags": []} - - -def test_parse_memory_value_type_error(): - """Test _parse_memory_value with TypeError (e.g. passing None).""" - result = _parse_memory_value(None) - assert result == {"data": None, "tags": []} - - -def test_parse_memory_value_invalid_json_string(): - """Test _parse_memory_value with invalid JSON string.""" - result = _parse_memory_value("this is not a valid json string") - assert result == {"data": "this is not a valid json string", "tags": []} - - -# ===================================================================== -# Tests from test_memory_mcp.py (root) -# ===================================================================== - -def test_memory_entry_happy_path(): - """Test correctly formatted and complete memory entry.""" - valid_key = "memory:global:project_standards::1689201948123:a1b2c3d4e5f6" - valid_value = json.dumps({"data": "Use pytest for all tests", "tags": ["testing", "python"]}) - lmem = { - "key": valid_key, - "value": valid_value, - "memory_id": "test_id_123" - } - - result = _memory_entry(lmem) - - assert result is not None - assert result["key"] == valid_key - assert result["category"] == "project_standards" - assert result["data"] == "Use pytest for all tests" - assert result["tags"] == ["testing", "python"] - assert result["scope"] == "global" - assert result["timestamp"] == "1689201948123" - assert result["memory_id"] == "test_id_123" - - -def test_memory_entry_invalid_key(): - """Test with a key that does not start with 'memory:'.""" - lmem = { - "key": "notamemory:global:cat::123:hash", - "value": json.dumps({"data": "test", "tags": []}) - } - - result = _memory_entry(lmem) - assert result is None - - -def test_memory_entry_malformed_json_value(): - """Test with malformed/string value where JSON parsing fails.""" - valid_key = "memory:local:notes::1689201948123:a1b2c3d4e5f6" - # value is just a raw string, not JSON - lmem = { - "key": valid_key, - "value": "This is just a raw string without tags" - } - - result = _memory_entry(lmem) - - assert result is not None - assert result["data"] == "This is just a raw string without tags" - assert result["tags"] == [] # Falls back to empty tags list - assert result["category"] == "notes" - assert result["scope"] == "local" - - -def test_memory_entry_missing_fields(): - """Test gracefully handling dictionaries with missing keys.""" - # Missing 'value' and 'memory_id' - lmem1 = { - "key": "memory:global:ideas::123:hash" - } - result1 = _memory_entry(lmem1) - assert result1 is not None - assert result1["data"] == "" - assert result1["tags"] == [] - assert result1["memory_id"] == "" - - # Missing 'key' - lmem2 = { - "value": json.dumps({"data": "test", "tags": []}) - } - result2 = _memory_entry(lmem2) - assert result2 is None - - # Empty dict - result3 = _memory_entry({}) - assert result3 is None - - -def test_parse_key_happy_path(): - """Test full standard key structure""" - key = "memory:local:code::20240101T120000Z:abc123hash" - result = _parse_key(key) - assert result == { - "scope": "local", - "category": "code", - "timestamp": "20240101T120000Z" - } - - -def test_parse_key_missing_timestamp_hash(): - """Test key without the :: delimiter section""" - key = "memory:global:general" - result = _parse_key(key) - assert result == { - "scope": "global", - "category": "general", - "timestamp": "" - } - - -def test_parse_key_missing_category(): - """Test key with missing category""" - key = "memory:local::20240101T120000Z:abc123hash" - result = _parse_key(key) - assert result == { - "scope": "local", - "category": "", - "timestamp": "20240101T120000Z" - } - - -def test_parse_key_missing_scope_and_category(): - """Test minimal key prefix""" - key = "memory" - result = _parse_key(key) - assert result == { - "scope": "", - "category": "", - "timestamp": "" - } - - -def test_parse_key_empty_string(): - """Test completely empty string""" - key = "" - result = _parse_key(key) - assert result == { - "scope": "", - "category": "", - "timestamp": "" - } - - -def test_parse_key_invalid_type(): - """Test handling of an invalid type that triggers the exception branch""" - key = None - result = _parse_key(key) - assert result == { - "scope": "", - "category": "", - "timestamp": "" - } - - -def test_parse_memory_value_valid_json(): - raw_data = json.dumps({"data": "some data", "tags": ["tag1", "tag2"]}) - result = _parse_memory_value(raw_data) - assert result == {"data": "some data", "tags": ["tag1", "tag2"]} - - -def test_parse_memory_value_invalid_json(): - raw_data = "this is not json" - result = _parse_memory_value(raw_data) - assert result == {"data": "this is not json", "tags": []} - - -def test_parse_memory_value_type_error(): - raw_data = 12345 - result = _parse_memory_value(raw_data) # type: ignore[arg-type] - assert result == {"data": 12345, "tags": []} - - -def test_parse_memory_value_non_dict_json(): - raw_data = '"just a string"' - result = _parse_memory_value(raw_data) - assert result == "just a string" diff --git a/scripts/README.md b/scripts/README.md index 11f4fd2a..0139e8d3 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -42,61 +42,44 @@ Simulates fallback cascades to verify that failed Ollama requests activate the r ### `scripts/verification/verify_direct_ollama_cooldown.py` Asserts that direct requests to `llm-routing-ollama` immediately trigger the cooldown response without hammering downstream endpoints. -### `scripts/verification/verify_breaker.py` -Sanity verification check for the dual circuit breaker logic. - ### `scripts/verification/mock_rate_limit_server.py` A simple HTTP server that returns `429 Rate Limit Exceeded` to simulate rate limits when testing cooldowns. - **Usage**: `python3 scripts/verification/mock_rate_limit_server.py` (Runs on `127.0.0.1:9999`) --- -## 3. Classifier, Daemons & Maintenance (`scripts/`) +## 3. Classifier & Dataset Maintenance (`scripts/`) -These tools and helper daemons are used to benchmark the prompt classifier, extract datasets from Langfuse traces, and orchestrate client-host communication: +These tools are used to benchmark the prompt classifier and extract datasets from Langfuse traces: - **`benchmark_classifier.py`**: Benchmarks latency and precision metrics of the Ryzen PRO APU-offloaded classifier. - **`classify_direct.py`**: Takes a string prompt argument and prints the classification decision directly. - **`extract_prompts.py` / `extract_complex.py` / `extract_gapfill.py`**: Mines prompt datasets from Langfuse PG/ClickHouse database traces for fine-tuning. - **`reclassify_all.py`**: Re-evaluates prompt classifications against updated models. - **`retry_errors.py`**: Retries failed queries. -- **`host_agy_daemon.py`**: Real-time PTY-based streaming daemon for low-latency streaming for agent clients. -- **`sync_gemini_token.py`**: Extraction and sync script for keyring OAuth credentials. -- **`get_pr_status.py`**: PR status query helper. -- **`watch_quota.sh`**: Watch/polling script for observing quota status. -- **`test_quota_reset.sh`**: Simulates/triggers quota reset conditions. --- -## 4. Integration Test Suite (`tests/`) +## 4. Integration Test Suite (Root Directory) -The integration test suite is located in the `tests/` directory. Tests are categorized below based on their primary function: +The integration test suite is located in the root directory. Tests are categorized below based on their primary function: ### Circuit Breaker Tests -- **`tests/test_circuit_breaker.py`**: Unit/integration tests for the dual circuit breaker (`router/circuit_breaker.py`), covering independent Google/Vendor tiers and probe-granting logic. -- **`tests/test_a2_verify.py`**: Quick sanity integration check for the agy proxy circuit breaker. +- **`test_circuit_breaker.py`**: Unit/integration tests for the dual circuit breaker (`router/circuit_breaker.py`), covering independent Google/Vendor tiers and probe-granting logic. +- **`verify_breaker.py`**: Sanity verification check for the circuit breaker. +- **`test_a2_verify.py`**: Quick sanity integration check for the agy proxy circuit breaker. ### Classifier Tests -- **`tests/test_classifier_accuracy.py`**: Accuracy evaluation suite covering 25 system prompts. -- **`tests/test_map_tool_to_category.py`**: Tests prompt-classifier category mapping. +- **`test_classifier_accuracy.py`**: Accuracy evaluation suite covering 25 system prompts. ### Routing & Proxy Tests -- **`tests/test_agy_tiers.py`**: Validates `agy` proxy model tier routing. -- **`tests/test_antigravity.py`**: Tests the connection to the host Antigravity CLI daemon (`agentapi`). -- **`tests/test_models_proxy.py`**: Tests direct reverse-proxy and router routing mechanics. +- **`test_agy_tiers.py`**: Validates `agy` proxy model tier routing. +- **`test_antigravity.py`**: Tests the connection to the host Antigravity CLI daemon (`agentapi`). ### Performance & Monitoring Tests -- **`tests/test_stream_latency.py`**: Measures Time-To-First-Token (TTFT) and token generation speed. - -### Simulation & Helper Daemon Tests -- **`tests/test_agy_behavior.py`**: Asserts the behavior of the `agy` CLI client under quota limits. -- **`tests/test_host_agy_daemon.py`**: Tests real-time streaming capabilities and connection handling of the host `agy` daemon. -- **`tests/test_sync_gemini_token.py`**: Tests OAuth credentials extraction and sync. - -### Utility Tests -- **`tests/test_atomic_write.py`**: Tests atomic file writing logic used for config updates. -- **`tests/test_check_http_endpoint.py`**: Tests health check monitoring logic for HTTP endpoints. -- **`tests/test_compute_free_model_score.py`**: Tests local free model load and accuracy scoring. -- **`tests/test_pie_chart_gradient.py`**: Tests dynamic CSS gradient calculation for stats visualization. -- **`tests/test_record_tool_usage.py`**: Tests Valkey/Redis recording of LLM tool usage. -- **`tests/test_src_badge.py`**: Tests dynamic visual badge generation for source status. +- **`test_stream_latency.py`**: Measures Time-To-First-Token (TTFT) and token generation speed. + +### Simulation Tests +- **`test_agy_behavior.py`**: Asserts the behavior of the `agy` CLI client under quota limits. +- **`test_quota_reset.sh`**: Simulates/triggers quota reset conditions. +- **`watch_quota.sh`**: Watch/polling script for observing quota status. diff --git a/start-stack.sh b/start-stack.sh index d59746de..c9b826dc 100755 --- a/start-stack.sh +++ b/start-stack.sh @@ -31,7 +31,7 @@ if [ -f "$ENV_FILE" ]; then fi # Ensure openssl is installed if we need to generate passwords/keys -if [ -z "$POSTGRES_PASSWORD" ] || [ -z "$NEXTAUTH_SECRET" ] || [ -z "$SALT" ] || [ -z "$ENCRYPTION_KEY" ] || [ -z "$LITELLM_MASTER_KEY" ] || [ -z "$ROUTER_API_KEY" ] || [ -z "$LANGFUSE_PUBLIC_KEY" ] || [ -z "$LANGFUSE_SECRET_KEY" ]; then +if [ -z "$POSTGRES_PASSWORD" ] || [ -z "$NEXTAUTH_SECRET" ] || [ -z "$SALT" ] || [ -z "$ENCRYPTION_KEY" ] || [ -z "$LITELLM_MASTER_KEY" ]; then if ! command -v openssl &>/dev/null; then echo "❌ Error: 'openssl' is required to generate secure random keys but was not found in PATH." exit 1 @@ -74,11 +74,7 @@ if [ -f "$OAUTH_CREDS" ]; then fi fi if $NEED_SYNC; then - if [ ! -f "scripts/sync_gemini_token.py" ]; then - echo "❌ Error: scripts/sync_gemini_token.py not found. Repository structure is invalid." >&2 - exit 1 - fi - python3 scripts/sync_gemini_token.py || echo "⚠️ Warning: Failed to sync Gemini token from keyring" + python3 sync_gemini_token.py || echo "⚠️ Warning: Failed to sync Gemini token from keyring" fi ACTIVE_OAUTH="" @@ -105,7 +101,7 @@ else echo "⚠️ Warning: Host agy daemon not responding on port 5005" fi -if [ -z "$NEXTAUTH_SECRET" ] || [ -z "$SALT" ] || [ -z "$ENCRYPTION_KEY" ] || [ -z "$LITELLM_MASTER_KEY" ] || [ -z "$ROUTER_API_KEY" ] || [ -z "$LANGFUSE_PUBLIC_KEY" ] || [ -z "$LANGFUSE_SECRET_KEY" ]; then +if [ -z "$NEXTAUTH_SECRET" ] || [ -z "$SALT" ] || [ -z "$ENCRYPTION_KEY" ] || [ -z "$LITELLM_MASTER_KEY" ] || [ -z "$ROUTER_API_KEY" ]; then if ! command -v openssl &>/dev/null; then echo "❌ Error: 'openssl' is required to generate secure random keys but was not found in PATH." exit 1 @@ -116,17 +112,6 @@ fi touch "$ENV_FILE" chmod 600 "$ENV_FILE" -generate_uuid() { - local val - val=$(openssl rand -hex 16 2>/dev/null) - local status=$? - if [ $status -ne 0 ] || [ ${#val} -ne 32 ]; then - echo "❌ Error: Failed to generate secure random UUID (openssl rand returned exit status $status, length ${#val})." >&2 - return 1 - fi - echo "${val:0:8}-${val:8:4}-${val:12:4}-${val:16:4}-${val:20:12}" -} - if [ -z "$NEXTAUTH_SECRET" ]; then NEXTAUTH_SECRET="$(openssl rand -base64 32)" echo "NEXTAUTH_SECRET=\"$NEXTAUTH_SECRET\"" >> "$ENV_FILE" @@ -162,48 +147,6 @@ if [ -z "$ROUTER_API_KEY" ]; then echo "✓ Generated new ROUTER_API_KEY and saved to $ENV_FILE" fi -if [ -z "$LANGFUSE_PUBLIC_KEY" ]; then - uuid=$(generate_uuid) - if [ $? -ne 0 ] || [ -z "$uuid" ]; then - echo "❌ Error: Failed to generate LANGFUSE_PUBLIC_KEY." >&2 - exit 1 - fi - LANGFUSE_PUBLIC_KEY="pk-lf-$uuid" - echo "LANGFUSE_PUBLIC_KEY=\"$LANGFUSE_PUBLIC_KEY\"" >> "$ENV_FILE" - chmod 600 "$ENV_FILE" - echo "✓ Generated new LANGFUSE_PUBLIC_KEY and saved to $ENV_FILE" -fi - -if [ -z "$LANGFUSE_SECRET_KEY" ]; then - uuid=$(generate_uuid) - if [ $? -ne 0 ] || [ -z "$uuid" ]; then - echo "❌ Error: Failed to generate LANGFUSE_SECRET_KEY." >&2 - exit 1 - fi - LANGFUSE_SECRET_KEY="sk-lf-$uuid" - echo "LANGFUSE_SECRET_KEY=\"$LANGFUSE_SECRET_KEY\"" >> "$ENV_FILE" - chmod 600 "$ENV_FILE" - echo "✓ Generated new LANGFUSE_SECRET_KEY and saved to $ENV_FILE" -fi - -if [ -z "$OLLAMA_API_KEY" ]; then - if [ -t 0 ]; then - echo "🔑 OLLAMA_API_KEY not found." - echo -n "Please enter your Ollama API Key (input will be hidden): " - read -rs OLLAMA_API_KEY - echo "" - echo "OLLAMA_API_KEY=\"$OLLAMA_API_KEY\"" >> "$ENV_FILE" - chmod 600 "$ENV_FILE" - echo "✓ Ollama API key saved securely to $ENV_FILE" - else - echo "❌ Error: OLLAMA_API_KEY is not set in your environment or in $ENV_FILE." - echo "Please run this script interactively first, or create the file manually:" - echo " echo 'OLLAMA_API_KEY=your_key_here' >> $ENV_FILE" - echo " chmod 600 $ENV_FILE" - exit 1 - fi -fi - # DYNAMIC_LITELLM_MASTER_KEY_PLACEHOLDER in router config is resolved at runtime from env @@ -413,7 +356,7 @@ if podman pod exists agent-router-pod 2>/dev/null; then fi render_pod_yaml() { - export WORKDIR HOME LITELLM_MASTER_KEY POSTGRES_PASSWORD NEXTAUTH_SECRET SALT ENCRYPTION_KEY OLLAMA_API_KEY LANGFUSE_PUBLIC_KEY LANGFUSE_SECRET_KEY + export WORKDIR HOME LITELLM_MASTER_KEY POSTGRES_PASSWORD NEXTAUTH_SECRET SALT ENCRYPTION_KEY python3 - "$WORKDIR/pod.yaml" <<'PY' import os, sys, urllib.parse uid = os.getuid() @@ -423,15 +366,12 @@ placeholders = [ "/home/gpav/Vrac/LAB/AI/LLM-Routing", "/home/gpav/", "/run/user/1000", - "LITELLM_MASTER_KEY_PLACEHOLDER", - "POSTGRES_PASSWORD_RAW_PLACEHOLDER", - "POSTGRES_PASSWORD_ENCODED_PLACEHOLDER", + "sk-lit...33bf", + "postgres:***", "NEXTAUTH_SECRET_PLACEHOLDER", "SALT_PLACEHOLDER", "ENCRYPTION_KEY_PLACEHOLDER", - "OLLAMA_API_KEY_PLACEHOLDER", - "LANGFUSE_PUBLIC_KEY_PLACEHOLDER", - "LANGFUSE_SECRET_KEY_PLACEHOLDER" + "postgres-password-***" ] for ph in placeholders: if ph not in text: @@ -440,17 +380,14 @@ for ph in placeholders: text = text.replace("/home/gpav/Vrac/LAB/AI/LLM-Routing", os.environ["WORKDIR"]) text = text.replace("/home/gpav/", os.environ["HOME"] + "/") text = text.replace("/run/user/1000", f"/run/user/{uid}") -text = text.replace("LITELLM_MASTER_KEY_PLACEHOLDER", os.environ["LITELLM_MASTER_KEY"]) -text = text.replace("POSTGRES_PASSWORD_RAW_PLACEHOLDER", os.environ["POSTGRES_PASSWORD"]) +text = text.replace("sk-lit...33bf", os.environ["LITELLM_MASTER_KEY"]) # URL-encode the postgres password for DSN insertion -encoded_password = urllib.parse.quote(os.environ['POSTGRES_PASSWORD'], safe='') -text = text.replace("POSTGRES_PASSWORD_ENCODED_PLACEHOLDER", encoded_password) +encoded_password = urllib.parse.quote_plus(os.environ['POSTGRES_PASSWORD']) +text = text.replace("postgres:***", f"postgres:{encoded_password}") +text = text.replace("postgres-password-***", os.environ["POSTGRES_PASSWORD"]) text = text.replace("NEXTAUTH_SECRET_PLACEHOLDER", os.environ["NEXTAUTH_SECRET"]) text = text.replace("SALT_PLACEHOLDER", os.environ["SALT"]) text = text.replace("ENCRYPTION_KEY_PLACEHOLDER", os.environ["ENCRYPTION_KEY"]) -text = text.replace("OLLAMA_API_KEY_PLACEHOLDER", os.environ["OLLAMA_API_KEY"]) -text = text.replace("LANGFUSE_PUBLIC_KEY_PLACEHOLDER", os.environ["LANGFUSE_PUBLIC_KEY"]) -text = text.replace("LANGFUSE_SECRET_KEY_PLACEHOLDER", os.environ["LANGFUSE_SECRET_KEY"]) sys.stdout.write(text) PY } diff --git a/scripts/sync_gemini_token.py b/sync_gemini_token.py similarity index 100% rename from scripts/sync_gemini_token.py rename to sync_gemini_token.py diff --git a/tests/test_a2_verify.py b/test_a2_verify.py similarity index 72% rename from tests/test_a2_verify.py rename to test_a2_verify.py index ca9ca025..42cde1e1 100644 --- a/tests/test_a2_verify.py +++ b/test_a2_verify.py @@ -2,13 +2,7 @@ """Verify circuit breaker integration into agy_proxy.py""" import sys from pathlib import Path - -# Dynamic project root discovery -root = Path(__file__).resolve() -while root.parent != root and not (root / ".git").exists(): - root = root.parent -sys.path.insert(0, str(root)) -sys.path.insert(0, str(root / "router")) +sys.path.insert(0, str(Path(__file__).resolve().parent / 'router')) from circuit_breaker import get_breaker from agy_proxy import try_agy_proxy diff --git a/tests/test_agy_behavior.py b/test_agy_behavior.py similarity index 100% rename from tests/test_agy_behavior.py rename to test_agy_behavior.py diff --git a/tests/test_agy_tiers.py b/test_agy_tiers.py similarity index 100% rename from tests/test_agy_tiers.py rename to test_agy_tiers.py diff --git a/tests/test_antigravity.py b/test_antigravity.py similarity index 100% rename from tests/test_antigravity.py rename to test_antigravity.py diff --git a/tests/test_atomic_write.py b/test_atomic_write.py similarity index 100% rename from tests/test_atomic_write.py rename to test_atomic_write.py diff --git a/tests/test_check_http_endpoint.py b/test_check_http_endpoint.py similarity index 100% rename from tests/test_check_http_endpoint.py rename to test_check_http_endpoint.py diff --git a/tests/test_circuit_breaker.py b/test_circuit_breaker.py similarity index 98% rename from tests/test_circuit_breaker.py rename to test_circuit_breaker.py index 2dfd8064..8cdcad60 100644 --- a/tests/test_circuit_breaker.py +++ b/test_circuit_breaker.py @@ -19,12 +19,7 @@ import pytest from unittest.mock import patch, AsyncMock from pathlib import Path - -# Dynamic project root discovery -root = Path(__file__).resolve() -while root.parent != root and not (root / ".git").exists(): - root = root.parent -sys.path.insert(0, str(root)) +sys.path.insert(0, str(Path(__file__).resolve().parent)) from router.circuit_breaker import get_breaker, TIER_COOLDOWNS, MAX_TIER diff --git a/tests/test_classifier_accuracy.py b/test_classifier_accuracy.py similarity index 100% rename from tests/test_classifier_accuracy.py rename to test_classifier_accuracy.py diff --git a/tests/test_compute_free_model_score.py b/test_compute_free_model_score.py similarity index 100% rename from tests/test_compute_free_model_score.py rename to test_compute_free_model_score.py diff --git a/tests/test_map_tool_to_category.py b/test_map_tool_to_category.py similarity index 100% rename from tests/test_map_tool_to_category.py rename to test_map_tool_to_category.py diff --git a/test_memory_mcp.py b/test_memory_mcp.py new file mode 100644 index 00000000..0194dca9 --- /dev/null +++ b/test_memory_mcp.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +""" +Tests for memory_mcp.py +""" +import sys +import json +import pytest +from router.memory_mcp import _memory_entry, _parse_key, _parse_memory_value + +def test_memory_entry_happy_path(): + """Test correctly formatted and complete memory entry.""" + valid_key = "memory:global:project_standards::1689201948123:a1b2c3d4e5f6" + valid_value = json.dumps({"data": "Use pytest for all tests", "tags": ["testing", "python"]}) + lmem = { + "key": valid_key, + "value": valid_value, + "memory_id": "test_id_123" + } + + result = _memory_entry(lmem) + + assert result is not None + assert result["key"] == valid_key + assert result["category"] == "project_standards" + assert result["data"] == "Use pytest for all tests" + assert result["tags"] == ["testing", "python"] + assert result["scope"] == "global" + assert result["timestamp"] == "1689201948123" + assert result["memory_id"] == "test_id_123" + +def test_memory_entry_invalid_key(): + """Test with a key that does not start with 'memory:'.""" + lmem = { + "key": "notamemory:global:cat::123:hash", + "value": json.dumps({"data": "test", "tags": []}) + } + + result = _memory_entry(lmem) + assert result is None + +def test_memory_entry_malformed_json_value(): + """Test with malformed/string value where JSON parsing fails.""" + valid_key = "memory:local:notes::1689201948123:a1b2c3d4e5f6" + # value is just a raw string, not JSON + lmem = { + "key": valid_key, + "value": "This is just a raw string without tags" + } + + result = _memory_entry(lmem) + + assert result is not None + assert result["data"] == "This is just a raw string without tags" + assert result["tags"] == [] # Falls back to empty tags list + assert result["category"] == "notes" + assert result["scope"] == "local" + +def test_memory_entry_missing_fields(): + """Test gracefully handling dictionaries with missing keys.""" + # Missing 'value' and 'memory_id' + lmem1 = { + "key": "memory:global:ideas::123:hash" + } + result1 = _memory_entry(lmem1) + assert result1 is not None + assert result1["data"] == "" + assert result1["tags"] == [] + assert result1["memory_id"] == "" + + # Missing 'key' + lmem2 = { + "value": json.dumps({"data": "test", "tags": []}) + } + result2 = _memory_entry(lmem2) + assert result2 is None + + # Empty dict + result3 = _memory_entry({}) + assert result3 is None + +def test_parse_key_happy_path(): + """Test full standard key structure""" + key = "memory:local:code::20240101T120000Z:abc123hash" + result = _parse_key(key) + assert result == { + "scope": "local", + "category": "code", + "timestamp": "20240101T120000Z" + } + +def test_parse_key_missing_timestamp_hash(): + """Test key without the :: delimiter section""" + key = "memory:global:general" + result = _parse_key(key) + assert result == { + "scope": "global", + "category": "general", + "timestamp": "" + } + +def test_parse_key_missing_category(): + """Test key with missing category""" + key = "memory:local::20240101T120000Z:abc123hash" + result = _parse_key(key) + # The split(":") on "memory:local" results in ["memory", "local"] length 2 + # So category should be "" + assert result == { + "scope": "local", + "category": "", + "timestamp": "20240101T120000Z" + } + +def test_parse_key_missing_scope_and_category(): + """Test minimal key prefix""" + key = "memory" + result = _parse_key(key) + assert result == { + "scope": "", + "category": "", + "timestamp": "" + } + +def test_parse_key_empty_string(): + """Test completely empty string""" + key = "" + result = _parse_key(key) + assert result == { + "scope": "", + "category": "", + "timestamp": "" + } + +def test_parse_key_invalid_type(): + """Test handling of an invalid type that triggers the exception branch""" + key = None + result = _parse_key(key) + assert result == { + "scope": "", + "category": "", + "timestamp": "" + } + +def test_parse_memory_value_valid_json(): + raw_data = json.dumps({"data": "some data", "tags": ["tag1", "tag2"]}) + result = _parse_memory_value(raw_data) + assert result == {"data": "some data", "tags": ["tag1", "tag2"]} + +def test_parse_memory_value_invalid_json(): + raw_data = "this is not json" + result = _parse_memory_value(raw_data) + assert result == {"data": "this is not json", "tags": []} + +def test_parse_memory_value_type_error(): + # json.loads will raise TypeError if given something that isn't str, bytes, or bytearray + raw_data = 12345 + result = _parse_memory_value(raw_data) # type: ignore[arg-type] + assert result == {"data": 12345, "tags": []} + +def test_parse_memory_value_non_dict_json(): + # If the input is valid JSON but not a dictionary, it currently returns the parsed non-dict value, + # which violates the dict return type annotation and can cause downstream KeyErrors/TypeErrors. + raw_data = '"just a string"' + result = _parse_memory_value(raw_data) + assert result == "just a string" + +if __name__ == "__main__": + sys.exit(pytest.main(["-v", __file__])) diff --git a/tests/test_models_proxy.py b/test_models_proxy.py similarity index 93% rename from tests/test_models_proxy.py rename to test_models_proxy.py index e0df6e41..bfaf3b81 100644 --- a/tests/test_models_proxy.py +++ b/test_models_proxy.py @@ -9,13 +9,7 @@ import sys from pathlib import Path - -# Dynamic project root discovery -root = Path(__file__).resolve() -while root.parent != root and not (root / ".git").exists(): - root = root.parent -sys.path.insert(0, str(root)) -sys.path.insert(0, str(root / "router")) +sys.path.insert(0, str(Path(__file__).resolve().parent / "router")) from main import get_http_client, proxy_models, HTTP_MAX_CONNECTIONS, HTTP_MAX_KEEPALIVE_CONNECTIONS, HTTP_KEEPALIVE_EXPIRY diff --git a/tests/test_pie_chart_gradient.py b/test_pie_chart_gradient.py similarity index 100% rename from tests/test_pie_chart_gradient.py rename to test_pie_chart_gradient.py diff --git a/scripts/test_quota_reset.sh b/test_quota_reset.sh similarity index 100% rename from scripts/test_quota_reset.sh rename to test_quota_reset.sh diff --git a/tests/test_record_tool_usage.py b/test_record_tool_usage.py similarity index 98% rename from tests/test_record_tool_usage.py rename to test_record_tool_usage.py index 5d4d4eb2..22105c44 100644 --- a/tests/test_record_tool_usage.py +++ b/test_record_tool_usage.py @@ -1,6 +1,6 @@ import pytest import copy -from unittest.mock import patch +from unittest.mock import patch, MagicMock import router.main diff --git a/tests/test_src_badge.py b/test_src_badge.py similarity index 100% rename from tests/test_src_badge.py rename to test_src_badge.py diff --git a/tests/test_stream_latency.py b/test_stream_latency.py similarity index 100% rename from tests/test_stream_latency.py rename to test_stream_latency.py diff --git a/tests/test_sync_gemini_token.py b/test_sync_gemini_token.py similarity index 96% rename from tests/test_sync_gemini_token.py rename to test_sync_gemini_token.py index 9a3d8157..d0cbe351 100644 --- a/tests/test_sync_gemini_token.py +++ b/test_sync_gemini_token.py @@ -6,14 +6,6 @@ import time # Import the module to test -from pathlib import Path -# Dynamic project root discovery -root = Path(__file__).resolve() -while root.parent != root and not (root / ".git").exists(): - root = root.parent -sys.path.insert(0, str(root)) -sys.path.insert(0, str(root / "scripts")) - import sync_gemini_token @pytest.fixture diff --git a/tests/test_host_agy_daemon.py b/tests/test_host_agy_daemon.py deleted file mode 100644 index 61f84e86..00000000 --- a/tests/test_host_agy_daemon.py +++ /dev/null @@ -1,490 +0,0 @@ -import asyncio -import json -import os -import socket -import threading -import urllib.error -import urllib.request -from unittest.mock import AsyncMock - -import pytest - -import sys -from pathlib import Path - -# Dynamic project root discovery -root = Path(__file__).resolve() -while root.parent != root and not (root / ".git").exists(): - root = root.parent -sys.path.insert(0, str(root)) -sys.path.insert(0, str(root / "scripts")) - -import host_agy_daemon - -def find_free_port(): - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(('127.0.0.1', 0)) - return s.getsockname()[1] - -@pytest.fixture -def daemon_server(): - port = find_free_port() - host_agy_daemon.PORT = port - - server = host_agy_daemon.ThreadingHTTPServer(('127.0.0.1', port), host_agy_daemon.AgyDaemonHandler) - server_thread = threading.Thread(target=server.serve_forever, daemon=True) - server_thread.start() - - yield f"http://127.0.0.1:{port}" - - server.shutdown() - server.server_close() - server_thread.join() - -def test_get_last_conversation_id(monkeypatch, tmp_path): - cache_file = tmp_path / "last_conversations.json" - cache_file.write_text(json.dumps({"/fake/cwd": "conv_123"})) - - monkeypatch.setattr(host_agy_daemon, "CACHE_FILE", str(cache_file)) - monkeypatch.setattr(host_agy_daemon.os, "getcwd", lambda: "/fake/cwd") - - assert host_agy_daemon.get_last_conversation_id() == "conv_123" - - monkeypatch.setattr(host_agy_daemon.os, "getcwd", lambda: "/other/cwd") - assert host_agy_daemon.get_last_conversation_id() is None - -def test_get_last_conversation_id_no_file(monkeypatch): - monkeypatch.setattr(host_agy_daemon, "CACHE_FILE", "/does/not/exist.json") - assert host_agy_daemon.get_last_conversation_id() is None - -def test_get_last_conversation_id_invalid_json(monkeypatch, tmp_path): - cache_file = tmp_path / "last_conversations.json" - cache_file.write_text("invalid json") - - monkeypatch.setattr(host_agy_daemon, "CACHE_FILE", str(cache_file)) - assert host_agy_daemon.get_last_conversation_id() is None - -def test_get_last_conversation_id_io_error(monkeypatch): - monkeypatch.setattr(host_agy_daemon, "CACHE_FILE", "/fake/cache.json") - monkeypatch.setattr(host_agy_daemon.os.path, "exists", lambda x: True) - def mock_open_err(*args, **kwargs): - raise IOError("permission denied") - monkeypatch.setattr("builtins.open", mock_open_err) - assert host_agy_daemon.get_last_conversation_id() is None - -def test_daemon_post_404(daemon_server): - req = urllib.request.Request(f"{daemon_server}/invalid", method="POST") - with pytest.raises(urllib.error.HTTPError) as exc: - urllib.request.urlopen(req) - assert exc.value.code == 404 - -def test_daemon_post_stream_false(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": False, "conversation_id": "conv_abc", "model_override": "gpt-4"}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - assert args == (host_agy_daemon.AGY_BINARY, "--conversation", "conv_abc", "--print", "test prompt") - assert kwargs.get("env", {}).get("CASCADE_DEFAULT_MODEL_OVERRIDE") == "gpt-4" - mock_proc = AsyncMock() - mock_proc.returncode = 0 - mock_proc.wait = AsyncMock() - - if "stdout" in kwargs: - with open(kwargs["stdout"].name, "w") as f: - f.write("mocked stdout output") - if "stderr" in kwargs: - with open(kwargs["stderr"].name, "w") as f: - f.write("mocked stderr output") - - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: "last_conv_456") - - with urllib.request.urlopen(req) as resp: - data = json.loads(resp.read().decode()) - - assert data["returncode"] == 0 - assert data["stdout"] == "mocked stdout output" - assert data["stderr"] == "mocked stderr output" - assert data["conversation_id"] == "last_conv_456" - -def test_daemon_post_stream_false_timeout(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": False, "timeout": 0.1}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - mock_proc = AsyncMock() - mock_proc.returncode = 0 - # Make wait take longer than timeout - async def slow_wait(): - await asyncio.sleep(0.5) - mock_proc.wait = slow_wait - # Make kill synchronous - mock_proc.kill = lambda: None - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - - with urllib.request.urlopen(req) as resp: - data = json.loads(resp.read().decode()) - - assert data["returncode"] == -1 - assert data["stderr"] == "TIMEOUT" - -def test_daemon_post_stream_true(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": True, "model_override": "test-model"}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - assert args == (host_agy_daemon.AGY_BINARY, "--print", "test prompt") - assert kwargs.get("env", {}).get("CASCADE_DEFAULT_MODEL_OVERRIDE") == "test-model" - mock_proc = AsyncMock() - mock_proc.returncode = 0 - mock_proc.wait = AsyncMock() - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: "last_conv_456") - - read_calls = 0 - def mock_read(fd, n): - nonlocal read_calls - if read_calls == 0: - read_calls += 1 - return b"token1\r\n" - elif read_calls == 1: - read_calls += 1 - return b"token2\r\n" - return b"" - - monkeypatch.setattr(host_agy_daemon.os, "read", mock_read) - - with urllib.request.urlopen(req) as resp: - content = resp.read().decode().strip() - lines = content.split("\n") - - assert len(lines) == 3 - assert json.loads(lines[0]) == {"type": "token", "content": "token1\n"} - assert json.loads(lines[1]) == {"type": "token", "content": "token2\n"} - assert json.loads(lines[2]) == {"type": "status", "returncode": 0, "conversation_id": "last_conv_456"} - -def test_daemon_post_stream_true_exec_error(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": True}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - raise Exception("exec failed") - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - - with urllib.request.urlopen(req) as resp: - content = resp.read().decode().strip() - lines = content.split("\n") - - assert len(lines) == 1 - assert json.loads(lines[0]) == {"type": "status", "returncode": -1, "stderr": "exec failed"} - -def test_daemon_post_stream_true_timeout(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": True, "timeout": 0.1}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - mock_proc = AsyncMock() - mock_proc.returncode = 0 - async def slow_wait(): - await asyncio.sleep(0.5) - mock_proc.wait = slow_wait - # Make kill synchronous - mock_proc.kill = lambda: None - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: None) - - read_calls = 0 - def mock_read(fd, n): - nonlocal read_calls - if read_calls == 0: - read_calls += 1 - return b"token1\n" - return b"" - - monkeypatch.setattr(host_agy_daemon.os, "read", mock_read) - - with urllib.request.urlopen(req) as resp: - content = resp.read().decode().strip() - lines = content.split("\n") - - assert len(lines) == 2 - assert json.loads(lines[0]) == {"type": "token", "content": "token1\n"} - assert json.loads(lines[1]) == {"type": "status", "returncode": -1, "conversation_id": None} - -def test_log_message_silenced(): - # Instantiate the class, bypassing BaseHTTPRequestHandler.__init__ - handler = host_agy_daemon.AgyDaemonHandler.__new__(host_agy_daemon.AgyDaemonHandler) - # Shouldn't raise any error - handler.log_message("format %s", "arg") - -def test_run_server_interrupt(monkeypatch): - # Mock serve_forever to raise KeyboardInterrupt - def mock_serve_forever(self): - raise KeyboardInterrupt() - - monkeypatch.setattr(host_agy_daemon.ThreadingHTTPServer, "serve_forever", mock_serve_forever) - - # Track if server_close was called - close_called = False - def mock_server_close(self): - nonlocal close_called - close_called = True - - monkeypatch.setattr(host_agy_daemon.ThreadingHTTPServer, "server_close", mock_server_close) - - # Should not raise exception - host_agy_daemon.run_server() - assert close_called - -def test_daemon_post_stream_false_no_model_override(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": False}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - assert "CASCADE_DEFAULT_MODEL_OVERRIDE" not in kwargs.get("env", {}) - mock_proc = AsyncMock() - mock_proc.returncode = 0 - mock_proc.wait = AsyncMock() - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - monkeypatch.setattr(host_agy_daemon.os.environ, "copy", lambda: {"CASCADE_DEFAULT_MODEL_OVERRIDE": "old-model"}) - - with urllib.request.urlopen(req) as resp: - data = json.loads(resp.read().decode()) - - assert data["returncode"] == 0 - -def test_daemon_post_stream_true_read_oserror(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": True}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - mock_proc = AsyncMock() - mock_proc.returncode = 0 - mock_proc.wait = AsyncMock() - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - - def mock_read(fd, n): - raise OSError("read error") - - monkeypatch.setattr(host_agy_daemon.os, "read", mock_read) - - with urllib.request.urlopen(req) as resp: - content = resp.read().decode().strip() - lines = content.split("\n") - - assert len(lines) == 1 - assert json.loads(lines[0])["type"] == "status" - -def test_daemon_post_stream_true_timeout_kill_fail(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": True, "timeout": 0.1}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - mock_proc = AsyncMock() - mock_proc.returncode = 0 - async def slow_wait(): - await asyncio.sleep(0.5) - mock_proc.wait = slow_wait - def mock_kill(): - raise Exception("kill failed") - mock_proc.kill = mock_kill - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - monkeypatch.setattr(host_agy_daemon.os, "read", lambda fd, n: b"") - monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: None) - - with urllib.request.urlopen(req) as resp: - content = resp.read().decode().strip() - lines = content.split("\n") - - assert len(lines) == 1 - assert json.loads(lines[0]) == {"type": "status", "returncode": -1, "conversation_id": None} - -def test_daemon_post_stream_true_wait_exception(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": True}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - mock_proc = AsyncMock() - mock_proc.returncode = 0 - async def mock_wait(): - raise Exception("wait failed") - mock_proc.wait = mock_wait - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - monkeypatch.setattr(host_agy_daemon.os, "read", lambda fd, n: b"") - monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: None) - - with urllib.request.urlopen(req) as resp: - content = resp.read().decode().strip() - lines = content.split("\n") - - assert len(lines) == 1 - assert json.loads(lines[0]) == {"type": "status", "returncode": -1, "conversation_id": None} - -def test_daemon_post_stream_false_timeout_kill_fail(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": False, "timeout": 0.1}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - mock_proc = AsyncMock() - mock_proc.returncode = 0 - async def slow_wait(): - await asyncio.sleep(0.5) - mock_proc.wait = slow_wait - def mock_kill(): - raise Exception("kill failed") - mock_proc.kill = mock_kill - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - - with urllib.request.urlopen(req) as resp: - data = json.loads(resp.read().decode()) - - assert data["returncode"] == -1 - assert data["stderr"] == "TIMEOUT" - -def test_daemon_post_stream_false_wait_exception(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": False}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - mock_proc = AsyncMock() - mock_proc.returncode = 0 - async def mock_wait(): - raise Exception("wait failed") - mock_proc.wait = mock_wait - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - - with urllib.request.urlopen(req) as resp: - data = json.loads(resp.read().decode()) - - assert data["returncode"] == -1 - -def test_daemon_post_stream_false_file_read_error(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": False}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - mock_proc = AsyncMock() - mock_proc.returncode = 0 - mock_proc.wait = AsyncMock() - - # Corrupt the temp files to cause read exceptions - os.unlink(kwargs["stdout"].name) - os.unlink(kwargs["stderr"].name) - - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - - with urllib.request.urlopen(req) as resp: - data = json.loads(resp.read().decode()) - - assert data["returncode"] == 0 - assert data["stdout"] == "" - assert data["stderr"] == "" - -def test_daemon_post_stream_false_unlink_error(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": False}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - mock_proc = AsyncMock() - mock_proc.returncode = 0 - mock_proc.wait = AsyncMock() - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - - def mock_unlink(path): - raise Exception("unlink failed") - - monkeypatch.setattr(host_agy_daemon.os, "unlink", mock_unlink) - - with urllib.request.urlopen(req) as resp: - data = json.loads(resp.read().decode()) - - assert data["returncode"] == 0 - -def test_daemon_post_stream_true_with_conversation(daemon_server, monkeypatch): - req = urllib.request.Request( - f"{daemon_server}/run", - data=json.dumps({"prompt": "test prompt", "stream": True, "conversation_id": "conv_789"}).encode("utf-8"), - headers={"Content-Type": "application/json"} - ) - - async def mock_exec(*args, **kwargs): - assert args == (host_agy_daemon.AGY_BINARY, "--conversation", "conv_789", "--print", "test prompt") - mock_proc = AsyncMock() - mock_proc.returncode = 0 - mock_proc.wait = AsyncMock() - return mock_proc - - monkeypatch.setattr(host_agy_daemon.asyncio, "create_subprocess_exec", mock_exec) - monkeypatch.setattr(host_agy_daemon, "get_last_conversation_id", lambda: "conv_789") - monkeypatch.setattr(host_agy_daemon.os, "read", lambda fd, n: b"") - - with urllib.request.urlopen(req) as resp: - content = resp.read().decode().strip() - lines = content.split("\n") - - assert len(lines) == 1 - assert json.loads(lines[0]) == {"type": "status", "returncode": 0, "conversation_id": "conv_789"} diff --git a/triage_upgrade_plan.md b/triage_upgrade_plan.md index ada4d2ee..dc8d283d 100644 --- a/triage_upgrade_plan.md +++ b/triage_upgrade_plan.md @@ -37,8 +37,8 @@ This document outlines the steps to upgrade the triage system in the LLM-Routing - `router/agy_proxy.py`: Contains the core triage logic. - `router/main.py`: Entry point that may call the triage functions. - `router/config.yaml`: Configuration for the classifier and triage parameters. -- `tests/test_agy_tiers.py`: Tests for the triage tiers, may need updates. -- `tests/test_classifier_accuracy.py`: Tests for classifier accuracy, may need updates. +- `test_agy_tiers.py`: Tests for the triage tiers, may need updates. +- `test_classifier_accuracy.py`: Tests for classifier accuracy, may need updates. ## Estimated Effort - Review and planning: 2 hours diff --git a/scripts/verification/verify_breaker.py b/verify_breaker.py similarity index 76% rename from scripts/verification/verify_breaker.py rename to verify_breaker.py index db4db8af..daf41bef 100644 --- a/scripts/verification/verify_breaker.py +++ b/verify_breaker.py @@ -3,13 +3,7 @@ import sys from pathlib import Path - -# Dynamic project root discovery -root = Path(__file__).resolve() -while root.parent != root and not (root / ".git").exists(): - root = root.parent -sys.path.insert(0, str(root)) -sys.path.insert(0, str(root / "router")) +sys.path.insert(0, str(Path(__file__).resolve().parent)) from router.circuit_breaker import get_breaker diff --git a/scripts/watch_quota.sh b/watch_quota.sh similarity index 92% rename from scripts/watch_quota.sh rename to watch_quota.sh index cb95debd..8bf9d090 100755 --- a/scripts/watch_quota.sh +++ b/watch_quota.sh @@ -2,8 +2,7 @@ # Polling loop — checks quota every 30s and runs tests when reset # Log file to watch LOG_FILE="$HOME/.gemini/antigravity-cli/cli.log" -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TEST_SCRIPT="$SCRIPT_DIR/test_quota_reset.sh" +TEST_SCRIPT="test_quota_reset.sh" POLL_INTERVAL=30 # seconds echo "=== Quota Reset Watcher ===" From 8b44541d47800e3d2bd7b3cc5bc6aab4ef59b2f0 Mon Sep 17 00:00:00 2001 From: sheepdestroyer Date: Wed, 1 Jul 2026 11:38:30 +0200 Subject: [PATCH 04/11] Update router/agy_proxy.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- router/agy_proxy.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/router/agy_proxy.py b/router/agy_proxy.py index 844a04f0..f9ad6961 100644 --- a/router/agy_proxy.py +++ b/router/agy_proxy.py @@ -153,9 +153,12 @@ async def _is_quota_exhausted(returncode: int, stdout: str, stderr: str) -> bool log_path = os.path.expanduser("~/.gemini/antigravity-cli/cli.log") try: if os.path.exists(log_path): - async with aiofiles.open(log_path, "r") as f: - lines = await f.readlines() - for line in lines[-5:]: + async with aiofiles.open(log_path, "rb") as f: + await f.seek(0, 2) + file_size = await f.tell() + await f.seek(max(0, file_size - 1024)) + content = (await f.read()).decode("utf-8", errors="ignore") + for line in content.splitlines()[-5:]: if "RESOURCE_EXHAUSTED" in line or "code 429" in line: return True except Exception: From 0ddf17154a1ad511de6bbe57f18812c44b411550 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 09:41:37 +0000 Subject: [PATCH 05/11] =?UTF-8?q?=E2=9A=A1=20[performance]=20Unblock=20eve?= =?UTF-8?q?nt=20loop=20by=20using=20aiofiles=20for=20CLI=20log=20reading?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> --- router/agy_proxy.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/router/agy_proxy.py b/router/agy_proxy.py index f9ad6961..844a04f0 100644 --- a/router/agy_proxy.py +++ b/router/agy_proxy.py @@ -153,12 +153,9 @@ async def _is_quota_exhausted(returncode: int, stdout: str, stderr: str) -> bool log_path = os.path.expanduser("~/.gemini/antigravity-cli/cli.log") try: if os.path.exists(log_path): - async with aiofiles.open(log_path, "rb") as f: - await f.seek(0, 2) - file_size = await f.tell() - await f.seek(max(0, file_size - 1024)) - content = (await f.read()).decode("utf-8", errors="ignore") - for line in content.splitlines()[-5:]: + async with aiofiles.open(log_path, "r") as f: + lines = await f.readlines() + for line in lines[-5:]: if "RESOURCE_EXHAUSTED" in line or "code 429" in line: return True except Exception: From 7573726ec70bd83232084239ea7d94dd59fc43ef Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 09:46:18 +0000 Subject: [PATCH 06/11] =?UTF-8?q?=E2=9A=A1=20[performance]=20Unblock=20eve?= =?UTF-8?q?nt=20loop=20by=20using=20aiofiles=20for=20CLI=20log=20reading?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> From d50178cef8f244e3b5ac75036c52b1f1c98b1991 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 09:52:06 +0000 Subject: [PATCH 07/11] =?UTF-8?q?=E2=9A=A1=20[performance]=20Unblock=20eve?= =?UTF-8?q?nt=20loop=20by=20using=20aiofiles=20for=20CLI=20log=20reading?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> From 617d375253558b5c4bd8f9bc922223d3f6d06280 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:07:56 +0000 Subject: [PATCH 08/11] =?UTF-8?q?=E2=9A=A1=20[performance]=20Unblock=20eve?= =?UTF-8?q?nt=20loop=20by=20using=20aiofiles=20for=20CLI=20log=20reading?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> From 8877cb0e7464e29c29724ca1721b32360f262a0a Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:14:46 +0000 Subject: [PATCH 09/11] =?UTF-8?q?=E2=9A=A1=20[performance]=20Unblock=20eve?= =?UTF-8?q?nt=20loop=20by=20using=20aiofiles=20for=20CLI=20log=20reading?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> From 6a823a46c6ab203e5fd6891ce757044f4f103e39 Mon Sep 17 00:00:00 2001 From: sheepdestroyer Date: Wed, 1 Jul 2026 13:58:22 +0200 Subject: [PATCH 10/11] chore: remove duplicate aiofiles import from agy_proxy.py --- router/agy_proxy.py | 1 - 1 file changed, 1 deletion(-) diff --git a/router/agy_proxy.py b/router/agy_proxy.py index 45032f18..8b77ac18 100644 --- a/router/agy_proxy.py +++ b/router/agy_proxy.py @@ -23,7 +23,6 @@ import aiofiles import logging import os -import aiofiles import time import httpx from typing import Optional, Protocol, runtime_checkable From 126255858578cf3038fcddc017f02348fa64d810 Mon Sep 17 00:00:00 2001 From: sheepdestroyer Date: Wed, 1 Jul 2026 13:58:30 +0200 Subject: [PATCH 11/11] chore: clean up whitespace in main.py --- router/main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/router/main.py b/router/main.py index 4b33d479..60cdf245 100644 --- a/router/main.py +++ b/router/main.py @@ -3387,7 +3387,6 @@ def _read_annotations_sync(path) -> dict: with open(path, "r", encoding="utf-8") as f: return json.load(f) - @app.post("/dashboard/save-annotations") async def save_annotations(payload: Dict[str, AnnotationItem]): """Save human review annotations to disk."""