diff --git a/README.md b/README.md index 3c5a3f2d..6e67eddc 100644 --- a/README.md +++ b/README.md @@ -213,6 +213,8 @@ Switch datasets or backends anytime: ```bash m4 use mimic-iv # Switch to full MIMIC-IV m4 backend bigquery # Switch to BigQuery (or duckdb) +m4 capabilities # Show available interfaces, datasets, tools, and policies +m4 doctor # Diagnose local, BigQuery, and MCP setup m4 status # Show active dataset and backend m4 status --all # List all available datasets m4 status --derived # Show per-table derived materialization status @@ -223,6 +225,9 @@ commands that do not mutate active configuration unless explicitly documented: ```bash m4 agent-env --dataset mimic-iv --backend duckdb --json +m4 capabilities --json +m4 doctor --json +m4 download mimic-iv --json m4 list-datasets --json --no-interactive m4 schema --dataset mimic-iv --backend duckdb --json --no-interactive m4 describe-table mimiciv_hosp.patients --dataset mimic-iv --json --no-interactive @@ -234,6 +239,17 @@ Machine-facing status and backend metadata hide local filesystem paths by default. Use `--paths` or `M4_PATH_DISCLOSURE=1` only when the caller is allowed to see raw local paths. +Long-running dataset setup can emit newline-delimited JSON progress events: + +```bash +m4 init mimic-iv --json --events ndjson --no-interactive --download \ + --physionet-credentials-file /path/to/physionet-credentials.json +``` + +When `--events ndjson` is used, stdout is an NDJSON stream instead of a single +JSON object. The final result is emitted as `operation_completed.result`; setup +failures are emitted as `operation_failed.error`. + **Derived concept tables** (MIMIC-IV only): ```bash m4 init-derived mimic-iv # Materialize ~63 derived tables (SOFA, sepsis3, KDIGO, etc.) @@ -247,21 +263,47 @@ After running `m4 init mimic-iv`, you are prompted whether to materialize derive 1. **Get PhysioNet credentials:** Complete the [credentialing process](https://physionet.org/settings/credentialing/) and sign the data use agreement for the dataset. -2. **Download the data:** +2. **Download the data with M4:** + ```bash + cat > physionet-credentials.json <<'JSON' + { + "username": "YOUR_USERNAME", + "password": "YOUR_PASSWORD" + } + JSON + + m4 init mimic-iv --download --physionet-credentials-file physionet-credentials.json + ``` + + Do not pass PhysioNet passwords as command-line flags. Use a scoped + credentials file with restrictive permissions and delete it after setup. + + M4 implements the same recursive, resumable pattern PhysioNet documents for + `wget -r -N -c -np` against `/files/...` dataset URLs, while preserving the + expected raw layout under `m4_data/raw_files//`. + + You can still download manually if needed: + ```bash + m4 download mimic-iv + ``` + + For credentialed datasets, `m4 download` validates the expected local layout + and prints a dataset-specific resumable `wget` command. + ```bash # For MIMIC-IV - wget -r -N -c -np --cut-dirs=2 -nH --user YOUR_USERNAME --ask-password \ + wget -r -N -c -np --cut-dirs=3 -nH --user YOUR_USERNAME --ask-password \ https://physionet.org/files/mimiciv/3.1/ \ -P m4_data/raw_files/mimic-iv # For eICU - wget -r -N -c -np --cut-dirs=2 -nH --user YOUR_USERNAME --ask-password \ + wget -r -N -c -np --cut-dirs=3 -nH --user YOUR_USERNAME --ask-password \ https://physionet.org/files/eicu-crd/2.0/ \ -P m4_data/raw_files/eicu ``` - The `--cut-dirs=2 -nH` flags ensure CSV files land directly in `m4_data/raw_files/mimic-iv/` rather than a nested `physionet.org/files/...` structure. + The `--cut-dirs=3 -nH` flags remove the PhysioNet `files///` prefix so CSV files land under `m4_data/raw_files//` with only dataset-internal folders preserved. -3. **Initialize:** +3. **Initialize after a manual download:** ```bash m4 init mimic-iv # or: m4 init eicu ``` diff --git a/docs/BIGQUERY.md b/docs/BIGQUERY.md index 1d2d4d52..e04a4f40 100644 --- a/docs/BIGQUERY.md +++ b/docs/BIGQUERY.md @@ -2,6 +2,10 @@ Use Google Cloud BigQuery to access full clinical datasets without downloading files locally. +For local DuckDB workflows, use `m4 download DATASET` followed by +`m4 init DATASET`. For BigQuery workflows, skip local download entirely and +configure credentials/project billing as described below. + ## Prerequisites 1. **Google Cloud account** with BigQuery access @@ -38,6 +42,13 @@ m4 config --backend bigquery --project-id YOUR_PROJECT_ID Replace `YOUR_PROJECT_ID` with your own billing project for BigQuery usage, not the PhysioNet dataset project. The variable is mandatory to ensure billing is correctly attributed. +You can also emit agent-ready environment guidance without changing config: + +```bash +m4 setup-agent --backend bigquery --project-id YOUR_PROJECT_ID --format dotenv +m4 doctor --json +``` + ### 4. Set the dataset ```bash @@ -100,7 +111,15 @@ BigQuery charges based on data scanned. Tips to minimize costs: **"Project not found" error:** - Check the project ID is correct - Ensure BigQuery API is enabled in your project +- Confirm `M4_PROJECT_ID` or `m4 config --project-id` refers to your billing project **Slow queries:** - BigQuery has network latency; consider local DuckDB for development - Use smaller `LIMIT` values while exploring + +**Local download layout problems:** +- Run `m4 download mimic-iv` or `m4 download eicu` to print dataset-specific + recovery guidance and a resumable `wget` command. +- If files landed under `physionet.org/files/...`, move the dataset contents up + to `m4_data/raw_files/DATASET` or rerun `wget` with the generated + `--cut-dirs` and `-nH` flags. diff --git a/docs/CUSTOM_DATASETS.md b/docs/CUSTOM_DATASETS.md index d11c5a5b..0dafbf44 100644 --- a/docs/CUSTOM_DATASETS.md +++ b/docs/CUSTOM_DATASETS.md @@ -147,10 +147,13 @@ For datasets requiring PhysioNet credentials (most full datasets): 1. Get credentialed access on PhysioNet 2. Download manually using wget: ```bash - wget -r -N -c -np --user YOUR_USERNAME --ask-password \ + wget -r -N -c -np --cut-dirs=3 -nH --user YOUR_USERNAME --ask-password \ https://physionet.org/files/dataset-name/version/ \ -P m4_data/raw_files/dataset-name ``` + Set `--cut-dirs` to the number of path components in the listing URL + (`files/dataset-name/version` is 3) so files land in the top-level raw + layout expected by `m4 init`. 3. Initialize: ```bash m4 init dataset-name diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md index df995c88..d327dd21 100644 --- a/docs/DEVELOPMENT.md +++ b/docs/DEVELOPMENT.md @@ -67,8 +67,12 @@ non-zero. { "name": "mimic-iv", "active": true, + "raw_present": true, "parquet_present": true, "db_present": true, + "requires_authentication": true, + "download_available": true, + "setup_state": "ready", "bigquery_available": true, "row_count": 431231, "parquet_size_gb": 8.5, @@ -86,7 +90,7 @@ non-zero. Raw local paths are hidden by default in machine-facing output. Pass `--paths` or set `M4_PATH_DISCLOSURE=1` to include path fields such as -`parquet_root` and `db_path`. +`raw_root`, `parquet_root`, and `db_path`. Dataset `warnings` is a list of stable warning codes. Currently documented status warnings: @@ -125,7 +129,11 @@ Command errors use the same envelope with `ok: false`: Stable command error codes are `dataset_not_found`, `backend_incompatible`, `invalid_backend`, `invalid_option`, `project_id_required`, and -`dataset_incompatible`. +`dataset_incompatible`. Dataset setup can also return `missing_credentials`, +`physionet_auth_failed`, `physionet_access_forbidden`, +`download_network_failed`, `download_filesystem_failed`, +`download_interrupted`, `raw_files_missing`, `conversion_failed`, +`duckdb_init_failed`, and `verification_failed`. `m4 init DATASET --json` uses the same result/error envelope and runs non-interactively. Human prompts, progress panels, and download output are @@ -142,17 +150,42 @@ states: "parquet_root": "/absolute/path/to/parquet/mimic-iv", "raw_root": "/absolute/path/to/raw_files/mimic-iv", "steps": [ - {"name": "raw_files", "status": "blocked", "message": "Download manually and rerun init."}, - {"name": "parquet", "status": "skipped"}, - {"name": "database", "status": "skipped"} + {"name": "raw_files", "status": "completed", "message": "Raw files are present."}, + {"name": "parquet", "status": "completed", "message": "Converted CSV to Parquet."}, + {"name": "database", "status": "completed", "message": "Created DuckDB views."} ], "warnings": [] } ``` Init step status is one of `skipped`, `completed`, `blocked`, or `failed`. -Credentialed/manual-download cases that the human CLI treats as informational -return `ok: true` with blocked or skipped steps. +For credentialed datasets, missing raw/parquet/database artifacts are an +`ok: false` `raw_files_missing` result unless `--download` is requested. + +For progress-aware wrappers, add `--events ndjson`: + +```bash +m4 init mimic-iv --json --events ndjson --no-interactive --download \ + --physionet-credentials-file /path/to/physionet-credentials.json +``` + +With `--events ndjson`, stdout changes from one JSON object to a newline-delimited +JSON stream. The final command payload appears in `operation_completed.result`; +setup failures appear in `operation_failed.error`. + +PhysioNet credentials files are JSON: + +```json +{ + "username": "YOUR_USERNAME", + "password": "YOUR_PASSWORD" +} +``` + +Do not expose passwords as command-line flags. M4 follows the same source basis +as PhysioNet's documented recursive resumable downloads (`wget -r -N -c -np` +from `/files/...`) but performs the download internally so wrappers can receive +structured progress and error events. Agent-oriented commands use a stable envelope with `version`, `ok`, `command`, `context`, `data`, `warnings`, and optional provenance fields: diff --git a/src/m4/__init__.py b/src/m4/__init__.py index 1caf13a5..f9ee30ad 100644 --- a/src/m4/__init__.py +++ b/src/m4/__init__.py @@ -13,7 +13,7 @@ For MCP server usage, run: m4 serve """ -__version__ = "0.4.5" +__version__ = "0.5.0" # Expose API functions at package level for easy imports from vitrine import show @@ -21,6 +21,7 @@ from m4.api import ( # Exceptions DatasetError, + M4Client, M4Error, ModalityError, QueryError, @@ -28,6 +29,7 @@ execute_query, # Dataset management get_active_dataset, + get_capabilities, # Clinical notes get_note, get_schema, @@ -43,12 +45,14 @@ __all__ = [ "DatasetError", + "M4Client", "M4Error", "ModalityError", "QueryError", "__version__", "execute_query", "get_active_dataset", + "get_capabilities", "get_note", "get_schema", "get_table_info", diff --git a/src/m4/api.py b/src/m4/api.py index fc147ec6..fedfe9cb 100644 --- a/src/m4/api.py +++ b/src/m4/api.py @@ -32,23 +32,14 @@ import pandas as pd +from m4.client import M4Client from m4.config import _ensure_custom_datasets_loaded from m4.config import get_active_dataset as _get_active_dataset from m4.config import set_active_dataset as _set_active_dataset from m4.core.datasets import DatasetRegistry from m4.core.exceptions import DatasetError, M4Error, ModalityError, QueryError -from m4.core.telemetry import invoke_tracked, set_interface -from m4.core.tools import ToolRegistry, ToolSelector, init_tools -from m4.core.tools.notes import ( - GetNoteInput, - ListPatientNotesInput, - SearchNotesInput, -) -from m4.core.tools.tabular import ( - ExecuteQueryInput, - GetDatabaseSchemaInput, - GetTableInfoInput, -) +from m4.core.telemetry import set_interface +from m4.core.tools import ToolSelector, init_tools # Initialize tools on module import init_tools() @@ -60,11 +51,13 @@ # Re-export exceptions for convenience __all__ = [ "DatasetError", + "M4Client", "M4Error", "ModalityError", "QueryError", "execute_query", "get_active_dataset", + "get_capabilities", "get_note", "get_schema", "get_table_info", @@ -138,6 +131,13 @@ def get_active_dataset() -> str: raise DatasetError(str(e)) from e +def get_capabilities() -> dict[str, Any]: + """Return the stable M4 capability manifest.""" + return M4Client.from_active( + interface="python_api", allow_missing_dataset=True + ).capabilities() + + # ============================================================================= # Tabular Data Tools # ============================================================================= @@ -157,9 +157,7 @@ def get_schema() -> dict[str, Any]: >>> print(schema['tables']) ['admissions', 'diagnoses_icd', 'patients', ...] """ - dataset = DatasetRegistry.get_active() - tool = ToolRegistry.get("get_database_schema") - return invoke_tracked(tool, dataset, GetDatabaseSchemaInput()) + return M4Client.from_active(interface="python_api").schema() def get_table_info(table_name: str, show_sample: bool = True) -> dict[str, Any]: @@ -184,10 +182,8 @@ def get_table_info(table_name: str, show_sample: bool = True) -> dict[str, Any]: >>> print(info['schema']) # DataFrame with column info >>> print(info['sample']) # DataFrame with sample rows """ - dataset = DatasetRegistry.get_active() - tool = ToolRegistry.get("get_table_info") - return invoke_tracked( - tool, dataset, GetTableInfoInput(table_name=table_name, show_sample=show_sample) + return M4Client.from_active(interface="python_api").table_info( + table_name, show_sample=show_sample ) @@ -211,9 +207,7 @@ def execute_query(sql: str) -> pd.DataFrame: 0 M 55 1 F 45 """ - dataset = DatasetRegistry.get_active() - tool = ToolRegistry.get("execute_query") - return invoke_tracked(tool, dataset, ExecuteQueryInput(sql_query=sql)) + return M4Client.from_active(interface="python_api").query(sql) # ============================================================================= @@ -266,17 +260,11 @@ def search_notes( """ _check_notes_compatibility("search_notes") - dataset = DatasetRegistry.get_active() - tool = ToolRegistry.get("search_notes") - return invoke_tracked( - tool, - dataset, - SearchNotesInput( - query=query, - note_type=note_type, - limit=limit, - snippet_length=snippet_length, - ), + return M4Client.from_active(interface="python_api").search_notes( + query=query, + note_type=note_type, + limit=limit, + snippet_length=snippet_length, ) @@ -306,12 +294,8 @@ def get_note(note_id: str, max_length: int | None = None) -> dict[str, Any]: """ _check_notes_compatibility("get_note") - dataset = DatasetRegistry.get_active() - tool = ToolRegistry.get("get_note") - return invoke_tracked( - tool, - dataset, - GetNoteInput(note_id=note_id, max_length=max_length), + return M4Client.from_active(interface="python_api").get_note( + note_id=note_id, max_length=max_length ) @@ -344,16 +328,10 @@ def list_patient_notes( """ _check_notes_compatibility("list_patient_notes") - dataset = DatasetRegistry.get_active() - tool = ToolRegistry.get("list_patient_notes") - return invoke_tracked( - tool, - dataset, - ListPatientNotesInput( - subject_id=subject_id, - note_type=note_type, - limit=limit, - ), + return M4Client.from_active(interface="python_api").list_patient_notes( + subject_id=subject_id, + note_type=note_type, + limit=limit, ) diff --git a/src/m4/apps/cohort_builder/tool.py b/src/m4/apps/cohort_builder/tool.py index 6a753b8e..66bcc4c1 100644 --- a/src/m4/apps/cohort_builder/tool.py +++ b/src/m4/apps/cohort_builder/tool.py @@ -17,6 +17,7 @@ build_gender_distribution_sql, ) from m4.core.backends import get_backend +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition, Modality from m4.core.exceptions import QueryError, SecurityError from m4.core.tools.base import ToolInput @@ -57,7 +58,10 @@ class CohortBuilderTool: supported_datasets: frozenset[str] | None = frozenset({"mimic-iv-demo", "mimic-iv"}) def invoke( - self, dataset: DatasetDefinition, params: CohortBuilderInput + self, + dataset: DatasetDefinition, + params: CohortBuilderInput, + context: M4ExecutionContext | None = None, ) -> dict[str, Any]: """Launch the cohort builder. @@ -119,7 +123,10 @@ class QueryCohortTool: supported_datasets: frozenset[str] | None = frozenset({"mimic-iv-demo", "mimic-iv"}) def invoke( - self, dataset: DatasetDefinition, params: QueryCohortInput + self, + dataset: DatasetDefinition, + params: QueryCohortInput, + context: M4ExecutionContext | None = None, ) -> dict[str, Any]: """Execute cohort query and return results. @@ -153,23 +160,35 @@ def invoke( ) # Execute queries - backend = get_backend() + backend = context.backend if context else get_backend() - count_result = backend.execute_query(count_sql, dataset) + count_result = ( + backend.execute_query(count_sql, dataset, context) + if context + else backend.execute_query(count_sql, dataset) + ) if not count_result.success: raise QueryError( count_result.error or "Count query failed", sql=count_sql, ) - demographics_result = backend.execute_query(demographics_sql, dataset) + demographics_result = ( + backend.execute_query(demographics_sql, dataset, context) + if context + else backend.execute_query(demographics_sql, dataset) + ) if not demographics_result.success: raise QueryError( demographics_result.error or "Demographics query failed", sql=demographics_sql, ) - gender_result = backend.execute_query(gender_sql, dataset) + gender_result = ( + backend.execute_query(gender_sql, dataset, context) + if context + else backend.execute_query(gender_sql, dataset) + ) if not gender_result.success: raise QueryError( gender_result.error or "Gender query failed", diff --git a/src/m4/cli.py b/src/m4/cli.py index 62ae24c1..5bf6724a 100644 --- a/src/m4/cli.py +++ b/src/m4/cli.py @@ -13,6 +13,7 @@ import pandas as pd import typer +from m4.client import M4Client from m4.config import ( get_active_backend, get_active_dataset, @@ -55,22 +56,20 @@ materialize_all, ) from m4.core.exceptions import DatasetError, M4Error -from m4.core.telemetry import invoke_tracked, set_interface -from m4.core.tools import ToolRegistry, init_tools -from m4.core.tools.tabular import ( - ExecuteQueryInput, - GetDatabaseSchemaInput, - GetTableInfoInput, -) +from m4.core.tools import init_tools from m4.data_io import ( + PhysioNetCredentials, convert_csv_to_parquet, download_dataset, init_duckdb_from_parquet, verify_table_rowcount, ) from m4.services.backend import set_active_backend_service +from m4.services.download import build_wget_command, download_dataset_service +from m4.services.events import NdjsonEventReporter from m4.services.init import initialize_dataset_service from m4.services.results import CommandError, CommandResult +from m4.services.setup import doctor_service, quickstart_service, setup_agent_service from m4.services.status import collect_status_snapshot from m4.services.use import set_active_dataset_service @@ -110,6 +109,16 @@ def _emit_command_json(result: CommandResult | CommandError) -> None: _emit_json(result.to_json_dict()) +def _dotenv_lines(values: dict[str, Any]) -> list[str]: + lines = [] + for key, value in values.items(): + if value is None: + continue + text = str(value).replace("\n", "\\n") + lines.append(f"{key}={text}") + return lines + + def _json_error(command: str, code: str, message: str, hint: str | None = None) -> None: _emit_command_json( CommandError(command=command, code=code, message=message, hint=hint) @@ -354,6 +363,302 @@ def main_callback( handler.setLevel(logging.INFO) +@app.command("capabilities") +def capabilities_cmd( + json_output: Annotated[ + bool, + typer.Option("--json", help="Print the stable capability manifest as JSON."), + ] = False, +): + """Show M4 interfaces, commands, tools, datasets, limits, and policies.""" + manifest = M4Client.from_active( + interface="cli", allow_missing_dataset=True + ).capabilities() + if json_output: + _emit_json(manifest) + return + + console.print("[bold]M4 capabilities[/bold]") + console.print(f"Schema version: {manifest['schema_version']}") + console.print( + "Interfaces: " + + ", ".join( + sorted(key for key in manifest["interfaces"] if key != "output_formats") + ) + ) + console.print("Datasets: " + ", ".join(ds["name"] for ds in manifest["datasets"])) + console.print("Tools: " + ", ".join(tool["name"] for tool in manifest["tools"])) + console.print("Machine output: [command]m4 capabilities --json[/command]") + + +@app.command("doctor") +def doctor_cmd( + json_output: Annotated[ + bool, + typer.Option("--json", help="Print diagnostics as JSON."), + ] = False, + include_paths: Annotated[ + bool, + typer.Option("--paths", help="Disclose raw local paths in diagnostics."), + ] = False, +): + """Run non-mutating diagnostics for local, BigQuery, and MCP setup.""" + result = doctor_service(include_paths=include_paths) + if json_output: + _emit_command_json(result) + return + + summary = result.data["summary"] + if summary["ok"]: + success("Doctor checks passed") + else: + warning("Doctor found setup issues") + for check in result.data["checks"]: + marker = "OK" if check["ok"] else "FAIL" + console.print(f"{marker} {check['name']}: {check['message']}") + if not check["ok"] and check.get("hint"): + console.print(f" Hint: {check['hint']}") + + +@app.command("download") +def download_cmd( + dataset_name: Annotated[ + str, + typer.Argument(help="Dataset to download or prepare download guidance for."), + ], + target: Annotated[ + str | None, + typer.Option("--target", help="Raw CSV.gz destination root."), + ] = None, + command_only: Annotated[ + bool, + typer.Option("--command-only", help="Only print/generated command guidance."), + ] = False, + json_output: Annotated[ + bool, + typer.Option("--json", help="Print result as JSON."), + ] = False, + events: Annotated[ + str | None, + typer.Option( + "--events", + help="Emit structured progress events. Supported: ndjson.", + ), + ] = None, + physionet_credentials_file: Annotated[ + str | None, + typer.Option( + "--physionet-credentials-file", + help="JSON file containing PhysioNet username and password fields.", + ), + ] = None, +): + """Download public datasets or generate credentialed PhysioNet commands.""" + if events and events != "ndjson": + _json_error( + "download", + "invalid_option", + f"Unsupported event format '{events}'.", + hint="Use: --events ndjson", + ) + if events and not json_output: + _json_error( + "download", + "invalid_option", + "--events requires --json.", + hint="Use: m4 download DATASET --json --events ndjson", + ) + + reporter = NdjsonEventReporter(command="download") if events == "ndjson" else None + credentials = None + if physionet_credentials_file: + try: + credentials = PhysioNetCredentials.from_json_file( + Path(physionet_credentials_file).expanduser().resolve() + ) + except Exception as exc: + result = CommandError( + command="download", + code="missing_credentials", + message=f"Could not read PhysioNet credentials file: {exc}", + ) + if reporter: + reporter.operation_failed(result.to_json_dict()["error"]) + elif json_output: + _emit_command_json(result) + else: + print_error_panel("Download Failed", result.message, hint=result.hint) + raise typer.Exit(code=1) + + if reporter: + reporter.operation_started( + dataset=dataset_name, + command_only=command_only, + credentials=bool(credentials), + ) + + with _silence_m4_logging() if json_output else nullcontext(): + result = download_dataset_service( + dataset_name, + target=target, + command_only=command_only, + physionet_credentials=credentials, + event_reporter=reporter, + ) + + if json_output: + if reporter: + if isinstance(result, CommandError): + reporter.operation_failed(result.to_json_dict()["error"]) + else: + reporter.operation_completed(result.to_json_dict()) + else: + _emit_command_json(result) + if isinstance(result, CommandError): + raise typer.Exit(code=1) + return + + if isinstance(result, CommandError): + print_error_panel("Download Failed", result.message, hint=result.hint) + raise typer.Exit(code=1) + + data = result.data + status = data.get("status") + if status == "completed": + success(f"Downloaded {data['dataset']} to {data['target']}") + elif status == "blocked": + warning(f"Credentialed dataset '{data['dataset']}' requires credentials") + else: + info(f"Download guidance for {data['dataset']}") + + if data.get("wget_command"): + console.print() + print_command(data["wget_command"]) + layout = data.get("layout", {}) + if layout.get("warnings") or layout.get("errors"): + console.print() + warning("Layout validation reported issues") + for item in layout.get("errors", []): + console.print(f" {item}") + for hint in layout.get("recovery", []): + console.print(f" Hint: {hint}") + + +@app.command("setup-agent") +def setup_agent_cmd( + mode: Annotated[ + str, typer.Option("--mode", help="Agent mode: local or protected.") + ] = "local", + client: Annotated[ + str, typer.Option("--client", help="Client: claude or generic.") + ] = "generic", + dataset_name: Annotated[ + str | None, typer.Option("--dataset", help="Default dataset.") + ] = None, + backend_name: Annotated[ + str | None, typer.Option("--backend", help="Backend: duckdb or bigquery.") + ] = None, + project_id: Annotated[ + str | None, typer.Option("--project-id", help="BigQuery billing project ID.") + ] = None, + output_format: Annotated[ + str, + typer.Option("--format", help="Output format: json, dotenv, or text."), + ] = "text", + apply_config: Annotated[ + bool, + typer.Option("--apply", help="Apply dataset/backend/project configuration."), + ] = False, +): + """Emit or apply agent environment and MCP client recommendations.""" + if output_format not in {"json", "dotenv", "text"}: + _json_error( + "setup-agent", + "invalid_format", + f"Unsupported format '{output_format}'.", + "Use --format json, dotenv, or text.", + ) + result = setup_agent_service( + mode=mode, + client=client, + dataset=dataset_name, + backend=backend_name, + project_id=project_id, + apply_config=apply_config, + ) + if isinstance(result, CommandError): + if output_format == "json": + _emit_command_json(result) + else: + print_error_panel("Setup Agent Failed", result.message, hint=result.hint) + raise typer.Exit(code=1) + + if output_format == "json": + _emit_command_json(result) + return + if output_format == "dotenv": + typer.echo("\n".join(_dotenv_lines(result.data["environment"]))) + return + + console.print("[bold]Agent environment[/bold]") + for line in _dotenv_lines(result.data["environment"]): + typer.echo(line) + console.print() + console.print("[bold]Recommended commands[/bold]") + for command in result.data["recommended_commands"]: + console.print(f" [command]{command}[/command]") + + +@app.command("quickstart") +def quickstart_cmd( + workflow: Annotated[ + str, + typer.Option("--workflow", help="Workflow: demo, local, or bigquery."), + ] = "demo", + dataset_name: Annotated[ + str | None, typer.Option("--dataset", help="Dataset name.") + ] = None, + backend_name: Annotated[ + str | None, typer.Option("--backend", help="Backend name.") + ] = None, + project_id: Annotated[ + str | None, typer.Option("--project-id", help="BigQuery billing project ID.") + ] = None, + apply_config: Annotated[ + bool, + typer.Option("--apply", help="Apply the quickstart configuration."), + ] = False, + json_output: Annotated[ + bool, + typer.Option("--json", help="Print result as JSON."), + ] = False, +): + """Show or run the guided happy path for demo, local, or BigQuery use.""" + result = quickstart_service( + workflow=workflow, + dataset=dataset_name, + backend=backend_name, + project_id=project_id, + apply_config=apply_config, + ) + if json_output: + _emit_command_json(result) + if isinstance(result, CommandError): + raise typer.Exit(code=1) + return + if isinstance(result, CommandError): + print_error_panel("Quickstart Failed", result.message, hint=result.hint) + raise typer.Exit(code=1) + + console.print(f"[bold]Quickstart: {result.data['workflow']}[/bold]") + for step in result.data["steps"]: + if "command" in step: + console.print(f" [command]{step['command']}[/command]") + if result.warnings: + for item in result.warnings: + warning(item) + + @app.command("init") def dataset_init_cmd( dataset_name: Annotated[ @@ -398,6 +703,34 @@ def dataset_init_cmd( help="Print result as JSON for scripts and automation.", ), ] = False, + events: Annotated[ + str | None, + typer.Option( + "--events", + help="Emit structured progress events. Supported: ndjson.", + ), + ] = None, + no_interactive: Annotated[ + bool, + typer.Option( + "--no-interactive", + help="Accepted for automation; JSON init never prompts.", + ), + ] = False, + download_requested: Annotated[ + bool, + typer.Option( + "--download", + help="Download missing raw files when a dataset listing URL is configured.", + ), + ] = False, + physionet_credentials_file: Annotated[ + str | None, + typer.Option( + "--physionet-credentials-file", + help="JSON file containing PhysioNet username and password fields.", + ), + ] = None, ): """ Initialize a local dataset in one step by detecting what's already present: @@ -409,15 +742,64 @@ def dataset_init_cmd( - Auto-download is based on the dataset definition URL. - For datasets without a download URL (e.g. mimic-iv-full), you must provide the --src path or place files in the expected location. """ + if events and events != "ndjson": + _json_error( + "init", + "invalid_option", + f"Unsupported event format '{events}'.", + hint="Use: --events ndjson", + ) + if events and not json_output: + _json_error( + "init", + "invalid_option", + "--events requires --json.", + hint="Use: m4 init DATASET --json --events ndjson", + ) + if json_output: + reporter = NdjsonEventReporter(command="init") if events == "ndjson" else None + credentials = None + if physionet_credentials_file: + try: + credentials = PhysioNetCredentials.from_json_file( + Path(physionet_credentials_file).expanduser().resolve() + ) + except Exception as exc: + result = CommandError( + command="init", + code="missing_credentials", + message=f"Could not read PhysioNet credentials file: {exc}", + ) + if reporter: + reporter.operation_failed(result.to_json_dict()["error"]) + else: + _emit_command_json(result) + raise typer.Exit(code=1) + + if reporter: + reporter.operation_started( + dataset=dataset_name, + download=download_requested, + no_interactive=no_interactive, + ) with _silence_m4_logging(): result = initialize_dataset_service( dataset_name, src=src, db_path_str=db_path_str, force=force, + download=download_requested, + physionet_credentials=credentials, + event_reporter=reporter, ) - _emit_command_json(result) + if reporter: + if isinstance(result, CommandError): + reporter.operation_failed(result.to_json_dict()["error"]) + else: + reporter.operation_completed(result.to_json_dict()) + else: + _emit_command_json(result) if isinstance(result, CommandError): raise typer.Exit(code=1) return @@ -508,7 +890,7 @@ def dataset_init_cmd( console.print() # Wget command tailored to the user's path - wget_cmd = f"wget -r -N -c -np --user YOUR_USERNAME --ask-password {base_url} -P {csv_root}" + wget_cmd = build_wget_command(ds, csv_root) print_command(wget_cmd) console.print() console.print( @@ -1199,33 +1581,35 @@ def schema_cmd( """List tables for a dataset/backend pair without mutating active config.""" command = "schema" init_tools() - set_interface("cli") + ctx = resolve_runtime_context(dataset=dataset_name, backend=backend_name) with _silence_m4_logging() if json_output else nullcontext(): - with _runtime_env_override(dataset=dataset_name, backend=backend_name): - ctx = resolve_runtime_context(dataset=dataset_name, backend=backend_name) - dataset = _resolve_agent_dataset( - command, dataset_name, ctx.public_context() + try: + client = M4Client( + dataset=dataset_name, + backend=backend_name, + interface="cli", + project_id=ctx.project_id, + path_disclosure=ctx.path_disclosure, ) - try: - tool = ToolRegistry.get("get_database_schema") - result = invoke_tracked(tool, dataset, GetDatabaseSchemaInput()) - except M4Error as exc: - if json_output: - _emit_agent_error( - command, - "schema_failed", - str(exc), - context=ctx.public_context(), - ) - error(str(exc)) - raise typer.Exit(code=1) + result = client.schema() + ctx_public = client.context.public_context() + except M4Error as exc: + if json_output: + _emit_agent_error( + command, + "schema_failed", + str(exc), + context=ctx.public_context(), + ) + error(str(exc)) + raise typer.Exit(code=1) if json_output: _emit_json( _agent_success_payload( command, {"tables": result.get("tables", [])}, - context=ctx.public_context(), + context=ctx_public, ) ) return @@ -1266,33 +1650,28 @@ def describe_table_cmd( """Describe a single table without mutating active config.""" command = "describe-table" init_tools() - set_interface("cli") + ctx = resolve_runtime_context(dataset=dataset_name, backend=backend_name) with _silence_m4_logging() if json_output else nullcontext(): - with _runtime_env_override(dataset=dataset_name, backend=backend_name): - ctx = resolve_runtime_context(dataset=dataset_name, backend=backend_name) - dataset = _resolve_agent_dataset( - command, dataset_name, ctx.public_context() + try: + client = M4Client( + dataset=dataset_name, + backend=backend_name, + interface="cli", + project_id=ctx.project_id, + path_disclosure=ctx.path_disclosure, ) - try: - tool = ToolRegistry.get("get_table_info") - result = invoke_tracked( - tool, - dataset, - GetTableInfoInput( - table_name=table_name, - show_sample=show_sample, - ), + result = client.table_info(table_name, show_sample=show_sample) + ctx_public = client.context.public_context() + except M4Error as exc: + if json_output: + _emit_agent_error( + command, + "describe_table_failed", + str(exc), + context=ctx.public_context(), ) - except M4Error as exc: - if json_output: - _emit_agent_error( - command, - "describe_table_failed", - str(exc), - context=ctx.public_context(), - ) - error(str(exc)) - raise typer.Exit(code=1) + error(str(exc)) + raise typer.Exit(code=1) if json_output: _emit_json( @@ -1303,7 +1682,7 @@ def describe_table_cmd( "schema": _dataframe_payload(result.get("schema")), "sample": _dataframe_payload(result.get("sample")), }, - context=ctx.public_context(), + context=ctx_public, ) ) return @@ -1343,26 +1722,28 @@ def query_cmd( """Execute a read-only SQL query without mutating active config.""" command = "query" init_tools() - set_interface("cli") + ctx = resolve_runtime_context(dataset=dataset_name, backend=backend_name) with _silence_m4_logging() if json_output else nullcontext(): - with _runtime_env_override(dataset=dataset_name, backend=backend_name): - ctx = resolve_runtime_context(dataset=dataset_name, backend=backend_name) - dataset = _resolve_agent_dataset( - command, dataset_name, ctx.public_context() + try: + client = M4Client( + dataset=dataset_name, + backend=backend_name, + interface="cli", + project_id=ctx.project_id, + path_disclosure=ctx.path_disclosure, ) - try: - tool = ToolRegistry.get("execute_query") - result = invoke_tracked(tool, dataset, ExecuteQueryInput(sql_query=sql)) - except M4Error as exc: - if json_output: - _emit_agent_error( - command, - "query_failed", - str(exc), - context=ctx.public_context(), - ) - error(str(exc)) - raise typer.Exit(code=1) + result = client.query(sql) + ctx_public = client.context.public_context() + except M4Error as exc: + if json_output: + _emit_agent_error( + command, + "query_failed", + str(exc), + context=ctx.public_context(), + ) + error(str(exc)) + raise typer.Exit(code=1) if json_output: _emit_json( @@ -1371,7 +1752,7 @@ def query_cmd( { "result": _dataframe_payload(result), }, - context=ctx.public_context(), + context=ctx_public, ) ) return @@ -1395,18 +1776,39 @@ def agent_env_cmd( str | None, typer.Option("--backend", help="Default backend for agent sessions."), ] = None, + project_id: Annotated[ + str | None, + typer.Option("--project-id", help="BigQuery billing project ID."), + ] = None, json_output: Annotated[ bool, typer.Option("--json", help="Print result as JSON."), ] = False, + output_format: Annotated[ + str, + typer.Option("--format", help="Output format: dotenv, json, or text."), + ] = "dotenv", include_paths: Annotated[ bool, typer.Option("--paths", help="Disclose raw local paths in metadata."), ] = False, ): """Return environment variables and command recommendations for agents.""" + if json_output: + output_format = "json" + if output_format not in {"dotenv", "json", "text"}: + if output_format == "json": + _emit_agent_error( + "agent-env", + "invalid_format", + f"Unsupported format '{output_format}'.", + hint="Use --format dotenv, json, or text.", + ) + error("Unsupported format. Use 'dotenv', 'json', or 'text'.") + raise typer.Exit(code=1) + if mode not in {"local", "protected"}: - if json_output: + if output_format == "json": _emit_agent_error( "agent-env", "invalid_mode", @@ -1428,6 +1830,8 @@ def agent_env_cmd( "M4_ACTOR": ctx.actor, "M4_TELEMETRY_DIR": str(ctx.telemetry_dir), } + if project_id or ctx.project_id: + env["M4_PROJECT_ID"] = project_id or ctx.project_id if mode == "local": env["M4_DATA_DIR"] = str(ctx.data_dir) else: @@ -1469,7 +1873,7 @@ def agent_env_cmd( ], } - if json_output: + if output_format == "json": _emit_json( _agent_success_payload( "agent-env", @@ -1480,8 +1884,17 @@ def agent_env_cmd( ) return - for key, value in data["environment"].items(): - typer.echo(f"{key}={value}") + if output_format == "dotenv": + typer.echo("\n".join(_dotenv_lines(data["environment"]))) + return + + console.print("[bold]Agent environment[/bold]") + for line in _dotenv_lines(data["environment"]): + typer.echo(line) + console.print() + console.print("[bold]Recommended commands[/bold]") + for command in data["recommended_commands"]: + console.print(f" [command]{command}[/command]") @provenance_app.command("export") diff --git a/src/m4/client.py b/src/m4/client.py new file mode 100644 index 00000000..4891054f --- /dev/null +++ b/src/m4/client.py @@ -0,0 +1,233 @@ +"""First-class Python client for M4 data access.""" + +import os +from dataclasses import replace +from pathlib import Path +from typing import Any + +import pandas as pd + +from m4.config import ( + _ensure_custom_datasets_loaded, + get_active_backend, + get_bigquery_project_id, +) +from m4.core.backends import Backend, get_backend +from m4.core.context import M4ExecutionContext +from m4.core.datasets import DatasetDefinition, DatasetRegistry +from m4.core.exceptions import DatasetError, ModalityError +from m4.core.telemetry import invoke_tracked +from m4.core.tools import ToolRegistry, ToolSelector, init_tools +from m4.core.tools.management import ListDatasetsInput +from m4.core.tools.notes import ( + GetNoteInput, + ListPatientNotesInput, + SearchNotesInput, +) +from m4.core.tools.tabular import ( + ExecuteQueryInput, + GetDatabaseSchemaInput, + GetTableInfoInput, +) + + +class M4Client: + """Resolved M4 client for Python, CLI, and MCP adapters.""" + + def __init__( + self, + dataset: str | DatasetDefinition | None = None, + backend: str | Backend | None = None, + study_id: str | None = None, + session_id: str | None = None, + actor: str | None = None, + interface: str = "python_api", + project_id: str | None = None, + db_path: str | Path | None = None, + path_disclosure: bool = False, + allow_missing_dataset: bool = False, + ) -> None: + init_tools() + _ensure_custom_datasets_loaded() + + self.dataset = self._resolve_dataset( + dataset, allow_missing_dataset=allow_missing_dataset + ) + self.backend_name, self.backend = self._resolve_backend(backend) + self.context = M4ExecutionContext( + dataset=self.dataset, + backend_name=self.backend_name, + backend=self.backend, + interface=interface, + study_id=study_id, + session_id=session_id, + actor=actor, + project_id=project_id, + db_path=Path(db_path).expanduser() if db_path else None, + path_disclosure=path_disclosure, + ) + self._tool_selector = ToolSelector() + + @classmethod + def from_active( + cls, interface: str = "python_api", allow_missing_dataset: bool = False + ) -> "M4Client": + """Create a client from active runtime configuration and environment.""" + path_disclosure = os.getenv("M4_PATH_DISCLOSURE", "").lower() in { + "1", + "true", + "yes", + "on", + "paths", + } + return cls( + interface=interface, + study_id=os.getenv("M4_STUDY_ID"), + session_id=os.getenv("M4_SESSION_ID"), + actor=os.getenv("M4_ACTOR"), + project_id=get_bigquery_project_id(), + db_path=os.getenv("M4_DB_PATH"), + path_disclosure=path_disclosure, + allow_missing_dataset=allow_missing_dataset, + ) + + def schema(self) -> dict[str, Any]: + """Return backend information and available table names.""" + return self._invoke("get_database_schema", GetDatabaseSchemaInput()) + + def table_info(self, table_name: str, show_sample: bool = True) -> dict[str, Any]: + """Return schema and optional sample rows for a table.""" + return self._invoke( + "get_table_info", + GetTableInfoInput(table_name=table_name, show_sample=show_sample), + ) + + def query(self, sql: str) -> pd.DataFrame: + """Execute a read-only SQL query.""" + return self._invoke("execute_query", ExecuteQueryInput(sql_query=sql)) + + def list_datasets(self) -> list[str]: + """Return registered dataset names.""" + _ensure_custom_datasets_loaded() + return [ds.name for ds in DatasetRegistry.list_all()] + + def dataset_status(self) -> dict[str, Any]: + """Return detailed dataset availability information.""" + return self._invoke("list_datasets", ListDatasetsInput()) + + def capabilities(self) -> dict[str, Any]: + """Return the stable M4 capability manifest.""" + from m4.services.capabilities import build_capabilities_manifest + + return build_capabilities_manifest() + + def search_notes( + self, + query: str, + note_type: str = "all", + limit: int = 5, + snippet_length: int = 300, + ) -> dict[str, Any]: + """Search clinical notes by keyword.""" + return self._invoke( + "search_notes", + SearchNotesInput( + query=query, + note_type=note_type, + limit=limit, + snippet_length=snippet_length, + ), + ) + + def get_note(self, note_id: str, max_length: int | None = None) -> dict[str, Any]: + """Retrieve a clinical note by note ID.""" + return self._invoke( + "get_note", + GetNoteInput(note_id=note_id, max_length=max_length), + ) + + def list_patient_notes( + self, + subject_id: int, + note_type: str = "all", + limit: int = 20, + ) -> dict[str, Any]: + """List clinical notes for a patient.""" + return self._invoke( + "list_patient_notes", + ListPatientNotesInput( + subject_id=subject_id, + note_type=note_type, + limit=limit, + ), + ) + + def cohort_builder(self) -> dict[str, Any]: + """Return cohort builder launch metadata.""" + from m4.apps.cohort_builder.tool import CohortBuilderInput + + return self._invoke("cohort_builder", CohortBuilderInput()) + + def query_cohort(self, **criteria: Any) -> dict[str, Any]: + """Query cohort counts and demographics.""" + from m4.apps.cohort_builder.query_builder import QueryCohortInput + + return self._invoke("query_cohort", QueryCohortInput(**criteria)) + + def invoke_tool(self, tool_name: str, params: Any) -> Any: + """Invoke a registered tool through this client's context.""" + return self._invoke(tool_name, params) + + def _invoke(self, tool_name: str, params: Any) -> Any: + tool = ToolRegistry.get(tool_name) + if tool is None: + raise DatasetError(f"Unknown tool: {tool_name}") + + compat = self._tool_selector.check_compatibility(tool_name, self.dataset) + if not compat.compatible: + raise ModalityError(compat.error_message) + + self._ensure_backend() + return invoke_tracked(tool, self.dataset, params, self.context) + + def _resolve_dataset( + self, + dataset: str | DatasetDefinition | None, + *, + allow_missing_dataset: bool = False, + ) -> DatasetDefinition: + if isinstance(dataset, DatasetDefinition): + return dataset + + if dataset is None: + try: + return DatasetRegistry.get_active() + except DatasetError: + if allow_missing_dataset: + return DatasetRegistry.list_all()[0] + raise + + resolved = DatasetRegistry.get(dataset.lower()) + if resolved is None: + supported = ", ".join(ds.name for ds in DatasetRegistry.list_all()) + raise DatasetError( + f"Dataset '{dataset}' not found. Available datasets: {supported}", + dataset_name=dataset, + ) + return resolved + + def _resolve_backend( + self, backend: str | Backend | None + ) -> tuple[str, Backend | None]: + if backend is not None and not isinstance(backend, str): + return backend.name, backend + + backend_name = (backend or get_active_backend()).lower() + return backend_name, None + + def _ensure_backend(self) -> None: + if self.backend is not None: + return + + self.backend = get_backend(self.backend_name) + self.context = replace(self.context, backend=self.backend) diff --git a/src/m4/config.py b/src/m4/config.py index 82942aae..31501955 100644 --- a/src/m4/config.py +++ b/src/m4/config.py @@ -165,11 +165,16 @@ def _get_project_data_dir() -> Path: # -------------------------------------------------- -def _ensure_custom_datasets_loaded(): +def ensure_custom_datasets_loaded(): """Ensure custom datasets are loaded from the custom datasets directory.""" DatasetRegistry.load_custom_datasets(_CUSTOM_DATASETS_DIR) +def _ensure_custom_datasets_loaded(): + """Backward-compatible private wrapper for existing internal callers.""" + ensure_custom_datasets_loaded() + + def resolve_runtime_context( *, dataset: str | None = None, @@ -291,6 +296,10 @@ def _has_parquet_files(path: Path | None) -> bool: return bool(path and path.exists() and any(path.rglob("*.parquet"))) +def _has_raw_files(path: Path | None) -> bool: + return bool(path and path.exists() and any(path.rglob("*.csv.gz"))) + + def detect_available_local_datasets() -> dict[str, dict[str, Any]]: """Return presence flags for all registered datasets.""" _ensure_custom_datasets_loaded() @@ -312,10 +321,13 @@ def detect_available_local_datasets() -> dict[str, dict[str, Any]]: db_path_str = cfg.get("duckdb_paths", {}).get(name) db_path = Path(db_path_str) if db_path_str else get_default_database_path(name) + raw_root = _PROJECT_DATA_DIR / "raw_files" / name results[name] = { + "raw_present": _has_raw_files(raw_root), "parquet_present": _has_parquet_files(parquet_root), "db_present": bool(db_path and db_path.exists()), + "raw_root": str(raw_root), "parquet_root": str(parquet_root) if parquet_root else "", "db_path": str(db_path) if db_path else "", } diff --git a/src/m4/core/backends/base.py b/src/m4/core/backends/base.py index 42f9535a..ae03ca6e 100644 --- a/src/m4/core/backends/base.py +++ b/src/m4/core/backends/base.py @@ -11,6 +11,7 @@ import pandas as pd from m4.config import logger +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition # Re-export exceptions from the central exceptions module for backwards compatibility @@ -143,7 +144,7 @@ class Backend(Protocol): Example: class DuckDBBackend: - def execute_query(self, sql, dataset): + def execute_query(self, sql, dataset, context): # DuckDB-specific implementation ... @@ -156,12 +157,18 @@ def get_table_list(self, dataset): result = backend.execute_query("SELECT * FROM patients LIMIT 5", mimic_demo) """ - def execute_query(self, sql: str, dataset: DatasetDefinition) -> QueryResult: + def execute_query( + self, + sql: str, + dataset: DatasetDefinition, + context: M4ExecutionContext, + ) -> QueryResult: """Execute a SQL query against the dataset. Args: sql: SQL query string (must be a safe SELECT or PRAGMA query) dataset: The dataset definition to query against + context: Resolved execution context Returns: QueryResult with the query output or error message @@ -172,7 +179,9 @@ def execute_query(self, sql: str, dataset: DatasetDefinition) -> QueryResult: """ ... - def get_table_list(self, dataset: DatasetDefinition) -> list[str]: + def get_table_list( + self, dataset: DatasetDefinition, context: M4ExecutionContext + ) -> list[str]: """Get list of available tables in the dataset. Args: @@ -184,7 +193,10 @@ def get_table_list(self, dataset: DatasetDefinition) -> list[str]: ... def get_table_info( - self, table_name: str, dataset: DatasetDefinition + self, + table_name: str, + dataset: DatasetDefinition, + context: M4ExecutionContext, ) -> QueryResult: """Get schema information for a specific table. @@ -198,7 +210,11 @@ def get_table_info( ... def get_sample_data( - self, table_name: str, dataset: DatasetDefinition, limit: int = 3 + self, + table_name: str, + dataset: DatasetDefinition, + limit: int, + context: M4ExecutionContext, ) -> QueryResult: """Get sample rows from a table. @@ -212,7 +228,9 @@ def get_sample_data( """ ... - def get_backend_info(self, dataset: DatasetDefinition) -> str: + def get_backend_info( + self, dataset: DatasetDefinition, context: M4ExecutionContext + ) -> str: """Get human-readable information about the current backend. Args: diff --git a/src/m4/core/backends/bigquery.py b/src/m4/core/backends/bigquery.py index f590048f..2df6ce31 100644 --- a/src/m4/core/backends/bigquery.py +++ b/src/m4/core/backends/bigquery.py @@ -13,6 +13,7 @@ TableNotFoundError, sanitize_error_message, ) +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition @@ -81,7 +82,7 @@ def _get_project_id(self, dataset: DatasetDefinition) -> str: # Priority 3: Default return "physionet-data" - def _get_client(self) -> Any: + def _get_client(self, context: M4ExecutionContext | None = None) -> Any: """Get or create a BigQuery client. Clients are cached to avoid re-initialization overhead. @@ -102,9 +103,12 @@ def _get_client(self) -> Any: backend=self.name, ) - from m4.config import get_bigquery_project_id + if context: + project_id = context.project_id + else: + from m4.config import get_bigquery_project_id - project_id = get_bigquery_project_id() + project_id = get_bigquery_project_id() # Check cache if ( @@ -159,7 +163,12 @@ def _translate_canonical_to_bq(self, sql: str, dataset: DatasetDefinition) -> st return sql - def execute_query(self, sql: str, dataset: DatasetDefinition) -> QueryResult: + def execute_query( + self, + sql: str, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, + ) -> QueryResult: """Execute a SQL query against BigQuery. Args: @@ -185,7 +194,7 @@ def execute_query(self, sql: str, dataset: DatasetDefinition) -> QueryResult: import pandas as pd from google.cloud import bigquery as bq - client = self._get_client() + client = self._get_client(context) job_config = bq.QueryJobConfig() query_job = client.query(sql, job_config=job_config) @@ -212,7 +221,11 @@ def execute_query(self, sql: str, dataset: DatasetDefinition) -> QueryResult: error=sanitize_error_message(e, self.name), ) - def get_table_list(self, dataset: DatasetDefinition) -> list[str]: + def get_table_list( + self, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, + ) -> list[str]: """Get list of available tables in the dataset. Returns canonical schema.table names (e.g., mimiciv_hosp.patients). @@ -237,7 +250,11 @@ def get_table_list(self, dataset: DatasetDefinition) -> list[str]: SELECT table_name FROM `{project_id}.{dataset_id}.INFORMATION_SCHEMA.TABLES` """ - result = self.execute_query(query, dataset) + result = ( + self.execute_query(query, dataset, context) + if context + else self.execute_query(query, dataset) + ) if result.error or result.dataframe is None: continue @@ -251,7 +268,10 @@ def get_table_list(self, dataset: DatasetDefinition) -> list[str]: return sorted(tables) def get_table_info( - self, table_name: str, dataset: DatasetDefinition + self, + table_name: str, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, ) -> QueryResult: """Get schema information for a specific table. @@ -312,7 +332,11 @@ def get_table_info( ORDER BY ordinal_position """ - result = self.execute_query(query, dataset) + result = ( + self.execute_query(query, dataset, context) + if context + else self.execute_query(query, dataset) + ) if result.error or result.dataframe is None or result.dataframe.empty: raise TableNotFoundError(table_name, backend=self.name) return result @@ -338,7 +362,11 @@ def get_table_info( ORDER BY ordinal_position """ - result = self.execute_query(query, dataset) + result = ( + self.execute_query(query, dataset, context) + if context + else self.execute_query(query, dataset) + ) if ( not result.error and result.dataframe is not None @@ -349,7 +377,11 @@ def get_table_info( raise TableNotFoundError(table_name, backend=self.name) def get_sample_data( - self, table_name: str, dataset: DatasetDefinition, limit: int = 3 + self, + table_name: str, + dataset: DatasetDefinition, + limit: int = 3, + context: M4ExecutionContext | None = None, ) -> QueryResult: """Get sample rows from a table. @@ -399,7 +431,11 @@ def get_sample_data( ) query = f"SELECT * FROM {full_name} LIMIT {limit}" - return self.execute_query(query, dataset) + return ( + self.execute_query(query, dataset, context) + if context + else self.execute_query(query, dataset) + ) # Simple name - find in configured datasets if not dataset.bigquery_dataset_ids: @@ -418,7 +454,11 @@ def get_sample_data( full_name = f"`{project_id}.{dataset_id}.{table_name}`" query = f"SELECT * FROM {full_name} LIMIT {limit}" - result = self.execute_query(query, dataset) + result = ( + self.execute_query(query, dataset, context) + if context + else self.execute_query(query, dataset) + ) if not result.error: return result @@ -427,7 +467,11 @@ def get_sample_data( error=f"Table '{table_name}' not found in any configured dataset", ) - def get_backend_info(self, dataset: DatasetDefinition) -> str: + def get_backend_info( + self, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, + ) -> str: """Get human-readable information about the current backend. Args: diff --git a/src/m4/core/backends/duckdb.py b/src/m4/core/backends/duckdb.py index 95a49955..a35c45e4 100644 --- a/src/m4/core/backends/duckdb.py +++ b/src/m4/core/backends/duckdb.py @@ -5,6 +5,7 @@ """ import os +import threading from pathlib import Path import duckdb @@ -16,6 +17,7 @@ TableNotFoundError, sanitize_error_message, ) +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition @@ -51,18 +53,23 @@ def __init__(self, db_path_override: str | Path | None = None): regardless of the dataset parameter. """ self._db_path_override = Path(db_path_override) if db_path_override else None + self._connection_lock = threading.Lock() @property def name(self) -> str: """Get the backend name.""" return "duckdb" - def _get_db_path(self, dataset: DatasetDefinition) -> Path: + def _get_db_path( + self, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, + ) -> Path: """Get the database path for a dataset. Priority: - 1. Instance override (db_path_override) - 2. Environment variable M4_DB_PATH + 1. Execution context override + 2. Instance override (db_path_override) 3. Default path based on dataset configuration Args: @@ -74,14 +81,18 @@ def _get_db_path(self, dataset: DatasetDefinition) -> Path: Raises: ConnectionError: If no valid database path can be determined """ - # Priority 1: Instance override + # Priority 1: Per-client execution context override + if context and context.db_path: + return context.db_path + + # Priority 2: Instance override if self._db_path_override: return self._db_path_override - # Priority 2: Environment variable - env_path = os.getenv("M4_DB_PATH") - if env_path: - return Path(env_path) + if context is None: + env_path = os.getenv("M4_DB_PATH") + if env_path: + return Path(env_path) # Priority 3: Default based on dataset db_path = get_default_database_path(dataset.name) @@ -93,7 +104,11 @@ def _get_db_path(self, dataset: DatasetDefinition) -> Path: return db_path - def _connect(self, dataset: DatasetDefinition) -> duckdb.DuckDBPyConnection: + def _connect( + self, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, + ) -> duckdb.DuckDBPyConnection: """Create a connection to the DuckDB database. Args: @@ -105,7 +120,7 @@ def _connect(self, dataset: DatasetDefinition) -> duckdb.DuckDBPyConnection: Raises: ConnectionError: If the database file doesn't exist or can't be opened """ - db_path = self._get_db_path(dataset) + db_path = self._get_db_path(dataset, context) if not db_path.exists(): raise ConnectionError( @@ -134,7 +149,12 @@ def _connect(self, dataset: DatasetDefinition) -> duckdb.DuckDBPyConnection: backend=self.name, ) from e - def execute_query(self, sql: str, dataset: DatasetDefinition) -> QueryResult: + def execute_query( + self, + sql: str, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, + ) -> QueryResult: """Execute a SQL query against the dataset. Args: @@ -145,28 +165,29 @@ def execute_query(self, sql: str, dataset: DatasetDefinition) -> QueryResult: QueryResult with query output as native DataFrame """ try: - conn = self._connect(dataset) - try: - df = conn.execute(sql).df() + with self._connection_lock: + conn = self._connect(dataset, context) + try: + df = conn.execute(sql).df() - if df.empty: - import pandas as pd + if df.empty: + import pandas as pd - return QueryResult( - dataframe=pd.DataFrame(), - row_count=0, - ) + return QueryResult( + dataframe=pd.DataFrame(), + row_count=0, + ) - row_count = len(df) - truncated = row_count > 50 + row_count = len(df) + truncated = row_count > 50 - return QueryResult( - dataframe=df, - row_count=row_count, - truncated=truncated, - ) - finally: - conn.close() + return QueryResult( + dataframe=df, + row_count=row_count, + truncated=truncated, + ) + finally: + conn.close() except ConnectionError: raise @@ -177,7 +198,11 @@ def execute_query(self, sql: str, dataset: DatasetDefinition) -> QueryResult: error=sanitize_error_message(e, self.name), ) - def get_table_list(self, dataset: DatasetDefinition) -> list[str]: + def get_table_list( + self, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, + ) -> list[str]: """Get list of available tables in the dataset. Returns schema-qualified names (e.g. ``mimiciv_hosp.patients``) when @@ -198,7 +223,7 @@ def get_table_list(self, dataset: DatasetDefinition) -> list[str]: WHERE table_schema NOT IN ('main', 'information_schema', 'pg_catalog') ORDER BY table_schema, table_name """ - result = self.execute_query(schema_query, dataset) + result = self.execute_query(schema_query, dataset, context) if ( result.error is None @@ -214,7 +239,7 @@ def get_table_list(self, dataset: DatasetDefinition) -> list[str]: WHERE table_schema = 'main' ORDER BY table_name """ - result = self.execute_query(fallback_query, dataset) + result = self.execute_query(fallback_query, dataset, context) if result.error or result.dataframe is None or result.dataframe.empty: return [] @@ -222,7 +247,10 @@ def get_table_list(self, dataset: DatasetDefinition) -> list[str]: return result.dataframe["table_name"].tolist() def get_table_info( - self, table_name: str, dataset: DatasetDefinition + self, + table_name: str, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, ) -> QueryResult: """Get schema information for a specific table. @@ -257,16 +285,17 @@ def get_table_info( query = f"PRAGMA table_info('{table_name}')" try: - conn = self._connect(dataset) - try: - df = conn.execute(query).df() + with self._connection_lock: + conn = self._connect(dataset, context) + try: + df = conn.execute(query).df() - if df.empty: - raise TableNotFoundError(table_name, backend=self.name) + if df.empty: + raise TableNotFoundError(table_name, backend=self.name) - return QueryResult(dataframe=df, row_count=len(df)) - finally: - conn.close() + return QueryResult(dataframe=df, row_count=len(df)) + finally: + conn.close() except TableNotFoundError: raise @@ -283,7 +312,11 @@ def get_table_info( ) def get_sample_data( - self, table_name: str, dataset: DatasetDefinition, limit: int = 3 + self, + table_name: str, + dataset: DatasetDefinition, + limit: int = 3, + context: M4ExecutionContext | None = None, ) -> QueryResult: """Get sample rows from a table. @@ -306,9 +339,13 @@ def get_sample_data( query = f'SELECT * FROM {schema}."{table}" LIMIT {limit}' else: query = f'SELECT * FROM "{table_name}" LIMIT {limit}' - return self.execute_query(query, dataset) + return self.execute_query(query, dataset, context) - def get_backend_info(self, dataset: DatasetDefinition) -> str: + def get_backend_info( + self, + dataset: DatasetDefinition, + context: M4ExecutionContext | None = None, + ) -> str: """Get human-readable information about the current backend. Args: @@ -321,9 +358,15 @@ def get_backend_info(self, dataset: DatasetDefinition) -> str: f"**Current Backend:** DuckDB (local database)\n" f"**Active Dataset:** {dataset.name}" ) - if os.getenv("M4_PATH_DISCLOSURE", "").lower() in {"1", "true", "yes", "on"}: + disclose = ( + context.path_disclosure + if context + else os.getenv("M4_PATH_DISCLOSURE", "").lower() + in {"1", "true", "yes", "on"} + ) + if disclose: try: - db_path = self._get_db_path(dataset) + db_path = self._get_db_path(dataset, context) except ConnectionError: db_path = "unknown" info += f"\n**Database Path:** {db_path}" diff --git a/src/m4/core/context.py b/src/m4/core/context.py new file mode 100644 index 00000000..aade22f1 --- /dev/null +++ b/src/m4/core/context.py @@ -0,0 +1,34 @@ +"""Execution context passed through M4 tool and backend calls.""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from m4.core.datasets import DatasetDefinition + + +@dataclass(frozen=True) +class M4ExecutionContext: + """Resolved execution context for one M4 client call path.""" + + dataset: DatasetDefinition + backend_name: str + backend: Any + interface: str + study_id: str | None = None + session_id: str | None = None + actor: str | None = None + project_id: str | None = None + db_path: Path | None = None + path_disclosure: bool = False + + def public_context(self) -> dict[str, str | None]: + """Return non-path context fields suitable for JSON envelopes.""" + return { + "dataset": self.dataset.name, + "backend": self.backend_name, + "study_id": self.study_id, + "session_id": self.session_id, + "actor": self.actor, + "project_id": self.project_id, + } diff --git a/src/m4/core/datasets.py b/src/m4/core/datasets.py index 57220484..e8e9a8bd 100644 --- a/src/m4/core/datasets.py +++ b/src/m4/core/datasets.py @@ -82,6 +82,14 @@ class DatasetDefinition: # Format: {"dataset-name": "Description of how to link"} related_datasets: dict[str, str] = field(default_factory=dict) + # Access and local layout guidance. These fields are metadata only; M4 does + # not collect credentialed dataset passwords. + dataset_page_url: str | None = None + dua_url: str | None = None + bigquery_access_url: str | None = None + expected_raw_subdirectories: list[str] = field(default_factory=list) + recommended_local_target_root: str | None = None + # Filesystem directory -> canonical schema name # e.g. {"hosp": "mimiciv_hosp", "icu": "mimiciv_icu"} # Root-level files use empty string key: {"": "eicu_crd"} @@ -232,8 +240,11 @@ def _register_builtins(cls): mimic_iv_demo = DatasetDefinition( name="mimic-iv-demo", description="MIMIC-IV Clinical Database Demo", + dataset_page_url="https://physionet.org/content/mimic-iv-demo/", file_listing_url="https://physionet.org/files/mimic-iv-demo/2.2/", subdirectories_to_scan=["hosp", "icu"], + expected_raw_subdirectories=["hosp", "icu"], + recommended_local_target_root="m4_data/raw_files/mimic-iv-demo", primary_verification_table="mimiciv_hosp.admissions", bigquery_project_id=None, bigquery_dataset_ids=[], @@ -244,8 +255,13 @@ def _register_builtins(cls): mimic_iv = DatasetDefinition( name="mimic-iv", description="MIMIC-IV Clinical Database", + dataset_page_url="https://physionet.org/content/mimiciv/", + dua_url="https://physionet.org/content/mimiciv/", + bigquery_access_url="https://physionet.org/content/mimiciv/view-required-training/3.1/", file_listing_url="https://physionet.org/files/mimiciv/3.1/", subdirectories_to_scan=["hosp", "icu"], + expected_raw_subdirectories=["hosp", "icu"], + recommended_local_target_root="m4_data/raw_files/mimic-iv", primary_verification_table="mimiciv_hosp.admissions", bigquery_project_id="physionet-data", bigquery_dataset_ids=[ @@ -276,8 +292,13 @@ def _register_builtins(cls): mimic_iv_note = DatasetDefinition( name="mimic-iv-note", description="MIMIC-IV Clinical Notes (discharge summaries, radiology reports)", + dataset_page_url="https://physionet.org/content/mimic-iv-note/", + dua_url="https://physionet.org/content/mimic-iv-note/", + bigquery_access_url="https://physionet.org/content/mimic-iv-note/view-required-training/2.2/", file_listing_url="https://physionet.org/files/mimic-iv-note/2.2/", subdirectories_to_scan=["note"], + expected_raw_subdirectories=["note"], + recommended_local_target_root="m4_data/raw_files/mimic-iv-note", primary_verification_table="mimiciv_note.discharge", bigquery_project_id="physionet-data", bigquery_dataset_ids=["mimiciv_note"], @@ -296,8 +317,13 @@ def _register_builtins(cls): eicu = DatasetDefinition( name="eicu", description="eICU Collaborative Research Database", + dataset_page_url="https://physionet.org/content/eicu-crd/", + dua_url="https://physionet.org/content/eicu-crd/", + bigquery_access_url="https://physionet.org/content/eicu-crd/view-required-training/2.0/", file_listing_url="https://physionet.org/files/eicu-crd/2.0/", subdirectories_to_scan=[], + expected_raw_subdirectories=[], + recommended_local_target_root="m4_data/raw_files/eicu", primary_verification_table="eicu_crd.patient", bigquery_project_id="physionet-data", bigquery_dataset_ids=["eicu_crd"], diff --git a/src/m4/core/telemetry.py b/src/m4/core/telemetry.py index d3153393..fcc34b73 100644 --- a/src/m4/core/telemetry.py +++ b/src/m4/core/telemetry.py @@ -20,6 +20,7 @@ from logging.handlers import RotatingFileHandler from typing import Any +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition from m4.core.tools.base import Tool, ToolInput @@ -198,14 +199,19 @@ def reset(self) -> None: # --------------------------------------------------------------------------- -def invoke_tracked(tool: Tool, dataset: DatasetDefinition, params: ToolInput) -> Any: +def invoke_tracked( + tool: Tool, + dataset: DatasetDefinition, + params: ToolInput, + context: M4ExecutionContext, +) -> Any: """Invoke a tool with telemetry tracking. Wraps tool.invoke() to record timing, success/failure, and context. Records are written to JSONL and logged at INFO level. On failure, the original exception is re-raised. """ - interface = _interface_var.get() + interface = context.interface agent_id = _agent_id_var.get() terminal_session = _get_terminal_session() dataset_name = getattr(dataset, "name", None) @@ -234,7 +240,7 @@ def invoke_tracked(tool: Tool, dataset: DatasetDefinition, params: ToolInput) -> row_count = None try: - result = tool.invoke(dataset, params) + result = tool.invoke(dataset, params, context) import pandas as pd row_count = len(result) if isinstance(result, pd.DataFrame) else None @@ -260,9 +266,9 @@ def invoke_tracked(tool: Tool, dataset: DatasetDefinition, params: ToolInput) -> error_message=error_message, params_summary=params_summary, row_count=row_count, - study_id=os.environ.get("M4_STUDY_ID"), - session_id=os.environ.get("M4_SESSION_ID"), - actor=os.environ.get("M4_ACTOR") or agent_id, + study_id=context.study_id, + session_id=context.session_id, + actor=context.actor or agent_id, query_hash=query_hash, ) diff --git a/src/m4/core/tools/base.py b/src/m4/core/tools/base.py index 42f7f176..2ea28c1c 100644 --- a/src/m4/core/tools/base.py +++ b/src/m4/core/tools/base.py @@ -14,6 +14,7 @@ from dataclasses import dataclass from typing import Any, Protocol, runtime_checkable +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition, Modality @@ -85,9 +86,9 @@ class ExecuteQueryTool: required_modalities = frozenset({Modality.TABULAR}) supported_datasets = None - def invoke(self, dataset, params) -> pd.DataFrame: + def invoke(self, dataset, params, context) -> pd.DataFrame: # Returns DataFrame directly - result = backend.execute_query(sql, dataset) + result = backend.execute_query(sql, dataset, context) if not result.success: raise QueryError(result.error) return result.dataframe @@ -108,12 +109,18 @@ def is_compatible(self, dataset): required_modalities: AbstractSet[Modality] supported_datasets: AbstractSet[str] | None # None = all compatible datasets - def invoke(self, dataset: DatasetDefinition, params: ToolInput) -> Any: + def invoke( + self, + dataset: DatasetDefinition, + params: ToolInput, + context: M4ExecutionContext, + ) -> Any: """Execute the tool with given parameters on the specified dataset. Args: dataset: The dataset definition to query params: Tool-specific input parameters + context: Resolved execution context for this invocation Returns: Native Python type appropriate for the tool: diff --git a/src/m4/core/tools/management.py b/src/m4/core/tools/management.py index ae2ff379..6982bae8 100644 --- a/src/m4/core/tools/management.py +++ b/src/m4/core/tools/management.py @@ -21,6 +21,7 @@ get_active_dataset, set_active_dataset, ) +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition, DatasetRegistry, Modality from m4.core.derived.builtins import has_derived_support, list_builtins from m4.core.derived.materializer import get_derived_table_count @@ -61,7 +62,10 @@ class ListDatasetsTool: supported_datasets: frozenset[str] | None = None # Always available def invoke( - self, dataset: DatasetDefinition, params: ListDatasetsInput + self, + dataset: DatasetDefinition, + params: ListDatasetsInput, + context: M4ExecutionContext | None = None, ) -> dict[str, Any]: """List all available datasets with their status. @@ -71,9 +75,12 @@ def invoke( - backend: str - Backend type (duckdb or bigquery) - datasets: dict[str, dict] - Dataset availability info """ - active = get_active_dataset() + try: + active = get_active_dataset() + except Exception: + active = context.dataset.name if context else None availability = detect_available_local_datasets() - backend_name = get_active_backend() + backend_name = context.backend_name if context else get_active_backend() datasets_info: dict[str, dict] = {} @@ -135,7 +142,10 @@ class SetDatasetTool: supported_datasets: frozenset[str] | None = None # Always available def invoke( - self, dataset: DatasetDefinition, params: SetDatasetInput + self, + dataset: DatasetDefinition, + params: SetDatasetInput, + context: M4ExecutionContext | None = None, ) -> dict[str, Any]: """Switch to a different dataset. @@ -152,7 +162,7 @@ def invoke( """ dataset_name = params.dataset_name.lower() availability = detect_available_local_datasets() - backend_name = get_active_backend() + backend_name = context.backend_name if context else get_active_backend() if dataset_name not in availability: supported = ", ".join(availability.keys()) diff --git a/src/m4/core/tools/notes.py b/src/m4/core/tools/notes.py index 5e46e592..bec574f8 100644 --- a/src/m4/core/tools/notes.py +++ b/src/m4/core/tools/notes.py @@ -20,6 +20,7 @@ import pandas as pd from m4.core.backends import get_backend +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition, Modality from m4.core.exceptions import QueryError from m4.core.tools.base import ToolInput @@ -83,7 +84,10 @@ class SearchNotesTool: supported_datasets: frozenset[str] | None = None def invoke( - self, dataset: DatasetDefinition, params: SearchNotesInput + self, + dataset: DatasetDefinition, + params: SearchNotesInput, + context: M4ExecutionContext | None = None, ) -> dict[str, Any]: """Search notes and return snippets around matches. @@ -97,7 +101,7 @@ def invoke( Raises: QueryError: If note_type is invalid """ - backend = get_backend() + backend = context.backend if context else get_backend() # Determine which tables to search tables_to_search = self._get_tables_for_type(params.note_type) @@ -134,14 +138,22 @@ def invoke( LIMIT {params.limit} """ - result = backend.execute_query(sql, dataset) + result = ( + backend.execute_query(sql, dataset, context) + if context + else backend.execute_query(sql, dataset) + ) if result.success and result.dataframe is not None: results[table] = result.dataframe elif result.error: errors.append(f"{table}: {result.error}") response: dict[str, Any] = { - "backend_info": backend.get_backend_info(dataset), + "backend_info": ( + backend.get_backend_info(dataset, context) + if context + else backend.get_backend_info(dataset) + ), "query": params.query, "snippet_length": params.snippet_length, "results": results, @@ -191,7 +203,10 @@ class GetNoteTool: supported_datasets: frozenset[str] | None = None def invoke( - self, dataset: DatasetDefinition, params: GetNoteInput + self, + dataset: DatasetDefinition, + params: GetNoteInput, + context: M4ExecutionContext | None = None, ) -> dict[str, Any]: """Retrieve a single note by ID. @@ -207,7 +222,7 @@ def invoke( Raises: QueryError: If note not found """ - backend = get_backend() + backend = context.backend if context else get_backend() # Note IDs contain the note type (e.g., "10000032_DS-1" for discharge) note_id = params.note_id.replace("'", "''") @@ -226,7 +241,11 @@ def invoke( LIMIT 1 """ - result = backend.execute_query(sql, dataset) + result = ( + backend.execute_query(sql, dataset, context) + if context + else backend.execute_query(sql, dataset) + ) if result.error: errors.append(f"{table}: {result.error}") continue @@ -246,7 +265,11 @@ def invoke( truncated = True return { - "backend_info": backend.get_backend_info(dataset), + "backend_info": ( + backend.get_backend_info(dataset, context) + if context + else backend.get_backend_info(dataset) + ), "note_id": str(row["note_id"]), "subject_id": int(row["subject_id"]), "text": text, @@ -294,7 +317,10 @@ class ListPatientNotesTool: supported_datasets: frozenset[str] | None = None def invoke( - self, dataset: DatasetDefinition, params: ListPatientNotesInput + self, + dataset: DatasetDefinition, + params: ListPatientNotesInput, + context: M4ExecutionContext | None = None, ) -> dict[str, Any]: """List notes for a patient without returning full text. @@ -307,7 +333,7 @@ def invoke( Raises: QueryError: If note_type is invalid """ - backend = get_backend() + backend = context.backend if context else get_backend() tables_to_query = self._get_tables_for_type(params.note_type) @@ -334,14 +360,22 @@ def invoke( LIMIT {params.limit} """ - result = backend.execute_query(sql, dataset) + result = ( + backend.execute_query(sql, dataset, context) + if context + else backend.execute_query(sql, dataset) + ) if result.success and result.dataframe is not None: notes[table] = result.dataframe elif result.error: errors.append(f"{table}: {result.error}") response: dict[str, Any] = { - "backend_info": backend.get_backend_info(dataset), + "backend_info": ( + backend.get_backend_info(dataset, context) + if context + else backend.get_backend_info(dataset) + ), "subject_id": params.subject_id, "notes": notes, } diff --git a/src/m4/core/tools/tabular.py b/src/m4/core/tools/tabular.py index f6bc8181..5fdbae7b 100644 --- a/src/m4/core/tools/tabular.py +++ b/src/m4/core/tools/tabular.py @@ -19,6 +19,7 @@ import pandas as pd from m4.core.backends import get_backend +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition, Modality from m4.core.exceptions import QueryError, SecurityError from m4.core.tools.base import ToolInput @@ -71,7 +72,10 @@ class GetDatabaseSchemaTool: supported_datasets: frozenset[str] | None = None def invoke( - self, dataset: DatasetDefinition, params: GetDatabaseSchemaInput + self, + dataset: DatasetDefinition, + params: GetDatabaseSchemaInput, + context: M4ExecutionContext | None = None, ) -> dict[str, Any]: """List available tables using the backend. @@ -80,11 +84,19 @@ def invoke( - backend_info: str - Backend description - tables: list[str] - List of table names """ - backend = get_backend() - tables = backend.get_table_list(dataset) + backend = context.backend if context else get_backend() + tables = ( + backend.get_table_list(dataset, context) + if context + else backend.get_table_list(dataset) + ) return { - "backend_info": backend.get_backend_info(dataset), + "backend_info": ( + backend.get_backend_info(dataset, context) + if context + else backend.get_backend_info(dataset) + ), "tables": tables, } @@ -115,7 +127,10 @@ class GetTableInfoTool: supported_datasets: frozenset[str] | None = None def invoke( - self, dataset: DatasetDefinition, params: GetTableInfoInput + self, + dataset: DatasetDefinition, + params: GetTableInfoInput, + context: M4ExecutionContext | None = None, ) -> dict[str, Any]: """Get table structure and sample data using the backend. @@ -129,19 +144,27 @@ def invoke( Raises: QueryError: If table doesn't exist or query fails """ - backend = get_backend() + backend = context.backend if context else get_backend() # Validate table name if not validate_table_name(params.table_name): raise QueryError(f"Invalid table name '{params.table_name}'") # Get table schema - schema_result = backend.get_table_info(params.table_name, dataset) + schema_result = ( + backend.get_table_info(params.table_name, dataset, context) + if context + else backend.get_table_info(params.table_name, dataset) + ) if not schema_result.success: raise QueryError(schema_result.error or "Failed to get table info") result = { - "backend_info": backend.get_backend_info(dataset), + "backend_info": ( + backend.get_backend_info(dataset, context) + if context + else backend.get_backend_info(dataset) + ), "table_name": params.table_name, "schema": schema_result.dataframe, "sample": None, @@ -149,7 +172,13 @@ def invoke( # Get sample data if requested if params.show_sample: - sample_result = backend.get_sample_data(params.table_name, dataset, limit=3) + sample_result = ( + backend.get_sample_data( + params.table_name, dataset, limit=3, context=context + ) + if context + else backend.get_sample_data(params.table_name, dataset, limit=3) + ) if sample_result.success: result["sample"] = sample_result.dataframe @@ -182,7 +211,10 @@ class ExecuteQueryTool: supported_datasets: frozenset[str] | None = None def invoke( - self, dataset: DatasetDefinition, params: ExecuteQueryInput + self, + dataset: DatasetDefinition, + params: ExecuteQueryInput, + context: M4ExecutionContext | None = None, ) -> pd.DataFrame: """Execute a SQL query with safety validation. @@ -198,8 +230,12 @@ def invoke( if not safe: raise SecurityError(msg, query=params.sql_query) - backend = get_backend() - result = backend.execute_query(params.sql_query, dataset) + backend = context.backend if context else get_backend() + result = ( + backend.execute_query(params.sql_query, dataset, context) + if context + else backend.execute_query(params.sql_query, dataset) + ) if not result.success: raise QueryError(result.error or "Unknown error", sql=params.sql_query) diff --git a/src/m4/data_io.py b/src/m4/data_io.py index 28b33bb8..f8821c4e 100644 --- a/src/m4/data_io.py +++ b/src/m4/data_io.py @@ -1,6 +1,7 @@ import os import time from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass from pathlib import Path from urllib.parse import urljoin, urlparse @@ -21,6 +22,8 @@ success, ) from m4.core.datasets import DatasetRegistry +from m4.services.events import EventReporter, get_event_reporter +from m4.services.redaction import redact_sensitive, register_sensitive_value ######################################################## # Download functionality @@ -32,19 +35,122 @@ ) +@dataclass(frozen=True) +class PhysioNetCredentials: + username: str + password: str + + @classmethod + def from_json_file(cls, path: Path) -> "PhysioNetCredentials": + import json + + payload = json.loads(path.read_text()) + username = payload.get("username") or payload.get("user") + password = payload.get("password") + if not isinstance(username, str) or not isinstance(password, str): + raise ValueError( + "PhysioNet credentials file must contain string username and password fields." + ) + register_sensitive_value(password) + return cls(username=username, password=password) + + +class DatasetDownloadError(RuntimeError): + def __init__(self, code: str, message: str): + super().__init__(redact_sensitive(message)) + self.code = code + self.message = str(self) + + +def _download_error_for_response( + response: requests.Response, url: str +) -> DatasetDownloadError: + if response.status_code == 401: + return DatasetDownloadError( + "physionet_auth_failed", + f"PhysioNet authentication failed while accessing {url}.", + ) + if response.status_code == 403: + return DatasetDownloadError( + "physionet_access_forbidden", + f"PhysioNet access is forbidden for {url}. Confirm DUA access.", + ) + return DatasetDownloadError( + "download_network_failed", + f"HTTP {response.status_code} while downloading {url}: {response.reason}", + ) + + +def _remote_content_length( + url: str, session: requests.Session +) -> tuple[int | None, bool]: + try: + response = session.head(url, allow_redirects=True, timeout=30) + except requests.exceptions.RequestException: + return None, False + if response.status_code in {401, 403}: + raise _download_error_for_response(response, url) + if not (200 <= response.status_code < 300): + return None, False + content_length = response.headers.get("content-length") + accept_ranges = response.headers.get("accept-ranges", "").lower() == "bytes" + return ( + int(content_length) if content_length and content_length.isdigit() else None + ), accept_ranges + + def _download_single_file( url: str, target_filepath: Path, session: requests.Session, progress=None, task_id=None, + event_reporter: EventReporter | None = None, ) -> bool: """Downloads a single file with progress tracking.""" + reporter = get_event_reporter(event_reporter) logger.debug(f"Attempting to download {url} to {target_filepath}...") + part_path = target_filepath.with_name(f"{target_filepath.name}.part") try: - response = session.get(url, stream=True, timeout=60) - response.raise_for_status() - total_size = int(response.headers.get("content-length", 0)) + remote_size, range_supported = _remote_content_length(url, session) + if remote_size is not None and target_filepath.exists(): + if target_filepath.stat().st_size == remote_size: + reporter.emit( + "download_file_skipped", + url=url, + path=str(target_filepath), + bytes_total=remote_size, + reason="complete", + ) + return True + + resume_from = part_path.stat().st_size if part_path.exists() else 0 + headers = {} + mode = "wb" + if resume_from and range_supported: + headers["Range"] = f"bytes={resume_from}-" + mode = "ab" + + response = session.get(url, stream=True, timeout=60, headers=headers) + if response.status_code in {401, 403}: + raise _download_error_for_response(response, url) + if ( + response.status_code == 416 + and remote_size is not None + and resume_from >= remote_size + ): + part_path.replace(target_filepath) + return True + if response.status_code == 200 and headers.get("Range"): + resume_from = 0 + mode = "wb" + if not (200 <= response.status_code < 300): + raise _download_error_for_response(response, url) + + response_size = int(response.headers.get("content-length", 0)) + total_size = remote_size or ( + resume_from + response_size if response_size else 0 + ) file_display_name = target_filepath.name target_filepath.parent.mkdir(parents=True, exist_ok=True) @@ -53,27 +159,63 @@ def _download_single_file( if progress and task_id is not None: progress.update(task_id, total=total_size, description=file_display_name) - with open(target_filepath, "wb") as f: + reporter.emit( + "download_file_started", + url=url, + path=str(target_filepath), + bytes_downloaded=resume_from, + bytes_total=total_size or None, + ) + downloaded = resume_from + last_emit = 0.0 + with open(part_path, mode) as f: for chunk in response.iter_content(chunk_size=8192): if chunk: f.write(chunk) + downloaded += len(chunk) if progress and task_id is not None: progress.update(task_id, advance=len(chunk)) - + now = time.monotonic() + if now - last_emit >= 0.15: + reporter.emit( + "download_file_progress", + path=str(target_filepath), + bytes_downloaded=downloaded, + bytes_total=total_size or None, + ) + last_emit = now + + part_path.replace(target_filepath) + reporter.emit( + "download_file_completed", + path=str(target_filepath), + bytes_downloaded=downloaded, + bytes_total=total_size or downloaded, + ) logger.info(f"Successfully downloaded: {file_display_name}") return True + except KeyboardInterrupt as exc: + raise DatasetDownloadError( + "download_interrupted", "Download interrupted." + ) from exc + except DatasetDownloadError: + raise except requests.exceptions.HTTPError as e: - status = e.response.status_code - if status == 404: - logger.error(f"Download failed (404 Not Found): {url}.") - else: - logger.error(f"HTTP error {status} downloading {url}: {e.response.reason}") + raise _download_error_for_response(e.response, url) from e except requests.exceptions.Timeout: - logger.error(f"Timeout occurred while downloading {url}.") + raise DatasetDownloadError( + "download_network_failed", f"Timeout occurred while downloading {url}." + ) except requests.exceptions.RequestException as e: - logger.error(f"A network or request error occurred downloading {url}: {e}") + raise DatasetDownloadError( + "download_network_failed", + f"A network or request error occurred downloading {url}: {e}", + ) except OSError as e: - logger.error(f"File system error writing {target_filepath}: {e}") + raise DatasetDownloadError( + "download_filesystem_failed", + f"File system error writing {target_filepath}: {e}", + ) # If download failed, attempt to remove partially downloaded file if target_filepath.exists(): @@ -92,7 +234,10 @@ def _scrape_urls_from_html_page( logger.debug(f"Scraping for '{file_suffix}' links on page: {page_url}") try: page_response = session.get(page_url, timeout=30) - page_response.raise_for_status() + if page_response.status_code in {401, 403}: + raise _download_error_for_response(page_response, page_url) + if not (200 <= page_response.status_code < 300): + raise _download_error_for_response(page_response, page_url) soup = BeautifulSoup(page_response.content, "html.parser") for link_tag in soup.find_all("a", href=True): href_path = link_tag["href"] @@ -104,15 +249,26 @@ def _scrape_urls_from_html_page( ): absolute_url = urljoin(page_url, href_path) found_urls.append(absolute_url) + except DatasetDownloadError: + raise except requests.exceptions.RequestException as e: - logger.error(f"Could not access or parse page {page_url} for scraping: {e}") + raise DatasetDownloadError( + "download_network_failed", + f"Could not access or parse page {page_url} for scraping: {e}", + ) from e return found_urls def _download_dataset_files( - dataset_name: str, dataset_config: dict, raw_files_root_dir: Path + dataset_name: str, + dataset_config: dict, + raw_files_root_dir: Path, + *, + credentials: PhysioNetCredentials | None = None, + event_reporter: EventReporter | None = None, ) -> bool: """Downloads all relevant files for a dataset based on its configuration.""" + reporter = get_event_reporter(event_reporter) base_listing_url = dataset_config["file_listing_url"] subdirs_to_scan = dataset_config.get("subdirectories_to_scan", []) @@ -121,6 +277,8 @@ def _download_dataset_files( ) session = requests.Session() session.headers.update({"User-Agent": COMMON_USER_AGENT}) + if credentials: + session.auth = (credentials.username, credentials.password) all_files_to_process = [] # List of (url, local_target_path) @@ -137,7 +295,14 @@ def _download_dataset_files( for subdir_name, listing_url in scan_targets: logger.info(f"Scanning for CSVs: {listing_url}") + reporter.emit("download_listing_started", dataset=dataset_name, url=listing_url) csv_urls_in_subdir = _scrape_urls_from_html_page(listing_url, session) + reporter.emit( + "download_listing_completed", + dataset=dataset_name, + url=listing_url, + file_count=len(csv_urls_in_subdir), + ) if not csv_urls_in_subdir: logger.warning(f"No .csv.gz files found in location: {listing_url}") @@ -179,8 +344,10 @@ def _download_dataset_files( all_files_to_process.append((file_url, local_target_path)) if not all_files_to_process: - logger.error(f"No '.csv.gz' download links found for dataset '{dataset_name}'.") - return False + raise DatasetDownloadError( + "raw_files_missing", + f"No '.csv.gz' download links found for dataset '{dataset_name}'.", + ) # Deduplicate and sort for consistent processing order unique_files_to_process = sorted( @@ -189,6 +356,7 @@ def _download_dataset_files( total_files = len(unique_files_to_process) info(f"Found {total_files} files to download") + reporter.emit("download_started", dataset=dataset_name, file_count=total_files) downloaded_count = 0 with create_download_progress() as progress: @@ -201,7 +369,12 @@ def _download_dataset_files( for file_url, target_filepath in unique_files_to_process: if not _download_single_file( - file_url, target_filepath, session, progress, file_task + file_url, + target_filepath, + session, + progress, + file_task, + event_reporter=reporter, ): logger.error( f"Critical download failed for '{target_filepath.name}'. " @@ -218,10 +391,19 @@ def _download_dataset_files( progress.reset(file_task) # Success only if all identified files were downloaded + reporter.emit( + "download_completed", dataset=dataset_name, file_count=downloaded_count + ) return downloaded_count == len(unique_files_to_process) -def download_dataset(dataset_name: str, output_root: Path) -> bool: +def download_dataset( + dataset_name: str, + output_root: Path, + *, + credentials: PhysioNetCredentials | None = None, + event_reporter: EventReporter | None = None, +) -> bool: """ Public wrapper to download a supported dataset's CSV files. - Currently intended for 'mimic-iv-demo' (public demo); extendable for others. @@ -229,23 +411,25 @@ def download_dataset(dataset_name: str, output_root: Path) -> bool: """ ds = DatasetRegistry.get(dataset_name.lower()) if not ds: - logger.error(f"Unsupported dataset: {dataset_name}") - return False + raise DatasetDownloadError( + "dataset_not_found", f"Unsupported dataset: {dataset_name}" + ) # Prevent accidental scraping of credentialed datasets - if ds.requires_authentication: - logger.error( - f"Dataset '{dataset_name}' requires authentication and cannot be auto-downloaded. " - "Please download files manually." + if ds.requires_authentication and credentials is None: + raise DatasetDownloadError( + "missing_credentials", + ( + f"Dataset '{dataset_name}' requires PhysioNet credentials. " + "Provide --physionet-credentials-file." + ), ) - return False if not ds.file_listing_url: - logger.error( - f"Dataset '{dataset_name}' does not have a configured listing URL. " - "This version only supports public demo download." + raise DatasetDownloadError( + "raw_files_missing", + f"Dataset '{dataset_name}' does not have a configured listing URL.", ) - return False output_root.mkdir(parents=True, exist_ok=True) @@ -254,7 +438,13 @@ def download_dataset(dataset_name: str, output_root: Path) -> bool: "file_listing_url": ds.file_listing_url, "subdirectories_to_scan": ds.subdirectories_to_scan, } - return _download_dataset_files(dataset_name, dataset_config, output_root) + return _download_dataset_files( + dataset_name, + dataset_config, + output_root, + credentials=credentials, + event_reporter=event_reporter, + ) ######################################################## @@ -262,7 +452,9 @@ def download_dataset(dataset_name: str, output_root: Path) -> bool: ######################################################## -def _csv_to_parquet_all(src_root: Path, parquet_root: Path) -> bool: +def _csv_to_parquet_all( + src_root: Path, parquet_root: Path, event_reporter: EventReporter | None = None +) -> bool: """ Convert all CSV files in the source directory to Parquet files. - Streams via DuckDB COPY to keep memory low @@ -277,6 +469,7 @@ def _csv_to_parquet_all(src_root: Path, parquet_root: Path) -> bool: if not csv_files: logger.error(f"No CSV files found in {src_root}") return False + reporter = get_event_reporter(event_reporter) # Optional: process small files first so progress moves smoothly try: @@ -326,6 +519,12 @@ def _convert_one(csv_gz: Path) -> tuple[Path | None, float, str]: logger.info( f"Converting {total_files} CSV files to Parquet using {max_workers} workers..." ) + reporter.emit( + "conversion_started", + source=str(src_root), + destination=str(parquet_root), + file_count=total_files, + ) console.print() with create_task_progress() as progress: @@ -337,6 +536,7 @@ def _convert_one(csv_gz: Path) -> tuple[Path | None, float, str]: futures = {ex.submit(_convert_one, f): f for f in csv_files} for fut in as_completed(futures): + csv_file = futures[fut] try: result_path, _, filename = fut.result() if result_path is not None: @@ -348,19 +548,34 @@ def _convert_one(csv_gz: Path) -> tuple[Path | None, float, str]: description=f"Converted {filename} ({max_workers} workers)", ) logger.debug(f"Converted: {filename}") + reporter.emit( + "conversion_file_completed", + path=str(result_path), + source=str(csv_file), + completed=completed, + file_count=total_files, + ) except Exception as e: - csv_file = futures[fut] logger.error(f"Parquet conversion failed for {csv_file}: {e}") ex.shutdown(cancel_futures=True) return False elapsed_time = time.time() - start_time success(f"Converted {len(parquet_paths)} files in {elapsed_time:.1f}s") + reporter.emit( + "conversion_completed", + destination=str(parquet_root), + file_count=len(parquet_paths), + elapsed_seconds=elapsed_time, + ) return True def convert_csv_to_parquet( - dataset_name: str, csv_root: Path, parquet_root: Path + dataset_name: str, + csv_root: Path, + parquet_root: Path, + event_reporter: EventReporter | None = None, ) -> bool: """ Public wrapper to convert CSV.gz files to Parquet for a dataset. @@ -371,7 +586,7 @@ def convert_csv_to_parquet( logger.error(f"CSV root not found: {csv_root}") return False parquet_root.mkdir(parents=True, exist_ok=True) - return _csv_to_parquet_all(csv_root, parquet_root) + return _csv_to_parquet_all(csv_root, parquet_root, event_reporter=event_reporter) ######################################################## @@ -379,7 +594,11 @@ def convert_csv_to_parquet( ######################################################## -def init_duckdb_from_parquet(dataset_name: str, db_target_path: Path) -> bool: +def init_duckdb_from_parquet( + dataset_name: str, + db_target_path: Path, + event_reporter: EventReporter | None = None, +) -> bool: """ Initialize or refresh a DuckDB for the dataset by creating views over Parquet. @@ -403,13 +622,16 @@ def init_duckdb_from_parquet(dataset_name: str, db_target_path: Path) -> bool: f"Creating or refreshing views in {db_target_path} for Parquet under {parquet_root}" ) mapping = ds.schema_mapping if ds.schema_mapping else None - return _create_duckdb_with_views(db_target_path, parquet_root, mapping) + return _create_duckdb_with_views( + db_target_path, parquet_root, mapping, event_reporter=event_reporter + ) def _create_duckdb_with_views( db_path: Path, parquet_root: Path, schema_mapping: dict[str, str] | None = None, + event_reporter: EventReporter | None = None, ) -> bool: """ Create a DuckDB database and define one view per Parquet file. @@ -437,6 +659,7 @@ def _create_duckdb_with_views( raise try: + reporter = get_event_reporter(event_reporter) # Find all parquet files parquet_files = list(parquet_root.rglob("*.parquet")) if not parquet_files: @@ -456,6 +679,12 @@ def _create_duckdb_with_views( con.execute(f'CREATE SCHEMA IF NOT EXISTS "{schema_name}"') logger.info(f"Creating {len(parquet_files)} views in DuckDB...") + reporter.emit( + "duckdb_init_started", + database=str(db_path), + parquet_root=str(parquet_root), + file_count=len(parquet_files), + ) start_time = time.time() created = 0 @@ -515,6 +744,13 @@ def _create_duckdb_with_views( task, advance=1, description=f"Created view: {view_name}" ) logger.debug(f"Created view: {view_name}") + reporter.emit( + "duckdb_view_created", + database=str(db_path), + view=view_name, + completed=created, + file_count=len(parquet_files), + ) except Exception as e: logger.error(f"Failed to create view {view_name} from {pq}: {e}") raise @@ -522,6 +758,12 @@ def _create_duckdb_with_views( con.commit() elapsed_time = time.time() - start_time success(f"Created {created} views in {elapsed_time:.1f}s") + reporter.emit( + "duckdb_init_completed", + database=str(db_path), + view_count=created, + elapsed_seconds=elapsed_time, + ) # List all created views for verification views_result = con.execute( diff --git a/src/m4/mcp_server.py b/src/m4/mcp_server.py index 1a6de543..ba757365 100644 --- a/src/m4/mcp_server.py +++ b/src/m4/mcp_server.py @@ -29,25 +29,14 @@ from m4.apps import init_apps from m4.apps.cohort_builder import RESOURCE_URI as COHORT_BUILDER_URI from m4.apps.cohort_builder import get_ui_html -from m4.apps.cohort_builder.query_builder import QueryCohortInput -from m4.apps.cohort_builder.tool import CohortBuilderInput from m4.auth import init_oauth2, require_oauth2 +from m4.client import M4Client from m4.core.datasets import DatasetRegistry from m4.core.exceptions import M4Error from m4.core.serialization import serialize_for_mcp -from m4.core.telemetry import invoke_tracked, set_interface -from m4.core.tools import ToolRegistry, ToolSelector, init_tools -from m4.core.tools.management import ListDatasetsInput, SetDatasetInput -from m4.core.tools.notes import ( - GetNoteInput, - ListPatientNotesInput, - SearchNotesInput, -) -from m4.core.tools.tabular import ( - ExecuteQueryInput, - GetDatabaseSchemaInput, - GetTableInfoInput, -) +from m4.core.telemetry import set_interface +from m4.core.tools import ToolSelector, init_tools +from m4.core.tools.management import SetDatasetInput # Create FastMCP server instance mcp = FastMCP("m4") @@ -74,6 +63,7 @@ "list_patient_notes", "cohort_builder", "query_cohort", + "capabilities", } ) @@ -270,6 +260,20 @@ def _serialize_list_patient_notes_result(result: dict[str, Any]) -> str: # ========================================== +@mcp.resource("m4://capabilities", mime_type="application/json") +def capabilities_resource() -> str: + """Return the M4 capability manifest as JSON.""" + client = M4Client.from_active(interface="mcp", allow_missing_dataset=True) + return json.dumps(client.capabilities(), indent=2) + + +@mcp.tool() +def capabilities() -> str: + """Return the M4 capability manifest as JSON.""" + client = M4Client.from_active(interface="mcp", allow_missing_dataset=True) + return json.dumps(client.capabilities(), indent=2) + + @mcp.tool() def list_datasets() -> str: """📋 List all available datasets and their status. @@ -279,9 +283,8 @@ def list_datasets() -> str: and showing availability of local database and BigQuery support. """ try: - tool = ToolRegistry.get("list_datasets") - dataset = DatasetRegistry.get_active() - result = invoke_tracked(tool, dataset, ListDatasetsInput()) + client = M4Client.from_active(interface="mcp") + result = client.dataset_status() return _serialize_datasets_result(result) except M4Error as e: return f"**Error:** {e}" @@ -301,10 +304,9 @@ def set_dataset(dataset_name: str) -> str: # Check if target dataset exists before switching target_dataset_def = DatasetRegistry.get(dataset_name.lower()) - tool = ToolRegistry.get("set_dataset") - dataset = DatasetRegistry.get_active() - result = invoke_tracked( - tool, dataset, SetDatasetInput(dataset_name=dataset_name) + client = M4Client.from_active(interface="mcp") + result = client.invoke_tool( + "set_dataset", SetDatasetInput(dataset_name=dataset_name) ) output = _serialize_set_dataset_result(result) @@ -330,7 +332,8 @@ def get_database_schema() -> str: List of all available tables in the database with current backend info. """ try: - dataset = DatasetRegistry.get_active() + client = M4Client.from_active(interface="mcp") + dataset = client.dataset # Proactive capability check compat_result = _tool_selector.check_compatibility( @@ -339,8 +342,7 @@ def get_database_schema() -> str: if not compat_result.compatible: return compat_result.error_message - tool = ToolRegistry.get("get_database_schema") - result = invoke_tracked(tool, dataset, GetDatabaseSchemaInput()) + result = client.schema() return _serialize_schema_result(result) except M4Error as e: return f"**Error:** {e}" @@ -361,18 +363,17 @@ def get_table_info(table_name: str, show_sample: bool = True) -> str: Table structure with column names, types, and sample data. """ try: - dataset = DatasetRegistry.get_active() + client = M4Client.from_active(interface="mcp") + dataset = client.dataset # Proactive capability check compat_result = _tool_selector.check_compatibility("get_table_info", dataset) if not compat_result.compatible: return compat_result.error_message - tool = ToolRegistry.get("get_table_info") - result = invoke_tracked( - tool, - dataset, - GetTableInfoInput(table_name=table_name, show_sample=show_sample), + result = client.table_info( + table_name=table_name, + show_sample=show_sample, ) return _serialize_table_info_result(result) except M4Error as e: @@ -396,15 +397,15 @@ def execute_query(sql_query: str) -> str: Query results or helpful error messages. """ try: - dataset = DatasetRegistry.get_active() + client = M4Client.from_active(interface="mcp") + dataset = client.dataset # Proactive capability check compat_result = _tool_selector.check_compatibility("execute_query", dataset) if not compat_result.compatible: return compat_result.error_message - tool = ToolRegistry.get("execute_query") - result = invoke_tracked(tool, dataset, ExecuteQueryInput(sql_query=sql_query)) + result = client.query(sql_query) # Result is a DataFrame - serialize it return serialize_for_mcp(result) except M4Error as e: @@ -441,22 +442,18 @@ def search_notes( Matching snippets with note IDs for follow-up retrieval. """ try: - dataset = DatasetRegistry.get_active() + client = M4Client.from_active(interface="mcp") + dataset = client.dataset compat_result = _tool_selector.check_compatibility("search_notes", dataset) if not compat_result.compatible: return compat_result.error_message - tool = ToolRegistry.get("search_notes") - result = invoke_tracked( - tool, - dataset, - SearchNotesInput( - query=query, - note_type=note_type, - limit=limit, - snippet_length=snippet_length, - ), + result = client.search_notes( + query=query, + note_type=note_type, + limit=limit, + snippet_length=snippet_length, ) return _serialize_search_notes_result(result) except M4Error as e: @@ -480,17 +477,16 @@ def get_note(note_id: str, max_length: int | None = None) -> str: Full note text, or truncated version if max_length specified. """ try: - dataset = DatasetRegistry.get_active() + client = M4Client.from_active(interface="mcp") + dataset = client.dataset compat_result = _tool_selector.check_compatibility("get_note", dataset) if not compat_result.compatible: return compat_result.error_message - tool = ToolRegistry.get("get_note") - result = invoke_tracked( - tool, - dataset, - GetNoteInput(note_id=note_id, max_length=max_length), + result = client.get_note( + note_id=note_id, + max_length=max_length, ) return _serialize_get_note_result(result) except M4Error as e: @@ -521,7 +517,8 @@ def list_patient_notes( List of available notes with metadata for the patient. """ try: - dataset = DatasetRegistry.get_active() + client = M4Client.from_active(interface="mcp") + dataset = client.dataset compat_result = _tool_selector.check_compatibility( "list_patient_notes", dataset @@ -529,15 +526,10 @@ def list_patient_notes( if not compat_result.compatible: return compat_result.error_message - tool = ToolRegistry.get("list_patient_notes") - result = invoke_tracked( - tool, - dataset, - ListPatientNotesInput( - subject_id=subject_id, - note_type=note_type, - limit=limit, - ), + result = client.list_patient_notes( + subject_id=subject_id, + note_type=note_type, + limit=limit, ) return _serialize_list_patient_notes_result(result) except M4Error as e: @@ -571,15 +563,15 @@ def cohort_builder() -> str: interactive cohort builder interface. """ try: - dataset = DatasetRegistry.get_active() + client = M4Client.from_active(interface="mcp") + dataset = client.dataset # Proactive capability check compat_result = _tool_selector.check_compatibility("cohort_builder", dataset) if not compat_result.compatible: return compat_result.error_message - tool = ToolRegistry.get("cohort_builder") - result = invoke_tracked(tool, dataset, CohortBuilderInput()) + result = client.cohort_builder() return serialize_for_mcp(result) except M4Error as e: return f"**Error:** {e}" @@ -614,7 +606,8 @@ def query_cohort( JSON with patient_count, admission_count, demographics, and SQL. """ try: - dataset = DatasetRegistry.get_active() + client = M4Client.from_active(interface="mcp") + dataset = client.dataset # Proactive capability check compat_result = _tool_selector.check_compatibility("query_cohort", dataset) @@ -622,19 +615,14 @@ def query_cohort( # Return JSON error for UI compatibility return json.dumps({"error": compat_result.error_message}) - tool = ToolRegistry.get("query_cohort") - result = invoke_tracked( - tool, - dataset, - QueryCohortInput( - age_min=age_min, - age_max=age_max, - gender=gender, - icd_codes=icd_codes, - icd_match_all=icd_match_all, - has_icu_stay=has_icu_stay, - in_hospital_mortality=in_hospital_mortality, - ), + result = client.query_cohort( + age_min=age_min, + age_max=age_max, + gender=gender, + icd_codes=icd_codes, + icd_match_all=icd_match_all, + has_icu_stay=has_icu_stay, + in_hospital_mortality=in_hospital_mortality, ) # Return JSON directly for MCP App UI compatibility return json.dumps(result) diff --git a/src/m4/services/capabilities.py b/src/m4/services/capabilities.py new file mode 100644 index 00000000..596b1010 --- /dev/null +++ b/src/m4/services/capabilities.py @@ -0,0 +1,275 @@ +from __future__ import annotations + +from dataclasses import MISSING, fields, is_dataclass +from pathlib import Path +from typing import Any + +from m4.config import ( + ensure_custom_datasets_loaded, + get_active_backend, + get_active_dataset, + get_bigquery_project_id, +) +from m4.core.datasets import DatasetDefinition, DatasetRegistry +from m4.core.derived.builtins import has_derived_support, list_builtins +from m4.core.tools import ToolRegistry, init_tools + +CAPABILITIES_SCHEMA_VERSION = 1 + + +def _field_default(field: Any) -> Any: + if field.default is not MISSING: + return field.default + if field.default_factory is not MISSING: # type: ignore[attr-defined] + return field.default_factory() # type: ignore[misc] + return None + + +def _input_fields(input_model: type) -> list[dict[str, Any]]: + if not is_dataclass(input_model): + return [] + result = [] + for field in fields(input_model): + result.append( + { + "name": field.name, + "type": str(field.type), + "required": ( + field.default is MISSING and field.default_factory is MISSING # type: ignore[attr-defined] + ), + "default": _field_default(field), + } + ) + return result + + +def _dataset_payload(ds: DatasetDefinition) -> dict[str, Any]: + return { + "name": ds.name, + "description": ds.description, + "version": ds.version, + "requires_authentication": ds.requires_authentication, + "modalities": sorted(modality.name for modality in ds.modalities), + "dataset_page_url": ds.dataset_page_url, + "dua_url": ds.dua_url, + "file_listing_url": ds.file_listing_url, + "bigquery": { + "available": bool(ds.bigquery_dataset_ids), + "project_id": ds.bigquery_project_id, + "dataset_ids": list(ds.bigquery_dataset_ids), + "schema_mapping": dict(ds.bigquery_schema_mapping), + "access_url": ds.bigquery_access_url, + }, + "verification_table": ds.primary_verification_table, + "schema_mapping": dict(ds.schema_mapping), + "expected_local_layout": { + "recommended_raw_root": ds.recommended_local_target_root, + "raw_subdirectories": list(ds.expected_raw_subdirectories), + "parquet_root": f"m4_data/parquet/{ds.name}", + "duckdb_filename": ds.default_duckdb_filename, + }, + "related_datasets": dict(ds.related_datasets), + } + + +def _tool_payloads() -> list[dict[str, Any]]: + init_tools() + datasets = DatasetRegistry.list_all() + result = [] + for tool in ToolRegistry.list_all(): + required_modalities = sorted( + modality.name for modality in getattr(tool, "required_modalities", []) + ) + compatible = [] + for ds in datasets: + try: + if tool.is_compatible(ds): + compatible.append(ds.name) + except Exception: + continue + result.append( + { + "name": tool.name, + "description": tool.description, + "input_fields": _input_fields(tool.input_model), + "required_modalities": required_modalities, + "compatible_datasets": sorted(compatible), + "supported_datasets": ( + sorted(tool.supported_datasets) + if tool.supported_datasets is not None + else None + ), + } + ) + return sorted(result, key=lambda item: item["name"]) + + +def _skill_inventory() -> list[dict[str, Any]]: + skills_dir = Path(__file__).resolve().parents[1] / "skills" + if not skills_dir.exists(): + return [] + + inventory = [] + for skill_file in sorted(skills_dir.glob("*/*/SKILL.md")): + skill_dir = skill_file.parent + category = skill_dir.parent.name + description = "" + try: + for line in skill_file.read_text(encoding="utf-8").splitlines()[:40]: + if line.startswith("description:"): + description = line.split(":", 1)[1].strip() + break + except OSError: + pass + inventory.append( + { + "name": skill_dir.name, + "category": category, + "description": description, + "packaged": True, + } + ) + return inventory + + +def _derived_inventory() -> dict[str, Any]: + datasets: dict[str, Any] = {} + for ds in DatasetRegistry.list_all(): + if not has_derived_support(ds.name): + datasets[ds.name] = {"available": False, "tables": []} + continue + try: + tables = list_builtins(ds.name) + except Exception: + tables = [] + datasets[ds.name] = {"available": True, "tables": tables} + return datasets + + +def build_capabilities_manifest() -> dict[str, Any]: + """Return the stable M4 capability manifest.""" + ensure_custom_datasets_loaded() + try: + active_dataset = get_active_dataset() + except Exception: + active_dataset = None + + commands = [ + {"name": "capabilities", "flags": ["--json"], "mutates": False}, + {"name": "doctor", "flags": ["--json", "--paths"], "mutates": False}, + {"name": "status", "flags": ["--all", "--derived", "--json"], "mutates": False}, + { + "name": "schema", + "flags": ["--dataset", "--backend", "--json"], + "mutates": False, + }, + { + "name": "query", + "flags": ["--dataset", "--backend", "--sql", "--json"], + "mutates": False, + }, + { + "name": "download", + "flags": [ + "--target", + "--json", + "--command-only", + "--physionet-credentials-file", + "--events", + ], + "mutates": True, + }, + { + "name": "init", + "flags": [ + "--src", + "--db-path", + "--force", + "--json", + "--download", + "--physionet-credentials-file", + "--events", + ], + "mutates": True, + }, + {"name": "config", "flags": ["--backend", "--project-id"], "mutates": True}, + { + "name": "agent-env", + "flags": [ + "--mode", + "--dataset", + "--backend", + "--project-id", + "--json", + "--format", + "--paths", + ], + "mutates": False, + }, + { + "name": "setup-agent", + "flags": [ + "--mode", + "--client", + "--dataset", + "--backend", + "--project-id", + "--format", + "--apply", + ], + "mutates": False, + "mutates_with": ["--apply"], + }, + { + "name": "quickstart", + "flags": [ + "--workflow", + "--dataset", + "--backend", + "--project-id", + "--apply", + "--json", + ], + "mutates": False, + "mutates_with": ["--apply"], + }, + ] + + return { + "schema_version": CAPABILITIES_SCHEMA_VERSION, + "interfaces": { + "cli": {"entrypoint": "m4"}, + "python_api": {"function": "m4.get_capabilities"}, + "mcp": {"tool": "capabilities", "resource": "m4://capabilities"}, + "apps": ["cohort_builder"], + "output_formats": ["text", "json", "dotenv"], + }, + "runtime": { + "active_dataset": active_dataset, + "backend": get_active_backend(), + "bigquery_project_id_configured": bool(get_bigquery_project_id()), + }, + "commands": commands, + "tools": _tool_payloads(), + "datasets": [_dataset_payload(ds) for ds in DatasetRegistry.list_all()], + "limits": { + "query_row_limit_default": 100, + "path_redaction_default": True, + "supported_backends": ["duckdb", "bigquery"], + "conversion_env": [ + "M4_CONVERT_MAX_WORKERS", + "M4_DUCKDB_MEM", + "M4_DUCKDB_THREADS", + ], + }, + "concepts": { + "derived_tables": _derived_inventory(), + "skills": _skill_inventory(), + }, + "provenance_policy": { + "telemetry_destination": "M4_TELEMETRY_DIR or /telemetry", + "path_redaction": "Machine-facing output hides raw paths unless --paths or M4_PATH_DISCLOSURE=1 is used.", + "event_export_command": "m4 provenance export --json", + "non_phi_policy": "M4 telemetry is intended for operational provenance only and must not include PHI.", + }, + } diff --git a/src/m4/services/download.py b/src/m4/services/download.py new file mode 100644 index 00000000..a57aac4e --- /dev/null +++ b/src/m4/services/download.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +from pathlib import Path +from shlex import quote +from typing import Any +from urllib.parse import urlparse + +from m4.config import ensure_custom_datasets_loaded, resolve_runtime_context +from m4.console import console +from m4.core.datasets import DatasetDefinition, DatasetRegistry +from m4.data_io import DatasetDownloadError, PhysioNetCredentials, download_dataset +from m4.services.events import EventReporter, get_event_reporter +from m4.services.results import ( + ERROR_DATASET_NOT_FOUND, + ERROR_INVALID_OPTION, + CommandError, + CommandResult, +) + + +class _quiet_console: + def __enter__(self) -> None: + self.previous_quiet = console.quiet + console.quiet = True + + def __exit__(self, *args: object) -> None: + console.quiet = self.previous_quiet + + +def default_raw_root(dataset_name: str) -> Path: + ctx = resolve_runtime_context() + return ctx.data_dir / "raw_files" / dataset_name.lower() + + +def expected_raw_subdirectories(ds: DatasetDefinition) -> list[str]: + return list(ds.expected_raw_subdirectories or ds.subdirectories_to_scan) + + +def build_wget_command(ds: DatasetDefinition, target: Path) -> str: + if not ds.file_listing_url: + return "" + path_parts = [ + part for part in urlparse(ds.file_listing_url).path.split("/") if part + ] + cut_dirs = len(path_parts) + return ( + f"wget -r -N -c -np --cut-dirs={cut_dirs} -nH " + "--user YOUR_USERNAME --ask-password " + f"{quote(ds.file_listing_url)} -P {quote(str(target))}" + ) + + +def validate_raw_layout(dataset_name: str, root: Path) -> dict[str, Any]: + ensure_custom_datasets_loaded() + ds = DatasetRegistry.get(dataset_name.lower()) + warnings: list[str] = [] + errors: list[str] = [] + recovery: list[str] = [] + + if not ds: + return { + "ok": False, + "warnings": [], + "errors": [f"Unknown dataset: {dataset_name}"], + "csv_gz_count": 0, + "empty_csv_gz": [], + "recovery": [], + } + + if not root.exists(): + return { + "ok": False, + "warnings": [], + "errors": [f"Raw root does not exist: {root}"], + "csv_gz_count": 0, + "empty_csv_gz": [], + "recovery": ["Run the generated wget command, then retry m4 download."], + } + + nested_markers = [ + root / "physionet.org" / "files", + root / "files" / "mimiciv", + root / "files" / "mimic-iv-note", + root / "files" / "eicu-crd", + ] + if any(path.exists() for path in nested_markers): + warnings.append("nested_physionet_layout") + recovery.append( + "Move the dataset contents up to the raw root or rerun wget with the generated --cut-dirs and -nH flags." + ) + + csv_files = sorted(root.rglob("*.csv.gz")) + empty_files = [str(path) for path in csv_files if path.stat().st_size == 0] + if empty_files: + warnings.append("empty_csv_gz") + recovery.append("Delete empty *.csv.gz files and rerun the resumable download.") + + expected_dirs = expected_raw_subdirectories(ds) + missing_dirs = [name for name in expected_dirs if not (root / name).is_dir()] + if missing_dirs: + warnings.append("missing_required_subdirectories") + errors.append( + "Missing required raw subdirectories: " + ", ".join(sorted(missing_dirs)) + ) + recovery.append( + "Confirm the target root and rerun the dataset-specific wget command." + ) + + if ds.name == "eicu": + root_csv_count = len(list(root.glob("*.csv.gz"))) + if root_csv_count == 0 and csv_files: + warnings.append("wrong_eicu_root_layout") + recovery.append( + "eICU CSV files should be directly under the eicu raw root, not nested." + ) + + if not csv_files: + errors.append("No *.csv.gz files found.") + recovery.append("Download the raw CSV files before initializing DuckDB.") + + if expected_dirs and csv_files and missing_dirs: + warnings.append("partial_download") + + return { + "ok": not errors and not empty_files, + "warnings": sorted(set(warnings)), + "errors": errors, + "csv_gz_count": len(csv_files), + "empty_csv_gz": empty_files, + "recovery": recovery, + } + + +def _download_guidance_data( + ds: DatasetDefinition, dataset_key: str, target_root: Path +) -> dict[str, Any]: + access_url = ds.dua_url or ds.dataset_page_url or ds.file_listing_url + return { + "dataset": dataset_key, + "target": str(target_root), + "requires_authentication": ds.requires_authentication, + "file_listing_url": ds.file_listing_url, + "wget_command": build_wget_command(ds, target_root) or None, + "layout": validate_raw_layout(dataset_key, target_root), + "recovery_hints": [ + "If conversion fails, rerun m4 download and then m4 init with --force.", + "If DuckDB is locked, stop MCP servers or notebooks using the database.", + "For BigQuery errors, verify gcloud application-default credentials and M4_PROJECT_ID.", + ], + "access_url": access_url, + } + + +def download_dataset_service( + dataset_name: str, + *, + target: str | None = None, + command_only: bool = False, + physionet_credentials: PhysioNetCredentials | None = None, + event_reporter: EventReporter | None = None, +) -> CommandResult | CommandError: + dataset_key = dataset_name.lower() + ensure_custom_datasets_loaded() + ds = DatasetRegistry.get(dataset_key) + if not ds: + supported = ", ".join(ds.name for ds in DatasetRegistry.list_all()) + return CommandError( + command="download", + code=ERROR_DATASET_NOT_FOUND, + message=f"Dataset '{dataset_name}' is not supported or not configured.", + hint=f"Supported datasets: {supported}", + ) + + target_root = ( + Path(target).expanduser().resolve() if target else default_raw_root(dataset_key) + ) + data = _download_guidance_data(ds, dataset_key, target_root) + + if command_only: + data["status"] = "command_only" + return CommandResult(command="download", data=data) + + if ds.requires_authentication and physionet_credentials is None: + access_url = data["access_url"] or "the dataset provider" + data["status"] = "blocked" + data["next_steps"] = [ + f"Confirm PhysioNet access: {access_url}", + "Run the generated wget command yourself, or pass --physionet-credentials-file to let M4 download.", + f"Then run: m4 init {dataset_key}", + ] + return CommandResult( + command="download", data=data, warnings=["credentialed_dataset"] + ) + + if not ds.file_listing_url: + return CommandError( + command="download", + code=ERROR_INVALID_OPTION, + message=f"Dataset '{dataset_key}' does not have a configured download URL.", + ) + + target_root.mkdir(parents=True, exist_ok=True) + reporter = get_event_reporter(event_reporter) + try: + with _quiet_console(): + downloaded = download_dataset( + dataset_key, + target_root, + credentials=physionet_credentials, + event_reporter=reporter if event_reporter is not None else None, + ) + except DatasetDownloadError as exc: + return CommandError( + command="download", + code=exc.code, + message=exc.message, + ) + + if not downloaded: + return CommandError( + command="download", + code=ERROR_INVALID_OPTION, + message="Download failed. Please check logs for details.", + hint="Retry the command; downloads are resumable.", + ) + + data["status"] = "completed" + data["layout"] = validate_raw_layout(dataset_key, target_root) + return CommandResult(command="download", data=data) diff --git a/src/m4/services/events.py b/src/m4/services/events.py new file mode 100644 index 00000000..75ad9de0 --- /dev/null +++ b/src/m4/services/events.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import json +import sys +import time +from dataclasses import dataclass, field +from typing import Any, TextIO + +from m4.services.redaction import redact_sensitive + + +@dataclass +class EventReporter: + command: str = "init" + enabled: bool = False + + def emit(self, event: str, **fields: Any) -> None: + return None + + def operation_started(self, **fields: Any) -> None: + self.emit("operation_started", command=self.command, **fields) + + def operation_completed(self, result: dict[str, Any]) -> None: + self.emit("operation_completed", command=self.command, result=result) + + def operation_failed(self, error: dict[str, Any]) -> None: + self.emit("operation_failed", command=self.command, error=error) + + +@dataclass +class NoopEventReporter(EventReporter): + enabled: bool = False + + +@dataclass +class NdjsonEventReporter(EventReporter): + stream: TextIO = field(default_factory=lambda: sys.stdout) + enabled: bool = True + _sequence: int = field(default=0, init=False) + + def emit(self, event: str, **fields: Any) -> None: + self._sequence += 1 + payload = { + "version": 1, + "sequence": self._sequence, + "time": time.time(), + "event": event, + **fields, + } + self.stream.write(json.dumps(redact_sensitive(payload), allow_nan=False) + "\n") + self.stream.flush() + + +def get_event_reporter(reporter: EventReporter | None) -> EventReporter: + return reporter if reporter is not None else NoopEventReporter() diff --git a/src/m4/services/init.py b/src/m4/services/init.py index b8b07c37..9395ee3b 100644 --- a/src/m4/services/init.py +++ b/src/m4/services/init.py @@ -14,11 +14,14 @@ from m4.core.derived.builtins import has_derived_support from m4.core.derived.materializer import get_derived_table_count, materialize_all from m4.data_io import ( + DatasetDownloadError, + PhysioNetCredentials, convert_csv_to_parquet, download_dataset, init_duckdb_from_parquet, verify_table_rowcount, ) +from m4.services.events import EventReporter, get_event_reporter from m4.services.results import ( ERROR_DATASET_NOT_FOUND, ERROR_INVALID_OPTION, @@ -73,8 +76,12 @@ def initialize_dataset_service( src: str | None = None, db_path_str: str | None = None, force: bool = False, + download: bool = False, + physionet_credentials: PhysioNetCredentials | None = None, + event_reporter: EventReporter | None = None, ) -> CommandResult | CommandError: """Run the non-interactive dataset initialization workflow.""" + reporter = get_event_reporter(event_reporter) dataset_key = dataset_name.lower() ds = DatasetRegistry.get(dataset_key) if not ds: @@ -109,34 +116,26 @@ def initialize_dataset_service( if not raw_present and not parquet_present: if ds.requires_authentication: - steps.extend( - [ - _step( - "raw_files", - "blocked", - ( - f"Files not found for credentialed dataset " - f"'{dataset_key}'. Download manually and rerun init." - ), - ), - _step("parquet", "skipped", "Raw files are not available."), - _step( - "database", "skipped", "Parquet files are not available." + if not download: + return CommandError( + command="init", + code="raw_files_missing", + message=( + f"Files not found for credentialed dataset '{dataset_key}'. " + "Pass --download with --physionet-credentials-file or place raw " + "CSV.gz files in the expected location." ), - _step( - "derived", - "skipped", - "Database initialization did not run.", + ) + if physionet_credentials is None: + return CommandError( + command="init", + code="missing_credentials", + message=( + f"Dataset '{dataset_key}' requires PhysioNet credentials " + "for download." ), - ] - ) - return CommandResult( - command="init", - data=_build_data( - dataset_key, final_db_path, pq_root, csv_root, steps - ), - warnings=[], - ) + hint="Provide --physionet-credentials-file with username and password fields.", + ) listing_url = ds.file_listing_url if not listing_url: @@ -170,10 +169,23 @@ def initialize_dataset_service( ) csv_root_default.mkdir(parents=True, exist_ok=True) - if not download_dataset(dataset_key, csv_root_default): + try: + downloaded = download_dataset( + dataset_key, + csv_root_default, + credentials=physionet_credentials, + event_reporter=reporter if event_reporter is not None else None, + ) + except DatasetDownloadError as exc: return CommandError( command="init", - code=ERROR_INVALID_OPTION, + code=exc.code, + message=exc.message, + ) + if not downloaded: + return CommandError( + command="init", + code="download_network_failed", message="Download failed. Please check logs for details.", ) csv_root = csv_root_default @@ -191,10 +203,15 @@ def initialize_dataset_service( steps.append( _step("parquet", "skipped", "Raw files are not available.") ) - elif not convert_csv_to_parquet(dataset_key, csv_root, pq_root): + elif not convert_csv_to_parquet( + dataset_key, + csv_root, + pq_root, + event_reporter=reporter if event_reporter is not None else None, + ): return CommandError( command="init", - code=ERROR_INVALID_OPTION, + code="conversion_failed", message="Conversion failed. Please check logs for details.", ) else: @@ -221,12 +238,16 @@ def initialize_dataset_service( message=f"Parquet directory not found at {pq_root}", ) - if not init_duckdb_from_parquet( - dataset_name=dataset_key, db_target_path=final_db_path - ): + init_kwargs = { + "dataset_name": dataset_key, + "db_target_path": final_db_path, + } + if event_reporter is not None: + init_kwargs["event_reporter"] = reporter + if not init_duckdb_from_parquet(**init_kwargs): return CommandError( command="init", - code=ERROR_INVALID_OPTION, + code="duckdb_init_failed", message=( f"Dataset '{dataset_name}' initialization FAILED. " "Please check logs for details." @@ -240,6 +261,11 @@ def initialize_dataset_service( record_count = verify_table_rowcount( final_db_path, verification_table_name ) + reporter.emit( + "verification_completed", + table=verification_table_name, + row_count=record_count, + ) steps.append( _step( "verification", @@ -251,6 +277,11 @@ def initialize_dataset_service( steps.append( _step("verification", "failed", f"Verification failed: {exc}") ) + return CommandError( + command="init", + code="verification_failed", + message=f"Verification failed: {exc}", + ) else: steps.append( _step("verification", "skipped", "No verification table configured.") diff --git a/src/m4/services/redaction.py b/src/m4/services/redaction.py new file mode 100644 index 00000000..513fa8c2 --- /dev/null +++ b/src/m4/services/redaction.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +_sensitive_values: set[str] = set() + + +def register_sensitive_value(value: str | None) -> None: + if value: + _sensitive_values.add(value) + + +def redact_sensitive(value: Any) -> Any: + if isinstance(value, dict): + return {str(key): redact_sensitive(item) for key, item in value.items()} + if isinstance(value, list | tuple | set): + return [redact_sensitive(item) for item in value] + if isinstance(value, Path): + return redact_sensitive(str(value)) + if isinstance(value, str): + redacted = value + for sensitive in sorted(_sensitive_values, key=len, reverse=True): + if sensitive: + redacted = redacted.replace(sensitive, "") + return redacted + return value diff --git a/src/m4/services/setup.py b/src/m4/services/setup.py new file mode 100644 index 00000000..22f59733 --- /dev/null +++ b/src/m4/services/setup.py @@ -0,0 +1,334 @@ +from __future__ import annotations + +import importlib.util +import os +import platform +import sys +from pathlib import Path +from typing import Any + +from m4.config import ( + ensure_custom_datasets_loaded, + get_active_backend, + get_active_dataset, + get_bigquery_project_id, + resolve_runtime_context, + set_active_backend, + set_active_dataset, + set_bigquery_project_id, +) +from m4.core.datasets import DatasetRegistry +from m4.services.download import default_raw_root, validate_raw_layout +from m4.services.init import initialize_dataset_service +from m4.services.results import ( + ERROR_INVALID_BACKEND, + ERROR_INVALID_OPTION, + CommandError, + CommandResult, +) +from m4.services.status import collect_status_snapshot + + +def _check( + name: str, ok: bool, message: str, hint: str | None = None +) -> dict[str, Any]: + result = {"name": name, "ok": ok, "message": message} + if hint: + result["hint"] = hint + return result + + +def doctor_service(*, include_paths: bool = False) -> CommandResult: + ensure_custom_datasets_loaded() + checks: list[dict[str, Any]] = [] + warnings: list[str] = [] + + checks.append( + _check( + "python_version", + sys.version_info >= (3, 10), + platform.python_version(), + "Use Python 3.10 or newer.", + ) + ) + checks.append( + _check( + "m4_import", + importlib.util.find_spec("m4") is not None, + "m4 package is importable", + "Run commands through uv run from the M4 project environment.", + ) + ) + checks.append( + _check( + "duckdb_import", + importlib.util.find_spec("duckdb") is not None, + "duckdb package is importable", + "Install project dependencies before using local DuckDB.", + ) + ) + + try: + active_dataset = get_active_dataset() + checks.append(_check("active_dataset", True, active_dataset)) + except Exception as exc: + active_dataset = None + checks.append( + _check( + "active_dataset", + False, + str(exc), + "Run m4 quickstart or m4 use .", + ) + ) + warnings.append("no_active_dataset") + + backend = get_active_backend() + checks.append( + _check( + "backend", + backend in {"duckdb", "bigquery"}, + backend, + "Run m4 backend duckdb or m4 backend bigquery.", + ) + ) + + snapshot = collect_status_snapshot(show_all=True, include_paths=include_paths) + if backend == "duckdb" and active_dataset: + active_status = next( + (ds for ds in snapshot["datasets"] if ds["name"] == active_dataset), None + ) + checks.append( + _check( + f"duckdb:{active_dataset}", + bool(active_status and active_status["db_present"]), + "local DuckDB present" + if active_status and active_status["db_present"] + else "local DuckDB missing", + f"Run m4 init {active_dataset}.", + ) + ) + + for ds in snapshot["datasets"]: + warnings.extend(ds.get("warnings", [])) + + if backend == "bigquery": + project_id = get_bigquery_project_id() + checks.append( + _check( + "bigquery_project_id", + bool(project_id), + "configured" if project_id else "missing", + "Run m4 setup-agent --backend bigquery --project-id YOUR_PROJECT_ID.", + ) + ) + checks.append( + _check( + "google_application_credentials", + bool(os.getenv("GOOGLE_APPLICATION_CREDENTIALS")) + or Path.home() + .joinpath(".config/gcloud/application_default_credentials.json") + .exists(), + "ambient credentials detected", + "Run gcloud auth application-default login.", + ) + ) + + ctx = resolve_runtime_context(path_disclosure=include_paths) + checks.append( + _check( + "mcp_config_hint", + True, + "Use m4 setup-agent --client claude or m4 config for MCP client setup.", + ) + ) + + return CommandResult( + command="doctor", + data={ + "summary": { + "ok": all(check["ok"] for check in checks), + "failed": [check["name"] for check in checks if not check["ok"]], + }, + "context": ctx.public_context(), + "data_dir": str(ctx.data_dir) if include_paths else None, + "checks": checks, + "status": snapshot, + }, + warnings=sorted(set(warnings)), + ) + + +def setup_agent_service( + *, + mode: str, + client: str, + dataset: str | None, + backend: str | None, + project_id: str | None, + apply_config: bool = False, +) -> CommandResult | CommandError: + ensure_custom_datasets_loaded() + if mode not in {"local", "protected"}: + return CommandError( + command="setup-agent", + code=ERROR_INVALID_OPTION, + message=f"Unsupported mode '{mode}'.", + hint="Use --mode local or --mode protected.", + ) + if client not in {"claude", "generic"}: + return CommandError( + command="setup-agent", + code=ERROR_INVALID_OPTION, + message=f"Unsupported client '{client}'.", + hint="Use --client claude or --client generic.", + ) + + resolved_backend = (backend or get_active_backend()).lower() + if resolved_backend not in {"duckdb", "bigquery"}: + return CommandError( + command="setup-agent", + code=ERROR_INVALID_BACKEND, + message=f"Unsupported backend '{resolved_backend}'.", + ) + if dataset and not DatasetRegistry.get(dataset): + supported = ", ".join(ds.name for ds in DatasetRegistry.list_all()) + return CommandError( + command="setup-agent", + code=ERROR_INVALID_OPTION, + message=f"Dataset '{dataset}' is not registered.", + hint=f"Supported datasets: {supported}", + ) + + if apply_config: + if dataset: + set_active_dataset(dataset) + if backend: + set_active_backend(resolved_backend) + if project_id: + set_bigquery_project_id(project_id) + + ctx = resolve_runtime_context(dataset=dataset, backend=resolved_backend) + env = { + "M4_HOME": str(ctx.home), + "M4_BACKEND": resolved_backend, + "M4_DATASET": dataset or ctx.dataset, + "M4_TELEMETRY_DIR": str(ctx.telemetry_dir), + } + if mode == "local": + env["M4_DATA_DIR"] = str(ctx.data_dir) + if resolved_backend == "bigquery" and (project_id or ctx.project_id): + env["M4_PROJECT_ID"] = project_id or ctx.project_id + + commands = [ + "m4 doctor", + f"m4 status --all{' --json' if client == 'generic' else ''}", + ] + if client == "claude": + commands.append( + "m4 config claude" + + (f" --backend {resolved_backend}" if resolved_backend else "") + + (f" --project-id {project_id}" if project_id else "") + ) + else: + commands.append("m4-infra") + + warnings = [] + if mode == "protected": + warnings.append("protected_mode_omits_data_dir") + if resolved_backend == "bigquery" and not (project_id or ctx.project_id): + warnings.append("bigquery_project_id_missing") + + return CommandResult( + command="setup-agent", + data={ + "mode": mode, + "client": client, + "applied": apply_config, + "environment": { + key: value for key, value in env.items() if value is not None + }, + "recommended_commands": commands, + "notes": [ + "Use protected mode when an agent should not see local data paths.", + "M4 telemetry is non-PHI operational provenance.", + ], + }, + warnings=warnings, + ) + + +def quickstart_service( + *, + workflow: str, + dataset: str | None = None, + backend: str | None = None, + project_id: str | None = None, + apply_config: bool = False, +) -> CommandResult | CommandError: + ensure_custom_datasets_loaded() + if workflow not in {"demo", "local", "bigquery"}: + return CommandError( + command="quickstart", + code=ERROR_INVALID_OPTION, + message=f"Unsupported workflow '{workflow}'.", + hint="Use --workflow demo, local, or bigquery.", + ) + + resolved_dataset = dataset or ( + "mimic-iv-demo" if workflow == "demo" else "mimic-iv" + ) + resolved_backend = backend or ("bigquery" if workflow == "bigquery" else "duckdb") + steps: list[dict[str, Any]] = [] + + if workflow == "demo": + steps.append({"command": "m4 init mimic-iv-demo", "mutates": True}) + if apply_config: + init_result = initialize_dataset_service("mimic-iv-demo") + steps.append({"result": init_result.to_json_dict()}) + elif workflow == "local": + raw_root = default_raw_root(resolved_dataset) + steps.extend( + [ + {"command": f"m4 download {resolved_dataset}", "mutates": False}, + { + "command": f"m4 init {resolved_dataset}", + "mutates": True, + "layout": validate_raw_layout(resolved_dataset, raw_root), + }, + ] + ) + else: + steps.extend( + [ + {"command": "gcloud auth application-default login", "mutates": True}, + { + "command": f"m4 backend bigquery --project-id {project_id or 'YOUR_PROJECT_ID'}", + "mutates": True, + }, + {"command": f"m4 use {resolved_dataset}", "mutates": True}, + ] + ) + + if apply_config: + set_active_backend(resolved_backend) + if project_id: + set_bigquery_project_id(project_id) + if DatasetRegistry.get(resolved_dataset): + set_active_dataset(resolved_dataset) + + return CommandResult( + command="quickstart", + data={ + "workflow": workflow, + "dataset": resolved_dataset, + "backend": resolved_backend, + "applied": apply_config, + "steps": steps, + }, + warnings=( + ["bigquery_project_id_missing"] + if workflow == "bigquery" and not (project_id or get_bigquery_project_id()) + else [] + ), + ) diff --git a/src/m4/services/status.py b/src/m4/services/status.py index 0a702d95..60f9d560 100644 --- a/src/m4/services/status.py +++ b/src/m4/services/status.py @@ -36,11 +36,15 @@ def _collect_dataset_status( include_paths: bool, ) -> dict[str, Any]: ds_def = DatasetRegistry.get(name) + raw_present = bool(ds_info.get("raw_present")) parquet_present = bool(ds_info.get("parquet_present")) db_present = bool(ds_info.get("db_present")) + raw_root = _absolute_path_or_none(ds_info.get("raw_root")) parquet_root = _absolute_path_or_none(ds_info.get("parquet_root")) db_path = _absolute_path_or_none(ds_info.get("db_path")) bigquery_available = bool(ds_def and ds_def.bigquery_dataset_ids) + requires_authentication = bool(ds_def and ds_def.requires_authentication) + download_available = bool(ds_def and ds_def.file_listing_url) warnings: list[str] = [] parquet_size_gb = None @@ -85,11 +89,28 @@ def _collect_dataset_status( except Exception: pass + if parquet_present and db_present: + setup_state = "ready" + elif requires_authentication and not raw_present and not parquet_present: + setup_state = "credentials_required" + elif not raw_present and not parquet_present: + setup_state = "missing_raw_files" + elif not parquet_present: + setup_state = "missing_parquet" + elif not db_present: + setup_state = "missing_database" + else: + setup_state = "ready" + result = { "name": name, "active": name == active_dataset, + "raw_present": raw_present, "parquet_present": parquet_present, "db_present": db_present, + "requires_authentication": requires_authentication, + "download_available": download_available, + "setup_state": setup_state, "bigquery_available": bigquery_available, "row_count": row_count, "parquet_size_gb": parquet_size_gb, @@ -103,6 +124,7 @@ def _collect_dataset_status( } if include_paths: + result["raw_root"] = raw_root result["parquet_root"] = parquet_root result["db_path"] = db_path diff --git a/tests/apps/cohort_builder/test_mcp_integration.py b/tests/apps/cohort_builder/test_mcp_integration.py index 69e0cec7..732082c4 100644 --- a/tests/apps/cohort_builder/test_mcp_integration.py +++ b/tests/apps/cohort_builder/test_mcp_integration.py @@ -299,7 +299,7 @@ async def test_query_cohort_empty_criteria(self, cohort_test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.apps.cohort_builder.tool.get_backend") as mock_backend: + with patch("m4.client.get_backend") as mock_backend: from m4.core.backends.duckdb import DuckDBBackend mock_backend.return_value = DuckDBBackend( @@ -338,7 +338,7 @@ async def test_query_cohort_with_age_filter(self, cohort_test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.apps.cohort_builder.tool.get_backend") as mock_backend: + with patch("m4.client.get_backend") as mock_backend: from m4.core.backends.duckdb import DuckDBBackend mock_backend.return_value = DuckDBBackend( @@ -377,7 +377,7 @@ async def test_query_cohort_with_gender_filter(self, cohort_test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.apps.cohort_builder.tool.get_backend") as mock_backend: + with patch("m4.client.get_backend") as mock_backend: from m4.core.backends.duckdb import DuckDBBackend mock_backend.return_value = DuckDBBackend( @@ -415,7 +415,7 @@ async def test_query_cohort_with_icd_filter(self, cohort_test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.apps.cohort_builder.tool.get_backend") as mock_backend: + with patch("m4.client.get_backend") as mock_backend: from m4.core.backends.duckdb import DuckDBBackend mock_backend.return_value = DuckDBBackend( @@ -455,7 +455,7 @@ async def test_query_cohort_with_icu_filter(self, cohort_test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.apps.cohort_builder.tool.get_backend") as mock_backend: + with patch("m4.client.get_backend") as mock_backend: from m4.core.backends.duckdb import DuckDBBackend mock_backend.return_value = DuckDBBackend( @@ -504,7 +504,7 @@ async def test_query_cohort_with_mortality_filter(self, cohort_test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.apps.cohort_builder.tool.get_backend") as mock_backend: + with patch("m4.client.get_backend") as mock_backend: from m4.core.backends.duckdb import DuckDBBackend mock_backend.return_value = DuckDBBackend( @@ -544,7 +544,7 @@ async def test_query_cohort_combined_filters(self, cohort_test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.apps.cohort_builder.tool.get_backend") as mock_backend: + with patch("m4.client.get_backend") as mock_backend: from m4.core.backends.duckdb import DuckDBBackend mock_backend.return_value = DuckDBBackend( diff --git a/tests/core/test_telemetry.py b/tests/core/test_telemetry.py index 4124cf0b..0b8487ef 100644 --- a/tests/core/test_telemetry.py +++ b/tests/core/test_telemetry.py @@ -9,6 +9,7 @@ import pandas as pd import pytest +from m4.core.context import M4ExecutionContext from m4.core.datasets import DatasetDefinition, Modality from m4.core.exceptions import M4Error from m4.core.telemetry import ( @@ -45,7 +46,7 @@ def __init__(self, return_value="ok", side_effect=None): self._return_value = return_value self._side_effect = side_effect - def invoke(self, dataset, params): + def invoke(self, dataset, params, context): if self._side_effect: raise self._side_effect return self._return_value @@ -75,12 +76,32 @@ def reset_telemetry_state(): _agent_id_var.reset(token_agent) -def _capture_record(mock_dataset, tool=None, params=None): +def _make_context( + mock_dataset, + *, + interface="unknown", + study_id=None, + session_id=None, + actor=None, +): + return M4ExecutionContext( + dataset=mock_dataset, + backend_name="mock", + backend=object(), + interface=interface, + study_id=study_id, + session_id=session_id, + actor=actor, + ) + + +def _capture_record(mock_dataset, tool=None, params=None, context=None): """Helper: invoke a tool and return the parsed telemetry record.""" tool = tool or MockTool(return_value="ok") params = params or MockInput() + context = context or _make_context(mock_dataset) with patch("m4.core.telemetry.logger") as mock_logger: - invoke_tracked(tool, mock_dataset, params) + invoke_tracked(tool, mock_dataset, params, context) return json.loads(mock_logger.info.call_args[0][0]) @@ -92,7 +113,9 @@ def _capture_record(mock_dataset, tool=None, params=None): class TestInvokeTracked: def test_returns_tool_result(self, mock_dataset): tool = MockTool(return_value={"data": [1, 2, 3]}) - result = invoke_tracked(tool, mock_dataset, MockInput()) + result = invoke_tracked( + tool, mock_dataset, MockInput(), _make_context(mock_dataset) + ) assert result == {"data": [1, 2, 3]} def test_success_record(self, mock_dataset): @@ -110,7 +133,9 @@ def test_failure_record_reraises(self, mock_dataset): with patch("m4.core.telemetry.logger") as mock_logger: with pytest.raises(M4Error, match="something broke"): - invoke_tracked(tool, mock_dataset, MockInput()) + invoke_tracked( + tool, mock_dataset, MockInput(), _make_context(mock_dataset) + ) record = json.loads(mock_logger.info.call_args[0][0]) @@ -119,10 +144,12 @@ def test_failure_record_reraises(self, mock_dataset): assert record["error_message"] == "something broke" def test_context_vars_in_record(self, mock_dataset): - set_interface("mcp") set_agent_id("agent-42") - record = _capture_record(mock_dataset) + record = _capture_record( + mock_dataset, + context=_make_context(mock_dataset, interface="mcp"), + ) assert record["interface"] == "mcp" assert record["agent_id"] == "agent-42" @@ -136,16 +163,16 @@ def test_params_captured(self, mock_dataset): assert record["params_summary"]["limit"] == 5 assert record["query_hash"] is not None - def test_env_attribution_fields(self, mock_dataset): - with patch.dict( - os.environ, - { - "M4_STUDY_ID": "study-1", - "M4_SESSION_ID": "session-1", - "M4_ACTOR": "actor-1", - }, - ): - record = _capture_record(mock_dataset) + def test_explicit_attribution_fields(self, mock_dataset): + record = _capture_record( + mock_dataset, + context=_make_context( + mock_dataset, + study_id="study-1", + session_id="session-1", + actor="actor-1", + ), + ) assert record["study_id"] == "study-1" assert record["session_id"] == "session-1" @@ -167,8 +194,13 @@ def test_jsonl_written(self, mock_dataset, tmp_path): """JSONL file is written with valid JSON per line.""" with patch("m4.config.get_telemetry_dir", return_value=tmp_path): tool = MockTool(return_value="ok") - invoke_tracked(tool, mock_dataset, MockInput()) - invoke_tracked(tool, mock_dataset, MockInput(sql_query="SELECT 2")) + invoke_tracked(tool, mock_dataset, MockInput(), _make_context(mock_dataset)) + invoke_tracked( + tool, + mock_dataset, + MockInput(sql_query="SELECT 2"), + _make_context(mock_dataset), + ) jsonl_path = tmp_path / "tool_calls.jsonl" assert jsonl_path.exists() @@ -186,7 +218,9 @@ def test_telemetry_off_suppresses_file(self, mock_dataset, tmp_path): with patch.dict(os.environ, {"M4_TELEMETRY": "off"}): with patch("m4.config.get_telemetry_dir", return_value=tmp_path): tool = MockTool(return_value="ok") - invoke_tracked(tool, mock_dataset, MockInput()) + invoke_tracked( + tool, mock_dataset, MockInput(), _make_context(mock_dataset) + ) jsonl_path = tmp_path / "tool_calls.jsonl" assert not jsonl_path.exists() @@ -195,7 +229,9 @@ def test_event_log_env_overrides_default_path(self, mock_dataset, tmp_path): event_log = tmp_path / "events.jsonl" with patch.dict(os.environ, {"M4_EVENT_LOG": str(event_log)}): with patch("m4.config.get_telemetry_dir", return_value=tmp_path / "unused"): - invoke_tracked(MockTool(), mock_dataset, MockInput()) + invoke_tracked( + MockTool(), mock_dataset, MockInput(), _make_context(mock_dataset) + ) assert event_log.exists() assert not (tmp_path / "unused" / "tool_calls.jsonl").exists() @@ -260,7 +296,9 @@ def test_row_count_none_on_error(self, mock_dataset): tool = MockTool(side_effect=M4Error("fail")) with patch("m4.core.telemetry.logger") as mock_logger: with pytest.raises(M4Error): - invoke_tracked(tool, mock_dataset, MockInput()) + invoke_tracked( + tool, mock_dataset, MockInput(), _make_context(mock_dataset) + ) record = json.loads(mock_logger.info.call_args[0][0]) assert record["row_count"] is None @@ -277,7 +315,9 @@ def test_constant_value(self): def test_jsonl_uses_constant(self, mock_dataset, tmp_path): with patch("m4.config.get_telemetry_dir", return_value=tmp_path): - invoke_tracked(MockTool(), mock_dataset, MockInput()) + invoke_tracked( + MockTool(), mock_dataset, MockInput(), _make_context(mock_dataset) + ) assert (tmp_path / TELEMETRY_FILENAME).exists() diff --git a/tests/test_api.py b/tests/test_api.py index 70b714a6..a0f3c25b 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -31,9 +31,9 @@ from m4.core.exceptions import SecurityError from m4.core.tools import init_tools -# Patch at the point of use in tool modules, not where defined -TABULAR_BACKEND_PATCH = "m4.core.tools.tabular.get_backend" -NOTES_BACKEND_PATCH = "m4.core.tools.notes.get_backend" +# Patch at the client backend factory boundary. +TABULAR_BACKEND_PATCH = "m4.client.get_backend" +NOTES_BACKEND_PATCH = "m4.client.get_backend" @pytest.fixture(autouse=True) diff --git a/tests/test_capabilities_download_setup.py b/tests/test_capabilities_download_setup.py new file mode 100644 index 00000000..a8abb553 --- /dev/null +++ b/tests/test_capabilities_download_setup.py @@ -0,0 +1,466 @@ +import json +from pathlib import Path +from unittest.mock import patch + +import pytest +from typer.testing import CliRunner + +import m4.config as cfg_mod +from m4 import M4Client, get_capabilities +from m4.cli import app +from m4.core.datasets import DatasetRegistry +from m4.data_io import DatasetDownloadError, PhysioNetCredentials +from m4.services.capabilities import build_capabilities_manifest +from m4.services.download import ( + build_wget_command, + download_dataset_service, + validate_raw_layout, +) +from m4.services.results import CommandError, CommandResult +from m4.services.setup import doctor_service, setup_agent_service + +runner = CliRunner() + + +def _install_custom_dataset(tmp_path, monkeypatch, name="custom-ed"): + datasets_dir = tmp_path / "m4_data" / "datasets" + datasets_dir.mkdir(parents=True) + (datasets_dir / f"{name}.json").write_text( + json.dumps( + { + "name": name, + "description": "Custom test dataset", + "file_listing_url": "https://physionet.org/files/custom-ed/1.0/", + "requires_authentication": True, + "modalities": ["TABULAR"], + "schema_mapping": {"": "custom_ed"}, + } + ) + ) + monkeypatch.setattr(cfg_mod, "_CUSTOM_DATASETS_DIR", datasets_dir) + DatasetRegistry.reset() + + +def teardown_function(): + DatasetRegistry.reset() + + +def test_capabilities_manifest_shape(): + manifest = build_capabilities_manifest() + + assert manifest["schema_version"] == 1 + assert "cli" in manifest["interfaces"] + assert "mcp" in manifest["interfaces"] + assert any(dataset["name"] == "mimic-iv" for dataset in manifest["datasets"]) + assert any(tool["name"] == "execute_query" for tool in manifest["tools"]) + assert ( + manifest["provenance_policy"]["event_export_command"] + == "m4 provenance export --json" + ) + commands = {command["name"]: command for command in manifest["commands"]} + assert "agent-env" in commands + assert "--physionet-credentials-file" in commands["download"]["flags"] + assert "--events" in commands["download"]["flags"] + assert commands["setup-agent"]["mutates"] is False + assert commands["setup-agent"]["mutates_with"] == ["--apply"] + assert "--apply" in commands["quickstart"]["flags"] + + +def test_capabilities_manifest_structural_contract(): + manifest = build_capabilities_manifest() + + assert { + "schema_version", + "interfaces", + "runtime", + "commands", + "tools", + "datasets", + "limits", + "concepts", + "provenance_policy", + }.issubset(manifest) + + for command in manifest["commands"]: + assert {"name", "flags", "mutates"}.issubset(command) + + for dataset in manifest["datasets"]: + assert { + "name", + "requires_authentication", + "modalities", + "bigquery", + "verification_table", + "schema_mapping", + "expected_local_layout", + }.issubset(dataset) + assert { + "available", + "project_id", + "dataset_ids", + "schema_mapping", + }.issubset(dataset["bigquery"]) + assert { + "recommended_raw_root", + "raw_subdirectories", + "parquet_root", + "duckdb_filename", + }.issubset(dataset["expected_local_layout"]) + + for tool in manifest["tools"]: + assert { + "name", + "description", + "input_fields", + "required_modalities", + "compatible_datasets", + "supported_datasets", + }.issubset(tool) + + assert { + "query_row_limit_default", + "path_redaction_default", + "supported_backends", + "conversion_env", + }.issubset(manifest["limits"]) + assert {"derived_tables", "skills"}.issubset(manifest["concepts"]) + assert { + "telemetry_destination", + "path_redaction", + "event_export_command", + "non_phi_policy", + }.issubset(manifest["provenance_policy"]) + + +def test_capabilities_manifest_agent_command_contract(): + manifest = build_capabilities_manifest() + commands = {command["name"]: command for command in manifest["commands"]} + expected = { + "download", + "init", + "setup-agent", + "quickstart", + "doctor", + "capabilities", + "schema", + "query", + } + + assert expected.issubset(commands) + for name in expected: + assert isinstance(commands[name]["flags"], list) + assert isinstance(commands[name]["mutates"], bool) + assert commands["setup-agent"]["mutates_with"] == ["--apply"] + assert commands["quickstart"]["mutates_with"] == ["--apply"] + + +def test_capabilities_manifest_builtin_dataset_identity_contract(): + manifest = build_capabilities_manifest() + datasets = {dataset["name"]: dataset for dataset in manifest["datasets"]} + expected = {"mimic-iv-demo", "mimic-iv", "mimic-iv-note", "eicu"} + + assert expected.issubset(datasets) + for name in expected: + dataset = datasets[name] + assert dataset["name"] == name + assert isinstance(dataset["requires_authentication"], bool) + assert dataset["modalities"] + assert dataset["expected_local_layout"]["recommended_raw_root"] + assert dataset["expected_local_layout"]["duckdb_filename"] + assert "available" in dataset["bigquery"] + assert "dataset_ids" in dataset["bigquery"] + assert "verification_table" in dataset + + +def test_python_get_capabilities_exports_manifest(): + assert get_capabilities()["schema_version"] == 1 + + +def test_m4client_capabilities_exports_manifest(): + client = M4Client(dataset="mimic-iv-demo", backend="duckdb") + + assert client.capabilities()["schema_version"] == 1 + + +def test_capabilities_cli_json(): + result = runner.invoke(app, ["capabilities", "--json"]) + + assert result.exit_code == 0 + payload = json.loads(result.stdout) + assert payload["schema_version"] == 1 + assert "datasets" in payload + + +def test_agent_env_formats(): + dotenv = runner.invoke(app, ["agent-env", "--dataset", "mimic-iv-demo"]) + assert dotenv.exit_code == 0 + assert "M4_BACKEND=" in dotenv.stdout + + json_result = runner.invoke( + app, ["agent-env", "--dataset", "mimic-iv-demo", "--format", "json"] + ) + assert json_result.exit_code == 0 + payload = json.loads(json_result.stdout) + assert payload["command"] == "agent-env" + + text = runner.invoke( + app, ["agent-env", "--dataset", "mimic-iv-demo", "--format", "text"] + ) + assert text.exit_code == 0 + assert "Recommended commands" in text.stdout + + +def test_agent_env_invalid_format(): + result = runner.invoke(app, ["agent-env", "--format", "yaml"]) + + assert result.exit_code == 1 + assert "Unsupported format" in result.stdout + + +def test_credentialed_download_returns_wget_guidance(tmp_path): + result = download_dataset_service("mimic-iv", target=str(tmp_path)) + + assert isinstance(result, CommandResult) + assert result.ok is True + assert result.data["status"] == "blocked" + assert "--cut-dirs=3 -nH" in result.data["wget_command"] + assert "--ask-password" in result.data["wget_command"] + assert result.data["next_steps"][0].endswith("mimiciv/") + + +def test_eicu_download_returns_top_level_raw_layout_guidance(tmp_path): + result = download_dataset_service("eicu", target=str(tmp_path)) + + assert isinstance(result, CommandResult) + assert result.data["status"] == "blocked" + assert "--cut-dirs=3 -nH" in result.data["wget_command"] + assert "https://physionet.org/files/eicu-crd/2.0/" in result.data["wget_command"] + + +def test_download_service_loads_custom_dataset_and_quotes_target(tmp_path, monkeypatch): + _install_custom_dataset(tmp_path, monkeypatch) + target = tmp_path / "raw data" + + result = download_dataset_service("custom-ed", target=str(target)) + + assert isinstance(result, CommandResult) + assert result.data["status"] == "blocked" + assert result.data["dataset"] == "custom-ed" + assert "--cut-dirs=3 -nH" in result.data["wget_command"] + assert f"-P '{target}'" in result.data["wget_command"] + assert result.data["next_steps"][0].endswith("/custom-ed/1.0/") + + +def test_wget_command_cut_dirs_follows_listing_url_path(tmp_path, monkeypatch): + _install_custom_dataset(tmp_path, monkeypatch) + cfg_mod.ensure_custom_datasets_loaded() + dataset = DatasetRegistry.get("custom-ed") + + command = build_wget_command(dataset, tmp_path / "target with spaces") + + assert "--cut-dirs=3 -nH" in command + assert f"-P '{tmp_path / 'target with spaces'}'" in command + + +def test_public_download_service_uses_downloader(tmp_path): + with patch( + "m4.services.download.download_dataset", return_value=True + ) as mock_download: + result = download_dataset_service("mimic-iv-demo", target=str(tmp_path)) + + assert isinstance(result, CommandResult) + assert result.data["status"] == "completed" + mock_download.assert_called_once() + + +def test_public_download_cli_json_uses_downloader(tmp_path): + with patch( + "m4.services.download.download_dataset", return_value=True + ) as mock_download: + result = runner.invoke( + app, ["download", "mimic-iv-demo", "--target", str(tmp_path), "--json"] + ) + + assert result.exit_code == 0 + payload = json.loads(result.stdout) + assert payload["command"] == "download" + assert payload["status"] == "completed" + mock_download.assert_called_once() + + +def test_credentialed_download_cli_without_credentials_returns_guidance(): + result = runner.invoke(app, ["download", "mimic-iv", "--json"]) + + assert result.exit_code == 0 + payload = json.loads(result.stdout) + assert payload["status"] == "blocked" + assert payload["wget_command"] + assert "None" not in "\n".join(payload["next_steps"]) + + +def test_init_credentialed_guidance_matches_download_command(tmp_path): + pq_root = tmp_path / "m4_data" / "parquet" / "mimic-iv" + pq_root.mkdir(parents=True) + + with ( + patch("m4.config._find_project_root_from_cwd", return_value=Path.cwd()), + patch("m4.cli.get_dataset_parquet_root", return_value=pq_root), + ): + result = runner.invoke(app, ["init", "mimic-iv", "--no-interactive"]) + + assert result.exit_code == 0 + assert "--cut-dirs=3 -nH" in result.output + assert "raw_files/mimic-iv" in result.output.replace("\n", "") + + +def test_credentialed_download_with_credentials_delegates(tmp_path): + creds = PhysioNetCredentials(username="alice", password="secret") + with patch( + "m4.services.download.download_dataset", return_value=True + ) as mock_download: + result = download_dataset_service( + "mimic-iv", + target=str(tmp_path), + physionet_credentials=creds, + ) + + assert isinstance(result, CommandResult) + assert result.data["status"] == "completed" + mock_download.assert_called_once() + assert mock_download.call_args.kwargs["credentials"] == creds + + +def test_credentialed_download_cli_with_credentials_file_delegates(tmp_path): + creds_file = tmp_path / "physionet.json" + creds_file.write_text(json.dumps({"username": "alice", "password": "secret"})) + + with patch( + "m4.services.download.download_dataset", return_value=True + ) as mock_download: + result = runner.invoke( + app, + [ + "download", + "mimic-iv", + "--target", + str(tmp_path), + "--json", + "--physionet-credentials-file", + str(creds_file), + ], + ) + + assert result.exit_code == 0 + payload = json.loads(result.stdout) + assert payload["status"] == "completed" + mock_download.assert_called_once() + assert mock_download.call_args.kwargs["credentials"].username == "alice" + + +@pytest.mark.parametrize( + ("code", "message"), + [ + ("physionet_auth_failed", "bad credentials"), + ("download_network_failed", "network down"), + ], +) +def test_download_errors_become_stable_command_errors(tmp_path, code, message): + with patch( + "m4.services.download.download_dataset", + side_effect=DatasetDownloadError(code, message), + ): + result = download_dataset_service( + "mimic-iv", + target=str(tmp_path), + physionet_credentials=PhysioNetCredentials("alice", "secret"), + ) + + assert isinstance(result, CommandError) + assert result.code == code + assert result.message == message + + +def test_layout_validation_detects_nested_and_missing_dirs(tmp_path): + nested = tmp_path / "physionet.org" / "files" / "mimiciv" / "3.1" + nested.mkdir(parents=True) + (nested / "patients.csv.gz").write_bytes(b"") + + result = validate_raw_layout("mimic-iv", tmp_path) + + assert result["ok"] is False + assert "nested_physionet_layout" in result["warnings"] + assert "missing_required_subdirectories" in result["warnings"] + assert "empty_csv_gz" in result["warnings"] + + +def test_doctor_setup_agent_quickstart_json(): + doctor = runner.invoke(app, ["doctor", "--json"]) + assert doctor.exit_code == 0 + assert json.loads(doctor.stdout)["command"] == "doctor" + + setup = runner.invoke( + app, + [ + "setup-agent", + "--dataset", + "mimic-iv-demo", + "--backend", + "duckdb", + "--format", + "json", + ], + ) + assert setup.exit_code == 0 + assert json.loads(setup.stdout)["command"] == "setup-agent" + + quickstart = runner.invoke(app, ["quickstart", "--workflow", "demo", "--json"]) + assert quickstart.exit_code == 0 + assert json.loads(quickstart.stdout)["command"] == "quickstart" + + +def test_setup_agent_service_loads_custom_dataset(tmp_path, monkeypatch): + _install_custom_dataset(tmp_path, monkeypatch) + + result = setup_agent_service( + mode="local", + client="generic", + dataset="custom-ed", + backend="duckdb", + project_id=None, + ) + + assert isinstance(result, CommandResult) + assert result.data["environment"]["M4_DATASET"] == "custom-ed" + + +@patch("m4.services.setup.get_active_backend", return_value="duckdb") +@patch("m4.services.setup.get_active_dataset", return_value="mimic-iv-demo") +@patch("m4.services.setup.collect_status_snapshot") +def test_doctor_only_requires_active_duckdb_dataset( + mock_status, mock_dataset, mock_backend +): + mock_status.return_value = { + "version": 1, + "active_dataset": "mimic-iv-demo", + "backend": "duckdb", + "bigquery_project_id": None, + "datasets": [ + { + "name": "mimic-iv-demo", + "db_present": True, + "warnings": [], + }, + { + "name": "mimic-iv", + "db_present": False, + "warnings": [], + }, + ], + } + + result = doctor_service() + + assert result.data["summary"]["ok"] is True + check_names = [check["name"] for check in result.data["checks"]] + assert "duckdb:mimic-iv-demo" in check_names + assert "duckdb:mimic-iv" not in check_names diff --git a/tests/test_cli.py b/tests/test_cli.py index 5e83f603..eb61d3bc 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -8,10 +8,15 @@ from m4.cli import app from m4.core.exceptions import DatasetError +from m4.services.results import CommandResult runner = CliRunner() +def _ndjson_lines(output: str) -> list[dict]: + return [json.loads(line) for line in output.splitlines() if line.strip()] + + @pytest.fixture(autouse=True) def inject_version(monkeypatch): # Patch __version__ in the console module where print_logo imports it @@ -82,6 +87,129 @@ def test_init_command_duckdb_custom_path(tmp_path): mock_rowcount.assert_called() +@patch("m4.cli.initialize_dataset_service") +def test_init_json_preserves_single_object_output(mock_init): + mock_init.return_value = CommandResult( + command="init", + data={ + "dataset": "mimic-iv-demo", + "db_path": None, + "parquet_root": None, + "raw_root": None, + "steps": [], + }, + ) + + result = runner.invoke(app, ["init", "mimic-iv-demo", "--json"]) + + assert result.exit_code == 0 + payload = json.loads(result.stdout) + assert payload["ok"] is True + assert payload["dataset"] == "mimic-iv-demo" + assert len(result.stdout.strip().splitlines()) > 1 + assert not result.stdout.lstrip().startswith('{"version"') + + +@patch("m4.cli.initialize_dataset_service") +def test_init_json_events_ndjson_wraps_final_result(mock_init): + mock_init.return_value = CommandResult( + command="init", + data={ + "dataset": "mimic-iv-demo", + "db_path": None, + "parquet_root": None, + "raw_root": None, + "steps": [], + }, + ) + + result = runner.invoke( + app, + ["init", "mimic-iv-demo", "--json", "--events", "ndjson", "--no-interactive"], + ) + + assert result.exit_code == 0 + events = _ndjson_lines(result.stdout) + assert events[0]["event"] == "operation_started" + assert events[-1]["event"] == "operation_completed" + assert events[-1]["result"]["ok"] is True + assert events[-1]["result"]["dataset"] == "mimic-iv-demo" + + +@patch("m4.services.init.get_default_database_path") +@patch("m4.services.init.get_dataset_parquet_root") +def test_init_download_missing_credentials_returns_structured_error( + mock_parquet_root, mock_db_path, tmp_path +): + pq_root = tmp_path / "parquet" / "mimic-iv" + pq_root.mkdir(parents=True) + mock_parquet_root.return_value = pq_root + mock_db_path.return_value = tmp_path / "mimic.duckdb" + + result = runner.invoke( + app, + ["init", "mimic-iv", "--json", "--no-interactive", "--download"], + ) + + assert result.exit_code == 1 + payload = json.loads(result.stdout) + assert payload["ok"] is False + assert payload["error"]["code"] == "missing_credentials" + + +@patch("m4.services.init.get_default_database_path") +@patch("m4.services.init.get_dataset_parquet_root") +def test_init_events_invalid_physionet_credentials_redacts_password( + mock_parquet_root, mock_db_path, tmp_path +): + class UnauthorizedResponse: + def __init__(self): + self.status_code = 401 + self.reason = "Unauthorized" + self.content = b"" + self.headers = {} + + class Session: + def __init__(self): + self.headers = {} + self.auth = None + + def get(self, *args, **kwargs): + return UnauthorizedResponse() + + pq_root = tmp_path / "parquet" / "mimic-iv" + pq_root.mkdir(parents=True) + mock_parquet_root.return_value = pq_root + mock_db_path.return_value = tmp_path / "mimic.duckdb" + credentials_path = tmp_path / "physionet.json" + credentials_path.write_text( + json.dumps({"username": "alice", "password": "do-not-print"}) + ) + + with patch("m4.data_io.requests.Session", Session): + result = runner.invoke( + app, + [ + "init", + "mimic-iv", + "--json", + "--events", + "ndjson", + "--no-interactive", + "--download", + "--physionet-credentials-file", + str(credentials_path), + ], + ) + + assert result.exit_code == 1 + events = _ndjson_lines(result.stdout) + assert events[-1]["event"] == "operation_failed" + assert events[-1]["error"]["code"] == "physionet_auth_failed" + assert "do-not-print" not in result.stdout + assert "do-not-print" not in result.stderr + + def test_config_validation_bigquery_with_db_path(): """Test that bigquery backend rejects db-path parameter.""" result = runner.invoke( @@ -507,6 +635,52 @@ def test_status_json_excludes_secret_values_but_allows_project_id( assert "super-secret-password" not in result.stdout +@patch("m4.services.status.get_bigquery_project_id", return_value=None) +@patch("m4.services.status.get_active_backend", return_value="duckdb") +@patch("m4.services.status.detect_available_local_datasets") +@patch("m4.services.status.get_active_dataset", return_value="mimic-iv") +def test_status_json_redacts_paths_by_default_and_exposes_with_paths( + mock_active, + mock_detect, + mock_backend, + mock_project, +): + mock_detect.return_value = { + "mimic-iv": { + "raw_present": False, + "parquet_present": False, + "db_present": False, + "raw_root": "/tmp/m4/raw_files/mimic-iv", + "parquet_root": "/tmp/m4/parquet/mimic-iv", + "db_path": "/tmp/m4/databases/mimic.duckdb", + } + } + + default_result = runner.invoke(app, ["status", "--json", "--no-interactive"]) + paths_result = runner.invoke( + app, ["status", "--json", "--paths", "--no-interactive"] + ) + + assert default_result.exit_code == 0 + default_dataset = json.loads(default_result.stdout)["datasets"][0] + assert default_dataset["setup_state"] == "credentials_required" + assert default_dataset["requires_authentication"] is True + assert "raw_root" not in default_dataset + assert "/tmp/m4" not in default_result.stdout + + assert paths_result.exit_code == 0 + paths_dataset = json.loads(paths_result.stdout)["datasets"][0] + assert paths_dataset["raw_root"] == str( + Path("/tmp/m4/raw_files/mimic-iv").resolve() + ) + assert paths_dataset["parquet_root"] == str( + Path("/tmp/m4/parquet/mimic-iv").resolve() + ) + assert paths_dataset["db_path"] == str( + Path("/tmp/m4/databases/mimic.duckdb").resolve() + ) + + @patch("m4.services.status.get_bigquery_project_id", return_value=None) @patch("m4.services.status.get_active_backend", return_value="duckdb") @patch( diff --git a/tests/test_cli_subprocess.py b/tests/test_cli_subprocess.py index 54525b4f..5f16690f 100644 --- a/tests/test_cli_subprocess.py +++ b/tests/test_cli_subprocess.py @@ -265,12 +265,11 @@ def test_init_json_subprocess_error_is_parseable(tmp_path): assert payload["error"]["code"] == "dataset_not_found" -def test_init_json_subprocess_blocked_state_is_parseable(tmp_path): +def test_init_json_subprocess_missing_credentialed_raw_files_is_parseable(tmp_path): result = _run_m4(["init", "mimic-iv", "--json"], tmp_path) - assert result.returncode == 0 + assert result.returncode != 0 payload = _assert_single_json_stdout(result) - assert payload["ok"] is True + assert payload["ok"] is False assert payload["command"] == "init" - assert payload["dataset"] == "mimic-iv" - assert payload["steps"][0]["status"] == "blocked" + assert payload["error"]["code"] == "raw_files_missing" diff --git a/tests/test_client.py b/tests/test_client.py new file mode 100644 index 00000000..6bf6cd32 --- /dev/null +++ b/tests/test_client.py @@ -0,0 +1,135 @@ +"""Tests for the first-class M4Client API.""" + +import json +import os +from unittest.mock import MagicMock, patch + +import pandas as pd +import pytest + +from m4 import M4Client +from m4.core.backends.base import QueryResult +from m4.core.datasets import DatasetDefinition, Modality +from m4.core.exceptions import ModalityError +from m4.core.tools import init_tools + + +@pytest.fixture(autouse=True) +def initialized_tools(): + init_tools() + + +@pytest.fixture +def tabular_dataset(): + return DatasetDefinition( + name="test-tabular", + modalities=frozenset({Modality.TABULAR}), + ) + + +@pytest.fixture +def notes_dataset(): + return DatasetDefinition( + name="test-notes", + modalities=frozenset({Modality.NOTES}), + ) + + +@pytest.fixture +def backend(): + mock = MagicMock() + mock.name = "mock" + mock.get_backend_info.return_value = "Mock backend" + return mock + + +def test_constructor_resolves_dataset_and_backend(tabular_dataset, backend): + client = M4Client(dataset=tabular_dataset, backend=backend, interface="python_api") + + assert client.dataset is tabular_dataset + assert client.backend is backend + assert client.context.backend_name == "mock" + assert client.context.interface == "python_api" + + +def test_schema_table_info_and_query_delegate_through_context(tabular_dataset, backend): + schema_df = pd.DataFrame({"name": ["subject_id"], "type": ["INTEGER"]}) + sample_df = pd.DataFrame({"subject_id": [1]}) + query_df = pd.DataFrame({"count": [1]}) + + backend.get_table_list.return_value = ["patients"] + backend.get_table_info.return_value = QueryResult(dataframe=schema_df, row_count=1) + backend.get_sample_data.return_value = QueryResult(dataframe=sample_df, row_count=1) + backend.execute_query.return_value = QueryResult(dataframe=query_df, row_count=1) + + client = M4Client(dataset=tabular_dataset, backend=backend) + + assert client.schema()["tables"] == ["patients"] + assert client.table_info("patients")["schema"].equals(schema_df) + assert client.query("SELECT COUNT(*) FROM patients").equals(query_df) + + backend.get_table_list.assert_called_once_with(tabular_dataset, client.context) + backend.get_table_info.assert_called_once_with( + "patients", tabular_dataset, client.context + ) + backend.execute_query.assert_called_once_with( + "SELECT COUNT(*) FROM patients", tabular_dataset, client.context + ) + + +def test_explicit_backend_selection_does_not_mutate_environment(tabular_dataset): + previous = os.environ.get("M4_BACKEND") + os.environ["M4_BACKEND"] = "duckdb" + + try: + with patch("m4.client.get_backend") as mock_get_backend: + mock_backend = MagicMock() + mock_backend.name = "bigquery" + mock_backend.get_table_list.return_value = [] + mock_backend.get_backend_info.return_value = "BigQuery" + mock_get_backend.return_value = mock_backend + + client = M4Client(dataset=tabular_dataset, backend="bigquery") + client.schema() + + assert client.context.backend_name == "bigquery" + assert os.environ.get("M4_BACKEND") == "duckdb" + mock_get_backend.assert_called_once_with("bigquery") + finally: + if previous is None: + os.environ.pop("M4_BACKEND", None) + else: + os.environ["M4_BACKEND"] = previous + + +def test_explicit_attribution_recorded_without_environment(tabular_dataset, backend): + os.environ.pop("M4_STUDY_ID", None) + os.environ.pop("M4_SESSION_ID", None) + os.environ.pop("M4_ACTOR", None) + + backend.execute_query.return_value = QueryResult( + dataframe=pd.DataFrame({"x": [1]}), row_count=1 + ) + client = M4Client( + dataset=tabular_dataset, + backend=backend, + study_id="study-1", + session_id="session-1", + actor="actor-1", + ) + + with patch("m4.core.telemetry.logger") as mock_logger: + client.query("SELECT 1") + + record = json.loads(mock_logger.info.call_args[0][0]) + assert record["interface"] == "python_api" + assert record["study_id"] == "study-1" + assert record["session_id"] == "session-1" + assert record["actor"] == "actor-1" + + +def test_notes_methods_raise_for_tabular_dataset(tabular_dataset, backend): + client = M4Client(dataset=tabular_dataset, backend=backend) + + with pytest.raises(ModalityError): + client.search_notes("pneumonia") diff --git a/tests/test_data_io.py b/tests/test_data_io.py index 5f3986f2..c43afd04 100644 --- a/tests/test_data_io.py +++ b/tests/test_data_io.py @@ -8,10 +8,12 @@ from m4.core.datasets import DatasetDefinition, Modality from m4.data_io import ( COMMON_USER_AGENT, + PhysioNetCredentials, _create_duckdb_with_views, _scrape_urls_from_html_page, compute_parquet_dir_size, convert_csv_to_parquet, + download_dataset, init_duckdb_from_parquet, verify_table_rowcount, ) @@ -87,6 +89,55 @@ def test_common_user_agent_header(): assert "Mozilla/" in COMMON_USER_AGENT +def test_download_dataset_uses_credentials_and_preserves_listing_layout( + tmp_path, monkeypatch +): + class Response: + def __init__(self, content=b"", status_code=200, headers=None): + self.content = content + self.status_code = status_code + self.headers = headers or {} + self.reason = "OK" + + def iter_content(self, chunk_size=8192): + for index in range(0, len(self.content), chunk_size): + yield self.content[index : index + chunk_size] + + class Session: + def __init__(self): + self.headers = {} + self.auth = None + + def head(self, url, **kwargs): + return Response(headers={"content-length": "9", "accept-ranges": "bytes"}) + + def get(self, url, **kwargs): + if url.endswith("/hosp/"): + return Response( + b'admissions', + headers={"content-length": "9"}, + ) + if url.endswith("/icu/"): + return Response(b"", headers={"content-length": "0"}) + return Response(b"csv-bytes", headers={"content-length": "9"}) + + sessions = [] + + def make_session(): + session = Session() + sessions.append(session) + return session + + monkeypatch.setattr("m4.data_io.requests.Session", make_session) + + credentials = PhysioNetCredentials(username="alice", password="secret") + ok = download_dataset("mimic-iv", tmp_path, credentials=credentials) + + assert ok is True + assert (tmp_path / "hosp" / "admissions.csv.gz").read_bytes() == b"csv-bytes" + assert sessions[0].auth == ("alice", "secret") + + # ------------------------------------------------------------ # CSV -> Parquet conversion and DuckDB init tests # ------------------------------------------------------------ diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index adbd29d7..e582f197 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -39,6 +39,22 @@ def test_server_instance_exists(self): assert mcp is not None assert mcp.name == "m4" + @pytest.mark.asyncio + async def test_capabilities_tool_via_client(self): + async with Client(mcp) as client: + result = await client.call_tool("capabilities", {}) + result_text = str(result) + assert "schema_version" in result_text + assert "m4://capabilities" in result_text + + @pytest.mark.asyncio + async def test_capabilities_resource_via_client(self): + async with Client(mcp) as client: + result = await client.read_resource("m4://capabilities") + result_text = str(result) + assert "schema_version" in result_text + assert "m4://capabilities" in result_text + class TestMCPTools: """Test MCP tools functionality.""" @@ -121,7 +137,7 @@ async def test_tools_via_client(self, test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.core.tools.tabular.get_backend") as mock_get_backend: + with patch("m4.client.get_backend") as mock_get_backend: from m4.core.backends.duckdb import DuckDBBackend # Use real DuckDB backend with test database @@ -209,7 +225,7 @@ async def test_invalid_sql(self, test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.core.tools.tabular.get_backend") as mock_get_backend: + with patch("m4.client.get_backend") as mock_get_backend: mock_get_backend.return_value = DuckDBBackend( db_path_override=test_db ) @@ -249,7 +265,7 @@ async def test_empty_results(self, test_db): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.core.tools.tabular.get_backend") as mock_get_backend: + with patch("m4.client.get_backend") as mock_get_backend: mock_get_backend.return_value = DuckDBBackend( db_path_override=test_db ) @@ -333,7 +349,7 @@ async def test_bigquery_tools(self): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=mock_ds ): - with patch("m4.core.tools.tabular.get_backend") as mock_get_backend: + with patch("m4.client.get_backend") as mock_get_backend: # Mock the backend mock_backend = Mock() mock_backend.name = "bigquery" @@ -399,7 +415,7 @@ async def test_incompatible_tool_returns_proactive_error(self): "m4.mcp_server.DatasetRegistry.get_active", return_value=notes_only_ds ): # Mock backend that should NOT be called - with patch("m4.core.tools.tabular.get_backend") as mock_backend: + with patch("m4.client.get_backend") as mock_backend: async with Client(mcp) as client: # Call execute_query which requires TABULAR modality result = await client.call_tool( @@ -454,7 +470,7 @@ async def test_compatible_tool_executes_successfully(self, tmp_path): with patch( "m4.mcp_server.DatasetRegistry.get_active", return_value=tabular_ds ): - with patch("m4.core.tools.tabular.get_backend") as mock_get_backend: + with patch("m4.client.get_backend") as mock_get_backend: mock_get_backend.return_value = DuckDBBackend( db_path_override=str(db_path) ) diff --git a/tests/test_services_init.py b/tests/test_services_init.py index 7a2a7403..23a3286f 100644 --- a/tests/test_services_init.py +++ b/tests/test_services_init.py @@ -94,7 +94,7 @@ def test_init_service_raw_to_parquet_conversion_path( @patch("m4.services.init.get_default_database_path") @patch("m4.services.init.get_dataset_parquet_root") -def test_init_service_credentialed_dataset_returns_blocked_state( +def test_init_service_credentialed_dataset_missing_files_is_error( mock_parquet_root, mock_db_path, tmp_path ): pq_root = tmp_path / "parquet" / "mimic-iv" @@ -105,9 +105,8 @@ def test_init_service_credentialed_dataset_returns_blocked_state( result = initialize_dataset_service("mimic-iv") - assert isinstance(result, CommandResult) - assert _step_by_name(result, "raw_files")["status"] == "blocked" - assert _step_by_name(result, "database")["status"] == "skipped" + assert isinstance(result, CommandError) + assert result.code == "raw_files_missing" @patch("m4.services.init.has_derived_support", return_value=False)