{llamacpp_models_html}
@@ -3310,7 +3779,7 @@ async def get_dashboard():
-
{src_badge('GOOSE', '#fbbf24')} Session Directory
+
{src_badge("GOOSE", "#fbbf24")} Session Directory
{goose_html}
@@ -3322,21 +3791,21 @@ async def get_dashboard():
@@ -3352,6 +3821,7 @@ async def get_dashboard():
"""
return html_content
+
# --- Static files (visualizer, data files) ---
STATIC_DIR = Path(__file__).resolve().parent / "static"
DATA_DIR = Path(__file__).resolve().parent / "data"
@@ -3359,6 +3829,7 @@ async def get_dashboard():
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
app.mount("/data", StaticFiles(directory=str(DATA_DIR)), name="data")
+
@app.get("/visualizer", response_class=HTMLResponse)
async def get_visualizer():
"""Serve the dataset visualizer for human review."""
@@ -3369,13 +3840,52 @@ async def get_visualizer():
return HTMLResponse("
Visualizer not found
", status_code=404)
+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."""
- tier: Union[int, str, None] = None
- note: str = ""
- ts: Optional[str] = None
+ model_config = ConfigDict(extra="forbid")
-VALID_TIERS = {"agent-simple-core", "agent-medium-core", "agent-complex-core", "agent-reasoning-core", "agent-advanced-core"}
+ tier: Union[int, str, None] = None
+ note: Optional[str] = Field(default=None, max_length=1000)
+ ts: Optional[str] = Field(default=None, max_length=100)
+
+ @field_validator("tier")
+ @classmethod
+ def validate_tier(cls, v):
+ if v is None:
+ return v
+ if isinstance(v, int):
+ if v < 0 or v > 4:
+ raise ValueError(f"Invalid tier index {v}: must be between 0 and 4")
+ elif isinstance(v, str):
+ if v not in VALID_TIERS and v != "?":
+ raise ValueError(f"Invalid tier string '{v}'")
+ else:
+ raise ValueError("Tier must be int, str, or null")
+ return v
+
+class AnnotationPayload(RootModel):
+ root: Dict[str, AnnotationItem]
+
+ @model_validator(mode="after")
+ 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, 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
# across different workers can still race. Eventual consistency is maintained via
@@ -3400,73 +3910,44 @@ def _read_annotations_sync(path) -> dict:
return copy.deepcopy(_annotations_cache[path]["data"])
+
@app.post("/dashboard/save-annotations")
-async def save_annotations(payload: Dict[str, AnnotationItem]):
+async def save_annotations(payload: AnnotationPayload):
"""Save human review annotations to disk."""
- if len(payload) > 1000:
- raise HTTPException(
- status_code=400,
- detail="Payload size limit exceeded: maximum of 1000 annotations allowed per request."
- )
- for k, item in payload.items():
- # Allow numeric strings (dataset indexes) or stable hash keys starting with 'h' (hexadecimal)
- 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 HTTPException(
- status_code=400,
- detail=f"Invalid payload key '{k}': keys must be numeric strings or stable hash keys (e.g., 'h12345abc')."
- )
-
- t = item.tier
- if t is not None:
- if isinstance(t, int):
- if t < 0 or t > 4:
- raise HTTPException(
- status_code=400,
- detail=f"Invalid tier index {t} for index {k}: must be between 0 and 4."
- )
- elif isinstance(t, str):
- if t not in VALID_TIERS and t != "?":
- raise HTTPException(
- status_code=400,
- detail=f"Invalid tier string '{t}' for index {k}."
- )
- else:
- raise HTTPException(
- status_code=400,
- detail=f"Invalid tier type for index {k}: must be int, str, or null."
- )
-
- if item.note and len(item.note) > 1000:
- raise HTTPException(
- status_code=400,
- detail=f"Note length limit exceeded at index {k}: maximum of 1000 characters allowed."
- )
try:
+ data = payload.root
ann_path = DATA_DIR / "annotations.json"
existing = {}
async with annotations_lock:
if ann_path.exists():
try:
- existing = await asyncio.to_thread(_read_annotations_sync, str(ann_path))
+ existing = await asyncio.to_thread(
+ _read_annotations_sync, str(ann_path)
+ )
except Exception as read_err:
- logger.warning(f"Could not read existing annotations: {read_err}. Overwriting.")
+ logger.warning(
+ f"Could not read existing annotations: {read_err}. Overwriting."
+ )
# Merge new annotations into existing
- for k, item in payload.items():
- existing[k] = item.model_dump() if hasattr(item, "model_dump") else item.dict()
-
+ for k, item in data.items():
+ # For partial updates, merge only fields provided in the request
+ update_data = item.model_dump(exclude_unset=True)
+ if k in existing and isinstance(existing[k], dict):
+ existing[k].update(update_data)
+ else:
+ existing[k] = item.model_dump()
await _atomic_write_json_async(str(ann_path), existing)
- return JSONResponse({"status": "ok", "saved": len(payload)})
+ return JSONResponse({"status": "ok", "saved": len(data)})
except Exception as e:
logger.error(f"Failed to save annotations: {e}")
raise HTTPException(status_code=500, detail="Failed to save annotations")
+
if __name__ == "__main__":
import uvicorn
+
logger.info(f"Starting LLM Triage Router on {host}:{port}...")
uvicorn.run(app, host=host, port=port)
diff --git a/router/test_memory_mcp.py b/router/test_memory_mcp.py
deleted file mode 100644
index 24d23a34..00000000
--- a/router/test_memory_mcp.py
+++ /dev/null
@@ -1,131 +0,0 @@
-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 404059eb..3358439e 100644
--- a/router/tests/test_agy_proxy.py
+++ b/router/tests/test_agy_proxy.py
@@ -1,3 +1,13 @@
+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
from router.agy_proxy import _wrap_response, _is_quota_exhausted
diff --git a/router/tests/test_dashboard_data.py b/router/tests/test_dashboard_data.py
index 643425f8..6c960d4d 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, MagicMock
+from unittest.mock import AsyncMock, patch # Removed unused MagicMock import
import sys
import os
diff --git a/router/tests/test_detect_active_tool.py b/router/tests/test_detect_active_tool.py
deleted file mode 100644
index 3105ab1e..00000000
--- a/router/tests/test_detect_active_tool.py
+++ /dev/null
@@ -1,114 +0,0 @@
-import pytest
-import os
-import sys
-from pathlib import Path
-
-# Set CONFIG_PATH for import
-os.environ["CONFIG_PATH"] = str(Path(__file__).resolve().parent.parent / "config.yaml")
-
-# Add the parent directory to the path so we can import from router
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-
-from router.main import detect_active_tool
-
-def test_detect_active_tool_empty():
- assert detect_active_tool({}) == "none"
- assert detect_active_tool({"messages": []}) == "none"
- assert detect_active_tool({"messages": [{"role": "system", "content": "hello"}]}) == "none"
-
-def test_detect_active_tool_role_tool_with_name():
- # Write tool name mapped to write
- body = {
- "messages": [
- {"role": "tool", "name": "edit_file", "content": "..."}
- ]
- }
- assert detect_active_tool(body) == "write"
-
- # View tool mapped to view
- body = {
- "messages": [
- {"role": "tool", "name": "cat_file", "content": "..."}
- ]
- }
- assert detect_active_tool(body) == "view"
-
-def test_detect_active_tool_role_tool_without_name_but_matched_tool_call_id():
- body = {
- "messages": [
- {"role": "assistant", "tool_calls": [{"id": "call_123", "function": {"name": "read_file"}}]},
- {"role": "tool", "tool_call_id": "call_123", "content": "success"}
- ]
- }
- assert detect_active_tool(body) == "view"
-
-def test_detect_active_tool_role_tool_without_name_unmatched_tool_call_id():
- body = {
- "messages": [
- {"role": "assistant", "tool_calls": [{"id": "call_999", "function": {"name": "read_file"}}]},
- {"role": "tool", "tool_call_id": "call_123", "content": "success"}
- ]
- }
- assert detect_active_tool(body) == "other"
-
-def test_detect_active_tool_role_assistant_with_tool_calls():
- body = {
- "messages": [
- {"role": "user", "content": "do something"},
- {"role": "assistant", "tool_calls": [{"id": "call_1", "function": {"name": "write_to_file"}}]}
- ]
- }
- assert detect_active_tool(body) == "write"
-
-def test_detect_active_tool_fallback_user_keyword():
- # Tests matching "tree"
- body = {
- "messages": [
- {"role": "user", "content": "show me the tree"}
- ]
- }
- assert detect_active_tool(body) == "tree"
-
- # Tests matching "shell"
- body = {
- "messages": [
- {"role": "user", "content": "run this in shell"}
- ]
- }
- assert detect_active_tool(body) == "shell"
-
- # Tests matching "write"
- body = {
- "messages": [
- {"role": "user", "content": "create file test.py"}
- ]
- }
- assert detect_active_tool(body) == "write"
-
- # Tests matching "view"
- body = {
- "messages": [
- {"role": "user", "content": "cat main.py"}
- ]
- }
- assert detect_active_tool(body) == "view"
-
-def test_detect_active_tool_ignores_invalid_message_formats():
- body = {
- "messages": [
- "this is not a dict",
- {"role": "user", "content": "read this"}
- ]
- }
- assert detect_active_tool(body) == "view"
-
-def test_detect_active_tool_precedence():
- # If there are multiple tools, it processes from the last message backwards
- body = {
- "messages": [
- {"role": "tool", "name": "edit_file", "content": "done"},
- {"role": "tool", "name": "cat_file", "content": "..."}
- ]
- }
- # It starts from the end, so it sees "cat_file" first, which maps to "view"
- assert detect_active_tool(body) == "view"
diff --git a/router/tests/test_estimate_prompt_tokens.py b/router/tests/test_estimate_prompt_tokens.py
index ae74a9e5..2a12f97c 100644
--- a/router/tests/test_estimate_prompt_tokens.py
+++ b/router/tests/test_estimate_prompt_tokens.py
@@ -1,4 +1,3 @@
-import pytest
import sys
import os
from pathlib import Path
@@ -23,27 +22,25 @@ def test_estimate_prompt_tokens_empty_messages():
def test_estimate_prompt_tokens_string_content():
body = {
"messages": [
- {"content": "word " * 4}, # 4 * 1.2 = 4.8
- {"content": "word " * 8} # 8 * 1.2 = 9.6
+ {"content": "1234"}, # 1 token
+ {"content": "12345678"} # 2 tokens
]
}
- # Total is int(round(4.8 + 9.6)) + 50 = int(round(14.4)) + 50 = 14 + 50 = 64
- assert estimate_prompt_tokens(body) == 50 + 14
+ assert estimate_prompt_tokens(body) == 50 + 1 + 2
def test_estimate_prompt_tokens_list_content():
body = {
"messages": [
{
"content": [
- {"type": "text", "text": "word " * 4}, # 4.8 tokens
+ {"type": "text", "text": "1234"}, # 1 token
{"type": "image_url", "url": "ignored"}, # 0 tokens
- {"type": "text", "text": "word " * 8} # 9.6 tokens
+ {"type": "text", "text": "12345678"} # 2 tokens
]
}
]
}
- # Total is int(round(4.8 + 9.6)) + 50 = 64
- assert estimate_prompt_tokens(body) == 50 + 14
+ assert estimate_prompt_tokens(body) == 50 + 1 + 2
def test_estimate_prompt_tokens_mixed_and_invalid_msgs():
body = {
@@ -54,11 +51,10 @@ def test_estimate_prompt_tokens_mixed_and_invalid_msgs():
"invalid_block_type", # Should be skipped
{"type": "text", "text": None} # None text, handled as empty string
]},
- {"content": "word " * 4} # 4.8 tokens
+ {"content": "1234"} # 1 token
]
}
- # Total is int(round(4.8)) + 50 = 5 + 50 = 55
- assert estimate_prompt_tokens(body) == 50 + 5
+ assert estimate_prompt_tokens(body) == 50 + 1
def test_estimate_prompt_tokens_missing_content():
body = {
diff --git a/router/tests/test_load_persisted_stats.py b/router/tests/test_load_persisted_stats.py
index 76edbf51..43cf52ff 100644
--- a/router/tests/test_load_persisted_stats.py
+++ b/router/tests/test_load_persisted_stats.py
@@ -1,61 +1,99 @@
-import json
import pytest
+import os
+import json
+import sys
from unittest.mock import patch, mock_open
-import router.main
-from router.main import load_persisted_stats
+# 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)
-@pytest.fixture
-def mock_stats():
- # Setup a clean stats dictionary for testing
- clean_stats = {
- "total_requests": 0,
- "nested_dict": {"a": 1, "b": 2},
- "existing_key": "value"
- }
- with patch.dict(router.main.stats, clean_stats, clear=True):
- yield router.main.stats
-
-def test_load_persisted_stats_file_not_exists(mock_stats):
- with patch("router.main.os.path.exists", return_value=False) as mock_exists:
- load_persisted_stats()
- mock_exists.assert_called_once_with(router.main.STATS_JSON_PATH)
- # Stats should remain unchanged
- assert mock_stats["total_requests"] == 0
-
-def test_load_persisted_stats_success(mock_stats):
- mock_data = {
+import main
+
+def test_load_persisted_stats_success():
+ mock_stats = {
"total_requests": 100,
- "nested_dict": {"b": 3, "c": 4},
+ "some_dict": {"a": 1, "b": 2},
"new_key": "new_value"
}
- mock_json = json.dumps(mock_data)
-
- with patch("router.main.os.path.exists", return_value=True):
- with patch("router.main.open", mock_open(read_data=mock_json)):
- with patch("router.main.logger.info") as mock_logger:
- load_persisted_stats()
-
- # Assert simple value updated via else block
- assert mock_stats["total_requests"] == 100
- # Assert nested_dict updated via if block (b updated, c added, a unchanged)
- assert mock_stats["nested_dict"] == {"a": 1, "b": 3, "c": 4}
- # Assert new_key added via else block
- assert mock_stats["new_key"] == "new_value"
- # Assert existing_key unchanged
- assert mock_stats["existing_key"] == "value"
-
- mock_logger.assert_called_once_with("✓ Successfully loaded persisted gateway statistics from disk.")
-
-def test_load_persisted_stats_exception(mock_stats):
- with patch("router.main.os.path.exists", return_value=True):
- with patch("router.main.open", side_effect=Exception("Mock read error")):
- with patch("router.main.logger.error") as mock_logger:
- load_persisted_stats()
-
- # Stats should remain unchanged
- assert mock_stats["total_requests"] == 0
-
- # Error should be logged
- mock_logger.assert_called_once()
- assert "Failed to load persisted stats: Mock read error" in mock_logger.call_args[0][0]
+
+ 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
new file mode 100644
index 00000000..92c5bc96
--- /dev/null
+++ b/router/tests/test_memory_mcp.py
@@ -0,0 +1,327 @@
+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 0139e8d3..11f4fd2a 100644
--- a/scripts/README.md
+++ b/scripts/README.md
@@ -42,44 +42,61 @@ 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 & Dataset Maintenance (`scripts/`)
+## 3. Classifier, Daemons & Maintenance (`scripts/`)
-These tools are used to benchmark the prompt classifier and extract datasets from Langfuse traces:
+These tools and helper daemons are used to benchmark the prompt classifier, extract datasets from Langfuse traces, and orchestrate client-host communication:
- **`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 (Root Directory)
+## 4. Integration Test Suite (`tests/`)
-The integration test suite is located in the root directory. Tests are categorized below based on their primary function:
+The integration test suite is located in the `tests/` directory. Tests are categorized below based on their primary function:
### Circuit Breaker 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.
-- **`verify_breaker.py`**: Sanity verification check for the circuit breaker.
-- **`test_a2_verify.py`**: Quick sanity integration check for the agy proxy circuit breaker.
+- **`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.
### Classifier Tests
-- **`test_classifier_accuracy.py`**: Accuracy evaluation suite covering 25 system prompts.
+- **`tests/test_classifier_accuracy.py`**: Accuracy evaluation suite covering 25 system prompts.
+- **`tests/test_map_tool_to_category.py`**: Tests prompt-classifier category mapping.
### Routing & Proxy Tests
-- **`test_agy_tiers.py`**: Validates `agy` proxy model tier routing.
-- **`test_antigravity.py`**: Tests the connection to the host Antigravity CLI daemon (`agentapi`).
+- **`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.
### Performance & Monitoring Tests
-- **`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.
+- **`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.
diff --git a/scripts/benchmark_tokens.py b/scripts/benchmark_tokens.py
deleted file mode 100644
index d701b4eb..00000000
--- a/scripts/benchmark_tokens.py
+++ /dev/null
@@ -1,76 +0,0 @@
-import sys
-import os
-from pathlib import Path
-
-# Set CONFIG_PATH and ROUTER_API_KEY for import
-os.environ["CONFIG_PATH"] = str(Path(__file__).resolve().parent.parent / "router" / "config.yaml")
-os.environ["ROUTER_API_KEY"] = "local-token"
-# Add the parent directory and the router directory to the path
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "router"))
-
-from router.main import estimate_prompt_tokens
-
-def verify_accuracy():
- """Benchmarking utility to verify token estimation accuracy across content types."""
- # Test cases inspired by the problem description
- test_cases = [
- {
- "name": "English prose",
- "content": "This is a standard English prose sentence intended to evaluate the accuracy of the token estimation heuristic for typical content. " * 5,
- "actual_tokens": 110,
- },
- {
- "name": "Python code",
- "content": """
-def calculate_factorial(n):
- if n == 0:
- return 1
- else:
- return n * calculate_factorial(n-1)
-
-for i in range(10):
- print(f"Factorial of {i} is {calculate_factorial(i)}")
-""" * 3,
- "actual_tokens": 150,
- },
- {
- "name": "CJK text",
- "content": "这是一个测试，用于验证中文字符的令牌估算逻辑。它应该比字符计数更准确。" * 5,
- "actual_tokens": 60,
- },
- {
- "name": "Whitespace-padded JSON",
- "content": '{\n "key": "value",\n "nested": {\n "inner": "data"\n }\n}\n' * 5,
- "actual_tokens": 60,
- },
- {
- "name": "Emoji",
- "content": "🚀🔥-🤖✨-📈💎-🚨🛠️-🌐" * 5,
- "actual_tokens": 25,
- }
- ]
-
- print(f"{'Case':<25} | {'Actual':<7} | {'Estimated':<9} | {'Error':<7}")
- print("-" * 55)
-
- all_passed = True
- for case in test_cases:
- body = {"messages": [{"content": case["content"]}]}
- est = estimate_prompt_tokens(body) - 50 # Subtract metadata overhead
- error = abs(est - case["actual_tokens"]) / case["actual_tokens"]
- print(f"{case['name']:<25} | {case['actual_tokens']:<7} | {est:<9} | {error:.1%}")
- # Acceptance criteria: within ±25% for these rough heuristics
- if error > 0.25:
- print(f" --> FAILURE: {case['name']} error exceeds target threshold")
- all_passed = False
-
- assert all_passed, "Token estimation accuracy benchmark failed"
-
-if __name__ == "__main__":
- try:
- verify_accuracy()
- sys.exit(0)
- except AssertionError as e:
- print(f"\nERROR: {e}")
- sys.exit(1)
diff --git a/get_pr_status.py b/scripts/get_pr_status.py
similarity index 52%
rename from get_pr_status.py
rename to scripts/get_pr_status.py
index 6214fbd5..0e042465 100644
--- a/get_pr_status.py
+++ b/scripts/get_pr_status.py
@@ -1,10 +1,11 @@
import subprocess
-import shlex
+from typing import Sequence
-def run_cmd(cmd):
+
+def run_cmd(argv: Sequence[str]) -> str:
# 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, check=True, timeout=30)
+ # 3. Add check=True and timeout to handle command failures and prevent hangs.
+ result = subprocess.run(argv, shell=False, capture_output=True, text=True, check=True, timeout=30)
return result.stdout.strip()
diff --git a/host_agy_daemon.py b/scripts/host_agy_daemon.py
similarity index 73%
rename from host_agy_daemon.py
rename to scripts/host_agy_daemon.py
index 09162856..25194a70 100755
--- a/host_agy_daemon.py
+++ b/scripts/host_agy_daemon.py
@@ -65,67 +65,91 @@ 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:
- os.close(slave_fd)
- os.close(master_fd)
+ try:
+ os.close(slave_fd)
+ except OSError:
+ pass
+ try:
+ os.close(master_fd)
+ except OSError:
+ pass
# Write failure details as status
- err_msg = json.dumps({"type": "status", "returncode": -1, "stderr": str(e)}) + "\n"
- self.wfile.write(err_msg.encode('utf-8'))
- self.wfile.flush()
+ 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
return
-
- loop_ref = asyncio.get_running_loop()
-
- def read_bytes():
+ finally:
+ # Always close the slave end in the parent process
try:
- return os.read(master_fd, 1024)
+ os.close(slave_fd)
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')
- # 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()
+ pass
+ 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
- os.close(master_fd)
-
- # Retrieve last conversation ID
- result_conv_id = get_last_conversation_id()
+ # Ensure master FD is closed
+ try:
+ os.close(master_fd)
+ except OSError:
+ pass
- # 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()
+ # 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
loop.run_until_complete(run_stream())
loop.close()
diff --git a/sync_gemini_token.py b/scripts/sync_gemini_token.py
similarity index 100%
rename from sync_gemini_token.py
rename to scripts/sync_gemini_token.py
diff --git a/test_quota_reset.sh b/scripts/test_quota_reset.sh
similarity index 100%
rename from test_quota_reset.sh
rename to scripts/test_quota_reset.sh
diff --git a/verify_breaker.py b/scripts/verification/verify_breaker.py
similarity index 76%
rename from verify_breaker.py
rename to scripts/verification/verify_breaker.py
index daf41bef..db4db8af 100644
--- a/verify_breaker.py
+++ b/scripts/verification/verify_breaker.py
@@ -3,7 +3,13 @@
import sys
from pathlib import Path
-sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+# 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 router.circuit_breaker import get_breaker
diff --git a/watch_quota.sh b/scripts/watch_quota.sh
similarity index 92%
rename from watch_quota.sh
rename to scripts/watch_quota.sh
index 8bf9d090..cb95debd 100755
--- a/watch_quota.sh
+++ b/scripts/watch_quota.sh
@@ -2,7 +2,8 @@
# Polling loop — checks quota every 30s and runs tests when reset
# Log file to watch
LOG_FILE="$HOME/.gemini/antigravity-cli/cli.log"
-TEST_SCRIPT="test_quota_reset.sh"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+TEST_SCRIPT="$SCRIPT_DIR/test_quota_reset.sh"
POLL_INTERVAL=30 # seconds
echo "=== Quota Reset Watcher ==="
diff --git a/start-stack.sh b/start-stack.sh
index 8abdb259..d59746de 100755
--- a/start-stack.sh
+++ b/start-stack.sh
@@ -31,6 +31,12 @@ 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 ! 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
@@ -68,7 +74,11 @@ if [ -f "$OAUTH_CREDS" ]; then
fi
fi
if $NEED_SYNC; then
- python3 sync_gemini_token.py || echo "⚠️ Warning: Failed to sync Gemini token from keyring"
+ 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"
fi
ACTIVE_OAUTH=""
@@ -95,7 +105,7 @@ else
echo "⚠️ Warning: Host agy daemon not responding on port 5005"
fi
-if [ -z "$POSTGRES_PASSWORD" ] || [ -z "$NEXTAUTH_SECRET" ] || [ -z "$SALT" ] || [ -z "$ENCRYPTION_KEY" ] || [ -z "$LITELLM_MASTER_KEY" ] || [ -z "$ROUTER_API_KEY" ] || [ -z "$MINIO_ROOT_USER" ] || [ -z "$MINIO_ROOT_PASSWORD" ]; then
+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 ! command -v openssl &>/dev/null; then
echo "❌ Error: 'openssl' is required to generate secure random keys but was not found in PATH."
exit 1
@@ -106,6 +116,17 @@ 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"
@@ -135,43 +156,53 @@ if [ -z "$LITELLM_MASTER_KEY" ]; then
exit 1
fi
-if [ -z "$LANGFUSE_INIT_USER_PASSWORD" ]; then
- LANGFUSE_INIT_USER_PASSWORD="$(openssl rand -hex 16)"
- echo "LANGFUSE_INIT_USER_PASSWORD=\"$LANGFUSE_INIT_USER_PASSWORD\"" >> "$ENV_FILE"
- echo "✓ Generated new LANGFUSE_INIT_USER_PASSWORD and saved to $ENV_FILE"
-fi
-
-if [ -z "$REDIS_AUTH" ]; then
- REDIS_AUTH="$(openssl rand -hex 16)"
- echo "REDIS_AUTH=\"$REDIS_AUTH\"" >> "$ENV_FILE"
- echo "✓ Generated new REDIS_AUTH and saved to $ENV_FILE"
-fi
-
-if [ -z "$CLICKHOUSE_PASSWORD" ]; then
- CLICKHOUSE_PASSWORD="$(openssl rand -hex 16)"
- echo "CLICKHOUSE_PASSWORD=\"$CLICKHOUSE_PASSWORD\"" >> "$ENV_FILE"
- echo "✓ Generated new CLICKHOUSE_PASSWORD and saved to $ENV_FILE"
-fi
-
-
if [ -z "$ROUTER_API_KEY" ]; then
ROUTER_API_KEY="$(openssl rand -hex 32)"
echo "ROUTER_API_KEY=\"$ROUTER_API_KEY\"" >> "$ENV_FILE"
echo "✓ Generated new ROUTER_API_KEY and saved to $ENV_FILE"
fi
-if [ -z "$MINIO_ROOT_USER" ]; then
- MINIO_ROOT_USER="minio-$(openssl rand -hex 4)"
- echo "MINIO_ROOT_USER=\"$MINIO_ROOT_USER\"" >> "$ENV_FILE"
- echo "✓ Generated new MINIO_ROOT_USER and saved to $ENV_FILE"
+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 "$MINIO_ROOT_PASSWORD" ]; then
- MINIO_ROOT_PASSWORD="$(openssl rand -hex 16)"
- echo "MINIO_ROOT_PASSWORD=\"$MINIO_ROOT_PASSWORD\"" >> "$ENV_FILE"
- echo "✓ Generated new MINIO_ROOT_PASSWORD and saved to $ENV_FILE"
+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
@@ -281,7 +312,7 @@ setup_minio_buckets() {
# Ensure mc alias points to the correct MinIO S3 API port (9002, not 9000)
# The default 'local' alias in the MinIO image points to :9000 which is ClickHouse,
# not MinIO. We must override it.
- podman exec agent-router-pod-minio-s3 mc alias set local http://127.0.0.1:9002 "$MINIO_ROOT_USER" "$MINIO_ROOT_PASSWORD" 2>/dev/null
+ podman exec agent-router-pod-minio-s3 mc alias set local http://127.0.0.1:9002 minioadmin minioadmin 2>/dev/null
# Create required buckets (idempotent)
local BUCKETS=("langfuse-events" "proj-triage-gateway-id")
@@ -382,7 +413,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 LANGFUSE_INIT_USER_PASSWORD MINIO_ROOT_USER MINIO_ROOT_PASSWORD REDIS_AUTH CLICKHOUSE_PASSWORD
+ export WORKDIR HOME LITELLM_MASTER_KEY POSTGRES_PASSWORD NEXTAUTH_SECRET SALT ENCRYPTION_KEY OLLAMA_API_KEY LANGFUSE_PUBLIC_KEY LANGFUSE_SECRET_KEY
python3 - "$WORKDIR/pod.yaml" <<'PY'
import os, sys, urllib.parse
uid = os.getuid()
@@ -392,38 +423,34 @@ placeholders = [
"/home/gpav/Vrac/LAB/AI/LLM-Routing",
"/home/gpav/",
"/run/user/1000",
- "sk-lit...33bf",
- "postgres:***",
+ "LITELLM_MASTER_KEY_PLACEHOLDER",
+ "POSTGRES_PASSWORD_RAW_PLACEHOLDER",
+ "POSTGRES_PASSWORD_ENCODED_PLACEHOLDER",
"NEXTAUTH_SECRET_PLACEHOLDER",
"SALT_PLACEHOLDER",
"ENCRYPTION_KEY_PLACEHOLDER",
- "postgres-password-***",
- "MINIO_USER_PLACEHOLDER",
- "MINIO_PASSWORD_PLACEHOLDER",
- "LANGFUSE_INIT_USER_PASSWORD_PLACEHOLDER",
- "REDIS_AUTH_PLACEHOLDER",
- "CLICKHOUSE_PASSWORD_PLACEHOLDER"
+ "OLLAMA_API_KEY_PLACEHOLDER",
+ "LANGFUSE_PUBLIC_KEY_PLACEHOLDER",
+ "LANGFUSE_SECRET_KEY_PLACEHOLDER"
]
for ph in placeholders:
if ph not in text:
- sys.stderr.write(f"Error: Required placeholder '{ph}' not found in pod.yaml. Ensure you are using the latest version of the template.\n")
+ sys.stderr.write(f"Error: Required placeholder '{ph}' not found in pod.yaml\n")
sys.exit(1)
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("LITELLM_MASTER_KEY_PLACEHOLDER", os.environ["LITELLM_MASTER_KEY"])
+text = text.replace("POSTGRES_PASSWORD_RAW_PLACEHOLDER", 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"])
+encoded_password = urllib.parse.quote(os.environ['POSTGRES_PASSWORD'], safe='')
+text = text.replace("POSTGRES_PASSWORD_ENCODED_PLACEHOLDER", encoded_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("MINIO_USER_PLACEHOLDER", os.environ["MINIO_ROOT_USER"])
-text = text.replace("MINIO_PASSWORD_PLACEHOLDER", os.environ["MINIO_ROOT_PASSWORD"])
-text = text.replace("LANGFUSE_INIT_USER_PASSWORD_PLACEHOLDER", os.environ["LANGFUSE_INIT_USER_PASSWORD"])
-text = text.replace("REDIS_AUTH_PLACEHOLDER", os.environ["REDIS_AUTH"])
-text = text.replace("CLICKHOUSE_PASSWORD_PLACEHOLDER", os.environ["CLICKHOUSE_PASSWORD"])
+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/test_host_agy_daemon.py b/test_host_agy_daemon.py
deleted file mode 100644
index e0cae61c..00000000
--- a/test_host_agy_daemon.py
+++ /dev/null
@@ -1,44 +0,0 @@
-import os
-import json
-import pytest
-from unittest.mock import patch, mock_open
-
-import host_agy_daemon
-
-def test_get_last_conversation_id_success(monkeypatch):
- test_data = {"/test/path": "conv_123"}
- monkeypatch.setattr(os, "getcwd", lambda: "/test/path")
- monkeypatch.setattr(os.path, "exists", lambda x: True)
-
- m_open = mock_open(read_data=json.dumps(test_data))
- with patch("builtins.open", m_open):
- assert host_agy_daemon.get_last_conversation_id() == "conv_123"
-
-def test_get_last_conversation_id_not_found(monkeypatch):
- test_data = {"/other/path": "conv_123"}
- monkeypatch.setattr(os, "getcwd", lambda: "/test/path")
- monkeypatch.setattr(os.path, "exists", lambda x: True)
-
- m_open = mock_open(read_data=json.dumps(test_data))
- with patch("builtins.open", m_open):
- assert host_agy_daemon.get_last_conversation_id() is None
-
-def test_get_last_conversation_id_file_missing(monkeypatch):
- monkeypatch.setattr(os.path, "exists", lambda x: False)
- assert host_agy_daemon.get_last_conversation_id() is None
-
-def test_get_last_conversation_id_exception(monkeypatch):
- monkeypatch.setattr(os.path, "exists", lambda x: True)
- # Invalid JSON to trigger JSONDecodeError
- m_open = mock_open(read_data="{invalid json")
- with patch("builtins.open", m_open):
- assert host_agy_daemon.get_last_conversation_id() is None
-
-def test_get_last_conversation_id_io_error(monkeypatch):
- monkeypatch.setattr(os.path, "exists", lambda x: True)
-
- def raise_error(*args, **kwargs):
- raise IOError("permission denied")
-
- with patch("builtins.open", side_effect=raise_error):
- assert host_agy_daemon.get_last_conversation_id() is None
diff --git a/test_memory_mcp.py b/test_memory_mcp.py
deleted file mode 100644
index 0194dca9..00000000
--- a/test_memory_mcp.py
+++ /dev/null
@@ -1,167 +0,0 @@
-#!/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/test_a2_verify.py b/tests/test_a2_verify.py
similarity index 72%
rename from test_a2_verify.py
rename to tests/test_a2_verify.py
index 42cde1e1..ca9ca025 100644
--- a/test_a2_verify.py
+++ b/tests/test_a2_verify.py
@@ -2,7 +2,13 @@
"""Verify circuit breaker integration into agy_proxy.py"""
import sys
from pathlib import Path
-sys.path.insert(0, str(Path(__file__).resolve().parent / 'router'))
+
+# 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 circuit_breaker import get_breaker
from agy_proxy import try_agy_proxy
diff --git a/test_agy_behavior.py b/tests/test_agy_behavior.py
similarity index 100%
rename from test_agy_behavior.py
rename to tests/test_agy_behavior.py
diff --git a/test_agy_tiers.py b/tests/test_agy_tiers.py
similarity index 100%
rename from test_agy_tiers.py
rename to tests/test_agy_tiers.py
diff --git a/test_antigravity.py b/tests/test_antigravity.py
similarity index 100%
rename from test_antigravity.py
rename to tests/test_antigravity.py
diff --git a/test_atomic_write.py b/tests/test_atomic_write.py
similarity index 100%
rename from test_atomic_write.py
rename to tests/test_atomic_write.py
diff --git a/test_check_http_endpoint.py b/tests/test_check_http_endpoint.py
similarity index 100%
rename from test_check_http_endpoint.py
rename to tests/test_check_http_endpoint.py
diff --git a/test_circuit_breaker.py b/tests/test_circuit_breaker.py
similarity index 98%
rename from test_circuit_breaker.py
rename to tests/test_circuit_breaker.py
index 8cdcad60..2dfd8064 100644
--- a/test_circuit_breaker.py
+++ b/tests/test_circuit_breaker.py
@@ -19,7 +19,12 @@
import pytest
from unittest.mock import patch, AsyncMock
from pathlib import Path
-sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+# 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))
from router.circuit_breaker import get_breaker, TIER_COOLDOWNS, MAX_TIER
diff --git a/test_classifier_accuracy.py b/tests/test_classifier_accuracy.py
similarity index 100%
rename from test_classifier_accuracy.py
rename to tests/test_classifier_accuracy.py
diff --git a/test_compute_free_model_score.py b/tests/test_compute_free_model_score.py
similarity index 100%
rename from test_compute_free_model_score.py
rename to tests/test_compute_free_model_score.py
diff --git a/tests/test_host_agy_daemon.py b/tests/test_host_agy_daemon.py
new file mode 100644
index 00000000..61f84e86
--- /dev/null
+++ b/tests/test_host_agy_daemon.py
@@ -0,0 +1,490 @@
+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/test_map_tool_to_category.py b/tests/test_map_tool_to_category.py
similarity index 100%
rename from test_map_tool_to_category.py
rename to tests/test_map_tool_to_category.py
diff --git a/test_models_proxy.py b/tests/test_models_proxy.py
similarity index 93%
rename from test_models_proxy.py
rename to tests/test_models_proxy.py
index bfaf3b81..e0df6e41 100644
--- a/test_models_proxy.py
+++ b/tests/test_models_proxy.py
@@ -9,7 +9,13 @@
import sys
from pathlib import Path
-sys.path.insert(0, str(Path(__file__).resolve().parent / "router"))
+
+# 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 main import get_http_client, proxy_models, HTTP_MAX_CONNECTIONS, HTTP_MAX_KEEPALIVE_CONNECTIONS, HTTP_KEEPALIVE_EXPIRY
diff --git a/test_pie_chart_gradient.py b/tests/test_pie_chart_gradient.py
similarity index 68%
rename from test_pie_chart_gradient.py
rename to tests/test_pie_chart_gradient.py
index c9e85499..1d9a33bd 100644
--- a/test_pie_chart_gradient.py
+++ b/tests/test_pie_chart_gradient.py
@@ -4,26 +4,22 @@
@pytest.fixture
def mock_stats():
- # Patch router.main.stats with a real dictionary containing tool_tokens
- # to ensure the function under test reads from the correct key.
- test_stats = {
- "tool_tokens": {
- "tree": 0,
- "shell": 0,
- "write": 0,
- "view": 0,
- "other": 0
- }
- }
- with patch("router.main.stats", test_stats):
- yield test_stats
+ with patch("router.main.stats") as mock_stats_obj:
+ yield mock_stats_obj
def test_get_pie_chart_gradient_empty(mock_stats):
+ mock_stats.__getitem__.return_value = {
+ "tree": 0,
+ "shell": 0,
+ "write": 0,
+ "view": 0,
+ "other": 0
+ }
result = get_pie_chart_gradient()
assert result == "background: rgba(255, 255, 255, 0.05);"
def test_get_pie_chart_gradient_one_tool(mock_stats):
- mock_stats["tool_tokens"] = {
+ mock_stats.__getitem__.return_value = {
"tree": 100,
"shell": 0,
"write": 0,
@@ -34,7 +30,7 @@ def test_get_pie_chart_gradient_one_tool(mock_stats):
assert result == "background: conic-gradient(#34d399 0.0% 100.0%);"
def test_get_pie_chart_gradient_multiple_tools(mock_stats):
- mock_stats["tool_tokens"] = {
+ mock_stats.__getitem__.return_value = {
"tree": 50,
"shell": 25,
"write": 25,
@@ -45,7 +41,7 @@ def test_get_pie_chart_gradient_multiple_tools(mock_stats):
assert result == "background: conic-gradient(#34d399 0.0% 50.0%, #fbbf24 50.0% 75.0%, #a78bfa 75.0% 100.0%);"
def test_get_pie_chart_gradient_unrecognized_tool(mock_stats):
- mock_stats["tool_tokens"] = {
+ mock_stats.__getitem__.return_value = {
"unknown_tool": 100
}
result = get_pie_chart_gradient()
diff --git a/test_record_tool_usage.py b/tests/test_record_tool_usage.py
similarity index 98%
rename from test_record_tool_usage.py
rename to tests/test_record_tool_usage.py
index 22105c44..5d4d4eb2 100644
--- a/test_record_tool_usage.py
+++ b/tests/test_record_tool_usage.py
@@ -1,6 +1,6 @@
import pytest
import copy
-from unittest.mock import patch, MagicMock
+from unittest.mock import patch
import router.main
diff --git a/test_src_badge.py b/tests/test_src_badge.py
similarity index 100%
rename from test_src_badge.py
rename to tests/test_src_badge.py
diff --git a/test_stream_latency.py b/tests/test_stream_latency.py
similarity index 100%
rename from test_stream_latency.py
rename to tests/test_stream_latency.py
diff --git a/test_sync_gemini_token.py b/tests/test_sync_gemini_token.py
similarity index 96%
rename from test_sync_gemini_token.py
rename to tests/test_sync_gemini_token.py
index d0cbe351..9a3d8157 100644
--- a/test_sync_gemini_token.py
+++ b/tests/test_sync_gemini_token.py
@@ -6,6 +6,14 @@
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/triage_upgrade_plan.md b/triage_upgrade_plan.md
index dc8d283d..ada4d2ee 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.
-- `test_agy_tiers.py`: Tests for the triage tiers, may need updates.
-- `test_classifier_accuracy.py`: Tests for classifier accuracy, may need updates.
+- `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.
## Estimated Effort
- Review and planning: 2 hours
From 1249025ceb58a651a35666f4cd0e4e5af74b4cc1 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:15:32 +0000
Subject: [PATCH 4/4] Merge master changes & Resolve all conflicts before
continuing
Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com>
---
.github/dependabot.yml | 4 +-
.github/workflows/test.yml | 6 +-
.gitignore | 3 -
README.md | 84 +-
scripts/get_pr_status.py => get_pr_status.py | 9 +-
.../host_agy_daemon.py => host_agy_daemon.py | 104 +-
litellm/entrypoint.py | 20 -
pod.yaml | 54 +-
router/main.py | 1319 ++++++-----------
router/test_memory_mcp.py | 131 ++
router/tests/test_agy_proxy.py | 10 -
router/tests/test_dashboard_data.py | 2 +-
router/tests/test_detect_active_tool.py | 114 ++
router/tests/test_estimate_prompt_tokens.py | 20 +-
router/tests/test_get_goose_sessions.py | 46 +
router/tests/test_load_persisted_stats.py | 146 +-
router/tests/test_memory_mcp.py | 327 ----
scripts/README.md | 49 +-
scripts/benchmark_tokens.py | 76 +
start-stack.sh | 109 +-
...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
test_host_agy_daemon.py | 44 +
...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 | 28 +-
...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 +-
44 files changed, 1252 insertions(+), 2158 deletions(-)
rename scripts/get_pr_status.py => get_pr_status.py (52%)
rename scripts/host_agy_daemon.py => host_agy_daemon.py (73%)
create mode 100644 router/test_memory_mcp.py
create mode 100644 router/tests/test_detect_active_tool.py
create mode 100644 router/tests/test_get_goose_sessions.py
delete mode 100644 router/tests/test_memory_mcp.py
create mode 100644 scripts/benchmark_tokens.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%)
create mode 100644 test_host_agy_daemon.py
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 (68%)
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 4250afa9..0b8e5878 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
- 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/.gitignore b/.gitignore
index 5ccf9013..77450f74 100644
--- a/.gitignore
+++ b/.gitignore
@@ -32,11 +32,8 @@ router/free_models_roster.json
# agent artifacts
.hermes/
.jules/
-.local/
workdirs/
pr_description.txt
# Dataset work in progress
data/
-.cache/
-test_output*.log
diff --git a/README.md b/README.md
index 00673fe8..38ca0461 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,15 +76,15 @@ 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 |
-| **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 |
+| **valkey-cache** (9.1.0-alpine) | `valkey-cli PING` every 10s | `valkey-cli PING` every 5s |
+| **litellm-gateway** | Python `urllib` GET `/health` (port 4000, accepts 200/401) every 15s | Same, every 10s |
+| **llm-triage-router** | Python `urllib` GET `/dashboard` (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 |
-| **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 |
+| **clickhouse-db** | `clickhouse-client --user clickhouse --password clickhouse --query "SELECT 1"` every 15s | Same, every 10s |
+| **valkey-lf** | `redis-cli -p 6380 -a langfuse-redis-2026 PING` every 10s | Same, every 5s |
+| **langfuse-web** | `wget -qO /dev/null http://127.0.0.1:3001/` every 15s | Same, every 10s |
+| **langfuse-worker** | `pgrep -f langfuse-worker` every 15s | — |
+| **minio-s3** | TCP socket check on port 9002 every 15s | Same, every 10s |
The pod-level `restartPolicy: Always` combined with these probes means Podman will restart any container that fails its health check or exits unexpectedly, enabling true self-healing for the entire stack.
@@ -210,8 +212,8 @@ 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/
-├── .env # Environment file for API keys, passwords, and generated secrets (ignored by git)
+/home/gpav/Vrac/LAB/AI/LLM-Routing/
+├── .env # Environment file for OpenRouter API Key (ignored by git)
├── .gitignore # Git ignore policy protecting secrets & database files
├── README.md # In-depth system and operational guide
├── pod.yaml # Podman Kubernetes template defining the 10-container stack
@@ -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
@@ -388,7 +379,7 @@ Run the startup script from the root of the repository:
# health probes, env vars, containers — no rebuild)
./start-stack.sh --full-rebuild # Full reset: rebuild image + recreate pod
```
-*Note: If running for the first time, the script will prompt you for your `OpenRouter API Key`, securely saving it inside `.env` with restrictive permissions (`chmod 600`). The script also automatically generates and persists secure random secrets (`LITELLM_MASTER_KEY`, `POSTGRES_PASSWORD`, `NEXTAUTH_SECRET`, `SALT`, `ENCRYPTION_KEY`, and `ROUTER_API_KEY`) to this file on startup if they are missing.*
+*Note: If running for the first time, the script will prompt you for your `OpenRouter API Key`, securely saving it inside `.env` with restrictive permissions (`chmod 600`).*
### 2. Verify Container Status
Check that all **10 containers** inside `agent-router-pod` are up and running:
@@ -584,25 +575,19 @@ Without Minio, Langfuse v3 **will not start** — it validates S3 connectivity a
|----------|-------|
| `LANGFUSE_S3_EVENT_UPLOAD_BUCKET` | `langfuse-events` |
| `LANGFUSE_S3_EVENT_UPLOAD_ENDPOINT` | `http://127.0.0.1:9002` |
-| `LANGFUSE_S3_EVENT_UPLOAD_ACCESS_KEY_ID` | `minioadmin` |
-| `LANGFUSE_S3_EVENT_UPLOAD_SECRET_ACCESS_KEY` | `minioadmin` |
+| `LANGFUSE_S3_EVENT_UPLOAD_ACCESS_KEY_ID` | (auto-generated in `.env`) |
+| `LANGFUSE_S3_EVENT_UPLOAD_SECRET_ACCESS_KEY` | (auto-generated in `.env`) |
| `S3_FORCE_PATH_STYLE` | `true` |
-Minio runs on ports **9001** (web console) and **9002** (S3 API). Credentials: `minioadmin` / `minioadmin`. Image pinned to `docker.io/minio/minio:RELEASE.2025-10-15T17-29-55Z`.
+Minio runs on ports **9001** (web console) and **9002** (S3 API). Credentials are automatically generated and stored in `.env`. Image pinned to `docker.io/minio/minio:RELEASE.2025-10-15T17-29-55Z`.
### Health Check
-MinIO's health is monitored using its native structured endpoints `/minio/health/live` (liveness) and `/minio/health/ready` (readiness) on port 9002:
+Minio's minimal Go image has no HTTP client tools. The probe uses a raw TCP socket check:
```yaml
-livenessProbe:
- httpGet:
- path: /minio/health/live
- port: 9002
-readinessProbe:
- httpGet:
- path: /minio/health/ready
- port: 9002
+exec:
+ command: [sh, -c, "exec 3<>/dev/tcp/127.0.0.1/9002 && echo ok"]
```
---
@@ -704,6 +689,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 +713,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 +721,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 +783,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 |
@@ -797,7 +791,7 @@ For auto-routing modes, the Triage Router handles failures by silently falling b
This project is supported by a dedicated NotebookLM companion notebook:
* **Notebook Name:** `TriageGate-Architect-KB`
-* **Notebook ID:** llm-triage-gateway
+* **Notebook ID:** `llm-triage-gateway`
* **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/scripts/get_pr_status.py b/get_pr_status.py
similarity index 52%
rename from scripts/get_pr_status.py
rename to get_pr_status.py
index 0e042465..6214fbd5 100644
--- a/scripts/get_pr_status.py
+++ b/get_pr_status.py
@@ -1,11 +1,10 @@
import subprocess
-from typing import Sequence
+import shlex
-
-def run_cmd(argv: Sequence[str]) -> str:
+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
- # 3. Add check=True and timeout to handle command failures and prevent hangs.
- result = subprocess.run(argv, shell=False, capture_output=True, text=True, check=True, timeout=30)
+ args = shlex.split(cmd)
+ result = subprocess.run(args, shell=False, capture_output=True, text=True, check=True, timeout=30)
return result.stdout.strip()
diff --git a/scripts/host_agy_daemon.py b/host_agy_daemon.py
similarity index 73%
rename from scripts/host_agy_daemon.py
rename to host_agy_daemon.py
index 25194a70..9233ea65 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
- try:
- os.close(slave_fd)
- except OSError:
- pass
- returncode = -1
- try:
- loop_ref = asyncio.get_running_loop()
+ loop_ref = asyncio.get_running_loop()
- def read_bytes():
- try:
- return os.read(master_fd, 1024)
- except OSError:
- return b""
+ 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()
+ 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()
- # Wait for subprocess
+ try:
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
+ os.close(master_fd)
- # 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
+ # 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..188f78cd 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
@@ -273,9 +273,9 @@ spec:
- name: LANGFUSE_S3_EVENT_UPLOAD_REGION
value: us-east-1
- name: LANGFUSE_S3_EVENT_UPLOAD_ACCESS_KEY_ID
- value: minioadmin
+ value: MINIO_USER_PLACEHOLDER
- name: LANGFUSE_S3_EVENT_UPLOAD_SECRET_ACCESS_KEY
- value: minioadmin
+ value: MINIO_PASSWORD_PLACEHOLDER
- name: LANGFUSE_S3_EVENT_UPLOAD_ENDPOINT
value: http://127.0.0.1:9002
- name: S3_FORCE_PATH_STYLE
@@ -307,10 +307,10 @@ spec:
- name: LANGFUSE_INIT_USER_EMAIL
value: admin@local.dev
- name: LANGFUSE_INIT_USER_PASSWORD
- value: admin-local-pw-2026
+ value: LANGFUSE_INIT_USER_PASSWORD_PLACEHOLDER
- 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
@@ -363,16 +363,16 @@ spec:
- name: LANGFUSE_S3_EVENT_UPLOAD_REGION
value: us-east-1
- name: LANGFUSE_S3_EVENT_UPLOAD_ACCESS_KEY_ID
- value: minioadmin
+ value: MINIO_USER_PLACEHOLDER
- name: S3_FORCE_PATH_STYLE
value: 'true'
- name: LANGFUSE_S3_EVENT_UPLOAD_ENDPOINT
value: http://127.0.0.1:9002
- name: LANGFUSE_S3_EVENT_UPLOAD_SECRET_ACCESS_KEY
- value: minioadmin
+ value: MINIO_PASSWORD_PLACEHOLDER
- name: LANGFUSE_LOG_LEVEL
value: warn
- image: langfuse/langfuse-worker:3.202.1
+ image: docker.io/langfuse/langfuse-worker:3
livenessProbe:
exec:
command:
@@ -393,10 +393,10 @@ spec:
- ":9001"
env:
- name: MINIO_ROOT_USER
- value: minioadmin
+ value: MINIO_USER_PLACEHOLDER
- name: MINIO_ROOT_PASSWORD
- value: minioadmin
- image: minio/minio:RELEASE.2025-09-07T16-13-09Z
+ value: MINIO_PASSWORD_PLACEHOLDER
+ image: docker.io/minio/minio:latest
livenessProbe:
httpGet:
path: /minio/health/live
diff --git a/router/main.py b/router/main.py
index 408772ba..c5ec94f9 100644
--- a/router/main.py
+++ b/router/main.py
@@ -1,7 +1,9 @@
import os
+import re
import sys
import json
import time
+import socket
import asyncio
import logging
import copy
@@ -10,27 +12,15 @@
import httpx
import redis.asyncio as aioredis
from contextlib import asynccontextmanager
-
from fastapi import FastAPI, Request, HTTPException, Response
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pathlib import Path
from circuit_breaker import get_breaker
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator, RootModel
+from pydantic import BaseModel
from typing import Dict, Optional, Union
-
LITELLM_URL = (os.getenv("LITELLM_ADMIN_URL") or "http://127.0.0.1:4000").rstrip("/")
-LLAMA_SERVER_URL = (os.getenv("LLAMA_SERVER_URL") or "http://127.0.0.1:8080").rstrip(
- "/"
-)
-
-
-_redis_client = None
-_redis_last_init_attempt = 0.0
-_REDIS_RETRY_INTERVAL_SECONDS = 5.0
-
-
-
+LLAMA_SERVER_URL = (os.getenv("LLAMA_SERVER_URL") or "http://127.0.0.1:8080").rstrip("/")
_redis_client = None
_redis_last_init_attempt = 0.0
@@ -63,14 +53,11 @@ def get_redis():
# Connection pool limits configuration for the shared HTTP client
HTTP_MAX_CONNECTIONS = int(os.getenv("HTTP_MAX_CONNECTIONS") or "1000")
-HTTP_MAX_KEEPALIVE_CONNECTIONS = int(
- os.getenv("HTTP_MAX_KEEPALIVE_CONNECTIONS") or "500"
-)
+HTTP_MAX_KEEPALIVE_CONNECTIONS = int(os.getenv("HTTP_MAX_KEEPALIVE_CONNECTIONS") or "500")
HTTP_KEEPALIVE_EXPIRY = float(os.getenv("HTTP_KEEPALIVE_EXPIRY") or "5.0")
_http_client = None
-
def get_http_client():
"""Return the shared global httpx.AsyncClient singleton with configured limits."""
global _http_client
@@ -84,24 +71,60 @@ def get_http_client():
return _http_client
+# Compiled regular expressions for token estimation heuristics
+WORD_RE = re.compile(r'[a-zA-Z0-9]+')
+NON_ASCII_RE = re.compile(r'[^\s\x00-\x7F]')
+PUNC_RE = re.compile(r'[\x21-\x2f\x3a-\x40\x5b-\x60\x7b-\x7e]')
+
+
+def _count_tokens_heuristic(text: str) -> float:
+ """Heuristically estimate token count using weighted categories and optimized regex splitting.
+
+ This replaces the naive character-count logic with a more granular approach that
+ balances English words, technical identifiers, punctuation, and multi-byte characters.
+
+ Returns a float to prevent intermediate rounding errors when summing across multiple
+ message blocks. Callers should round the total sum to convert it to an integer.
+ """
+ if not text:
+ return 0.0
+
+ # 1. Alphanumeric runs (Words/Identifiers/Hashes/Base64)
+ # Use a length-aware heuristic to avoid under-counting technical content.
+ word_matches = WORD_RE.findall(text)
+ word_total = sum(1.2 if len(w) <= 8 else len(w) / 4.0 for w in word_matches)
+
+ # 2. Non-ASCII characters (CJK/Emoji)
+ # Each character is weighted at 0.35 tokens.
+ non_ascii_count = len(NON_ASCII_RE.findall(text))
+
+ # 3. ASCII Punctuation/Symbols
+ # Characters that are ASCII but not alphanumeric or whitespace.
+ punc_count = len(PUNC_RE.findall(text))
+
+ return word_total + (non_ascii_count * 0.35) + (punc_count * 0.4)
+
+
def estimate_prompt_tokens(body: dict) -> int:
- """Estimate prompt tokens by counting characters in message contents (1 token ~= 4 chars)
- to avoid inflating metrics with large tool/schema declarations.
+ """Estimate prompt tokens using a regex-based weighted heuristic for mixed content.
"""
- tokens = 0
+ total = 0.0
for msg in body.get("messages", []):
if not isinstance(msg, dict):
continue
content = msg.get("content") or ""
if isinstance(content, str):
- tokens += len(content) // 4
+ total += _count_tokens_heuristic(content)
elif isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "text":
- tokens += len(block.get("text") or "") // 4
- # Include a flat estimate for system prompt / metadata overhead
- tokens += 50
- return max(1, tokens)
+ text = block.get("text")
+ if isinstance(text, str):
+ total += _count_tokens_heuristic(text)
+
+ # Include a flat estimate for system prompt / metadata overhead.
+ # Use rounding to avoid truncation bias (e.g., 1.9 -> 1).
+ return max(1, int(round(total)) + 50)
async def sync_cooldowns_from_valkey() -> None:
@@ -159,7 +182,6 @@ async def save_cooldowns_to_valkey() -> None:
class ValkeyCooldownPersistence:
"""Persistence provider mapping Valkey/Redis client synchronization to the global handlers."""
-
async def sync(self) -> None:
await sync_cooldowns_from_valkey()
@@ -167,6 +189,7 @@ async def save(self) -> None:
await save_cooldowns_to_valkey()
+
# Configure logging — respect LOG_LEVEL env var (default: WARNING)
_log_level_str = os.getenv("LOG_LEVEL", "WARNING").upper()
_log_level = getattr(logging, _log_level_str, logging.WARNING)
@@ -177,7 +200,6 @@ async def save(self) -> None:
# Langfuse observability — per-request traces + aggregate score pushes
_langfuse_client = None
-
def get_langfuse():
"""Return the Langfuse client singleton, lazily initialized.
Returns None if Langfuse is unreachable (non-fatal)."""
@@ -185,7 +207,6 @@ def get_langfuse():
if _langfuse_client is None:
try:
import langfuse
-
_langfuse_client = langfuse.Langfuse(
public_key=os.getenv("LANGFUSE_PUBLIC_KEY", ""),
secret_key=os.getenv("LANGFUSE_SECRET_KEY", ""),
@@ -194,13 +215,10 @@ def get_langfuse():
)
logger.info("Langfuse client initialized")
except (ImportError, ValueError, TypeError) as e:
- logger.warning(
- f"Langfuse client initialization failed: {e} — traces disabled"
- )
+ logger.warning(f"Langfuse client initialization failed: {e} — traces disabled")
_langfuse_client = False # sentinel to avoid retry
return _langfuse_client if _langfuse_client is not False else None
-
async def push_aggregate_scores():
"""Push aggregate KPIs as Langfuse scores every 5 minutes."""
while True:
@@ -214,53 +232,18 @@ async def push_aggregate_scores():
continue
router = get_breaker()
scores = [
- {
- "name": "simple_ratio_pct",
- "value": stats.get("simple_requests", 0) / total * 100,
- },
- {
- "name": "medium_ratio_pct",
- "value": stats.get("medium_requests", 0) / total * 100,
- },
- {
- "name": "complex_ratio_pct",
- "value": stats.get("complex_requests", 0) / total * 100,
- },
- {
- "name": "reasoning_ratio_pct",
- "value": stats.get("reasoning_requests", 0) / total * 100,
- },
- {
- "name": "advanced_ratio_pct",
- "value": stats.get("advanced_requests", 0) / total * 100,
- },
- {
- "name": "cache_hit_rate_pct",
- "value": stats["cache_hits"] / total * 100,
- },
- {
- "name": "avg_triage_latency_ms",
- "value": stats["avg_triage_latency_ms"],
- },
- {
- "name": "avg_proxy_latency_ms",
- "value": stats["avg_proxy_latency_ms"],
- },
+ {"name": "simple_ratio_pct", "value": stats.get("simple_requests", 0) / total * 100},
+ {"name": "medium_ratio_pct", "value": stats.get("medium_requests", 0) / total * 100},
+ {"name": "complex_ratio_pct", "value": stats.get("complex_requests", 0) / total * 100},
+ {"name": "reasoning_ratio_pct", "value": stats.get("reasoning_requests", 0) / total * 100},
+ {"name": "advanced_ratio_pct", "value": stats.get("advanced_requests", 0) / total * 100},
+ {"name": "cache_hit_rate_pct", "value": stats["cache_hits"] / total * 100},
+ {"name": "avg_triage_latency_ms", "value": stats["avg_triage_latency_ms"]},
+ {"name": "avg_proxy_latency_ms", "value": stats["avg_proxy_latency_ms"]},
{"name": "total_requests", "value": float(total)},
- {
- "name": "circuit_breaker_google_tier",
- "value": float(router.google.tier),
- },
- {
- "name": "circuit_breaker_vendor_tier",
- "value": float(router.vendor.tier),
- },
- {
- "name": "google_oauth_direct_ratio_pct",
- "value": stats["routing_paths"]["google_oauth_direct"]
- / total
- * 100,
- },
+ {"name": "circuit_breaker_google_tier", "value": float(router.google.tier)},
+ {"name": "circuit_breaker_vendor_tier", "value": float(router.vendor.tier)},
+ {"name": "google_oauth_direct_ratio_pct", "value": stats["routing_paths"]["google_oauth_direct"] / total * 100},
]
trace_id = lf.create_trace_id(seed=f"aggregate_scores_{int(time.time())}")
lf.start_observation(
@@ -269,15 +252,16 @@ async def push_aggregate_scores():
level="DEFAULT",
)
for s in scores:
- lf.create_score(name=s["name"], value=s["value"], trace_id=trace_id)
+ lf.create_score(
+ name=s["name"],
+ value=s["value"],
+ trace_id=trace_id
+ )
lf.flush()
- logger.info(
- f"Pushed {len(scores)} aggregate scores to Langfuse (trace_id={trace_id})"
- )
+ logger.info(f"Pushed {len(scores)} aggregate scores to Langfuse (trace_id={trace_id})")
except Exception as e:
logger.warning(f"Langfuse score push failed (non-fatal): {e}")
-
# Load configuration
CONFIG_PATH = os.getenv("CONFIG_PATH", "/config/config.yaml")
try:
@@ -292,17 +276,10 @@ 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")
-if not router_api_key:
- raise RuntimeError("Configuration error: 'api_key' is missing from router_model configuration.")
+router_api_key = router_model_conf.get("api_key", "local-token")
if router_api_key.startswith("os.environ/"):
env_var = router_api_key.split("/", 1)[1]
- 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_api_key = os.environ.get(env_var, "local-token")
router_model_name = router_model_conf.get("model", "qwen-0.8b-routing")
system_prompt = config.get("classification_rules", {}).get("system_prompt", "")
@@ -333,9 +310,18 @@ async def push_aggregate_scores():
"total_proxy_time_ms": 0.0,
"prompt_tokens": 0,
"completion_tokens": 0,
- "tool_tokens": {"tree": 0, "shell": 0, "write": 0, "view": 0, "other": 0},
- "routing_paths": {"google_oauth_direct": 0, "litellm_fallback": 0},
- "timeline": [],
+ "tool_tokens": {
+ "tree": 0,
+ "shell": 0,
+ "write": 0,
+ "view": 0,
+ "other": 0
+ },
+ "routing_paths": {
+ "google_oauth_direct": 0,
+ "litellm_fallback": 0
+ },
+ "timeline": []
}
# ---------------------------------------------------------------------------
@@ -346,11 +332,9 @@ async def push_aggregate_scores():
# triage router tracks Ollama failures itself and returns 429 immediately
# during the cooldown window, skipping the LiteLLM call entirely.
# ---------------------------------------------------------------------------
-_ollama_cooldown_until: float = 0.0 # monotonic timestamp when cooldown expires
+_ollama_cooldown_until: float = 0.0 # monotonic timestamp when cooldown expires
try:
- OLLAMA_COOLDOWN_SECONDS: int = int(
- os.getenv("OLLAMA_COOLDOWN_SECONDS", "300")
- ) # 5 min default
+ OLLAMA_COOLDOWN_SECONDS: int = int(os.getenv("OLLAMA_COOLDOWN_SECONDS", "300")) # 5 min default
if OLLAMA_COOLDOWN_SECONDS <= 0:
raise ValueError("OLLAMA_COOLDOWN_SECONDS must be positive")
except (TypeError, ValueError) as e:
@@ -363,7 +347,6 @@ async def push_aggregate_scores():
# preventing premature garbage collection before the task completes (Ruff RUF006).
_background_tasks: set = set()
-
def load_persisted_stats():
"""Loads persisted statistics from disk on startup to prevent resets on pod redeployment."""
global stats
@@ -378,20 +361,9 @@ def load_persisted_stats():
else:
stats[k] = v
logger.info("✓ Successfully loaded persisted gateway statistics from disk.")
- # Load timeline from disk (may be stale after pod restart, but better than empty)
- timeline_path = os.path.join(
- os.path.dirname(CONFIG_PATH), "router_timeline.json"
- )
- if os.path.exists(timeline_path):
- try:
- with open(timeline_path, "r") as f:
- stats["timeline"] = json.load(f)
- except Exception:
- pass # stale/broken timeline file → start fresh
except Exception as e:
logger.error(f"Failed to load persisted stats: {e}")
-
def _atomic_write_json_sync(path: str, data) -> None:
"""Synchronously write JSON data to path using atomic temp-file + os.replace."""
os.makedirs(os.path.dirname(path), exist_ok=True)
@@ -427,7 +399,6 @@ async def _atomic_write_json_async(path: str, data) -> None:
_last_stats_save = 0.0
-
async def save_persisted_stats(force=False):
"""Persists current statistics in-memory structure to disk securely (non-blocking).
@@ -449,7 +420,6 @@ async def save_persisted_stats(force=False):
_last_stats_save = 0.0 # Reset on failure to allow immediate retry
logger.error(f"Failed to persist stats to disk: {e}")
-
# Load initial stats from persistent storage
load_persisted_stats()
@@ -458,20 +428,18 @@ async def save_persisted_stats(force=False):
CACHE_TTL_SECONDS = 86400 # Decisions cached for 24 hours
classification_lock = asyncio.Lock()
-
async def _purge_stale_deployments(db_url: str, pattern: str):
"""Purge stale deployments matching the pattern from LiteLLM's DB."""
import asyncpg
-
conn = await asyncpg.connect(db_url)
try:
await conn.execute(
- 'DELETE FROM "LiteLLM_ProxyModelTable" WHERE model_name LIKE $1', pattern
+ 'DELETE FROM "LiteLLM_ProxyModelTable" WHERE model_name LIKE $1',
+ pattern
)
finally:
await conn.close()
-
async def sync_adaptive_router_roster(master_key: str):
"""Fetch free OpenRouter models and register them as deployments in LiteLLM."""
if not master_key:
@@ -499,8 +467,8 @@ async def sync_adaptive_router_roster(master_key: str):
continue
# 1. Enforce Tool/Function Calling Support
- supported_params = m.get("supported_parameters") or []
- if "tools" not in supported_params:
+ supported_params = m.get('supported_parameters') or []
+ if 'tools' not in supported_params:
logger.info(f"🚫 Skipping {mid} — Model does not support tool calling.")
continue
@@ -508,19 +476,14 @@ async def sync_adaptive_router_roster(master_key: str):
# llama-3.3-70b reports 131K ctx but actual endpoint enforces 65K → context_limit errors.
# All meta-llama and llama-derived models are too old and unreliable on free tier.
_denylist_prefixes = (
- "meta-llama/",
- "nousresearch/hermes-3-llama",
+ "meta-llama/", "nousresearch/hermes-3-llama",
)
if any(mid.startswith(p) for p in _denylist_prefixes):
- logger.info(
- f"🚫 Skipping {mid} — denylisted (stale/unreliable free tier model)"
- )
+ logger.info(f"🚫 Skipping {mid} — denylisted (stale/unreliable free tier model)")
continue
pricing = m.get("pricing", {})
- if pricing.get("prompt") in ("0", 0, "0.0", 0.0) and pricing.get(
- "completion"
- ) in ("0", 0, "0.0", 0.0):
+ if pricing.get("prompt") in ("0", 0, "0.0", 0.0) and pricing.get("completion") in ("0", 0, "0.0", 0.0):
try:
score = compute_free_model_score(m)
except Exception:
@@ -533,10 +496,8 @@ async def sync_adaptive_router_roster(master_key: str):
logger.warning("No free models found — skipping roster sync")
return
tier_assignments = {
- "agent-simple-core": [],
- "agent-medium-core": [],
- "agent-complex-core": [],
- "agent-reasoning-core": [],
+ "agent-simple-core": [], "agent-medium-core": [],
+ "agent-complex-core": [], "agent-reasoning-core": [],
"agent-advanced-core": [],
}
# Normalize scores to 0-100 scale based on the actual max score in this roster.
@@ -548,28 +509,16 @@ async def sync_adaptive_router_roster(master_key: str):
max_score = max(raw_scores) if raw_scores else 55.0
if max_score < 1.0:
max_score = 55.0 # safety floor
-
def norm(s: float) -> float:
"""Helper to scale raw model index score against max score in roster to 0-100 range."""
return (s / max_score) * 100.0
-
- for (
- score,
- mid,
- ) in (
- free_models
- ): # include all models — top 2 are also assigned to their correct tier
+ for score, mid in free_models: # include all models — top 2 are also assigned to their correct tier
n = norm(score)
- if n >= 80:
- tier_assignments["agent-advanced-core"].append(mid)
- elif n >= 75:
- tier_assignments["agent-reasoning-core"].append(mid)
- elif n >= 68:
- tier_assignments["agent-complex-core"].append(mid)
- elif n >= 60:
- tier_assignments["agent-medium-core"].append(mid)
- else:
- tier_assignments["agent-simple-core"].append(mid)
+ if n >= 80: tier_assignments["agent-advanced-core"].append(mid)
+ elif n >= 75: tier_assignments["agent-reasoning-core"].append(mid)
+ elif n >= 68: tier_assignments["agent-complex-core"].append(mid)
+ elif n >= 60: tier_assignments["agent-medium-core"].append(mid)
+ else: tier_assignments["agent-simple-core"].append(mid)
# Cascading: models capable of higher tiers also serve lower tiers.
# A model that qualifies for advanced should be available for reasoning,
# complex, and medium requests too — not just advanced. Without this,
@@ -605,11 +554,9 @@ def norm(s: float) -> float:
try:
db_url = os.getenv("DATABASE_URL")
if not db_url:
- logger.warning(
- "DATABASE_URL is not set; skipping purge of stale agent-* deployments"
- )
+ logger.warning("DATABASE_URL is not set; skipping purge of stale agent-* deployments")
else:
- await _purge_stale_deployments(db_url, "agent-%")
+ await _purge_stale_deployments(db_url, 'agent-%')
logger.info("🧹 Purged stale agent-* deployments before roster sync")
except Exception as e:
logger.warning(f"Failed to purge stale deployments (non-fatal): {e}")
@@ -630,30 +577,20 @@ def norm(s: float) -> float:
"mode": "chat",
"max_tokens": ctx_len,
"max_input_tokens": ctx_len,
- "is_public_model_group": True,
- },
+ "is_public_model_group": True
+ }
}
try:
- r = await client.post(
- f"{admin_url}/model/new",
- headers=headers,
- json=payload,
- timeout=10.0,
- )
+ r = await client.post(f"{admin_url}/model/new", headers=headers, json=payload, timeout=10.0)
if r.status_code in (200, 201):
registered += 1
else:
failed += 1
- logger.warning(
- f"model/new {mid} → {tier_name}: HTTP {r.status_code} — {r.text[:200]}"
- )
+ logger.warning(f"model/new {mid} → {tier_name}: HTTP {r.status_code} — {r.text[:200]}")
except Exception as e:
failed += 1
logger.warning(f"Failed to register {mid} under {tier_name}: {e}")
- logger.info(
- f"📊 Roster sync: registered {registered} deployments ({failed} failed) across 5 tiers — {sum(len(v) for v in tier_assignments.values())} attempted"
- )
-
+ logger.info(f"📊 Roster sync: registered {registered} deployments ({failed} failed) across 5 tiers — {sum(len(v) for v in tier_assignments.values())} attempted")
async def _register_ollama_models_in_db(master_key: str):
"""Register static ollama models via /model/new so they become DB models.
@@ -664,23 +601,19 @@ async def _register_ollama_models_in_db(master_key: str):
as null/false. Registering them as DB models ensures our model_info wins.
"""
if not master_key:
- logger.warning(
- "No LiteLLM master key provided — skipping Ollama DB registration"
- )
+ logger.warning("No LiteLLM master key provided — skipping Ollama DB registration")
return
admin_url = LITELLM_URL
headers = {"Authorization": f"Bearer {master_key}", "Content-Type": "application/json"}
ollama_models = []
- litellm_config_path = os.getenv(
- "LITELLM_CONFIG_PATH", "/config/litellm_dir/config.yaml"
- )
+ litellm_config_path = os.getenv("LITELLM_CONFIG_PATH", "/config/litellm_dir/config.yaml")
config_paths_to_try = [
litellm_config_path,
str(Path(__file__).resolve().parent.parent / "litellm" / "config.yaml"),
- "./litellm/config.yaml",
+ "./litellm/config.yaml"
]
def _load_yaml(p):
@@ -696,24 +629,18 @@ def _load_yaml(p):
for item in litellm_config["model_list"]:
if isinstance(item, dict):
model_name = item.get("model_name", "")
- if isinstance(model_name, str) and model_name.startswith(
- "ollama-deepseek-"
- ):
+ if isinstance(model_name, str) and model_name.startswith("ollama-deepseek-"):
# Create a clean deep copy to avoid mutating configuration structures
ollama_models.append(copy.deepcopy(item))
if ollama_models:
- logger.info(
- f"Loaded {len(ollama_models)} Ollama model configurations dynamically from {path}"
- )
+ logger.info(f"Loaded {len(ollama_models)} Ollama model configurations dynamically from {path}")
loaded_from_config = True
break
except Exception as e:
logger.warning(f"Failed to load/parse LiteLLM config at {path}: {e}")
if not loaded_from_config:
- logger.warning(
- "Could not load Ollama models from config.yaml, falling back to static definitions"
- )
+ logger.warning("Could not load Ollama models from config.yaml, falling back to static definitions")
ollama_models = [
{
"model_name": "ollama-deepseek-v4-pro",
@@ -762,14 +689,10 @@ def _load_yaml(p):
try:
db_url = os.getenv("DATABASE_URL")
if not db_url:
- logger.warning(
- "DATABASE_URL is not set; skipping purge of stale ollama-deepseek-* DB entries"
- )
+ logger.warning("DATABASE_URL is not set; skipping purge of stale ollama-deepseek-* DB entries")
else:
- await _purge_stale_deployments(db_url, "ollama-deepseek-%")
- logger.info(
- "🧹 Purged stale ollama-deepseek-* DB entries before registration"
- )
+ await _purge_stale_deployments(db_url, 'ollama-deepseek-%')
+ logger.info("🧹 Purged stale ollama-deepseek-* DB entries before registration")
except Exception as e:
logger.warning(f"Failed to purge stale ollama DB entries (non-fatal): {e}")
@@ -778,16 +701,12 @@ def _load_yaml(p):
failed = 0
for payload in ollama_models:
try:
- r = await client.post(
- f"{admin_url}/model/new", headers=headers, json=payload, timeout=10.0
- )
+ r = await client.post(f"{admin_url}/model/new", headers=headers, json=payload, timeout=10.0)
if r.status_code in (200, 201):
registered += 1
else:
failed += 1
- logger.warning(
- f"model/new {payload['model_name']}: HTTP {r.status_code} — {r.text[:200]}"
- )
+ logger.warning(f"model/new {payload['model_name']}: HTTP {r.status_code} — {r.text[:200]}")
except Exception as e:
failed += 1
logger.warning(f"Failed to register {payload['model_name']}: {e}")
@@ -810,15 +729,13 @@ async def lifespan(app: FastAPI):
try:
r = await client.get(litellm_ready_url, timeout=2.0)
if r.status_code == 200:
- logger.info(f"✅ LiteLLM ready after {i + 1}s")
+ logger.info(f"✅ LiteLLM ready after {i+1}s")
break
except Exception:
pass
await asyncio.sleep(1)
else:
- logger.warning(
- "⚠️ LiteLLM not ready within timeout — proceeding without roster sync"
- )
+ logger.warning("⚠️ LiteLLM not ready within timeout — proceeding without roster sync")
# Sync free-model roster into LiteLLM (non-fatal if it fails)
if litellm_master_key:
@@ -864,17 +781,13 @@ async def lifespan(app: FastAPI):
# Flush any buffered stats/timeline on clean shutdown (always runs)
await save_persisted_stats(force=True)
try:
- timeline_path = os.path.join(
- os.path.dirname(CONFIG_PATH), "router_timeline.json"
- )
+ timeline_path = os.path.join(os.path.dirname(CONFIG_PATH), "router_timeline.json")
await _atomic_write_json_async(timeline_path, stats["timeline"])
except Exception as e:
logger.warning(f"Failed to persist timeline on shutdown: {e}")
-
app = FastAPI(title="LLM Triage Router", lifespan=lifespan)
-
async def check_tcp_port(ip: str, port: int) -> bool:
"""Verifies if a TCP port is open locally asynchronously."""
try:
@@ -885,7 +798,6 @@ async def check_tcp_port(ip: str, port: int) -> bool:
except Exception:
return False
-
async def check_http_endpoint(url: str) -> bool:
"""Verifies if an HTTP endpoint is responsive."""
try:
@@ -895,10 +807,7 @@ async def check_http_endpoint(url: str) -> bool:
except Exception:
return False
-
-async def classify_request(
- prompt: str, bypass_cache: bool = False, langfuse_trace_id: str | None = None
-) -> tuple[str, float, bool, str]:
+async def classify_request(prompt: str, bypass_cache: bool = False, langfuse_trace_id: str | None = None) -> tuple[str, float, bool, str]:
"""Queries the local fast Qwen instance to classify request complexity with TTL caching.
When langfuse_trace_id is provided, the classifier HTTP call is wrapped in a child
@@ -913,9 +822,7 @@ async def classify_request(
if not bypass_cache and normalized_prompt in triage_cache:
cached_decision, cached_time = triage_cache[normalized_prompt]
if time.time() - cached_time < CACHE_TTL_SECONDS:
- logger.info(
- f"⚡ Triage Cache Hit for prompt: '{normalized_prompt[:50]}...' -> routed to '{cached_decision}'"
- )
+ logger.info(f"⚡ Triage Cache Hit for prompt: '{normalized_prompt[:50]}...' -> routed to '{cached_decision}'")
stats["cache_hits"] = stats.get("cache_hits", 0) + 1
await save_persisted_stats()
return cached_decision, 0.0, True, cached_decision # was_cache_hit=True
@@ -928,9 +835,7 @@ async def classify_request(
if not bypass_cache and normalized_prompt in triage_cache:
cached_decision, cached_time = triage_cache[normalized_prompt]
if time.time() - cached_time < CACHE_TTL_SECONDS:
- logger.info(
- f"⚡ Triage Cache Hit (post-queue) for prompt: '{normalized_prompt[:50]}...' -> routed to '{cached_decision}'"
- )
+ logger.info(f"⚡ Triage Cache Hit (post-queue) for prompt: '{normalized_prompt[:50]}...' -> routed to '{cached_decision}'")
stats["cache_hits"] = stats.get("cache_hits", 0) + 1
await save_persisted_stats()
return cached_decision, 0.0, True, cached_decision
@@ -939,15 +844,15 @@ async def classify_request(
client = get_http_client()
payload = {
"model": router_model_name,
- "messages": [{"role": "user", "content": system_prompt + prompt}],
+ "messages": [
+ {"role": "user", "content": system_prompt + prompt}
+ ],
"temperature": 0.0,
"max_tokens": 15,
}
headers = {"Authorization": f"Bearer {router_api_key}"}
- logger.info(
- f"Classifying intent via {router_api_base} using model {router_model_name}..."
- )
+ logger.info(f"Classifying intent via {router_api_base} using model {router_model_name}...")
# --- Langfuse child span: classifier call ---
class_span_obj = None
@@ -969,7 +874,7 @@ async def classify_request(
f"{router_api_base}/chat/completions",
json=payload,
headers=headers,
- timeout=120.0,
+ timeout=120.0
)
latency = (time.time() - start_time) * 1000.0
@@ -978,17 +883,12 @@ async def classify_request(
if class_span_obj:
try:
class_span_obj.end(
- output={
- "status": response.status_code,
- "error": "classification_failed",
- },
+ output={"status": response.status_code, "error": "classification_failed"},
metadata={"latency_ms": latency},
)
except Exception:
pass
- logger.error(
- f"Classification failed with status {response.status_code}: {response.text}"
- )
+ logger.error(f"Classification failed with status {response.status_code}: {response.text}")
return "agent-advanced-core", latency, False, "advanced (fallback)"
result = response.json()
@@ -1000,11 +900,8 @@ async def classify_request(
# 5-tier grammar parsing (was 3-tier, missed medium + advanced)
valid_tiers = {
- "agent-simple-core",
- "agent-medium-core",
- "agent-complex-core",
- "agent-reasoning-core",
- "agent-advanced-core",
+ "agent-simple-core", "agent-medium-core", "agent-complex-core",
+ "agent-reasoning-core", "agent-advanced-core"
}
if content_clean in valid_tiers:
decision = content_clean
@@ -1030,7 +927,6 @@ async def classify_request(
logger.error(f"Exception during classification: {e}")
return "agent-advanced-core", latency, False, "advanced (exception)"
-
def get_live_gemini_oauth_token() -> str | None:
"""Retrieve the current valid Gemini OAuth access token from local storage if not expired."""
try:
@@ -1043,42 +939,29 @@ def get_live_gemini_oauth_token() -> str | None:
# Convert current time to milliseconds
current_ms = int(time.time() * 1000)
if access_token and current_ms < expiry_ms:
- logger.info(
- "🔑 Found valid, unexpired Gemini OAuth token from host!"
- )
+ logger.info("🔑 Found valid, unexpired Gemini OAuth token from host!")
return access_token
else:
# agy CLI uses the OS system keyring (GNOME Keyring), not this
# stale disk file. The file being expired is expected — don't warn.
- logger.debug(
- "Gemini OAuth token on disk is expired — agy uses system keyring instead."
- )
+ logger.debug("Gemini OAuth token on disk is expired — agy uses system keyring instead.")
except Exception as e:
logger.error(f"Failed to read live OAuth token: {e}")
return None
-
def get_gemini_oauth_status() -> dict:
"""Returns structured OAuth status for the dashboard banner."""
creds_path = "/config/gemini_auth/oauth_creds.json"
try:
if not os.path.exists(creds_path):
- return {
- "status": "missing",
- "detail": "No oauth_creds.json found",
- "expiry_ms": 0,
- }
+ return {"status": "missing", "detail": "No oauth_creds.json found", "expiry_ms": 0}
with open(creds_path, "r") as f:
data = json.load(f)
access_token = data.get("access_token")
expiry_ms = data.get("expiry_date", 0)
current_ms = int(time.time() * 1000)
if not access_token:
- return {
- "status": "missing",
- "detail": "No access token in file",
- "expiry_ms": 0,
- }
+ return {"status": "missing", "detail": "No access token in file", "expiry_ms": 0}
diff_sec = (expiry_ms - current_ms) / 1000.0
if diff_sec > 0:
# Token is valid — compute human-readable remaining time
@@ -1088,11 +971,7 @@ def get_gemini_oauth_status() -> dict:
remaining = f"{int(diff_sec // 60)}m {int(diff_sec % 60)}s"
else:
remaining = f"{int(diff_sec // 3600)}h {int((diff_sec % 3600) // 60)}m"
- return {
- "status": "valid",
- "detail": f"Expires in {remaining}",
- "expiry_ms": expiry_ms,
- }
+ return {"status": "valid", "detail": f"Expires in {remaining}", "expiry_ms": expiry_ms}
else:
# Token is expired — compute human-readable elapsed time
elapsed = abs(diff_sec)
@@ -1102,15 +981,10 @@ def get_gemini_oauth_status() -> dict:
ago = f"{int(elapsed // 3600)} hours ago"
else:
ago = f"{int(elapsed // 86400)} days ago"
- return {
- "status": "expired",
- "detail": f"Expired {ago}",
- "expiry_ms": expiry_ms,
- }
+ return {"status": "expired", "detail": f"Expired {ago}", "expiry_ms": expiry_ms}
except Exception as e:
return {"status": "error", "detail": str(e), "expiry_ms": 0}
-
def map_tool_to_category(tool_name: str) -> str:
"""Groups low-level developer tool names into the five high-level dashboard metrics."""
name = tool_name.lower().strip()
@@ -1119,35 +993,14 @@ def map_tool_to_category(tool_name: str) -> str:
if "tree" in name or "list_dir" in name or "list-dir" in name:
return "tree"
- elif (
- "shell" in name
- or "command" in name
- or "cmd" in name
- or "execute" in name
- or "run" in name
- ):
+ elif "shell" in name or "command" in name or "cmd" in name or "execute" in name or "run" in name:
return "shell"
- elif (
- "write" in name
- or "edit" in name
- or "create" in name
- or "patch" in name
- or "replace" in name
- or "save" in name
- ):
+ elif "write" in name or "edit" in name or "create" in name or "patch" in name or "replace" in name or "save" in name:
return "write"
- elif (
- "view" in name
- or "read" in name
- or "cat" in name
- or "grep" in name
- or "search" in name
- or "find" in name
- ):
+ elif "view" in name or "read" in name or "cat" in name or "grep" in name or "search" in name or "find" in name:
return "view"
return "other"
-
def detect_active_tool(body: dict) -> str:
"""Inspects request payload messages to identify which developer tool is currently being invoked."""
messages = body.get("messages", [])
@@ -1170,15 +1023,8 @@ def detect_active_tool(body: dict) -> str:
tcalls = prev_msg.get("tool_calls") or []
if isinstance(tcalls, list):
for tc in tcalls:
-
-
- if (
- isinstance(tc, dict)
- and tc.get("id") == tool_call_id
- ):
+ if isinstance(tc, dict) and tc.get("id") == tool_call_id:
fn = tc.get("function")
-
-
if isinstance(fn, dict):
name = fn.get("name")
break
@@ -1193,9 +1039,7 @@ def detect_active_tool(body: dict) -> str:
for tc in tool_calls:
if isinstance(tc, dict):
fn = tc.get("function")
- name = (
- fn.get("name") if isinstance(fn, dict) else None
- ) or "other"
+ name = (fn.get("name") if isinstance(fn, dict) else None) or "other"
return map_tool_to_category(name)
# Fallback to keyphrase scanning in the user message
@@ -1214,15 +1058,7 @@ def detect_active_tool(body: dict) -> str:
return "view"
return "none"
-
-def record_tool_usage(
- tool_name: str,
- prompt_tokens: int,
- completion_tokens: int,
- model: str,
- latency_ms: float,
- route: str = "litellm_fallback",
-):
+def record_tool_usage(tool_name: str, prompt_tokens: int, completion_tokens: int, model: str, latency_ms: float, route: str = "litellm_fallback"):
"""Accumulates token counts in memory for active tools and tracks request timelines.
File writes are offloaded to a thread pool executor to avoid blocking the
@@ -1251,7 +1087,7 @@ def record_tool_usage(
"model": model,
"route": route,
"tokens": total,
- "latency_ms": int(latency_ms),
+ "latency_ms": int(latency_ms)
}
stats["timeline"].append(event)
if len(stats["timeline"]) > 15:
@@ -1284,7 +1120,7 @@ def record_tool_usage(
None,
_atomic_write_json_sync,
timeline_path,
- copy.deepcopy(list(stats["timeline"])),
+ copy.deepcopy(list(stats["timeline"]))
)
record_tool_usage._last_save = now
@@ -1306,7 +1142,6 @@ def done_callback(f):
except Exception as e:
logger.warning(f"Failed to persist timeline: {e}")
-
def get_goose_sessions() -> list:
"""Queries the live mounted SQLite goose database to fetch the latest agentic sessions."""
sessions_list = []
@@ -1315,25 +1150,22 @@ def get_goose_sessions() -> list:
return []
try:
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
-
async def get_llamacpp_metrics() -> dict:
"""Fetches live model inventory and slot statistics from the local llama-server."""
result = {"models": [], "slots": [], "build": "unknown"}
@@ -1346,16 +1178,14 @@ async def get_llamacpp_metrics() -> dict:
for m in data.get("data", []):
meta = m.get("meta", {})
status_obj = m.get("status", {})
- result["models"].append(
- {
- "id": m.get("id", "?"),
- "status": status_obj.get("value", "unknown"),
- "n_params": meta.get("n_params"),
- "n_ctx": meta.get("n_ctx"),
- "size_bytes": meta.get("size"),
- "n_embd": meta.get("n_embd"),
- }
- )
+ result["models"].append({
+ "id": m.get("id", "?"),
+ "status": status_obj.get("value", "unknown"),
+ "n_params": meta.get("n_params"),
+ "n_ctx": meta.get("n_ctx"),
+ "size_bytes": meta.get("size"),
+ "n_embd": meta.get("n_embd"),
+ })
# Fetch props for build info
r2 = await client.get(f"{LLAMA_SERVER_URL}/props", timeout=3.0)
if r2.status_code == 200:
@@ -1363,15 +1193,9 @@ async def get_llamacpp_metrics() -> dict:
result["build"] = props.get("build_info", "unknown")
# Fetch slots for the loaded model, falling back to the first available model if all are unloaded
loaded = [m["id"] for m in result["models"] if m["status"] == "loaded"]
- slot_model = (
- loaded[0]
- if loaded
- else (result["models"][0]["id"] if result["models"] else None)
- )
+ slot_model = loaded[0] if loaded else (result["models"][0]["id"] if result["models"] else None)
if slot_model:
- r3 = await client.get(
- f"{LLAMA_SERVER_URL}/slots?model={slot_model}", timeout=3.0
- )
+ r3 = await client.get(f"{LLAMA_SERVER_URL}/slots?model={slot_model}", timeout=3.0)
if r3.status_code == 200:
slots_data = r3.json()
for s in slots_data:
@@ -1396,16 +1220,17 @@ async def get_llamacpp_metrics() -> dict:
logger.warning(f"Failed to fetch llama.cpp metrics: {e}")
return result
-
# In-Memory Cache for OpenRouter Free Model list to prevent slow page renders
-free_model_cache = {"data": None, "last_fetched": 0.0}
+free_model_cache = {
+ "data": None,
+ "last_fetched": 0.0
+}
FREE_MODEL_CACHE_TTL = 3600 # Refresh cache every 1 hour
# --- Artificial Analysis Agentic Index scores cache ---
_AA_SCORES_CACHE: dict[str, float] = {}
_AA_SCORES_LOADED = False
-
def _load_aa_scores():
"""Load the Artificial Analysis agentic scores cache from local config."""
global _AA_SCORES_CACHE, _AA_SCORES_LOADED
@@ -1413,27 +1238,22 @@ def _load_aa_scores():
return
try:
import json
-
scores_path = os.path.join(os.path.dirname(__file__), "aa_scores.json")
with open(scores_path) as f:
data = json.load(f)
_AA_SCORES_CACHE = data.get("scores", {})
_AA_SCORES_LOADED = True
- logger.info(
- f"📊 Loaded {len(_AA_SCORES_CACHE)} AA agentic index scores from {scores_path}"
- )
+ logger.info(f"📊 Loaded {len(_AA_SCORES_CACHE)} AA agentic index scores from {scores_path}")
except Exception as e:
logger.warning(f"Could not load AA scores cache: {e}")
_AA_SCORES_LOADED = True # don't retry
-
def compute_free_model_score(m: dict) -> float:
"""Return AA agentic index score, or a low default for unknown models."""
_load_aa_scores()
mid = m.get("id", "")
return _AA_SCORES_CACHE.get(mid, 25.0)
-
def _save_free_models_roster(free_models: list[dict]) -> None:
"""Persist the full sorted free model list so Ralph can try alternatives."""
import json as _json
@@ -1450,7 +1270,7 @@ def _save_free_models_roster(free_models: list[dict]) -> None:
pass
-async def _save_best_model_to_disk(best_model: dict) -> None:
+def _save_best_model_to_disk(best_model: dict) -> None:
"""Persist the best free model to a JSON file Ralph can read."""
import json as _json
import datetime as _dt
@@ -1477,7 +1297,7 @@ async def get_best_free_model() -> dict:
"name": "MoonshotAI: Kimi K2.6 (free)",
"score": 82.5,
"context_length": 131072,
- "is_fallback": True,
+ "is_fallback": True
}
try:
@@ -1493,8 +1313,7 @@ async def get_best_free_model() -> dict:
mid = m.get("id", "")
# Denylist: skip stale/unreliable free tier models
_denylist_prefixes = (
- "meta-llama/",
- "nousresearch/hermes-3-llama",
+ "meta-llama/", "nousresearch/hermes-3-llama",
)
if any(mid.startswith(p) for p in _denylist_prefixes):
continue
@@ -1531,7 +1350,6 @@ async def get_best_free_model() -> dict:
await asyncio.to_thread(_save_best_model_to_disk, fallback_best)
return fallback_best
-
def get_pie_chart_gradient() -> str:
"""Computes a CSS conic-gradient representing the dynamic token distribution across developer tools."""
total_tokens = sum(stats["tool_tokens"].values())
@@ -1554,7 +1372,6 @@ def get_pie_chart_gradient() -> str:
return f"background: conic-gradient({', '.join(gradient_parts)});"
-
@app.api_route("/v1/memory{path:path}", methods=["GET", "POST", "DELETE", "PUT"])
async def proxy_memory(request: Request, path: str = ""):
"""Proxies memory API calls to the LiteLLM gateway on port 4000."""
@@ -1573,12 +1390,10 @@ async def proxy_memory(request: Request, path: str = ""):
litellm_key = os.getenv("LITELLM_MASTER_KEY")
headers = {
"Authorization": f"Bearer {litellm_key}",
- "Content-Type": request.headers.get("content-type", "application/json"),
+ "Content-Type": request.headers.get("content-type", "application/json")
}
- logger.info(
- f"Proxying memory request: {request.method} {url} with params {query_params}"
- )
+ logger.info(f"Proxying memory request: {request.method} {url} with params {query_params}")
try:
client = get_http_client()
@@ -1588,27 +1403,23 @@ async def proxy_memory(request: Request, path: str = ""):
params=query_params,
content=body,
headers=headers,
- timeout=30.0,
+ timeout=30.0
)
# Return response matching status and headers
response_headers = dict(r.headers)
# Exclude standard headers that FastAPI/uvicorn will manage
- for h in [
- "content-encoding",
- "content-length",
- "transfer-encoding",
- "connection",
- ]:
+ for h in ["content-encoding", "content-length", "transfer-encoding", "connection"]:
response_headers.pop(h, None)
return Response(
- content=r.content, status_code=r.status_code, headers=response_headers
+ content=r.content,
+ status_code=r.status_code,
+ headers=response_headers
)
except Exception as e:
logger.error(f"Failed to proxy memory request: {e}")
- raise HTTPException(status_code=502, detail="Memory proxy failed")
-
+ raise HTTPException(status_code=502, detail=f"Memory proxy failed: {e}")
@app.get("/v1/models")
async def proxy_models():
@@ -1620,7 +1431,7 @@ async def proxy_models():
r = await client.get(
f"{LITELLM_URL}/v1/models",
headers={"Authorization": auth_header},
- timeout=10.0,
+ timeout=10.0
)
if r.status_code == 200:
@@ -1634,73 +1445,31 @@ async def proxy_models():
# - auto-ollama / auto-agy-ollama / llm-routing-ollama: 524288 (512K)
# - llm-routing-agy: 1048576 (1M)
routing_models = [
- {
- "id": "llm-routing-auto-free",
- "object": "model",
- "created": 0,
- "owned_by": "llm-routing",
- "context_length": 262144,
- },
- {
- "id": "llm-routing-auto-agy",
- "object": "model",
- "created": 0,
- "owned_by": "llm-routing",
- "context_length": 262144,
- },
- {
- "id": "llm-routing-auto-ollama",
- "object": "model",
- "created": 0,
- "owned_by": "llm-routing",
- "context_length": 524288,
- },
- {
- "id": "llm-routing-auto-agy-ollama",
- "object": "model",
- "created": 0,
- "owned_by": "llm-routing",
- "context_length": 524288,
- },
- {
- "id": "llm-routing-agy",
- "object": "model",
- "created": 0,
- "owned_by": "llm-routing",
- "context_length": 1048576,
- },
- {
- "id": "llm-routing-ollama",
- "object": "model",
- "created": 0,
- "owned_by": "llm-routing",
- "context_length": 524288,
- },
+ {"id": "llm-routing-auto-free", "object": "model", "created": 0, "owned_by": "llm-routing", "context_length": 262144},
+ {"id": "llm-routing-auto-agy", "object": "model", "created": 0, "owned_by": "llm-routing", "context_length": 262144},
+ {"id": "llm-routing-auto-ollama", "object": "model", "created": 0, "owned_by": "llm-routing", "context_length": 524288},
+ {"id": "llm-routing-auto-agy-ollama", "object": "model", "created": 0, "owned_by": "llm-routing", "context_length": 524288},
+ {"id": "llm-routing-agy", "object": "model", "created": 0, "owned_by": "llm-routing", "context_length": 1048576},
+ {"id": "llm-routing-ollama", "object": "model", "created": 0, "owned_by": "llm-routing", "context_length": 524288},
]
data["data"] = routing_models + data["data"]
return JSONResponse(content=data, status_code=200)
except Exception as parse_err:
- logger.warning(
- f"Failed to parse /v1/models JSON despite status 200: {parse_err}"
- )
+ logger.warning(f"Failed to parse /v1/models JSON despite status 200: {parse_err}")
# If not 200, or parsing failed, return the raw response with appropriate headers
response_headers = dict(r.headers)
- for h in [
- "content-encoding",
- "content-length",
- "transfer-encoding",
- "connection",
- ]:
+ for h in ["content-encoding", "content-length", "transfer-encoding", "connection"]:
response_headers.pop(h, None)
return Response(
- content=r.content, status_code=r.status_code, headers=response_headers
+ content=r.content,
+ status_code=r.status_code,
+ headers=response_headers
)
except Exception as e:
logger.error(f"Failed to proxy /v1/models: {e}")
- raise HTTPException(status_code=502, detail="Model proxy failed")
-
+ raise HTTPException(status_code=502, detail=f"Model proxy failed: {e}")
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
@@ -1730,29 +1499,21 @@ async def chat_completions(request: Request):
if msg.get("role") == "user":
content = msg.get("content") or ""
if isinstance(content, list):
- content = "".join(
- block.get("text") or ""
- for block in content
- if isinstance(block, dict) and block.get("type") == "text"
- )
+ content = "".join(block.get("text") or "" for block in content if isinstance(block, dict) and block.get("type") == "text")
last_user_message = str(content)
break
# Known tier names that can be routed directly (bypass classifier)
DIRECT_TIERS = {
- "agent-simple-core",
- "agent-medium-core",
- "agent-complex-core",
- "agent-reasoning-core",
+ "agent-simple-core", "agent-medium-core",
+ "agent-complex-core", "agent-reasoning-core",
"agent-advanced-core",
"llm-routing-agy",
}
AUTO_MODELS = {
- "llm-routing-auto-free",
- "llm-routing-auto-agy",
- "llm-routing-auto-ollama",
- "llm-routing-auto-agy-ollama",
+ "llm-routing-auto-free", "llm-routing-auto-agy",
+ "llm-routing-auto-ollama", "llm-routing-auto-agy-ollama",
}
client_model = body.get("model", "llm-routing-auto-free")
@@ -1763,9 +1524,7 @@ async def chat_completions(request: Request):
lf = get_langfuse()
if lf:
try:
- langfuse_trace_id = lf.create_trace_id(
- seed=f"triage_{stats['total_requests']}"
- )
+ langfuse_trace_id = lf.create_trace_id(seed=f"triage_{stats['total_requests']}")
parent_obs = lf.start_observation(
trace_context={"trace_id": langfuse_trace_id},
name=f"triage-{client_model}",
@@ -1780,15 +1539,8 @@ async def chat_completions(request: Request):
if client_model in AUTO_MODELS or client_model == "llm-routing-ollama":
# Full pipeline: classify → route to best tier
bypass_cache = request.headers.get("x-bypass-cache") == "true"
- (
- target_model,
- triage_latency,
- was_cache_hit,
- raw_classification,
- ) = await classify_request(
- last_user_message,
- bypass_cache=bypass_cache,
- langfuse_trace_id=langfuse_trace_id,
+ target_model, triage_latency, was_cache_hit, raw_classification = await classify_request(
+ last_user_message, bypass_cache=bypass_cache, langfuse_trace_id=langfuse_trace_id
)
logger.info(f"Triage decision (auto/gated): Routing to -> '{target_model}'")
elif client_model in DIRECT_TIERS:
@@ -1797,23 +1549,19 @@ async def chat_completions(request: Request):
triage_latency = 0.0
was_cache_hit = False
raw_classification = f"direct ({client_model})"
- logger.info(
- f"Direct routing: Client requested '{client_model}', skipping classifier"
- )
+ logger.info(f"Direct routing: Client requested '{client_model}', skipping classifier")
else:
raise HTTPException(
status_code=400,
detail=f"Unknown model '{client_model}'. Use 'llm-routing-auto-free' for automatic routing, "
- f"or one of: {', '.join(sorted(DIRECT_TIERS))}",
+ f"or one of: {', '.join(sorted(DIRECT_TIERS))}"
)
# Update in-memory statistics
stats["total_requests"] += 1
stats["last_triage_decision"] = target_model
stats["total_triage_time_ms"] += triage_latency
- stats["avg_triage_latency_ms"] = (
- stats["total_triage_time_ms"] / stats["total_requests"]
- )
+ stats["avg_triage_latency_ms"] = stats["total_triage_time_ms"] / stats["total_requests"]
if target_model == "agent-simple-core":
stats["simple_requests"] = stats.get("simple_requests", 0) + 1
@@ -1861,19 +1609,11 @@ async def chat_completions(request: Request):
should_try_agy = (
client_model == "llm-routing-agy" # direct — always try
- or (
- client_model in ("llm-routing-auto-agy", "llm-routing-auto-agy-ollama")
- and target_model in ("agent-advanced-core", "agent-reasoning-core")
- )
+ or (client_model in ("llm-routing-auto-agy", "llm-routing-auto-agy-ollama") and target_model in ("agent-advanced-core", "agent-reasoning-core"))
)
should_try_ollama = (
- client_model
- == "llm-routing-ollama" # always try (will map to flash for complex/below)
- or (
- client_model in ("llm-routing-auto-ollama", "llm-routing-auto-agy-ollama")
- and target_model
- in ("agent-advanced-core", "agent-reasoning-core", "agent-complex-core")
- )
+ client_model == "llm-routing-ollama" # always try (will map to flash for complex/below)
+ or (client_model in ("llm-routing-auto-ollama", "llm-routing-auto-agy-ollama") and target_model in ("agent-advanced-core", "agent-reasoning-core", "agent-complex-core"))
)
# --- AGY PROXY ---
@@ -1889,11 +1629,7 @@ async def chat_completions(request: Request):
if msg.get("role") == "user":
content = msg.get("content") or ""
if isinstance(content, list):
- content = "".join(
- block.get("text") or ""
- for block in content
- if isinstance(block, dict) and block.get("type") == "text"
- )
+ content = "".join(block.get("text") or "" for block in content if isinstance(block, dict) and block.get("type") == "text")
last_prompt = str(content)
break
@@ -1930,7 +1666,7 @@ async def chat_completions(request: Request):
stream=is_stream_requested,
target_tier=target_model,
client=get_http_client(),
- cooldown_persistence=ValkeyCooldownPersistence(),
+ cooldown_persistence=ValkeyCooldownPersistence()
)
if agy_response:
model_name = agy_response.get("model", "gemini-3.5-flash (via agy)")
@@ -1940,7 +1676,6 @@ async def chat_completions(request: Request):
async def native_agy_stream_generator(stream_gen, model_name):
"""Asynchronous generator yielding native OpenAI-compatible streaming chunks from the real agy daemon."""
import uuid
-
created_time = int(time.time())
chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
token_count = 0
@@ -1954,17 +1689,13 @@ async def native_agy_stream_generator(stream_gen, model_name):
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
- "choices": [
- {
- "index": 0,
- "delta": {"content": token},
- "finish_reason": None,
- }
- ],
+ "choices": [{
+ "index": 0,
+ "delta": {"content": token},
+ "finish_reason": None
+ }]
}
- yield f"data: {json.dumps(chunk_data)}\n\n".encode(
- "utf-8"
- )
+ yield f"data: {json.dumps(chunk_data)}\n\n".encode("utf-8")
# End of stream chunk
finish_data = {
@@ -1972,17 +1703,13 @@ async def native_agy_stream_generator(stream_gen, model_name):
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
- "choices": [
- {
- "index": 0,
- "delta": {},
- "finish_reason": "stop",
- }
- ],
+ "choices": [{
+ "index": 0,
+ "delta": {},
+ "finish_reason": "stop"
+ }]
}
- yield f"data: {json.dumps(finish_data)}\n\n".encode(
- "utf-8"
- )
+ yield f"data: {json.dumps(finish_data)}\n\n".encode("utf-8")
yield b"data: [DONE]\n\n"
# Success telemetry
@@ -1990,114 +1717,74 @@ async def native_agy_stream_generator(stream_gen, model_name):
approx_prompt_tokens = estimate_prompt_tokens(body)
record_tool_usage(
- active_tool,
- approx_prompt_tokens,
- token_count,
- model_name,
- latency_ms,
- route="google_oauth_direct",
- )
- logger.info(
- f"✅ native agy stream succeeded: {model_name}, {latency_ms:.0f}ms"
+ active_tool, approx_prompt_tokens, token_count,
+ model_name, latency_ms, route="google_oauth_direct"
)
+ logger.info(f"✅ native agy stream succeeded: {model_name}, {latency_ms:.0f}ms")
if agy_span_obj:
try:
agy_span_obj.end(
- output={
- "model": model_name,
- "tokens": token_count,
- },
- metadata={
- "latency_ms": latency_ms,
- "tier": target_model,
- },
+ output={"model": model_name, "tokens": token_count},
+ metadata={"latency_ms": latency_ms, "tier": target_model},
)
except Exception:
pass
except Exception as stream_err:
- logger.error(
- f"Error during native agy stream generation: {type(stream_err).__name__}"
- )
+ logger.error(f"Error during native agy stream generation: {stream_err}")
if agy_span_obj:
try:
agy_span_obj.end(
- output={"error": type(stream_err).__name__},
+ output={"error": str(stream_err)[:200]},
metadata={"status": "failed"},
)
except Exception:
pass
raise
-
- return StreamingResponse(
- native_agy_stream_generator(
- agy_response["stream"], model_name
- ),
- media_type="text/event-stream",
- )
+ return StreamingResponse(native_agy_stream_generator(agy_response["stream"], model_name), media_type="text/event-stream")
else:
latency_ms = (time.time() - start_time) * 1000.0
usage = agy_response.get("usage") or {}
prompt_tokens = usage.get("prompt_tokens") or 0
completion_tokens = usage.get("completion_tokens") or 0
record_tool_usage(
- active_tool,
- prompt_tokens,
- completion_tokens,
- model_name,
- latency_ms,
- route="google_oauth_direct",
- )
- logger.info(
- f"✅ agy proxy succeeded: {model_name}, {latency_ms:.0f}ms"
+ active_tool, prompt_tokens, completion_tokens,
+ model_name, latency_ms, route="google_oauth_direct"
)
+ logger.info(f"✅ agy proxy succeeded: {model_name}, {latency_ms:.0f}ms")
# Finalize agy span
if agy_span_obj:
try:
agy_span_obj.end(
- output={
- "model": model_name,
- "tokens": completion_tokens,
- },
- metadata={
- "latency_ms": latency_ms,
- "tier": target_model,
- },
+ output={"model": model_name, "tokens": completion_tokens},
+ metadata={"latency_ms": latency_ms, "tier": target_model},
)
except Exception:
pass
if is_stream_requested:
# Robust fallback: simulate stream if we requested stream but got buffered response
- content = (agy_response.get("choices") or [{}])[0].get(
- "message", {}
- ).get("content") or ""
-
+ content = (agy_response.get("choices") or [{}])[0].get("message", {}).get("content") or ""
async def agy_stream_generator():
"""Asynchronous generator yielding simulated OpenAI-compatible streaming chunks from a static agy response."""
import uuid
-
created_time = int(time.time())
chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
chunk_size = 40
for i in range(0, len(content), chunk_size):
- chunk_text = content[i : i + chunk_size]
+ chunk_text = content[i:i+chunk_size]
chunk_data = {
"id": chunk_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
- "choices": [
- {
- "index": 0,
- "delta": {"content": chunk_text},
- "finish_reason": None,
- }
- ],
+ "choices": [{
+ "index": 0,
+ "delta": {"content": chunk_text},
+ "finish_reason": None
+ }]
}
- yield f"data: {json.dumps(chunk_data)}\n\n".encode(
- "utf-8"
- )
+ yield f"data: {json.dumps(chunk_data)}\n\n".encode("utf-8")
await asyncio.sleep(0.005)
finish_data = {
@@ -2105,22 +1792,15 @@ async def agy_stream_generator():
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
- "choices": [
- {
- "index": 0,
- "delta": {},
- "finish_reason": "stop",
- }
- ],
+ "choices": [{
+ "index": 0,
+ "delta": {},
+ "finish_reason": "stop"
+ }]
}
- yield f"data: {json.dumps(finish_data)}\n\n".encode(
- "utf-8"
- )
+ yield f"data: {json.dumps(finish_data)}\n\n".encode("utf-8")
yield b"data: [DONE]\n\n"
-
- return StreamingResponse(
- agy_stream_generator(), media_type="text/event-stream"
- )
+ return StreamingResponse(agy_stream_generator(), media_type="text/event-stream")
else:
return agy_response
except ImportError:
@@ -2137,12 +1817,12 @@ async def agy_stream_generator():
if agy_span_obj:
try:
agy_span_obj.end(
- output={"error": type(e).__name__},
+ output={"error": str(e)[:200]},
metadata={"status": "failed"},
)
except Exception:
pass
- logger.error(f"agy proxy failed: {type(e).__name__}, falling back to LiteLLM")
+ logger.error(f"agy proxy failed: {e}, falling back to LiteLLM")
original_target_model = target_model
@@ -2174,9 +1854,7 @@ async def execute_proxy(model_name: str):
backend_conf = backends.get(model_name)
if not backend_conf:
logger.error(f"Backend '{model_name}' not found in configuration backends.")
- raise HTTPException(
- status_code=500, detail=f"Backend {model_name} misconfigured"
- )
+ raise HTTPException(status_code=500, detail=f"Backend {model_name} misconfigured")
backend_api_base = backend_conf["api_base"]
backend_api_key = backend_conf["api_key"]
@@ -2231,7 +1909,7 @@ async def execute_proxy(model_name: str):
if _safe_max < 1024:
raise HTTPException(
status_code=400,
- detail=f"Context window exceeded. Estimated input tokens ({_est_input}) plus safety margin (2048) exceeds model context limit ({_min_ctx}).",
+ detail=f"Context window exceeded. Estimated input tokens ({_est_input}) plus safety margin (2048) exceeds model context limit ({_min_ctx})."
)
if requested_max_tokens > _safe_max:
logger.warning(
@@ -2245,27 +1923,18 @@ async def execute_proxy(model_name: str):
logger.warning(f"Pre-screening failed (non-fatal): {e}")
body_to_send = body.copy()
body_to_send["model"] = model_name
- if "metadata" not in body_to_send or not isinstance(
- body_to_send["metadata"], dict
- ):
+ if "metadata" not in body_to_send or not isinstance(body_to_send["metadata"], dict):
body_to_send["metadata"] = {}
body_to_send["metadata"]["trace_name"] = "agent-completion"
if body.get("stream", False):
logger.info(f"Proxying streaming to LiteLLM as model={model_name}")
- req = client.build_request(
- "POST",
- f"{backend_api_base}/chat/completions",
- json=body_to_send,
- headers=headers,
- )
+ req = client.build_request("POST", f"{backend_api_base}/chat/completions", json=body_to_send, headers=headers)
r = await client.send(req, stream=True)
if r.status_code == 200:
-
async def stream_generator():
"""Asynchronous generator that yields streaming chunks from LiteLLM completions response and logs usage stats on completion."""
import codecs
-
completion_chars = 0
request_tokens = estimate_prompt_tokens(body_to_send)
sse_buffer = ""
@@ -2285,14 +1954,10 @@ async def stream_generator():
try:
data_json = json.loads(data_str)
choices = data_json.get("choices", [])
- if choices and isinstance(
- choices[0], dict
- ):
+ if choices and isinstance(choices[0], dict):
delta = choices[0].get("delta")
if isinstance(delta, dict):
- content = (
- delta.get("content") or ""
- )
+ content = delta.get("content") or ""
completion_chars += len(content)
except Exception:
pass
@@ -2300,26 +1965,14 @@ async def stream_generator():
pass
proxy_latency = (time.time() - proxy_start) * 1000.0
stats["total_proxy_time_ms"] += proxy_latency
- stats["avg_proxy_latency_ms"] = (
- stats["total_proxy_time_ms"] / stats["total_requests"]
- )
- record_tool_usage(
- active_tool,
- request_tokens,
- completion_chars // 4,
- model_name,
- proxy_latency,
- route="litellm_fallback",
- )
+ stats["avg_proxy_latency_ms"] = stats["total_proxy_time_ms"] / stats["total_requests"]
+ record_tool_usage(active_tool, request_tokens, completion_chars // 4, model_name, proxy_latency, route="litellm_fallback")
# Finalize LiteLLM span (streaming path)
if litellm_span_obj:
try:
litellm_span_obj.end(
output={"model": model_name, "stream": True},
- metadata={
- "latency_ms": proxy_latency,
- "tokens": completion_chars // 4,
- },
+ metadata={"latency_ms": proxy_latency, "tokens": completion_chars // 4},
)
except Exception:
pass
@@ -2327,97 +1980,58 @@ async def stream_generator():
logger.error(f"Stream error: {ex}")
if model_name.startswith("ollama-"):
global _ollama_cooldown_until
- _ollama_cooldown_until = (
- time.monotonic() + OLLAMA_COOLDOWN_SECONDS
- )
+ _ollama_cooldown_until = time.monotonic() + OLLAMA_COOLDOWN_SECONDS
try:
await save_cooldowns_to_valkey()
logger.error(
f"🧊 Ollama failed midway through stream, activating {OLLAMA_COOLDOWN_SECONDS}s cooldown"
)
except Exception as save_err:
- logger.warning(
- f"Failed to save cooldowns to Valkey: {save_err}"
- )
+ logger.warning(f"Failed to save cooldowns to Valkey: {save_err}")
finally:
await r.aclose()
-
- return StreamingResponse(
- stream_generator(), media_type="text/event-stream"
- )
+ return StreamingResponse(stream_generator(), media_type="text/event-stream")
else:
error_body = await r.aread() if r else b""
- logger.warning(
- f"LiteLLM stream failed ({r.status_code}): {error_body[:300]}"
- )
+ logger.warning(f"LiteLLM stream failed ({r.status_code}): {error_body[:300]}")
await r.aclose()
- raise HTTPException(
- status_code=r.status_code,
- detail="LiteLLM upstream request failed",
- )
+ raise HTTPException(status_code=r.status_code, detail="LiteLLM upstream request failed")
else:
logger.info(f"Proxying to LiteLLM as model={model_name}")
- response = await client.post(
- f"{backend_api_base}/chat/completions",
- json=body_to_send,
- headers=headers,
- )
+ response = await client.post(f"{backend_api_base}/chat/completions", json=body_to_send, headers=headers)
if response.status_code == 200:
proxy_latency = (time.time() - proxy_start) * 1000.0
stats["total_proxy_time_ms"] += proxy_latency
- stats["avg_proxy_latency_ms"] = (
- stats["total_proxy_time_ms"] / stats["total_requests"]
- )
+ stats["avg_proxy_latency_ms"] = stats["total_proxy_time_ms"] / stats["total_requests"]
resp_json = response.json()
usage = resp_json.get("usage") or {}
- prompt_tokens = usage.get(
- "prompt_tokens"
- ) or estimate_prompt_tokens(body_to_send)
+ prompt_tokens = usage.get("prompt_tokens") or estimate_prompt_tokens(body_to_send)
choices = resp_json.get("choices") or []
fallback_completion = 0
if choices and isinstance(choices[0], dict):
msg = choices[0].get("message")
if isinstance(msg, dict):
fallback_completion = len(msg.get("content") or "") // 4
- completion_tokens = (
- usage.get("completion_tokens") or fallback_completion
- )
- record_tool_usage(
- active_tool,
- prompt_tokens,
- completion_tokens,
- model_name,
- proxy_latency,
- route="litellm_fallback",
- )
+ completion_tokens = usage.get("completion_tokens") or fallback_completion
+ record_tool_usage(active_tool, prompt_tokens, completion_tokens, model_name, proxy_latency, route="litellm_fallback")
# Finalize LiteLLM span (non-streaming path)
if litellm_span_obj:
try:
litellm_span_obj.end(
- output={
- "model": model_name,
- "tokens": completion_tokens,
- },
+ output={"model": model_name, "tokens": completion_tokens},
metadata={"latency_ms": proxy_latency},
)
except Exception:
pass
return resp_json
else:
- logger.warning(
- f"LiteLLM failed ({response.status_code}): {response.text[:300]}"
- )
- raise HTTPException(
- status_code=response.status_code,
- detail="LiteLLM upstream request failed",
- )
+ logger.warning(f"LiteLLM failed ({response.status_code}): {response.text[:300]}")
+ raise HTTPException(status_code=response.status_code, detail="LiteLLM upstream request failed")
except HTTPException:
raise
except Exception as exc:
logger.error(f"httpx call failed: {exc}")
- raise HTTPException(
- status_code=502, detail="Proxy call failed"
- ) from exc
+ raise HTTPException(status_code=502, detail=f"Proxy call failed: {exc}") from exc
if should_try_ollama:
# Sync state from Valkey first
@@ -2432,21 +2046,16 @@ async def stream_generator():
f"⏳ Ollama cooldown active ({remaining}s remaining), "
f"skipping {target_model}"
)
- if client_model in (
- "llm-routing-auto-ollama",
- "llm-routing-auto-agy-ollama",
- ):
+ if client_model in ("llm-routing-auto-ollama", "llm-routing-auto-agy-ollama"):
# Auto mode: silently fall through to the free tier
- logger.info(
- f"Auto-mode fallback: {target_model} → {original_target_model} (Ollama cooled down)"
- )
+ logger.info(f"Auto-mode fallback: {target_model} → {original_target_model} (Ollama cooled down)")
return await execute_proxy(original_target_model)
else:
# Direct/fallback llm-routing-ollama: return 429 so LiteLLM
# skips this model group and moves to openrouter-auto
raise HTTPException(
status_code=429,
- detail=f"Ollama backend cooled down ({remaining}s remaining)",
+ detail=f"Ollama backend cooled down ({remaining}s remaining)"
)
try:
@@ -2461,26 +2070,19 @@ async def stream_generator():
logger.error(
f"🧊 Ollama failed ({e.status_code}), activating {OLLAMA_COOLDOWN_SECONDS}s cooldown"
)
- if client_model in (
- "llm-routing-auto-ollama",
- "llm-routing-auto-agy-ollama",
- ):
+ if client_model in ("llm-routing-auto-ollama", "llm-routing-auto-agy-ollama"):
if is_transient:
- logger.warning(
- f"Ollama proxy failed ({e.detail}), falling back to free tier {original_target_model}"
- )
+ logger.warning(f"Ollama proxy failed ({e.detail}), falling back to free tier {original_target_model}")
return await execute_proxy(original_target_model)
else:
raise e
else:
# Direct/fallback llm-routing-ollama request
if is_transient:
- logger.error(
- f"Ollama proxy failed ({e.detail}) for direct/fallback request, returning 429"
- )
+ logger.error(f"Ollama proxy failed ({e.detail}) for direct/fallback request, returning 429")
raise HTTPException(
status_code=429,
- detail="Ollama backend rate limited/unavailable",
+ detail="Ollama backend rate limited/unavailable"
) from e
else:
raise e
@@ -2491,22 +2093,17 @@ async def stream_generator():
logger.error(
f"🧊 Ollama unexpected error ({e}), activating {OLLAMA_COOLDOWN_SECONDS}s cooldown"
)
- if client_model in (
- "llm-routing-auto-ollama",
- "llm-routing-auto-agy-ollama",
- ):
- logger.warning(
- f"Ollama proxy error ({e}), falling back to free tier {original_target_model}"
- )
+ if client_model in ("llm-routing-auto-ollama", "llm-routing-auto-agy-ollama"):
+ logger.warning(f"Ollama proxy error ({e}), falling back to free tier {original_target_model}")
return await execute_proxy(original_target_model)
else:
raise HTTPException(
- status_code=429, detail="Ollama backend rate limited/unavailable"
+ status_code=429,
+ detail="Ollama backend rate limited/unavailable"
) from e
else:
return await execute_proxy(target_model)
-
@app.get("/metrics")
async def metrics():
"""Expose triage and circuit breaker metrics in Prometheus format."""
@@ -2565,50 +2162,36 @@ async def metrics():
# Circuit breaker metrics — dual breaker (google + vendor)
google = breaker_status["google"]
vendor = breaker_status["vendor"]
- lines.append(
- "# HELP circuit_breaker_google_tier Google breaker cooldown tier (0=open, 3=max)"
- )
+ lines.append("# HELP circuit_breaker_google_tier Google breaker cooldown tier (0=open, 3=max)")
lines.append("# TYPE circuit_breaker_google_tier gauge")
lines.append(f"circuit_breaker_google_tier {google['tier']}")
- lines.append(
- "# HELP circuit_breaker_vendor_tier Vendor breaker cooldown tier (0=open, 3=max)"
- )
+ lines.append("# HELP circuit_breaker_vendor_tier Vendor breaker cooldown tier (0=open, 3=max)")
lines.append("# TYPE circuit_breaker_vendor_tier gauge")
lines.append(f"circuit_breaker_vendor_tier {vendor['tier']}")
- lines.append(
- "# HELP circuit_breaker_agy_allowed Whether EITHER breaker allows agy (backward-compat)"
- )
+ lines.append("# HELP circuit_breaker_agy_allowed Whether EITHER breaker allows agy (backward-compat)")
lines.append("# TYPE circuit_breaker_agy_allowed gauge")
lines.append(f"circuit_breaker_agy_allowed {int(breaker.is_allowed_peek())}")
lines.append("# HELP circuit_breaker_total_trips Total trips across both breakers")
lines.append("# TYPE circuit_breaker_total_trips counter")
- lines.append(
- f"circuit_breaker_total_trips {google['total_trips'] + vendor['total_trips']}"
- )
+ lines.append(f"circuit_breaker_total_trips {google['total_trips'] + vendor['total_trips']}")
# Ollama router-side cooldown metrics
_now_mono = time.monotonic()
_ollama_remaining = max(0.0, _ollama_cooldown_until - _now_mono)
- lines.append(
- "# HELP ollama_cooldown_active Whether Ollama is in router-side cooldown (1=active)"
- )
+ lines.append("# HELP ollama_cooldown_active Whether Ollama is in router-side cooldown (1=active)")
lines.append("# TYPE ollama_cooldown_active gauge")
lines.append(f"ollama_cooldown_active {int(_ollama_remaining > 0)}")
- lines.append(
- "# HELP ollama_cooldown_remaining_seconds Seconds remaining in Ollama cooldown"
- )
+ lines.append("# HELP ollama_cooldown_remaining_seconds Seconds remaining in Ollama cooldown")
lines.append("# TYPE ollama_cooldown_remaining_seconds gauge")
lines.append(f"ollama_cooldown_remaining_seconds {_ollama_remaining:.0f}")
return Response(content="\n".join(lines), media_type="text/plain; version=0.0.4")
-
# Source badge helper: generates a colored inline source tag
def src_badge(label, color):
"""Generate inline HTML span styled as a colored status/category badge."""
return f"{label}"
-
async def get_dashboard_data():
"""Fetch all metrics and pre-compute HTML snippets for the dashboard."""
# Run ALL independent I/O concurrently with protective timeouts
@@ -2721,31 +2304,11 @@ async def get_dashboard_data():
# 3. Calculative metrics — 5-tier triage table
tier_data = [
- {
- "tier": "agent-simple-core",
- "count": stats.get("simple_requests", 0),
- "color": "#34d399",
- },
- {
- "tier": "agent-medium-core",
- "count": stats.get("medium_requests", 0),
- "color": "#fbbf24",
- },
- {
- "tier": "agent-complex-core",
- "count": stats.get("complex_requests", 0),
- "color": "#a78bfa",
- },
- {
- "tier": "agent-reasoning-core",
- "count": stats.get("reasoning_requests", 0),
- "color": "#60a5fa",
- },
- {
- "tier": "agent-advanced-core",
- "count": stats.get("advanced_requests", 0),
- "color": "#f472b6",
- },
+ {"tier": "agent-simple-core", "count": stats.get("simple_requests", 0), "color": "#34d399"},
+ {"tier": "agent-medium-core", "count": stats.get("medium_requests", 0), "color": "#fbbf24"},
+ {"tier": "agent-complex-core", "count": stats.get("complex_requests", 0), "color": "#a78bfa"},
+ {"tier": "agent-reasoning-core", "count": stats.get("reasoning_requests", 0), "color": "#60a5fa"},
+ {"tier": "agent-advanced-core", "count": stats.get("advanced_requests", 0), "color": "#f472b6"},
]
total_tier = sum(t["count"] for t in tier_data)
for t in tier_data:
@@ -2756,12 +2319,12 @@ async def get_dashboard_data():
for t in tier_data:
tier_table_rows += f"""
- |
-
- {t["tier"]}
+ |
+
+ {t['tier']}
|
- {t["count"]} |
- {t["ratio"]:.1f}% |
+ {t['count']} |
+ {t['ratio']:.1f}% |
"""
tier_table_html = f"""
@@ -2790,6 +2353,7 @@ async def get_dashboard_data():
pct = (token_count / max_tool_val) * 100.0
overall_pct = (token_count / total_tool_tokens * 100.0) if total_tool_tokens > 0 else 0.0
color = TOOL_COLORS.get(tool_name, "#94a3b8")
+
# Horizontal meters
tool_tokens_html += f"""
@@ -2818,26 +2382,22 @@ async def get_dashboard_data():
timeline_html = "
Waiting for active tool executions...
"
else:
for ev in reversed(stats["timeline"]):
- route_label = ev.get("route", "litellm_fallback")
- route_color = (
- "#fbbf24" if route_label == "google_oauth_direct" else "#818cf8"
- )
- route_short = (
- "GOOGLE" if route_label == "google_oauth_direct" else "LITELLM"
- )
+ route_label = ev.get('route', 'litellm_fallback')
+ route_color = '#fbbf24' if route_label == 'google_oauth_direct' else '#818cf8'
+ route_short = 'GOOGLE' if route_label == 'google_oauth_direct' else 'LITELLM'
timeline_html += f"""
- 🔧 {ev["tool"]} {route_short}
- {ev["timestamp"]}
+ 🔧 {ev['tool']} {route_short}
+ {ev['timestamp']}
- Processed {ev["tokens"]:,} tokens on {ev["model"]}
+ Processed {ev['tokens']:,} tokens on {ev['model']}
- Latency: {ev["latency_ms"]} ms
+ Latency: {ev['latency_ms']} ms
@@ -2853,51 +2413,44 @@ async def get_dashboard_data():
"""
else:
for idx, sess in enumerate(goose_sessions):
- is_active = idx == 0
- badge_style = (
- "background: rgba(129, 140, 248, 0.15); color: #c084fc; border: 1px solid rgba(129, 140, 248, 0.3);"
- if is_active
- else "background: rgba(255,255,255,0.03); color: #fff; border: 1px solid rgba(255,255,255,0.05);"
- )
- active_label = (
- "
ACTIVE"
- if is_active
- else ""
- )
+ is_active = (idx == 0)
+ badge_style = "background: rgba(129, 140, 248, 0.15); color: #c084fc; border: 1px solid rgba(129, 140, 248, 0.3);" if is_active else "background: rgba(255,255,255,0.03); color: #fff; border: 1px solid rgba(255,255,255,0.05);"
+ active_label = "
ACTIVE" if is_active else ""
- desc = sess.get("description") or sess.get("name") or "Interactive session"
- tokens = sess.get("accumulated_total_tokens", 0) or 0
+ desc = sess.get('description') or sess.get('name') or "Interactive session"
+ tokens = sess.get('accumulated_total_tokens', 0) or 0
goose_html += f"""
{active_label}
- Session {sess["id"]}
+ Session {sess['id']}
-
{sess.get("goose_mode", "auto").upper()}
+
{sess.get('goose_mode', 'auto').upper()}
{desc}
- 📅 {sess["updated_at"]}
+ 📅 {sess['updated_at']}
{tokens:,} total tokens
"""
# 8. Routing Paths pie chart & legend
- routing_paths = stats.get(
- "routing_paths", {"google_oauth_direct": 0, "litellm_fallback": 0}
- )
+ routing_paths = stats.get("routing_paths", {"google_oauth_direct": 0, "litellm_fallback": 0})
total_routed = sum(routing_paths.values())
routing_pie_gradient = "background: rgba(255, 255, 255, 0.05);"
routing_legend_html = ""
- routing_colors = {"google_oauth_direct": "#fbbf24", "litellm_fallback": "#818cf8"}
+ routing_colors = {
+ "google_oauth_direct": "#fbbf24",
+ "litellm_fallback": "#818cf8"
+ }
routing_labels = {
"google_oauth_direct": "Google OAuth Direct",
- "litellm_fallback": "LiteLLM Fallback",
+ "litellm_fallback": "LiteLLM Fallback"
}
if total_routed > 0:
current_angle = 0.0
@@ -2915,9 +2468,7 @@ async def get_dashboard_data():
- {m["id"]}
- {m["status"].upper()}
+ {m['id']}
+ {m['status'].upper()}
{params_str}{ctx_str}{size_str}
@@ -2967,18 +2504,14 @@ async def get_dashboard_data():
if llamacpp["slots"]:
slot_items = ""
for sl in llamacpp["slots"]:
- dot_style = (
- "background: #34d399; box-shadow: 0 0 8px #34d399;"
- if sl["is_processing"]
- else "background: rgba(255,255,255,0.15);"
- )
+ dot_style = "background: #34d399; box-shadow: 0 0 8px #34d399;" if sl["is_processing"] else "background: rgba(255,255,255,0.15);"
slot_items += f"""
-
Slot {sl["id"]}
+
Slot {sl['id']}
- Prompt: {sl["n_prompt_processed"]} tok
- Decoded: {sl["n_decoded"]} tok
+ Prompt: {sl['n_prompt_processed']} tok
+ Decoded: {sl['n_decoded']} tok
"""
@@ -3017,16 +2550,14 @@ async def get_dashboard_data():
"avg_proxy_latency_ms": stats["avg_proxy_latency_ms"],
"cache_hits": stats["cache_hits"],
"total_requests": stats["total_requests"],
- "last_triage_decision": stats["last_triage_decision"],
+ "last_triage_decision": stats["last_triage_decision"]
}
-
@app.get("/api/dashboard-stats")
async def get_dashboard_stats():
"""Return dashboard metrics and pre-computed HTML as JSON for asynchronous UI updates."""
return await get_dashboard_data()
-
@app.get("/dashboard", response_class=HTMLResponse)
async def get_dashboard():
"""Render the router main dashboard HTML showing system metrics, health checks, and recent token usage."""
@@ -3568,7 +3099,7 @@ async def get_dashboard():
- {src_badge("ROUTER", "#818cf8")} Gateway Performance Telemetry
+ {src_badge('ROUTER', '#818cf8')} Gateway Performance Telemetry
Persistent telemetry
@@ -3596,7 +3127,7 @@ async def get_dashboard():
-
{src_badge("ROUTER", "#818cf8")} Triage Routing Split
+
{src_badge('ROUTER', '#818cf8')} Triage Routing Split
{tier_table_html}
@@ -3606,7 +3137,7 @@ async def get_dashboard():
- {src_badge("ROUTER", "#818cf8")} Tool Token Distribution
+ {src_badge('ROUTER', '#818cf8')} Tool Token Distribution
Live conic-gradient pie
@@ -3639,7 +3170,7 @@ async def get_dashboard():
- {src_badge("ROUTER", "#818cf8")} Routing Path Distribution
+ {src_badge('ROUTER', '#818cf8')} Routing Path Distribution
% requests per path
@@ -3655,7 +3186,7 @@ async def get_dashboard():
- {src_badge("LITELLM", "#34d399")} Model Usage
+ {src_badge('LITELLM', '#34d399')} Model Usage
Full traces in Langfuse
@@ -3667,7 +3198,7 @@ async def get_dashboard():
- {src_badge("GOOSE", "#fbbf24")} Live Tool Token Meters
+ {src_badge('GOOSE', '#fbbf24')} Live Tool Token Meters
Token meters per extension tool
@@ -3757,8 +3288,8 @@ async def get_dashboard():
Langfuse Traces
:3001
-
- {"Online" if langfuse_status else "Offline"}
+
+ {'Online' if langfuse_status else 'Offline'}
@@ -3766,8 +3297,8 @@ async def get_dashboard():
- {src_badge("LLAMA.CPP", "#fb923c")} Engine Metrics
- build {data["llamacpp_build"]}
+ {src_badge('LLAMA.CPP', '#fb923c')} Engine Metrics
+ build {data['llamacpp_build']}
{llamacpp_models_html}
@@ -3779,7 +3310,7 @@ async def get_dashboard():
-
{src_badge("GOOSE", "#fbbf24")} Session Directory
+
{src_badge('GOOSE', '#fbbf24')} Session Directory
{goose_html}
@@ -3791,21 +3322,21 @@ async def get_dashboard():
@@ -3821,7 +3352,6 @@ async def get_dashboard():
"""
return html_content
-
# --- Static files (visualizer, data files) ---
STATIC_DIR = Path(__file__).resolve().parent / "static"
DATA_DIR = Path(__file__).resolve().parent / "data"
@@ -3829,7 +3359,6 @@ async def get_dashboard():
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
app.mount("/data", StaticFiles(directory=str(DATA_DIR)), name="data")
-
@app.get("/visualizer", response_class=HTMLResponse)
async def get_visualizer():
"""Serve the dataset visualizer for human review."""
@@ -3840,52 +3369,13 @@ async def get_visualizer():
return HTMLResponse("
Visualizer not found
", status_code=404)
-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="forbid")
-
tier: Union[int, str, None] = None
- note: Optional[str] = Field(default=None, max_length=1000)
- ts: Optional[str] = Field(default=None, max_length=100)
-
- @field_validator("tier")
- @classmethod
- def validate_tier(cls, v):
- if v is None:
- return v
- if isinstance(v, int):
- if v < 0 or v > 4:
- raise ValueError(f"Invalid tier index {v}: must be between 0 and 4")
- elif isinstance(v, str):
- if v not in VALID_TIERS and v != "?":
- raise ValueError(f"Invalid tier string '{v}'")
- else:
- raise ValueError("Tier must be int, str, or null")
- return v
-
-class AnnotationPayload(RootModel):
- root: Dict[str, AnnotationItem]
-
- @model_validator(mode="after")
- 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, 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: str = ""
+ ts: Optional[str] = None
+
+VALID_TIERS = {"agent-simple-core", "agent-medium-core", "agent-complex-core", "agent-reasoning-core", "agent-advanced-core"}
# NOTE: annotations_lock (asyncio.Lock) only provides concurrency protection within
# a single Python process. In multi-worker uvicorn deployments, concurrent requests
# across different workers can still race. Eventual consistency is maintained via
@@ -3910,44 +3400,73 @@ def _read_annotations_sync(path) -> dict:
return copy.deepcopy(_annotations_cache[path]["data"])
-
@app.post("/dashboard/save-annotations")
-async def save_annotations(payload: AnnotationPayload):
+async def save_annotations(payload: Dict[str, AnnotationItem]):
"""Save human review annotations to disk."""
+ if len(payload) > 1000:
+ raise HTTPException(
+ status_code=400,
+ detail="Payload size limit exceeded: maximum of 1000 annotations allowed per request."
+ )
+ for k, item in payload.items():
+ # Allow numeric strings (dataset indexes) or stable hash keys starting with 'h' (hexadecimal)
+ 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 HTTPException(
+ status_code=400,
+ detail=f"Invalid payload key '{k}': keys must be numeric strings or stable hash keys (e.g., 'h12345abc')."
+ )
+
+ t = item.tier
+ if t is not None:
+ if isinstance(t, int):
+ if t < 0 or t > 4:
+ raise HTTPException(
+ status_code=400,
+ detail=f"Invalid tier index {t} for index {k}: must be between 0 and 4."
+ )
+ elif isinstance(t, str):
+ if t not in VALID_TIERS and t != "?":
+ raise HTTPException(
+ status_code=400,
+ detail=f"Invalid tier string '{t}' for index {k}."
+ )
+ else:
+ raise HTTPException(
+ status_code=400,
+ detail=f"Invalid tier type for index {k}: must be int, str, or null."
+ )
+
+ if item.note and len(item.note) > 1000:
+ raise HTTPException(
+ status_code=400,
+ detail=f"Note length limit exceeded at index {k}: maximum of 1000 characters allowed."
+ )
try:
- data = payload.root
ann_path = DATA_DIR / "annotations.json"
existing = {}
async with annotations_lock:
if ann_path.exists():
try:
- existing = await asyncio.to_thread(
- _read_annotations_sync, str(ann_path)
- )
+ existing = await asyncio.to_thread(_read_annotations_sync, str(ann_path))
except Exception as read_err:
- logger.warning(
- f"Could not read existing annotations: {read_err}. Overwriting."
- )
+ logger.warning(f"Could not read existing annotations: {read_err}. Overwriting.")
# Merge new annotations into existing
- for k, item in data.items():
- # For partial updates, merge only fields provided in the request
- update_data = item.model_dump(exclude_unset=True)
- if k in existing and isinstance(existing[k], dict):
- existing[k].update(update_data)
- else:
- existing[k] = item.model_dump()
+ for k, item in payload.items():
+ existing[k] = item.model_dump() if hasattr(item, "model_dump") else item.dict()
+
await _atomic_write_json_async(str(ann_path), existing)
- return JSONResponse({"status": "ok", "saved": len(data)})
+ return JSONResponse({"status": "ok", "saved": len(payload)})
except Exception as e:
logger.error(f"Failed to save annotations: {e}")
raise HTTPException(status_code=500, detail="Failed to save annotations")
-
if __name__ == "__main__":
import uvicorn
-
logger.info(f"Starting LLM Triage Router on {host}:{port}...")
uvicorn.run(app, host=host, port=port)
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 3358439e..404059eb 100644
--- a/router/tests/test_agy_proxy.py
+++ b/router/tests/test_agy_proxy.py
@@ -1,13 +1,3 @@
-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
from router.agy_proxy import _wrap_response, _is_quota_exhausted
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_detect_active_tool.py b/router/tests/test_detect_active_tool.py
new file mode 100644
index 00000000..3105ab1e
--- /dev/null
+++ b/router/tests/test_detect_active_tool.py
@@ -0,0 +1,114 @@
+import pytest
+import os
+import sys
+from pathlib import Path
+
+# Set CONFIG_PATH for import
+os.environ["CONFIG_PATH"] = str(Path(__file__).resolve().parent.parent / "config.yaml")
+
+# Add the parent directory to the path so we can import from router
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from router.main import detect_active_tool
+
+def test_detect_active_tool_empty():
+ assert detect_active_tool({}) == "none"
+ assert detect_active_tool({"messages": []}) == "none"
+ assert detect_active_tool({"messages": [{"role": "system", "content": "hello"}]}) == "none"
+
+def test_detect_active_tool_role_tool_with_name():
+ # Write tool name mapped to write
+ body = {
+ "messages": [
+ {"role": "tool", "name": "edit_file", "content": "..."}
+ ]
+ }
+ assert detect_active_tool(body) == "write"
+
+ # View tool mapped to view
+ body = {
+ "messages": [
+ {"role": "tool", "name": "cat_file", "content": "..."}
+ ]
+ }
+ assert detect_active_tool(body) == "view"
+
+def test_detect_active_tool_role_tool_without_name_but_matched_tool_call_id():
+ body = {
+ "messages": [
+ {"role": "assistant", "tool_calls": [{"id": "call_123", "function": {"name": "read_file"}}]},
+ {"role": "tool", "tool_call_id": "call_123", "content": "success"}
+ ]
+ }
+ assert detect_active_tool(body) == "view"
+
+def test_detect_active_tool_role_tool_without_name_unmatched_tool_call_id():
+ body = {
+ "messages": [
+ {"role": "assistant", "tool_calls": [{"id": "call_999", "function": {"name": "read_file"}}]},
+ {"role": "tool", "tool_call_id": "call_123", "content": "success"}
+ ]
+ }
+ assert detect_active_tool(body) == "other"
+
+def test_detect_active_tool_role_assistant_with_tool_calls():
+ body = {
+ "messages": [
+ {"role": "user", "content": "do something"},
+ {"role": "assistant", "tool_calls": [{"id": "call_1", "function": {"name": "write_to_file"}}]}
+ ]
+ }
+ assert detect_active_tool(body) == "write"
+
+def test_detect_active_tool_fallback_user_keyword():
+ # Tests matching "tree"
+ body = {
+ "messages": [
+ {"role": "user", "content": "show me the tree"}
+ ]
+ }
+ assert detect_active_tool(body) == "tree"
+
+ # Tests matching "shell"
+ body = {
+ "messages": [
+ {"role": "user", "content": "run this in shell"}
+ ]
+ }
+ assert detect_active_tool(body) == "shell"
+
+ # Tests matching "write"
+ body = {
+ "messages": [
+ {"role": "user", "content": "create file test.py"}
+ ]
+ }
+ assert detect_active_tool(body) == "write"
+
+ # Tests matching "view"
+ body = {
+ "messages": [
+ {"role": "user", "content": "cat main.py"}
+ ]
+ }
+ assert detect_active_tool(body) == "view"
+
+def test_detect_active_tool_ignores_invalid_message_formats():
+ body = {
+ "messages": [
+ "this is not a dict",
+ {"role": "user", "content": "read this"}
+ ]
+ }
+ assert detect_active_tool(body) == "view"
+
+def test_detect_active_tool_precedence():
+ # If there are multiple tools, it processes from the last message backwards
+ body = {
+ "messages": [
+ {"role": "tool", "name": "edit_file", "content": "done"},
+ {"role": "tool", "name": "cat_file", "content": "..."}
+ ]
+ }
+ # It starts from the end, so it sees "cat_file" first, which maps to "view"
+ assert detect_active_tool(body) == "view"
diff --git a/router/tests/test_estimate_prompt_tokens.py b/router/tests/test_estimate_prompt_tokens.py
index 2a12f97c..ae74a9e5 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
@@ -22,25 +23,27 @@ def test_estimate_prompt_tokens_empty_messages():
def test_estimate_prompt_tokens_string_content():
body = {
"messages": [
- {"content": "1234"}, # 1 token
- {"content": "12345678"} # 2 tokens
+ {"content": "word " * 4}, # 4 * 1.2 = 4.8
+ {"content": "word " * 8} # 8 * 1.2 = 9.6
]
}
- assert estimate_prompt_tokens(body) == 50 + 1 + 2
+ # Total is int(round(4.8 + 9.6)) + 50 = int(round(14.4)) + 50 = 14 + 50 = 64
+ assert estimate_prompt_tokens(body) == 50 + 14
def test_estimate_prompt_tokens_list_content():
body = {
"messages": [
{
"content": [
- {"type": "text", "text": "1234"}, # 1 token
+ {"type": "text", "text": "word " * 4}, # 4.8 tokens
{"type": "image_url", "url": "ignored"}, # 0 tokens
- {"type": "text", "text": "12345678"} # 2 tokens
+ {"type": "text", "text": "word " * 8} # 9.6 tokens
]
}
]
}
- assert estimate_prompt_tokens(body) == 50 + 1 + 2
+ # Total is int(round(4.8 + 9.6)) + 50 = 64
+ assert estimate_prompt_tokens(body) == 50 + 14
def test_estimate_prompt_tokens_mixed_and_invalid_msgs():
body = {
@@ -51,10 +54,11 @@ def test_estimate_prompt_tokens_mixed_and_invalid_msgs():
"invalid_block_type", # Should be skipped
{"type": "text", "text": None} # None text, handled as empty string
]},
- {"content": "1234"} # 1 token
+ {"content": "word " * 4} # 4.8 tokens
]
}
- assert estimate_prompt_tokens(body) == 50 + 1
+ # Total is int(round(4.8)) + 50 = 5 + 50 = 55
+ assert estimate_prompt_tokens(body) == 50 + 5
def test_estimate_prompt_tokens_missing_content():
body = {
diff --git a/router/tests/test_get_goose_sessions.py b/router/tests/test_get_goose_sessions.py
new file mode 100644
index 00000000..52640fef
--- /dev/null
+++ b/router/tests/test_get_goose_sessions.py
@@ -0,0 +1,46 @@
+import pytest
+import sys
+import os
+from pathlib import Path
+from unittest.mock import patch, MagicMock
+
+os.environ["CONFIG_PATH"] = str(Path(__file__).resolve().parent.parent / "config.yaml")
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from main import get_goose_sessions
+
+def test_get_goose_sessions_no_db():
+ with patch('os.path.exists', return_value=False):
+ assert get_goose_sessions() == []
+
+def test_get_goose_sessions_success():
+ mock_sqlite3 = MagicMock()
+ mock_conn = MagicMock()
+ mock_cursor = MagicMock()
+
+ mock_sqlite3.connect.return_value = mock_conn
+ mock_conn.cursor.return_value = mock_cursor
+
+ mock_cursor.fetchall.return_value = [
+ {"id": 1, "name": "s1"},
+ {"id": 2, "name": "s2"}
+ ]
+
+ with patch('os.path.exists', return_value=True):
+ with patch.dict(sys.modules, {'sqlite3': mock_sqlite3}):
+ result = get_goose_sessions()
+
+ assert len(result) == 2
+ assert result[0] == {"id": 1, "name": "s1"}
+ mock_sqlite3.connect.assert_called_once_with("/config/goose_sessions/sessions/sessions.db", timeout=1.0)
+ mock_cursor.execute.assert_called_once()
+ mock_conn.close.assert_called_once()
+
+def test_get_goose_sessions_exception():
+ mock_sqlite3 = MagicMock()
+ mock_sqlite3.connect.side_effect = Exception("DB error")
+
+ with patch('os.path.exists', return_value=True):
+ with patch.dict(sys.modules, {'sqlite3': mock_sqlite3}):
+ result = get_goose_sessions()
+ assert result == []
diff --git a/router/tests/test_load_persisted_stats.py b/router/tests/test_load_persisted_stats.py
index 43cf52ff..76edbf51 100644
--- a/router/tests/test_load_persisted_stats.py
+++ b/router/tests/test_load_persisted_stats.py
@@ -1,99 +1,61 @@
-import pytest
-import os
import json
-import sys
+import pytest
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
+import router.main
+from router.main import load_persisted_stats
-def test_load_persisted_stats_success():
- mock_stats = {
+@pytest.fixture
+def mock_stats():
+ # Setup a clean stats dictionary for testing
+ clean_stats = {
+ "total_requests": 0,
+ "nested_dict": {"a": 1, "b": 2},
+ "existing_key": "value"
+ }
+ with patch.dict(router.main.stats, clean_stats, clear=True):
+ yield router.main.stats
+
+def test_load_persisted_stats_file_not_exists(mock_stats):
+ with patch("router.main.os.path.exists", return_value=False) as mock_exists:
+ load_persisted_stats()
+ mock_exists.assert_called_once_with(router.main.STATS_JSON_PATH)
+ # Stats should remain unchanged
+ assert mock_stats["total_requests"] == 0
+
+def test_load_persisted_stats_success(mock_stats):
+ mock_data = {
"total_requests": 100,
- "some_dict": {"a": 1, "b": 2},
+ "nested_dict": {"b": 3, "c": 4},
"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]
+ mock_json = json.dumps(mock_data)
+
+ with patch("router.main.os.path.exists", return_value=True):
+ with patch("router.main.open", mock_open(read_data=mock_json)):
+ with patch("router.main.logger.info") as mock_logger:
+ load_persisted_stats()
+
+ # Assert simple value updated via else block
+ assert mock_stats["total_requests"] == 100
+ # Assert nested_dict updated via if block (b updated, c added, a unchanged)
+ assert mock_stats["nested_dict"] == {"a": 1, "b": 3, "c": 4}
+ # Assert new_key added via else block
+ assert mock_stats["new_key"] == "new_value"
+ # Assert existing_key unchanged
+ assert mock_stats["existing_key"] == "value"
+
+ mock_logger.assert_called_once_with("✓ Successfully loaded persisted gateway statistics from disk.")
+
+def test_load_persisted_stats_exception(mock_stats):
+ with patch("router.main.os.path.exists", return_value=True):
+ with patch("router.main.open", side_effect=Exception("Mock read error")):
+ with patch("router.main.logger.error") as mock_logger:
+ load_persisted_stats()
+
+ # Stats should remain unchanged
+ assert mock_stats["total_requests"] == 0
+
+ # Error should be logged
+ mock_logger.assert_called_once()
+ assert "Failed to load persisted stats: Mock read error" 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/scripts/benchmark_tokens.py b/scripts/benchmark_tokens.py
new file mode 100644
index 00000000..d701b4eb
--- /dev/null
+++ b/scripts/benchmark_tokens.py
@@ -0,0 +1,76 @@
+import sys
+import os
+from pathlib import Path
+
+# Set CONFIG_PATH and ROUTER_API_KEY for import
+os.environ["CONFIG_PATH"] = str(Path(__file__).resolve().parent.parent / "router" / "config.yaml")
+os.environ["ROUTER_API_KEY"] = "local-token"
+# Add the parent directory and the router directory to the path
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "router"))
+
+from router.main import estimate_prompt_tokens
+
+def verify_accuracy():
+ """Benchmarking utility to verify token estimation accuracy across content types."""
+ # Test cases inspired by the problem description
+ test_cases = [
+ {
+ "name": "English prose",
+ "content": "This is a standard English prose sentence intended to evaluate the accuracy of the token estimation heuristic for typical content. " * 5,
+ "actual_tokens": 110,
+ },
+ {
+ "name": "Python code",
+ "content": """
+def calculate_factorial(n):
+ if n == 0:
+ return 1
+ else:
+ return n * calculate_factorial(n-1)
+
+for i in range(10):
+ print(f"Factorial of {i} is {calculate_factorial(i)}")
+""" * 3,
+ "actual_tokens": 150,
+ },
+ {
+ "name": "CJK text",
+ "content": "这是一个测试，用于验证中文字符的令牌估算逻辑。它应该比字符计数更准确。" * 5,
+ "actual_tokens": 60,
+ },
+ {
+ "name": "Whitespace-padded JSON",
+ "content": '{\n "key": "value",\n "nested": {\n "inner": "data"\n }\n}\n' * 5,
+ "actual_tokens": 60,
+ },
+ {
+ "name": "Emoji",
+ "content": "🚀🔥-🤖✨-📈💎-🚨🛠️-🌐" * 5,
+ "actual_tokens": 25,
+ }
+ ]
+
+ print(f"{'Case':<25} | {'Actual':<7} | {'Estimated':<9} | {'Error':<7}")
+ print("-" * 55)
+
+ all_passed = True
+ for case in test_cases:
+ body = {"messages": [{"content": case["content"]}]}
+ est = estimate_prompt_tokens(body) - 50 # Subtract metadata overhead
+ error = abs(est - case["actual_tokens"]) / case["actual_tokens"]
+ print(f"{case['name']:<25} | {case['actual_tokens']:<7} | {est:<9} | {error:.1%}")
+ # Acceptance criteria: within ±25% for these rough heuristics
+ if error > 0.25:
+ print(f" --> FAILURE: {case['name']} error exceeds target threshold")
+ all_passed = False
+
+ assert all_passed, "Token estimation accuracy benchmark failed"
+
+if __name__ == "__main__":
+ try:
+ verify_accuracy()
+ sys.exit(0)
+ except AssertionError as e:
+ print(f"\nERROR: {e}")
+ sys.exit(1)
diff --git a/start-stack.sh b/start-stack.sh
index d59746de..381b864c 100755
--- a/start-stack.sh
+++ b/start-stack.sh
@@ -31,12 +31,6 @@ 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 ! 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
@@ -74,11 +68,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 +95,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 "$POSTGRES_PASSWORD" ] || [ -z "$NEXTAUTH_SECRET" ] || [ -z "$SALT" ] || [ -z "$ENCRYPTION_KEY" ] || [ -z "$LITELLM_MASTER_KEY" ] || [ -z "$ROUTER_API_KEY" ] || [ -z "$MINIO_ROOT_USER" ] || [ -z "$MINIO_ROOT_PASSWORD" ]; 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 +106,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"
@@ -156,53 +135,31 @@ if [ -z "$LITELLM_MASTER_KEY" ]; then
exit 1
fi
+if [ -z "$LANGFUSE_INIT_USER_PASSWORD" ]; then
+ LANGFUSE_INIT_USER_PASSWORD="$(openssl rand -hex 16)"
+ echo "LANGFUSE_INIT_USER_PASSWORD=\"$LANGFUSE_INIT_USER_PASSWORD\"" >> "$ENV_FILE"
+ echo "✓ Generated new LANGFUSE_INIT_USER_PASSWORD and saved to $ENV_FILE"
+fi
+
+
if [ -z "$ROUTER_API_KEY" ]; then
ROUTER_API_KEY="$(openssl rand -hex 32)"
echo "ROUTER_API_KEY=\"$ROUTER_API_KEY\"" >> "$ENV_FILE"
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"
+if [ -z "$MINIO_ROOT_USER" ]; then
+ MINIO_ROOT_USER="minio-$(openssl rand -hex 4)"
+ echo "MINIO_ROOT_USER=\"$MINIO_ROOT_USER\"" >> "$ENV_FILE"
+ echo "✓ Generated new MINIO_ROOT_USER 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"
+if [ -z "$MINIO_ROOT_PASSWORD" ]; then
+ MINIO_ROOT_PASSWORD="$(openssl rand -hex 16)"
+ echo "MINIO_ROOT_PASSWORD=\"$MINIO_ROOT_PASSWORD\"" >> "$ENV_FILE"
+ echo "✓ Generated new MINIO_ROOT_PASSWORD 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
@@ -312,7 +269,7 @@ setup_minio_buckets() {
# Ensure mc alias points to the correct MinIO S3 API port (9002, not 9000)
# The default 'local' alias in the MinIO image points to :9000 which is ClickHouse,
# not MinIO. We must override it.
- podman exec agent-router-pod-minio-s3 mc alias set local http://127.0.0.1:9002 minioadmin minioadmin 2>/dev/null
+ podman exec agent-router-pod-minio-s3 mc alias set local http://127.0.0.1:9002 "$MINIO_ROOT_USER" "$MINIO_ROOT_PASSWORD" 2>/dev/null
# Create required buckets (idempotent)
local BUCKETS=("langfuse-events" "proj-triage-gateway-id")
@@ -413,7 +370,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 LANGFUSE_INIT_USER_PASSWORD MINIO_ROOT_USER MINIO_ROOT_PASSWORD
python3 - "$WORKDIR/pod.yaml" <<'PY'
import os, sys, urllib.parse
uid = os.getuid()
@@ -423,34 +380,34 @@ 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-***",
+ "MINIO_USER_PLACEHOLDER",
+ "MINIO_PASSWORD_PLACEHOLDER"
+ "LANGFUSE_INIT_USER_PASSWORD_PLACEHOLDER"
]
for ph in placeholders:
if ph not in text:
- sys.stderr.write(f"Error: Required placeholder '{ph}' not found in pod.yaml\n")
+ sys.stderr.write(f"Error: Required placeholder '{ph}' not found in pod.yaml. Ensure you are using the latest version of the template.\n")
sys.exit(1)
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"])
+text = text.replace("MINIO_USER_PLACEHOLDER", os.environ["MINIO_ROOT_USER"])
+text = text.replace("MINIO_PASSWORD_PLACEHOLDER", os.environ["MINIO_ROOT_PASSWORD"])
+text = text.replace("LANGFUSE_INIT_USER_PASSWORD_PLACEHOLDER", os.environ["LANGFUSE_INIT_USER_PASSWORD"])
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/test_host_agy_daemon.py b/test_host_agy_daemon.py
new file mode 100644
index 00000000..e0cae61c
--- /dev/null
+++ b/test_host_agy_daemon.py
@@ -0,0 +1,44 @@
+import os
+import json
+import pytest
+from unittest.mock import patch, mock_open
+
+import host_agy_daemon
+
+def test_get_last_conversation_id_success(monkeypatch):
+ test_data = {"/test/path": "conv_123"}
+ monkeypatch.setattr(os, "getcwd", lambda: "/test/path")
+ monkeypatch.setattr(os.path, "exists", lambda x: True)
+
+ m_open = mock_open(read_data=json.dumps(test_data))
+ with patch("builtins.open", m_open):
+ assert host_agy_daemon.get_last_conversation_id() == "conv_123"
+
+def test_get_last_conversation_id_not_found(monkeypatch):
+ test_data = {"/other/path": "conv_123"}
+ monkeypatch.setattr(os, "getcwd", lambda: "/test/path")
+ monkeypatch.setattr(os.path, "exists", lambda x: True)
+
+ m_open = mock_open(read_data=json.dumps(test_data))
+ with patch("builtins.open", m_open):
+ assert host_agy_daemon.get_last_conversation_id() is None
+
+def test_get_last_conversation_id_file_missing(monkeypatch):
+ monkeypatch.setattr(os.path, "exists", lambda x: False)
+ assert host_agy_daemon.get_last_conversation_id() is None
+
+def test_get_last_conversation_id_exception(monkeypatch):
+ monkeypatch.setattr(os.path, "exists", lambda x: True)
+ # Invalid JSON to trigger JSONDecodeError
+ m_open = mock_open(read_data="{invalid json")
+ with patch("builtins.open", m_open):
+ assert host_agy_daemon.get_last_conversation_id() is None
+
+def test_get_last_conversation_id_io_error(monkeypatch):
+ monkeypatch.setattr(os.path, "exists", lambda x: True)
+
+ def raise_error(*args, **kwargs):
+ raise IOError("permission denied")
+
+ with patch("builtins.open", side_effect=raise_error):
+ assert host_agy_daemon.get_last_conversation_id() is None
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 68%
rename from tests/test_pie_chart_gradient.py
rename to test_pie_chart_gradient.py
index 1d9a33bd..c9e85499 100644
--- a/tests/test_pie_chart_gradient.py
+++ b/test_pie_chart_gradient.py
@@ -4,22 +4,26 @@
@pytest.fixture
def mock_stats():
- with patch("router.main.stats") as mock_stats_obj:
- yield mock_stats_obj
+ # Patch router.main.stats with a real dictionary containing tool_tokens
+ # to ensure the function under test reads from the correct key.
+ test_stats = {
+ "tool_tokens": {
+ "tree": 0,
+ "shell": 0,
+ "write": 0,
+ "view": 0,
+ "other": 0
+ }
+ }
+ with patch("router.main.stats", test_stats):
+ yield test_stats
def test_get_pie_chart_gradient_empty(mock_stats):
- mock_stats.__getitem__.return_value = {
- "tree": 0,
- "shell": 0,
- "write": 0,
- "view": 0,
- "other": 0
- }
result = get_pie_chart_gradient()
assert result == "background: rgba(255, 255, 255, 0.05);"
def test_get_pie_chart_gradient_one_tool(mock_stats):
- mock_stats.__getitem__.return_value = {
+ mock_stats["tool_tokens"] = {
"tree": 100,
"shell": 0,
"write": 0,
@@ -30,7 +34,7 @@ def test_get_pie_chart_gradient_one_tool(mock_stats):
assert result == "background: conic-gradient(#34d399 0.0% 100.0%);"
def test_get_pie_chart_gradient_multiple_tools(mock_stats):
- mock_stats.__getitem__.return_value = {
+ mock_stats["tool_tokens"] = {
"tree": 50,
"shell": 25,
"write": 25,
@@ -41,7 +45,7 @@ def test_get_pie_chart_gradient_multiple_tools(mock_stats):
assert result == "background: conic-gradient(#34d399 0.0% 50.0%, #fbbf24 50.0% 75.0%, #a78bfa 75.0% 100.0%);"
def test_get_pie_chart_gradient_unrecognized_tool(mock_stats):
- mock_stats.__getitem__.return_value = {
+ mock_stats["tool_tokens"] = {
"unknown_tool": 100
}
result = get_pie_chart_gradient()
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 ==="