diff --git a/README.md b/README.md index 2b3f7a4..8447682 100644 --- a/README.md +++ b/README.md @@ -79,8 +79,8 @@ The configuration system uses a flexible, hierarchical approach supporting diffe ### Configuration Files Structure -- **Main pipeline config**: Controls manuscript processing, AI model selection, and pipeline steps. -- **QC config (`config.qc.yaml`)**: Controls all quality control tests, test metadata, and versioning. Example: +- **Main pipeline config** (`config.dev.yaml`, `config.sample.yaml`): Controls manuscript processing and pipeline step prompts. Each AI step uses a **flat** block with only `model` and shared `prompts` — no per-provider nesting and no sampling parameters in YAML (API defaults apply at runtime). +- **QC config (`config.qc.yaml`)**: Separate schema for quality control tests (still uses provider-specific blocks under `default.openai` / `default.anthropic`). Example: ```yaml qc_version: "2" @@ -230,10 +230,9 @@ test_runs: - Panel source assignment - Data availability extraction -2. **Model Configuration**: Configure different models and parameters: - - Multiple providers (OpenAI, Anthropic) - - Various models per provider - - Temperature and top_p parameter combinations +2. **Model Configuration**: Configure models to benchmark (provider inferred from model name): + - OpenAI models (`gpt-*`, `o*-mini`, etc.) + - Anthropic models (`claude-*`) when supported by the benchmark runner - Multiple runs per configuration 3. **Output and Metrics**: @@ -263,6 +262,73 @@ The benchmark system generates several output files: - Scoring details - Error information +### Main pipeline YAML shape (v3.3+) + +Each AI step declares a **model** and shared **prompts** only. The runtime provider is inferred from the model name (`gpt*` / `o1*` / `o3*` / `o4*` → OpenAI, `claude*` → Anthropic). Top-level `ai_provider` is optional and, when set, must match the inferred provider. + +```yaml +# Optional cross-check (inferred from models when omitted) +ai_provider: "openai" + +default: + pipeline: + extract_sections: + model: "gpt-5.4-mini" + prompts: + system: | + ... + user: | + ... + + extract_caption_title: + model: "gpt-5.4-mini" + prompts: + system: | + ... + user: | + ... + + extract_panel_sequence: + model: "gpt-5.4-mini" + prompts: + system: | + ... + user: | + ... + + extract_data_sources: + model: "gpt-5.4-mini" + prompts: + system: | + ... + user: | + ... + + assign_panel_source: + model: "gpt-5.4-mini" + prompts: + system: | + ... + user: | + ... + + match_caption_panel: + model: "gpt-5.4-mini" + prompts: + system: | + ... + user: | + ... + + object_detection: # non-AI; unchanged + model_path: "data/models/panel_detection_model_no_labels.pt" + confidence_threshold: 0.25 +``` + +See `config.sample.yaml` for a minimal template and `config.dev.yaml` for the full development prompts. + +**Note:** `config.qc.yaml` keeps its own provider-specific layout; only the main curation pipeline uses the flat step schema above. + ## Pipeline Steps The soda-curation pipeline processes scientific manuscripts through the following detailed steps: @@ -857,6 +923,12 @@ For any questions or issues, please open an issue on the GitHub repository. We a ## Changelog +### 3.4.1 (2026-07-06) +- **mmqc-utils from PyPI**: Depend on published `mmqc-utils ^0.3.3` instead of a local wheel path so CI and Docker builds resolve the package without a sibling checkout. + +### 3.4.0 (2026-07-06) +- **Unified pipeline config**: Main pipeline steps use flat `model` + `prompts` blocks; provider is inferred from the model name. Per-step `openai`/`anthropic` nesting and YAML sampling fields are removed from the main config (QC config unchanged). + ### 3.3.0 (2026-06-10) - **mmqc_utils integration**: Manuscript HTML cleanup and figure image conversion now use `mmqc_utils` (`document_to_html`, `convert_to_bounded_jpeg`) for parity with soda-mmqc; removed legacy HTML normalization and PDF fallback paths. - **Python 3.12+**: Raised minimum Python to 3.12; Docker and CI updated accordingly. diff --git a/config.dev.yaml b/config.dev.yaml index a740655..42e9d96 100644 --- a/config.dev.yaml +++ b/config.dev.yaml @@ -3,7 +3,8 @@ ########################################################## include_qc: true environment: "dev" -# AI provider: "openai" (default) or "anthropic" +# Optional: cross-check runtime provider against models in pipeline steps. +# When omitted, provider is inferred from each step's model name (gpt* → openai, claude* → anthropic). ai_provider: "openai" default: &default pipeline: @@ -12,16 +13,8 @@ default: &default ########################################################## # Updated extract_sections prompts for smolagents implementation extract_sections: - openai: - # OpenAI-specific parameters remain the same - model: "gpt-5.4-mini" # Can be "gpt-4o", "gpt-4o-mini", or "gpt-5" - temperature: 0.1 - top_p: 1.0 - max_tokens: 2048 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-5.4-mini" + prompts: system: | You are a scientific text analyzer. Your job is EXTRACTIVE COPYING from the manuscript — not summarizing, paraphrasing, rewriting, or generating new text. @@ -98,15 +91,8 @@ default: &default # 2) Extract Individual Captions Step ########################################################## extract_caption_title: - openai: - model: "gpt-5.4-mini" - temperature: 0.1 - top_p: 1.0 - max_tokens: 4096 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-5.4-mini" + prompts: system: | You are an AI assistant specializing in EXTRACTIVE copying of figure captions from scientific manuscripts. This is NOT a generative or summarization task: you must copy text exactly as it appears in the input. @@ -151,15 +137,8 @@ default: &default Verify with the verification tool that `figure_caption` and `caption_title` are exact substrings of the input HTML before returning. extract_panel_sequence: - openai: - model: "gpt-5.4-mini" - temperature: 0.1 - top_p: 1.0 - max_tokens: 4096 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-5.4-mini" + prompts: system: | You extract figure panel structure from the figure caption text ONLY Output JSON with one object per panel, in strict left-to-right reading order of the caption. @@ -197,15 +176,8 @@ default: &default ########################################################## extract_data_sources: - openai: - model: "gpt-5.4-mini" - temperature: 0.5 - top_p: 0.5 - max_tokens: 2048 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-5.4-mini" + prompts: system: | You are an expert at extracting structured data source information from scientific Data Availability sections. Your task is to identify and extract information about databases, accession numbers, and URLs into a specific JSON format. @@ -308,15 +280,8 @@ default: &default # 4) Assign Panel Source Step ########################################################## assign_panel_source: - openai: - model: "gpt-5.4-mini" - temperature: 0.3 - top_p: 1.0 - max_tokens: 2048 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-5.4-mini" + prompts: system: | You are an expert AI assistant for analyzing scientific data organization, particularly for matching source data files to specific panels within scientific figures. Your task is to analyze file lists from source data ZIP files and determine which files correspond to which figure panels. @@ -382,15 +347,8 @@ default: &default # 5) Match Caption Panel Step ########################################################## match_caption_panel: - openai: - model: "gpt-5.4-mini" - temperature: 0.3 - top_p: 1.0 - max_tokens: 2048 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-5.4-mini" + prompts: system: | You match ONE cropped panel image to the correct entry from a fixed catalog of panels. Panel labels and panel_caption text were already taken from the figure caption in an earlier step — you must NOT rewrite or improve them. diff --git a/config.sample.yaml b/config.sample.yaml index d3ee58e..32b0c10 100644 --- a/config.sample.yaml +++ b/config.sample.yaml @@ -1,5 +1,8 @@ ########################################################## -# Default Profile +# Default Profile — main pipeline template (v3.3+ flat step schema) +# +# Each AI step: model + prompts only. Provider inferred from model name. +# See README "Main pipeline YAML shape" and config.dev.yaml for full prompts. ########################################################## default: &default pipeline: @@ -7,15 +10,8 @@ default: &default # 1) Extract Figure Legends and Data Availability Sections ########################################################## extract_sections: - openai: - model: "gpt-4o" - temperature: 0.1 - top_p: 1.0 - max_tokens: 2048 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-4o" + prompts: system: | # EXTRACTIVE task: copy contiguous spans from cleaned HTML manuscript input. # figure_legends and data_availability must be exact character-for-character @@ -32,15 +28,8 @@ default: &default # 2) Extract Individual Captions ########################################################## extract_caption_title: - openai: - model: "gpt-4o" - temperature: 0.1 - top_p: 1.0 - max_tokens: 4096 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-4o" + prompts: system: | # EXTRACTIVE task: copy figure caption verbatim from HTML input. # figure_caption and caption_title must be exact character-for-character @@ -54,15 +43,8 @@ default: &default # Reinforces extractive one-to-one HTML copy. extract_panel_sequence: - openai: - model: "gpt-4o" - temperature: 0.1 - top_p: 1.0 - max_tokens: 4096 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-4o" + prompts: system: | # System prompt for panel sequence extraction # Panel captions must be exact verbatim HTML substrings of the figure caption @@ -76,15 +58,8 @@ default: &default # 3) Extract Data Sources ########################################################## extract_data_sources: - openai: - model: "gpt-4o-mini" - temperature: 0.1 - top_p: 0.5 - max_tokens: 2048 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-4o-mini" + prompts: system: | # System prompt for data source extraction user: | @@ -95,15 +70,8 @@ default: &default # 4) Assign Panel Source ########################################################## assign_panel_source: - openai: - model: "gpt-4o" - temperature: 0.3 - top_p: 1.0 - max_tokens: 2048 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-4o" + prompts: system: | # System prompt for panel source assignment # Note: any caption text in the inputs is HTML-formatted @@ -116,15 +84,8 @@ default: &default # 5) Match Caption Panel ########################################################## match_caption_panel: - openai: - model: "gpt-4o" - temperature: 0.3 - top_p: 1.0 - max_tokens: 2048 - frequency_penalty: 0.0 - presence_penalty: 0.0 - json_mode: true - prompts: + model: "gpt-4o" + prompts: system: | # System prompt for matching panel images to captions # Caption and panel_caption strings are HTML; copy them exactly, including tags diff --git a/poetry.lock b/poetry.lock index 575c20d..76f8337 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1584,41 +1584,6 @@ files = [ [package.extras] all = ["flake8 (>=7.1.1)", "mypy (>=1.11.2)", "pytest (>=8.3.2)", "ruff (>=0.6.2)"] -[[package]] -name = "imagecodecs" -version = "2026.6.6" -description = "Image transformation, compression, and decompression codecs" -optional = false -python-versions = ">=3.12" -groups = ["main"] -files = [ - {file = "imagecodecs-2026.6.6-cp312-abi3-macosx_10_15_x86_64.whl", hash = "sha256:f2d35d6771251d128556fe144b9f8492a7189c4525d5fd0d7c8318d6fcbfb390"}, - {file = "imagecodecs-2026.6.6-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:c3052aed81b3d348c77163b475ab00252964a5f7caf74ae9c62d4ca90e247fe8"}, - {file = "imagecodecs-2026.6.6-cp312-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:907bb9b858c3dc1d5c4ff91385ea210ae90fa2fb00b44206643a24e20db742d2"}, - {file = "imagecodecs-2026.6.6-cp312-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:bca94547db076322c606f37d501f6e7ab3fa3f9eee3fcd6f233854c67014e9fb"}, - {file = "imagecodecs-2026.6.6-cp312-abi3-win32.whl", hash = "sha256:00596506fcea638063491db1baeef9a2445b3e0147bcbd76633561df702b1eb1"}, - {file = "imagecodecs-2026.6.6-cp312-abi3-win_amd64.whl", hash = "sha256:3ba2bcee53dabfb7d4dbd6f3c3ae731f3bfa43fa9c936d969886dff2f50e2160"}, - {file = "imagecodecs-2026.6.6-cp312-abi3-win_arm64.whl", hash = "sha256:031c5ea4d0c8bfacc87a5ca0ffdfee9a06b64265960d8b0b5d3f4a32434404c3"}, - {file = "imagecodecs-2026.6.6-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:5ee77053841dc7a8fbe59df774c0456fbe604edf2cd12c795c29deee1384ce8d"}, - {file = "imagecodecs-2026.6.6-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1fe0a7c8fb67bdd1bbb406c0e830dd1652acf6dce62620071a47e20f7d00619a"}, - {file = "imagecodecs-2026.6.6-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:2673fc78fc6e2178b13957eccaa2e97c716ea03607af6a2a44a2f09d00919f3f"}, - {file = "imagecodecs-2026.6.6-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:ed0760b9d428408798c010593ef204f88d949a9f57fa037623544e811c462c67"}, - {file = "imagecodecs-2026.6.6-cp314-cp314t-win32.whl", hash = "sha256:16540ed96854312d235bb86f885cb5f6ca2f4f1acde8e283d13a2bda2fb01e26"}, - {file = "imagecodecs-2026.6.6-cp314-cp314t-win_amd64.whl", hash = "sha256:325f72db99e2678c9ad0ccca8ac9f274b6ac7b4a1a5df71db0f96e23f964aeab"}, - {file = "imagecodecs-2026.6.6-cp314-cp314t-win_arm64.whl", hash = "sha256:fcd8b838bbe10c6b312b8aa67add4bb375bd1537ae4d92815c1ea522bef43cbc"}, - {file = "imagecodecs-2026.6.6-cp315-cp315t-win32.whl", hash = "sha256:b20be458a4d4a0283acaae17c94ad8128dfc6ce62fefdd55f1882c8e3a67850c"}, - {file = "imagecodecs-2026.6.6-cp315-cp315t-win_amd64.whl", hash = "sha256:51703586bb553d3b3ef9a0ab5c5b26432611c2bbb262f5f8fdeb19095c5e130f"}, - {file = "imagecodecs-2026.6.6-cp315-cp315t-win_arm64.whl", hash = "sha256:e4a059c2f57df75f99024bf9499b17b10d5a3916c5b5addba31317a523f120cd"}, - {file = "imagecodecs-2026.6.6.tar.gz", hash = "sha256:c7b014f5c0e73ff455700b1db271415d3c35a3bddedca6527cb2eefcd2fa8a10"}, -] - -[package.dependencies] -numpy = ">=2.1" - -[package.extras] -all = ["matplotlib", "numcodecs", "tifffile"] -test = ["backports.zstd", "bitshuffle", "blosc", "blosc2", "brotli", "czifile", "kerchunk", "liffile", "lz4", "numcodecs", "pyliblzfse", "pytest", "pytest-run-parallel", "python-lzf", "python-snappy", "tifffile", "zarr", "zopflipy"] - [[package]] name = "importlib-metadata" version = "7.0.0" @@ -2527,29 +2492,23 @@ files = [ [[package]] name = "mmqc-utils" -version = "0.3.5" +version = "0.3.3" description = "Reusable conversion utilities for images, documents, text, and related data types." optional = false python-versions = ">=3.12" groups = ["main"] files = [ - {file = "mmqc_utils-0.3.5-py3-none-any.whl", hash = "sha256:dde47167111a0cd5c3b3662960d3021c61944be24e65b47b0438eba52d27d25c"}, + {file = "mmqc_utils-0.3.3-py3-none-any.whl", hash = "sha256:4b7031a7ca7d1cdfc2c9beb599eec8076e31980e6fadf3883f2a110c63a49d37"}, + {file = "mmqc_utils-0.3.3.tar.gz", hash = "sha256:65d7826e5b0fbc457b87309d29dfacc80e855c6cc7ca297335d31ce767c944ef"}, ] [package.dependencies] beautifulsoup4 = ">=4.15.0" -imagecodecs = ">=2026.6.6" -pillow = ">=12.2.0" pypandoc-binary = ">=1.17" pypdf = ">=6.10.2" rapidfuzz = ">=3.14.3" -tifffile = ">=2026.6.1" Wand = ">=0.6.13" -[package.source] -type = "file" -url = "../soda-mmQC/mmqc_utils/dist/mmqc_utils-0.3.5-py3-none-any.whl" - [[package]] name = "more-itertools" version = "10.8.0" @@ -5463,29 +5422,6 @@ files = [ doc = ["reno", "sphinx"] test = ["pytest", "tornado (>=4.5)", "typeguard"] -[[package]] -name = "tifffile" -version = "2026.6.1" -description = "Read and write TIFF files" -optional = false -python-versions = ">=3.12" -groups = ["main"] -files = [ - {file = "tifffile-2026.6.1-py3-none-any.whl", hash = "sha256:0d7382d2769b855b81ce358528e2b40c16d48aa39031746efa81215205332a8d"}, - {file = "tifffile-2026.6.1.tar.gz", hash = "sha256:626c892c0e899d959b9438e7c0e1491dc154a7fead1f1f37a991724a50eceba9"}, -] - -[package.dependencies] -numpy = ">=2.1" - -[package.extras] -all = ["fsspec", "imagecodecs (>=2026.5.10)", "kerchunk", "lxml", "matplotlib", "xarray", "zarr (>=3.2.0)"] -codecs = ["imagecodecs (>=2026.5.10)"] -plot = ["matplotlib"] -test = ["cmapfile", "czifile", "dask", "fsspec", "imagecodecs", "kerchunk", "lfdfiles", "lxml", "ndtiff", "oiffile", "psdtags", "pytest", "requests", "roifile", "xarray", "zarr (>=3.2.0)"] -xml = ["lxml"] -zarr = ["fsspec", "kerchunk", "zarr (>=3.2.0)"] - [[package]] name = "tiktoken" version = "0.12.0" @@ -6220,4 +6156,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = ">=3.12,<3.14" -content-hash = "ea1fc0d42a2859ce72596a8c1409f45616d65bb6181d8f295c58e28dd22a06aa" +content-hash = "64613c8e54e53ceb093657aa234ab432eb2b46459588fc635e27cfa27f40f547" diff --git a/pyproject.toml b/pyproject.toml index 4362595..a795e4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api" [project] name = "soda-curation" -version = "3.3.0" +version = "3.4.1" description = "A professional Python package for data curation with AI capabilities" readme = "README.md" requires-python = ">=3.12,<3.14" @@ -25,8 +25,7 @@ packages = [{include = "soda_curation", from = "src"}] [tool.poetry.dependencies] python = ">=3.12,<3.14" -# TODO: switch to the PyPI release of mmqc-utils once it is published. -mmqc-utils = {path = "/Users/jabreu/PycharmProjects/soda-mmQC/mmqc_utils/dist/mmqc_utils-0.3.5-py3-none-any.whl"} +mmqc-utils = "^0.3.3" openai = "^1.43.0" pyyaml = "^6.0.2" python-docx = "^0.8.11" diff --git a/src/soda_curation/config.py b/src/soda_curation/config.py index 2b80cbc..d016b96 100644 --- a/src/soda_curation/config.py +++ b/src/soda_curation/config.py @@ -116,17 +116,12 @@ def get_pipeline_config(self, step: PipelineStep) -> Dict[str, Any]: if step == PipelineStep.OBJECT_DETECTION: return cast(Dict[str, Any], self.config["pipeline"]["object_detection"]) - # Get provider-specific configuration for AI steps - provider = os.getenv("MODEL_PROVIDER", "openai") - step_config = self.config["pipeline"][step.value] + from .pipeline.step_config import resolve_step_config - if provider not in step_config: - raise ConfigurationError( - f"Provider {provider} not configured for step {step.value}" - ) - - # Merge provider configuration with API key - config = step_config[provider].copy() + step_config = self.config["pipeline"][step.value] + resolved = resolve_step_config(step_config) + provider = resolved["provider"] + config = {"model": resolved["model"], "prompts": resolved["prompts"]} config["api_key"] = os.getenv(f"{provider.upper()}_API_KEY") return cast(Dict[str, Any], config) diff --git a/src/soda_curation/main.py b/src/soda_curation/main.py index 26bd797..ecf7b8a 100755 --- a/src/soda_curation/main.py +++ b/src/soda_curation/main.py @@ -56,6 +56,7 @@ MatchPanelCaptionOpenAI, ) from .pipeline.prompt_handler import PromptHandler +from .pipeline.step_config import AI_PIPELINE_STEPS, resolve_pipeline_provider # Import QC module (to be implemented) from .qc.qc_pipeline import QCPipeline @@ -140,51 +141,39 @@ def _execute_pipeline_step( return None -def _validate_ai_provider_config(config: dict, ai_provider: str, run_id: str) -> None: - """Ensure runtime provider selection is explicit and configuration is consistent.""" - if ai_provider not in SUPPORTED_AI_PROVIDERS: - raise ValueError( - f"Unsupported ai_provider '{ai_provider}'. Expected one of {sorted(SUPPORTED_AI_PROVIDERS)}." - ) +def _resolve_ai_provider(config: dict, run_id: str) -> str: + """ + Infer the runtime AI provider from pipeline step model names. + Optional top-level ``ai_provider`` must match the inferred value when set. + """ pipeline_config = config.get("pipeline", {}) if not isinstance(pipeline_config, dict): raise ValueError("Invalid configuration: missing `pipeline` section.") - ignored_provider = "anthropic" if ai_provider == "openai" else "openai" - for step_name in AI_PROVIDER_STEPS: - step_config = pipeline_config.get(step_name, {}) - if not isinstance(step_config, dict): + inferred = resolve_pipeline_provider(pipeline_config, steps=AI_PIPELINE_STEPS) + configured = config.get("ai_provider") + if configured is not None: + configured = str(configured).lower() + if configured not in SUPPORTED_AI_PROVIDERS: raise ValueError( - f"Invalid configuration for step '{step_name}': expected mapping." + f"Unsupported ai_provider '{configured}'. " + f"Expected one of {sorted(SUPPORTED_AI_PROVIDERS)}." ) - - has_openai = "openai" in step_config - has_anthropic = "anthropic" in step_config - assert has_openai or has_anthropic, ( - f"Configuration error in step '{step_name}': no provider block found. " - "Define exactly one of 'openai' or 'anthropic'." - ) - assert not (has_openai and has_anthropic), ( - f"Configuration error in step '{step_name}': both 'openai' and 'anthropic' " - "are defined. Define only one provider per step." - ) - - configured_provider = "openai" if has_openai else "anthropic" - if configured_provider != ai_provider: + if configured != inferred: raise ValueError( - "Selected provider does not match step configuration. " - f"step={step_name}, configured_provider={configured_provider}, " - f"ai_provider={ai_provider}, ignored_provider={ignored_provider}" + "Configured ai_provider does not match models in pipeline steps. " + f"ai_provider={configured}, inferred_from_models={inferred}" ) logger.info( - "AI provider configuration validated", + "AI provider inferred from pipeline models", extra={ "run_id": run_id, - "ai_provider": ai_provider, + "ai_provider": inferred, }, ) + return inferred def run_qc_pipeline_async( @@ -267,13 +256,11 @@ def main(zip_path: str, config_path: str, output_path: Optional[str] = None) -> zip_structure.manuscript_text = manuscript_content prompt_handler = PromptHandler(config_loader.config["pipeline"]) - # Select AI provider (default: openai) - ai_provider = config_loader.config.get("ai_provider", "openai").lower() + ai_provider = _resolve_ai_provider(config_loader.config, run_id) logger.info( f"Using AI provider: {ai_provider}", extra={"run_id": run_id, "ai_provider": ai_provider}, ) - _validate_ai_provider_config(config_loader.config, ai_provider, run_id) zip_structure.ai_provider = ai_provider # Extract relevant sections for the pipeline diff --git a/src/soda_curation/pipeline/assign_panel_source/assign_panel_source_anthropic.py b/src/soda_curation/pipeline/assign_panel_source/assign_panel_source_anthropic.py index 8d85f28..36bd592 100644 --- a/src/soda_curation/pipeline/assign_panel_source/assign_panel_source_anthropic.py +++ b/src/soda_curation/pipeline/assign_panel_source/assign_panel_source_anthropic.py @@ -11,6 +11,7 @@ from ..anthropic_utils import call_anthropic, validate_anthropic_model from ..cost_tracking import update_token_usage from ..prompt_handler import PromptHandler +from ..step_config import resolve_step_config from .assign_panel_source_base import ( AsignedFiles, AsignedFilesList, @@ -30,9 +31,12 @@ def __init__( self.client = anthropic.Anthropic() def _validate_config(self) -> None: - """Validate Anthropic configuration parameters.""" - config_ = self.config["pipeline"]["assign_panel_source"]["anthropic"] - validate_anthropic_model(config_.get("model", "claude-sonnet-4-6")) + resolved = resolve_step_config(self.config["pipeline"]["assign_panel_source"]) + if resolved["provider"] != "anthropic": + raise ValueError( + f"assign_panel_source model '{resolved['model']}' requires OpenAI." + ) + validate_anthropic_model(resolved["model"]) def call_ai_service(self, prompt: str, allowed_files: List) -> AsignedFilesList: """Call Claude with the given prompt and return assigned files.""" @@ -52,16 +56,15 @@ def call_ai_service(self, prompt: str, allowed_files: List) -> AsignedFilesList: {"role": "user", "content": prompt}, ] - config_ = self.config["pipeline"]["assign_panel_source"]["anthropic"] - model_ = config_.get("model", "claude-sonnet-4-6") + model_ = resolve_step_config(self.config["pipeline"]["assign_panel_source"])[ + "model" + ] response = call_anthropic( client=self.client, model=model_, messages=messages, response_format=AsignedFilesList, - temperature=config_.get("temperature", 0.3), - max_tokens=config_.get("max_tokens", 2048), operation="main.assign_panel_source", request_metadata={ "provider": "anthropic", diff --git a/src/soda_curation/pipeline/assign_panel_source/assign_panel_source_openai.py b/src/soda_curation/pipeline/assign_panel_source/assign_panel_source_openai.py index ad9f35b..a4c7cc4 100644 --- a/src/soda_curation/pipeline/assign_panel_source/assign_panel_source_openai.py +++ b/src/soda_curation/pipeline/assign_panel_source/assign_panel_source_openai.py @@ -8,8 +8,9 @@ from ..ai_observability import summarize_text from ..cost_tracking import update_token_usage -from ..openai_utils import DEFAULT_OPENAI_MODEL, call_openai, validate_model_config +from ..openai_utils import call_openai from ..prompt_handler import PromptHandler +from ..step_config import resolve_step_config from .assign_panel_source_base import ( AsignedFiles, AsignedFilesList, @@ -28,10 +29,11 @@ def __init__( self.client = openai.OpenAI() def _validate_config(self) -> None: - """Validate OpenAI configuration parameters.""" - config_ = self.config["pipeline"]["assign_panel_source"]["openai"] - model = config_.get("model", DEFAULT_OPENAI_MODEL) - validate_model_config(model, config_) + resolved = resolve_step_config(self.config["pipeline"]["assign_panel_source"]) + if resolved["provider"] != "openai": + raise ValueError( + f"assign_panel_source model '{resolved['model']}' requires Anthropic." + ) def call_ai_service(self, prompt: str, allowed_files: List) -> AsignedFilesList: """Call OpenAI service with the given prompt.""" @@ -53,18 +55,15 @@ def call_ai_service(self, prompt: str, allowed_files: List) -> AsignedFilesList: {"role": "user", "content": prompt}, ] - config_ = self.config["pipeline"]["assign_panel_source"]["openai"] - model_ = config_.get("model", "gpt-4o") + model_ = resolve_step_config(self.config["pipeline"]["assign_panel_source"])[ + "model" + ] response = call_openai( client=self.client, model=model_, messages=messages, response_format=AsignedFilesList, # Ensure the response is in JSON format - temperature=config_.get("temperature", 0.3), - top_p=config_.get("top_p", 1.0), - frequency_penalty=config_.get("frequency_penalty", 0), - presence_penalty=config_.get("presence_penalty", 0), operation="main.assign_panel_source", request_metadata={ "provider": "openai", diff --git a/src/soda_curation/pipeline/data_availability/data_availability_anthropic.py b/src/soda_curation/pipeline/data_availability/data_availability_anthropic.py index 5c64a30..ef5b664 100644 --- a/src/soda_curation/pipeline/data_availability/data_availability_anthropic.py +++ b/src/soda_curation/pipeline/data_availability/data_availability_anthropic.py @@ -11,6 +11,7 @@ from ..anthropic_utils import call_anthropic, validate_anthropic_model from ..cost_tracking import update_token_usage from ..manuscript_structure.manuscript_structure import ZipStructure +from ..step_config import resolve_step_config from .data_availability_base import DataAvailabilityExtractor, ExtractDataSources logger = logging.getLogger(__name__) @@ -37,9 +38,12 @@ def _load_database_registry(self) -> Dict[Any, Any]: return {"databases": []} def _validate_config(self) -> None: - """Validate Anthropic configuration parameters.""" - config_ = self.config["pipeline"]["extract_data_sources"]["anthropic"] - validate_anthropic_model(config_.get("model", "claude-sonnet-4-6")) + resolved = resolve_step_config(self.config["pipeline"]["extract_data_sources"]) + if resolved["provider"] != "anthropic": + raise ValueError( + f"extract_data_sources model '{resolved['model']}' requires OpenAI." + ) + validate_anthropic_model(resolved["model"]) def extract_data_sources( self, section_text: str, zip_structure: ZipStructure @@ -71,16 +75,15 @@ def extract_data_sources( {"role": "user", "content": prompts["user"]}, ] - config_ = self.config["pipeline"]["extract_data_sources"]["anthropic"] - model_ = config_.get("model", "claude-sonnet-4-6") + model_ = resolve_step_config(self.config["pipeline"]["extract_data_sources"])[ + "model" + ] response = call_anthropic( client=self.client, model=model_, messages=messages, response_format=ExtractDataSources, - temperature=config_.get("temperature", 0.1), - max_tokens=config_.get("max_tokens", 2048), operation="main.extract_data_sources", request_metadata={ "registry_database_count": len( diff --git a/src/soda_curation/pipeline/data_availability/data_availability_openai.py b/src/soda_curation/pipeline/data_availability/data_availability_openai.py index 3afe325..01a7710 100644 --- a/src/soda_curation/pipeline/data_availability/data_availability_openai.py +++ b/src/soda_curation/pipeline/data_availability/data_availability_openai.py @@ -11,7 +11,8 @@ from ..ai_observability import summarize_text from ..cost_tracking import update_token_usage from ..manuscript_structure.manuscript_structure import ZipStructure -from ..openai_utils import DEFAULT_OPENAI_MODEL, call_openai, validate_model_config +from ..openai_utils import call_openai +from ..step_config import resolve_step_config from .data_availability_base import DataAvailabilityExtractor, ExtractDataSources logger = logging.getLogger(__name__) @@ -53,10 +54,11 @@ def _load_database_registry(self) -> Dict[Any, Any]: return {"databases": []} def _validate_config(self) -> None: - """Validate OpenAI configuration parameters.""" - config_ = self.config["pipeline"]["extract_data_sources"]["openai"] - model = config_.get("model", DEFAULT_OPENAI_MODEL) - validate_model_config(model, config_) + resolved = resolve_step_config(self.config["pipeline"]["extract_data_sources"]) + if resolved["provider"] != "openai": + raise ValueError( + f"extract_data_sources model '{resolved['model']}' requires Anthropic." + ) def extract_data_sources( self, section_text: str, zip_structure: ZipStructure @@ -92,18 +94,15 @@ def extract_data_sources( {"role": "user", "content": prompts["user"]}, ] - config_ = self.config["pipeline"]["extract_data_sources"]["openai"] - model_ = config_.get("model", "gpt-4o") + model_ = resolve_step_config(self.config["pipeline"]["extract_data_sources"])[ + "model" + ] response = call_openai( client=self.client, model=model_, messages=messages, response_format=ExtractDataSources, - temperature=config_.get("temperature", 0.1), - top_p=config_.get("top_p", 1.0), - frequency_penalty=config_.get("frequency_penalty", 0), - presence_penalty=config_.get("presence_penalty", 0), operation="main.extract_data_sources", request_metadata={ "registry_database_count": len( diff --git a/src/soda_curation/pipeline/extract_captions/extract_captions_anthropic.py b/src/soda_curation/pipeline/extract_captions/extract_captions_anthropic.py index 6fe9bd7..0630740 100644 --- a/src/soda_curation/pipeline/extract_captions/extract_captions_anthropic.py +++ b/src/soda_curation/pipeline/extract_captions/extract_captions_anthropic.py @@ -10,6 +10,7 @@ from ..anthropic_utils import call_anthropic, validate_anthropic_model from ..cost_tracking import update_token_usage from ..manuscript_structure.manuscript_structure import TokenUsage, ZipStructure +from ..step_config import resolve_step_config from .extract_captions_base import ( CaptionExtraction, FigureCaptionExtractor, @@ -26,13 +27,15 @@ class FigureCaptionExtractorAnthropic(FigureCaptionExtractor): def __init__(self, config: Dict[str, Any], prompt_handler): super().__init__(config, prompt_handler) self.client = anthropic.Anthropic() - self.caption_config = config["pipeline"]["extract_caption_title"]["anthropic"] - self.panel_config = config["pipeline"]["extract_panel_sequence"]["anthropic"] def _validate_config(self) -> None: - for step in ["extract_caption_title", "extract_panel_sequence"]: - config_ = self.config["pipeline"][step]["anthropic"] - validate_anthropic_model(config_.get("model", "claude-sonnet-4-6")) + for step in ("extract_caption_title", "extract_panel_sequence"): + resolved = resolve_step_config(self.config["pipeline"][step]) + if resolved["provider"] != "anthropic": + raise ValueError( + f"{step} model '{resolved['model']}' requires the OpenAI extractor." + ) + validate_anthropic_model(resolved["model"]) def extract_figure_caption( self, figure_label: str, all_captions: str, zip_structure: ZipStructure @@ -65,16 +68,15 @@ def extract_figure_caption( {"role": "user", "content": prompts["user"]}, ] - config_ = self.caption_config - model_ = config_.get("model", "claude-sonnet-4-6") + model_ = resolve_step_config(self.config["pipeline"]["extract_caption_title"])[ + "model" + ] response = call_anthropic( client=self.client, model=model_, messages=messages, response_format=CaptionExtraction, - temperature=config_.get("temperature", 0.1), - max_tokens=config_.get("max_tokens", 4096), operation="main.extract_caption_title", request_metadata={"figure_label": figure_label}, ) @@ -122,16 +124,15 @@ def extract_figure_panels( {"role": "user", "content": prompts["user"]}, ] - config_ = self.panel_config - model_ = config_.get("model", "claude-sonnet-4-6") + model_ = resolve_step_config(self.config["pipeline"]["extract_panel_sequence"])[ + "model" + ] response = call_anthropic( client=self.client, model=model_, messages=messages, response_format=PanelExtraction, - temperature=config_.get("temperature", 0.1), - max_tokens=config_.get("max_tokens", 4096), operation="main.extract_panel_sequence", request_metadata={"figure_label": figure_label}, ) diff --git a/src/soda_curation/pipeline/extract_captions/extract_captions_openai.py b/src/soda_curation/pipeline/extract_captions/extract_captions_openai.py index de77eb5..15be859 100644 --- a/src/soda_curation/pipeline/extract_captions/extract_captions_openai.py +++ b/src/soda_curation/pipeline/extract_captions/extract_captions_openai.py @@ -10,7 +10,8 @@ from ..ai_observability import summarize_text from ..cost_tracking import update_token_usage from ..manuscript_structure.manuscript_structure import TokenUsage, ZipStructure -from ..openai_utils import DEFAULT_OPENAI_MODEL, call_openai, validate_model_config +from ..openai_utils import call_openai +from ..step_config import resolve_step_config from .extract_captions_base import ( CaptionExtraction, FigureCaptionExtractor, @@ -32,14 +33,14 @@ def __init__(self, config: Dict[str, Any], prompt_handler): raise ValueError("OPENAI_API_KEY environment variable is not set") self.client = openai.OpenAI(api_key=api_key) - self.caption_config = config["pipeline"]["extract_caption_title"]["openai"] - self.panel_config = config["pipeline"]["extract_panel_sequence"]["openai"] def _validate_config(self) -> None: - for step in ["extract_caption_title", "extract_panel_sequence"]: - config_ = self.config["pipeline"][step]["openai"] - model = config_.get("model", DEFAULT_OPENAI_MODEL) - validate_model_config(model, config_) + for step in ("extract_caption_title", "extract_panel_sequence"): + resolved = resolve_step_config(self.config["pipeline"][step]) + if resolved["provider"] != "openai": + raise ValueError( + f"{step} model '{resolved['model']}' requires the Anthropic extractor." + ) def extract_figure_caption( self, figure_label: str, all_captions: str, zip_structure: ZipStructure @@ -71,18 +72,15 @@ def extract_figure_caption( {"role": "user", "content": prompts["user"]}, ] - config_ = self.caption_config - model_ = config_.get("model", DEFAULT_OPENAI_MODEL) + model_ = resolve_step_config(self.config["pipeline"]["extract_caption_title"])[ + "model" + ] response = call_openai( client=self.client, model=model_, messages=messages, response_format=CaptionExtraction, - temperature=config_.get("temperature", 0.1), - top_p=config_.get("top_p", 1.0), - frequency_penalty=config_.get("frequency_penalty", 0), - presence_penalty=config_.get("presence_penalty", 0), operation="main.extract_caption_title", request_metadata={"figure_label": figure_label}, ) @@ -142,18 +140,15 @@ def extract_figure_panels( {"role": "user", "content": prompts["user"]}, ] - config_ = self.panel_config - model_ = config_.get("model", DEFAULT_OPENAI_MODEL) + model_ = resolve_step_config(self.config["pipeline"]["extract_panel_sequence"])[ + "model" + ] response = call_openai( client=self.client, model=model_, messages=messages, response_format=PanelExtraction, - temperature=config_.get("temperature", 0.1), - top_p=config_.get("top_p", 1.0), - frequency_penalty=config_.get("frequency_penalty", 0), - presence_penalty=config_.get("presence_penalty", 0), operation="main.extract_panel_sequence", request_metadata={"figure_label": figure_label}, ) diff --git a/src/soda_curation/pipeline/extract_sections/extract_sections_anthropic.py b/src/soda_curation/pipeline/extract_sections/extract_sections_anthropic.py index ae53cdd..c49353f 100644 --- a/src/soda_curation/pipeline/extract_sections/extract_sections_anthropic.py +++ b/src/soda_curation/pipeline/extract_sections/extract_sections_anthropic.py @@ -11,6 +11,7 @@ from ..anthropic_utils import call_anthropic, validate_anthropic_model from ..cost_tracking import update_token_usage from ..manuscript_structure.manuscript_structure import ZipStructure +from ..step_config import resolve_step_config from .extract_sections_base import ExtractedSections, SectionExtractor logger = logging.getLogger(__name__) @@ -24,9 +25,13 @@ def __init__(self, config: Dict, prompt_handler): self.client = anthropic.Anthropic() def _validate_config(self) -> None: - """Validate Anthropic configuration parameters.""" - config_ = self.config["pipeline"]["extract_sections"]["anthropic"] - validate_anthropic_model(config_.get("model", "claude-sonnet-4-6")) + """Validate step has a model configured for Anthropic.""" + resolved = resolve_step_config(self.config["pipeline"]["extract_sections"]) + if resolved["provider"] != "anthropic": + raise ValueError( + f"extract_sections model '{resolved['model']}' requires the OpenAI extractor." + ) + validate_anthropic_model(resolved["model"]) def extract_sections( self, @@ -59,16 +64,15 @@ def extract_sections( {"role": "user", "content": prompts["user"]}, ] - config_ = self.config["pipeline"]["extract_sections"]["anthropic"] - model_ = config_.get("model", "claude-sonnet-4-6") + model_ = resolve_step_config(self.config["pipeline"]["extract_sections"])[ + "model" + ] response = call_anthropic( client=self.client, model=model_, messages=messages, response_format=ExtractedSections, - temperature=config_.get("temperature", 0.1), - max_tokens=config_.get("max_tokens", 2048), operation="main.extract_sections", request_metadata={"figure_count": len(zip_structure.figures)}, ) diff --git a/src/soda_curation/pipeline/extract_sections/extract_sections_openai.py b/src/soda_curation/pipeline/extract_sections/extract_sections_openai.py index 7fb4e6c..f769b15 100644 --- a/src/soda_curation/pipeline/extract_sections/extract_sections_openai.py +++ b/src/soda_curation/pipeline/extract_sections/extract_sections_openai.py @@ -11,7 +11,8 @@ from ..ai_observability import summarize_text from ..cost_tracking import update_token_usage from ..manuscript_structure.manuscript_structure import ZipStructure -from ..openai_utils import DEFAULT_OPENAI_MODEL, call_openai, validate_model_config +from ..openai_utils import call_openai +from ..step_config import resolve_step_config from .extract_sections_base import ExtractedSections, SectionExtractor logger = logging.getLogger(__name__) @@ -32,10 +33,12 @@ def __init__(self, config: Dict, prompt_handler): self.client = openai.OpenAI(api_key=api_key) def _validate_config(self) -> None: - """Validate OpenAI configuration parameters.""" - config_ = self.config["pipeline"]["extract_sections"]["openai"] - model = config_.get("model", DEFAULT_OPENAI_MODEL) - validate_model_config(model, config_) + """Validate step has a model configured for OpenAI.""" + resolved = resolve_step_config(self.config["pipeline"]["extract_sections"]) + if resolved["provider"] != "openai": + raise ValueError( + f"extract_sections model '{resolved['model']}' requires the Anthropic extractor." + ) def extract_sections( self, @@ -70,18 +73,15 @@ def extract_sections( {"role": "user", "content": prompts["user"]}, ] - config_ = self.config["pipeline"]["extract_sections"]["openai"] - model_ = config_.get("model", DEFAULT_OPENAI_MODEL) + model_ = resolve_step_config(self.config["pipeline"]["extract_sections"])[ + "model" + ] response = call_openai( client=self.client, model=model_, messages=messages, response_format=ExtractedSections, - temperature=config_.get("temperature", 0.1), - top_p=config_.get("top_p", 1.0), - frequency_penalty=config_.get("frequency_penalty", 0), - presence_penalty=config_.get("presence_penalty", 0), operation="main.extract_sections", request_metadata={ "figure_count": len(zip_structure.figures), diff --git a/src/soda_curation/pipeline/match_caption_panel/match_caption_panel_anthropic.py b/src/soda_curation/pipeline/match_caption_panel/match_caption_panel_anthropic.py index 1ebf2e3..85519ab 100644 --- a/src/soda_curation/pipeline/match_caption_panel/match_caption_panel_anthropic.py +++ b/src/soda_curation/pipeline/match_caption_panel/match_caption_panel_anthropic.py @@ -11,6 +11,7 @@ from ..anthropic_utils import call_anthropic, validate_anthropic_model from ..cost_tracking import update_token_usage from ..manuscript_structure.manuscript_structure import Panel +from ..step_config import resolve_step_config from .match_caption_panel_base import MatchPanelCaption, PanelObject logger = logging.getLogger(__name__) @@ -22,12 +23,14 @@ class MatchPanelCaptionAnthropic(MatchPanelCaption): def __init__(self, config: Dict[str, Any], prompt_handler: Any, extract_dir: Path): super().__init__(config, prompt_handler, extract_dir) self.client = anthropic.Anthropic() - self.anthropic_config = config["pipeline"]["match_caption_panel"]["anthropic"] def _validate_config(self) -> None: - """Validate Anthropic configuration parameters.""" - config_ = self.config["pipeline"]["match_caption_panel"]["anthropic"] - validate_anthropic_model(config_.get("model", "claude-sonnet-4-6")) + resolved = resolve_step_config(self.config["pipeline"]["match_caption_panel"]) + if resolved["provider"] != "anthropic": + raise ValueError( + f"match_caption_panel model '{resolved['model']}' requires OpenAI." + ) + validate_anthropic_model(resolved["model"]) def _match_panel_caption( self, @@ -58,7 +61,9 @@ def _match_panel_caption( } prompts = self.prompt_handler.get_prompt("match_caption_panel", variables) - model = self.anthropic_config.get("model", "claude-sonnet-4-6") + model = resolve_step_config(self.config["pipeline"]["match_caption_panel"])[ + "model" + ] logger.info( "Preparing panel-caption vision request", extra={ @@ -89,8 +94,6 @@ def _match_panel_caption( }, ], response_format=PanelObject, - temperature=self.anthropic_config.get("temperature", 0.1), - max_tokens=self.anthropic_config.get("max_tokens", 512), operation="main.match_caption_panel", request_metadata={ "provider": "anthropic", diff --git a/src/soda_curation/pipeline/match_caption_panel/match_caption_panel_openai.py b/src/soda_curation/pipeline/match_caption_panel/match_caption_panel_openai.py index 7f5c061..c34e322 100644 --- a/src/soda_curation/pipeline/match_caption_panel/match_caption_panel_openai.py +++ b/src/soda_curation/pipeline/match_caption_panel/match_caption_panel_openai.py @@ -16,7 +16,8 @@ from ..ai_observability import summarize_text from ..cost_tracking import update_token_usage from ..manuscript_structure.manuscript_structure import Panel, ZipStructure -from ..openai_utils import DEFAULT_OPENAI_MODEL, call_openai, validate_model_config +from ..openai_utils import call_openai +from ..step_config import resolve_step_config from .match_caption_panel_base import MatchPanelCaption, PanelObject logger = logging.getLogger(__name__) @@ -45,14 +46,12 @@ def __init__(self, config: Dict[str, Any], prompt_handler: Any, extract_dir: Pat # Initialize OpenAI client self.client = openai.OpenAI() - # Get OpenAI specific config - self.openai_config = config["pipeline"]["match_caption_panel"]["openai"] - def _validate_config(self) -> None: - """Validate OpenAI configuration parameters.""" - config_ = self.config["pipeline"]["match_caption_panel"]["openai"] - model = config_.get("model", DEFAULT_OPENAI_MODEL) - validate_model_config(model, config_) + resolved = resolve_step_config(self.config["pipeline"]["match_caption_panel"]) + if resolved["provider"] != "openai": + raise ValueError( + f"match_caption_panel model '{resolved['model']}' requires Anthropic." + ) def _match_panel_caption( self, @@ -82,7 +81,9 @@ def _match_panel_caption( "allowed_panel_catalog_json": json.dumps(catalog, ensure_ascii=False), } prompts = self.prompt_handler.get_prompt("match_caption_panel", variables) - model = self.openai_config.get("model", DEFAULT_OPENAI_MODEL) + model = resolve_step_config(self.config["pipeline"]["match_caption_panel"])[ + "model" + ] logger.info( "Preparing panel-caption vision request", extra={ @@ -113,11 +114,6 @@ def _match_panel_caption( }, ], response_format=PanelObject, - temperature=self.openai_config.get("temperature", 0.1), - top_p=self.openai_config.get("top_p", 1.0), - frequency_penalty=self.openai_config.get("frequency_penalty", 0), - presence_penalty=self.openai_config.get("presence_penalty", 0), - max_tokens=self.openai_config.get("max_tokens", 512), operation="main.match_caption_panel", request_metadata={ "provider": "openai", diff --git a/src/soda_curation/pipeline/prompt_handler.py b/src/soda_curation/pipeline/prompt_handler.py index 99a841a..85a4584 100644 --- a/src/soda_curation/pipeline/prompt_handler.py +++ b/src/soda_curation/pipeline/prompt_handler.py @@ -4,17 +4,17 @@ from string import Template from typing import Dict -logger = logging.getLogger(__name__) +from .step_config import get_step_prompts, is_ai_step -SUPPORTED_PROVIDERS = ["openai", "anthropic"] +logger = logging.getLogger(__name__) class PromptHandler: """ Handle prompt loading and template substitution for pipeline components. - Each pipeline step has a single set of 'system' and 'user' prompts, - under a provider-specific config block (e.g., `openai`, `anthropic`). + Each AI pipeline step has a single shared ``prompts`` block (``system`` + ``user``) + alongside a ``model`` field. Provider-specific nesting is no longer required. """ def __init__(self, pipeline_config: Dict): @@ -22,20 +22,8 @@ def __init__(self, pipeline_config: Dict): Initialize with pipeline configuration. Args: - pipeline_config (Dict): A dict where keys are step names, - and values contain provider config, e.g.: - { - "extract_individual_captions": { - "openai": { - "model": "...", - "prompts": { - "system": "some system prompt", - "user": "some user prompt" - } - } - }, - ... - } + pipeline_config (Dict): ``config["pipeline"]`` mapping step names to + step config (``model`` + ``prompts`` for AI steps). """ self.pipeline_config = pipeline_config or {} self.validate_prompts() @@ -45,36 +33,23 @@ def get_prompt(self, step: str, variables: Dict) -> Dict[str, str]: Retrieve the system/user prompts for a given step, substituting variables. Args: - step (str): Pipeline step name (e.g., "extract_individual_captions"). + step (str): Pipeline step name (e.g., ``extract_sections``). variables (Dict): Template variables to substitute. Returns: - Dict[str, str]: { "system": "...", "user": "..." } + Dict[str, str]: ``{"system": "...", "user": "..."}`` Raises: KeyError: If step not found or if prompts are missing. """ - # Get the step configuration step_config = self.pipeline_config.get(step) if not step_config: raise KeyError(f"No configuration found for step: '{step}'") - # Identify the provider config (e.g., "openai" or "anthropic") - provider_config = self._get_provider_config(step_config) - if not provider_config: - raise KeyError(f"No recognized provider config found for step '{step}'") - - # Get the single system/user prompt pair - prompts = provider_config.get("prompts", {}) - if not prompts: - raise KeyError(f"No prompts found in step '{step}' under provider config") - - system_str = prompts.get("system", "") - user_str = prompts.get("user", "") - + prompts = get_step_prompts(step_config) return { - "system": Template(system_str).safe_substitute(variables), - "user": Template(user_str).safe_substitute(variables), + "system": Template(prompts["system"]).safe_substitute(variables), + "user": Template(prompts["user"]).safe_substitute(variables), } def validate_prompts(self) -> None: @@ -88,41 +63,14 @@ def validate_prompts(self) -> None: raise ValueError("No pipeline configuration provided") for step_name, step_config in self.pipeline_config.items(): - # Skip steps that have no recognized provider (like object_detection, etc.) - provider_config = self._get_provider_config(step_config) - if not provider_config: - continue # Not an AI step or no recognized provider, skip - - # Check there's a 'prompts' dictionary with at least system and user - prompts = provider_config.get("prompts") - if not isinstance(prompts, dict): - raise ValueError( - f"No valid 'prompts' dict in step '{step_name}' under provider config" - ) - - # We expect exactly "system" and "user" - if "system" not in prompts or "user" not in prompts: - raise ValueError( - f"Missing required 'system'/'user' prompts in step '{step_name}'" - ) - - # Validate templates by trying to compile them - for prompt_type in ["system", "user"]: - template_str = prompts[prompt_type] + if not is_ai_step(step_config): + continue + + prompts = get_step_prompts(step_config) + for prompt_type in ("system", "user"): try: - Template(template_str) + Template(prompts[prompt_type]) except Exception as e: raise ValueError( f"Invalid template in step '{step_name}' -> '{prompt_type}': {e}" - ) - - @staticmethod - def _get_provider_config(step_config: Dict) -> Dict: - """ - Return the first recognized provider config (e.g., 'openai', 'anthropic') - within the step config. If none found, returns an empty dict. - """ - for provider in SUPPORTED_PROVIDERS: - if provider in step_config: - return step_config[provider] - return {} + ) from e diff --git a/src/soda_curation/pipeline/step_config.py b/src/soda_curation/pipeline/step_config.py new file mode 100644 index 0000000..b16ea8d --- /dev/null +++ b/src/soda_curation/pipeline/step_config.py @@ -0,0 +1,145 @@ +"""Resolve unified pipeline step configuration (model + shared prompts).""" + +from __future__ import annotations + +from typing import Any, Dict, Iterable, Optional + +SUPPORTED_PROVIDERS = ("openai", "anthropic") + +# Steps that call an LLM in the main curation pipeline. +AI_PIPELINE_STEPS = ( + "extract_sections", + "extract_caption_title", + "extract_panel_sequence", + "extract_data_sources", + "match_caption_panel", + "assign_panel_source", +) + + +def infer_provider_from_model(model: str) -> str: + """ + Infer API provider from a model identifier. + + - ``gpt*``, ``o1*``, ``o3*``, ``o4*`` → OpenAI + - ``claude*`` → Anthropic + """ + if not model or not str(model).strip(): + raise ValueError("Model name is required to infer AI provider.") + + name = str(model).strip().lower() + openai_prefixes = ("gpt", "o1", "o3", "o4") + if any( + name == prefix or name.startswith(f"{prefix}-") for prefix in openai_prefixes + ): + return "openai" + if name.startswith("claude"): + return "anthropic" + raise ValueError( + f"Cannot infer AI provider from model '{model}'. " + "Use an OpenAI (gpt/o-series) or Anthropic (claude) model name." + ) + + +def is_ai_step(step_config: Dict[str, Any]) -> bool: + """Return True when a pipeline step configures an LLM (flat or legacy nested).""" + if not isinstance(step_config, dict): + return False + if "model" in step_config and "prompts" in step_config: + return True + return any(provider in step_config for provider in SUPPORTED_PROVIDERS) + + +def _legacy_provider_block(step_config: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Return the single legacy nested provider block, if present.""" + found = [step_config[p] for p in SUPPORTED_PROVIDERS if p in step_config] + if not found: + return None + if len(found) > 1: + raise ValueError( + "Step defines both 'openai' and 'anthropic' blocks. " + "Use a single flat step with one 'model' and shared 'prompts'." + ) + block = found[0] + if not isinstance(block, dict): + raise ValueError("Legacy provider block must be a mapping.") + return block + + +def resolve_step_config(step_config: Dict[str, Any]) -> Dict[str, Any]: + """ + Normalize a pipeline step to ``{model, prompts, provider}``. + + Supports the flat shape:: + + model: gpt-5.4-mini + prompts: {system: ..., user: ...} + + and the legacy nested shape (backward compatible):: + + openai: + model: ... + prompts: ... + """ + if not isinstance(step_config, dict): + raise ValueError("Step configuration must be a mapping.") + + if "model" in step_config: + model = step_config.get("model") + prompts = step_config.get("prompts") + if not model: + raise ValueError("Step is missing required 'model'.") + if not isinstance(prompts, dict): + raise ValueError("Step is missing required 'prompts' mapping.") + provider = infer_provider_from_model(str(model)) + return {"model": str(model), "prompts": prompts, "provider": provider} + + legacy = _legacy_provider_block(step_config) + if legacy is None: + raise ValueError("Step has no 'model'/'prompts' and no legacy provider block.") + + model = legacy.get("model") + prompts = legacy.get("prompts") + if not model: + raise ValueError("Legacy provider block is missing 'model'.") + if not isinstance(prompts, dict): + raise ValueError("Legacy provider block is missing 'prompts' mapping.") + + provider = next(p for p in SUPPORTED_PROVIDERS if p in step_config) + return {"model": str(model), "prompts": prompts, "provider": provider} + + +def get_step_prompts(step_config: Dict[str, Any]) -> Dict[str, str]: + """Return the ``system`` and ``user`` prompt templates for a step.""" + resolved = resolve_step_config(step_config) + prompts = resolved["prompts"] + system = prompts.get("system", "") + user = prompts.get("user", "") + if not system or not user: + raise ValueError("Step prompts must include non-empty 'system' and 'user'.") + return {"system": system, "user": user} + + +def resolve_pipeline_provider( + pipeline_config: Dict[str, Any], + steps: Iterable[str] = AI_PIPELINE_STEPS, +) -> str: + """ + Infer the single runtime provider for all AI steps from their model names. + + Raises when steps disagree (e.g. one gpt model and one claude model). + """ + providers: set[str] = set() + for step_name in steps: + step_config = pipeline_config.get(step_name) + if not isinstance(step_config, dict): + raise ValueError(f"Missing pipeline configuration for step '{step_name}'.") + resolved = resolve_step_config(step_config) + providers.add(resolved["provider"]) + + if len(providers) != 1: + raise ValueError( + "All AI pipeline steps must use models from the same provider. " + f"Found providers: {sorted(providers)}." + ) + return providers.pop() diff --git a/tests/test_pipeline/benchmark/test_assign_panel.py b/tests/test_pipeline/benchmark/test_assign_panel.py index 1a32419..d132225 100644 --- a/tests/test_pipeline/benchmark/test_assign_panel.py +++ b/tests/test_pipeline/benchmark/test_assign_panel.py @@ -12,6 +12,7 @@ MatchPanelCaptionOpenAI, ) from src.soda_curation.pipeline.prompt_handler import PromptHandler +from src.soda_curation.pipeline.step_config import resolve_step_config from .base_runner import BaseBenchmarkRunner @@ -81,26 +82,15 @@ def _run_panel_assignment_test( pipeline_config = dev_config.get("default", {}).get("pipeline", {}) if not pipeline_config: raise ValueError("No pipeline configuration found in default profile") - # Provider-specific configuration - provider = test_case["provider"] - # Create extractor configuration + match_prompts = resolve_step_config(pipeline_config["match_caption_panel"])[ + "prompts" + ] extractor_config = { - "api_key": None, # We don't need to pass the API key here + "api_key": None, "pipeline": { "match_caption_panel": { - provider: { - "model": test_case["model"], - "temperature": test_case["temperature"], - "top_p": test_case["top_p"], - "prompts": { - "system": pipeline_config["match_caption_panel"][provider][ - "prompts" - ]["system"], - "user": pipeline_config["match_caption_panel"][provider][ - "prompts" - ]["user"], - }, - } + "model": test_case["model"], + "prompts": match_prompts, } }, "extraction_dir": str(temp_extract_dir), diff --git a/tests/test_pipeline/benchmark/test_extract_individual_captions.py b/tests/test_pipeline/benchmark/test_extract_individual_captions.py index 408dad8..88b7580 100644 --- a/tests/test_pipeline/benchmark/test_extract_individual_captions.py +++ b/tests/test_pipeline/benchmark/test_extract_individual_captions.py @@ -14,6 +14,7 @@ ZipStructure, ) from src.soda_curation.pipeline.prompt_handler import PromptHandler +from src.soda_curation.pipeline.step_config import resolve_step_config from .base_runner import BaseBenchmarkRunner @@ -56,67 +57,23 @@ def _run_extract_captions_test( if not pipeline_config: raise ValueError("No pipeline configuration found in default profile") - # Provider-specific configuration provider = test_case["provider"] + caption_prompts = resolve_step_config( + pipeline_config["extract_caption_title"] + )["prompts"] + panel_prompts = resolve_step_config( + pipeline_config["extract_panel_sequence"] + )["prompts"] - # Create extractor configuration - Updated for new structure with split extraction steps extractor_config = { "pipeline": { - # Extract caption title configuration "extract_caption_title": { - provider: { - "model": test_case["model"], - "temperature": test_case["temperature"], - "top_p": test_case["top_p"], - "max_tokens": pipeline_config["extract_caption_title"][ - provider - ].get("max_tokens", 2048), - "frequency_penalty": pipeline_config[ - "extract_caption_title" - ][provider].get("frequency_penalty", 0.0), - "presence_penalty": pipeline_config[ - "extract_caption_title" - ][provider].get("presence_penalty", 0.0), - "json_mode": pipeline_config["extract_caption_title"][ - provider - ].get("json_mode", True), - "prompts": { - "system": pipeline_config["extract_caption_title"][ - provider - ]["prompts"]["system"], - "user": pipeline_config["extract_caption_title"][ - provider - ]["prompts"]["user"], - }, - } + "model": test_case["model"], + "prompts": caption_prompts, }, - # Extract panel sequence configuration "extract_panel_sequence": { - provider: { - "model": test_case["model"], - "temperature": test_case["temperature"], - "top_p": test_case["top_p"], - "max_tokens": pipeline_config["extract_panel_sequence"][ - provider - ].get("max_tokens", 2048), - "frequency_penalty": pipeline_config[ - "extract_panel_sequence" - ][provider].get("frequency_penalty", 0.0), - "presence_penalty": pipeline_config[ - "extract_panel_sequence" - ][provider].get("presence_penalty", 0.0), - "json_mode": pipeline_config["extract_panel_sequence"][ - provider - ].get("json_mode", True), - "prompts": { - "system": pipeline_config["extract_panel_sequence"][ - provider - ]["prompts"]["system"], - "user": pipeline_config["extract_panel_sequence"][ - provider - ]["prompts"]["user"], - }, - } + "model": test_case["model"], + "prompts": panel_prompts, }, }, } diff --git a/tests/test_pipeline/benchmark/test_extract_sections.py b/tests/test_pipeline/benchmark/test_extract_sections.py index fad1fdc..aaad826 100644 --- a/tests/test_pipeline/benchmark/test_extract_sections.py +++ b/tests/test_pipeline/benchmark/test_extract_sections.py @@ -14,6 +14,7 @@ XMLStructureExtractor, ) from src.soda_curation.pipeline.prompt_handler import PromptHandler +from src.soda_curation.pipeline.step_config import resolve_step_config from .base_runner import BaseBenchmarkRunner @@ -99,27 +100,17 @@ def _run_extract_sections_test( if not pipeline_config: raise ValueError("No pipeline configuration found in default profile") - # Provider-specific configuration provider = test_case["provider"] + section_prompts = resolve_step_config(pipeline_config["extract_sections"])[ + "prompts" + ] - # Create extractor configuration extractor_config = { - "api_key": None, # We don't need to pass the API key here + "api_key": None, "pipeline": { "extract_sections": { - provider: { - "model": test_case["model"], - "temperature": test_case["temperature"], - "top_p": test_case["top_p"], - "prompts": { - "system": pipeline_config["extract_sections"][provider][ - "prompts" - ]["system"], - "user": pipeline_config["extract_sections"][provider][ - "prompts" - ]["user"], - }, - } + "model": test_case["model"], + "prompts": section_prompts, } }, } diff --git a/tests/test_pipeline/benchmark/utils.py b/tests/test_pipeline/benchmark/utils.py index d351d12..f1c4512 100644 --- a/tests/test_pipeline/benchmark/utils.py +++ b/tests/test_pipeline/benchmark/utils.py @@ -14,6 +14,7 @@ XMLStructureExtractor, ) from src.soda_curation.pipeline.prompt_handler import PromptHandler +from src.soda_curation.pipeline.step_config import resolve_step_config from .base_runner import BaseBenchmarkRunner @@ -73,27 +74,17 @@ def _run_data_availability_test( if not pipeline_config: raise ValueError("No pipeline configuration found in default profile") - # Provider-specific configuration provider = test_case["provider"] + data_prompts = resolve_step_config(pipeline_config["extract_data_sources"])[ + "prompts" + ] - # Create extractor configuration extractor_config = { - "api_key": None, # We don't need to pass the API key here + "api_key": None, "pipeline": { "extract_data_sources": { - provider: { - "model": test_case["model"], - "temperature": test_case["temperature"], - "top_p": test_case["top_p"], - "prompts": { - "system": pipeline_config["extract_data_sources"][ - provider - ]["prompts"]["system"], - "user": pipeline_config["extract_data_sources"][ - provider - ]["prompts"]["user"], - }, - } + "model": test_case["model"], + "prompts": data_prompts, } }, } diff --git a/tests/test_suite/test_match_caption_panel.py b/tests/test_suite/test_match_caption_panel.py index be5a47e..94fa39b 100644 --- a/tests/test_suite/test_match_caption_panel.py +++ b/tests/test_suite/test_match_caption_panel.py @@ -53,12 +53,11 @@ def mock_config(): "extraction_dir": "/tmp", "pipeline": { "match_caption_panel": { - "openai": { - "model": "gpt-4o", - "temperature": 0.3, - "top_p": 1.0, - "max_tokens": 512, - } + "model": "gpt-4o", + "prompts": { + "system": "Match panel", + "user": "Caption: $figure_caption", + }, } }, } diff --git a/tests/test_suite/test_step_config.py b/tests/test_suite/test_step_config.py new file mode 100644 index 0000000..a58821a --- /dev/null +++ b/tests/test_suite/test_step_config.py @@ -0,0 +1,111 @@ +"""Tests for unified pipeline step configuration resolution.""" + +import unittest + +from src.soda_curation.pipeline.step_config import ( + infer_provider_from_model, + resolve_pipeline_provider, + resolve_step_config, +) + + +class TestInferProviderFromModel(unittest.TestCase): + def test_openai_gpt(self): + self.assertEqual(infer_provider_from_model("gpt-5.4-mini"), "openai") + self.assertEqual(infer_provider_from_model("gpt-4o"), "openai") + + def test_openai_o_series(self): + self.assertEqual(infer_provider_from_model("o1-preview"), "openai") + self.assertEqual(infer_provider_from_model("o3-mini"), "openai") + + def test_anthropic_claude(self): + self.assertEqual(infer_provider_from_model("claude-sonnet-4-6"), "anthropic") + + def test_unknown_model_raises(self): + with self.assertRaises(ValueError): + infer_provider_from_model("llama-3") + + +class TestResolveStepConfig(unittest.TestCase): + def test_flat_shape(self): + step = { + "model": "gpt-5.4-mini", + "prompts": {"system": "sys", "user": "usr"}, + } + resolved = resolve_step_config(step) + self.assertEqual(resolved["model"], "gpt-5.4-mini") + self.assertEqual(resolved["provider"], "openai") + self.assertEqual(resolved["prompts"]["system"], "sys") + + def test_legacy_nested_openai(self): + step = { + "openai": { + "model": "gpt-4o", + "temperature": 0.1, + "prompts": {"system": "sys", "user": "usr"}, + } + } + resolved = resolve_step_config(step) + self.assertEqual(resolved["model"], "gpt-4o") + self.assertEqual(resolved["provider"], "openai") + + +class TestResolvePipelineProvider(unittest.TestCase): + def test_single_provider_across_steps(self): + pipeline = { + "extract_sections": { + "model": "gpt-5.4-mini", + "prompts": {"system": "s", "user": "u"}, + }, + "extract_caption_title": { + "model": "gpt-4o", + "prompts": {"system": "s", "user": "u"}, + }, + "extract_panel_sequence": { + "model": "gpt-4o", + "prompts": {"system": "s", "user": "u"}, + }, + "extract_data_sources": { + "model": "gpt-4o-mini", + "prompts": {"system": "s", "user": "u"}, + }, + "match_caption_panel": { + "model": "gpt-4o", + "prompts": {"system": "s", "user": "u"}, + }, + "assign_panel_source": { + "model": "gpt-4o", + "prompts": {"system": "s", "user": "u"}, + }, + } + self.assertEqual(resolve_pipeline_provider(pipeline), "openai") + + def test_mixed_providers_raises(self): + pipeline = { + "extract_sections": { + "model": "gpt-4o", + "prompts": {"system": "s", "user": "u"}, + }, + "extract_caption_title": { + "model": "claude-sonnet-4-6", + "prompts": {"system": "s", "user": "u"}, + }, + "extract_panel_sequence": { + "model": "claude-sonnet-4-6", + "prompts": {"system": "s", "user": "u"}, + }, + "extract_data_sources": { + "model": "claude-sonnet-4-6", + "prompts": {"system": "s", "user": "u"}, + }, + "match_caption_panel": { + "model": "claude-sonnet-4-6", + "prompts": {"system": "s", "user": "u"}, + }, + "assign_panel_source": { + "model": "claude-sonnet-4-6", + "prompts": {"system": "s", "user": "u"}, + }, + } + with self.assertRaises(ValueError): + resolve_pipeline_provider(pipeline)