diff --git a/src/casecrawler/generation/blueprint_judge.py b/src/casecrawler/generation/blueprint_judge.py new file mode 100644 index 0000000..f5ce680 --- /dev/null +++ b/src/casecrawler/generation/blueprint_judge.py @@ -0,0 +1,220 @@ +from __future__ import annotations + +import hashlib +import json +import logging +from collections.abc import Callable +from uuid import uuid4 + +from casecrawler.llm.base import BaseLLMProvider +from casecrawler.llm.factory import get_provider +from casecrawler.models.blueprint import ( + BlueprintGenerationRequest, + ClinicalBlueprint, + GenerationAttempt, + GenerationAttemptStatus, + GenerationRole, + GenerationRolePolicy, + JudgeReport, +) +from casecrawler.storage.dataset_store import DatasetStore + + +ProviderFactory = Callable[[str, str], BaseLLMProvider] +logger = logging.getLogger(__name__) + + +class BlueprintJudge: + def __init__(self, provider_factory: ProviderFactory = get_provider) -> None: + self._provider_factory = provider_factory + + async def evaluate( + self, + request: BlueprintGenerationRequest, + blueprint: ClinicalBlueprint, + *, + store: DatasetStore | None = None, + ) -> JudgeReport: + policy = request.policy_for(GenerationRole.JUDGE) + if policy is None: + if store is not None: + self._save_failed_attempt_best_effort( + store, + self._missing_policy_attempt(request=request, blueprint=blueprint), + ) + raise ValueError("A judge role policy is required for blueprint judging.") + + provider = self._provider_factory(policy.provider, policy.model) + prompt = self._build_prompt(request, blueprint) + prompt_hash = self._prompt_hash(prompt, policy) + + try: + result = await provider.generate_structured( + prompt, + JudgeReport, + system=_JUDGE_SYSTEM_PROMPT, + temperature=policy.temperature, + ) + report = self._canonicalize_report( + JudgeReport.model_validate(result.data), + blueprint=blueprint, + ) + except Exception as err: + if store is not None: + self._save_failed_attempt_best_effort( + store, + self._attempt( + blueprint=blueprint, + policy=policy, + status=GenerationAttemptStatus.FAILED, + prompt_hash=prompt_hash, + errors=[str(err)], + ), + ) + raise + + if store is not None: + store.save_judge_report_with_attempt( + report, + self._attempt( + blueprint=blueprint, + policy=policy, + status=GenerationAttemptStatus.SUCCEEDED, + prompt_hash=prompt_hash, + input_tokens=result.input_tokens, + output_tokens=result.output_tokens, + ) + ) + return report + + def _canonicalize_report( + self, + raw_report: JudgeReport, + *, + blueprint: ClinicalBlueprint, + ) -> JudgeReport: + return JudgeReport.model_validate( + { + **raw_report.model_dump(), + "report_id": f"judge-{uuid4()}", + "dataset_id": blueprint.dataset_id, + "artifact_id": blueprint.blueprint_id, + "role": GenerationRole.JUDGE, + } + ) + + def _build_prompt( + self, + request: BlueprintGenerationRequest, + blueprint: ClinicalBlueprint, + ) -> str: + blueprint_json = json.dumps( + blueprint.model_dump(mode="json"), + sort_keys=True, + separators=(",", ":"), + ) + return "\n".join( + [ + "Evaluate this clinical blueprint before synthetic case generation.", + f"User request: {request.request}", + f"Blueprint id: {blueprint.blueprint_id}", + f"Dataset id: {blueprint.dataset_id}", + "Judge for clinical plausibility, internal consistency, grounding, " + "safety, and usefulness as a source-of-truth case plan.", + ( + "Return a JudgeReport with a calibrated score, pass/fail decision, " + "rubric name, and concrete findings for each material issue." + ), + f"Blueprint JSON: {blueprint_json}", + ] + ) + + def _prompt_hash( + self, + prompt: str, + policy: GenerationRolePolicy, + ) -> str: + payload = { + "model": policy.model, + "provider": policy.provider, + "schema": JudgeReport.__name__, + "system": _JUDGE_SYSTEM_PROMPT, + "temperature": policy.temperature, + "user": prompt, + } + serialized = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest() + + def _attempt( + self, + *, + blueprint: ClinicalBlueprint, + policy: GenerationRolePolicy, + status: GenerationAttemptStatus, + prompt_hash: str, + input_tokens: int = 0, + output_tokens: int = 0, + errors: list[str] | None = None, + ) -> GenerationAttempt: + return GenerationAttempt( + attempt_id=f"attempt-{uuid4()}", + dataset_id=blueprint.dataset_id, + role=GenerationRole.JUDGE, + status=status, + provider=policy.provider, + model=policy.model, + prompt_hash=prompt_hash, + input_tokens=input_tokens, + output_tokens=output_tokens, + errors=errors or [], + artifact_id=blueprint.blueprint_id, + ) + + def _missing_policy_attempt( + self, + *, + request: BlueprintGenerationRequest, + blueprint: ClinicalBlueprint, + ) -> GenerationAttempt: + prompt_hash_payload = { + "artifact_id": blueprint.blueprint_id, + "reason": "missing_policy", + "request": request.request, + "role": GenerationRole.JUDGE.value, + } + prompt_hash = hashlib.sha256( + json.dumps( + prompt_hash_payload, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest() + return GenerationAttempt( + attempt_id=f"attempt-{uuid4()}", + dataset_id=blueprint.dataset_id, + role=GenerationRole.JUDGE, + status=GenerationAttemptStatus.FAILED, + provider="unconfigured", + model="unconfigured", + prompt_hash=prompt_hash, + errors=["missing judge role policy"], + artifact_id=blueprint.blueprint_id, + metadata={"reason": "missing_policy"}, + ) + + def _save_failed_attempt_best_effort( + self, + store: DatasetStore, + attempt: GenerationAttempt, + ) -> None: + try: + store.save_generation_attempt(attempt) + except Exception: + logger.exception("Failed to persist judge failure audit.") + + +_JUDGE_SYSTEM_PROMPT = ( + "You are an independent clinical QA judge for medical AI synthetic-data " + "blueprints. Evaluate only the supplied structured blueprint. Do not add " + "new patient facts, final training examples, or patient-facing advice." +) diff --git a/src/casecrawler/storage/dataset_store.py b/src/casecrawler/storage/dataset_store.py index 8b97c64..893f12c 100644 --- a/src/casecrawler/storage/dataset_store.py +++ b/src/casecrawler/storage/dataset_store.py @@ -396,6 +396,45 @@ def save_judge_report(self, report: JudgeReport) -> None: ) self._conn.commit() + def save_judge_report_with_attempt( + self, + report: JudgeReport, + attempt: GenerationAttempt, + ) -> None: + with self._write_lock: + try: + self._conn.execute( + """INSERT OR REPLACE INTO judge_reports + (report_id, dataset_id, artifact_id, role, passed, score, report_json) + VALUES (?, ?, ?, ?, ?, ?, ?)""", + ( + report.report_id, + report.dataset_id, + report.artifact_id, + report.role.value, + int(report.passed), + report.score, + report.model_dump_json(), + ), + ) + self._conn.execute( + """INSERT OR REPLACE INTO generation_attempts + (attempt_id, dataset_id, role, status, artifact_id, attempt_json) + VALUES (?, ?, ?, ?, ?, ?)""", + ( + attempt.attempt_id, + attempt.dataset_id, + attempt.role.value, + attempt.status.value, + attempt.artifact_id, + attempt.model_dump_json(), + ), + ) + self._conn.commit() + except Exception: + self._conn.rollback() + raise + def get_judge_report(self, report_id: str) -> JudgeReport | None: row = self._conn.execute( "SELECT report_json FROM judge_reports WHERE report_id = ?", diff --git a/tests/test_blueprint_judge.py b/tests/test_blueprint_judge.py new file mode 100644 index 0000000..d356ebf --- /dev/null +++ b/tests/test_blueprint_judge.py @@ -0,0 +1,208 @@ +import pytest + +from casecrawler.llm.base import StructuredGenerationResult +from casecrawler.models.blueprint import ( + BlueprintEvidence, + BlueprintGenerationRequest, + ClinicalBlueprint, + GenerationAttemptStatus, + GenerationRole, + GenerationRolePolicy, + JudgeReport, +) +from casecrawler.storage.dataset_store import DatasetStore + + +class FakeJudgeProvider: + def __init__(self, report: JudgeReport) -> None: + self.report = report + self.calls = [] + + async def generate_structured(self, prompt, schema, system="", **kwargs): + self.calls.append( + { + "prompt": prompt, + "schema": schema, + "system": system, + "kwargs": kwargs, + } + ) + return StructuredGenerationResult( + data=self.report, + input_tokens=80, + output_tokens=30, + model="judge-model", + ) + + +class FailingJudgeProvider: + async def generate_structured(self, prompt, schema, system="", **kwargs): + raise RuntimeError("judge boom") + + +class BrokenAttemptStore: + def save_generation_attempt(self, attempt): + raise RuntimeError("storage boom") + + +def _request() -> BlueprintGenerationRequest: + return BlueprintGenerationRequest( + request="Judge generated cardiology blueprints before release.", + target_count=1, + role_policies=[ + GenerationRolePolicy( + role=GenerationRole.JUDGE, + provider="openai", + model="gpt-4.1-mini", + temperature=0.0, + ) + ], + ) + + +def _blueprint() -> ClinicalBlueprint: + return ClinicalBlueprint( + blueprint_id="bp-1", + dataset_id="ds-1", + cohort_plan_id="plan-1", + archetype_name="anticoagulation decision", + organ_system="cardiovascular", + setting="outpatient", + patient={"age": 72, "sex": "female"}, + chief_concern="Atrial fibrillation anticoagulation follow-up.", + diagnoses=[ + { + "name": "atrial fibrillation", + "supporting_findings": ["ECG confirms AF"], + } + ], + clinical_reasoning_targets=["Review renal dosing and bleeding risk."], + safety_constraints=["Review bleeding risk before anticoagulation."], + evidence=BlueprintEvidence( + supported_claims=["AF anticoagulation requires renal-dose review."], + citations=[{"source": "dailymed", "claim": "renal-dose review"}], + ), + ) + + +def _raw_report() -> JudgeReport: + return JudgeReport( + report_id="model-controlled-id", + dataset_id="wrong-dataset", + artifact_id="wrong-artifact", + role=GenerationRole.REPAIR, + score=0.91, + passed=True, + rubric="blueprint_plausibility", + findings=[{"criterion": "diagnostic_support", "passed": True}], + ) + + +@pytest.mark.asyncio +async def test_blueprint_judge_uses_role_policy_and_persists_report(tmp_path): + from casecrawler.generation.blueprint_judge import BlueprintJudge + + store = DatasetStore(db_path=str(tmp_path / "datasets.db")) + provider = FakeJudgeProvider(_raw_report()) + judge = BlueprintJudge(provider_factory=lambda provider_name, model: provider) + + report = await judge.evaluate(_request(), _blueprint(), store=store) + + assert report.report_id.startswith("judge-") + assert report.report_id != "model-controlled-id" + assert report.dataset_id == "ds-1" + assert report.artifact_id == "bp-1" + assert report.role == GenerationRole.JUDGE + assert report.score == 0.91 + assert report.passed is True + assert report.rubric == "blueprint_plausibility" + assert store.list_judge_reports(artifact_id="bp-1") == [report] + + attempts = store.list_generation_attempts(dataset_id="ds-1") + assert len(attempts) == 1 + assert attempts[0].role == GenerationRole.JUDGE + assert attempts[0].status == GenerationAttemptStatus.SUCCEEDED + assert attempts[0].provider == "openai" + assert attempts[0].model == "gpt-4.1-mini" + assert attempts[0].artifact_id == "bp-1" + assert attempts[0].total_tokens == 110 + assert attempts[0].prompt_hash + assert provider.calls[0]["schema"] is JudgeReport + assert provider.calls[0]["kwargs"]["temperature"] == 0.0 + assert "bp-1" in provider.calls[0]["prompt"] + + +@pytest.mark.asyncio +async def test_blueprint_judge_requires_judge_policy(): + from casecrawler.generation.blueprint_judge import BlueprintJudge + + request = BlueprintGenerationRequest( + request="Judge generated cardiology blueprints before release.", + target_count=1, + ) + + with pytest.raises(ValueError, match="judge role policy"): + await BlueprintJudge( + provider_factory=lambda provider_name, model: None + ).evaluate(request, _blueprint()) + + +@pytest.mark.asyncio +async def test_blueprint_judge_audits_missing_judge_policy(tmp_path): + from casecrawler.generation.blueprint_judge import BlueprintJudge + + store = DatasetStore(db_path=str(tmp_path / "datasets.db")) + request = BlueprintGenerationRequest( + request="Judge generated cardiology blueprints before release.", + target_count=1, + ) + + with pytest.raises(ValueError, match="judge role policy"): + await BlueprintJudge( + provider_factory=lambda provider_name, model: None + ).evaluate(request, _blueprint(), store=store) + + attempts = store.list_generation_attempts(dataset_id="ds-1") + assert len(attempts) == 1 + assert attempts[0].role == GenerationRole.JUDGE + assert attempts[0].status == GenerationAttemptStatus.FAILED + assert attempts[0].artifact_id == "bp-1" + assert attempts[0].provider == "unconfigured" + assert attempts[0].model == "unconfigured" + assert attempts[0].errors == ["missing judge role policy"] + assert attempts[0].metadata["reason"] == "missing_policy" + assert attempts[0].prompt_hash + + +@pytest.mark.asyncio +async def test_blueprint_judge_persists_failed_attempt(tmp_path): + from casecrawler.generation.blueprint_judge import BlueprintJudge + + store = DatasetStore(db_path=str(tmp_path / "datasets.db")) + judge = BlueprintJudge( + provider_factory=lambda provider_name, model: FailingJudgeProvider() + ) + + with pytest.raises(RuntimeError, match="judge boom"): + await judge.evaluate(_request(), _blueprint(), store=store) + + attempts = store.list_generation_attempts(dataset_id="ds-1") + assert len(attempts) == 1 + assert attempts[0].role == GenerationRole.JUDGE + assert attempts[0].status == GenerationAttemptStatus.FAILED + assert attempts[0].artifact_id == "bp-1" + assert attempts[0].errors == ["judge boom"] + assert attempts[0].total_tokens == 0 + assert attempts[0].prompt_hash + + +@pytest.mark.asyncio +async def test_blueprint_judge_preserves_provider_error_when_failed_audit_fails(): + from casecrawler.generation.blueprint_judge import BlueprintJudge + + judge = BlueprintJudge( + provider_factory=lambda provider_name, model: FailingJudgeProvider() + ) + + with pytest.raises(RuntimeError, match="judge boom"): + await judge.evaluate(_request(), _blueprint(), store=BrokenAttemptStore()) diff --git a/tests/test_blueprint_storage.py b/tests/test_blueprint_storage.py index 44c7528..4b7f5c4 100644 --- a/tests/test_blueprint_storage.py +++ b/tests/test_blueprint_storage.py @@ -146,6 +146,37 @@ def test_dataset_store_tracks_attempts_and_judge_reports_by_artifact(tmp_path): assert store.get_blueprint_validation_report("bp-1") == validation +def test_dataset_store_saves_judge_report_with_attempt(tmp_path): + store = DatasetStore(db_path=str(tmp_path / "datasets.db")) + attempt = GenerationAttempt( + attempt_id="attempt-1", + dataset_id="ds-1", + role=GenerationRole.JUDGE, + status=GenerationAttemptStatus.SUCCEEDED, + provider="openai", + model="gpt-4.1-mini", + prompt_hash="abc123", + input_tokens=100, + output_tokens=75, + artifact_id="bp-1", + ) + judge_report = JudgeReport( + report_id="judge-1", + dataset_id="ds-1", + artifact_id="bp-1", + role=GenerationRole.JUDGE, + score=0.92, + passed=True, + rubric="blueprint_plausibility", + findings=[{"criterion": "diagnostic_support", "passed": True}], + ) + + store.save_judge_report_with_attempt(judge_report, attempt) + + assert store.get_judge_report("judge-1") == judge_report + assert store.get_generation_attempt("attempt-1") == attempt + + def test_dataset_manifest_includes_blueprint_persistence_counts(tmp_path): store = DatasetStore(db_path=str(tmp_path / "datasets.db")) store.save_record(_record())