120 lines
5.8 KiB
Python
120 lines
5.8 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_harness_requires_two_distinct_validators(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):
|
|
with self.assertRaises(api.HTTPException) as context:
|
|
api._harness_models(body)
|
|
self.assertEqual(context.exception.status_code, 400)
|
|
|
|
def test_chat_route_returns_one_compiled_answer(self):
|
|
body = api.ChatRequest(
|
|
primary_model="primary",
|
|
validator_models=["validator-a", "validator-b"],
|
|
harness=True,
|
|
message="Answer this once and validate it.",
|
|
)
|
|
persisted_messages = []
|
|
|
|
def fake_stream(payload, request_id, cancel_event=None, parent_id=None):
|
|
prompt = payload["messages"][-1]["content"]
|
|
if payload["model"] == "primary" and "VALIDATION REPORTS:" not in prompt:
|
|
content = "draft"
|
|
elif payload["model"].startswith("validator"):
|
|
content = "No material issues found."
|
|
else:
|
|
content = "single compiled 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, "_harness_models", return_value=("primary", ["validator-a", "validator-b"], True)), patch.object(
|
|
api, "_require_installed_model", side_effect=lambda name: name
|
|
), patch.object(api, "_ensure_conversation"), patch.object(api, "_persist_message", side_effect=lambda *args, **kwargs: persisted_messages.append(args)), patch.object(
|
|
api, "_persist_metric", side_effect=fake_metric
|
|
), patch.object(api, "_stream_chat_request", side_effect=fake_stream), patch.object(
|
|
api, "_runtime_snapshot", return_value={}
|
|
):
|
|
response = api.chat(body)
|
|
|
|
self.assertEqual(response["mode"], "harness")
|
|
self.assertEqual(response["message"]["content"], "single compiled answer")
|
|
self.assertEqual(response["primary_model"], "primary")
|
|
self.assertEqual(response["validator_models"], ["validator-a", "validator-b"])
|
|
assistant_messages = [args for args in persisted_messages if len(args) >= 3 and args[2] == "assistant"]
|
|
self.assertEqual(len(assistant_messages), 1)
|
|
self.assertEqual(assistant_messages[0][3], "single compiled answer")
|
|
self.assertEqual(len(response["validation_reports"]), 2)
|
|
|
|
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):
|
|
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))
|
|
with api._chat_requests_lock:
|
|
child_states = [state for key, state in api._chat_requests.items() if key != "root-request"]
|
|
self.assertGreaterEqual(len(child_states), 4)
|
|
self.assertTrue(all(state.get("parent_id") == "root-request" for state in child_states[-4:]))
|
|
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()
|