mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-08-05 02:45:28 +00:00
277 lines
9.5 KiB
Python
277 lines
9.5 KiB
Python
"""Regression coverage for strict foreground model selection."""
|
|
|
|
import asyncio
|
|
import json
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import src.agent_loop as agent_loop
|
|
import src.endpoint_resolver as endpoint_resolver
|
|
import src.foreground_model_routing as foreground_model_routing
|
|
import routes.chat_routes as chat_routes
|
|
from src.foreground_model_routing import (
|
|
build_foreground_model_candidates,
|
|
resolve_foreground_fallback_candidates,
|
|
)
|
|
|
|
|
|
def _collect(gen):
|
|
async def _run():
|
|
return [chunk async for chunk in gen]
|
|
|
|
return asyncio.run(_run())
|
|
|
|
|
|
class _EmptyQuery:
|
|
def filter(self, *args, **kwargs):
|
|
return self
|
|
|
|
def order_by(self, *args, **kwargs):
|
|
return self
|
|
|
|
def first(self):
|
|
return None
|
|
|
|
|
|
class _EmptyDb:
|
|
def query(self, *args, **kwargs):
|
|
return _EmptyQuery()
|
|
|
|
def close(self):
|
|
return None
|
|
|
|
|
|
class _RouteRequest:
|
|
def __init__(self, mode):
|
|
self.headers = {}
|
|
self.app = SimpleNamespace(state=SimpleNamespace(auth_manager=None))
|
|
self._form = {
|
|
"message": "hello",
|
|
"session": "session-1",
|
|
"mode": mode,
|
|
"compare_mode": "true",
|
|
}
|
|
|
|
async def form(self):
|
|
return self._form
|
|
|
|
|
|
def _chat_stream_endpoint(monkeypatch, mode, captured):
|
|
session = SimpleNamespace(
|
|
endpoint_url="https://selected.example/v1",
|
|
model="selected-model",
|
|
headers={"Authorization": "Bearer selected"},
|
|
name="test",
|
|
history=[],
|
|
add_message=lambda message: None,
|
|
)
|
|
session_manager = SimpleNamespace(
|
|
get_session=lambda session_id: session,
|
|
save_sessions=lambda: None,
|
|
)
|
|
context = SimpleNamespace(
|
|
user="alice",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
preprocessed=SimpleNamespace(attachment_meta=[]),
|
|
auto_opened_docs=[],
|
|
rag_sources=[],
|
|
web_sources=[],
|
|
used_memories=[],
|
|
uploaded_files=[],
|
|
uprefs={},
|
|
was_compacted=False,
|
|
context_trimmed=False,
|
|
context_length=4096,
|
|
context_messages_before_trim=1,
|
|
context_messages_after_trim=1,
|
|
context_tokens_before_trim=10,
|
|
context_tokens_after_trim=10,
|
|
preset=SimpleNamespace(temperature=0.2, max_tokens=128, character_name=None),
|
|
)
|
|
|
|
async def fake_build_context(*args, **kwargs):
|
|
return context
|
|
|
|
async def fake_chat_stream(candidates, messages, **kwargs):
|
|
captured["chat"] = candidates
|
|
yield f'data: {json.dumps({"delta": "done"})}\n\n'
|
|
yield "data: [DONE]\n\n"
|
|
|
|
async def fake_agent_stream(endpoint_url, model, messages, **kwargs):
|
|
captured["agent"] = {
|
|
"primary": (endpoint_url, model, kwargs.get("headers")),
|
|
"fallbacks": kwargs.get("fallbacks"),
|
|
}
|
|
yield f'data: {json.dumps({"delta": "done"})}\n\n'
|
|
yield "data: [DONE]\n\n"
|
|
|
|
monkeypatch.setattr(chat_routes, "coerce_message_and_session", lambda *args, **kwargs: ("hello", "session-1"))
|
|
monkeypatch.setattr(chat_routes, "_verify_session_owner", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(chat_routes, "effective_user", lambda request: "alice")
|
|
monkeypatch.setattr(chat_routes, "_clear_orphaned_session_endpoint", lambda *args, **kwargs: False)
|
|
monkeypatch.setattr(chat_routes, "_recover_empty_session_model", lambda *args, **kwargs: False)
|
|
monkeypatch.setattr(chat_routes, "_enforce_chat_privileges", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(chat_routes, "resolve_session_auth", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(chat_routes, "get_session_mode", lambda session_id: "chat")
|
|
monkeypatch.setattr(chat_routes, "set_session_mode", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(chat_routes, "build_chat_context", fake_build_context)
|
|
monkeypatch.setattr(chat_routes, "SessionLocal", _EmptyDb)
|
|
monkeypatch.setattr(chat_routes, "_is_image_generation_session", lambda *args, **kwargs: False)
|
|
monkeypatch.setattr(chat_routes, "stream_llm_with_fallback", fake_chat_stream)
|
|
monkeypatch.setattr(chat_routes, "stream_agent_loop", fake_agent_stream)
|
|
monkeypatch.setattr(chat_routes, "save_assistant_response", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(chat_routes, "run_post_response_tasks", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(chat_routes, "estimate_tokens", lambda messages: 10)
|
|
monkeypatch.setattr(
|
|
endpoint_resolver,
|
|
"resolve_chat_fallback_candidates",
|
|
lambda owner=None: [("https://legacy.example/v1", "legacy-model", {})],
|
|
)
|
|
|
|
import src.settings as settings
|
|
|
|
monkeypatch.setattr(
|
|
settings,
|
|
"get_setting",
|
|
lambda key, default=None: default,
|
|
)
|
|
monkeypatch.setattr(
|
|
settings,
|
|
"get_user_setting",
|
|
lambda key, owner="", default=None: (
|
|
[{"endpoint_id": "legacy", "model": "legacy-model"}]
|
|
if key == "default_model_fallbacks"
|
|
else default
|
|
),
|
|
)
|
|
|
|
router = chat_routes.setup_chat_routes(
|
|
session_manager,
|
|
SimpleNamespace(),
|
|
SimpleNamespace(),
|
|
SimpleNamespace(),
|
|
SimpleNamespace(),
|
|
SimpleNamespace(),
|
|
)
|
|
return next(route.endpoint for route in router.routes if route.path == "/api/chat_stream")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("mode", ["chat", "agent"])
|
|
async def test_chat_stream_route_keeps_selected_model_strict_with_legacy_data(monkeypatch, mode):
|
|
captured = {}
|
|
endpoint = _chat_stream_endpoint(monkeypatch, mode, captured)
|
|
|
|
response = await endpoint(_RouteRequest(mode))
|
|
async for _ in response.body_iterator:
|
|
pass
|
|
|
|
selected = (
|
|
"https://selected.example/v1",
|
|
"selected-model",
|
|
{"Authorization": "Bearer selected"},
|
|
)
|
|
if mode == "chat":
|
|
assert captured == {"chat": [selected]}
|
|
else:
|
|
assert captured == {"agent": {"primary": selected, "fallbacks": []}}
|
|
|
|
|
|
def test_candidate_builder_appends_only_policy_authorized_fallbacks(monkeypatch):
|
|
"""Chat and Agent share the same candidate-building policy boundary."""
|
|
|
|
authorized = [("https://opt-in.example/v1", "opt-in-model", {})]
|
|
monkeypatch.setattr(
|
|
foreground_model_routing,
|
|
"resolve_foreground_fallback_candidates",
|
|
lambda owner=None: authorized,
|
|
)
|
|
|
|
assert build_foreground_model_candidates(
|
|
"https://selected.example/v1",
|
|
"selected-model",
|
|
{"Authorization": "Bearer selected"},
|
|
owner="alice",
|
|
) == [
|
|
("https://selected.example/v1", "selected-model", {"Authorization": "Bearer selected"}),
|
|
*authorized,
|
|
]
|
|
|
|
|
|
def test_strict_policy_builds_only_the_selected_chat_candidate():
|
|
candidates = build_foreground_model_candidates(
|
|
"https://selected.example/v1",
|
|
"selected-model",
|
|
{"Authorization": "Bearer selected"},
|
|
owner="alice",
|
|
)
|
|
|
|
assert candidates == [
|
|
("https://selected.example/v1", "selected-model", {"Authorization": "Bearer selected"})
|
|
]
|
|
|
|
|
|
def test_legacy_chat_resolver_is_disconnected():
|
|
assert endpoint_resolver.resolve_chat_fallback_candidates(owner="alice") == []
|
|
|
|
|
|
def test_utility_resolver_does_not_inherit_legacy_chat_fallbacks(monkeypatch):
|
|
seen_keys = []
|
|
|
|
def fake_resolve(setting_key, owner=None):
|
|
seen_keys.append((setting_key, owner))
|
|
return [("https://utility.example/v1", "utility-model", {})]
|
|
|
|
monkeypatch.setattr(endpoint_resolver, "_resolve_fallback_candidates", fake_resolve)
|
|
|
|
assert endpoint_resolver.resolve_utility_fallback_candidates(owner="alice") == [
|
|
("https://utility.example/v1", "utility-model", {})
|
|
]
|
|
assert seen_keys == [("utility_model_fallbacks", "alice")]
|
|
|
|
|
|
def test_multi_round_agent_uses_only_selected_model(monkeypatch):
|
|
"""Every Agent round receives only the selected foreground candidate."""
|
|
|
|
seen_candidates = []
|
|
round_number = 0
|
|
|
|
monkeypatch.setattr(agent_loop, "get_setting", lambda key, default=None: default)
|
|
monkeypatch.setattr(agent_loop, "get_mcp_manager", lambda: None)
|
|
monkeypatch.setattr(agent_loop, "estimate_tokens", lambda *args, **kwargs: 10)
|
|
|
|
async def fake_stream(candidates, messages, **kwargs):
|
|
nonlocal round_number
|
|
round_number += 1
|
|
seen_candidates.append([(url, model) for url, model, _headers in candidates])
|
|
if round_number == 1:
|
|
call = {"name": "bash", "arguments": json.dumps({"command": "printf ok"})}
|
|
yield f'data: {json.dumps({"type": "tool_calls", "calls": [call]})}\n\n'
|
|
else:
|
|
yield f'data: {json.dumps({"delta": "done"})}\n\n'
|
|
yield "data: [DONE]\n\n"
|
|
|
|
async def fake_execute(block, *args, **kwargs):
|
|
return "bash", {"output": "ok", "exit_code": 0}
|
|
|
|
monkeypatch.setattr(agent_loop, "stream_llm_with_fallback", fake_stream)
|
|
monkeypatch.setattr(agent_loop, "execute_tool_block", fake_execute)
|
|
|
|
fallbacks = resolve_foreground_fallback_candidates(owner="alice")
|
|
chunks = _collect(
|
|
agent_loop.stream_agent_loop(
|
|
"https://selected.example/v1",
|
|
"selected-model",
|
|
[{"role": "user", "content": "Run one tool and report back."}],
|
|
max_rounds=3,
|
|
relevant_tools={"bash"},
|
|
fallbacks=fallbacks,
|
|
_is_teacher_run=True,
|
|
)
|
|
)
|
|
|
|
assert seen_candidates == [
|
|
[("https://selected.example/v1", "selected-model")],
|
|
[("https://selected.example/v1", "selected-model")],
|
|
]
|
|
assert any('"delta": "done"' in chunk for chunk in chunks)
|