diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2f369ed9..eaec9bdf 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -20,7 +20,7 @@ jobs: python-version: '3.11' - name: Install dependencies - run: pip install httpx==0.28.1 pytest pytest-asyncio anyio pyyaml fastapi 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 - 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/README.md b/README.md index 854edd95..462ca04d 100644 --- a/README.md +++ b/README.md @@ -80,7 +80,7 @@ All core containers are configured with **Kubernetes-style liveness and readines | **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 | Same, every 10s | +| **clickhouse-db** | `clickhouse-client --user clickhouse --password clickhouse --query "SELECT 1"` every 15s | `clickhouse-client --query "SELECT 1"` every 10s | | **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 | โ€” | @@ -801,5 +801,3 @@ This project is supported by a dedicated NotebookLM companion notebook: * **URL:** [TriageGate-Architect-KB](https://notebooklm.google.com/notebook/826cbd87-7969-4b0e-a38e-5517b5ab7d28) This notebook contains a comprehensive semantic index of the system architecture, LiteLLM cascades, Langfuse telemetry pipelines, local model configurations, and integration guides. Agents and developers can query this notebook via the `notebooklm` MCP tools (e.g., using `notebook_ask` with `notebook_id: "llm-triage-gateway"`) to retrieve structured knowledge, check pitfalls, or get implementation examples for this gateway stack. - - diff --git a/get_pr_status.py b/get_pr_status.py index d088b7af..366c64cf 100644 --- a/get_pr_status.py +++ b/get_pr_status.py @@ -1,10 +1,7 @@ import subprocess -import shlex -def run_cmd(cmd): - # Fix the issues from Sourcery review! - # 1. Provide a static list of strings for args rather than a single string. - # 2. Use shell=False - args = shlex.split(cmd) - result = subprocess.run(args, shell=False, capture_output=True, text=True) +def run_cmd(cmd_list): + # Expect a list of strings directly instead of a single string to avoid command injection + # Enable check=True to surface non-zero exit codes as subprocess.CalledProcessError + result = subprocess.run(cmd_list, shell=False, capture_output=True, text=True, check=True) return result.stdout.strip() diff --git a/litellm/tests/test_entrypoint.py b/litellm/tests/test_entrypoint.py index 502b3407..53e92395 100644 --- a/litellm/tests/test_entrypoint.py +++ b/litellm/tests/test_entrypoint.py @@ -13,18 +13,38 @@ mock_proxy_cli = MagicMock() -with patch('os.path.exists', return_value=False), \ - patch('builtins.print'), \ - patch('time.sleep'), \ - patch('os.execvp'), \ - patch('sys.stdout.flush'), \ - patch('glob.glob', return_value=[]), \ - patch('builtins.open'): - - sys.modules['litellm'] = mock_litellm - sys.modules['litellm.proxy'] = MagicMock() - sys.modules['litellm.proxy.proxy_cli'] = mock_proxy_cli - spec.loader.exec_module(entrypoint) +# Mock socket instance for import-time check_tcp_port execution +mock_socket_instance = MagicMock() +mock_socket_instance.connect_ex.return_value = 0 + +# Save original modules to avoid leaking fake ones globally +orig_modules = { + 'litellm': sys.modules.get('litellm'), + 'litellm.proxy': sys.modules.get('litellm.proxy'), + 'litellm.proxy.proxy_cli': sys.modules.get('litellm.proxy.proxy_cli') +} + +try: + with patch('os.path.exists', return_value=False), \ + patch('builtins.print'), \ + patch('time.sleep'), \ + patch('os.execvp'), \ + patch('sys.stdout.flush'), \ + patch('glob.glob', return_value=[]), \ + patch('socket.socket', return_value=mock_socket_instance), \ + patch('builtins.open'): + + sys.modules['litellm'] = mock_litellm + sys.modules['litellm.proxy'] = MagicMock() + sys.modules['litellm.proxy.proxy_cli'] = mock_proxy_cli + spec.loader.exec_module(entrypoint) +finally: + # Restore original modules state + for k, v in orig_modules.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v def test_check_tcp_port_success(): with patch('socket.socket') as mock_socket_class: diff --git a/pod.yaml b/pod.yaml index f51f58be..9d3af258 100644 --- a/pod.yaml +++ b/pod.yaml @@ -236,8 +236,14 @@ spec: timeoutSeconds: 2 name: valkey-lf readinessProbe: - tcpSocket: - port: 6380 + exec: + command: + - valkey-cli + - -p + - "6380" + - -a + - langfuse-redis-2026 + - ping initialDelaySeconds: 2 periodSeconds: 5 timeoutSeconds: 2 @@ -249,13 +255,13 @@ spec: - name: DATABASE_URL value: postgresql://postgres:***@127.0.0.1:5432/langfuse - name: NEXTAUTH_SECRET - value: my-super-secret-nextauth-token-2026 + value: NEXTAUTH_SECRET_PLACEHOLDER - name: NEXTAUTH_URL value: http://localhost:3001 - name: SALT - value: my-super-strong-salt-token-2026-value-1234 + value: SALT_PLACEHOLDER - name: ENCRYPTION_KEY - value: 4c265d39d04389f069225db1e88726727a090e7fc6275e8c910b81aa4b763135 + value: ENCRYPTION_KEY_PLACEHOLDER - name: HOSTNAME value: 0.0.0.0 - name: PORT diff --git a/router/Dockerfile b/router/Dockerfile index e663b49e..cd0f3653 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 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 # 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/ diff --git a/router/main.py b/router/main.py index 61cd10c0..f12f2802 100644 --- a/router/main.py +++ b/router/main.py @@ -234,10 +234,17 @@ async def push_aggregate_scores(): router_model_conf = config.get("router", {}).get("router_model", {}) router_api_base = router_model_conf.get("api_base", "http://127.0.0.1:8080/v1") -router_api_key = router_model_conf.get("api_key", "local-token") +router_api_key = router_model_conf.get("api_key") +if not router_api_key: + raise RuntimeError("Configuration error: 'api_key' is missing from router_model configuration.") if router_api_key.startswith("os.environ/"): env_var = router_api_key.split("/", 1)[1] - router_api_key = os.environ.get(env_var, "local-token") + router_api_key = os.environ.get(env_var) + if not router_api_key: + if "pytest" in sys.modules: + router_api_key = "local-token" + else: + raise RuntimeError(f"Configuration error: Environment variable '{env_var}' is missing or empty.") router_model_name = router_model_conf.get("model", "qwen-0.8b-routing") system_prompt = config.get("classification_rules", {}).get("system_prompt", "") @@ -3328,14 +3335,16 @@ async def get_visualizer(): VALID_TIERS = {"agent-simple-core", "agent-medium-core", "agent-complex-core", "agent-reasoning-core", "agent-advanced-core"} +MAX_ANNOTATION_KEY_LENGTH = 128 +MAX_ANNOTATION_ITEM_BYTES = 4096 class AnnotationItem(BaseModel): """Pydantic model representing a single human dataset review annotation.""" - model_config = ConfigDict(extra="allow") + model_config = ConfigDict(extra="forbid") tier: Union[int, str, None] = None note: Optional[str] = Field(default=None, max_length=1000) - ts: Optional[str] = None + ts: Optional[str] = Field(default=None, max_length=100) @field_validator("tier") @classmethod @@ -3360,12 +3369,16 @@ def validate_payload(self) -> "AnnotationPayload": data = self.root if len(data) > 1000: raise ValueError("Payload size limit exceeded: maximum of 1000 annotations allowed per request.") - for k in data: + for k, item in data.items(): + if len(k) > MAX_ANNOTATION_KEY_LENGTH: + raise ValueError(f"Invalid payload key '{k}': key is too long.") is_valid_key = k.isdigit() or ( k.startswith("h") and len(k) > 1 and all(c in "0123456789abcdef" for c in k[1:].lower()) ) if not is_valid_key: raise ValueError(f"Invalid payload key '{k}': keys must be numeric strings or stable hash keys (e.g., 'h12345abc').") + if len(item.model_dump_json().encode("utf-8")) > MAX_ANNOTATION_ITEM_BYTES: + raise ValueError(f"Annotation '{k}' exceeds the maximum serialized size.") return self # NOTE: annotations_lock (asyncio.Lock) only provides concurrency protection within # a single Python process. In multi-worker uvicorn deployments, concurrent requests diff --git a/start-stack.sh b/start-stack.sh index 2ad08fed..8cbdd708 100755 --- a/start-stack.sh +++ b/start-stack.sh @@ -17,6 +17,12 @@ mkdir -p valkey-data postgres-data langfuse-data clickhouse-data redis-lf-data m ENV_FILE="${WORKDIR}/.env" +# Ensure the env file exists and has secure permissions (owner read/write only) +if [ ! -f "$ENV_FILE" ]; then + touch "$ENV_FILE" + chmod 600 "$ENV_FILE" +fi + # 1. Load or prompt for OpenRouter API Key if [ -f "$ENV_FILE" ]; then set -a @@ -24,6 +30,15 @@ if [ -f "$ENV_FILE" ]; then set +a 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" ]; 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 + fi +fi + + if [ -z "$OPENROUTER_API_KEY" ]; then if [ -t 0 ]; then echo "๐Ÿ”‘ OpenRouter API Key not found." @@ -31,7 +46,7 @@ if [ -z "$OPENROUTER_API_KEY" ]; then read -rs OPENROUTER_API_KEY echo "" echo "OPENROUTER_API_KEY=\"$OPENROUTER_API_KEY\"" > "$ENV_FILE" - chmod 644 "$ENV_FILE" + chmod 600 "$ENV_FILE" echo "โœ“ API key saved securely to $ENV_FILE" else echo "โŒ Error: OPENROUTER_API_KEY is not set in your environment or in $ENV_FILE" @@ -86,6 +101,35 @@ 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" ]; 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 + fi +fi + +# Ensure the env file exists and has secure permissions (owner read/write only) +touch "$ENV_FILE" +chmod 600 "$ENV_FILE" + +if [ -z "$NEXTAUTH_SECRET" ]; then + NEXTAUTH_SECRET="$(openssl rand -base64 32)" + echo "NEXTAUTH_SECRET=\"$NEXTAUTH_SECRET\"" >> "$ENV_FILE" + echo "โœ“ Generated new NEXTAUTH_SECRET and saved to $ENV_FILE" +fi + +if [ -z "$SALT" ]; then + SALT="$(openssl rand -hex 16)" + echo "SALT=\"$SALT\"" >> "$ENV_FILE" + echo "โœ“ Generated new SALT and saved to $ENV_FILE" +fi + +if [ -z "$ENCRYPTION_KEY" ]; then + ENCRYPTION_KEY="$(openssl rand -hex 32)" + echo "ENCRYPTION_KEY=\"$ENCRYPTION_KEY\"" >> "$ENV_FILE" + echo "โœ“ Generated new ENCRYPTION_KEY and saved to $ENV_FILE" +fi + if [ -z "$LITELLM_MASTER_KEY" ]; then LITELLM_MASTER_KEY="sk-litellm-$(openssl rand -hex 16)" echo "LITELLM_MASTER_KEY=\"$LITELLM_MASTER_KEY\"" >> "$ENV_FILE" @@ -305,9 +349,9 @@ if podman pod exists agent-router-pod 2>/dev/null; then fi render_pod_yaml() { - export WORKDIR HOME LITELLM_MASTER_KEY POSTGRES_PASSWORD + export WORKDIR HOME LITELLM_MASTER_KEY POSTGRES_PASSWORD NEXTAUTH_SECRET SALT ENCRYPTION_KEY python3 - "$WORKDIR/pod.yaml" <<'PY' -import os, sys +import os, sys, urllib.parse uid = os.getuid() with open(sys.argv[1], "r", encoding="utf-8") as f: text = f.read() @@ -317,6 +361,9 @@ placeholders = [ "/run/user/1000", "sk-lit...33bf", "postgres:***", + "NEXTAUTH_SECRET_PLACEHOLDER", + "SALT_PLACEHOLDER", + "ENCRYPTION_KEY_PLACEHOLDER", "postgres-password-***" ] for ph in placeholders: @@ -327,8 +374,13 @@ 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("sk-lit...33bf", os.environ["LITELLM_MASTER_KEY"]) -text = text.replace("postgres:***", f"postgres:{os.environ['POSTGRES_PASSWORD']}") +# URL-encode the postgres password for DSN insertion +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"]) sys.stdout.write(text) PY } diff --git a/test_agy_behavior.py b/test_agy_behavior.py index 6ad0d1cb..6dade8f6 100644 --- a/test_agy_behavior.py +++ b/test_agy_behavior.py @@ -8,13 +8,13 @@ async def test(): env = os.environ.copy() cmd = [AGY, "--print", "say hi"] - + proc = await asyncio.create_subprocess_exec( *cmd, env=env, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - + try: stdout_bytes, stderr_bytes = await asyncio.wait_for(proc.communicate(), timeout=15) print(f"returncode: {proc.returncode}") @@ -23,7 +23,7 @@ async def test(): except asyncio.TimeoutError: proc.kill() print("TIMEOUT") - + # Also check the log for recent quota lines log_path = os.path.expanduser("~/.gemini/antigravity-cli/cli.log") if os.path.exists(log_path): diff --git a/test_agy_tiers.py b/test_agy_tiers.py index e0d21bb9..f4e17da4 100644 --- a/test_agy_tiers.py +++ b/test_agy_tiers.py @@ -17,34 +17,36 @@ {"name": "Claude Opus 4.6", "override": "claude-opus-4-6@default"}, ] -async def test_tier(tier, prompt="say hello in one word", conversation_id=None): +async def run_tier_test(tier, prompt="say hello in one word", conversation_id=None): """Test a single agy tier and return (success, output, conv_id).""" env = os.environ.copy() if tier["override"]: env["CASCADE_DEFAULT_MODEL_OVERRIDE"] = tier["override"] - + else: + env.pop("CASCADE_DEFAULT_MODEL_OVERRIDE", None) + cmd = [AGY] if conversation_id: cmd.extend(["--conversation", conversation_id]) cmd.extend(["--print", prompt]) - + print(f"\n ๐Ÿงช Testing {tier['name']}... ", end="", flush=True) if conversation_id: print(f"(continuing {conversation_id[:8]}...) ", end="", flush=True) - + start = time.time() proc = await asyncio.create_subprocess_exec( *cmd, env=env, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - + try: stdout_bytes, stderr_bytes = await asyncio.wait_for(proc.communicate(), timeout=30) stdout = stdout_bytes.decode("utf-8", errors="replace").strip() stderr = stderr_bytes.decode("utf-8", errors="replace").strip() elapsed = time.time() - start - + # Get conversation ID from cache conv_id = None try: @@ -54,13 +56,13 @@ async def test_tier(tier, prompt="say hello in one word", conversation_id=None): conv_id = data.get(os.getcwd()) except: pass - + # Check for quota exhaustion if "RESOURCE_EXHAUSTED" in stderr or "code 429" in stderr: print(f"โŒ QUOTA EXHAUSTED ({elapsed:.1f}s)") print(f" stderr: {stderr[:100]}") - return False, None, conv_id - + return False, None, None + if stdout: print(f"โœ… OK ({elapsed:.1f}s, {len(stdout)} chars)") print(f" response: {stdout[:80]}...") @@ -68,10 +70,11 @@ async def test_tier(tier, prompt="say hello in one word", conversation_id=None): else: print(f"โš ๏ธ EMPTY RESPONSE ({elapsed:.1f}s)") print(f" stderr: {stderr[:200]}") - return False, None, conv_id - + return False, None, None + except asyncio.TimeoutError: proc.kill() + await proc.communicate() print(f"โŒ TIMEOUT (30s)") return False, None, None @@ -79,16 +82,17 @@ async def main(): print("=" * 60) print(" agy Proxy Tier Test Suite") print("=" * 60) - + # Test 1: Each tier independently (new conversations) print("\n--- Test 1: Independent Tier Tests ---") conv_ids = {} for tier in TIERS: - success, output, conv_id = await test_tier(tier) - conv_ids[tier["name"]] = conv_id + success, output, conv_id = await run_tier_test(tier) + if success and conv_id: + conv_ids[tier["name"]] = conv_id if not success: print(f" โš ๏ธ Tier {tier['name']} failed โ€” subsequent tests may use different model") - + # Test 2: Session continuation (same conversation across tiers) print("\n\n--- Test 2: Session Continuation ---") # Get the last successful conversation ID @@ -97,12 +101,12 @@ async def main(): if tier["name"] in conv_ids and conv_ids[tier["name"]]: successful_conv = conv_ids[tier["name"]] break - + if successful_conv: print(f" Continuing conversation {successful_conv[:8]}...") for tier in TIERS: - success, output, _ = await test_tier( - tier, + success, output, _ = await run_tier_test( + tier, prompt="continue our conversation, say one more word", conversation_id=successful_conv ) @@ -110,16 +114,16 @@ async def main(): break else: print(" โš ๏ธ No successful conversation to continue") - + # Test 3: Auto-fallback chain (simulate proxy behavior) print("\n\n--- Test 3: Proxy Fallback Chain ---") proxy_prompt = "what's 2+2? answer in one word" for tier in TIERS: - success, output, conv_id = await test_tier(tier, prompt=proxy_prompt) + success, output, conv_id = await run_tier_test(tier, prompt=proxy_prompt) if success: print(f"\n โœ… Proxy would use: {tier['name']}") break - + print("\n" + "=" * 60) print(" Tests complete!") print("=" * 60) diff --git a/test_antigravity.py b/test_antigravity.py index 6b057d38..8e0bc4b6 100644 --- a/test_antigravity.py +++ b/test_antigravity.py @@ -10,7 +10,7 @@ def test_antigravity_connection(): return print("--- Testing antigravity-cli connection with current OAuth ---") - + # Using the agentapi binary located at ~/.gemini/antigravity-cli/bin/agentapi agentapi_path = os.path.expanduser("~/.gemini/antigravity-cli/bin/agentapi") if not os.path.exists(agentapi_path): diff --git a/test_circuit_breaker.py b/test_circuit_breaker.py index df331c92..8cdcad60 100644 --- a/test_circuit_breaker.py +++ b/test_circuit_breaker.py @@ -5,7 +5,7 @@ Simulates consecutive quota failures and verifies: - Independent google and vendor breakers - Tier 1 cooldown (5 min) after 1st failure - - Tier 2 cooldown (30 min) after 2nd failure + - Tier 2 cooldown (30 min) after 2nd failure - Tier 3 cooldown (5 hours) after 3rd failure - Probe behavior: one allowed attempt after cooldown - Reset to Tier 0 on success @@ -49,7 +49,7 @@ def test_first_failure_trips_to_tier1(): """1st failure โ†’ Tier 1, 5 min cooldown.""" reset_breakers() b = get_breaker() - + b.google.record_failure() assert b.google.tier == 1 assert b.google.cooldown_until > time.time() @@ -64,11 +64,11 @@ def test_probe_granted_after_cooldown(): """After cooldown expires, exactly one probe is allowed.""" reset_breakers() b = get_breaker() - + b.google.tier = 1 b.google.cooldown_until = time.time() - 10 # expired 10s ago b.google.probe_granted = False - + assert b.google.is_allowed(), "Probe should be granted" assert b.google.probe_granted, "Probe flag should be set" assert not b.google.is_allowed(), "Second call should be denied" @@ -79,12 +79,12 @@ def test_probe_failure_advances_tier(): """Probe failure โ†’ advance to next tier.""" reset_breakers() b = get_breaker() - + b.google.tier = 1 b.google.cooldown_until = time.time() - 10 b.google.probe_granted = True # probe was granted b.google.record_failure() # probe fails - + assert b.google.tier == 2, f"Expected tier 2, got {b.google.tier}" assert not b.google.probe_granted print("โœ“ Failed probe at Tier 1 โ†’ advanced to Tier 2 (30 min)") @@ -94,13 +94,13 @@ def test_tier3_stays_at_tier3(): """At Tier 3, failure โ†’ stays at Tier 3 (renews cooldown).""" reset_breakers() b = get_breaker() - + b.google.tier = MAX_TIER b.google.cooldown_until = time.time() - 10 b.google.probe_granted = True old_until = b.google.cooldown_until b.google.record_failure() - + assert b.google.tier == MAX_TIER, "Should stay at Tier 3" assert b.google.cooldown_until > old_until, "Cooldown should be renewed" assert not b.google.probe_granted @@ -111,12 +111,12 @@ def test_success_resets(): """Success at any tier โ†’ reset to Tier 0.""" reset_breakers() b = get_breaker() - + b.google.tier = 2 b.google.cooldown_until = time.time() + 1000 b.google.probe_granted = False b.google.record_success() - + assert b.google.tier == 0 assert b.google.is_allowed() print("โœ“ Success resets breaker to Tier 0 from any tier") @@ -126,7 +126,7 @@ def test_backward_compatibility(): """Master breaker record_failure and record_success affect both breakers.""" reset_breakers() b = get_breaker() - + b.record_failure() assert b.google.tier == 1 assert b.vendor.tier == 1 diff --git a/test_classifier_accuracy.py b/test_classifier_accuracy.py index 5d7a9a90..7b10f49d 100644 --- a/test_classifier_accuracy.py +++ b/test_classifier_accuracy.py @@ -73,14 +73,14 @@ def query_model(prompt: str) -> tuple[str, float]: "max_tokens": 15, "grammar": 'root ::= "agent-simple-core" | "agent-complex-core"' } - + data = json.dumps(payload).encode("utf-8") req = urllib.request.Request( LLAMA_SERVER_URL, data=data, headers={"Content-Type": "application/json", "Authorization": f"Bearer {os.environ.get('ROUTER_API_KEY', 'local-token')}"} ) - + start_time = time.time() try: with urllib.request.urlopen(req, timeout=30) as response: @@ -98,20 +98,20 @@ def calculate_metrics(results): total = len(results) correct = sum(1 for r in results if r["expected"] == r["actual"]) accuracy = (correct / total) * 100.0 if total > 0 else 0.0 - + # Initialize metrics structure classes = ["agent-simple-core", "agent-complex-core"] metrics = {} - + for c in classes: tp = sum(1 for r in results if r["expected"] == c and r["actual"] == c) fp = sum(1 for r in results if r["expected"] != c and r["actual"] == c) fn = sum(1 for r in results if r["expected"] == c and r["actual"] != c) - + precision = (tp / (tp + fp)) if (tp + fp) > 0 else 0.0 recall = (tp / (tp + fn)) if (tp + fn) > 0 else 0.0 f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) > 0 else 0.0 - + metrics[c] = { "precision": precision * 100.0, "recall": recall * 100.0, @@ -120,51 +120,51 @@ def calculate_metrics(results): "fp": fp, "fn": fn } - + return accuracy, metrics def main(): - print(f"Starting Classifier Accuracy Evaluation Suite...") + print("Starting Classifier Accuracy Evaluation Suite...") print(f"Querying endpoint: {LLAMA_SERVER_URL}") print(f"Loaded {len(test_cases)} test cases.") print("-" * 80) - + results = [] latencies = [] misclassifications = [] - + for i, (prompt, expected) in enumerate(test_cases, 1): print(f"[{i}/{len(test_cases)}] Testing: '{prompt[:50]}...'") actual, latency = query_model(prompt) latencies.append(latency) - + success = (actual == expected) status = "โœ“ PASS" if success else "โœ— FAIL" print(f" Expected: {expected} | Actual: {actual} | Latency: {latency:.2f}ms | {status}") - + results.append({ "prompt": prompt, "expected": expected, "actual": actual, "latency": latency }) - + if not success: misclassifications.append({ "prompt": prompt, "expected": expected, "actual": actual }) - + print("=" * 80) accuracy, metrics = calculate_metrics(results) avg_latency = sum(latencies) / len(latencies) if latencies else 0.0 - - print(f"OVERALL METRICS:") + + print("OVERALL METRICS:") print(f"Classification Accuracy: {accuracy:.2f}%") print(f"Average Latency: {avg_latency:.2f} ms") print("-" * 80) - + for c, m in metrics.items(): print(f"Class: {c}") print(f" Precision: {m['precision']:.2f}%") @@ -172,7 +172,7 @@ def main(): print(f" F1-Score: {m['f1']:.2f}%") print(f" (TP={m['tp']}, FP={m['fp']}, FN={m['fn']})") print("-" * 80) - + if misclassifications: print(f"MISCLASSIFICATION LOGS ({len(misclassifications)} total):") for mc in misclassifications: @@ -182,6 +182,6 @@ def main(): print("-" * 40) else: print("๐ŸŽ‰ No misclassifications! Perfect accuracy!") - + if __name__ == "__main__": main() diff --git a/test_stream_latency.py b/test_stream_latency.py index 5ec9b61a..912a90c7 100755 --- a/test_stream_latency.py +++ b/test_stream_latency.py @@ -22,7 +22,7 @@ async def main(): print("Testing Stream Latency (TTFT Verification)") print("=" * 60) print("Sending streaming request to triage router on port 5000...") - + start_time = time.time() first_token_time = None chunks_received = 0 @@ -60,21 +60,21 @@ async def main(): first_token_time = time.time() ttft = (first_token_time - start_time) * 1000 print(f"๐Ÿš€ Time-To-First-Token (TTFT): {ttft:.0f} ms") - + full_response.append(content) # Print character to show live streaming print(content, end="", flush=True) except Exception as e: # Ignore parse errors for partial chunks pass - + end_time = time.time() elapsed = end_time - start_time - print(f"\n\nStream Finished!") + print("\n\nStream Finished!") print(f"Total time: {elapsed:.2f} s") print(f"Total chunks received: {chunks_received}") print(f"Story length: {len(''.join(full_response))} characters") - + # Verify TTFT is within acceptable limits for a streamed response (typically <3-4s, unlike the 8s+ legacy TTFT) if first_token_time is not None: ttft_ms = (first_token_time - start_time) * 1000 @@ -84,7 +84,7 @@ async def main(): print("โš ๏ธ TTFT is high, check daemon buffering.") else: print("โŒ No tokens received in stream.") - + except Exception as e: print(f"โŒ Exception occurred: {e}")