diff --git a/src/casecrawler/cli.py b/src/casecrawler/cli.py index 2732b9b..97b80fb 100644 --- a/src/casecrawler/cli.py +++ b/src/casecrawler/cli.py @@ -4,6 +4,7 @@ import json import subprocess import time +from uuid import uuid4 import click @@ -914,6 +915,92 @@ def generate_dataset( click.echo(f"Approved: {result['approved']}") +@cli.command("generate-blueprints") +@click.argument("request_text") +@click.option("--count", default=1, type=click.IntRange(1), help="Number of blueprints to generate") +@click.option("--dataset-id", default=None, help="Optional dataset id for persisted blueprints") +@click.option("--domain", "domains", multiple=True, help="Broad clinical domain or organ system") +@click.option("--setting", "settings", multiple=True, help="Care setting such as outpatient or ICU") +@click.option("--planner-provider", required=True, help="BYOK provider for cohort planning") +@click.option("--planner-model", required=True, help="Model id for cohort planning") +@click.option( + "--planner-temperature", + default=0.2, + type=float, + show_default=True, + help="Planner model temperature", +) +@click.option("--blueprint-provider", required=True, help="BYOK provider for blueprint generation") +@click.option("--blueprint-model", required=True, help="Model id for blueprint generation") +@click.option( + "--blueprint-temperature", + default=0.3, + type=float, + show_default=True, + help="Blueprint model temperature", +) +@click.option("--no-grounding", is_flag=True, help="Do not require grounding in the generated plan") +def generate_blueprints( + request_text: str, + count: int, + dataset_id: str | None, + domains: tuple[str, ...], + settings: tuple[str, ...], + planner_provider: str, + planner_model: str, + planner_temperature: float, + blueprint_provider: str, + blueprint_model: str, + blueprint_temperature: float, + no_grounding: bool, +) -> None: + """Generate model-planned clinical blueprints without topic recipes.""" + from casecrawler.generation import blueprint_pipeline as blueprint_pipeline_module + from casecrawler.models.blueprint import ( + BlueprintGenerationRequest, + GenerationRole, + GenerationRolePolicy, + ) + from casecrawler.storage.dataset_store import DatasetStore + + resolved_dataset_id = dataset_id or f"blueprint-ds-{uuid4()}" + req = BlueprintGenerationRequest( + request=request_text, + target_count=count, + domains=list(domains), + settings=list(settings), + required_grounding=not no_grounding, + role_policies=[ + GenerationRolePolicy( + role=GenerationRole.PLANNER, + provider=planner_provider, + model=planner_model, + temperature=planner_temperature, + ), + GenerationRolePolicy( + role=GenerationRole.BLUEPRINT_GENERATOR, + provider=blueprint_provider, + model=blueprint_model, + temperature=blueprint_temperature, + ), + ], + ) + try: + store = DatasetStore() + result = asyncio.run( + blueprint_pipeline_module.BlueprintPipeline().generate( + req, + dataset_id=resolved_dataset_id, + store=store, + ) + ) + except ValueError as exc: + raise click.ClickException(str(exc)) from exc + click.echo(f"Dataset: {result.dataset_id}") + click.echo(f"Plan: {result.plan.plan_id}") + click.echo(f"Blueprints: {result.generated_count}") + + @cli.command("generate-release-package") @click.argument("topic") @click.option("--output-dir", required=True, help="Output directory for the release package") diff --git a/tests/test_cli_synthetic.py b/tests/test_cli_synthetic.py index 9b23478..6291e5d 100644 --- a/tests/test_cli_synthetic.py +++ b/tests/test_cli_synthetic.py @@ -1,6 +1,54 @@ from click.testing import CliRunner from casecrawler.cli import cli +from casecrawler.generation import blueprint_pipeline as blueprint_pipeline_module +from casecrawler.generation.blueprint_pipeline import BlueprintPipelineResult +from casecrawler.models.blueprint import ( + BlueprintEvidence, + ClinicalBlueprint, + CohortArchetype, + CohortPlan, + GenerationRole, +) +from casecrawler.models.synthetic import Modality +from casecrawler.storage.dataset_store import DatasetStore + + +def _blueprint_cli_result(dataset_id: str) -> BlueprintPipelineResult: + archetype = CohortArchetype( + name="anticoagulation decision", + organ_system="cardiovascular", + setting="outpatient", + target_count=1, + required_modalities=[Modality.STRUCTURED_EHR, Modality.CLINICAL_TEXT], + ) + plan = CohortPlan( + plan_id="plan-1", + request="Generate anticoagulation decision cases.", + target_count=1, + archetypes=[archetype], + created_by=GenerationRole.PLANNER, + ) + blueprint = ClinicalBlueprint( + blueprint_id="bp-1", + dataset_id=dataset_id, + 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"], + } + ], + evidence=BlueprintEvidence( + supported_claims=["AF anticoagulation requires renal-dose review."], + ), + ) + return BlueprintPipelineResult(dataset_id=dataset_id, plan=plan, blueprints=[blueprint]) def test_generate_dataset_command_smoke(tmp_path, monkeypatch): @@ -21,3 +69,54 @@ def test_generate_dataset_invalid_complexity_fails(): assert result.exit_code != 0 assert "Invalid value for '--complexity'" in result.output + + +def test_generate_blueprints_command_uses_model_driven_request(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + captured = [] + + class FakeBlueprintPipeline: + async def generate(self, req, *, dataset_id, store): + captured.append((req, dataset_id, store)) + store.save_cohort_plan(_blueprint_cli_result(dataset_id).plan) + for blueprint in _blueprint_cli_result(dataset_id).blueprints: + store.save_blueprint(blueprint) + return _blueprint_cli_result(dataset_id) + + monkeypatch.setattr( + blueprint_pipeline_module, + "BlueprintPipeline", + FakeBlueprintPipeline, + ) + runner = CliRunner() + + result = runner.invoke( + cli, + [ + "generate-blueprints", + "Generate anticoagulation decision cases.", + "--count", + "1", + "--planner-provider", + "openrouter", + "--planner-model", + "planner-model", + "--blueprint-provider", + "openrouter", + "--blueprint-model", + "blueprint-model", + ], + ) + + assert result.exit_code == 0 + assert "Dataset: blueprint-ds-" in result.output + assert "Plan: plan-1" in result.output + assert "Blueprints: 1" in result.output + assert captured[0][0].request == "Generate anticoagulation decision cases." + assert captured[0][0].target_count == 1 + assert captured[0][0].policy_for(GenerationRole.PLANNER).model == "planner-model" + assert ( + captured[0][0].policy_for(GenerationRole.BLUEPRINT_GENERATOR).model + == "blueprint-model" + ) + assert isinstance(captured[0][2], DatasetStore)