146 lines
6.6 KiB
Python
146 lines
6.6 KiB
Python
import threading
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from dashboard import plugin_api as api
|
|
|
|
|
|
class ValidationHarnessTests(unittest.TestCase):
|
|
def test_legacy_multiple_models_map_to_primary_and_validators(self):
|
|
body = api.ChatRequest(models=["primary", "validator-a", "validator-b"])
|
|
with patch.object(api, "_require_installed_model", side_effect=lambda name: name):
|
|
primary, validators, harness = api._harness_models(body)
|
|
self.assertEqual(primary, "primary")
|
|
self.assertEqual(validators, ["validator-a", "validator-b"])
|
|
self.assertTrue(harness)
|
|
|
|
def test_sqlite_is_default_even_when_postgres_is_installed(self):
|
|
with patch.object(api, "_storage_config", return_value={"backend": "sqlite"}), patch.object(
|
|
api, "_postgres_configured", return_value=True
|
|
):
|
|
self.assertEqual(api._storage_backend(), "sqlite")
|
|
self.assertFalse(api._postgres_enabled())
|
|
|
|
def test_harness_requires_one_distinct_validator(self):
|
|
body = api.ChatRequest(primary_model="primary", validator_models=[], harness=True)
|
|
with patch.object(api, "_require_installed_model", side_effect=lambda name: name):
|
|
with self.assertRaises(api.HTTPException) as context:
|
|
api._harness_models(body)
|
|
self.assertEqual(context.exception.status_code, 400)
|
|
|
|
def test_harness_accepts_one_validator(self):
|
|
body = api.ChatRequest(primary_model="primary", validator_models=["validator-a"], harness=True)
|
|
with patch.object(api, "_require_installed_model", side_effect=lambda name: name):
|
|
primary, validators, harness = api._harness_models(body)
|
|
self.assertEqual(primary, "primary")
|
|
self.assertEqual(validators, ["validator-a"])
|
|
self.assertTrue(harness)
|
|
|
|
def test_chat_route_queues_server_owned_job(self):
|
|
body = api.ChatRequest(primary_model="primary", message="Queue this request")
|
|
fake_job = {
|
|
"request_id": "request",
|
|
"conversation_id": "conversation",
|
|
"status": "queued",
|
|
"mode": "direct",
|
|
"primary_model": "primary",
|
|
"validator_models_json": "[]",
|
|
"result_json": "{}",
|
|
"error": "",
|
|
"updated_at": 1.0,
|
|
}
|
|
with patch.object(api, "_harness_models", return_value=("primary", [], False)), patch.object(
|
|
api, "_get_chat_job", side_effect=[None, fake_job]
|
|
), patch.object(api, "_ensure_conversation"), patch.object(api, "_persist_message"), patch.object(
|
|
api, "_chat_state"
|
|
), patch.object(api, "_create_chat_job"), patch.object(api, "_submit_chat_job"):
|
|
response = api.chat(body)
|
|
self.assertEqual(response["status"], "queued")
|
|
self.assertFalse(response["done"])
|
|
self.assertEqual(response["mode"], "direct")
|
|
|
|
def test_job_status_includes_durable_heartbeat_fields(self):
|
|
job = {
|
|
"request_id": "request",
|
|
"conversation_id": "conversation",
|
|
"status": "running",
|
|
"mode": "direct",
|
|
"primary_model": "primary",
|
|
"validator_models_json": "[]",
|
|
"result_json": "{}",
|
|
"error": "",
|
|
"attempt": 2,
|
|
"started_at": 100.0,
|
|
"heartbeat_at": 110.0,
|
|
"finished_at": None,
|
|
"updated_at": 110.0,
|
|
}
|
|
with patch.object(api.time, "time", return_value=112.5):
|
|
response = api._job_status_response(job)
|
|
self.assertFalse(response["done"])
|
|
self.assertEqual(response["attempt"], 2)
|
|
self.assertEqual(response["heartbeat_at"], 110.0)
|
|
self.assertEqual(response["heartbeat_age"], 2.5)
|
|
|
|
def test_active_jobs_route_returns_server_owned_jobs(self):
|
|
active = [{"request_id": "request", "status": "running"}]
|
|
with patch.object(api, "_list_chat_jobs", return_value=active) as listed:
|
|
response = api.chat_jobs(active=True, conversation_id="conversation", limit=20)
|
|
self.assertEqual(response, {"jobs": active})
|
|
listed.assert_called_once_with(active_only=True, conversation_id="conversation", limit=20)
|
|
|
|
def test_primary_draft_validators_and_primary_compilation_produce_one_answer(self):
|
|
body = api.ChatRequest(
|
|
primary_model="primary",
|
|
validator_models=["validator-a", "validator-b"],
|
|
harness=True,
|
|
message="What is the verified answer?",
|
|
)
|
|
calls = []
|
|
|
|
def fake_stream(payload, request_id, cancel_event=None, parent_id=None):
|
|
model = payload["model"]
|
|
prompt = payload["messages"][-1]["content"]
|
|
calls.append((model, prompt, parent_id))
|
|
if model == "primary" and "VALIDATION REPORTS:" not in prompt:
|
|
content = "primary draft"
|
|
elif model.startswith("validator"):
|
|
content = f"{model} found no material issue"
|
|
else:
|
|
content = "one compiled final answer"
|
|
return {"message": {"role": "assistant", "content": content}, "done": True}
|
|
|
|
def fake_metric(conversation_id, request_id, model, state, status=None, error=""):
|
|
return {"request_id": request_id, "model": model, "status": status}
|
|
|
|
with patch.object(api, "_require_installed_model", side_effect=lambda name: name), patch.object(
|
|
api, "_stream_chat_request", side_effect=fake_stream
|
|
), patch.object(api, "_persist_metric", side_effect=fake_metric), patch.object(
|
|
api, "_persist_chat_stage"
|
|
), patch.object(api, "_persist_chat_event"):
|
|
final, metrics, reports = api._run_validation_harness(
|
|
body,
|
|
"root-request",
|
|
"conversation",
|
|
"primary",
|
|
["validator-a", "validator-b"],
|
|
threading.Event(),
|
|
[],
|
|
)
|
|
|
|
self.assertEqual(final, "one compiled final answer")
|
|
self.assertEqual([item["model"] for item in reports], ["validator-a", "validator-b"])
|
|
self.assertEqual(len(metrics), 4)
|
|
self.assertEqual(len(calls), 4)
|
|
self.assertEqual(calls[0][0], "primary")
|
|
self.assertTrue(all(call[2] == "root-request" for call in calls))
|
|
compiler_prompt = calls[-1][1]
|
|
self.assertIn("PRIMARY DRAFT:", compiler_prompt)
|
|
self.assertIn("VALIDATOR 1", compiler_prompt)
|
|
self.assertIn("VALIDATOR 2", compiler_prompt)
|
|
self.assertNotIn("validator-a found no material issue\n\nvalidator-b found no material issue", final)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|