diff --git a/.config.schema.yaml b/.config.schema.yaml index 5da62f8..9fca799 100644 --- a/.config.schema.yaml +++ b/.config.schema.yaml @@ -55,7 +55,7 @@ properties: type: array items: type: string - enum: ["React-to-Me", "Cross-Database Prototype"] + enum: ["React-to-Me", "Cross-Database Prototype", "Plant Reactome"] usage_limits: type: object properties: diff --git a/.dockerignore b/.dockerignore index 5b3c643..c454cfe 100644 --- a/.dockerignore +++ b/.dockerignore @@ -2,5 +2,12 @@ data/ embeddings/ embeddings_bak/ csv_files/ +records/ .venv/ .env +.env.* +.git/ +__pycache__/ +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ diff --git a/.github/actions/install_python_poetry/action.yml b/.github/actions/install_python_poetry/action.yml index 3e51b57..64059ca 100644 --- a/.github/actions/install_python_poetry/action.yml +++ b/.github/actions/install_python_poetry/action.yml @@ -1,16 +1,34 @@ name: 'Setup Python and Poetry' -description: 'Setup Python environment and install Poetry' +description: 'Setup Python environment and install Poetry, with the venv cached' runs: using: 'composite' steps: - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: '3.12' - - name: Install dependencies + # Deliberately `pip install` into the interpreter set up above, not pipx. + # pipx gives Poetry its own interpreter -- on the macOS runner that is + # Homebrew's Python 3.14 -- and Poetry then builds the project venv with it. + # pydantic-core publishes no 3.14 wheel yet, so the build falls back to + # compiling and fails: PyO3 0.22 does not support 3.14. + - name: Install Poetry shell: bash run: | - pip install --upgrade pip - pip install poetry==1.8.4 - poetry install --no-root + python -m pip install --upgrade pip + python -m pip install poetry==1.8.4 + + # poetry.toml sets virtualenvs.in-project, so the environment lands in + # ./.venv and can be cached wholesale. Without this every job reinstalls the + # dependency set from scratch -- which includes a ~200 MB torch wheel -- and + # there are several jobs per run. + - name: Cache the virtualenv + uses: actions/cache@v4 + with: + path: .venv + key: venv-${{ runner.os }}-${{ runner.arch }}-py3.12-${{ hashFiles('poetry.lock') }} + + - name: Install dependencies + shell: bash + run: poetry install --no-root diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c806cff..2c6bf08 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -10,13 +10,19 @@ on: branches: - main +# A second push to a PR makes the first run's result irrelevant; without this +# they both run to completion and queue behind each other. +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +# Least privilege by default. id-token: write is granted only to docker-push, +# which needs it to assume the AWS role. permissions: - id-token: write contents: read jobs: lint: - if: ${{ github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch' }} runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 @@ -24,11 +30,27 @@ jobs: - name: Set up Python and Poetry uses: ./.github/actions/install_python_poetry - - name: Run linters - run: | - poetry run ruff check . - poetry run mypy . - poetry run isort --check . + # ruff replaces black + isort; its `I` rules sort imports and + # `ruff format` is black-compatible. Config lives in pyproject.toml. + - name: Lint + run: poetry run ruff check . + + - name: Check formatting + run: poetry run ruff format --check . + + - name: Type check + run: poetry run mypy + + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python and Poetry + uses: ./.github/actions/install_python_poetry + + - name: Run tests + run: poetry run pytest poetry-check: if: ${{ github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch' }} @@ -39,19 +61,31 @@ jobs: steps: - uses: actions/checkout@v4 + with: + fetch-depth: 0 # need the base branch to diff against + # Plain git rather than a third-party action: this workflow can reach + # AWS, so every extra action is supply-chain surface for a one-line check. - name: Check poetry.lock for changes id: check-poetry-lock - uses: tj-actions/changed-files@v45 - with: - files: poetry.lock + shell: bash + run: | + base="${{ github.base_ref }}" + if [ -z "$base" ]; then + # manual run: no base to compare against, so always verify + echo "changed=true" >> "$GITHUB_OUTPUT" + elif git diff --name-only "origin/$base...HEAD" -- poetry.lock | grep -q .; then + echo "changed=true" >> "$GITHUB_OUTPUT" + else + echo "changed=false" >> "$GITHUB_OUTPUT" + fi - name: Set up Python and Poetry - if: steps.check-poetry-lock.outputs.any_changed == 'true' + if: steps.check-poetry-lock.outputs.changed == 'true' uses: ./.github/actions/install_python_poetry - name: Verify Python imports - if: steps.check-poetry-lock.outputs.any_changed == 'true' + if: steps.check-poetry-lock.outputs.changed == 'true' env: PYTHONPATH: ./bin:./src run: | @@ -68,7 +102,7 @@ jobs: uses: docker/setup-buildx-action@v3 - name: Build and push Docker image - uses: docker/build-push-action@v5 + uses: docker/build-push-action@v6 with: context: . file: ./Dockerfile @@ -84,6 +118,9 @@ jobs: if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} needs: docker-build runs-on: ubuntu-latest + permissions: + id-token: write # assume the AWS role via OIDC + contents: read steps: - uses: actions/download-artifact@v4 diff --git a/.gitignore b/.gitignore index f881919..dad05ec 100644 --- a/.gitignore +++ b/.gitignore @@ -121,6 +121,10 @@ celerybeat.pid # Environments .env +# any per-environment variant: .env.beta, .env.local, ... +.env.* +!.env.example +!env_template .venv env/ venv/ @@ -157,7 +161,7 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +.idea/ .files/ .ruff_cache/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..e9ac8e1 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,20 @@ +# Install once: pipx install pre-commit && pre-commit install +# Run manually: pre-commit run --all-files +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.7.4 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: check-yaml + - id: check-toml + - id: check-merge-conflict + - id: end-of-file-fixer + - id: trailing-whitespace + - id: check-added-large-files + args: [--maxkb=1024] diff --git a/README.md b/README.md index f38a4d3..2c2c5c8 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,8 @@ The ChatBot's knowledge of a given data source is generated using the latest dat In the case of Reactome, embeddings bundles are generated once per release from [reactome/graphdb](https://hub.docker.com/r/reactome/graphdb) releases from DockerHub and uploaded to AWS S3 for easy retrieval. +User guide embeddings are generated separately from Reactome website documentation and use a date-based version identifier (for example, `userguide/2025-06`). See [Embeddings Manager documentation](docs/embeddings_manager.md) for details. + ### Embeddings Manager Script All aspects of generating, managing, uploading, and retrieving embeddings bundles are handled by the `./bin/embeddings_manager` script. @@ -109,24 +111,35 @@ All aspects of generating, managing, uploading, and retrieving embeddings bundle ### Code Quality -To do main consistency checks +All tool configuration lives in `pyproject.toml`. Ruff handles linting, import +sorting, and formatting (it replaces `black` and `isort`). + ```bash -poetry run ruff check . +poetry run ruff check . # lint (add --fix to autofix) +poetry run ruff format . # format +poetry run mypy # type check +poetry run pytest # tests ``` -To make style consistent +CI runs all four on every pull request and on pushes to `main`. + +Optionally, run the same checks on every commit: ```bash -poetry run black . +pipx install pre-commit && pre-commit install ``` -To make sure imports are organized - +### Tests ```bash -poetry run isort . +poetry run pytest +poetry run pytest -m "not requires_retrieval_stack" # no ML deps needed ``` +Tests that need an installed embeddings bundle are marked `requires_embeddings` +and skip themselves when none is present. See `tests/README.md` for what is +covered and why coverage is currently narrow. + ### Contributing Contributions to the Reactome ChatBot project are welcome! If you encounter any issues or have suggestions for improvements, feel free to open an issue or submit a pull request. diff --git a/analysis/expert_survey/README.md b/analysis/expert_survey/README.md new file mode 100644 index 0000000..db9a866 --- /dev/null +++ b/analysis/expert_survey/README.md @@ -0,0 +1,106 @@ +# Expert Survey Analysis Pipeline + +This directory contains a modular R workflow that reproduces the expert survey evaluation used to compare ChatGPT (baseline) vs React-to-Me (intervention) across three expert-rated metrics: + +- Factual accuracy +- Level of Granularity +- Relational Depth + +Each participant rated 15 questions (9 Query, 6 Reasoning) on a 4-point ordinal scale, yielding 450 paired comparisons (900 total observations). Place the evaluation data at `analysis/expert_survey/survey_results.csv` with the following columns: + +| Column | Description | +|----------------------|------------------------------------------------------------------| +| `participant_id` | Participant identifier (string) | +| `question_id` | Question identifier (numeric) | +| `question_type` | `Query` or `Reasoning` | +| `metric` | `Factual accuracy`, `Level of Granularity`, `Relational Depth` | +| `React-to-Me score` | React-to-Me rating (integer 1–4) | +| `GPT-baseline score` | ChatGPT rating (integer 1–4) | + +## Running the Pipeline + +```bash +poetry run Rscript analysis/expert_survey/run_all_analysis.R +``` + +This orchestrates all numbered scripts (`00_setup.R` → `08_tables_and_text.R`) and writes outputs under `analysis/expert_survey/results/`. Install dependencies first: + +```r +install.packages(c( + "tidyverse", + "ordinal", + "lme4", + "lmerTest", + "effsize", + "irr", + "boot", + "broom", + "broom.mixed", + "pwr" +)) +``` + +The workflow has been developed and tested with **R 4.3.x** (latest 4.x release); earlier versions may fail to install one or more packages. If you prefer to run outside Poetry, call the scripts with `Rscript` after installing these packages. + +## Processing & Stratifications + +- **00_setup.R**: Loads the reference dataset (`survey_results.csv`), trims whitespace, validates counts, converts to long format, and creates a paired wide dataset with score differences. +- **01_descriptives.R**: Generates descriptive tables: + - Score distributions by system × metric × score. + - Summary statistics by system, by metric, and stratified by question type. + - Paired-difference summaries overall and by question type. +- **03_ordinal_models.R**: Fits cumulative link mixed models (CLMM) with logit link, flexible thresholds, and random intercept/slope per participant and question. + - Stratifications: overall, by metric, by question type, and metric × question type. + - Interaction models (`system * question_type`) are attempted for each metric; likelihood-ratio tests are logged but skipped if they cannot be evaluated. +- **04_nonparametrics.R**: + - Wilcoxon signed-rank tests (one-sided, React-to-Me > ChatGPT) with Hodges–Lehmann estimates and confidence intervals. + - Binomial sign tests at question-level and participant-level. + - Effect sizes: Cliff’s Delta, Cohen’s d, strict/common-language effect size (CLES) with bootstrap CIs. + - Stratifications mirror the CLMM step (overall, metrics, question types, metric × question type). + +## Inferential Adjustments + +- **05_multiple_testing.R**: Holm corrections applied separately to: + - Primary family: three per-metric comparisons. + - Secondary family: two question-type comparisons. + - Exploratory family: remaining strata (metric × question type, overall) reported without adjustment. +- **06_power_analysis.R**: Post-hoc power estimates using `pwr.t.test` for paired differences (Cohen’s d, α = 0.05 one-sided). + +## Results Assembly & Publication Outputs + +- **07_results_compilation.R**: Merges CLMM odds ratios, nonparametric results, Holm adjustments, power, and sign-test summaries into `COMPREHENSIVE_RESULTS_ALL_STRATIFICATIONS.csv`. Each row is labeled as: + - `1. Primary (By Metric)` + - `2. Secondary (By Question Type)` + - `3. Exploratory (Overall)` + - `4. Exploratory (Metric × Question Type)` +- **08_tables_and_text.R**: Produces publication-ready assets: + - `TABLE_MAIN_RESULTS_FORMATTED.csv` (primary metrics with one-sided p-values and Holm-adjusted p-values). + - `TABLE_STRATIFIED_RESULTS_FORMATTED.csv` (secondary and exploratory strata). + - `KEY_STATISTICS_FOR_TEXT.txt` (concise summary lines per stratum). + +## Statistical Methods Summary + +- **Ordinal Mixed Models (CLMM)**: `ordinal::clmm`, logit link, flexible thresholds, random intercepts and system-specific slopes at participant and question levels. +- **Nonparametric Tests**: One-sided Wilcoxon signed-rank, sign tests (question-level and participant-level), bootstrap effect-size CIs. +- **Effect Sizes**: Odds ratios with Wald CIs, Cliff’s Delta, Cohen’s d, strict/inclusive common language effect sizes. +- **Multiple Testing**: Holm correction by analysis family. +- **Power Analysis**: Paired t-test approximation via Cohen’s d. + +All outputs are written under `analysis/expert_survey/results/` with separate folders for tables, stratified summaries, raw processed data, and sensitivity analyses. The pipeline is self-contained; rerunning it will regenerate every table from scratch given the reference dataset. + +## Key Outputs + +| File/Folder | Description | +|-----------------------------------------------------------------------|-------------------------------------------------------------------| +| `results/raw_data/survey_long.csv` / `paired_differences.csv` | Processed datasets (long-form and paired differences) | +| `results/tables/summary_statistics_*.csv` | Descriptive summaries | +| `results/tables/ordinal_models_all_stratifications.csv` | CLMM odds ratios and p-values | +| `results/tables/nonparametric_tests_all_stratifications.csv` | Wilcoxon results, effect sizes, sign-test p-values | +| `results/tables/multiple_testing_correction_comprehensive.csv` | Holm-adjusted p-values by family | +| `results/tables/power_analysis_all_stratifications.csv` | Post-hoc power estimates | +| `results/tables/sign_test_question_level.csv` / `_participant_level.csv` | Sign-test summaries | +| `results/tables/COMPREHENSIVE_RESULTS_ALL_STRATIFICATIONS.csv` | Unified table combining all statistics | +| `results/tables/TABLE_MAIN_RESULTS_FORMATTED.csv` | Publication-ready table for primary metrics | +| `results/tables/TABLE_STRATIFIED_RESULTS_FORMATTED.csv` | Secondary and exploratory formatted table | +| `results/tables/KEY_STATISTICS_FOR_TEXT.txt` | Key statistics lines for manuscript text | + diff --git a/analysis/expert_survey/config/analysis_parameters.R b/analysis/expert_survey/config/analysis_parameters.R new file mode 100644 index 0000000..c2ee319 --- /dev/null +++ b/analysis/expert_survey/config/analysis_parameters.R @@ -0,0 +1,18 @@ +# Analysis constants for expert survey workflow + +ALPHA_ONESIDED <- 0.05 + +PRIMARY_METRICS <- c( + "Factual accuracy", + "Level of Granularity", + "Relational Depth" +) + +QUESTION_TYPES <- c("Query", "Reasoning") + +NI_MARGINS <- list( + OR_margin = 1.10, + win_prob_margin = 0.55, + median_diff_margin = 0.25 +) + diff --git a/analysis/expert_survey/config/file_paths.R b/analysis/expert_survey/config/file_paths.R new file mode 100644 index 0000000..62bb292 --- /dev/null +++ b/analysis/expert_survey/config/file_paths.R @@ -0,0 +1,15 @@ +# File path configuration for expert survey analysis + +ANALYSIS_ROOT <- "analysis/expert_survey" + +DATA_PATH <- file.path(ANALYSIS_ROOT, "survey_results.csv") + +RESULTS_DIR <- file.path(ANALYSIS_ROOT, "results") +TABLES_DIR <- file.path(RESULTS_DIR, "tables") +RAW_DATA_DIR <- file.path(RESULTS_DIR, "raw_data") +SENSITIVITY_DIR <- file.path(RESULTS_DIR, "sensitivity") +STRATIFIED_DIR <- file.path(RESULTS_DIR, "stratified") + +PROCESSED_LONG_PATH <- file.path(RAW_DATA_DIR, "survey_long.csv") +PROCESSED_PAIRED_PATH <- file.path(RAW_DATA_DIR, "paired_differences.csv") + diff --git a/analysis/expert_survey/run_all_analysis.R b/analysis/expert_survey/run_all_analysis.R new file mode 100644 index 0000000..d83e39f --- /dev/null +++ b/analysis/expert_survey/run_all_analysis.R @@ -0,0 +1,44 @@ +#!/usr/bin/env Rscript + +# Orchestrate expert survey analysis pipeline (scripts 00–08). + +scripts <- sprintf( + "analysis/expert_survey/scripts/%02d_%s.R", + 0:8, + c( + "setup", + "descriptives", + "ordinal_models", # note: placeholder, actual 02 skipped (reliability) + "ordinal_models", + "nonparametrics", + "multiple_testing", + "power_analysis", + "results_compilation", + "tables_and_text" + ) +) + +# Adjust script names to match existing files +scripts <- c( + "analysis/expert_survey/scripts/00_setup.R", + "analysis/expert_survey/scripts/01_descriptives.R", + "analysis/expert_survey/scripts/03_ordinal_models.R", + "analysis/expert_survey/scripts/04_nonparametrics.R", + "analysis/expert_survey/scripts/05_multiple_testing.R", + "analysis/expert_survey/scripts/06_power_analysis.R", + "analysis/expert_survey/scripts/07_results_compilation.R", + "analysis/expert_survey/scripts/08_tables_and_text.R" +) + +for (script in scripts) { + message("\n=== Running ", script, " ===") + tryCatch( + source(script, echo = TRUE, max.deparse.length = Inf), + error = function(e) { + stop(sprintf("Error running %s: %s", script, e$message), call. = FALSE) + } + ) +} + +message("\nAll analysis steps completed successfully.") + diff --git a/analysis/expert_survey/scripts/00_setup.R b/analysis/expert_survey/scripts/00_setup.R new file mode 100644 index 0000000..6ef2514 --- /dev/null +++ b/analysis/expert_survey/scripts/00_setup.R @@ -0,0 +1,47 @@ +#!/usr/bin/env Rscript + +# Setup script: load data, ensure output directories, and write processed datasets. + +suppressPackageStartupMessages({ + library(dplyr) + library(readr) + library(tidyr) +}) + +source("analysis/expert_survey/config/file_paths.R") +source("analysis/expert_survey/config/analysis_parameters.R") +source("analysis/expert_survey/utils/data_preparation.R") + +# Create output directories ----------------------------------------------------- +dirs_to_create <- c( + RESULTS_DIR, + TABLES_DIR, + RAW_DATA_DIR, + SENSITIVITY_DIR, + STRATIFIED_DIR +) + +invisible(lapply(dirs_to_create, dir.create, recursive = TRUE, showWarnings = FALSE)) + +cat("Output directories verified.\n") + +# Load and prepare data --------------------------------------------------------- +cat("Loading survey data from:", DATA_PATH, "\n") +survey_wide <- load_and_validate_data(DATA_PATH) + +cat("Converting to long format...\n") +survey_long <- to_long_format(survey_wide) + +cat("Creating paired dataset...\n") +survey_paired <- create_paired_dataset(survey_long) + +# Write processed datasets ------------------------------------------------------ +write_csv(survey_long, PROCESSED_LONG_PATH) +write_csv(survey_paired, PROCESSED_PAIRED_PATH) + +cat("Processed datasets saved:\n") +cat(" - Long format:", PROCESSED_LONG_PATH, "\n") +cat(" - Paired differences:", PROCESSED_PAIRED_PATH, "\n") + +cat("Setup complete.\n") + diff --git a/analysis/expert_survey/scripts/01_descriptives.R b/analysis/expert_survey/scripts/01_descriptives.R new file mode 100644 index 0000000..202f32c --- /dev/null +++ b/analysis/expert_survey/scripts/01_descriptives.R @@ -0,0 +1,131 @@ +#!/usr/bin/env Rscript + +# Descriptive statistics and summary tables. + +suppressPackageStartupMessages({ + library(dplyr) + library(readr) + library(tidyr) + library(stringr) +}) + +source("analysis/expert_survey/config/file_paths.R") +source("analysis/expert_survey/config/analysis_parameters.R") + +cat("Loading processed datasets...\n") +survey_long <- read_csv(PROCESSED_LONG_PATH, show_col_types = FALSE) %>% + mutate( + system = factor(system, levels = c("ChatGPT", "React-to-Me")), + metric = factor(metric, levels = PRIMARY_METRICS), + question_type = factor(question_type, levels = QUESTION_TYPES), + participant_id = factor(participant_id), + question_id = factor(question_id) + ) + +survey_paired <- read_csv(PROCESSED_PAIRED_PATH, show_col_types = FALSE) %>% + mutate( + metric = factor(metric, levels = PRIMARY_METRICS), + question_type = factor(question_type, levels = QUESTION_TYPES), + participant_id = factor(participant_id), + question_id = factor(question_id) + ) + +# Score distributions ----------------------------------------------------------- +cat("Generating score distribution tables...\n") +score_dist <- survey_long %>% + group_by(system, metric, score) %>% + summarise(n = n(), .groups = "drop") %>% + group_by(system, metric) %>% + mutate( + total = sum(n), + percent = 100 * n / total + ) %>% + arrange(system, metric, score) + +write_csv(score_dist, file.path(TABLES_DIR, "score_distributions.csv")) + +# Summary statistics by system and metric --------------------------------------- +summary_stats <- survey_long %>% + mutate(score_num = as.numeric(score)) %>% + group_by(system, metric) %>% + summarise( + n = n(), + mean = mean(score_num), + sd = sd(score_num), + median = median(score_num), + q25 = quantile(score_num, 0.25), + q75 = quantile(score_num, 0.75), + pct_ge3 = 100 * mean(score_num >= 3), + .groups = "drop" + ) %>% + arrange(metric, system) + +write_csv(summary_stats, file.path(TABLES_DIR, "summary_statistics_by_system.csv")) + +# Stratified summary by question type ------------------------------------------- +summary_stats_stratified <- survey_long %>% + mutate(score_num = as.numeric(score)) %>% + group_by(system, metric, question_type) %>% + summarise( + n = n(), + mean = mean(score_num), + sd = sd(score_num), + median = median(score_num), + pct_ge3 = 100 * mean(score_num >= 3), + .groups = "drop" + ) %>% + arrange(metric, question_type, system) + +write_csv( + summary_stats_stratified, + file.path(TABLES_DIR, "summary_statistics_stratified.csv") +) + +# Paired differences summary ---------------------------------------------------- +diff_summary <- survey_paired %>% + group_by(metric) %>% + summarise( + n_pairs = n(), + n_improve = sum(diff_numeric > 0), + n_worsen = sum(diff_numeric < 0), + n_ties = sum(diff_numeric == 0), + pct_improve = 100 * n_improve / n_pairs, + pct_worsen = 100 * n_worsen / n_pairs, + pct_ties = 100 * n_ties / n_pairs, + mean_diff = mean(diff_numeric), + sd_diff = sd(diff_numeric), + median_diff = median(diff_numeric), + .groups = "drop" + ) %>% + arrange(metric) + +write_csv(diff_summary, file.path(TABLES_DIR, "paired_differences_summary.csv")) + +# Paired differences by question type ------------------------------------------ +diff_summary_stratified <- survey_paired %>% + group_by(metric, question_type) %>% + summarise( + n_pairs = n(), + n_improve = sum(diff_numeric > 0), + n_worsen = sum(diff_numeric < 0), + n_ties = sum(diff_numeric == 0), + pct_improve = 100 * n_improve / n_pairs, + mean_diff = mean(diff_numeric), + median_diff = median(diff_numeric), + .groups = "drop" + ) %>% + arrange(metric, question_type) + +write_csv( + diff_summary_stratified, + file.path(STRATIFIED_DIR, "paired_differences_by_question_type.csv") +) + +cat("Descriptive tables created:\n") +cat(" - score_distributions.csv\n") +cat(" - summary_statistics_by_system.csv\n") +cat(" - summary_statistics_stratified.csv\n") +cat(" - paired_differences_summary.csv\n") +cat(" - paired_differences_by_question_type.csv\n") +cat("Descriptive analysis complete.\n") + diff --git a/analysis/expert_survey/scripts/03_ordinal_models.R b/analysis/expert_survey/scripts/03_ordinal_models.R new file mode 100644 index 0000000..5e5c057 --- /dev/null +++ b/analysis/expert_survey/scripts/03_ordinal_models.R @@ -0,0 +1,207 @@ +#!/usr/bin/env Rscript + +# Ordinal mixed-effects models across stratifications. + +suppressPackageStartupMessages({ + library(dplyr) + library(readr) + library(tidyr) + library(purrr) + library(ordinal) +}) + +source("analysis/expert_survey/config/file_paths.R") +source("analysis/expert_survey/config/analysis_parameters.R") + +# ------------------------------------------------------------------------- +# Helper to fit CLMM and extract odds ratios +# ------------------------------------------------------------------------- +fit_and_extract_ordinal <- function(data, stratum_name, test_interaction = FALSE) { + message("\n--- Fitting model for: ", stratum_name, " ---") + + data <- droplevels(data) + data$system <- relevel(data$system, ref = "ChatGPT") + + formula_str <- if (test_interaction) { + "score ~ system * question_type + (1 + system|participant_id) + (1 + system|question_id)" + } else { + "score ~ system + (1 + system|participant_id) + (1 + system|question_id)" + } + + message("Formula: ", formula_str) + + model <- tryCatch( + clmm( + as.formula(formula_str), + data = data, + link = "logit", + threshold = "flexible" + ), + error = function(e) { + message("Model failed for ", stratum_name, ": ", e$message) + return(NULL) + } + ) + + if (is.null(model)) { + return(NULL) + } + + message("Convergence: ", model$convergence$code == 0) + + coef_df <- tryCatch( + { + tmp <- as.data.frame(summary(model)$coefficients) + tmp$parameter <- rownames(tmp) + tmp + }, + error = function(e) { + message("Coefficient extraction failed for ", stratum_name, ": ", e$message) + return(NULL) + } + ) + + if (is.null(coef_df)) { + return(NULL) + } + + results <- coef_df %>% + as_tibble() %>% + filter(grepl("system", parameter)) %>% + mutate( + OR = exp(Estimate), + OR_lower_95 = exp(Estimate - 1.96 * `Std. Error`), + OR_upper_95 = exp(Estimate + 1.96 * `Std. Error`), + stratum = stratum_name, + interaction = test_interaction + ) %>% + select( + stratum, + interaction, + parameter, + Estimate, + `Std. Error`, + `z value`, + p_value = `Pr(>|z|)`, + OR, + OR_lower_95, + OR_upper_95 + ) + + list(model = model, results = results) +} + +# ------------------------------------------------------------------------- +# Load processed data +# ------------------------------------------------------------------------- + +cat("Loading long-format dataset...\n") +survey_long <- read_csv(PROCESSED_LONG_PATH, show_col_types = FALSE) %>% + mutate( + system = factor(system, levels = c("ChatGPT", "React-to-Me")), + metric = factor(metric, levels = PRIMARY_METRICS), + question_type = factor(question_type, levels = QUESTION_TYPES), + participant_id = factor(participant_id), + question_id = factor(question_id), + score = factor(score, levels = 1:4, ordered = TRUE) + ) + +# ------------------------------------------------------------------------- +# Overall model +# ------------------------------------------------------------------------- + +ordinal_results <- list() + +overall_model <- fit_and_extract_ordinal(survey_long, "Overall", test_interaction = FALSE) +ordinal_results[["Overall"]] <- overall_model$results + +# ------------------------------------------------------------------------- +# By metric (with interaction tests) +# ------------------------------------------------------------------------- + +metric_models <- lapply(PRIMARY_METRICS, function(metric_name) { + data_met <- filter(survey_long, metric == metric_name) + + main_model <- fit_and_extract_ordinal(data_met, metric_name, test_interaction = FALSE) + if (is.null(main_model)) { + message("Main model unavailable for ", metric_name, "; skipping metric.") + return(list(metric = metric_name, main = NULL, interaction = NULL)) + } + + int_model <- fit_and_extract_ordinal( + data_met, + paste0(metric_name, " (interaction)"), + test_interaction = TRUE + ) + + interaction_sig <- FALSE + if (!is.null(int_model)) { + lr_result <- tryCatch( + anova(main_model$model, int_model$model), + error = function(e) { + message("LR test failed for ", metric_name, ": ", e$message) + NULL + } + ) + if (!is.null(lr_result) && nrow(lr_result) >= 2) { + interaction_sig <- lr_result$`Pr(>Chisq)`[2] < 0.05 + } + } + + list( + metric = metric_name, + main = main_model$results, + interaction = if (interaction_sig && !is.null(int_model)) int_model$results else NULL + ) +}) + +for (mm in metric_models) { + if (is.null(mm$main)) next + ordinal_results[[mm$metric]] <- mm$main + if (!is.null(mm$interaction)) { + ordinal_results[[paste0(mm$metric, " (interaction)")]] <- mm$interaction + } +} + +# ------------------------------------------------------------------------- +# By question type +# ------------------------------------------------------------------------- + +for (qtype in QUESTION_TYPES) { + data_qtype <- filter(survey_long, question_type == qtype) + q_model <- fit_and_extract_ordinal(data_qtype, qtype, test_interaction = FALSE) + if (!is.null(q_model)) { + ordinal_results[[qtype]] <- q_model$results + } +} + +# ------------------------------------------------------------------------- +# Metric × question type combinations +# ------------------------------------------------------------------------- + +for (metric_name in PRIMARY_METRICS) { + for (qtype in QUESTION_TYPES) { + data_strat <- filter(survey_long, metric == metric_name, question_type == qtype) + + if (nrow(data_strat) < 10) { + message("Skipping ", metric_name, " - ", qtype, ": insufficient data (n = ", nrow(data_strat), ")") + next + } + + stratum_name <- paste0(metric_name, " - ", qtype) + strat_model <- fit_and_extract_ordinal(data_strat, stratum_name, test_interaction = FALSE) + if (!is.null(strat_model)) { + ordinal_results[[stratum_name]] <- strat_model$results + } + } +} + +# ------------------------------------------------------------------------- +# Combine and write results +# ------------------------------------------------------------------------- + +ordinal_results_table <- bind_rows(ordinal_results) +write_csv(ordinal_results_table, file.path(TABLES_DIR, "ordinal_models_all_stratifications.csv")) + +cat("Ordinal model results written to ordinal_models_all_stratifications.csv\n") + diff --git a/analysis/expert_survey/scripts/04_nonparametrics.R b/analysis/expert_survey/scripts/04_nonparametrics.R new file mode 100644 index 0000000..2d6e69b --- /dev/null +++ b/analysis/expert_survey/scripts/04_nonparametrics.R @@ -0,0 +1,361 @@ +#!/usr/bin/env Rscript + +# Nonparametric analyses: Wilcoxon, sign tests, and effect sizes. + +suppressPackageStartupMessages({ + library(dplyr) + library(readr) + library(tidyr) + library(purrr) + library(effsize) + library(boot) +}) + +source("analysis/expert_survey/config/file_paths.R") +source("analysis/expert_survey/config/analysis_parameters.R") + +PROCESSED_PAIRED_PATH + +cat("Loading paired dataset...\n") +survey_paired <- read_csv(PROCESSED_PAIRED_PATH, show_col_types = FALSE) %>% + mutate( + metric = factor(metric, levels = PRIMARY_METRICS), + question_type = factor(question_type, levels = QUESTION_TYPES), + participant_id = factor(participant_id), + question_id = factor(question_id) + ) + +# ------------------------------------------------------------------------- +# Wilcoxon + effect sizes helper +# ------------------------------------------------------------------------- +run_nonparametric_comprehensive <- function(data, stratum_name) { + cat("\n", strrep("-", 70), "\n") + cat("Stratum:", stratum_name, "\n") + cat(strrep("-", 70), "\n") + + wilcox_result <- wilcox.test( + data$diff_numeric, + alternative = "greater", + conf.int = TRUE, + conf.level = 0.95, + exact = FALSE + ) + + n_pos <- sum(data$diff_numeric > 0) + n_neg <- sum(data$diff_numeric < 0) + n_ties <- sum(data$diff_numeric == 0) + n_total <- nrow(data) + n_nonzero <- n_pos + n_neg + + if (n_nonzero > 0) { + sign_result <- binom.test(n_pos, n_nonzero, p = 0.5, alternative = "greater") + sign_p <- sign_result$p.value + } else { + sign_p <- NA_real_ + } + + cliff_result <- cliff.delta(data$react_score_num, data$base_score_num, paired = TRUE) + cohens_d <- mean(data$diff_numeric) / sd(data$diff_numeric) + + cles_strict <- n_pos / n_total + cles_inclusive <- n_pos / n_total + 0.5 * (n_ties / n_total) + + boot_effect_sizes <- function(df, indices) { + d <- df[indices, ] + n_wins <- sum(d$diff_numeric > 0) + n_ties <- sum(d$diff_numeric == 0) + n_tot <- nrow(d) + c( + strict = n_wins / n_tot, + inclusive = n_wins / n_tot + 0.5 * (n_ties / n_tot) + ) + } + + boot_result <- boot( + data, + statistic = boot_effect_sizes, + R = 10000, + strata = data$participant_id + ) + + boot_ci_strict <- tryCatch({ + ci <- boot.ci(boot_result, index = 1, type = "perc", conf = 0.95) + c(ci$percent[4], ci$percent[5]) + }, error = function(e) { + quantile(boot_result$t[, 1], c(0.025, 0.975), na.rm = TRUE) + }) + + boot_ci_inclusive <- tryCatch({ + ci <- boot.ci(boot_result, index = 2, type = "perc", conf = 0.95) + c(ci$percent[4], ci$percent[5]) + }, error = function(e) { + quantile(boot_result$t[, 2], c(0.025, 0.975), na.rm = TRUE) + }) + + tibble( + stratum = stratum_name, + n_pairs = n_total, + n_improve = n_pos, + n_worsen = n_neg, + n_ties = n_ties, + pct_improve = 100 * n_pos / n_total, + wilcox_V = as.numeric(wilcox_result$statistic), + wilcox_p = wilcox_result$p.value, + hl_estimate = as.numeric(wilcox_result$estimate), + hl_ci_lower = wilcox_result$conf.int[1], + hl_ci_upper = wilcox_result$conf.int[2], + sign_p = sign_p, + cohens_d = cohens_d, + cliff_delta = cliff_result$estimate, + cliff_magnitude = cliff_result$magnitude, + cles_strict = cles_strict, + cles_strict_ci_lower = boot_ci_strict[1], + cles_strict_ci_upper = boot_ci_strict[2], + cles_inclusive = cles_inclusive, + cles_inclusive_ci_lower = boot_ci_inclusive[1], + cles_inclusive_ci_upper = boot_ci_inclusive[2] + ) +} + +# ------------------------------------------------------------------------- +# Stratifications to evaluate +# ------------------------------------------------------------------------- + +strata_list <- list( + "Overall" = survey_paired, + "Factual accuracy" = filter(survey_paired, metric == "Factual accuracy"), + "Level of Granularity" = filter(survey_paired, metric == "Level of Granularity"), + "Relational Depth" = filter(survey_paired, metric == "Relational Depth"), + "Query" = filter(survey_paired, question_type == "Query"), + "Reasoning" = filter(survey_paired, question_type == "Reasoning") +) + +for (metric_name in PRIMARY_METRICS) { + for (qtype in QUESTION_TYPES) { + data_strat <- filter(survey_paired, metric == metric_name, question_type == qtype) + if (nrow(data_strat) >= 10) { + stratum_name <- paste0(metric_name, " - ", qtype) + strata_list[[stratum_name]] <- data_strat + } else { + message("Skipping ", metric_name, " - ", qtype, ": insufficient data (n = ", nrow(data_strat), ")") + } + } +} + +# ------------------------------------------------------------------------- +# Run analyses for each stratum +# ------------------------------------------------------------------------- + +nonparam_results <- map2_dfr( + strata_list, + names(strata_list), + ~ run_nonparametric_comprehensive(.x, .y) +) + +write_csv(nonparam_results, file.path(TABLES_DIR, "nonparametric_tests_all_stratifications.csv")) + +# ------------------------------------------------------------------------- +# Detailed sign-test outputs +# ------------------------------------------------------------------------- + +run_sign_test_question_level <- function(data, stratum_name) { + df_wide <- data %>% + select(participant_id, question_id, metric, question_type, `React-to-Me`, ChatGPT) %>% + pivot_longer( + cols = c(`React-to-Me`, ChatGPT), + names_to = "system", + values_to = "score" + ) %>% + mutate(system = factor(system, levels = c("ChatGPT", "React-to-Me"))) %>% + pivot_wider( + id_cols = c(participant_id, question_id, metric, question_type), + names_from = system, + values_from = score, + names_prefix = "score_" + ) %>% + filter(!is.na(score_ChatGPT) & !is.na(`score_React-to-Me`)) %>% + rename( + score_react = `score_React-to-Me`, + score_chatgpt = score_ChatGPT + ) + + if (nrow(df_wide) == 0) { + return(tibble( + stratum = stratum_name, + analysis_type = "question_level", + n_pairs = 0, + n_positive = 0, + n_negative = 0, + n_zero = 0, + n_nonzero = 0, + sign_test_statistic = NA_real_, + sign_test_p_two_sided = NA_real_, + sign_test_p_one_sided = NA_real_, + median_diff = NA_real_, + mean_diff = NA_real_ + )) + } + + df_wide <- df_wide %>% + mutate( + diff = score_react - score_chatgpt, + diff_sign = case_when( + diff > 0 ~ "positive", + diff < 0 ~ "negative", + TRUE ~ "zero" + ) + ) + + sign_counts <- df_wide %>% + count(diff_sign) %>% + pivot_wider(names_from = diff_sign, values_from = n, values_fill = 0) + + n_positive <- sign_counts$positive %||% 0 + n_negative <- if ("negative" %in% names(sign_counts)) sign_counts$negative else 0 + n_zero <- if ("zero" %in% names(sign_counts)) sign_counts$zero else 0 + n_nonzero <- n_positive + n_negative + + median_diff <- median(df_wide$diff) + mean_diff <- mean(df_wide$diff) + + if (n_nonzero > 0) { + sign_test_p_two_sided <- binom.test( + min(n_positive, n_negative), + n_nonzero, + p = 0.5, + alternative = "two.sided" + )$p.value + + sign_test_p_one_sided <- binom.test( + n_positive, + n_nonzero, + p = 0.5, + alternative = "greater" + )$p.value + + sign_test_statistic <- min(n_positive, n_negative) + } else { + sign_test_p_two_sided <- NA_real_ + sign_test_p_one_sided <- NA_real_ + sign_test_statistic <- NA_real_ + } + + tibble( + stratum = stratum_name, + analysis_type = "question_level", + n_pairs = nrow(df_wide), + n_positive = n_positive, + n_negative = n_negative, + n_zero = n_zero, + n_nonzero = n_nonzero, + sign_test_statistic = sign_test_statistic, + sign_test_p_two_sided = sign_test_p_two_sided, + sign_test_p_one_sided = sign_test_p_one_sided, + median_diff = median_diff, + mean_diff = mean_diff + ) +} + +run_sign_test_participant_level <- function(data, stratum_name) { + participant_preferences <- data %>% + mutate(diff = `React-to-Me` - ChatGPT) %>% + group_by(participant_id) %>% + summarise( + n_obs = n(), + sum_diff = sum(diff), + mean_diff = mean(diff), + n_positive = sum(diff > 0), + n_negative = sum(diff < 0), + n_zero = sum(diff == 0), + participant_preference = case_when( + sum_diff > 0 ~ "positive", + sum_diff < 0 ~ "negative", + TRUE ~ "zero" + ), + .groups = "drop" + ) + + if (nrow(participant_preferences) == 0) { + return(tibble( + stratum = stratum_name, + analysis_type = "participant_level", + n_participants = 0, + n_positive = 0, + n_negative = 0, + n_zero = 0, + sign_test_statistic = NA_real_, + sign_test_p_two_sided = NA_real_, + sign_test_p_one_sided = NA_real_ + )) + } + + pref_counts <- participant_preferences %>% + count(participant_preference) %>% + pivot_wider(names_from = participant_preference, values_from = n, values_fill = 0) + + n_positive <- pref_counts$positive %||% 0 + n_negative <- if ("negative" %in% names(pref_counts)) pref_counts$negative else 0 + n_zero <- if ("zero" %in% names(pref_counts)) pref_counts$zero else 0 + n_nonzero <- n_positive + n_negative + n_participants <- nrow(participant_preferences) + + if (n_nonzero > 0) { + sign_test_p_two_sided <- binom.test( + min(n_positive, n_negative), + n_nonzero, + p = 0.5, + alternative = "two.sided" + )$p.value + + sign_test_p_one_sided <- binom.test( + n_positive, + n_nonzero, + p = 0.5, + alternative = "greater" + )$p.value + + sign_test_statistic <- min(n_positive, n_negative) + } else { + sign_test_p_two_sided <- NA_real_ + sign_test_p_one_sided <- NA_real_ + sign_test_statistic <- NA_real_ + } + + tibble( + stratum = stratum_name, + analysis_type = "participant_level", + n_participants = n_participants, + n_positive = n_positive, + n_negative = n_negative, + n_zero = n_zero, + sign_test_statistic = sign_test_statistic, + sign_test_p_two_sided = sign_test_p_two_sided, + sign_test_p_one_sided = sign_test_p_one_sided + ) +} + +sign_results <- map2( + strata_list, + names(strata_list), + function(data, stratum_name) { + list( + question_level = run_sign_test_question_level(data, stratum_name), + participant_level = run_sign_test_participant_level(data, stratum_name) + ) + } +) + +sign_question_level <- map_dfr(sign_results, ~ .x$question_level, .id = "stratum_id") %>% + select(-stratum_id) +sign_participant_level <- map_dfr(sign_results, ~ .x$participant_level, .id = "stratum_id") %>% + select(-stratum_id) + +write_csv(sign_question_level, file.path(TABLES_DIR, "sign_test_question_level.csv")) +write_csv(sign_participant_level, file.path(TABLES_DIR, "sign_test_participant_level.csv")) + +cat("Nonparametric analysis complete.\n") +cat("Results written to:\n") +cat(" - nonparametric_tests_all_stratifications.csv\n") +cat(" - sign_test_question_level.csv\n") +cat(" - sign_test_participant_level.csv\n") + diff --git a/analysis/expert_survey/scripts/05_multiple_testing.R b/analysis/expert_survey/scripts/05_multiple_testing.R new file mode 100644 index 0000000..034268c --- /dev/null +++ b/analysis/expert_survey/scripts/05_multiple_testing.R @@ -0,0 +1,76 @@ +#!/usr/bin/env Rscript + +# Multiple testing correction using Holm method across families. + +suppressPackageStartupMessages({ + library(dplyr) + library(readr) + library(stringr) +}) + +source("analysis/expert_survey/config/file_paths.R") +source("analysis/expert_survey/config/analysis_parameters.R") + +ordinal_results <- read_csv( + file.path(TABLES_DIR, "ordinal_models_all_stratifications.csv"), + show_col_types = FALSE +) %>% + mutate(interaction = as.logical(interaction)) + +primary_strata <- PRIMARY_METRICS +secondary_strata <- QUESTION_TYPES + +ordinal_filtered <- ordinal_results %>% + filter(str_detect(parameter, "systemReact-to-Me")) + +p_values_primary <- ordinal_filtered %>% + filter(stratum %in% primary_strata, !interaction) %>% + transmute( + stratum, + p_raw = p_value, + test_family = "Primary" + ) %>% + mutate( + p_holm = p.adjust(p_raw, method = "holm"), + significant_raw = p_raw < 0.05, + significant_holm = p_holm < 0.05 + ) + +p_values_secondary <- ordinal_filtered %>% + filter(stratum %in% secondary_strata) %>% + transmute( + stratum, + p_raw = p_value, + test_family = "Secondary" + ) %>% + mutate( + p_holm = p.adjust(p_raw, method = "holm"), + significant_raw = p_raw < 0.05, + significant_holm = p_holm < 0.05 + ) + +exploratory_strata <- setdiff(unique(ordinal_filtered$stratum), c(primary_strata, secondary_strata)) + +p_values_exploratory <- ordinal_filtered %>% + filter(stratum %in% exploratory_strata) %>% + transmute( + stratum, + p_raw = p_value, + test_family = "Exploratory" + ) %>% + mutate( + p_holm = p_raw, + significant_raw = p_raw < 0.05, + significant_holm = significant_raw + ) + +all_corrections <- bind_rows( + p_values_primary, + p_values_secondary, + p_values_exploratory +) + +write_csv(all_corrections, file.path(TABLES_DIR, "multiple_testing_correction_comprehensive.csv")) + +cat("Multiple testing corrections saved to multiple_testing_correction_comprehensive.csv\n") + diff --git a/analysis/expert_survey/scripts/06_power_analysis.R b/analysis/expert_survey/scripts/06_power_analysis.R new file mode 100644 index 0000000..21ca58f --- /dev/null +++ b/analysis/expert_survey/scripts/06_power_analysis.R @@ -0,0 +1,52 @@ +#!/usr/bin/env Rscript + +# Post-hoc power analysis for each stratum using paired t-test approximation. + +suppressPackageStartupMessages({ + library(dplyr) + library(readr) + library(purrr) + library(pwr) +}) + +source("analysis/expert_survey/config/file_paths.R") +source("analysis/expert_survey/config/analysis_parameters.R") + +nonparam_results <- read_csv( + file.path(TABLES_DIR, "nonparametric_tests_all_stratifications.csv"), + show_col_types = FALSE +) + +compute_power <- function(n, d) { + if (is.na(d) || n <= 2 || d == 0) { + return(NA_real_) + } + + res <- tryCatch( + pwr.t.test( + n = n, + d = abs(d), + sig.level = ALPHA_ONESIDED, + type = "paired", + alternative = "greater" + ), + error = function(e) NULL + ) + + if (is.null(res)) { + return(NA_real_) + } + + res$power +} + +power_results <- nonparam_results %>% + mutate( + achieved_power = map2_dbl(n_pairs, cohens_d, compute_power) + ) %>% + select(stratum, n_pairs, cohens_d, achieved_power) + +write_csv(power_results, file.path(TABLES_DIR, "power_analysis_all_stratifications.csv")) + +cat("Power analysis results saved to power_analysis_all_stratifications.csv\n") + diff --git a/analysis/expert_survey/scripts/07_results_compilation.R b/analysis/expert_survey/scripts/07_results_compilation.R new file mode 100644 index 0000000..dc5763e --- /dev/null +++ b/analysis/expert_survey/scripts/07_results_compilation.R @@ -0,0 +1,118 @@ +#!/usr/bin/env Rscript + +# Merge ordinal, nonparametric, Holm corrections, power, and sign tests. + +suppressPackageStartupMessages({ + library(dplyr) + library(readr) + library(tidyr) + library(stringr) +}) + +source("analysis/expert_survey/config/file_paths.R") +source("analysis/expert_survey/config/analysis_parameters.R") + +ordinal_results <- read_csv( + file.path(TABLES_DIR, "ordinal_models_all_stratifications.csv"), + show_col_types = FALSE +) + +nonparam_results <- read_csv( + file.path(TABLES_DIR, "nonparametric_tests_all_stratifications.csv"), + show_col_types = FALSE +) + +holm_results <- read_csv( + file.path(TABLES_DIR, "multiple_testing_correction_comprehensive.csv"), + show_col_types = FALSE +) + +power_results <- read_csv( + file.path(TABLES_DIR, "power_analysis_all_stratifications.csv"), + show_col_types = FALSE +) + +sign_question <- read_csv( + file.path(TABLES_DIR, "sign_test_question_level.csv"), + show_col_types = FALSE +) + +sign_participant <- read_csv( + file.path(TABLES_DIR, "sign_test_participant_level.csv"), + show_col_types = FALSE +) + +ordinal_filtered <- ordinal_results %>% + filter(str_detect(parameter, "systemReact-to-Me")) %>% + select( + stratum, + interaction, + OR, + OR_lower_95, + OR_upper_95, + p_ordinal = p_value + ) + +merge_table <- ordinal_filtered %>% + left_join( + nonparam_results %>% + select( + stratum, + n_pairs, + n_improve, + pct_improve, + cles_strict, + cles_strict_ci_lower, + cles_strict_ci_upper, + cles_inclusive, + cles_inclusive_ci_lower, + cles_inclusive_ci_upper, + cliff_delta, + cohens_d, + wilcox_p, + sign_p + ), + by = "stratum" + ) %>% + left_join( + holm_results %>% + select(stratum, test_family, p_holm, significant_holm), + by = "stratum" + ) %>% + left_join( + power_results %>% + rename( + cohens_d_power = cohens_d, + n_pairs_power = n_pairs + ), + by = "stratum" + ) %>% + left_join( + sign_question %>% + select( + stratum, + n_positive, + n_negative, + n_zero, + n_nonzero, + sign_test_p_two_sided, + sign_test_p_one_sided, + median_diff, + mean_diff + ), + by = "stratum" + ) %>% + mutate( + stratum_type = case_when( + stratum %in% PRIMARY_METRICS ~ "1. Primary (By Metric)", + stratum %in% QUESTION_TYPES ~ "2. Secondary (By Question Type)", + stratum == "Overall" ~ "3. Exploratory (Overall)", + TRUE ~ "4. Exploratory (Metric × Question Type)" + ) + ) %>% + arrange(stratum_type, stratum) + +write_csv(merge_table, file.path(TABLES_DIR, "COMPREHENSIVE_RESULTS_ALL_STRATIFICATIONS.csv")) + +cat("Comprehensive results saved to COMPREHENSIVE_RESULTS_ALL_STRATIFICATIONS.csv\n") + diff --git a/analysis/expert_survey/scripts/08_tables_and_text.R b/analysis/expert_survey/scripts/08_tables_and_text.R new file mode 100644 index 0000000..76e449b --- /dev/null +++ b/analysis/expert_survey/scripts/08_tables_and_text.R @@ -0,0 +1,152 @@ +#!/usr/bin/env Rscript + +# Publication-ready tables and key text snippets (no figures). + +suppressPackageStartupMessages({ + library(dplyr) + library(readr) + library(stringr) +}) + +source("analysis/expert_survey/config/file_paths.R") +source("analysis/expert_survey/config/analysis_parameters.R") + +comprehensive_results <- read_csv( + file.path(TABLES_DIR, "COMPREHENSIVE_RESULTS_ALL_STRATIFICATIONS.csv"), + show_col_types = FALSE +) + +question_strata <- filter(comprehensive_results, stratum %in% QUESTION_TYPES) + +main_results_formatted <- comprehensive_results %>% + filter(stratum_type == "1. Primary (By Metric)") %>% + mutate( + p_onesided = p_ordinal / 2, + p_holm_onesided = if_else(!is.na(p_holm), p_holm / 2, NA_real_), + `Odds Ratio (95% CI)` = sprintf("%.2f (%.2f-%.2f)", OR, OR_lower_95, OR_upper_95), + `P-value (one-sided)‡` = format.pval(p_onesided, digits = 3, eps = 0.001), + `Adjusted P†` = if_else( + !is.na(p_holm_onesided), + format.pval(p_holm_onesided, digits = 3, eps = 0.001), + "\u2014" + ), + `Strict Win Rate (95% CI)` = sprintf( + "%.2f (%.2f-%.2f)", + cles_strict, + cles_strict_ci_lower, + cles_strict_ci_upper + ), + `CLES with Ties (95% CI)` = sprintf( + "%.2f (%.2f-%.2f)", + cles_inclusive, + cles_inclusive_ci_lower, + cles_inclusive_ci_upper + ), + `% Improved` = sprintf("%.1f%%", pct_improve), + `Superiority Claim` = if_else(p_onesided < ALPHA_ONESIDED & OR > 1, "Supported", "Not Supported") + ) %>% + select( + Stratum = stratum, + `N Pairs` = n_pairs, + `% Improved`, + `Odds Ratio (95% CI)`, + `P-value (one-sided)‡`, + `Adjusted P†`, + `Strict Win Rate (95% CI)`, + `CLES with Ties (95% CI)`, + `Superiority Claim` + ) + +write_csv(main_results_formatted, file.path(TABLES_DIR, "TABLE_MAIN_RESULTS_FORMATTED.csv")) + +stratified_results_formatted <- comprehensive_results %>% + filter(stratum_type %in% c("2. Secondary (By Question Type)", "3. Exploratory (Overall)", "4. Exploratory (Metric × Question Type)")) %>% + mutate( + `Odds Ratio (95% CI)` = sprintf("%.2f (%.2f-%.2f)", OR, OR_lower_95, OR_upper_95), + `P-value` = format.pval(p_ordinal, digits = 3, eps = 0.001), + `Strict Win Rate (95% CI)` = sprintf( + "%.2f (%.2f-%.2f)", + cles_strict, + cles_strict_ci_lower, + cles_strict_ci_upper + ), + `CLES with Ties (95% CI)` = sprintf( + "%.2f (%.2f-%.2f)", + cles_inclusive, + cles_inclusive_ci_lower, + cles_inclusive_ci_upper + ), + `% Improved` = sprintf("%.1f%%", pct_improve) + ) %>% + select( + Category = stratum_type, + Stratum = stratum, + `N Pairs` = n_pairs, + `% Improved`, + `Odds Ratio (95% CI)`, + `P-value`, + `Strict Win Rate (95% CI)`, + `CLES with Ties (95% CI)`, + `Analysis Type` = test_family + ) + +write_csv( + stratified_results_formatted, + file.path(TABLES_DIR, "TABLE_STRATIFIED_RESULTS_FORMATTED.csv") +) + +key_stats <- comprehensive_results %>% + mutate( + p_onesided = p_ordinal / 2, + summary_line = case_when( + stratum %in% PRIMARY_METRICS ~ sprintf( + "%s: OR = %.2f [%.2f-%.2f], p(one-sided) = %s (Holm adj: %s), CLES = %.2f, Improvement = %.1f%%", + stratum, + OR, + OR_lower_95, + OR_upper_95, + format.pval(p_onesided, digits = 3), + format.pval(p_holm / 2, digits = 3), + cles_strict, + pct_improve + ), + stratum %in% QUESTION_TYPES ~ sprintf( + "%s questions: OR = %.2f [%.2f-%.2f], p = %s, CLES = %.2f, Improvement = %.1f%%", + stratum, + OR, + OR_lower_95, + OR_upper_95, + format.pval(p_ordinal, digits = 3), + cles_strict, + pct_improve + ), + stratum == "Overall" ~ sprintf( + "Overall: OR = %.2f [%.2f-%.2f], p = %s, CLES = %.2f, Improvement = %.1f%%", + OR, + OR_lower_95, + OR_upper_95, + format.pval(p_ordinal, digits = 3), + cles_strict, + pct_improve + ), + TRUE ~ sprintf( + "%s: OR = %.2f [%.2f-%.2f], p = %s, CLES = %.2f, Improvement = %.1f%%", + stratum, + OR, + OR_lower_95, + OR_upper_95, + format.pval(p_ordinal, digits = 3), + cles_strict, + pct_improve + ) + ) + ) %>% + pull(summary_line) + +writeLines(key_stats, file.path(TABLES_DIR, "KEY_STATISTICS_FOR_TEXT.txt")) + +cat("Publication tables saved:\n") +cat(" - TABLE_MAIN_RESULTS_FORMATTED.csv\n") +cat(" - TABLE_STRATIFIED_RESULTS_FORMATTED.csv\n") +cat(" - KEY_STATISTICS_FOR_TEXT.txt\n") + diff --git a/analysis/expert_survey/utils/data_preparation.R b/analysis/expert_survey/utils/data_preparation.R new file mode 100644 index 0000000..4163431 --- /dev/null +++ b/analysis/expert_survey/utils/data_preparation.R @@ -0,0 +1,134 @@ +# Data Preparation Utilities for Expert Survey Analysis + +library(readr) +library(dplyr) +library(tidyr) +library(stringr) + +#' Load and validate expert survey data +#' +#' @param file_path Path to the survey CSV. +#' @return A tibble with cleaned column names and string values trimmed. +load_and_validate_data <- function(file_path) { + df <- read_csv( + file_path, + col_types = cols( + participant_id = col_character(), + question_id = col_double(), + question_type = col_character(), + metric = col_character(), + `React-to-Me score` = col_double(), + `GPT-baseline score` = col_double() + ) + ) + + df <- df %>% + mutate( + participant_id = str_trim(participant_id), + question_type = str_trim(question_type), + metric = str_trim(metric) + ) %>% + rename( + react_score = `React-to-Me score`, + baseline_score = `GPT-baseline score` + ) + + expected_participants <- 10 + expected_questions <- 15 + expected_metrics <- 3 + expected_systems <- 2 + + n_participants <- n_distinct(df$participant_id) + n_questions <- n_distinct(df$question_id) + n_metrics <- n_distinct(df$metric) + + if (n_participants != expected_participants) { + stop( + sprintf( + "Unexpected number of participants: %s (expected %s)", + n_participants, + expected_participants + ) + ) + } + + if (n_questions != expected_questions) { + stop( + sprintf( + "Unexpected number of questions: %s (expected %s)", + n_questions, + expected_questions + ) + ) + } + + if (n_metrics != expected_metrics) { + stop( + sprintf( + "Unexpected number of metrics: %s (expected %s)", + n_metrics, + expected_metrics + ) + ) + } + + if (any(df$react_score < 1 | df$react_score > 4, na.rm = TRUE) || + any(df$baseline_score < 1 | df$baseline_score > 4, na.rm = TRUE)) { + stop("Scores must be within the 1-4 scale.") + } + + df +} + +#' Convert cleaned data to long format for modeling +#' +#' @param df Cleaned wide-format tibble. +#' @return Long-format tibble with system and score columns. +to_long_format <- function(df) { + df_long <- df %>% + pivot_longer( + cols = c(react_score, baseline_score), + names_to = "system", + values_to = "score" + ) %>% + mutate( + system = recode( + system, + react_score = "React-to-Me", + baseline_score = "ChatGPT" + ), + system = factor(system, levels = c("ChatGPT", "React-to-Me")), + metric = factor( + metric, + levels = c("Factual accuracy", "Level of Granularity", "Relational Depth") + ), + question_type = factor(question_type, levels = c("Query", "Reasoning")), + participant_id = factor(participant_id), + question_id = factor(question_id) + ) + + df_long +} + +#' Create paired dataset for nonparametric analyses +#' +#' @param df_long Long-format tibble produced by to_long_format(). +#' @return Tibble with paired scores and difference columns. +create_paired_dataset <- function(df_long) { + df_wide <- df_long %>% + select(participant_id, question_id, question_type, metric, system, score) %>% + pivot_wider( + names_from = system, + values_from = score + ) %>% + mutate( + diff_numeric = `React-to-Me` - ChatGPT, + react_score_num = `React-to-Me`, + base_score_num = ChatGPT, + participant_num = as.numeric(participant_id), + question_num = as.numeric(question_id) + ) + + df_wide +} + diff --git a/bin/chat-chainlit.py b/bin/chat-chainlit.py index fa4faf6..225e0a5 100644 --- a/bin/chat-chainlit.py +++ b/bin/chat-chainlit.py @@ -3,6 +3,7 @@ import chainlit as cl from chainlit.data.base import BaseDataLayer from chainlit.data.sql_alchemy import SQLAlchemyDataLayer +from chainlit.oauth_providers import providers from chainlit.types import ThreadDict from dotenv import load_dotenv from langchain_community.callbacks import OpenAICallbackHandler @@ -10,11 +11,17 @@ from agent.graph import AgentGraph from agent.profiles import ProfileName, get_chat_profiles from agent.profiles.base import OutputState -from util.chainlit_helpers import (PrefixedS3StorageClient, is_feature_enabled, - message_rate_limited, save_openai_metrics, - static_messages, update_search_results) +from util.chainlit_helpers import ( + PrefixedS3StorageClient, + is_feature_enabled, + message_rate_limited, + save_openai_metrics, + static_messages, + update_search_results, +) from util.config_yml import Config, TriggerEvent from util.logging import logging +from util.orcid_provider import ORCIDOAuthProvider load_dotenv() config: Config | None = Config.from_yaml() @@ -31,6 +38,7 @@ if POSTGRES_CHAINLIT_DB and POSTGRES_USER and POSTGRES_PASSWORD: CHAINLIT_DB_URI = f"postgresql+psycopg://{POSTGRES_USER}:{POSTGRES_PASSWORD}@postgres:5432/{POSTGRES_CHAINLIT_DB}?sslmode=disable" + storage_client: PrefixedS3StorageClient | None if S3_BUCKET and S3_CHAINLIT_PREFIX: storage_client = PrefixedS3StorageClient(S3_BUCKET, S3_CHAINLIT_PREFIX) else: @@ -46,6 +54,9 @@ def get_data_layer() -> BaseDataLayer: else: logging.warning("POSTGRES_CHAINLIT_DB undefined; Chainlit persistence disabled.") +if os.getenv("OAUTH_ORCID_CLIENT_ID") and not any(p.id == "orcid" for p in providers): + providers.append(ORCIDOAuthProvider()) + if os.getenv("CHAINLIT_AUTH_SECRET"): @cl.oauth_callback diff --git a/bin/chat-fastapi.py b/bin/chat-fastapi.py index 82dda31..a984cb7 100644 --- a/bin/chat-fastapi.py +++ b/bin/chat-fastapi.py @@ -1,7 +1,9 @@ import hashlib import hmac import os +from collections.abc import Awaitable, Callable from string import Template +from urllib.parse import urlsplit import requests from chainlit.utils import mount_chainlit @@ -59,28 +61,31 @@ def verify_secure_cookie(cookie_value: str) -> bool: @app.middleware("http") -async def verify_captcha_middleware(request: Request, call_next): - if ( - CHAINLIT_URI - and not request.url.path.startswith(CHAINLIT_URI) - and request.url.path[-1] != "/" - ): - return RedirectResponse(url=f"{request.url.path}/") +async def verify_captcha_middleware( + request: Request, call_next: Callable[[Request], Awaitable[Response]] +) -> Response: + path = request.url.path + if CHAINLIT_URI and path == CHAINLIT_URI and not path.endswith("/"): + # Safety: ensure the path is a clean, simple relative path with no + # scheme/host/dot-segments before echoing it back. + clean_path = urlsplit(path).path + if ".." not in clean_path: + return RedirectResponse(url=f"{clean_path}/") + # Allow access to CAPTCHA pages and static files if ( - request.url.path + path in [ "/chat/", f"{CHAINLIT_URI}/verify_captcha", f"{CHAINLIT_URI}/verify_captcha_page", f"{CHAINLIT_URI}/static", ] - or request.url.path.startswith("/static") + or path.startswith("/static") or not os.getenv("CLOUDFLARE_SECRET_KEY") - or (CHAINLIT_URI and not request.url.path.startswith(CHAINLIT_URI)) + or (CHAINLIT_URI and not path.startswith(CHAINLIT_URI)) ): - response = await call_next(request) - return response + return await call_next(request) host = request.headers.get("referer") if host and host.startswith("http:"): @@ -96,13 +101,12 @@ async def verify_captcha_middleware(request: Request, call_next): if not captcha_verified or not verify_secure_cookie(captcha_verified): return RedirectResponse(url=f"{CHAINLIT_URI}/verify_captcha_page") - response = await call_next(request) - return response + return await call_next(request) # Serve the CAPTCHA verification page (basic HTML form) @app.get(f"{CHAINLIT_URI}/verify_captcha_page") -async def captcha_page(): +async def captcha_page() -> Response: html_content = f""" @@ -128,7 +132,7 @@ async def captcha_page(): @app.post(f"{CHAINLIT_URI}/verify_captcha") -async def verify_captcha(request: Request): +async def verify_captcha(request: Request) -> Response: form_data = await request.form() cf_turnstile_response = form_data.get("cf-turnstile-response") if not isinstance(cf_turnstile_response, str): @@ -167,7 +171,7 @@ async def verify_captcha(request: Request): } # Perform request to Cloudflare Turnstile verification endpoint - response = requests.post(url, data=data) + response = requests.post(url, data=data, timeout=10) result = response.json() # If CAPTCHA validation fails, return an error @@ -199,7 +203,7 @@ async def verify_captcha(request: Request): @app.get("/chat/") -async def landing_page(): +async def landing_page() -> HTMLResponse: html_content = Template( """ diff --git a/bin/embeddings_manager b/bin/embeddings_manager index 385e315..a7c4d74 100755 --- a/bin/embeddings_manager +++ b/bin/embeddings_manager @@ -15,6 +15,7 @@ from botocore.client import Config from data_generation.alliance import generate_alliance_embeddings from data_generation.reactome import generate_reactome_embeddings from data_generation.uniprot import generate_uniprot_embeddings +from data_generation.userguide import generate_userguide_embeddings from util.embedding_environment import EM_ARCHIVE, EmbeddingEnvironment S3_BUCKET = "download.reactome.org" @@ -48,6 +49,7 @@ class EmbeddingSelection(NamedTuple): def pull(embedding: EmbeddingSelection): + EM_ARCHIVE.mkdir(parents=True, exist_ok=True) embedding_path:Path = embedding.path(check_exists=False) zip_tmpfile:Path = EM_ARCHIVE / "tmp.zip" s3 = boto3.resource("s3", config=Config(signature_version=UNSIGNED)) @@ -87,10 +89,16 @@ def make( os.environ["HUGGINGFACEHUB_API_TOKEN"] = hf_key if embedding.db == "reactome": generate_reactome_embeddings(str(embedding_path), hf_model=embedding.model, **kwargs) + elif embedding.db == "plantreactome": + generate_reactome_embeddings(str(embedding_path), hf_model=embedding.model, **kwargs) elif embedding.db == "uniprot": generate_uniprot_embeddings(embedding_path, hf_model=embedding.model, **kwargs) elif embedding.db == "alliance": generate_alliance_embeddings(str(embedding_path), hf_model=embedding.model, **kwargs) + elif embedding.db == "userguide": + generate_userguide_embeddings( + str(embedding_path), hf_model=embedding.model, **kwargs + ) else: raise NotImplementedError(f"db: {embedding.db}") use(embedding) diff --git a/bin/export_nologin_usage.py b/bin/export_nologin_usage.py index 17adc2f..e755627 100644 --- a/bin/export_nologin_usage.py +++ b/bin/export_nologin_usage.py @@ -12,7 +12,7 @@ def build_query() -> str: - query = """ + return """ SELECT thread_id, checkpoint_id, @@ -25,19 +25,17 @@ def build_query() -> str: ORDER BY checkpoint->'ts'; """ - return query -def main(records_dir: Path): +def main(records_dir: Path) -> None: records_dir.mkdir(exist_ok=True) query: str = build_query() - with psycopg.connect(LANGGRAPH_NOLOGIN_DB_URI) as conn: - with conn.cursor() as cur: - cur.execute(query) - header = [col.name for col in cur.description] if cur.description else None - records = cur.fetchall() + with psycopg.connect(LANGGRAPH_NOLOGIN_DB_URI) as conn, conn.cursor() as cur: + cur.execute(query) + header = [col.name for col in cur.description] if cur.description else None + records = cur.fetchall() if len(records) == 0: print("No new records found.") diff --git a/bin/export_records.py b/bin/export_records.py index fa091ec..f7d62dc 100644 --- a/bin/export_records.py +++ b/bin/export_records.py @@ -11,10 +11,14 @@ CHAINLIT_DB_URI = f"postgresql://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}@postgres:5432/{os.getenv('POSTGRES_CHAINLIT_DB')}?sslmode=disable" -def build_query(since_timestamp: str | None) -> str: - if since_timestamp is None: - since_timestamp = "" - query = f""" +def build_query() -> str: + """The since-timestamp is bound as a parameter, not interpolated. + + It comes from a previously exported CSV, i.e. from values the database + produced -- but building SQL by string interpolation is the wrong habit to + keep in a script that runs against the production chat history. + """ + return """ SELECT steps."threadId", steps."createdAt", @@ -31,7 +35,7 @@ def build_query(since_timestamp: str | None) -> str: threads ON steps."threadId" = threads.id WHERE steps.type IN ('user_message', 'assistant_message') AND - steps."createdAt" > '{since_timestamp}' + steps."createdAt" > %(since_timestamp)s ORDER BY ( SELECT MIN(s."createdAt") @@ -40,29 +44,26 @@ def build_query(since_timestamp: str | None) -> str: ), steps."createdAt"; """ - return query def last_record_timestamp(records_dir: Path) -> str | None: - record_names: list[str] = list(f.stem for f in records_dir.glob("records_*.csv")) + record_names: list[str] = [f.stem for f in records_dir.glob("records_*.csv")] if len(record_names) > 0: last_record: str = max(record_names) return last_record[len("records_") :] - else: - return None + return None -def main(records_dir: Path): +def main(records_dir: Path) -> None: records_dir.mkdir(exist_ok=True) since_timestamp: str | None = last_record_timestamp(records_dir) - query: str = build_query(since_timestamp) + query: str = build_query() - with psycopg.connect(CHAINLIT_DB_URI) as conn: - with conn.cursor() as cur: - cur.execute(query) - header = [col.name for col in cur.description] if cur.description else None - records = cur.fetchall() + with psycopg.connect(CHAINLIT_DB_URI) as conn, conn.cursor() as cur: + cur.execute(query, {"since_timestamp": since_timestamp or ""}) + header = [col.name for col in cur.description] if cur.description else None + records = cur.fetchall() if len(records) == 0: print("No new records found.") diff --git a/bin/retrieval_baseline b/bin/retrieval_baseline new file mode 100755 index 0000000..89e543c --- /dev/null +++ b/bin/retrieval_baseline @@ -0,0 +1,298 @@ +#!/usr/bin/env python +"""Record what each retriever returns, so a change to retrieval can be diffed. + +Retrieval quality has no right answer, only a "did this change". This captures +the documents each retriever returns for a fixed question set, per Chroma +collection, and writes them to JSON. Run it before a change and after, then +`compare` the two files. + +Two of the pipeline's steps call an LLM and so are not reproducible: the +multi-query expansion, and SelfQueryRetriever's translation of a question into a +metadata filter. This tool deliberately skips the expansion and feeds each +question to the retrievers directly. Measured across two identical runs, BM25 is +byte-identical on 80/80 question-collections and plain vector on 78/80. +SelfQuery turns out to be stable too (78/80), so its divergence from plain +vector search is a real difference in behaviour, not model noise. + + ./bin/retrieval_baseline capture --out before.json + ./bin/retrieval_baseline overlap before.json # selfquery vs plain vector + ./bin/retrieval_baseline compare before.json after.json + +Requires an installed reactome embeddings bundle and OPENAI_API_KEY. +""" + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any + +import nltk +from dotenv import load_dotenv +from langchain.retrievers.self_query.base import SelfQueryRetriever +from langchain_chroma.vectorstores import Chroma +from langchain_community.document_loaders.csv_loader import CSVLoader +from langchain_community.retrievers import BM25Retriever +from langchain_core.documents import Document +from langchain_core.retrievers import BaseRetriever +from nltk.tokenize import word_tokenize + +from agent.models import get_embedding, get_llm +from retrievers.csv_chroma import chroma_settings, list_chroma_subdirectories +from retrievers.reactome.metadata_info import ( + reactome_descriptions_info, + reactome_field_info, +) +from util.embedding_environment import EmbeddingEnvironment + +DEFAULT_QUESTIONS = Path(__file__).parent.parent / "tests" / "golden" / "questions.txt" +DEFAULT_K = 10 + +# {question: {collection: {retriever: [doc id, ...]}}} +Capture = dict[str, dict[str, dict[str, list[str]]]] + + +def read_questions(path: Path) -> list[str]: + lines = path.read_text().splitlines() + return [ln.strip() for ln in lines if ln.strip() and not ln.startswith("#")] + + +def doc_id(doc: Document) -> str: + """Reactome stable IDs survive a bundle rebuild; row numbers do not.""" + st_id = doc.metadata.get("st_id") + return str(st_id) if st_id else f"row:{doc.metadata.get('row', '?')}" + + +def build_retrievers( + embeddings_dir: Path, collection: str, k: int, *, with_selfquery: bool +) -> dict[str, BaseRetriever]: + llm = get_llm("openai", "gpt-4o-mini") + embedding = get_embedding("openai", "text-embedding-3-large") + + csv_path = embeddings_dir / "csv_files" / f"{collection}.csv" + bm25 = BM25Retriever.from_documents( + CSVLoader(file_path=str(csv_path)).load(), + preprocess_func=lambda text: word_tokenize(text.casefold(), language="english"), + ) + bm25.k = k + + vectordb = Chroma( + persist_directory=str(embeddings_dir / collection), + embedding_function=embedding, + client_settings=chroma_settings, + ) + + retrievers: dict[str, BaseRetriever] = { + "bm25": bm25, + # The plain semantic retriever: similarity search with no LLM in the loop. + "vector": vectordb.as_retriever(search_kwargs={"k": k}), + } + if with_selfquery: + retrievers["selfquery"] = SelfQueryRetriever.from_llm( + llm=llm, + vectorstore=vectordb, + document_contents=reactome_descriptions_info[collection], + metadata_field_info=reactome_field_info[collection], + search_kwargs={"k": k}, + ) + return retrievers + + +def capture(args: argparse.Namespace) -> None: + embeddings_dir: Path | None = args.embeddings_dir + if embeddings_dir is None: + raise SystemExit( + "No reactome embeddings installed. " + "Run ./bin/embeddings_manager install , or pass " + "--embeddings-dir." + ) + questions = read_questions(args.questions) + collections = sorted(list_chroma_subdirectories(embeddings_dir)) + if not collections: + raise SystemExit(f"No Chroma collections found under {embeddings_dir}") + + print( + f"{len(questions)} questions x {len(collections)} collections " + f"({', '.join(collections)}), k={args.k}", + file=sys.stderr, + ) + + result: Capture = {q: {} for q in questions} + for collection in collections: + print(f" {collection}: building retrievers...", file=sys.stderr) + retrievers = build_retrievers( + embeddings_dir, collection, args.k, with_selfquery=args.with_selfquery + ) + for i, question in enumerate(questions, start=1): + print(f" [{i}/{len(questions)}] {question[:60]}", file=sys.stderr) + per_retriever: dict[str, list[str]] = {} + for name, retriever in retrievers.items(): + docs = retriever.invoke(question) + per_retriever[name] = [doc_id(d) for d in docs] + result[question][collection] = per_retriever + + payload: dict[str, Any] = { + "bundle": str(embeddings_dir), + "k": args.k, + "collections": collections, + "results": result, + } + args.out.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + print(f"\nWrote {args.out}", file=sys.stderr) + + +def _jaccard(a: list[str], b: list[str]) -> float: + """Overlap of the DISTINCT documents returned. + + Note this collapses duplicates, and duplicates are themselves a finding -- + vector search on `reactions` returns ~5 distinct documents in its top 10. + So this measures agreement about *which* documents, saying nothing about + ordering or about how much of k is wasted. `_rank_agreement` covers ordering; + `distinct_ratio` in the report covers the waste. + """ + sa, sb = set(a), set(b) + if not sa and not sb: + return 1.0 + return len(sa & sb) / len(sa | sb) + + +def _rank_agreement(a: list[str], b: list[str]) -> float: + """Fraction of positions holding the same document, order included.""" + if not a and not b: + return 1.0 + n = max(len(a), len(b)) + same = sum(1 for i in range(min(len(a), len(b))) if a[i] == b[i]) + return same / n + + +def overlap(args: argparse.Namespace) -> None: + """How far does plain vector search land from SelfQuery on the same question? + + This is the evidence for whether SelfQueryRetriever can be replaced by a + plain semantic retriever: high overlap means the LLM-built metadata filter + is rarely changing which documents come back. + """ + payload = json.loads(args.capture.read_text()) + results: Capture = payload["results"] + + sets: dict[str, list[float]] = {} + ranks: dict[str, list[float]] = {} + distinct: dict[str, list[float]] = {} + for _question, collections in results.items(): + for collection, retrievers in collections.items(): + if "selfquery" not in retrievers or "vector" not in retrievers: + raise SystemExit( + "This capture has no selfquery results. " + "Re-run capture with --with-selfquery." + ) + sq, vec = retrievers["selfquery"], retrievers["vector"] + sets.setdefault(collection, []).append(_jaccard(sq, vec)) + ranks.setdefault(collection, []).append(_rank_agreement(sq, vec)) + distinct.setdefault(collection, []).append( + len(set(vec)) / len(vec) if vec else 1.0 + ) + + def mean(xs: list[float]) -> float: + return sum(xs) / len(xs) + + print("selfquery vs plain vector (1.0 = identical)\n") + print( + f" {'collection':<14}{'set overlap':>13}{'rank agree':>12}{'distinct/k':>13}" + ) + for collection in sorted(sets): + print( + f" {collection:<14}{mean(sets[collection]):>13.2f}" + f"{mean(ranks[collection]):>12.2f}{mean(distinct[collection]):>13.2f}" + ) + print( + f"\n overall {mean([s for v in sets.values() for s in v]):>13.2f}" + f"{mean([s for v in ranks.values() for s in v]):>12.2f}" + f"{mean([s for v in distinct.values() for s in v]):>13.2f}" + ) + print( + "\n set overlap ignores duplicates and order; rank agree counts a " + "document\n only if it is in the same position; distinct/k is the " + "share of the vector\n result that is not a repeat (see issue #169)." + ) + + +def compare(args: argparse.Namespace) -> None: + before: Capture = json.loads(args.before.read_text())["results"] + after: Capture = json.loads(args.after.read_text())["results"] + + changed = 0 + total = 0 + for question, collections in before.items(): + if question not in after: + print(f"! question missing from the newer capture: {question}") + continue + for collection, retrievers in collections.items(): + for name, ids in retrievers.items(): + new_ids = after[question].get(collection, {}).get(name) + if new_ids is None: + continue + total += 1 + if new_ids == ids: + continue + changed += 1 + same_set = set(new_ids) == set(ids) + kind = "REORDERED" if same_set else "DIFFERENT DOCS" + print(f"\n{kind} [{collection}/{name}] {question}") + print(f" before: {ids}") + print(f" after: {new_ids}") + if not same_set: + print(f" dropped: {sorted(set(ids) - set(new_ids))}") + print(f" added: {sorted(set(new_ids) - set(ids))}") + + print(f"\n{changed}/{total} retriever results changed.") + if changed: + sys.exit(1) + + +def main() -> None: + load_dotenv() + try: + nltk.data.find("tokenizers/punkt_tab") + except LookupError: + print("Downloading nltk punkt_tab...", file=sys.stderr) + nltk.download("punkt_tab", quiet=True) + if not os.getenv("OPENAI_API_KEY"): + raise SystemExit("OPENAI_API_KEY is not set.") + + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(required=True) + + p_capture = sub.add_parser("capture", help="Record retriever output to JSON") + p_capture.add_argument("--out", type=Path, required=True) + p_capture.add_argument("--questions", type=Path, default=DEFAULT_QUESTIONS) + p_capture.add_argument("--k", type=int, default=DEFAULT_K) + p_capture.add_argument( + "--embeddings-dir", + type=Path, + default=EmbeddingEnvironment.get_dir("reactome"), + ) + p_capture.add_argument( + "--with-selfquery", + action="store_true", + help="Also run SelfQueryRetriever (one LLM call per question per collection)", + ) + p_capture.set_defaults(func=capture) + + p_overlap = sub.add_parser( + "overlap", help="Compare selfquery against plain vector within one capture" + ) + p_overlap.add_argument("capture", type=Path) + p_overlap.set_defaults(func=overlap) + + p_compare = sub.add_parser("compare", help="Diff two captures") + p_compare.add_argument("before", type=Path) + p_compare.add_argument("after", type=Path) + p_compare.set_defaults(func=compare) + + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/chainlit.md b/chainlit.md index 908001a..0057c91 100644 --- a/chainlit.md +++ b/chainlit.md @@ -1,37 +1,37 @@ # React-to-me -Welcome to React-to-me, your interactive chatbot for exploring Reactome! +Welcome to React-to-me, your interactive chatbot for exploring Plant Reactome! ## About -React-to-me is a specialized chatbot designed to provide fast and reliable answers about biological pathways and processes from the Reactome knowledgebase. Whether you're a researcher, student, or just curious about biology, React-to-me can help you access and understand complex biological information quickly and efficiently. +React-to-me is a specialized chatbot designed to provide fast and reliable answers about biological pathways and processes from the Plant Reactome knowledgebase. Whether you're a researcher, student, or just curious about biology, React-to-me can help you access and understand complex biological information quickly and efficiently. ## Features -With Reactome Chatbot, you can: +With Plant Reactome Chatbot, you can: - Ask questions about biological pathways and get responses in real-time. -- Browse through Reactome's extensive collection of pathways and related content. +- Browse through Plant Reactome's extensive collection of pathways and related content. - Access information in multiple languages, allowing you to ask questions and receive information in your preferred language. ## How to Use -Simply type your question about any Reactome content directly into the chat window, and the chatbot will provide you with detailed answers accompanied by links to relevant pages within the Reactome portal. +Simply type your question about any Plant Reactome content directly into the chat window, and the chatbot will provide you with detailed answers accompanied by links to relevant pages within the Plant Reactome portal. -Feel free to explore any topic within Reactome's database. Whether your questions are broad or highly specific, the chatbot will do its best to provide you with accurate and helpful responses! +Feel free to explore any topic within Plant Reactome's database. Whether your questions are broad or highly specific, the chatbot will do its best to provide you with accurate and helpful responses! ## Pathway Recommendations Explore pathways such as: -- Cell Cycle -- Glycolysis -- Apoptosis -- Signal Transduction +- Circadian rhythm +- Mitosis +- Detoxification +- Root gravitropism ## Additional Resources -- [Reactome Website](https://reactome.org/) -- [Reactome GitHub Repository](https://github.com/reactome) -- [Reactome Twitter](https://twitter.com/reactome) +- [Plant Reactome Website](https://plantreactome.gramene.org/) +- [Plant Reactome GitHub Repository](https://github.com/plantreactome) +- [Planteome Twitter](https://twitter.com/planteome) Happy exploring with React-to-me! @@ -41,4 +41,4 @@ _This chatbot uses large language model (LLM) technology to assist with question _The information you provide may be retained in accordance with Reactome’s AI provider’s retention policy, which is located [here](https://openai.com/enterprise-privacy/). Do not share sensitive, personal or confidential information._ -_The chatbot does not substitute for expert curation or peer-reviewed sources and is not a suitable resource for clinical decisions. Users are responsible for validating any output before using it for research, publication, or medical decisions. Any use of this chatbot is subject to Reactome’s [disclaimer](https://reactome.org/about/disclaimer)._ +_The chatbot does not substitute for expert curation or peer-reviewed sources and is not a suitable resource for clinical decisions. Users are responsible for validating any output before using it for research, publication, or medical decisions. Any use of this chatbot is subject to Plant Reactome’s [disclaimer](https://plantreactome.gramene.org/index.php?option=com_content&view=article&id=17&Itemid=254&lang=en)._ diff --git a/deploy/beta/README.md b/deploy/beta/README.md new file mode 100644 index 0000000..2c2cdc5 --- /dev/null +++ b/deploy/beta/README.md @@ -0,0 +1,77 @@ +# Running the chatbot behind beta.reactome.org/chat + +First cut: **guest access only**. That needs no Google OAuth (whose redirect URIs +are registered for `reactome.org`, not beta) and no Postgres — the LangGraph +checkpointer falls back to `MemorySaver`, so conversations live in memory and are +lost on restart. Chat history and the `/chat/personal` route come later. + +Leaving `CLOUDFLARE_SECRET_KEY` unset makes the captcha middleware bypass itself, +which is what we want: the Turnstile site key is bound to `reactome.org`. + +## 1. Embeddings + +The container answers nothing without a bundle. Available on S3 (probe with +`curl -I`, anonymous listing is denied): + +| bundle | size | +|---|---| +| `openai/text-embedding-3-large/reactome/Release95` | 1.3 GB | +| `openai/text-embedding-3-large/reactome/Release94` | 1.4 GB | +| `openai/text-embedding-3-large/reactome/Release91` | 1.0 GB | +| `openai/text-embedding-3-large/reactome/Release90` | 2.0 GB | +| `openai/text-embedding-3-large/reactome/Release89` | 1.9 GB | + +```bash +mkdir -p embeddings +docker run --rm -v "$PWD/embeddings:/app/embeddings" \ + public.ecr.aws/reactome/reactome-chatbot:e398a37 \ + ./bin/embeddings_manager install openai/text-embedding-3-large/reactome/Release95 +``` + +Budget roughly 1.3 GB download plus 2–3 GB extracted, on top of a ~4–6 GB image. + +## 2. Config + +```bash +cp config_default.yml config.yml # prod's config.yml is byte-identical to this +``` + +## 3. Environment + +Copy `env.beta.template` to `.env.beta` and fill in the two keys. Do not reuse +prod's `CHAINLIT_URL`, `CHAINLIT_ROOT_PATH` or OAuth values — they point at +`reactome.org` and will break asset URLs and logins on beta. + +## 4. Run + +Bound to loopback: Apache is the only thing that should reach it. + +```bash +docker run -d --name biochat_beta_guest --restart unless-stopped \ + --env-file .env.beta \ + -v "$PWD/embeddings:/app/embeddings" \ + -v "$PWD/config.yml:/app/config.yml" \ + -p 127.0.0.1:8000:8000 \ + public.ecr.aws/reactome/reactome-chatbot:e398a37 + +curl -s localhost:8000/chat/ | grep -o React-to-Me # should print React-to-Me +``` + +The image tag matches what production runs today, so this is a like-for-like +baseline to compare against after the dependency upgrade. + +Note: the landing page shows both a **Guest Access** and a **Log In** button. Only +Guest Access works in this setup; wiring Log In needs a second container on :8001 +with `CHAINLIT_URI=/chat/personal`, plus Postgres and OAuth. + +## 5. Apache + +See `../../../WebsiteAngular/deploy/apache/install-beta-chat-proxy.sh`. + +```bash +sudo a2enmod proxy proxy_http proxy_wstunnel rewrite +sudo ~/git/WebsiteAngular/deploy/apache/install-beta-chat-proxy.sh +``` + +`proxy_wstunnel` is not enabled on this host today and Chainlit needs it — +without it the UI renders and then hangs with no replies. diff --git a/deploy/beta/env.beta.template b/deploy/beta/env.beta.template new file mode 100644 index 0000000..d3f2787 --- /dev/null +++ b/deploy/beta/env.beta.template @@ -0,0 +1,18 @@ +# Guest-only chatbot for beta.reactome.org. Copy to .env.beta and fill in. +# Deliberately omitted (do not copy from prod's .env): +# OAUTH_GOOGLE_* redirect URIs are registered for reactome.org +# CHAINLIT_AUTH_SECRET enables the login flow, which has nowhere to go here +# CLOUDFLARE_* Turnstile site key is bound to reactome.org; unset = +# captcha middleware bypasses itself +# POSTGRES_* unset = MemorySaver, no chat history, no DB to run + +OPENAI_API_KEY= +TAVILY_API_KEY= + +CHAT_ENV=reactome +LOG_LEVEL=info +UVICORN_LOG_LEVEL=info + +# Chainlit is mounted here; the landing page lives at /chat/ +CHAINLIT_URI=/chat/guest +CHAINLIT_URL=https://beta.reactome.org diff --git a/deploy/beta/reclaim-docker-space.sh b/deploy/beta/reclaim-docker-space.sh new file mode 100755 index 0000000..1dbe44f --- /dev/null +++ b/deploy/beta/reclaim-docker-space.sh @@ -0,0 +1,89 @@ +#!/usr/bin/env bash +# +# Frees disk on dev.reactome.org so the chatbot image + embeddings bundle fit. +# +# sudo ~/reclaim-docker-space.sh # audit only, changes NOTHING +# sudo ~/reclaim-docker-space.sh --safe # dangling images + build cache +# sudo ~/reclaim-docker-space.sh --all # the above, plus dangling volumes +# +# This host serves both dev.reactome.org and the origin behind beta.reactome.org +# on one 88G volume. Filling it takes both down, which is why the default mode +# changes nothing and why --all is a separate, deliberate flag. +# +# --safe touches only things Docker can rebuild: layers no image references any +# more, and build cache. --all additionally removes ANONYMOUS volumes that no +# container references. Read the audit before using it; a removed volume is gone. + +set -euo pipefail + +MODE="${1:-audit}" + +say() { printf '\n\033[1m%s\033[0m\n' "$*"; } +ok() { printf ' \033[32m✓\033[0m %s\n' "$*"; } +warn() { printf ' \033[33m!\033[0m %s\n' "$*"; } + +free_gb() { df -BG --output=avail / | tail -1 | tr -dc '0-9'; } + +BEFORE=$(free_gb) +say "Disk before: ${BEFORE}G free" +df -h / | tail -1 + +say "What is reclaimable" +docker system df + +say "Dangling volumes (candidates for --all)" +found=0 +DANGLING=() +for v in $(docker volume ls -qf dangling=true); do + found=1 + DANGLING+=("$v") + mp=$(docker volume inspect -f '{{.Mountpoint}}' "$v") + created=$(docker volume inspect -f '{{.CreatedAt}}' "$v") + size=$(du -sh "$mp" 2>/dev/null | cut -f1 || echo '?') + printf '\n %s\n created %s, %s\n' "$v" "$created" "$size" + printf ' top-level contents:\n' + ls -A "$mp" 2>/dev/null | head -8 | sed 's/^/ /' || true + n=$(ls -A "$mp" 2>/dev/null | wc -l) + [ "$n" -gt 8 ] && printf ' ... and %s more entries\n' "$((n - 8))" +done +[ "$found" -eq 0 ] && ok "none" + +if [ "$MODE" = "audit" ]; then + say "Audit only -- nothing was changed." + echo " Re-run with --safe (images + build cache) or --all (also the volumes above)." + exit 0 +fi + +say "Reclaiming: dangling images" +docker image prune -f +say "Reclaiming: build cache" +docker builder prune -f + +if [ "$MODE" = "--all" ]; then + say "Reclaiming: dangling volumes" + if [ ${#DANGLING[@]} -gt 0 ]; then + # Remove by id rather than `docker volume prune`, which only sweeps volumes + # Docker tagged as anonymous and silently leaves older unreferenced ones. + docker volume rm "${DANGLING[@]}" || warn "some volumes could not be removed" + else + ok "no dangling volumes" + fi +elif [ "$MODE" != "--safe" ]; then + warn "unknown mode '$MODE' -- treated as --safe" +fi + +AFTER=$(free_gb) +say "Disk after: ${AFTER}G free (was ${BEFORE}G)" +df -h / | tail -1 + +# Peak requirement: ~3.2G image + 1.3G zip + ~3.0G extracted bundle. +if [ "$AFTER" -lt 9 ]; then + warn "under 9G free; the pull + bundle extract needs roughly 7.5G at peak." + if [ "$MODE" = "--all" ]; then + warn "already ran --all; the remaining space has to come from outside Docker." + else + warn "try --all, or free space outside Docker." + fi +else + ok "enough headroom for the image and the Release95 bundle" +fi diff --git a/docker-compose.yml b/docker-compose.yml index bf353ad..d684451 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -25,8 +25,10 @@ services: - CHAINLIT_URL=${CHAINLIT_URL} - CHAINLIT_ROOT_PATH=${CHAINLIT_ROOT_PATH} - TAVILY_API_KEY=${TAVILY_API_KEY} + - OAUTH_ORCID_CLIENT_ID=${OAUTH_ORCID_CLIENT_ID} + - OAUTH_ORCID_CLIENT_SECRET=${OAUTH_ORCID_CLIENT_SECRET} ports: - - "8000:8000" + - "8002:8000" depends_on: postgres: condition: service_healthy @@ -70,7 +72,7 @@ services: - POSTGRES_USER=${POSTGRES_USER} - POSTGRES_PASSWORD=${POSTGRES_PASSWORD} ports: - - "5432:5432" + - "127.0.0.1:5432:5432" volumes: - ./initdb:/docker-entrypoint-initdb.d - ./data:/var/lib/postgresql/data diff --git a/docs/embeddings_manager.md b/docs/embeddings_manager.md index 20b84da..d36d248 100644 --- a/docs/embeddings_manager.md +++ b/docs/embeddings_manager.md @@ -93,6 +93,16 @@ Either specify `--hf-key` or environment variable `HUGGINGFACEHUB_API_TOKEN`. ./bin/embeddings_manager make /reactome/ --hf-key ``` +### User Guide: + +Fetches Reactome website user guide pages, chunks them by section, and embeds them into Chroma. Version is date-based (not tied to graph DB releases). + +```sh +./bin/embeddings_manager make openai/text-embedding-3-large/userguide/ --openai-key +``` + +Use `--force` to re-fetch HTML from reactome.org and rebuild the `sections/` Chroma collection from scratch. + ## Uploading to S3: `push` ⚠️ Requires S3 write access. diff --git a/env_template b/env_template index 4b3e13b..48f0ea2 100644 --- a/env_template +++ b/env_template @@ -6,7 +6,6 @@ POSTGRES_CHAINLIT_DB=chatbio_chainlit POSTGRES_LANGGRAPH_DB=chatbio_langgraph PGADMIN_DEFAULT_EMAIL=test@test.com PGADMIN_DEFAULT_PASSWORD=test -PYTHON_PATH=/app/src:/app: CHAT_ENV=reactome CLOUDFLARE_SECRET_KEY= CLOUDFLARE_SITE_KEY=0x4AAAAAAAkzFQ5GRs2toYuv diff --git a/export_csvs.py b/export_csvs.py new file mode 100644 index 0000000..b50a4fa --- /dev/null +++ b/export_csvs.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +import csv +import os +from pathlib import Path + +from neo4j import Driver, GraphDatabase + +NEO4J_URI = os.environ.get("NEO4J_URI", "bolt://localhost:7687") +NEO4J_USER = os.environ.get("NEO4J_USERNAME", "neo4j") +NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD") +if not NEO4J_PASSWORD: + # Without this the driver constructs fine and fails later at connect time + # with an error that does not mention the missing variable. + raise SystemExit( + "NEO4J_PASSWORD is not set. Export it before running this script; it is " + "no longer hardcoded here." + ) + +OUTPUT_DIR = Path("./embeddings/openai/bge-m3/plantreactome/Release68/csv_files") +OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + +QUERIES = { + "reactions": """ + MATCH (pathway:Pathway)-[:hasEvent]->(reaction:ReactionLikeEvent) + OPTIONAL MATCH (reaction)-[:input]->(input:PhysicalEntity) + OPTIONAL MATCH (reaction)-[:output]->(output:PhysicalEntity) + OPTIONAL MATCH (reaction)-[:catalystActivity]->(cat:CatalystActivity)-[:physicalEntity]->(catalyst:PhysicalEntity) + RETURN reaction.stId AS st_id, reaction.displayName AS display_name, + pathway.stId AS pathway_id, pathway.displayName AS pathway_name, + pathway.speciesName AS species, + COLLECT(DISTINCT input.stId) AS input_id, + COLLECT(DISTINCT input.displayName) AS input_name, + COLLECT(DISTINCT output.stId) AS output_id, + COLLECT(DISTINCT output.displayName) AS output_name, + COLLECT(DISTINCT catalyst.stId) AS catalyst_id, + COLLECT(DISTINCT catalyst.displayName) AS catalyst_name, + "https://plantreactome.gramene.org/content/detail/" + reaction.stId AS url + """, + "summations": """ + MATCH (e)-[:summation]->(s:Summation) + WHERE (e:Pathway OR e:ReactionLikeEvent) + RETURN e.stId AS st_id, e.displayName AS display_name, labels(e) AS labels, + e.speciesName AS species, + CASE WHEN size(s.text) > 10000 THEN LEFT(s.text, 10000) + '...' ELSE s.text END AS summation, + "https://plantreactome.gramene.org/content/detail/" + e.stId AS url + """, + "complexes": """ + MATCH (complex:Complex)-[:hasComponent]->(component) + RETURN complex.speciesName AS species, complex.stId AS st_id, + complex.name AS display_name, component.stId AS component_id, component.name AS component_name, + "https://plantreactome.gramene.org/content/detail/" + complex.stId AS url + """, + "ewas": """ + MATCH (db:ReferenceDatabase)<-[:referenceDatabase]-(gene:ReferenceEntity)<-[:referenceEntity]-(prot:PhysicalEntity) + RETURN DISTINCT + prot.stId AS st_id, + prot.displayName AS display_name, + gene.geneName AS canonical_gene_name, + '' AS synonyms_gene_name, + gene.url AS uniprot_link, + "https://plantreactome.gramene.org/content/detail/" + prot.stId AS url + """, +} + + +def clean_value(v: object) -> str: + if v is None: + return "" + if isinstance(v, list): + return "|".join(str(x) for x in v) + s = str(v).strip() + if s.startswith("[") and s.endswith("]"): + inner = s[1:-1] + items = [i.strip().strip('"').strip("'") for i in inner.split(",") if i.strip()] + return "|".join(items) + return s + + +def run_query(driver: Driver, query: str) -> list[dict[str, str]]: + with driver.session() as session: + result = session.run(query) + records = [r.data() for r in result] + cleaned = [] + for row in records: + new_row = {} + for k, v in row.items(): + if k is None: + continue + new_row[k.strip()] = clean_value(v) + cleaned.append(new_row) + return cleaned + + +def main() -> None: + driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD)) + for name, query in QUERIES.items(): + print(f"Exporting {name}...") + rows = run_query(driver, query) + if not rows: + print(f" WARNING: No rows returned for {name}") + continue + fieldnames = list(rows[0].keys()) + outfile = OUTPUT_DIR / f"{name}.csv" + with open(outfile, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + print(f" Saved {len(rows)} rows to {outfile}") + driver.close() + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/mypy.ini b/mypy.ini deleted file mode 100644 index 7ec92a3..0000000 --- a/mypy.ini +++ /dev/null @@ -1,11 +0,0 @@ -[mypy] -ignore_missing_imports = True -allow_untyped_calls = True -allow_untyped_defs = True -allow_untyped_globals = True -explicit_package_bases = True -exclude = data/ -files = bin/,src/ - -[mypy.plugins.pandas.*] -init_forbid_dynamic = False diff --git a/poetry.lock b/poetry.lock index e21a020..f2b64d1 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. [[package]] name = "aiofiles" @@ -335,6 +335,28 @@ files = [ tests = ["pytest (>=3.2.1,!=3.3.0)"] typecheck = ["mypy"] +[[package]] +name = "beautifulsoup4" +version = "4.15.0" +description = "Screen-scraping library" +optional = false +python-versions = ">=3.7.0" +files = [ + {file = "beautifulsoup4-4.15.0-py3-none-any.whl", hash = "sha256:d6f88de62e1d4e38ecb1077eb9724cd0eff29d2a08ca16a401e9b9e93f117cf9"}, + {file = "beautifulsoup4-4.15.0.tar.gz", hash = "sha256:288e3ca7d54b06f2ac191970bc275c1939cb46d450b255bf6718b04aa37ab4f7"}, +] + +[package.dependencies] +soupsieve = ">=1.6.1" +typing-extensions = ">=4.0.0" + +[package.extras] +cchardet = ["cchardet"] +chardet = ["chardet"] +charset-normalizer = ["charset-normalizer"] +html5lib = ["html5lib"] +lxml = ["lxml"] + [[package]] name = "bidict" version = "0.23.1" @@ -1858,23 +1880,24 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<10" [[package]] name = "langchain-core" -version = "0.3.63" +version = "0.3.84" description = "Building applications with LLMs through composability" optional = false -python-versions = ">=3.9" +python-versions = "<4.0.0,>=3.9.0" files = [ - {file = "langchain_core-0.3.63-py3-none-any.whl", hash = "sha256:f91db8221b1bc6808f70b2e72fded1a94d50ee3f1dff1636fb5a5a514c64b7f5"}, - {file = "langchain_core-0.3.63.tar.gz", hash = "sha256:e2e30cfbb7684a5a0319f6cbf065fc3c438bfd1060302f085a122527890fb01e"}, + {file = "langchain_core-0.3.84-py3-none-any.whl", hash = "sha256:d0b3a7b6473e30a2b3d4588ee09dc6471b8d38c46cd48f3e7c3d1ab6547f63cb"}, + {file = "langchain_core-0.3.84.tar.gz", hash = "sha256:814b75bfe67a8460a53f5839bae9505bbfffc7af6f1aa0a5155715563f5cc490"}, ] [package.dependencies] -jsonpatch = ">=1.33,<2.0" -langsmith = ">=0.1.126,<0.4" -packaging = ">=23.2,<25" -pydantic = ">=2.7.4" -PyYAML = ">=5.3" +jsonpatch = ">=1.33.0,<2.0.0" +langsmith = ">=0.3.45,<1.0.0" +packaging = ">=23.2.0,<26.0.0" +pydantic = ">=2.7.4,<3.0.0" +PyYAML = ">=5.3.0,<7.0.0" tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<10.0.0" -typing-extensions = ">=4.7" +typing-extensions = ">=4.7.0,<5.0.0" +uuid-utils = ">=0.12.0,<1.0" [[package]] name = "langchain-huggingface" @@ -2004,18 +2027,19 @@ orjson = ">=3.10.1" [[package]] name = "langsmith" -version = "0.3.2" +version = "0.3.45" description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform." optional = false -python-versions = "<4.0,>=3.9" +python-versions = ">=3.9" files = [ - {file = "langsmith-0.3.2-py3-none-any.whl", hash = "sha256:48ff6bc5eda62f4729596bb68d4f96166d2654728ac32970b69b1be874c61925"}, - {file = "langsmith-0.3.2.tar.gz", hash = "sha256:7724668e9705734ab25a7977fc34a9ee15a40ba4108987926c69293a05d40229"}, + {file = "langsmith-0.3.45-py3-none-any.whl", hash = "sha256:5b55f0518601fa65f3bb6b1a3100379a96aa7b3ed5e9380581615ba9c65ed8ed"}, + {file = "langsmith-0.3.45.tar.gz", hash = "sha256:1df3c6820c73ed210b2c7bc5cdb7bfa19ddc9126cd03fdf0da54e2e171e6094d"}, ] [package.dependencies] httpx = ">=0.23.0,<1" orjson = {version = ">=3.9.14,<4.0.0", markers = "platform_python_implementation != \"PyPy\""} +packaging = ">=23.2" pydantic = [ {version = ">=1,<3", markers = "python_full_version < \"3.12.4\""}, {version = ">=2.7.4,<3.0.0", markers = "python_full_version >= \"3.12.4\""}, @@ -2026,6 +2050,8 @@ zstandard = ">=0.23.0,<0.24.0" [package.extras] langsmith-pyo3 = ["langsmith-pyo3 (>=0.1.0rc2,<0.2.0)"] +openai-agents = ["openai-agents (>=0.0.3,<0.1)"] +otel = ["opentelemetry-api (>=1.30.0,<2.0.0)", "opentelemetry-exporter-otlp-proto-http (>=1.30.0,<2.0.0)", "opentelemetry-sdk (>=1.30.0,<2.0.0)"] pytest = ["pytest (>=7.0.0)", "rich (>=13.9.4,<14.0.0)"] [[package]] @@ -2072,6 +2098,154 @@ httpx = ">=0.23.0" packaging = ">=23.0" pydantic = ">=1,<3" +[[package]] +name = "lxml" +version = "5.4.0" +description = "Powerful and Pythonic XML processing library combining libxml2/libxslt with the ElementTree API." +optional = false +python-versions = ">=3.6" +files = [ + {file = "lxml-5.4.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:e7bc6df34d42322c5289e37e9971d6ed114e3776b45fa879f734bded9d1fea9c"}, + {file = "lxml-5.4.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:6854f8bd8a1536f8a1d9a3655e6354faa6406621cf857dc27b681b69860645c7"}, + {file = "lxml-5.4.0-cp310-cp310-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:696ea9e87442467819ac22394ca36cb3d01848dad1be6fac3fb612d3bd5a12cf"}, + {file = "lxml-5.4.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ef80aeac414f33c24b3815ecd560cee272786c3adfa5f31316d8b349bfade28"}, + {file = "lxml-5.4.0-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3b9c2754cef6963f3408ab381ea55f47dabc6f78f4b8ebb0f0b25cf1ac1f7609"}, + {file = "lxml-5.4.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:7a62cc23d754bb449d63ff35334acc9f5c02e6dae830d78dab4dd12b78a524f4"}, + {file = "lxml-5.4.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f82125bc7203c5ae8633a7d5d20bcfdff0ba33e436e4ab0abc026a53a8960b7"}, + {file = "lxml-5.4.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:b67319b4aef1a6c56576ff544b67a2a6fbd7eaee485b241cabf53115e8908b8f"}, + {file = "lxml-5.4.0-cp310-cp310-manylinux_2_28_ppc64le.whl", hash = "sha256:a8ef956fce64c8551221f395ba21d0724fed6b9b6242ca4f2f7beb4ce2f41997"}, + {file = "lxml-5.4.0-cp310-cp310-manylinux_2_28_s390x.whl", hash = "sha256:0a01ce7d8479dce84fc03324e3b0c9c90b1ece9a9bb6a1b6c9025e7e4520e78c"}, + {file = "lxml-5.4.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:91505d3ddebf268bb1588eb0f63821f738d20e1e7f05d3c647a5ca900288760b"}, + {file = "lxml-5.4.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:a3bcdde35d82ff385f4ede021df801b5c4a5bcdfb61ea87caabcebfc4945dc1b"}, + {file = "lxml-5.4.0-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:aea7c06667b987787c7d1f5e1dfcd70419b711cdb47d6b4bb4ad4b76777a0563"}, + {file = "lxml-5.4.0-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:a7fb111eef4d05909b82152721a59c1b14d0f365e2be4c742a473c5d7372f4f5"}, + {file = "lxml-5.4.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:43d549b876ce64aa18b2328faff70f5877f8c6dede415f80a2f799d31644d776"}, + {file = "lxml-5.4.0-cp310-cp310-win32.whl", hash = "sha256:75133890e40d229d6c5837b0312abbe5bac1c342452cf0e12523477cd3aa21e7"}, + {file = "lxml-5.4.0-cp310-cp310-win_amd64.whl", hash = "sha256:de5b4e1088523e2b6f730d0509a9a813355b7f5659d70eb4f319c76beea2e250"}, + {file = "lxml-5.4.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:98a3912194c079ef37e716ed228ae0dcb960992100461b704aea4e93af6b0bb9"}, + {file = "lxml-5.4.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0ea0252b51d296a75f6118ed0d8696888e7403408ad42345d7dfd0d1e93309a7"}, + {file = "lxml-5.4.0-cp311-cp311-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b92b69441d1bd39f4940f9eadfa417a25862242ca2c396b406f9272ef09cdcaa"}, + {file = "lxml-5.4.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:20e16c08254b9b6466526bc1828d9370ee6c0d60a4b64836bc3ac2917d1e16df"}, + {file = "lxml-5.4.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7605c1c32c3d6e8c990dd28a0970a3cbbf1429d5b92279e37fda05fb0c92190e"}, + {file = "lxml-5.4.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ecf4c4b83f1ab3d5a7ace10bafcb6f11df6156857a3c418244cef41ca9fa3e44"}, + {file = "lxml-5.4.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0cef4feae82709eed352cd7e97ae062ef6ae9c7b5dbe3663f104cd2c0e8d94ba"}, + {file = "lxml-5.4.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:df53330a3bff250f10472ce96a9af28628ff1f4efc51ccba351a8820bca2a8ba"}, + {file = "lxml-5.4.0-cp311-cp311-manylinux_2_28_ppc64le.whl", hash = "sha256:aefe1a7cb852fa61150fcb21a8c8fcea7b58c4cb11fbe59c97a0a4b31cae3c8c"}, + {file = "lxml-5.4.0-cp311-cp311-manylinux_2_28_s390x.whl", hash = "sha256:ef5a7178fcc73b7d8c07229e89f8eb45b2908a9238eb90dcfc46571ccf0383b8"}, + {file = "lxml-5.4.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:d2ed1b3cb9ff1c10e6e8b00941bb2e5bb568b307bfc6b17dffbbe8be5eecba86"}, + {file = "lxml-5.4.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:72ac9762a9f8ce74c9eed4a4e74306f2f18613a6b71fa065495a67ac227b3056"}, + {file = "lxml-5.4.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:f5cb182f6396706dc6cc1896dd02b1c889d644c081b0cdec38747573db88a7d7"}, + {file = "lxml-5.4.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:3a3178b4873df8ef9457a4875703488eb1622632a9cee6d76464b60e90adbfcd"}, + {file = "lxml-5.4.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:e094ec83694b59d263802ed03a8384594fcce477ce484b0cbcd0008a211ca751"}, + {file = "lxml-5.4.0-cp311-cp311-win32.whl", hash = "sha256:4329422de653cdb2b72afa39b0aa04252fca9071550044904b2e7036d9d97fe4"}, + {file = "lxml-5.4.0-cp311-cp311-win_amd64.whl", hash = "sha256:fd3be6481ef54b8cfd0e1e953323b7aa9d9789b94842d0e5b142ef4bb7999539"}, + {file = "lxml-5.4.0-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:b5aff6f3e818e6bdbbb38e5967520f174b18f539c2b9de867b1e7fde6f8d95a4"}, + {file = "lxml-5.4.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:942a5d73f739ad7c452bf739a62a0f83e2578afd6b8e5406308731f4ce78b16d"}, + {file = "lxml-5.4.0-cp312-cp312-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:460508a4b07364d6abf53acaa0a90b6d370fafde5693ef37602566613a9b0779"}, + {file = "lxml-5.4.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:529024ab3a505fed78fe3cc5ddc079464e709f6c892733e3f5842007cec8ac6e"}, + {file = "lxml-5.4.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7ca56ebc2c474e8f3d5761debfd9283b8b18c76c4fc0967b74aeafba1f5647f9"}, + {file = "lxml-5.4.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a81e1196f0a5b4167a8dafe3a66aa67c4addac1b22dc47947abd5d5c7a3f24b5"}, + {file = "lxml-5.4.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00b8686694423ddae324cf614e1b9659c2edb754de617703c3d29ff568448df5"}, + {file = "lxml-5.4.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:c5681160758d3f6ac5b4fea370495c48aac0989d6a0f01bb9a72ad8ef5ab75c4"}, + {file = "lxml-5.4.0-cp312-cp312-manylinux_2_28_ppc64le.whl", hash = "sha256:2dc191e60425ad70e75a68c9fd90ab284df64d9cd410ba8d2b641c0c45bc006e"}, + {file = "lxml-5.4.0-cp312-cp312-manylinux_2_28_s390x.whl", hash = "sha256:67f779374c6b9753ae0a0195a892a1c234ce8416e4448fe1e9f34746482070a7"}, + {file = "lxml-5.4.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:79d5bfa9c1b455336f52343130b2067164040604e41f6dc4d8313867ed540079"}, + {file = "lxml-5.4.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3d3c30ba1c9b48c68489dc1829a6eede9873f52edca1dda900066542528d6b20"}, + {file = "lxml-5.4.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:1af80c6316ae68aded77e91cd9d80648f7dd40406cef73df841aa3c36f6907c8"}, + {file = "lxml-5.4.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:4d885698f5019abe0de3d352caf9466d5de2baded00a06ef3f1216c1a58ae78f"}, + {file = "lxml-5.4.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:aea53d51859b6c64e7c51d522c03cc2c48b9b5d6172126854cc7f01aa11f52bc"}, + {file = "lxml-5.4.0-cp312-cp312-win32.whl", hash = "sha256:d90b729fd2732df28130c064aac9bb8aff14ba20baa4aee7bd0795ff1187545f"}, + {file = "lxml-5.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:1dc4ca99e89c335a7ed47d38964abcb36c5910790f9bd106f2a8fa2ee0b909d2"}, + {file = "lxml-5.4.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:773e27b62920199c6197130632c18fb7ead3257fce1ffb7d286912e56ddb79e0"}, + {file = "lxml-5.4.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:ce9c671845de9699904b1e9df95acfe8dfc183f2310f163cdaa91a3535af95de"}, + {file = "lxml-5.4.0-cp313-cp313-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9454b8d8200ec99a224df8854786262b1bd6461f4280064c807303c642c05e76"}, + {file = "lxml-5.4.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cccd007d5c95279e529c146d095f1d39ac05139de26c098166c4beb9374b0f4d"}, + {file = "lxml-5.4.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0fce1294a0497edb034cb416ad3e77ecc89b313cff7adbee5334e4dc0d11f422"}, + {file = "lxml-5.4.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:24974f774f3a78ac12b95e3a20ef0931795ff04dbb16db81a90c37f589819551"}, + {file = "lxml-5.4.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:497cab4d8254c2a90bf988f162ace2ddbfdd806fce3bda3f581b9d24c852e03c"}, + {file = "lxml-5.4.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:e794f698ae4c5084414efea0f5cc9f4ac562ec02d66e1484ff822ef97c2cadff"}, + {file = "lxml-5.4.0-cp313-cp313-manylinux_2_28_ppc64le.whl", hash = "sha256:2c62891b1ea3094bb12097822b3d44b93fc6c325f2043c4d2736a8ff09e65f60"}, + {file = "lxml-5.4.0-cp313-cp313-manylinux_2_28_s390x.whl", hash = "sha256:142accb3e4d1edae4b392bd165a9abdee8a3c432a2cca193df995bc3886249c8"}, + {file = "lxml-5.4.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:1a42b3a19346e5601d1b8296ff6ef3d76038058f311902edd574461e9c036982"}, + {file = "lxml-5.4.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4291d3c409a17febf817259cb37bc62cb7eb398bcc95c1356947e2871911ae61"}, + {file = "lxml-5.4.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:4f5322cf38fe0e21c2d73901abf68e6329dc02a4994e483adbcf92b568a09a54"}, + {file = "lxml-5.4.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:0be91891bdb06ebe65122aa6bf3fc94489960cf7e03033c6f83a90863b23c58b"}, + {file = "lxml-5.4.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:15a665ad90054a3d4f397bc40f73948d48e36e4c09f9bcffc7d90c87410e478a"}, + {file = "lxml-5.4.0-cp313-cp313-win32.whl", hash = "sha256:d5663bc1b471c79f5c833cffbc9b87d7bf13f87e055a5c86c363ccd2348d7e82"}, + {file = "lxml-5.4.0-cp313-cp313-win_amd64.whl", hash = "sha256:bcb7a1096b4b6b24ce1ac24d4942ad98f983cd3810f9711bcd0293f43a9d8b9f"}, + {file = "lxml-5.4.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:7be701c24e7f843e6788353c055d806e8bd8466b52907bafe5d13ec6a6dbaecd"}, + {file = "lxml-5.4.0-cp36-cp36m-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:fb54f7c6bafaa808f27166569b1511fc42701a7713858dddc08afdde9746849e"}, + {file = "lxml-5.4.0-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:97dac543661e84a284502e0cf8a67b5c711b0ad5fb661d1bd505c02f8cf716d7"}, + {file = "lxml-5.4.0-cp36-cp36m-manylinux_2_28_x86_64.whl", hash = "sha256:c70e93fba207106cb16bf852e421c37bbded92acd5964390aad07cb50d60f5cf"}, + {file = "lxml-5.4.0-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:9c886b481aefdf818ad44846145f6eaf373a20d200b5ce1a5c8e1bc2d8745410"}, + {file = "lxml-5.4.0-cp36-cp36m-musllinux_1_2_x86_64.whl", hash = "sha256:fa0e294046de09acd6146be0ed6727d1f42ded4ce3ea1e9a19c11b6774eea27c"}, + {file = "lxml-5.4.0-cp36-cp36m-win32.whl", hash = "sha256:61c7bbf432f09ee44b1ccaa24896d21075e533cd01477966a5ff5a71d88b2f56"}, + {file = "lxml-5.4.0-cp36-cp36m-win_amd64.whl", hash = "sha256:7ce1a171ec325192c6a636b64c94418e71a1964f56d002cc28122fceff0b6121"}, + {file = "lxml-5.4.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:795f61bcaf8770e1b37eec24edf9771b307df3af74d1d6f27d812e15a9ff3872"}, + {file = "lxml-5.4.0-cp37-cp37m-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:29f451a4b614a7b5b6c2e043d7b64a15bd8304d7e767055e8ab68387a8cacf4e"}, + {file = "lxml-5.4.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:891f7f991a68d20c75cb13c5c9142b2a3f9eb161f1f12a9489c82172d1f133c0"}, + {file = "lxml-5.4.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4aa412a82e460571fad592d0f93ce9935a20090029ba08eca05c614f99b0cc92"}, + {file = "lxml-5.4.0-cp37-cp37m-manylinux_2_28_aarch64.whl", hash = "sha256:ac7ba71f9561cd7d7b55e1ea5511543c0282e2b6450f122672a2694621d63b7e"}, + {file = "lxml-5.4.0-cp37-cp37m-manylinux_2_28_x86_64.whl", hash = "sha256:c5d32f5284012deaccd37da1e2cd42f081feaa76981f0eaa474351b68df813c5"}, + {file = "lxml-5.4.0-cp37-cp37m-musllinux_1_2_aarch64.whl", hash = "sha256:ce31158630a6ac85bddd6b830cffd46085ff90498b397bd0a259f59d27a12188"}, + {file = "lxml-5.4.0-cp37-cp37m-musllinux_1_2_x86_64.whl", hash = "sha256:31e63621e073e04697c1b2d23fcb89991790eef370ec37ce4d5d469f40924ed6"}, + {file = "lxml-5.4.0-cp37-cp37m-win32.whl", hash = "sha256:be2ba4c3c5b7900246a8f866580700ef0d538f2ca32535e991027bdaba944063"}, + {file = "lxml-5.4.0-cp37-cp37m-win_amd64.whl", hash = "sha256:09846782b1ef650b321484ad429217f5154da4d6e786636c38e434fa32e94e49"}, + {file = "lxml-5.4.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:eaf24066ad0b30917186420d51e2e3edf4b0e2ea68d8cd885b14dc8afdcf6556"}, + {file = "lxml-5.4.0-cp38-cp38-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2b31a3a77501d86d8ade128abb01082724c0dfd9524f542f2f07d693c9f1175f"}, + {file = "lxml-5.4.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0e108352e203c7afd0eb91d782582f00a0b16a948d204d4dec8565024fafeea5"}, + {file = "lxml-5.4.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a11a96c3b3f7551c8a8109aa65e8594e551d5a84c76bf950da33d0fb6dfafab7"}, + {file = "lxml-5.4.0-cp38-cp38-manylinux_2_28_aarch64.whl", hash = "sha256:ca755eebf0d9e62d6cb013f1261e510317a41bf4650f22963474a663fdfe02aa"}, + {file = "lxml-5.4.0-cp38-cp38-manylinux_2_28_x86_64.whl", hash = "sha256:4cd915c0fb1bed47b5e6d6edd424ac25856252f09120e3e8ba5154b6b921860e"}, + {file = "lxml-5.4.0-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:226046e386556a45ebc787871d6d2467b32c37ce76c2680f5c608e25823ffc84"}, + {file = "lxml-5.4.0-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:b108134b9667bcd71236c5a02aad5ddd073e372fb5d48ea74853e009fe38acb6"}, + {file = "lxml-5.4.0-cp38-cp38-win32.whl", hash = "sha256:1320091caa89805df7dcb9e908add28166113dcd062590668514dbd510798c88"}, + {file = "lxml-5.4.0-cp38-cp38-win_amd64.whl", hash = "sha256:073eb6dcdf1f587d9b88c8c93528b57eccda40209cf9be549d469b942b41d70b"}, + {file = "lxml-5.4.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:bda3ea44c39eb74e2488297bb39d47186ed01342f0022c8ff407c250ac3f498e"}, + {file = "lxml-5.4.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9ceaf423b50ecfc23ca00b7f50b64baba85fb3fb91c53e2c9d00bc86150c7e40"}, + {file = "lxml-5.4.0-cp39-cp39-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:664cdc733bc87449fe781dbb1f309090966c11cc0c0cd7b84af956a02a8a4729"}, + {file = "lxml-5.4.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:67ed8a40665b84d161bae3181aa2763beea3747f748bca5874b4af4d75998f87"}, + {file = "lxml-5.4.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9b4a3bd174cc9cdaa1afbc4620c049038b441d6ba07629d89a83b408e54c35cd"}, + {file = "lxml-5.4.0-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:b0989737a3ba6cf2a16efb857fb0dfa20bc5c542737fddb6d893fde48be45433"}, + {file = "lxml-5.4.0-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:dc0af80267edc68adf85f2a5d9be1cdf062f973db6790c1d065e45025fa26140"}, + {file = "lxml-5.4.0-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:639978bccb04c42677db43c79bdaa23785dc7f9b83bfd87570da8207872f1ce5"}, + {file = "lxml-5.4.0-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:5a99d86351f9c15e4a901fc56404b485b1462039db59288b203f8c629260a142"}, + {file = "lxml-5.4.0-cp39-cp39-win32.whl", hash = "sha256:3e6d5557989cdc3ebb5302bbdc42b439733a841891762ded9514e74f60319ad6"}, + {file = "lxml-5.4.0-cp39-cp39-win_amd64.whl", hash = "sha256:a8c9b7f16b63e65bbba889acb436a1034a82d34fa09752d754f88d708eca80e1"}, + {file = "lxml-5.4.0-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:1b717b00a71b901b4667226bba282dd462c42ccf618ade12f9ba3674e1fabc55"}, + {file = "lxml-5.4.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:27a9ded0f0b52098ff89dd4c418325b987feed2ea5cc86e8860b0f844285d740"}, + {file = "lxml-5.4.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4b7ce10634113651d6f383aa712a194179dcd496bd8c41e191cec2099fa09de5"}, + {file = "lxml-5.4.0-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:53370c26500d22b45182f98847243efb518d268374a9570409d2e2276232fd37"}, + {file = "lxml-5.4.0-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:c6364038c519dffdbe07e3cf42e6a7f8b90c275d4d1617a69bb59734c1a2d571"}, + {file = "lxml-5.4.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:b12cb6527599808ada9eb2cd6e0e7d3d8f13fe7bbb01c6311255a15ded4c7ab4"}, + {file = "lxml-5.4.0-pp37-pypy37_pp73-macosx_10_9_x86_64.whl", hash = "sha256:5f11a1526ebd0dee85e7b1e39e39a0cc0d9d03fb527f56d8457f6df48a10dc0c"}, + {file = "lxml-5.4.0-pp37-pypy37_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:48b4afaf38bf79109bb060d9016fad014a9a48fb244e11b94f74ae366a64d252"}, + {file = "lxml-5.4.0-pp37-pypy37_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:de6f6bb8a7840c7bf216fb83eec4e2f79f7325eca8858167b68708b929ab2172"}, + {file = "lxml-5.4.0-pp37-pypy37_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:5cca36a194a4eb4e2ed6be36923d3cffd03dcdf477515dea687185506583d4c9"}, + {file = "lxml-5.4.0-pp37-pypy37_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:b7c86884ad23d61b025989d99bfdd92a7351de956e01c61307cb87035960bcb1"}, + {file = "lxml-5.4.0-pp37-pypy37_pp73-win_amd64.whl", hash = "sha256:53d9469ab5460402c19553b56c3648746774ecd0681b1b27ea74d5d8a3ef5590"}, + {file = "lxml-5.4.0-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:56dbdbab0551532bb26c19c914848d7251d73edb507c3079d6805fa8bba5b706"}, + {file = "lxml-5.4.0-pp38-pypy38_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:14479c2ad1cb08b62bb941ba8e0e05938524ee3c3114644df905d2331c76cd57"}, + {file = "lxml-5.4.0-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:32697d2ea994e0db19c1df9e40275ffe84973e4232b5c274f47e7c1ec9763cdd"}, + {file = "lxml-5.4.0-pp38-pypy38_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:24f6df5f24fc3385f622c0c9d63fe34604893bc1a5bdbb2dbf5870f85f9a404a"}, + {file = "lxml-5.4.0-pp38-pypy38_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:151d6c40bc9db11e960619d2bf2ec5829f0aaffb10b41dcf6ad2ce0f3c0b2325"}, + {file = "lxml-5.4.0-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:4025bf2884ac4370a3243c5aa8d66d3cb9e15d3ddd0af2d796eccc5f0244390e"}, + {file = "lxml-5.4.0-pp39-pypy39_pp73-macosx_10_15_x86_64.whl", hash = "sha256:9459e6892f59ecea2e2584ee1058f5d8f629446eab52ba2305ae13a32a059530"}, + {file = "lxml-5.4.0-pp39-pypy39_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:47fb24cc0f052f0576ea382872b3fc7e1f7e3028e53299ea751839418ade92a6"}, + {file = "lxml-5.4.0-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:50441c9de951a153c698b9b99992e806b71c1f36d14b154592580ff4a9d0d877"}, + {file = "lxml-5.4.0-pp39-pypy39_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:ab339536aa798b1e17750733663d272038bf28069761d5be57cb4a9b0137b4f8"}, + {file = "lxml-5.4.0-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:9776af1aad5a4b4a1317242ee2bea51da54b2a7b7b48674be736d463c999f37d"}, + {file = "lxml-5.4.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:63e7968ff83da2eb6fdda967483a7a023aa497d85ad8f05c3ad9b1f2e8c84987"}, + {file = "lxml-5.4.0.tar.gz", hash = "sha256:d12832e1dbea4be280b22fd0ea7c9b87f0d8fc51ba06e92dc62d52f804f78ebd"}, +] + +[package.extras] +cssselect = ["cssselect (>=0.7)"] +html-clean = ["lxml_html_clean"] +html5 = ["html5lib"] +htmlsoup = ["BeautifulSoup4"] +source = ["Cython (>=3.0.11,<3.1.0)"] + [[package]] name = "markdown-it-py" version = "3.0.0" @@ -4769,6 +4943,17 @@ files = [ {file = "sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc"}, ] +[[package]] +name = "soupsieve" +version = "2.8.4" +description = "A modern CSS selector implementation for Beautiful Soup." +optional = false +python-versions = ">=3.9" +files = [ + {file = "soupsieve-2.8.4-py3-none-any.whl", hash = "sha256:e7e6b0769c8f51ed59acab6e994b00621096cfb1c640a7509295987388fbaf65"}, + {file = "soupsieve-2.8.4.tar.gz", hash = "sha256:e121fd02e975c695e4e9e8774a5ee35d74714b59307868dcc5319ad2d9e3328e"}, +] + [[package]] name = "sqlalchemy" version = "2.0.37" @@ -5474,6 +5659,37 @@ h2 = ["h2 (>=4,<5)"] socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"] zstd = ["zstandard (>=0.18.0)"] +[[package]] +name = "uuid-utils" +version = "0.14.1" +description = "Fast, drop-in replacement for Python's uuid module, powered by Rust." +optional = false +python-versions = ">=3.9" +files = [ + {file = "uuid_utils-0.14.1-cp39-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:93a3b5dc798a54a1feb693f2d1cb4cf08258c32ff05ae4929b5f0a2ca624a4f0"}, + {file = "uuid_utils-0.14.1-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:ccd65a4b8e83af23eae5e56d88034b2fe7264f465d3e830845f10d1591b81741"}, + {file = "uuid_utils-0.14.1-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b56b0cacd81583834820588378e432b0696186683b813058b707aedc1e16c4b1"}, + {file = "uuid_utils-0.14.1-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bb3cf14de789097320a3c56bfdfdd51b1225d11d67298afbedee7e84e3837c96"}, + {file = "uuid_utils-0.14.1-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:60e0854a90d67f4b0cc6e54773deb8be618f4c9bad98d3326f081423b5d14fae"}, + {file = "uuid_utils-0.14.1-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce6743ba194de3910b5feb1a62590cd2587e33a73ab6af8a01b642ceb5055862"}, + {file = "uuid_utils-0.14.1-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:043fb58fde6cf1620a6c066382f04f87a8e74feb0f95a585e4ed46f5d44af57b"}, + {file = "uuid_utils-0.14.1-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:c915d53f22945e55fe0d3d3b0b87fd965a57f5fd15666fd92d6593a73b1dd297"}, + {file = "uuid_utils-0.14.1-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:0972488e3f9b449e83f006ead5a0e0a33ad4a13e4462e865b7c286ab7d7566a3"}, + {file = "uuid_utils-0.14.1-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:1c238812ae0c8ffe77d8d447a32c6dfd058ea4631246b08b5a71df586ff08531"}, + {file = "uuid_utils-0.14.1-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:bec8f8ef627af86abf8298e7ec50926627e29b34fa907fcfbedb45aaa72bca43"}, + {file = "uuid_utils-0.14.1-cp39-abi3-win32.whl", hash = "sha256:b54d6aa6252d96bac1fdbc80d26ba71bad9f220b2724d692ad2f2310c22ef523"}, + {file = "uuid_utils-0.14.1-cp39-abi3-win_amd64.whl", hash = "sha256:fc27638c2ce267a0ce3e06828aff786f91367f093c80625ee21dad0208e0f5ba"}, + {file = "uuid_utils-0.14.1-cp39-abi3-win_arm64.whl", hash = "sha256:b04cb49b42afbc4ff8dbc60cf054930afc479d6f4dd7f1ec3bbe5dbfdde06b7a"}, + {file = "uuid_utils-0.14.1-pp311-pypy311_pp73-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:b197cd5424cf89fb019ca7f53641d05bfe34b1879614bed111c9c313b5574cd8"}, + {file = "uuid_utils-0.14.1-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:12c65020ba6cb6abe1d57fcbfc2d0ea0506c67049ee031714057f5caf0f9bc9c"}, + {file = "uuid_utils-0.14.1-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0b5d2ad28063d422ccc2c28d46471d47b61a58de885d35113a8f18cb547e25bf"}, + {file = "uuid_utils-0.14.1-pp311-pypy311_pp73-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:da2234387b45fde40b0fedfee64a0ba591caeea9c48c7698ab6e2d85c7991533"}, + {file = "uuid_utils-0.14.1-pp311-pypy311_pp73-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:50fffc2827348c1e48972eed3d1c698959e63f9d030aa5dd82ba451113158a62"}, + {file = "uuid_utils-0.14.1-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c1dbe718765f70f5b7f9b7f66b6a937802941b1cc56bcf642ce0274169741e01"}, + {file = "uuid_utils-0.14.1-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:258186964039a8e36db10810c1ece879d229b01331e09e9030bc5dcabe231bd2"}, + {file = "uuid_utils-0.14.1.tar.gz", hash = "sha256:9bfc95f64af80ccf129c604fb6b8ca66c6f256451e32bc4570f760e4309c9b69"}, +] + [[package]] name = "uvicorn" version = "0.34.0" @@ -6142,4 +6358,4 @@ cffi = ["cffi (>=1.11)"] [metadata] lock-version = "2.0" python-versions = ">=3.12, <4" -content-hash = "5acb48fa66fd8699daaf4fde50646599217adacc034c994b7dfd79d542dcf6c4" +content-hash = "44ddb6b6c435385a7e316dd81e05eb72bbc5b413d69d82f2359a3cfa086dcae3" diff --git a/pyproject.toml b/pyproject.toml index 9e89357..9992ab2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,9 @@ pyyaml = "^6.0.2" tavily-python = "^0.5.0" openpyxl = "^3.1.5" nltk = "^3.9.1" +beautifulsoup4 = "^4.12.0" +lxml = "^5.0.0" +requests = "^2.32.0" [tool.poetry.group.dev.dependencies] ruff = "^0.7.1" @@ -72,3 +75,116 @@ priority = "explicit" [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" + +[tool.ruff] +line-length = 88 +target-version = "py312" +src = ["src", "bin"] +extend-exclude = ["embeddings", "records", "data", ".venv"] + +[tool.ruff.lint] +# One rule set for the whole repo. E501 is the formatter's job -- it cannot split +# long string literals anyway -- and RUF001 fires on intentional typography inside +# LLM prompt text. +select = [ + "E", "F", "I", "UP", "B", "SIM", "RUF", # the core set + "S", # bandit: security + "C4", # comprehensions + "RET", # return-statement hygiene + "PIE", "PT", "ISC", "INT", "ICN", "TID", "A", "LOG", "G", "ERA", "N", +] +ignore = ["E501", "RUF001", "ISC001"] + +# Not yet enforced. Each needs a decision rather than a mechanical fix, so they +# are listed here instead of being silently dropped: +# PTH ~41 os.path -> pathlib rewrites; mechanical but touches files the +# retriever work will rewrite. Worth its own commit afterwards. +# DTZ datetime.now() without tzinfo. The rate limiter and static-message +# throttle deliberately store naive local timestamps in user metadata; +# changing that is a storage-format decision, not a lint fix. +# ARG unused arguments -- several are required by LangGraph node signatures. +# BLE blind `except Exception` in top-level scripts. +# C901 two functions over the complexity threshold. + +[tool.ruff.lint.isort] +known-first-party = [ + "agent", + "data_generation", + "evaluation", + "retrievers", + "tools", + "util", +] + +[tool.ruff.lint.per-file-ignores] +# Entry-point scripts intentionally configure the app at import time. +"bin/*" = ["E402"] +# assert is how pytest asserts. +"tests/*" = ["S101"] + +[tool.ruff.format] +docstring-code-format = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["src", "bin"] +addopts = "-q --strict-markers --strict-config" +markers = [ + "requires_embeddings: needs an installed embeddings bundle (see ./bin/embeddings_manager)", + "requires_retrieval_stack: needs the full langchain/chromadb dependency set installed", +] + +[tool.mypy] +python_version = "3.12" +files = ["bin", "src", "tests", ".github", "export_csvs.py"] +exclude = "^(data|embeddings|records)/" +mypy_path = "src" +explicit_package_bases = true +ignore_missing_imports = true +# --- ratchet: currently-enforced floor (the code already largely satisfies this) --- +check_untyped_defs = true +no_implicit_optional = true +strict_equality = true +warn_redundant_casts = true +# Annotation coverage reached 100%, so these are enforced repo-wide rather than +# per-module. There is now one standard for every file. +disallow_untyped_defs = true +disallow_incomplete_defs = true +warn_no_return = true +warn_return_any = true +extra_checks = true +# disallow_untyped_decorators stays off: chainlit's @cl.* decorators are +# untyped upstream, so it only ever fires on bin/chat-chainlit.py. + +# --- mypy baseline ------------------------------------------------------------- +# Modules that do not yet meet the floor above. Each entry is debt: fix the module, +# then delete its block. Do not add to this list without a TODO. +[[tool.mypy.overrides]] +# TODO(phase-2): every node returns a partial state; the TypedDict needs +# `total=False` or per-node Partial* types. +module = ["agent.profiles.cross_database"] +disable_error_code = ["typeddict-item", "no-any-return"] + +[[tool.mypy.overrides]] +# A module may appear in only one override section, so both TODOs live here. +# TODO(phase-2): partial-state returns, as for cross_database above. +# TODO(phase-2): preprocess() narrows BaseState to ReactToMeState, which is an LSP +# violation -- parameters may widen, not narrow. The fix is to make BaseGraphBuilder +# generic in its state type (BaseGraphBuilder[StateT]) so each profile declares its +# own. Deferred because base.py is also modified by the safetycheck and analysis +# branches, and a generics refactor there would collide with both. +module = ["agent.profiles.react_to_me"] +disable_error_code = ["typeddict-item", "override"] + +[[tool.mypy.overrides]] +# TODO(phase-2): `EmbeddingEnvironment.get_dir()` returns `Path | None` but the +# parameter is typed `Path`, so an uninstalled bundle passes None straight through. +# Fixed properly by moving the lookup out of the default argument (see ruff B008). +module = [ + "retrievers.reactome.rag", + "retrievers.uniprot.rag", + "retrievers.plantreactome.rag", + "retrievers.userguide.rag", +] +disable_error_code = ["assignment"] + diff --git a/src/agent/graph.py b/src/agent/graph.py index 012df27..f51666b 100644 --- a/src/agent/graph.py +++ b/src/agent/graph.py @@ -1,6 +1,6 @@ import asyncio import os -from typing import Any +from typing import Any, cast from langchain_core.callbacks.base import Callbacks from langchain_core.embeddings import Embeddings @@ -30,8 +30,16 @@ def __init__( profiles: list[ProfileName], ) -> None: # Get base models - llm: BaseChatModel = get_llm("openai", "gpt-4o-mini") - embedding: Embeddings = get_embedding("openai", "text-embedding-3-large") + embedding_model = os.getenv("EMBEDDING_MODEL", "bge-m3") + llm_model = os.getenv("LLM_MODEL", "gpt-4o-mini") + llm_base_url = os.getenv("LLM_BASE_URL", None) + llm: BaseChatModel = get_llm( + "openai", llm_model, base_url=llm_base_url, request_timeout=360.0 + ) + embedding_base_url = os.getenv("OPENAI_BASE_URL", None) + embedding: Embeddings = get_embedding( + "openai", embedding_model, base_url=embedding_base_url + ) self.uncompiled_graph: dict[str, StateGraph] = create_profile_graphs( profiles, llm, embedding @@ -87,14 +95,19 @@ async def ainvoke( self.graph = await self.initialize() if profile not in self.graph: return OutputState() - result: OutputState = await self.graph[profile].ainvoke( - InputState(user_input=user_input), - config=RunnableConfig( - callbacks=callbacks, - configurable={ - "thread_id": thread_id, - "enable_postprocess": enable_postprocess, - }, + # ainvoke is typed dict[str, Any] | Any; the graph's output schema is + # OutputState. + result: OutputState = cast( + "OutputState", + await self.graph[profile].ainvoke( + InputState(user_input=user_input), + config=RunnableConfig( + callbacks=callbacks, + configurable={ + "thread_id": thread_id, + "enable_postprocess": enable_postprocess, + }, + ), ), ) return result diff --git a/src/agent/models.py b/src/agent/models.py index 01b324c..904ca3e 100644 --- a/src/agent/models.py +++ b/src/agent/models.py @@ -2,8 +2,7 @@ from langchain_core.embeddings import Embeddings from langchain_core.language_models.chat_models import BaseChatModel -from langchain_huggingface import (HuggingFaceEmbeddings, - HuggingFaceEndpointEmbeddings) +from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpointEmbeddings from langchain_ollama.chat_models import ChatOllama from langchain_openai.chat_models.base import ChatOpenAI from langchain_openai.embeddings import OpenAIEmbeddings @@ -21,21 +20,21 @@ def get_embedding( model: str | None = None, *, device: str | None = "cpu", + base_url: str | None = None, ) -> Embeddings: if model is None: provider, model = provider.split("/", 1) if provider == "openai": - return OpenAIEmbeddings(model=model) - elif provider == "huggingfacehub": + return OpenAIEmbeddings(model=model, base_url=base_url) + if provider == "huggingfacehub": return HuggingFaceEndpointEmbeddings(model=model) - elif provider == "huggingfacelocal": + if provider == "huggingfacelocal": return HuggingFaceEmbeddings( model_name=model, model_kwargs={"device": device, "trust_remote_code": True}, encode_kwargs={"batch_size": 12, "normalize_embeddings": False}, ) - else: - raise ValueError(f"Unknown provider: {provider}") + raise ValueError(f"Unknown provider: {provider}") def get_llm( @@ -49,6 +48,7 @@ def get_llm( model: str | None = None, *, base_url: str | None = None, + request_timeout: float | None = None, ) -> BaseChatModel: if model is None: provider, model = provider.split("/", 1) @@ -57,12 +57,13 @@ def get_llm( model=model, temperature=0.0, base_url=base_url, + request_timeout=request_timeout, # type: ignore[call-arg] # pydantic-generated __init__ ) - elif provider == "ollama": + if provider == "ollama": return ChatOllama( model=model, temperature=0.0, base_url=base_url, + request_timeout=request_timeout, # type: ignore[call-arg] # pydantic-generated __init__ ) - else: - raise ValueError(f"Unknown provider: {provider}") + raise ValueError(f"Unknown provider: {provider}") diff --git a/src/agent/profile_names.py b/src/agent/profile_names.py new file mode 100644 index 0000000..402b3cc --- /dev/null +++ b/src/agent/profile_names.py @@ -0,0 +1,8 @@ +from enum import StrEnum + + +class ProfileName(StrEnum): + # These should exactly match names in .config.schema.yaml + React_to_Me = "React-to-Me" + Cross_Database_Prototype = "Cross-Database Prototype" + Plant_Reactome = "Plant Reactome" diff --git a/src/agent/profiles/__init__.py b/src/agent/profiles/__init__.py index 8061302..fbbcc7d 100644 --- a/src/agent/profiles/__init__.py +++ b/src/agent/profiles/__init__.py @@ -1,20 +1,16 @@ -from enum import StrEnum -from typing import Callable, NamedTuple +from collections.abc import Callable +from typing import NamedTuple from langchain_core.embeddings import Embeddings from langchain_core.language_models.chat_models import BaseChatModel from langgraph.graph.state import StateGraph +from agent.profile_names import ProfileName from agent.profiles.cross_database import create_cross_database_graph +from agent.profiles.plantreactome import create_plantreactome_graph from agent.profiles.react_to_me import create_reactome_graph -class ProfileName(StrEnum): - # These should exactly match names in .config.schema.yaml - React_to_Me = "React-to-Me" - Cross_Database_Prototype = "Cross-Database Prototype" - - class Profile(NamedTuple): name: ProfileName description: str @@ -32,6 +28,11 @@ class Profile(NamedTuple): description="Early version of an AI assistant with knowledge from multiple bio-databases (**Reactome** + **Uniprot**).", graph_builder=create_cross_database_graph, ), + ProfileName.Plant_Reactome.lower(): Profile( + name=ProfileName.Plant_Reactome, + description="An AI assistant specialized in exploring **Plant Reactome** biological pathways and processes.", + graph_builder=create_plantreactome_graph, + ), } diff --git a/src/agent/profiles/cross_database.py b/src/agent/profiles/cross_database.py index 31ab21a..79763f9 100644 --- a/src/agent/profiles/cross_database.py +++ b/src/agent/profiles/cross_database.py @@ -7,14 +7,19 @@ from langgraph.graph.state import StateGraph from agent.profiles.base import BaseGraphBuilder, BaseState -from agent.tasks.completeness_grader import (CompletenessGrade, - create_completeness_grader) -from agent.tasks.cross_database.rewrite_reactome_with_uniprot import \ - create_reactome_rewriter_w_uniprot -from agent.tasks.cross_database.rewrite_uniprot_with_reactome import \ - create_uniprot_rewriter_w_reactome -from agent.tasks.cross_database.summarize_reactome_uniprot import \ - create_reactome_uniprot_summarizer +from agent.tasks.completeness_grader import ( + CompletenessGrade, + create_completeness_grader, +) +from agent.tasks.cross_database.rewrite_reactome_with_uniprot import ( + create_reactome_rewriter_w_uniprot, +) +from agent.tasks.cross_database.rewrite_uniprot_with_reactome import ( + create_uniprot_rewriter_w_reactome, +) +from agent.tasks.cross_database.summarize_reactome_uniprot import ( + create_reactome_uniprot_summarizer, +) from retrievers.reactome.rag import create_reactome_rag from retrievers.uniprot.rag import create_uniprot_rag @@ -105,8 +110,7 @@ def check_question_safety( reactome_answer="", uniprot_answer="", ) - else: - return CrossDatabaseState() + return CrossDatabaseState() async def conduct_research( self, state: CrossDatabaseState, config: RunnableConfig @@ -203,7 +207,9 @@ async def assess_completeness( uniprot_completeness=uniprot_completeness.binary_score, ) - async def decide_next_steps(self, state: CrossDatabaseState) -> Literal[ + async def decide_next_steps( + self, state: CrossDatabaseState + ) -> Literal[ "generate_final_response", "perform_web_search", "rewrite_reactome_query", @@ -213,12 +219,11 @@ async def decide_next_steps(self, state: CrossDatabaseState) -> Literal[ uniprot_complete = state["uniprot_completeness"] != "No" if reactome_complete and uniprot_complete: return "generate_final_response" - elif not reactome_complete and uniprot_complete: + if not reactome_complete and uniprot_complete: return "rewrite_reactome_query" - elif reactome_complete and not uniprot_complete: + if reactome_complete and not uniprot_complete: return "rewrite_uniprot_query" - else: - return "perform_web_search" + return "perform_web_search" async def generate_final_response( self, state: CrossDatabaseState, config: RunnableConfig diff --git a/src/agent/profiles/plantreactome.py b/src/agent/profiles/plantreactome.py new file mode 100644 index 0000000..5475214 --- /dev/null +++ b/src/agent/profiles/plantreactome.py @@ -0,0 +1,100 @@ +from typing import Any + +from langchain_core.embeddings import Embeddings +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.runnables import Runnable, RunnableConfig +from langgraph.graph.state import StateGraph + +from agent.profiles.base import BaseGraphBuilder, BaseState +from agent.tasks.unsafe_question import create_unsafe_answer_generator +from retrievers.plantreactome.rag import create_plantreactome_rag + + +class PlantReactomeState(BaseState): + pass + + +class PlantReactomeGraphBuilder(BaseGraphBuilder): + def __init__( + self, + llm: BaseChatModel, + embedding: Embeddings, + ) -> None: + super().__init__(llm, embedding) + + # Create runnables (tasks & tools) + self.unsafe_answer_generator: Runnable = create_unsafe_answer_generator( + llm, streaming=True + ) + self.plantreactome_rag: Runnable = create_plantreactome_rag( + llm, embedding, streaming=True + ) + + # Create graph + state_graph = StateGraph(PlantReactomeState) + # Set up nodes + state_graph.add_node("preprocess", self.preprocess) + state_graph.add_node("model", self.call_model) + state_graph.add_node("generate_unsafe_response", self.generate_unsafe_response) + state_graph.add_node("postprocess", self.postprocess) + # Set up edges + state_graph.set_entry_point("preprocess") + state_graph.add_conditional_edges( + "preprocess", + self.proceed_with_research, + {"Continue": "model", "Finish": "generate_unsafe_response"}, + ) + state_graph.add_edge("model", "postprocess") + state_graph.add_edge("generate_unsafe_response", "postprocess") + state_graph.set_finish_point("postprocess") + + self.uncompiled_graph: StateGraph = state_graph + + async def generate_unsafe_response( + self, state: PlantReactomeState, config: RunnableConfig + ) -> PlantReactomeState: + answer: str = await self.unsafe_answer_generator.ainvoke( + { + "language": state["detected_language"], + "user_input": state["rephrased_input"], + "reason_unsafe": state["reason_unsafe"], + }, + config, + ) + return PlantReactomeState( + chat_history=[ + HumanMessage(state["user_input"]), + AIMessage(answer), + ], + answer=answer, + ) + + async def call_model( + self, state: PlantReactomeState, config: RunnableConfig + ) -> PlantReactomeState: + result: dict[str, Any] = await self.plantreactome_rag.ainvoke( + { + "input": state["rephrased_input"], + "chat_history": ( + state["chat_history"] + if state["chat_history"] + else [HumanMessage(state["user_input"])] + ), + }, + config, + ) + return PlantReactomeState( + chat_history=[ + HumanMessage(state["user_input"]), + AIMessage(result["answer"]), + ], + answer=result["answer"], + ) + + +def create_plantreactome_graph( + llm: BaseChatModel, + embedding: Embeddings, +) -> StateGraph: + return PlantReactomeGraphBuilder(llm, embedding).uncompiled_graph diff --git a/src/agent/profiles/react_to_me.py b/src/agent/profiles/react_to_me.py index dab20f0..9ebb0cb 100644 --- a/src/agent/profiles/react_to_me.py +++ b/src/agent/profiles/react_to_me.py @@ -1,3 +1,4 @@ +import logging from typing import Any from langchain_core.embeddings import Embeddings @@ -7,12 +8,23 @@ from langgraph.graph.state import StateGraph from agent.profiles.base import BaseGraphBuilder, BaseState +from agent.tasks.intent_classifier import ( + QueryIntent, + SourceName, + create_intent_classifier, + resolve_active_sources, +) +from agent.tasks.safety_checker import SafetyCheck from agent.tasks.unsafe_question import create_unsafe_answer_generator from retrievers.reactome.rag import create_reactome_rag +from retrievers.userguide.rag import create_userguide_rag +from util.embedding_environment import EmbeddingEnvironment + +logger = logging.getLogger(__name__) class ReactToMeState(BaseState): - pass + active_sources: list[SourceName] class ReactToMeGraphBuilder(BaseGraphBuilder): @@ -23,22 +35,22 @@ def __init__( ) -> None: super().__init__(llm, embedding) - # Create runnables (tasks & tools) + self.intent_classifier: Runnable = create_intent_classifier(llm) self.unsafe_answer_generator: Runnable = create_unsafe_answer_generator( llm, streaming=True ) - self.reactome_rag: Runnable = create_reactome_rag( - llm, embedding, streaming=True - ) - # Create graph + self.rags: dict[SourceName, Runnable] = { + "reactome": create_reactome_rag(llm, embedding, streaming=True), + } + self._available_sources: frozenset[SourceName] = frozenset({"reactome"}) + self._register_userguide_rag(llm, embedding) + state_graph = StateGraph(ReactToMeState) - # Set up nodes state_graph.add_node("preprocess", self.preprocess) - state_graph.add_node("model", self.call_model) + state_graph.add_node("model", self.generate_answer) state_graph.add_node("generate_unsafe_response", self.generate_unsafe_response) state_graph.add_node("postprocess", self.postprocess) - # Set up edges state_graph.set_entry_point("preprocess") state_graph.add_conditional_edges( "preprocess", @@ -51,6 +63,69 @@ def __init__( self.uncompiled_graph: StateGraph = state_graph + def _register_userguide_rag( + self, + llm: BaseChatModel, + embedding: Embeddings, + ) -> None: + userguide_dir = EmbeddingEnvironment.get_dir("userguide") + if userguide_dir is None: + logger.info( + "User guide embeddings not configured; routing will use reactome only." + ) + return + + chroma_path = userguide_dir / "sections" / "chroma.sqlite3" + if not chroma_path.exists(): + logger.warning( + "User guide embeddings directory exists but Chroma DB is missing at %s", + chroma_path, + ) + return + + try: + self.rags["userguide"] = create_userguide_rag( + llm, embedding, userguide_dir, streaming=True + ) + self._available_sources = frozenset(self.rags) + except (FileNotFoundError, ValueError) as exc: + logger.warning("User guide RAG unavailable: %s", exc) + + async def preprocess( + self, state: ReactToMeState, config: RunnableConfig + ) -> ReactToMeState: + rephrased_input: str = await self.rephrase_chain.ainvoke( + { + "user_input": state["user_input"], + "chat_history": state.get("chat_history", []), + }, + config, + ) + safety_check: SafetyCheck = await self.safety_checker.ainvoke( + {"rephrased_input": rephrased_input}, config + ) + detected_language: str = await self.language_detector.ainvoke( + {"user_input": state["user_input"]}, config + ) + intent: QueryIntent = await self.intent_classifier.ainvoke( + {"rephrased_input": rephrased_input}, config + ) + active_sources = resolve_active_sources(intent.source, self._available_sources) + if intent.source not in self._available_sources: + logger.info( + "Requested source %r unavailable; falling back to %r", + intent.source, + active_sources[0], + ) + + return ReactToMeState( + rephrased_input=rephrased_input, + safety=safety_check.safety, + reason_unsafe=safety_check.reason_unsafe, + detected_language=detected_language, + active_sources=active_sources, + ) + async def generate_unsafe_response( self, state: ReactToMeState, config: RunnableConfig ) -> ReactToMeState: @@ -70,10 +145,12 @@ async def generate_unsafe_response( answer=answer, ) - async def call_model( + async def generate_answer( self, state: ReactToMeState, config: RunnableConfig ) -> ReactToMeState: - result: dict[str, Any] = await self.reactome_rag.ainvoke( + source = state["active_sources"][0] + rag = self.rags[source] + result: dict[str, Any] = await rag.ainvoke( { "input": state["rephrased_input"], "chat_history": ( diff --git a/src/agent/tasks/completeness_grader.py b/src/agent/tasks/completeness_grader.py index 129866e..e2254ac 100644 --- a/src/agent/tasks/completeness_grader.py +++ b/src/agent/tasks/completeness_grader.py @@ -7,9 +7,11 @@ You are an expert grader with extensive knowledge in molecular biology and experience as a curator for both Reactome and UniProt knowledgebases. Your task is to evaluate whether a response generated by an LLM is complete, meaning it addresses the user’s question with necessary details, background information, and context. -Provide a binary output as either: - - Yes: The response answers the user question and provides enough details and background. - - No: The response is incomplete, missing key details, or lacking sufficient context. +IMPORTANT: You must respond with ONLY a valid JSON object using this exact format: +{{"binary_score": "Yes"}} or {{"binary_score":"No"}} + +Do NOT use XML tags, function calls, or any other format. +Do NOT include any text outside the JSON object. """ completeness_prompt = ChatPromptTemplate.from_messages( diff --git a/src/agent/tasks/intent_classifier.py b/src/agent/tasks/intent_classifier.py new file mode 100644 index 0000000..17aa2b0 --- /dev/null +++ b/src/agent/tasks/intent_classifier.py @@ -0,0 +1,61 @@ +from typing import Literal + +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.prompts import ChatPromptTemplate +from langchain_core.runnables import Runnable +from pydantic import BaseModel, Field + +SourceName = Literal["reactome", "userguide"] + +intent_classifier_message = """ +You route user questions for the React-to-Me assistant to the correct knowledge source. + +Choose exactly one source: + +- **reactome**: Questions about biology, molecular mechanisms, pathways, reactions, proteins, genes, + diseases, and other scientific content in the Reactome Knowledgebase. + Examples: "What is apoptosis?", "Which pathways involve TP53?", "What does CDK5 do?" + +- **userguide**: Questions about how to use the Reactome **website**, tools, or interface. + Examples: "How do I use the pathway browser?", "How do I search Reactome?", + "How do I run gene list analysis?", "What is the Details Panel?" + +Rules: +- If the user asks how to perform a task in Reactome or about UI features, choose **userguide**. +- If the user asks about biological facts or pathway content, choose **reactome**. +- When unsure, prefer **reactome** for science content and **userguide** for clear how-to or UI questions. +""" + +intent_classifier_prompt = ChatPromptTemplate.from_messages( + [ + ("system", intent_classifier_message), + ("human", "User question:\n\n{rephrased_input}"), + ] +) + + +class QueryIntent(BaseModel): + source: SourceName = Field( + description="The knowledge source that should answer this question: 'reactome' or 'userguide'." + ) + + +_FALLBACK_ORDER: tuple[SourceName, ...] = ("reactome", "userguide") + + +def resolve_active_sources( + source: SourceName, + available_sources: frozenset[SourceName], +) -> list[SourceName]: + if not available_sources: + raise ValueError("available_sources must not be empty") + if source in available_sources: + return [source] + for fallback in _FALLBACK_ORDER: + if fallback in available_sources: + return [fallback] + return [next(iter(available_sources))] + + +def create_intent_classifier(llm: BaseChatModel) -> Runnable: + return intent_classifier_prompt | llm.with_structured_output(QueryIntent) diff --git a/src/agent/tasks/rephrase.py b/src/agent/tasks/rephrase.py index 1851747..23aaeaf 100644 --- a/src/agent/tasks/rephrase.py +++ b/src/agent/tasks/rephrase.py @@ -4,15 +4,17 @@ from langchain_core.runnables import Runnable contextualize_q_system_prompt = """ -You are an expert in question formulation with deep expertise in molecular biology and experience as a Reactome curator. Your task is to analyze the conversation history and the user’s latest query to fully understand their intent and what they seek to learn. +You are an expert in question formulation with deep expertise in molecular biology and experience as a Reactome curator. Your task is to analyze the conversation history and the user's latest query to fully understand their intent and what they seek to learn. If the user's question is not in English, reformulate the question and translate it to English, ensuring the meaning and intent are preserved. -Reformulate the user’s question into a standalone version that retains its full meaning without requiring prior context. The reformulated question should be: +Reformulate the user's question into a standalone version that retains its full meaning without requiring prior context. The reformulated question should be: - Clear, concise, and precise - Optimized for both vector search (semantic meaning) and case-sensitive keyword search - - Faithful to the user’s intent and scientific accuracy + - Faithful to the user's intent and scientific accuracy + +If the question is about how to use the Reactome website or its tools (Pathway Browser, search, analysis tools, Details Panel, etc.), keep it as a how-to or UI question. Do not rewrite it into a biological pathway or mechanism question. the returned question should always be in English. -If the user’s question is already in English, self-contained and well-formed, return it as is. +If the user's question is already in English, self-contained and well-formed, return it as is. Do NOT answer the question or provide explanations. """ diff --git a/src/agent/tasks/safety_checker.py b/src/agent/tasks/safety_checker.py index 91e539f..c136013 100644 --- a/src/agent/tasks/safety_checker.py +++ b/src/agent/tasks/safety_checker.py @@ -17,7 +17,9 @@ - Treat hypothetical, fictional, or made-up scenarios with the same level of scrutiny as real-world questions. 2. Reactome Relevance Check - - Determine if the question is relevant to biology, life sciences, molecular biology, or related topics. + - Determine if the question is relevant to Reactome. Relevant topics include: + - Biology, life sciences, molecular biology, pathways, proteins, genes, and related scientific topics. + - How to use the Reactome website, Pathway Browser, search, analysis tools, and other Reactome features (user guide topics). - Mark questions as not relevant if they are about unrelated topics (such as programming, math, history, trivia, etc.). IMPORTANT: @@ -44,6 +46,10 @@ 4. Q: What is the role of the immune system in the treatment of cancer? "safety": "true", "reason_unsafe": "" + + 5. Q: How do I use the Reactome pathway browser? + "safety": "true", + "reason_unsafe": "" """ safety_check_prompt = ChatPromptTemplate.from_messages( diff --git a/src/data_generation/alliance/__init__.py b/src/data_generation/alliance/__init__.py index b0d2490..2c79c9b 100644 --- a/src/data_generation/alliance/__init__.py +++ b/src/data_generation/alliance/__init__.py @@ -1,12 +1,10 @@ import os -from typing import Optional import requests import torch from langchain_community.vectorstores import Chroma from langchain_core.embeddings import Embeddings -from langchain_huggingface import (HuggingFaceEmbeddings, - HuggingFaceEndpointEmbeddings) +from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpointEmbeddings from langchain_openai import OpenAIEmbeddings from data_generation.alliance.csv_generator import generate_all_csvs @@ -15,27 +13,25 @@ def get_release_version() -> str: url: str = "https://www.alliancegenome.org/api/releaseInfo" - response = requests.get(url) + response = requests.get(url, timeout=60) if response.status_code == 200: response_json = response.json() - release_version = response_json.get("releaseVersion") + release_version: str | None = response_json.get("releaseVersion") if release_version: return release_version - else: - raise ValueError("Release version not found in the response.") - else: - raise ConnectionError( - f"Failed to get the response. Status code: {response.status_code}" - ) + raise ValueError("Release version not found in the response.") + raise ConnectionError( + f"Failed to get the response. Status code: {response.status_code}" + ) def upload_to_chromadb( embeddings_dir: str, version: str, force: bool, # Changed from str to bool - hf_model: Optional[str] = None, - device: Optional[str] = None, -) -> Optional[Chroma]: + hf_model: str | None = None, + device: str | None = None, +) -> Chroma | None: metadata_columns: dict[str, list] = { "genes": [ "Your Input", @@ -108,7 +104,7 @@ def upload_to_chromadb( "AnatomyTermQualifierIDs", "AnatomyTermQualifierTermNames", "SourceURL", - "Source," "Reference", + "Source,Reference", ], "molecular_interaction": [ "ID(s) interactor A", @@ -138,7 +134,7 @@ def upload_to_chromadb( "Annotation(s) interactor A", "Annotation(s) interactor B", "Interaction annotation(s)", - "Host organism(s)" "Interaction parameter(s)", + "Host organism(s)Interaction parameter(s)", "Creation date", "Update date", "Checksum(s) interactor A", @@ -298,9 +294,9 @@ def upload_to_chromadb( def generate_alliance_embeddings( embeddings_dir: str, force: bool = False, - hf_model: Optional[str] = None, - device: Optional[str] = None, - **kwargs, + hf_model: str | None = None, + device: str | None = None, + **kwargs: object, ) -> None: release_version = get_release_version() print(f"Release Version: {release_version}") diff --git a/src/data_generation/alliance/csv_generator.py b/src/data_generation/alliance/csv_generator.py index 691d016..f159b65 100644 --- a/src/data_generation/alliance/csv_generator.py +++ b/src/data_generation/alliance/csv_generator.py @@ -1,12 +1,11 @@ import gzip import os import shutil -from typing import Optional import requests -def download_file(url: str, dest: str, force: bool) -> Optional[str]: +def download_file(url: str, dest: str, force: bool) -> str | None: # Create the directory if it doesn't exist os.makedirs(os.path.dirname(dest), exist_ok=True) @@ -16,7 +15,7 @@ def download_file(url: str, dest: str, force: bool) -> Optional[str]: return dest # Send the GET request - response = requests.get(url) + response = requests.get(url, timeout=60) # Check if the request was successful if response.status_code == 200: @@ -28,18 +27,16 @@ def download_file(url: str, dest: str, force: bool) -> Optional[str]: # Check if the file is gzipped and decompress if necessary if dest.endswith(".gz"): unzipped_dest = dest[:-3] # Remove '.gz' from the filename - with gzip.open(dest, "rb") as f_in: - with open(unzipped_dest, "wb") as f_out: - shutil.copyfileobj(f_in, f_out) + with gzip.open(dest, "rb") as f_in, open(unzipped_dest, "wb") as f_out: + shutil.copyfileobj(f_in, f_out) print(f"File unzipped successfully and saved to {unzipped_dest}.") os.remove(dest) # Remove the gzipped file after extraction return unzipped_dest return dest - else: - print( - f"Failed to download the file from {url}. Status code: {response.status_code}" - ) - return None + print( + f"Failed to download the file from {url}. Status code: {response.status_code}" + ) + return None def get_genes(version: str, force: bool) -> str: @@ -96,7 +93,7 @@ def get_genes(version: str, force: bool) -> str: } # Send the POST request - response = requests.post(url, data=form_data) + response = requests.post(url, data=form_data, timeout=60) if response.status_code == 200: # Save the response content to a file diff --git a/src/data_generation/metadata_csv_loader.py b/src/data_generation/metadata_csv_loader.py index 0933281..093ef50 100644 --- a/src/data_generation/metadata_csv_loader.py +++ b/src/data_generation/metadata_csv_loader.py @@ -1,6 +1,6 @@ import csv from io import TextIOWrapper -from typing import Any, Optional +from typing import Any from langchain_community.document_loaders.base import BaseLoader from langchain_community.document_loaders.helpers import detect_file_encodings @@ -31,11 +31,11 @@ class MetaDataCSVLoader(BaseLoader): def __init__( self, file_path: str, - source_column: Optional[str] = None, - metadata_columns: Optional[list[str]] = None, - content_columns: Optional[list[str]] = None, - csv_args: dict[str, Any] = dict(), - encoding: Optional[str] = None, + source_column: str | None = None, + metadata_columns: list[str] | None = None, + content_columns: list[str] | None = None, + csv_args: dict[str, Any] | None = None, + encoding: str | None = None, autodetect_encoding: bool = False, ) -> None: """ @@ -51,11 +51,11 @@ def __init__( autodetect_encoding: Whether to try to autodetect the file encoding. """ self.file_path: str = file_path - self.source_column: Optional[str] = source_column - self.metadata_columns: Optional[list[str]] = metadata_columns - self.content_columns: Optional[list[str]] = content_columns - self.encoding: Optional[str] = encoding - self.csv_args: dict[str, Any] = csv_args + self.source_column: str | None = source_column + self.metadata_columns: list[str] | None = metadata_columns + self.content_columns: list[str] | None = content_columns + self.encoding: str | None = encoding + self.csv_args: dict[str, Any] = csv_args if csv_args is not None else {} self.autodetect_encoding: bool = autodetect_encoding def load(self) -> list[Document]: @@ -99,10 +99,10 @@ def __read_file(self, csvfile: TextIOWrapper) -> list[Document]: if self.source_column is not None else self.file_path ) - except KeyError: + except KeyError as e: raise ValueError( f"Source column '{self.source_column}' not found in CSV file." - ) + ) from e # Construct content from content_columns if provided, otherwise use all columns if self.content_columns: @@ -122,10 +122,10 @@ def __read_file(self, csvfile: TextIOWrapper) -> list[Document]: for col in self.metadata_columns: try: metadata[col] = row[col] - except KeyError: + except KeyError as e: raise ValueError( f"Metadata column '{col}' not found in CSV file." - ) + ) from e doc = Document(page_content=content, metadata=metadata) docs.append(doc) diff --git a/src/data_generation/reactome/__init__.py b/src/data_generation/reactome/__init__.py index f62cd38..8b6b90f 100644 --- a/src/data_generation/reactome/__init__.py +++ b/src/data_generation/reactome/__init__.py @@ -1,11 +1,10 @@ import os -from typing import Optional +from pathlib import Path import torch from langchain_community.vectorstores import Chroma from langchain_core.embeddings import Embeddings -from langchain_huggingface import (HuggingFaceEmbeddings, - HuggingFaceEndpointEmbeddings) +from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpointEmbeddings from langchain_openai import OpenAIEmbeddings from data_generation.metadata_csv_loader import MetaDataCSVLoader @@ -17,8 +16,8 @@ def upload_to_chromadb( embeddings_dir: str, file: str, embedding_table: str, - hf_model: Optional[str] = None, - device: Optional[str] = None, + hf_model: str | None = None, + device: str | None = None, ) -> Chroma: metadata_columns: dict[str, list] = { "reactions": [ @@ -26,6 +25,7 @@ def upload_to_chromadb( "display_name", "pathway_id", "pathway_name", + "species", "input_id", "input_name", "output_id", @@ -33,8 +33,14 @@ def upload_to_chromadb( "catalyst_id", "catalyst_name", ], - "summations": ["st_id", "display_name", "summation"], - "complexes": ["st_id", "display_name", "component_id", "component_name"], + "summations": ["st_id", "display_name", "labels", "species", "summation"], + "complexes": [ + "st_id", + "display_name", + "component_id", + "component_name", + "species", + ], "ewas": [ "st_id", "display_name", @@ -53,13 +59,13 @@ def upload_to_chromadb( embeddings_instance: Embeddings if hf_model is None: # Use OpenAI embeddings_instance = OpenAIEmbeddings( - chunk_size=500, + chunk_size=400, show_progress_bar=True, ) - elif hf_model.startswith("openai/text-embedding-"): + elif hf_model.startswith("openai/"): embeddings_instance = OpenAIEmbeddings( model=hf_model[len("openai/") :], - chunk_size=500, + chunk_size=400, show_progress_bar=True, ) elif "HUGGINGFACEHUB_API_TOKEN" in os.environ: @@ -86,20 +92,33 @@ def upload_to_chromadb( def generate_reactome_embeddings( embeddings_dir: str, neo4j_uri: str = "bolt://localhost:7687", - neo4j_username: Optional[str] = None, - neo4j_password: Optional[str] = None, + neo4j_username: str | None = None, + neo4j_password: str | None = None, force: bool = False, - hf_model: Optional[str] = None, - device: Optional[str] = None, + hf_model: str | None = None, + device: str | None = None, ) -> None: - connector = Neo4jConnector( - uri=neo4j_uri, user=neo4j_username, password=neo4j_password - ) + csv_dir = Path(embeddings_dir) / "csv_files" + reactions_csv = str(csv_dir / "reactions.csv") + summations_csv = str(csv_dir / "summations.csv") + complexes_csv = str(csv_dir / "complexes.csv") + ewas_csv = str(csv_dir / "ewas.csv") - (reactions_csv, summations_csv, complexes_csv, ewas_csv) = generate_all_csvs( - connector, embeddings_dir, force + all_exist = not force and all( + Path(p).exists() + for p in [reactions_csv, summations_csv, complexes_csv, ewas_csv] ) - connector.close() + + if not all_exist: + connector = Neo4jConnector( + uri=neo4j_uri, user=neo4j_username, password=neo4j_password + ) + reactions_csv, summations_csv, complexes_csv, ewas_csv = generate_all_csvs( + connector, embeddings_dir, force + ) + connector.close() + else: + print("Using existing CSV files. Skipping Neo4j.") db = upload_to_chromadb( embeddings_dir, reactions_csv, "reactions", hf_model, device diff --git a/src/data_generation/reactome/csv_generator.py b/src/data_generation/reactome/csv_generator.py index 15e3689..7f94f05 100644 --- a/src/data_generation/reactome/csv_generator.py +++ b/src/data_generation/reactome/csv_generator.py @@ -1,14 +1,18 @@ +from collections.abc import Callable from pathlib import Path -from typing import Callable import pandas as pd -from data_generation.reactome.neo4j_connector import (Neo4jConnector, - Neo4jDict, get_complexes, - get_ewas, get_reactions, - get_summations) +from data_generation.reactome.neo4j_connector import ( + Neo4jConnector, + Neo4jDict, + get_complexes, + get_ewas, + get_reactions, + get_summations, +) -CSV_GENERATION_MAP: dict[str, Callable[[Neo4jConnector], Neo4jDict]] = { +CSV_GENERATION_MAP: dict[str, Callable[[Neo4jConnector], list[Neo4jDict]]] = { "reactions.csv": get_reactions, "summations.csv": get_summations, "complexes.csv": get_complexes, @@ -18,7 +22,7 @@ def generate_csv( connector: Neo4jConnector, - data_fetch_func: Callable[[Neo4jConnector], Neo4jDict], + data_fetch_func: Callable[[Neo4jConnector], list[Neo4jDict]], file_name: str, csv_dir: Path, force: bool = False, @@ -28,7 +32,7 @@ def generate_csv( if not force and csv_file_path.exists(): return str(csv_file_path) - data: Neo4jDict = data_fetch_func(connector) + data: list[Neo4jDict] = data_fetch_func(connector) df: pd.DataFrame = pd.DataFrame(data) df["url"] = "https://reactome.org/content/detail/" + df["st_id"] df.to_csv(csv_file_path, index=False, lineterminator="\n") diff --git a/src/data_generation/reactome/neo4j_connector.py b/src/data_generation/reactome/neo4j_connector.py index 2c163e4..6172e51 100644 --- a/src/data_generation/reactome/neo4j_connector.py +++ b/src/data_generation/reactome/neo4j_connector.py @@ -1,4 +1,4 @@ -from typing import Any, Optional +from typing import Any from neo4j import GraphDatabase @@ -6,7 +6,7 @@ class Neo4jConnector: - def __init__(self, uri: str, user: Optional[str], password: Optional[str]): + def __init__(self, uri: str, user: str | None, password: str | None) -> None: if user is None or password is None: self._driver = GraphDatabase.driver(uri) else: @@ -18,7 +18,8 @@ def close(self) -> None: def execute_query(self, query: str) -> list[Neo4jDict]: with self._driver.session() as session: result = session.run(query) - return result.data() + records: list[Neo4jDict] = result.data() + return records def get_reactions(connector: Neo4jConnector) -> list[Neo4jDict]: diff --git a/src/data_generation/uniprot/__init__.py b/src/data_generation/uniprot/__init__.py index f5b24fb..7b27538 100644 --- a/src/data_generation/uniprot/__init__.py +++ b/src/data_generation/uniprot/__init__.py @@ -1,12 +1,10 @@ import os from pathlib import Path -from typing import Optional import torch from langchain_community.vectorstores import Chroma from langchain_core.embeddings import Embeddings -from langchain_huggingface import (HuggingFaceEmbeddings, - HuggingFaceEndpointEmbeddings) +from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpointEmbeddings from langchain_openai import OpenAIEmbeddings from data_generation.metadata_csv_loader import MetaDataCSVLoader @@ -17,8 +15,8 @@ def upload_to_chromadb( embeddings_dir: str, file: str, embedding_table: str, - hf_model: Optional[str] = None, - device: Optional[str] = None, + hf_model: str | None = None, + device: str | None = None, ) -> Chroma: metadata_columns: dict[str, list] = { "uniprot_data": [ @@ -44,13 +42,13 @@ def upload_to_chromadb( print("Using OpenAI embeddings") embeddings_instance = OpenAIEmbeddings( model="text-embedding-3-large", - chunk_size=800, + chunk_size=500, show_progress_bar=True, ) elif hf_model.startswith("openai/text-embedding-"): embeddings_instance = OpenAIEmbeddings( model=hf_model[len("openai/") :], - chunk_size=800, + chunk_size=500, show_progress_bar=True, ) elif "HUGGINGFACEHUB_API_TOKEN" in os.environ: @@ -76,9 +74,9 @@ def upload_to_chromadb( def generate_uniprot_embeddings( embedding_path: Path, - hf_model: Optional[str] = None, - device: Optional[str] = None, - **_, + hf_model: str | None = None, + device: str | None = None, + **_: object, ) -> None: csv_path = generate_uniprot_csv(embedding_path) db = upload_to_chromadb( diff --git a/src/data_generation/uniprot/api_connector.py b/src/data_generation/uniprot/api_connector.py index 50d0844..76db77e 100644 --- a/src/data_generation/uniprot/api_connector.py +++ b/src/data_generation/uniprot/api_connector.py @@ -1,4 +1,5 @@ import re +from collections.abc import Iterator, Mapping import requests from requests.adapters import HTTPAdapter, Retry @@ -8,7 +9,7 @@ class UniProtAPIConnector: BASE_URL = "https://rest.uniprot.org/uniprotkb/stream" @staticmethod - def get_download_url(): + def get_download_url() -> str: """ Returns the UniProt API URL for downloading human-reviewed protein data. """ @@ -21,10 +22,10 @@ def get_download_url(): ) return UniProtAPIConnector.BASE_URL + query_params - def __init__(self): + def __init__(self) -> None: self.session = self._initialize_session() - def _initialize_session(self): + def _initialize_session(self) -> requests.Session: """Creates a session with retry logic for robust downloading.""" retries = Retry( total=5, backoff_factor=0.25, status_forcelist=[500, 502, 503, 504] @@ -33,7 +34,7 @@ def _initialize_session(self): session.mount("https://", HTTPAdapter(max_retries=retries)) return session - def get_next_link(self, headers): + def get_next_link(self, headers: Mapping[str, str]) -> str | None: """Parses the 'Link' header to find the URL for the next batch of data.""" re_next_link = re.compile(r'<(.+)>; rel="next"') if "Link" in headers: @@ -42,11 +43,14 @@ def get_next_link(self, headers): return match.group(1) return None - def get_batch(self, batch_url): + def get_batch(self, batch_url: str) -> Iterator[tuple[requests.Response, str]]: """Generator to download data in batches.""" - while batch_url: - response = self.session.get(batch_url) + # Local because get_next_link returns None at the end of pagination, + # while callers always pass a real URL in. + url: str | None = batch_url + while url: + response = self.session.get(url) response.raise_for_status() # Ensure we stop on HTTP errors - total = response.headers.get("x-total-results", 0) + total = response.headers.get("x-total-results", "0") yield response, total - batch_url = self.get_next_link(response.headers) + url = self.get_next_link(response.headers) diff --git a/src/data_generation/uniprot/csv_generator.py b/src/data_generation/uniprot/csv_generator.py index 5159226..9c414c0 100644 --- a/src/data_generation/uniprot/csv_generator.py +++ b/src/data_generation/uniprot/csv_generator.py @@ -7,14 +7,16 @@ class UniProtDataCleaner: - def __init__(self, csv_dir: Path): + def __init__(self, csv_dir: Path) -> None: self.download_url = UniProtAPIConnector.get_download_url() self.xlsx_path = csv_dir / "uniprot_data.xlsx" self.csv_path = self.xlsx_path.with_suffix(".csv") - self.df = None + # Declared, not assigned: populated by load_data(). Using it before then + # raises AttributeError, which is louder than an empty DataFrame. + self.df: pd.DataFrame self.api = UniProtAPIConnector() - def download_data(self): + def download_data(self) -> None: """Downloads data batch by batch using UniProt API connector.""" progress = 0 with open(self.xlsx_path, "wb") as f: @@ -24,12 +26,12 @@ def download_data(self): print(f"Downloaded {progress} batches; Total: {total}") print(f"✅ UniProt data downloaded successfully to {self.xlsx_path}") - def load_data(self): + def load_data(self) -> None: """Loads data from Excel file into a DataFrame.""" print(f"Loading data from {self.xlsx_path}") self.df = pd.read_excel(self.xlsx_path) - def clean_data(self): + def clean_data(self) -> None: """Cleans the UniProt data using predefined processing steps.""" self.load_data() self.remove_prefixes() @@ -42,7 +44,7 @@ def clean_data(self): self.df.to_csv(self.csv_path, index=False) print(f"Cleaned data saved to {self.csv_path}") - def remove_prefixes(self): + def remove_prefixes(self) -> None: """Remove prefixes from specified columns.""" prefix_map = { "Entry Name": "_HUMAN", @@ -63,12 +65,12 @@ def remove_prefixes(self): self.df[column].str.replace(prefix, "", regex=False).str.strip() ) - def add_url(self): + def add_url(self) -> None: """Replace 'Entry' column with URLs constructed from entry IDs.""" base_url = "https://www.uniprot.org/uniprotkb/" self.df["Entry"] = base_url + self.df["Entry"].astype(str) + "/entry" - def format_names(self): + def format_names(self) -> None: """Format gene synonyms and protein names with semicolons and proper punctuation.""" self.df["Gene Names"] = ( self.df["Gene Names"].str.replace(" ", "; ", regex=False).str.strip("; ") @@ -84,12 +86,12 @@ def format_names(self): ) ) - def format_mass(self): + def format_mass(self) -> None: """Format the 'Mass' column by appending ' Da' to each mass value.""" if "Mass" in self.df.columns: self.df["Mass"] = self.df["Mass"].apply(lambda x: f"{x} Da") - def clean_evidence_codes(self): + def clean_evidence_codes(self) -> None: """Remove citations and evidence codes from textual columns.""" patterns = [r"\{ECO:[^\}]*\}", r"\(PubMed:[^\)]*\)", r"\[MIM:[^\]]*\]", r" +"] for column in self.df.columns: @@ -98,10 +100,10 @@ def clean_evidence_codes(self): self.df[column].str.replace(pattern, "", regex=True).str.strip() ) - def clean_columns(self): + def clean_columns(self) -> None: """Reformat entries in the 'Motif' column.""" - def reformat_motif(entry): + def reformat_motif(entry: str) -> str: if pd.isna(entry): return entry pattern = r"MOTIF (\d+\.\.\d+); /note=\"([^\"]*)\"; /evidence=\"[^\"]*\"" @@ -113,7 +115,7 @@ def reformat_motif(entry): ] ) - def reformat_domain(entry): + def reformat_domain(entry: str) -> str: if pd.isna(entry): return entry pattern = r"DOMAIN (\d+\.\.\d+); /note=\"([^\"]*)\"; /evidence=\"[^\"]*\"" @@ -128,7 +130,7 @@ def reformat_domain(entry): self.df["Motif"] = self.df["Motif"].apply(reformat_motif) self.df["Domain [FT]"] = self.df["Domain [FT]"].apply(reformat_domain) - def rename_columns(self): + def rename_columns(self) -> None: """Rename columns as specified.""" new_column_names = { "Entry": "url", diff --git a/src/data_generation/userguide/__init__.py b/src/data_generation/userguide/__init__.py new file mode 100644 index 0000000..82de7eb --- /dev/null +++ b/src/data_generation/userguide/__init__.py @@ -0,0 +1,87 @@ +import os +from pathlib import Path +from shutil import rmtree + +import torch +from langchain_community.vectorstores import Chroma +from langchain_core.documents import Document +from langchain_core.embeddings import Embeddings +from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpointEmbeddings +from langchain_openai import OpenAIEmbeddings + +from data_generation.userguide.fetch import fetch_userguide_pages +from data_generation.userguide.html_loader import UserGuideHTMLLoader +from data_generation.userguide.urls import USER_GUIDE_URLS + +CHROMA_COLLECTION = "sections" +HTML_CACHE_DIR = "html_snapshots" + + +def upload_to_chromadb( + embeddings_dir: str, + docs: list[Document], + embedding_table: str, + hf_model: str | None = None, + device: str | None = None, +) -> Chroma: + embeddings_instance: Embeddings + if hf_model is None: # Use OpenAI + embeddings_instance = OpenAIEmbeddings( + chunk_size=500, + show_progress_bar=True, + ) + elif hf_model.startswith("openai/text-embedding-"): + embeddings_instance = OpenAIEmbeddings( + model=hf_model[len("openai/") :], + chunk_size=500, + show_progress_bar=True, + ) + elif "HUGGINGFACEHUB_API_TOKEN" in os.environ: + embeddings_instance = HuggingFaceEndpointEmbeddings( + huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"], + model=hf_model, + ) + else: + if device == "cuda": + torch.cuda.empty_cache() + embeddings_instance = HuggingFaceEmbeddings( + model_name=hf_model, + model_kwargs={"device": device, "trust_remote_code": True}, + encode_kwargs={"batch_size": 12, "normalize_embeddings": False}, + ) + + return Chroma.from_documents( + documents=docs, + embedding=embeddings_instance, + persist_directory=os.path.join(embeddings_dir, embedding_table), + ) + + +def generate_userguide_embeddings( + embeddings_dir: str, + force: bool = False, + hf_model: str | None = None, + device: str | None = None, + **_: object, +) -> None: + embeddings_path = Path(embeddings_dir) + chroma_dir = embeddings_path / CHROMA_COLLECTION + if force and chroma_dir.exists(): + rmtree(chroma_dir) + + cache_dir = embeddings_path / HTML_CACHE_DIR + html_paths = fetch_userguide_pages( + USER_GUIDE_URLS, + cache_dir=cache_dir, + force=force, + ) + + loader = UserGuideHTMLLoader(html_paths) + docs = loader.load() + print(f"Loaded {len(docs)} user guide sections from {len(html_paths)} pages") + + if not docs: + raise RuntimeError("No user guide documents were produced") + + db = upload_to_chromadb(embeddings_dir, docs, CHROMA_COLLECTION, hf_model, device) + print(db._collection.count()) diff --git a/src/data_generation/userguide/fetch.py b/src/data_generation/userguide/fetch.py new file mode 100644 index 0000000..c57f85a --- /dev/null +++ b/src/data_generation/userguide/fetch.py @@ -0,0 +1,54 @@ +import time +from collections.abc import Sequence +from pathlib import Path +from urllib.parse import urlparse + +import requests + +USER_AGENT = "ReactomeChatbot/1.0 (+https://github.com/reactome/reactome_chatbot)" +REQUEST_DELAY_SECONDS = 0.5 + + +def url_to_slug(url: str) -> str: + """Derive a filesystem-safe slug from a user guide URL path.""" + path = urlparse(url).path.strip("/") + return path.replace("/", "_") if path else "index" + + +def fetch_userguide_pages( + urls: Sequence[str], + cache_dir: Path, + *, + force: bool = False, +) -> dict[str, Path]: + """Download user guide HTML pages, using on-disk cache when available. + + Args: + urls: Canonical user guide URLs to download. + cache_dir: Directory for cached ``.html`` files. + force: When ``True``, re-download pages even if cached. + + Returns: + Mapping of each URL to its local cached HTML path. + """ + cache_dir = Path(cache_dir) + cache_dir.mkdir(parents=True, exist_ok=True) + session = requests.Session() + session.headers["User-Agent"] = USER_AGENT + + html_paths: dict[str, Path] = {} + for i, url in enumerate(urls): + cache_path = cache_dir / f"{url_to_slug(url)}.html" + if cache_path.exists() and not force: + html_paths[url] = cache_path + continue + + response = session.get(url, timeout=60) + response.raise_for_status() + cache_path.write_text(response.text, encoding=response.encoding or "utf-8") + html_paths[url] = cache_path + + if i < len(urls) - 1: + time.sleep(REQUEST_DELAY_SECONDS) + + return html_paths diff --git a/src/data_generation/userguide/html_loader.py b/src/data_generation/userguide/html_loader.py new file mode 100644 index 0000000..28ecc81 --- /dev/null +++ b/src/data_generation/userguide/html_loader.py @@ -0,0 +1,238 @@ +import re +from pathlib import Path + +from bs4 import BeautifulSoup, Tag +from langchain_community.document_loaders.base import BaseLoader +from langchain_core.documents import Document +from langchain_text_splitters import RecursiveCharacterTextSplitter + +SPLIT_CANDIDATES = ("h2", "h3", "h4") +MIN_SECTION_HEADINGS = 2 +MAX_CHUNK_CHARS = 4000 +MIN_CHUNK_CHARS = 80 +CHUNK_OVERLAP = 200 + +SPAMBOT_PATTERN = re.compile( + r"This email address is being protected from spambots.*", + re.IGNORECASE | re.DOTALL, +) +COLLECTIBLE_TAGS = frozenset( + { + "p", + "ul", + "ol", + "table", + "blockquote", + "pre", + "div", + "dl", + "h3", + "h4", + "h5", + "h6", + } +) + + +def choose_split_tag(article_body: Tag) -> str | None: + """Pick the shallowest heading level with enough sections for this page.""" + for tag in SPLIT_CANDIDATES: + if len(article_body.find_all(tag)) >= MIN_SECTION_HEADINGS: + return tag + return None + + +def _heading_title(heading: Tag) -> str: + title = heading.get_text(separator=" ", strip=True) + if title: + return title + image = heading.find("img", alt=True) + if image is not None: + alt_attr = image.get("alt") + if isinstance(alt_attr, str): + alt = alt_attr.strip() + if alt: + return alt + return "" + + +def split_article_body_into_sections( + article_body: Tag, +) -> list[tuple[str, int, list[Tag]]]: + """Split article body into titled sections using adaptive heading boundaries.""" + split_tag = choose_split_tag(article_body) + if split_tag is None: + return [("Introduction", 0, _collect_all_content(article_body))] + + headings = article_body.find_all(split_tag) + sections: list[tuple[str, int, list[Tag]]] = [] + + intro_nodes = _collect_intro(article_body, headings[0]) + if intro_nodes: + sections.append(("Introduction", 0, intro_nodes)) + + for index, heading in enumerate(headings): + title = _heading_title(heading) + level = int(split_tag[1]) + next_heading = headings[index + 1] if index + 1 < len(headings) else None + nodes = _collect_between(heading, next_heading, split_tag) + if not title and not nodes: + continue + sections.append((title or "Untitled", level, nodes)) + + return sections + + +def _should_collect_block(element: Tag, collected: list[Tag]) -> bool: + if element.name not in COLLECTIBLE_TAGS: + return False + if not element.get_text(strip=True): + return False + for other in collected: + if other in element.parents or element in other.parents: + return False + return True + + +def _collect_intro(article_body: Tag, first_heading: Tag) -> list[Tag]: + collected: list[Tag] = [] + for element in article_body.descendants: + if element is first_heading: + break + if isinstance(element, Tag) and _should_collect_block(element, collected): + collected.append(element) + return collected + + +def _collect_between( + start_heading: Tag, + end_heading: Tag | None, + split_tag: str, +) -> list[Tag]: + collected: list[Tag] = [] + for element in start_heading.next_elements: + if end_heading is not None and element is end_heading: + break + if isinstance(element, Tag) and element.name == split_tag: + break + if isinstance(element, Tag) and _should_collect_block(element, collected): + collected.append(element) + return collected + + +def _collect_all_content(article_body: Tag) -> list[Tag]: + collected: list[Tag] = [] + for element in article_body.descendants: + if isinstance(element, Tag) and _should_collect_block(element, collected): + collected.append(element) + return collected + + +class UserGuideHTMLLoader(BaseLoader): + """Loads Reactome user guide HTML pages into section-level documents. + + Each document represents one section of a user guide page. The loader picks + the shallowest heading level (``h2``, ``h3``, or ``h4``) that yields at + least two sections on that page. Deeper headings remain within their parent + section. Oversized sections are split for embedding. + + The ``source`` metadata field is set to the canonical page URL. Section + titles and page titles are included in both metadata and ``page_content``. + + Output Example: + .. code-block:: txt + + Page: Pathway Browser + Section: Event Hierarchy + + The order of reactions from top to bottom... + """ + + def __init__(self, html_paths: dict[str, Path]) -> None: + """ + Args: + html_paths: Mapping of canonical page URLs to local HTML file paths. + """ + self.html_paths = html_paths + self._splitter = RecursiveCharacterTextSplitter( + chunk_size=MAX_CHUNK_CHARS, + chunk_overlap=CHUNK_OVERLAP, + ) + + def load(self) -> list[Document]: + """Load data into document objects.""" + documents: list[Document] = [] + for url, path in self.html_paths.items(): + documents.extend(self._load_page(url, path)) + return documents + + def _load_page(self, url: str, path: Path) -> list[Document]: + html = path.read_text(encoding="utf-8") + soup = BeautifulSoup(html, "lxml") + page_title = self._extract_page_title(soup) + article_body = soup.select_one('[itemprop="articleBody"]') + if article_body is None: + raise ValueError(f"No article body found for {url}") + + sections = split_article_body_into_sections(article_body) + + documents: list[Document] = [] + for section_title, section_level, nodes in sections: + text = self._nodes_to_text(nodes) + text = SPAMBOT_PATTERN.sub("", text).strip() + if section_level == 0 and len(text) < MIN_CHUNK_CHARS: + continue + if section_level > 0 and len(text) < MIN_CHUNK_CHARS: + text = f"{section_title}\n\n{text}".strip() if text else section_title + + page_content_prefix = ( + f"URL: {url}\nPage: {page_title}\nSection: {section_title}\n\n" + ) + chunks = self._splitter.split_text(text) + for chunk_index, chunk in enumerate(chunks): + documents.append( + Document( + page_content=page_content_prefix + chunk, + metadata={ + "source": url, + "page_title": page_title, + "section_title": section_title, + "section_level": str(section_level), + "chunk_index": str(chunk_index), + }, + ) + ) + return documents + + def _extract_page_title(self, soup: BeautifulSoup) -> str: + header = soup.select_one(".page-header h2") + if header: + title = header.get_text(strip=True) + if title: + return title + if soup.title and soup.title.string: + return soup.title.string.replace(" - Reactome Pathway Database", "").strip() + return "Unknown" + + def _nodes_to_text(self, nodes: list[Tag]) -> str: + parts: list[str] = [] + for node in nodes: + if node.name == "ul": + for li in node.find_all("li", recursive=False): + item = li.get_text(separator=" ", strip=True) + if item: + parts.append(f"- {item}") + elif node.name == "ol": + for i, li in enumerate(node.find_all("li", recursive=False), start=1): + item = li.get_text(separator=" ", strip=True) + if item: + parts.append(f"{i}. {item}") + elif node.name == "table": + text = node.get_text(separator=" ", strip=True) + if text: + parts.append(text) + else: + text = node.get_text(separator="\n", strip=True) + if text: + parts.append(text) + return "\n\n".join(parts) diff --git a/src/data_generation/userguide/urls.py b/src/data_generation/userguide/urls.py new file mode 100644 index 0000000..94970e5 --- /dev/null +++ b/src/data_generation/userguide/urls.py @@ -0,0 +1,16 @@ +"""Canonical Reactome user guide URLs for ingestion.""" + +REACTOME_BASE = "https://reactome.org" + +USER_GUIDE_URLS: tuple[str, ...] = ( + f"{REACTOME_BASE}/userguide", + f"{REACTOME_BASE}/userguide/pathway-browser", + f"{REACTOME_BASE}/userguide/searching", + f"{REACTOME_BASE}/userguide/details-panel", + f"{REACTOME_BASE}/userguide/analysis", + f"{REACTOME_BASE}/userguide/analysis/gsa", + f"{REACTOME_BASE}/userguide/diseases", + f"{REACTOME_BASE}/userguide/cytomics", + f"{REACTOME_BASE}/userguide/review-status", + f"{REACTOME_BASE}/userguide/reactome-fiviz", +) diff --git a/src/evaluation/README.md b/src/evaluation/README.md new file mode 100644 index 0000000..36fded4 --- /dev/null +++ b/src/evaluation/README.md @@ -0,0 +1,46 @@ +# RAGAS Evaluation Toolkit + +This folder contains utility scripts used to benchmark Reactome RAG pipelines with Ragas. + +- `test_generator.py` synthesizes question/answer test sets from the example corpora. It uses the Ragas `TestsetGenerator` to create LangChain-based evaluation datasets. +- `evaluator.py` runs either the basic or advanced Reactome RAG chain over a test set and scores the outputs with Ragas metrics (answer relevancy, context utilization, faithfulness, context recall), saving both responses and evaluation reports. + +## Requirements + +- Python 3.12 (project default) with Poetry environment +- `ragas` (see poetry.lock for the pinned version) +- OpenAI access: `OPENAI_API_KEY` (and optional Azure configuration if required) +- An installed embeddings bundle (`./bin/embeddings_manager install ...`). + `evaluator.py` defaults to whichever bundle is active; override with + `--embeddings-dir`. + +Run `poetry install` to set up dependencies, then activate the virtual environment via `poetry shell` or use `poetry run` for individual commands. + +## Usage + +1. **Generate test sets** + + ```bash + poetry run python src/evaluation/test_generator.py \ + --path src/evaluation/example \ + --model gpt-4o-mini \ + --temperature 0.3 \ + --test_size 10 \ + --distributions simple=0.25 reasoning=0.25 multi_context=0.25 conditional=0.25 + ``` + + Outputs are stored in a `testsets/` directory of your choosing. + +2. **Evaluate a RAG configuration** + + ```bash + poetry run python src/evaluation/evaluator.py \ + --testset_dir \ + --rag_type advanced \ + --model gpt-4o-mini + ``` + + Responses and metric reports are written to `response//` and `evals//` inside the testset directory. + +Adjust the paths, model names, and distribution weights as needed for local experimentation. + diff --git a/src/evaluation/evaluator.py b/src/evaluation/evaluator.py index 9d4d6da..b4ecf35 100644 --- a/src/evaluation/evaluator.py +++ b/src/evaluation/evaluator.py @@ -1,5 +1,7 @@ import argparse import os +from pathlib import Path +from typing import Any import pandas as pd from datasets import Dataset @@ -9,20 +11,29 @@ from langchain_chroma.vectorstores import Chroma from langchain_community.document_loaders.csv_loader import CSVLoader from langchain_community.retrievers import BM25Retriever +from langchain_core.retrievers import BaseRetriever +from langchain_core.runnables import Runnable from langchain_openai import ChatOpenAI, OpenAIEmbeddings from ragas import evaluate -from ragas.metrics import (ContextUtilization, answer_relevancy, - context_recall, faithfulness) +from ragas.metrics import ( + ContextUtilization, + answer_relevancy, + context_recall, + faithfulness, +) from retrievers.rag_chain import create_rag_chain -from retrievers.reactome.metadata_info import (reactome_descriptions_info, - reactome_field_info) +from retrievers.reactome.metadata_info import ( + reactome_descriptions_info, + reactome_field_info, +) from retrievers.reactome.prompt import reactome_qa_prompt +from util.embedding_environment import EmbeddingEnvironment context_utilization = ContextUtilization() -def parse_arguments(): +def parse_arguments() -> argparse.Namespace: """Parse command line arguments for the script.""" parser = argparse.ArgumentParser( description="Load a directory of testsets and evaluate answers generated by a language model." @@ -33,6 +44,15 @@ def parse_arguments(): required=True, help="Path to the directory containing testset Excel (.xlsx) files", ) + parser.add_argument( + "--embeddings-dir", + type=Path, + default=EmbeddingEnvironment.get_dir("reactome"), + help=( + "Reactome embeddings bundle to evaluate against. Defaults to the " + "installed one (see ./bin/embeddings_manager which)." + ), + ) parser.add_argument( "--model", type=str, @@ -48,27 +68,34 @@ def parse_arguments(): return parser.parse_args() -def load_dataset(testset_path): +def load_dataset(testset_path: str) -> list[dict[str, Any]]: """Load the dataset from an Excel (.xlsx) file.""" + # pandas types record keys as Hashable; they are column names. try: df = pd.read_excel(testset_path) - return df.to_dict( - orient="records" - ) # Convert DataFrame to a list of dictionaries - except FileNotFoundError: - raise FileNotFoundError(f"The file {testset_path} does not exist.") + records: list[dict[str, Any]] = df.to_dict(orient="records") # type: ignore[assignment] + return records + except FileNotFoundError as e: + raise FileNotFoundError(f"The file {testset_path} does not exist.") from e except ValueError as e: - raise ValueError(f"Error reading the Excel file: {e}") + raise ValueError(f"Error reading the Excel file: {e}") from e + +def initialize_rag_chain_with_memory( + embeddings_dir: Path, model_name: str, rag_type: str +) -> Runnable: + """Initialize the RAGChainWithMemory system. -def initialize_rag_chain_with_memory(embeddings_directory, model_name, rag_type): - """Initialize the RAGChainWithMemory system.""" + `embeddings_dir` is the bundle root, e.g. + embeddings/openai/text-embedding-3-large/reactome/Release95 . Both the BM25 + source CSV and the Chroma collection are derived from it; they used to be + absolute paths into a developer's home directory, so this script could not + run anywhere else. + """ llm = ChatOpenAI(temperature=0.0, verbose=True, model=model_name) - retriever_list = [] + retriever_list: list[BaseRetriever] = [] - loader = CSVLoader( - "/Users/hmohammadi/Desktop/react_to_me_github/reactome_chatbot/embeddings/openai/text-embedding-3-large/reactome/summation_csv/summations.csv" - ) + loader = CSVLoader(str(embeddings_dir / "csv_files" / "summations.csv")) data = loader.load() bm25_retriever = BM25Retriever.from_documents(data) bm25_retriever.k = 7 @@ -76,7 +103,7 @@ def initialize_rag_chain_with_memory(embeddings_directory, model_name, rag_type) # Set up vectorstore SelfQuery retriever embedding = OpenAIEmbeddings(model="text-embedding-3-large") vectordb = Chroma( - persist_directory=embeddings_directory, + persist_directory=str(embeddings_dir / "summations"), embedding_function=embedding, ) @@ -99,23 +126,21 @@ def initialize_rag_chain_with_memory(embeddings_directory, model_name, rag_type) reactome_retriever = MergerRetriever(retrievers=retriever_list) - qa = create_rag_chain( + return create_rag_chain( retriever=reactome_retriever, llm=llm, qa_prompt=reactome_qa_prompt, ) - return qa def process_testset( - testset_path, - qa_system, - embeddings_directory, - response_dir, - eval_dir, - model_name, - rag_type, -): + testset_path: str, + qa_system: Runnable, + response_dir: str, + eval_dir: str, + model_name: str, + rag_type: str, +) -> None: """Process a single testset file.""" testset = load_dataset(testset_path) questions = [item["question"] for item in testset] @@ -125,7 +150,7 @@ def process_testset( contexts = [] for question in questions: - response = qa_system.get_context(question) + response = qa_system.invoke({"input": question}) answers.append(response["answer"]) contexts.append([context.page_content for context in response["context"]]) @@ -167,7 +192,7 @@ def process_testset( print(f"Evaluation results saved to {evaluation_filename}") -def main(): +def main() -> None: args = parse_arguments() model_name = args.model rag_type = args.rag_type @@ -177,10 +202,13 @@ def main(): os.makedirs(eval_dir, exist_ok=True) # Initialize RAG Chain - embeddings_directory = "/Users/hmohammadi/Desktop/react_to_me_github/reactome_chatbot/embeddings/openai/text-embedding-3-large/reactome/Release90/summations" - qa_system = initialize_rag_chain_with_memory( - embeddings_directory, model_name, rag_type - ) + embeddings_dir: Path | None = args.embeddings_dir + if embeddings_dir is None: + raise SystemExit( + "No reactome embeddings installed and --embeddings-dir not given. " + "Install one with ./bin/embeddings_manager install ." + ) + qa_system = initialize_rag_chain_with_memory(embeddings_dir, model_name, rag_type) # Iterate over all .xlsx files in the directory for filename in os.listdir(args.testset_dir): @@ -191,7 +219,6 @@ def main(): process_testset( testset_path, qa_system, - embeddings_directory, response_dir, eval_dir, model_name, diff --git a/src/evaluation/test_generator.py b/src/evaluation/test_generator.py index b7ea836..bd954ec 100644 --- a/src/evaluation/test_generator.py +++ b/src/evaluation/test_generator.py @@ -1,5 +1,6 @@ import argparse import os +from typing import Any import pandas as pd from dotenv import load_dotenv @@ -8,7 +9,7 @@ from ragas.testset.synthesizers.generate import TestsetGenerator -def parse_arguments(): +def parse_arguments() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Generate a test set based on given documents and distributions." ) @@ -46,7 +47,7 @@ def parse_arguments(): return parser.parse_args() -def save_testset(testset, filename): +def save_testset(testset: Any, filename: str) -> None: output_dir = "testsets" os.makedirs(output_dir, exist_ok=True) output_path = os.path.join(output_dir, f"{filename}_testset.xlsx") @@ -65,7 +66,7 @@ def save_testset(testset, filename): print(f"Filtered testset saved to {output_path}") -def main(): +def main() -> None: try: # Load environment variables load_dotenv() @@ -82,10 +83,22 @@ def main(): embeddings = OpenAIEmbeddings() # Parse test type distributions - distributions = { - eval(key): float(value) - for key, value in (dist.split("=") for dist in args.distributions) - } + # TODO: this targets the ragas 0.1 API. `distributions` took evolution + # objects (simple, reasoning, multi_context, conditional) imported from + # ragas.testset.evolutions, which 0.2 replaced with synthesizers, and the + # generate_with_langchain_docs keyword changed too. The names were + # previously resolved with eval() on argv -- removed here because that + # executes arbitrary input, and because those names are not imported, so + # it raised NameError rather than working. Porting this to 0.2 needs the + # current ragas API. + distributions: dict[str, float] = {} + for dist in args.distributions: + key, sep, value = dist.partition("=") + if not sep: + raise SystemExit( + f"Malformed --distributions entry {dist!r}; use name=weight" + ) + distributions[key] = float(value) # Setup and run test set generator generator = TestsetGenerator.from_langchain( diff --git a/src/retrievers/csv_chroma.py b/src/retrievers/csv_chroma.py index a792c93..43b5dba 100644 --- a/src/retrievers/csv_chroma.py +++ b/src/retrievers/csv_chroma.py @@ -1,15 +1,19 @@ import asyncio +from collections.abc import Coroutine from pathlib import Path -from typing import Annotated, Any, Coroutine, TypedDict +from typing import Annotated, Any, TypedDict import chromadb.config from langchain.chains.query_constructor.schema import AttributeInfo from langchain.retrievers import EnsembleRetriever, MultiQueryRetriever -from langchain.retrievers.merger_retriever import MergerRetriever from langchain.retrievers.self_query.base import SelfQueryRetriever from langchain_chroma.vectorstores import Chroma from langchain_community.document_loaders.csv_loader import CSVLoader from langchain_community.retrievers import BM25Retriever +from langchain_core.callbacks import ( + AsyncCallbackManagerForRetrieverRun, + CallbackManagerForRetrieverRun, +) from langchain_core.documents import Document from langchain_core.embeddings import Embeddings from langchain_core.language_models.chat_models import BaseChatModel @@ -57,11 +61,56 @@ ] +RESULTS_PER_RETRIEVER = 10 +# The vector store is asked for more than we intend to keep, because one +# Reactome entity can occupy several rows -- a reaction appears once per +# pathway/input/output/catalyst combination -- and those rows have distinct +# page_content, so nothing upstream collapses them. Without over-fetching, a +# request for 10 returns about 5 distinct reactions. See issue #169. +VECTOR_OVERFETCH = 3 + +# How many fused documents each collection contributes to the answer prompt. +# +# weighted_reciprocal_rank returns *every* unique document across the lists it is +# given, not a top-N, so without this the retriever ranked ~222 documents by +# relevance and then sent all of them -- roughly 32k tokens, a quarter of +# gpt-4o-mini's window, on every message -- which made the ranking decorative. +# +# The cap is per collection rather than global on purpose: reactions, summations, +# complexes and ewas hold different kinds of information, and one global top-N +# would let a single collection crowd the others out. Per collection guarantees +# each one contributes. +# +# This value is a starting point, not a tuned one. It matches what a single +# retriever returns. Changing it trades recall against the model's difficulty +# attending to the middle of a long context; the right number should come from an +# answer-quality evaluation rather than from taste. +MAX_DOCUMENTS_PER_COLLECTION = RESULTS_PER_RETRIEVER + + +def dedupe_by_entity(docs: list[Document], limit: int) -> list[Document]: + """Keep the highest-ranked row per Reactome stable ID, up to `limit`. + + Falls back to page_content for documents with no st_id, so a collection + without that metadata degrades to the previous behaviour rather than raising. + """ + seen: set[str] = set() + kept: list[Document] = [] + for doc in docs: + key = str(doc.metadata.get("st_id") or doc.page_content) + if key in seen: + continue + seen.add(key) + kept.append(doc) + if len(kept) == limit: + break + return kept + + def list_chroma_subdirectories(directory: Path) -> list[str]: - subdirectories = list( + return [ chroma_file.parent.name for chroma_file in directory.glob("*/chroma.sqlite3") - ) - return subdirectories + ] def create_bm25_chroma_ensemble_retriever( @@ -71,7 +120,7 @@ def create_bm25_chroma_ensemble_retriever( *, descriptions_info: dict[str, str], field_info: dict[str, list[AttributeInfo]], -) -> MergerRetriever: +) -> "HybridRetriever": return HybridRetriever.from_subdirectory( llm, embedding, @@ -100,8 +149,8 @@ def from_subdirectory( *, descriptions_info: dict[str, str], field_info: dict[str, list[AttributeInfo]], - include_original=False, - ): + include_original: bool = False, + ) -> "HybridRetriever": _retrievers: dict[str, RetrieverDict] = {} for subdirectory in list_chroma_subdirectories(embeddings_directory): # set up BM25 retriever @@ -115,7 +164,7 @@ def from_subdirectory( text.casefold(), language="english" ), ) - bm25_retriever.k = 10 + bm25_retriever.k = RESULTS_PER_RETRIEVER # set up vectorstore SelfQuery retriever vectordb = Chroma( @@ -129,7 +178,7 @@ def from_subdirectory( vectorstore=vectordb, document_contents=descriptions_info[subdirectory], metadata_field_info=field_info[subdirectory], - search_kwargs={"k": 10}, + search_kwargs={"k": RESULTS_PER_RETRIEVER * VECTOR_OVERFETCH}, ) _retrievers[subdirectory] = { @@ -154,7 +203,9 @@ def weighted_reciprocal_rank( retrievers=[], weights=[1 / len(doc_lists)] * len(doc_lists) ).weighted_reciprocal_rank(doc_lists) - def retrieve_documents(self, queries: list[str], run_manager) -> list[Document]: + def retrieve_documents( + self, queries: list[str], run_manager: CallbackManagerForRetrieverRun + ) -> list[Document]: subdirectory_docs: list[Document] = [] for subdirectory, retrievers in self._retrievers.items(): bm25_retriever = retrievers["bm25"] @@ -177,12 +228,22 @@ def retrieve_documents(self, queries: list[str], run_manager) -> list[Document]: ) }, ) - doc_lists.append(bm25_docs + vector_docs) - subdirectory_docs.extend(self.weighted_reciprocal_rank(doc_lists)) + # Separate lists, not `bm25_docs + vector_docs`. RRF scores by + # position, so concatenating put every vector result at rank 11+ + # and scored the best of them 1/71 against BM25's 1/61 -- and it + # meant the two retrievers were never fused against each other, + # only across query variants. See issue #170. + doc_lists.append(dedupe_by_entity(bm25_docs, RESULTS_PER_RETRIEVER)) + doc_lists.append(dedupe_by_entity(vector_docs, RESULTS_PER_RETRIEVER)) + subdirectory_docs.extend( + self.weighted_reciprocal_rank(doc_lists)[:MAX_DOCUMENTS_PER_COLLECTION] + ) return subdirectory_docs async def aretrieve_documents( - self, queries: list[str], run_manager + self, + queries: list[str], + run_manager: AsyncCallbackManagerForRetrieverRun, ) -> list[Document]: subdirectory_results: dict[str, list[Coroutine[Any, Any, list[Document]]]] = {} for subdirectory, retrievers in self._retrievers.items(): @@ -213,10 +274,13 @@ async def aretrieve_documents( ) subdirectory_docs: list[Document] = [] for subdir_results in subdirectory_results.values(): - results_iter = iter(await asyncio.gather(*subdir_results)) + # Separate lists, de-duplicated per entity, matching the synchronous + # path above. See issues #169 and #170. doc_lists: list[list[Document]] = [ - bm25_results + vector_results - for bm25_results, vector_results in zip(results_iter, results_iter) + dedupe_by_entity(docs, RESULTS_PER_RETRIEVER) + for docs in await asyncio.gather(*subdir_results) ] - subdirectory_docs.extend(self.weighted_reciprocal_rank(doc_lists)) + subdirectory_docs.extend( + self.weighted_reciprocal_rank(doc_lists)[:MAX_DOCUMENTS_PER_COLLECTION] + ) return subdirectory_docs diff --git a/src/retrievers/plantreactome/metadata_info.py b/src/retrievers/plantreactome/metadata_info.py new file mode 100644 index 0000000..3ff6b24 --- /dev/null +++ b/src/retrievers/plantreactome/metadata_info.py @@ -0,0 +1,150 @@ +from langchain.chains.query_constructor.base import AttributeInfo + +pathway_id_description = "A Plant Reactome Identifier unique to each pathway. A pathway name may appear multiple times in the dataset\ + This ID allows for the specific identification and exploration of each pathway's details within the Plant Reactome Database." +pathway_name_description = "The name of the biological pathway, indicating a specific series of interactions or processes within a cell.\ + A pathway name may appear multiple times in the dataset, reflecting the fact that several reactions (identified by 'reaction_name') contribute to a single pathway.\ + The relationship between 'reaction_name' and 'pathway_name' is foundational, with each reaction serving as a step or component within the overarching pathway, contributing to its completion and functional outcome.\ + This relationship is critical to understanding the biological processes and mechanisms within the Plant Reactome Database." + +plantreactome_descriptions_info: dict[str, str] = { + "ewas": "Contains data on proteins and nucleic acids with known sequences. Includes entity names, IDs, canonical and synonymous gene names, and functions.", + "complexes": "Catalogs biological complexes, listing complex names and IDs along with the names and IDs of their components. ", + "reactions": "Documents biological pathways and their constituent reactions, detailing pathway and reaction names and IDs. It includes information on the inputs, outputs, and catalysts for each reaction, emphasizing the interconnected nature of cellular processes. Inputs and outputs, critical to the initiation and conclusion of reactions, along with catalysts that facilitate these processes, are cataloged to highlight their roles across various reactions and pathways", + "summations": "Enumerates biological reactions, accompanied by concise summaries ('summations') of each reaction. These summations encapsulate the essence and biochemical significance of the reactions, offering insights into their roles within cellular processes and pathways.", +} + + +plantreactome_field_info: dict[str, list[AttributeInfo]] = { + "summations": [ + AttributeInfo( + name="st_id", + description=pathway_id_description, + type="string", + ), + AttributeInfo( + name="display_name", + description=pathway_name_description, + type="string", + ), + ], + "reactions": [ + AttributeInfo( + name="st_id", + description="The Reactome Identifier (ID) for each biological reaction, serving as a unique key.\ + This ID allows for the specific identification and exploration of each reaction's details within the Reactome Database.", + type="string", + ), + AttributeInfo( + name="display_name", + description="The name of the biological reaction, encapsulating the interaction between proteins or molecules.\ + Each reaction name is a unique entry, reflecting a specific biological process.\ + These names provide insight into the dynamic processes within cellular functions, highlighting the roles of various proteins and molecules in biological mechanisms", + type="string", + ), + AttributeInfo( + name="pathway_id", + description=pathway_id_description, + type="string", + ), + AttributeInfo( + name="pathway_name", + description=pathway_name_description, + type="string", + ), + AttributeInfo( + name="input_id", + description="The Reactome Identifier (ID) for each input.\ + Given that a single input can be involved in various reactionss, this ID may repeat across multiple rows, each associated with a different reaction.\ + This ID allows for the specific identification and exploration of each input's details within the Reactome Database.", + type="string", + ), + AttributeInfo( + name="input_name", + description="Identifies the inputs of a biological reaction ('reaction_name'), which can be either entities or part of complexes.\ + Inputs are crucial for initiating reactions, acting as the reactants that drive the biochemical processes. \ + Given their fundamental role, inputs may repeat across multiple reactions, reflecting their involvement in various parts of the cellular machinery.", + type="string", + ), + AttributeInfo( + name="output_id", + description=" A Reactome Identifier unique to each output of a reaction.\ + Given that a single input can be involved in various reactionss, this ID may repeat across multiple rows, each associated with a different reaction.\ + This ID allows for the specific identification and exploration of each output's details within the Reactome Database.", + type="string", + ), + AttributeInfo( + name="output_name", + description="Represents the outputs of a biological reaction ('reaction_name'), denoting the products generated as a result of the biochemical interactions. \ + Outputs can be entities or complexes and may appear in multiple reactions, highlighting their multifunctional role in cellular pathways. \ + This repetition underscores the interconnected nature of biological processes, where one reaction's output can serve as another's input.", + type="string", + ), + AttributeInfo( + name="catalyst_id", + description="The Reactome Identifier (ID) for each biological catalyst, serving as a unique key.\ + Given that a single catalyst can be involved in various reactions, this ID may repeat across multiple rows, each associated with a different reaction.\ + This ID allows for the specific identification and exploration of each catalyst's details within the Reactome Database.", + type="string", + ), + AttributeInfo( + name="catalyst_name", + description="Specifies the catalysts that facilitate a biological reaction, potentially speeding up the process without being consumed.\ + Catalysts are crucial for modulating reaction rates and guiding the direction of the reaction, ensuring the efficient progression of biological pathways.\ + Catalysts can be proteins, enzymes, or molecular compounds, underscoring their vital role in cellular operations.", + type="string", + ), + ], + "complexes": [ + AttributeInfo( + name="st_id", + description="The Reactome Identifier (ID) for each biological complex, serving as a unique key.\ + Given that a single complex can consist of various components, this ID may repeat across multiple rows, each associated with a different component of the same complex.", + type="string", + ), + AttributeInfo( + name="display_name", + description="The name of the biological complex.\ + This field provides a reference to the complex itself, which may be listed across several rows to account for its multiple components.", + type="string", + ), + AttributeInfo( + name="component_id", + description=" A Reactome Identifier unique to each component within a complex.\ + This ID allows for the specific identification and exploration of each component's details within the Reactome Database.", + type="string", + ), + AttributeInfo( + name="component_name", + description="The name of the individual component associated with the complex in that row.\ + This reveals the specific protein or molecule constituting part of the complex, emphasizing the diversity of components within a single biological entity.", + type="string", + ), + ], + "ewas": [ + AttributeInfo( + name="st_id", + description="The Reactome Identifier (ID) for each biological complex, serving as a unique key.\ + Given that a single complex can consist of various components, this ID may repeat across multiple rows, each associated with a different component of the same complex.", + type="string", + ), + AttributeInfo( + name="display_name", + description="The name of the biological complex.\ + This field provides a reference to the complex itself, which may be listed across several rows to account for its multiple components.", + type="string", + ), + AttributeInfo( + name="canonical_geneName", + description=" A Reactome Identifier unique to each component within a complex.\ + This ID allows for the specific identification and exploration of each component's details within the Reactome Database.", + type="string", + ), + AttributeInfo( + name="synonyms_geneName", + description="The name of the individual component associated with the complex in that row.\ + This reveals the specific protein or molecule constituting part of the complex, emphasizing the diversity of components within a single biological entity.", + type="string", + ), + ], +} diff --git a/src/retrievers/plantreactome/prompt.py b/src/retrievers/plantreactome/prompt.py new file mode 100644 index 0000000..50fcd1f --- /dev/null +++ b/src/retrievers/plantreactome/prompt.py @@ -0,0 +1,40 @@ +from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder + +plantreactome_system_prompt = """ +You are an expert in molecular biology with access to the **Plant Reactome Knowledgebase**. +Your primary responsibility is to answer the user's questions **comprehensively, mechanistically, and with precision**, drawing strictly from the **Plant Reactome Knowledgebase**. + +Your output must emphasize biological processes, molecular complexes, regulatory mechanisms, and interactions most relevant to the user’s question. +Provide an information-rich narrative that explains not only what is happening but also how and why, based only on PlantReactome context. + + +## **Answering Guidelines** +1. Strict source discipline: Use only the information explicitly provided from Plant Reactome. Do not invent, infer, or draw from external knowledge. + - Use only information directly found in Plant Reactome. + - Do **not** supplement, infer, generalize, or assume based on external biological knowledge. + - If no relevant information exists in Plant Reactome, explain the information is not currently available in Plant Reactome. Do **not** answer the question. +2. Inline citations required: Every factual statement must include ≥1 inline anchor citation in the format: display_name + - If multiple entries support the same fact, cite them together (space-separated). +3. Comprehensiveness: Capture all mechanistically relevant details available in PlantReactome, focusing on processes, complexes, regulations, and interactions. +4. Tone & Style: + - Write in a clear, engaging, and conversational tone. + - Use accessible language while maintaining technical precision. + - Ensure the narrative flows logically, presenting background, mechanisms, and significance +5. Source list at the end: After the main narrative, provide a bullet-point list of each unique citation anchor exactly once, in the same Node Name format. + - Examples: + - Mitosis + - Cell Cycle + +## Internal QA (silent) +- All factual claims are cited correctly. +- No unverified claims or background knowledge are added. +- The Sources list is complete and de-duplicated. +""" + +plantreactome_qa_prompt = ChatPromptTemplate.from_messages( + [ + ("system", plantreactome_system_prompt), + MessagesPlaceholder(variable_name="chat_history"), + ("user", "Context:\n{context}\n\nQuestion: {input}"), + ] +) diff --git a/src/retrievers/plantreactome/rag.py b/src/retrievers/plantreactome/rag.py new file mode 100644 index 0000000..22a9537 --- /dev/null +++ b/src/retrievers/plantreactome/rag.py @@ -0,0 +1,37 @@ +from pathlib import Path + +from langchain_core.embeddings import Embeddings +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.runnables import Runnable + +from retrievers.csv_chroma import create_bm25_chroma_ensemble_retriever +from retrievers.plantreactome.metadata_info import ( + plantreactome_descriptions_info, + plantreactome_field_info, +) +from retrievers.plantreactome.prompt import plantreactome_qa_prompt +from retrievers.rag_chain import create_rag_chain +from util.embedding_environment import EmbeddingEnvironment + + +def create_plantreactome_rag( + llm: BaseChatModel, + embedding: Embeddings, + # TODO(phase-2): resolved at import time, so importing this module requires an + # installed embeddings bundle. Blocks unit-testing; fix with the agent-API refactor. + embeddings_directory: Path = EmbeddingEnvironment.get_dir("plantreactome"), # noqa: B008 + *, + streaming: bool = False, +) -> Runnable: + plantreactome_retriever = create_bm25_chroma_ensemble_retriever( + llm, + embedding, + embeddings_directory, + descriptions_info=plantreactome_descriptions_info, + field_info=plantreactome_field_info, + ) + + if streaming: + llm = llm.model_copy(update={"streaming": True}) + + return create_rag_chain(llm, plantreactome_retriever, plantreactome_qa_prompt) diff --git a/src/retrievers/rag_chain.py b/src/retrievers/rag_chain.py index 3e5df8e..11c0051 100644 --- a/src/retrievers/rag_chain.py +++ b/src/retrievers/rag_chain.py @@ -1,7 +1,7 @@ from langchain.chains.combine_documents import create_stuff_documents_chain from langchain.chains.retrieval import create_retrieval_chain from langchain_core.language_models.chat_models import BaseChatModel -from langchain_core.prompts import ChatPromptTemplate +from langchain_core.prompts import BasePromptTemplate, ChatPromptTemplate from langchain_core.retrievers import BaseRetriever from langchain_core.runnables import Runnable @@ -10,11 +10,14 @@ def create_rag_chain( llm: BaseChatModel, retriever: BaseRetriever, qa_prompt: ChatPromptTemplate, + *, + document_prompt: BasePromptTemplate | None = None, ) -> Runnable: # Create the documents chain question_answer_chain: Runnable = create_stuff_documents_chain( llm=llm, prompt=qa_prompt, + document_prompt=document_prompt, ) # Create the retrieval chain diff --git a/src/retrievers/reactome/rag.py b/src/retrievers/reactome/rag.py index 485b6e5..0611e27 100644 --- a/src/retrievers/reactome/rag.py +++ b/src/retrievers/reactome/rag.py @@ -6,8 +6,10 @@ from retrievers.csv_chroma import create_bm25_chroma_ensemble_retriever from retrievers.rag_chain import create_rag_chain -from retrievers.reactome.metadata_info import (reactome_descriptions_info, - reactome_field_info) +from retrievers.reactome.metadata_info import ( + reactome_descriptions_info, + reactome_field_info, +) from retrievers.reactome.prompt import reactome_qa_prompt from util.embedding_environment import EmbeddingEnvironment @@ -15,7 +17,9 @@ def create_reactome_rag( llm: BaseChatModel, embedding: Embeddings, - embeddings_directory: Path = EmbeddingEnvironment.get_dir("reactome"), + # TODO(phase-2): resolved at import time, so importing this module requires an + # installed embeddings bundle. Blocks unit-testing; fix with the agent-API refactor. + embeddings_directory: Path = EmbeddingEnvironment.get_dir("reactome"), # noqa: B008 *, streaming: bool = False, ) -> Runnable: diff --git a/src/retrievers/uniprot/rag.py b/src/retrievers/uniprot/rag.py index 99702d7..676d345 100644 --- a/src/retrievers/uniprot/rag.py +++ b/src/retrievers/uniprot/rag.py @@ -6,8 +6,10 @@ from retrievers.csv_chroma import create_bm25_chroma_ensemble_retriever from retrievers.rag_chain import create_rag_chain -from retrievers.uniprot.metadata_info import (uniprot_descriptions_info, - uniprot_field_info) +from retrievers.uniprot.metadata_info import ( + uniprot_descriptions_info, + uniprot_field_info, +) from retrievers.uniprot.prompt import uniprot_qa_prompt from util.embedding_environment import EmbeddingEnvironment @@ -15,7 +17,9 @@ def create_uniprot_rag( llm: BaseChatModel, embedding: Embeddings, - embeddings_directory: Path = EmbeddingEnvironment.get_dir("uniprot"), + # TODO(phase-2): resolved at import time, so importing this module requires an + # installed embeddings bundle. Blocks unit-testing; fix with the agent-API refactor. + embeddings_directory: Path = EmbeddingEnvironment.get_dir("uniprot"), # noqa: B008 *, streaming: bool = False, ) -> Runnable: diff --git a/src/retrievers/userguide/prompt.py b/src/retrievers/userguide/prompt.py new file mode 100644 index 0000000..cbb3efb --- /dev/null +++ b/src/retrievers/userguide/prompt.py @@ -0,0 +1,37 @@ +from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder + +userguide_system_prompt = """ +You are a helpful guide to the **Reactome website** and its tools. +Your primary responsibility is to answer questions about **how to use Reactome** — the Pathway Browser, search, analysis tools, Details Panel, and related features — using only the user guide excerpts provided in the context. + +## Answering Guidelines +1. Strict source discipline: Use only the information explicitly provided from the Reactome user guide. Do not invent steps, buttons, menus, or workflows. + - If the context does not contain enough information to answer, say the user guide does not currently cover that topic. Do **not** guess. +2. Inline citations required: Every factual statement must include ≥1 inline anchor citation in the format: display_name + - Use the **exact** URL from the context (the line starting with `URL:`). Copy it verbatim. + - Never guess, shorten, or construct URLs from page titles (for example, do not turn "ReactomeGSA" into `/userguide/reactomegsa`). + - Use a clear display name (page title or section title). + - If multiple excerpts support the same fact, cite them together (space-separated). +3. How-to focus: Give clear, actionable steps when the user asks how to perform a task. Name UI elements accurately (buttons, panels, tabs) as they appear in the context. +4. Tone and style: + - Write in a clear, friendly, and conversational tone. + - Use accessible language; avoid unnecessary jargon. + - Prefer numbered steps for multi-step procedures. +5. Source list at the end: After the main answer, provide a bullet-point list of each unique citation anchor exactly once, in the same display_name format. + - Examples: + - Pathway Browser + - Searching Reactome + +## Internal QA (silent) +- All factual claims are cited correctly. +- No UI steps or features are invented beyond the provided context. +- The Sources list is complete and de-duplicated. +""" + +userguide_qa_prompt = ChatPromptTemplate.from_messages( + [ + ("system", userguide_system_prompt), + MessagesPlaceholder(variable_name="chat_history"), + ("user", "Context:\n{context}\n\nQuestion: {input}"), + ] +) diff --git a/src/retrievers/userguide/rag.py b/src/retrievers/userguide/rag.py new file mode 100644 index 0000000..153910b --- /dev/null +++ b/src/retrievers/userguide/rag.py @@ -0,0 +1,27 @@ +from pathlib import Path + +from langchain_core.embeddings import Embeddings +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.runnables import Runnable + +from retrievers.rag_chain import create_rag_chain +from retrievers.userguide.prompt import userguide_qa_prompt +from retrievers.userguide.retriever import create_userguide_retriever +from util.embedding_environment import EmbeddingEnvironment + + +def create_userguide_rag( + llm: BaseChatModel, + embedding: Embeddings, + # TODO(phase-2): resolved at import time, so importing this module requires an + # installed embeddings bundle. Blocks unit-testing; fix with the agent-API refactor. + embeddings_directory: Path = EmbeddingEnvironment.get_dir("userguide"), # noqa: B008 + *, + streaming: bool = False, +) -> Runnable: + userguide_retriever = create_userguide_retriever(embedding, embeddings_directory) + + if streaming: + llm = llm.model_copy(update={"streaming": True}) + + return create_rag_chain(llm, userguide_retriever, userguide_qa_prompt) diff --git a/src/retrievers/userguide/retriever.py b/src/retrievers/userguide/retriever.py new file mode 100644 index 0000000..8500cb8 --- /dev/null +++ b/src/retrievers/userguide/retriever.py @@ -0,0 +1,37 @@ +from pathlib import Path + +from langchain_chroma.vectorstores import Chroma +from langchain_core.embeddings import Embeddings +from langchain_core.retrievers import BaseRetriever + +from retrievers.csv_chroma import chroma_settings + +CHROMA_COLLECTION = "sections" +DEFAULT_SEARCH_K = 6 + + +def create_userguide_retriever( + embedding: Embeddings, + embeddings_directory: Path | None, + *, + k: int = DEFAULT_SEARCH_K, +) -> BaseRetriever: + if embeddings_directory is None: + raise ValueError( + "User guide embeddings are not configured. " + "Run ./bin/embeddings_manager use /userguide/." + ) + + chroma_path = Path(embeddings_directory) / CHROMA_COLLECTION + if not (chroma_path / "chroma.sqlite3").is_file(): + raise FileNotFoundError( + f"User guide Chroma collection not found at {chroma_path}. " + "Run ./bin/embeddings_manager make /userguide/." + ) + + vectordb = Chroma( + persist_directory=str(chroma_path), + embedding_function=embedding, + client_settings=chroma_settings, + ) + return vectordb.as_retriever(search_kwargs={"k": k}) diff --git a/src/tools/external_search/tavily_wrapper.py b/src/tools/external_search/tavily_wrapper.py index 54e373b..96ab800 100644 --- a/src/tools/external_search/tavily_wrapper.py +++ b/src/tools/external_search/tavily_wrapper.py @@ -16,7 +16,7 @@ def __init__( search_depth: Literal["basic", "advanced"] = "advanced", max_results: int = 5, rate_limit: int = 100, # requests per minute - ): + ) -> None: self.tavily_client: AsyncTavilyClient | None = None self.search_depth = search_depth self.max_results = max_results diff --git a/src/tools/external_search/workflow.py b/src/tools/external_search/workflow.py index 0a409c8..038a44d 100644 --- a/src/tools/external_search/workflow.py +++ b/src/tools/external_search/workflow.py @@ -6,8 +6,10 @@ from langgraph.graph.state import CompiledStateGraph from langgraph.utils.runnable import RunnableLike -from agent.tasks.completeness_grader import (CompletenessGrade, - create_completeness_grader) +from agent.tasks.completeness_grader import ( + CompletenessGrade, + create_completeness_grader, +) from tools.external_search.state import SearchState from tools.external_search.tavily_wrapper import TavilyWrapper @@ -15,11 +17,10 @@ def decide_next_steps(state: SearchState) -> Literal["perform_web_search", "no_search"]: if state["complete"] == "No": return "perform_web_search" - else: - return "no_search" + return "no_search" -def no_search(_) -> SearchState: +def no_search(_: SearchState) -> SearchState: return SearchState(search_results=[]) diff --git a/src/util/chainlit_helpers.py b/src/util/chainlit_helpers.py index 4c1b4fd..3d72c3c 100644 --- a/src/util/chainlit_helpers.py +++ b/src/util/chainlit_helpers.py @@ -1,17 +1,19 @@ import os +from collections.abc import Iterable from datetime import datetime from pathlib import PurePosixPath -from typing import Any, Iterable +from typing import Any import chainlit as cl from chainlit.data import get_data_layer from chainlit.data.storage_clients.s3 import S3StorageClient from langchain_community.callbacks import OpenAICallbackHandler +from tools.external_search.state import WebSearchResult from util.config_yml import Config, TriggerEvent from util.config_yml.usage_limits import MessageRate -guest_user_metadata: dict[str, Any] = {} +_GUEST_METADATA_KEY = "_guest_metadata" class PrefixedS3StorageClient(S3StorageClient): @@ -43,6 +45,14 @@ def get_user_id() -> str | None: return user.identifier if user else None +def _get_guest_metadata() -> dict[str, Any]: + """Get the per-session guest metadata dict, creating it if needed.""" + metadata: dict[str, Any] = cl.user_session.get(_GUEST_METADATA_KEY, {}) + if not metadata: + cl.user_session.set(_GUEST_METADATA_KEY, metadata) + return metadata + + def get_user_metadata( key: Any, default: Any | None = None, @@ -51,10 +61,9 @@ def get_user_metadata( user: cl.User | None = cl.user_session.get("user") if user: return user.metadata.get(key, default) - elif use_guest: - return guest_user_metadata.get(key, default) - else: - return default + if use_guest: + return _get_guest_metadata().get(key, default) + return default def is_feature_enabled(config: Config | None, feature_id: str) -> bool: @@ -82,12 +91,11 @@ def save_openai_metrics(message_id: str, openai_cb: OpenAICallbackHandler) -> No def set_user_metadata(key: Any, value: Any, use_guest: bool = True) -> None: - global guest_user_metadata # not ideal, but works for now user: cl.User | None = cl.user_session.get("user") if user: user.metadata[key] = value elif use_guest: - guest_user_metadata[key] = value + _get_guest_metadata()[key] = value async def message_rate_limited(config: Config | None) -> bool: @@ -144,18 +152,18 @@ async def static_messages( chat_profile: str = cl.user_session.get("chat_profile") - messages_formatted: Iterable[str] = map( - lambda msg: msg.format( + messages_formatted: Iterable[str] = ( + msg.format( chat_profile=chat_profile, user_id=user_id, - ), - messages.values(), + ) + for msg in messages.values() ) await send_messages(messages_formatted) async def update_search_results( - search_results: list[dict[str, str]], + search_results: list[WebSearchResult], message: cl.Message, ) -> None: search_results_element = cl.CustomElement( diff --git a/src/util/config_yml/__init__.py b/src/util/config_yml/__init__.py index e6d57e9..cccded1 100644 --- a/src/util/config_yml/__init__.py +++ b/src/util/config_yml/__init__.py @@ -4,15 +4,19 @@ import yaml from pydantic import BaseModel, ValidationError -from agent.profiles import ProfileName +from agent.profile_names import ProfileName from util.config_yml.features import Feature, Features from util.config_yml.messages import Message, TriggerEvent from util.config_yml.usage_limits import MessageRate, UsageLimits from util.config_yml.user_matching import match_user from util.logging import logging -CONFIG_YML = Path("config.yml") -CONFIG_DEFAULT_YML = Path("config_default.yml") +# Anchored to the repo rather than the working directory, so these resolve the +# same whether the process starts from the repo root, a subdirectory, or /app in +# the container. Matches util.embedding_environment.REPO_ROOT. +REPO_ROOT: Path = Path(__file__).parent.parent.parent.parent +CONFIG_YML = REPO_ROOT / "config.yml" +CONFIG_DEFAULT_YML = REPO_ROOT / "config_default.yml" class Config(BaseModel): @@ -29,16 +33,16 @@ def get_feature( if feature_id in self.features.model_fields: feature: Feature = getattr(self.features, feature_id) return feature.enabled and feature.matches_user_group(user_id) - else: - return True + return True def get_messages( self, user_id: str | None = None, event: TriggerEvent | None = None, after_messages: int | None = None, - last_messages: dict[str, str] = {}, + last_messages: dict[str, str] | None = None, ) -> dict[str, str]: + last_messages = last_messages if last_messages is not None else {} return { message_id: message.message for message_id, message in self.messages.items() @@ -46,7 +50,7 @@ def get_messages( message.enabled and match_user(message.recipients, user_id) and message.trigger.match_trigger( - event, after_messages, last_messages.get(message_id, None) + event, after_messages, last_messages.get(message_id) ) ) } @@ -54,8 +58,11 @@ def get_messages( def get_message_rate_usage_limited( self, user_id: str | None = None, - message_times_queue: list[str] = [], + message_times_queue: list[str] | None = None, ) -> MessageRate | None: + message_times_queue = ( + message_times_queue if message_times_queue is not None else [] + ) message_rate: MessageRate for message_rate in self.usage_limits.message_rates: if match_user(message_rate.users, user_id): @@ -63,16 +70,53 @@ def get_message_rate_usage_limited( return None # not rate limited @classmethod - def from_yaml(cls, config_yml: Path = CONFIG_YML) -> Self | None: - if not config_yml.exists(): - logging.warning( - f"Config file not found: {config_yml} ; falling back to {CONFIG_DEFAULT_YML}" - ) - config_yml = CONFIG_DEFAULT_YML + def _load(cls, config_yml: Path) -> Self: with open(config_yml) as f: - yaml_data: dict = yaml.safe_load(f) + yaml_data = yaml.safe_load(f) + if not isinstance(yaml_data, dict): + raise ValueError(f"{config_yml} is empty or is not a YAML mapping") + return cls(**yaml_data) + + @classmethod + def from_yaml(cls, config_yml: Path = CONFIG_YML) -> Self | None: + """Load config.yml, or the shipped defaults when there is no config.yml. + + A *present but invalid* config.yml raises. It used to be swallowed, and + both of the quiet options are wrong: + + - returning None disables every config-driven feature including the + message quota, so one typo removed rate limiting for everybody; + - falling back to config_default.yml silently applies settings nobody + chose -- it would re-enable `postprocessing` (external web search, and + its per-message cost) for an operator who had deliberately turned it + off, and replace their quota with the default 100. + + Refusing to start is the only option that cannot quietly do the wrong + thing: the typo surfaces at deploy time rather than in a bill. An absent + config.yml is a different case and still falls back, because running + with documented defaults is what a fresh checkout expects. + """ + if config_yml != CONFIG_DEFAULT_YML: + try: + return cls._load(config_yml) + except (FileNotFoundError, IsADirectoryError): + # docker-compose bind-mounts ./config.yml; when the host file is + # missing Docker creates a directory in its place, so both mean + # "no config supplied". + logging.warning( + f"No config at {config_yml}; using {CONFIG_DEFAULT_YML}" + ) + except (ValidationError, ValueError, yaml.YAMLError) as e: + raise SystemExit( + f"Invalid config {config_yml}:\n{e}\n\n" + "Refusing to start. Fix the file, or remove it to run with " + f"the defaults in {CONFIG_DEFAULT_YML}." + ) from e + try: - return cls(**yaml_data) - except ValidationError as e: - logging.warning(e) - return None + return cls._load(CONFIG_DEFAULT_YML) + except Exception as e: + raise SystemExit( + f"The shipped default config {CONFIG_DEFAULT_YML} is unusable, " + f"which should not happen in a working checkout:\n{e}" + ) from e diff --git a/src/util/config_yml/features.py b/src/util/config_yml/features.py index eb6c760..f1db928 100644 --- a/src/util/config_yml/features.py +++ b/src/util/config_yml/features.py @@ -15,8 +15,7 @@ class Feature(BaseModel): def matches_user_group(self, user_id: str | None) -> bool: if self.user_group == UserGroup.logged_in: return user_id is not None - else: - return True + return True class Features(BaseModel): diff --git a/src/util/config_yml/intervals.py b/src/util/config_yml/intervals.py index f0ec989..9cd6517 100644 --- a/src/util/config_yml/intervals.py +++ b/src/util/config_yml/intervals.py @@ -1,6 +1,9 @@ import re from datetime import timedelta +# Kept in sync with the `interval` and `freq_max` patterns in .config.schema.yaml. +INTERVAL_PATTERN = r"^[0-9]+[smhdw]$" + interval_units = { "s": "seconds", "m": "minutes", @@ -11,9 +14,22 @@ def parse_interval(interval_str: str) -> timedelta: + """Parse an interval such as "3h" into a timedelta. + + Raises ValueError on anything malformed. It used to return timedelta(0), + which silently disabled rate limiting: a zero-length window means every + queued timestamp is already outside it, so the queue drained on every call + and no user was ever limited. + + Callers reach this only through fields that carry INTERVAL_PATTERN, so a bad + value is rejected when config.yml is loaded rather than here. + """ re_match = re.fullmatch(r"([0-9]+)([smhdw])", interval_str) if not re_match: - return timedelta(0) + raise ValueError( + f"malformed interval {interval_str!r}: expected a number followed by " + "one of s/m/h/d/w, e.g. '30s', '3h', '7d'" + ) value = int(re_match.group(1)) unit = interval_units[re_match.group(2)] return timedelta(**{unit: value}) diff --git a/src/util/config_yml/messages.py b/src/util/config_yml/messages.py index f8b45ac..937d00a 100644 --- a/src/util/config_yml/messages.py +++ b/src/util/config_yml/messages.py @@ -1,9 +1,9 @@ -from datetime import datetime +from datetime import UTC, datetime from enum import StrEnum, auto -from pydantic import BaseModel +from pydantic import BaseModel, Field -from util.config_yml.intervals import parse_interval +from util.config_yml.intervals import INTERVAL_PATTERN, parse_interval class TriggerEvent(StrEnum): @@ -13,38 +13,49 @@ class TriggerEvent(StrEnum): on_message = auto() +def _as_utc(value: datetime) -> datetime: + """Treat a naive datetime as UTC; convert an aware one to UTC.""" + if value.tzinfo is None: + return value.replace(tzinfo=UTC) + return value.astimezone(UTC) + + class Trigger(BaseModel): event: TriggerEvent | None = None after_messages: int | None = None start: datetime | None = None end: datetime | None = None - freq_max: str | None = None + freq_max: str | None = Field(default=None, pattern=INTERVAL_PATTERN) def match_trigger( self, event: TriggerEvent | None = None, after_messages: int | None = None, last_message: str | None = None, - ): - now = datetime.now() + ) -> bool: if self.event and self.event != event: return False if self.after_messages and self.after_messages != after_messages: return False - if self.start and self.start.replace(tzinfo=None) > now: + # start/end come from config.yml and are usually written with an offset + # ("2025-01-01T00:00:00Z"). These used to be compared by stripping tzinfo, + # which discards the offset instead of converting, shifting the window by + # the host's UTC offset. + now_utc = datetime.now(UTC) + if self.start and _as_utc(self.start) > now_utc: return False - if self.end and self.end.replace(tzinfo=None) < now: + if self.end and _as_utc(self.end) < now_utc: return False - if ( + # last_message is written by chainlit_helpers as a naive local + # datetime.now().isoformat(), so it keeps its own naive local clock. + return not ( self.freq_max and last_message and ( parse_interval(self.freq_max) - > now - datetime.fromisoformat(last_message) + > datetime.now() - datetime.fromisoformat(last_message) ) - ): - return False - return True + ) class Message(BaseModel): diff --git a/src/util/config_yml/usage_limits.py b/src/util/config_yml/usage_limits.py index 81460a8..11a553e 100644 --- a/src/util/config_yml/usage_limits.py +++ b/src/util/config_yml/usage_limits.py @@ -1,15 +1,17 @@ from datetime import datetime from typing import Self -from pydantic import BaseModel +from pydantic import BaseModel, Field -from util.config_yml.intervals import parse_interval +from util.config_yml.intervals import INTERVAL_PATTERN, parse_interval class MessageRate(BaseModel): users: list[str] - max_messages: int - interval: str + max_messages: int = Field(gt=0) + # Validated here so a typo like '3hr' is rejected when config.yml loads, + # rather than reaching parse_interval and disabling the limiter. + interval: str = Field(pattern=INTERVAL_PATTERN) def check_rate(self, message_times_queue: list[str]) -> Self | None: now = datetime.now() @@ -23,8 +25,7 @@ def check_rate(self, message_times_queue: list[str]) -> Self | None: if len(message_times_queue) < self.max_messages: message_times_queue.append(now.isoformat()) return None # not rate limited - else: - return self + return self class UsageLimits(BaseModel): diff --git a/src/util/config_yml/user_matching.py b/src/util/config_yml/user_matching.py index dc3ac36..6240dff 100644 --- a/src/util/config_yml/user_matching.py +++ b/src/util/config_yml/user_matching.py @@ -6,6 +6,8 @@ def match_user(users_spec: list[str] | None, user_id: str | None) -> bool: if not users_spec: return True for entry in users_spec: + if not entry: # an empty entry matches nobody; entry[0] below would raise + continue if entry == "all": return True if user_id is None: @@ -14,7 +16,7 @@ def match_user(users_spec: list[str] | None, user_id: str | None) -> bool: else: if entry == "logged_in": return True - elif entry[0] == "/" and entry[-1] == "/": + if entry[0] == "/" and entry[-1] == "/": if re.search(entry[1:-1], user_id): return True else: diff --git a/src/util/embedding_environment.py b/src/util/embedding_environment.py index ab4a43f..140de60 100644 --- a/src/util/embedding_environment.py +++ b/src/util/embedding_environment.py @@ -6,15 +6,15 @@ class EmbeddingEnvironment: - def __init__(self, env_path: str): - self.embeddings: dict[str, Path] = dict() + def __init__(self, env_path: str) -> None: + self.embeddings: dict[str, Path] = {} if env_path != "": for embedding_path in map(Path, env_path.split(":")): db: str = embedding_path.parent.name self.embeddings[db] = embedding_path @classmethod - def _get(cls): # -> Self + def _get(cls) -> "EmbeddingEnvironment": if EM_CURRENT.exists(): with EM_CURRENT.open("r") as current_fp: env_path = current_fp.read() @@ -30,8 +30,7 @@ def get_dict(cls) -> dict[str, Path]: def get_dir(cls, key: str) -> Path | None: if key in cls._get().embeddings: return EM_ARCHIVE / cls._get().embeddings[key] - else: - return None + return None @classmethod def get_model(cls, key: str) -> str: diff --git a/src/util/orcid_provider.py b/src/util/orcid_provider.py new file mode 100644 index 0000000..c1402b7 --- /dev/null +++ b/src/util/orcid_provider.py @@ -0,0 +1,70 @@ +import os +from typing import Any + +import httpx +from chainlit.oauth_providers import OAuthProvider +from chainlit.user import User +from fastapi import HTTPException + + +class ORCIDOAuthProvider(OAuthProvider): + id = "orcid" + # RUF012 is silenced at the site below: the base class declares this as an + # instance variable, so it cannot be narrowed to ClassVar here. + env = ["OAUTH_ORCID_CLIENT_ID", "OAUTH_ORCID_CLIENT_SECRET"] # noqa: RUF012 + + def __init__(self) -> None: + self.client_id = os.environ.get("OAUTH_ORCID_CLIENT_ID", "") + self.client_secret = os.environ.get("OAUTH_ORCID_CLIENT_SECRET", "") + self.authorize_url = "https://orcid.org/oauth/authorize" + self.token_url = "https://orcid.org/oauth/token" # noqa: S105 (a URL, not a secret) + self.user_info_url = "https://orcid.org/oauth/userinfo" + self.authorize_params = { + "response_type": "code", + "scope": "/authenticate", + } + + if prompt := self.get_prompt(): + self.authorize_params["prompt"] = prompt + + async def get_raw_token_response(self, code: str, url: str) -> dict: + payload = { + "client_id": self.client_id, + "client_secret": self.client_secret, + "code": code, + "grant_type": "authorization_code", + "redirect_uri": url, + } + async with httpx.AsyncClient() as client: + response = await client.post(self.token_url, data=payload) + response.raise_for_status() + token_response: dict[str, Any] = response.json() + return token_response + + async def get_token(self, code: str, url: str) -> str: + json = await self.get_raw_token_response(code, url) + token = json.get("access_token") + if not token: + raise HTTPException( + status_code=400, detail="Access token missing in the response" + ) + return str(token) + + async def get_user_info(self, token: str) -> tuple[dict[str, str], User]: + async with httpx.AsyncClient() as client: + response = await client.get( + self.user_info_url, + headers={"Authorization": f"Bearer {token}"}, + ) + response.raise_for_status() + + orcid_user = response.json() + + # ORCiD /userinfo returns the ORCID iD under "sub" (a stable identifier). + # Use it so chat history persists per ORCID account. + orcid_id = orcid_user.get("sub") or orcid_user.get("orcid") + user = User( + identifier=orcid_id or orcid_user.get("email", "orcid-user"), + metadata={"provider": "orcid"}, + ) + return (orcid_user, user) diff --git a/template.env b/template.env deleted file mode 100644 index e570b8b..0000000 --- a/template.env +++ /dev/null @@ -1 +0,0 @@ -OPENAI_API_KEY= diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..a8f9469 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,28 @@ +# Tests + +These are **characterization tests**: they pin down what the code does *today*, so that +the dependency upgrade (LangChain 0.3 -> 1.x, Chroma <1.0 -> 1.x) has a tripwire. Where +current behaviour looks wrong, the test asserts the current behaviour and carries a +`BUG:` comment rather than asserting the desired behaviour — change the test and the code +together, deliberately. + +## Running + + poetry run pytest # everything importable + poetry run pytest -m "not requires_retrieval_stack" + +## Markers + +- `requires_retrieval_stack` — needs langchain/chromadb/torch importable. +- `requires_embeddings` — additionally needs an installed bundle (`./bin/embeddings_manager ls`). + +## Why so little is covered + +Most of `src/` cannot be imported without a provisioned environment: `load_dotenv()` and +`AgentGraph(...)` run at import time in the entry points, and +`EmbeddingEnvironment.get_dir(...)` is a *default argument* in the retriever modules, so +importing them requires an embeddings bundle on disk. Even `util.config_yml`, which is +pure configuration logic, transitively imports torch because it pulls one enum from +`agent.profiles`. + +Breaking that import-time coupling is what unlocks real coverage here. diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..1ee9f59 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,17 @@ +import importlib.util + +import pytest + +RETRIEVAL_STACK_MODULES = ("langchain", "langchain_chroma", "chromadb", "nltk") + + +def _stack_available() -> bool: + return all(importlib.util.find_spec(m) is not None for m in RETRIEVAL_STACK_MODULES) + + +def pytest_runtest_setup(item: pytest.Item) -> None: + if ( + list(item.iter_markers(name="requires_retrieval_stack")) + and not _stack_available() + ): + pytest.skip("retrieval stack (langchain/chromadb/nltk) not installed") diff --git a/tests/golden/questions.txt b/tests/golden/questions.txt new file mode 100644 index 0000000..2955d53 --- /dev/null +++ b/tests/golden/questions.txt @@ -0,0 +1,28 @@ +# Fixed question set for the retrieval baseline. One per line; blank lines and +# `#` comments are ignored. +# +# Chosen to span the four Chroma collections and several query shapes: named +# entities, processes, complexes, gene/protein lookups, disease context, and a +# couple of vaguer questions where retrieval has to work harder. Keep this list +# stable -- changing it invalidates comparison against an earlier capture. + +What does CDK5 phosphorylate in Alzheimer's disease? +How does TP53 regulate PTEN transcription? +What is the role of CDK12 in DNA repair gene expression? +Which proteins are in the RNA polymerase II elongation complex? +What happens during Golgi fragmentation in neurodegeneration? +How is oxidative stress handled by peroxiredoxins? +What reactions involve EGFR autophosphorylation? +Describe the components of the proteasome +What is the function of BRCA1 in homologous recombination? +Which pathways involve insulin receptor signalling? +What role does p53 play in apoptosis? +How does the electron transport chain generate ATP? +What is the mechanism of ubiquitin-mediated protein degradation? +Which complexes contain histone deacetylases? +What are the steps of glycolysis? +How do Wnt signalling pathways control gene transcription? +What is the role of mTOR in cell growth? +Which reactions produce nitric oxide? +What proteins interact with beta-catenin? +How does autophagy degrade cellular components? diff --git a/tests/retrievers/test_hybrid_retriever.py b/tests/retrievers/test_hybrid_retriever.py new file mode 100644 index 0000000..f24b20d --- /dev/null +++ b/tests/retrievers/test_hybrid_retriever.py @@ -0,0 +1,202 @@ +"""Characterization of the hybrid retrieval fusion. + +`HybridRetriever` fuses BM25 and vector hits per Chroma subdirectory using +LangChain's Reciprocal Rank Fusion. RRF's constant and its de-duplication key are +LangChain *implementation details* that this repo depends on, so an upgrade across +the 0.3 -> 1.x boundary can silently change result ordering. These tests exist to +make that change loud. + +Skipped unless the retrieval stack is installed; run them on both sides of the +upgrade and diff the output. +""" + +from pathlib import Path +from typing import Any, cast + +import pytest + +pytest.importorskip("langchain", reason="retrieval stack not installed") +pytest.importorskip("chromadb", reason="retrieval stack not installed") + +from langchain_core.documents import Document # noqa: E402 + +from retrievers.csv_chroma import ( # noqa: E402 + MAX_DOCUMENTS_PER_COLLECTION, + HybridRetriever, + dedupe_by_entity, + list_chroma_subdirectories, +) + +pytestmark = pytest.mark.requires_retrieval_stack + + +def _doc(text: str) -> Document: + return Document(page_content=text) + + +def _fuse(doc_lists: list[list[Document]]) -> list[Document]: + """`weighted_reciprocal_rank` never touches `self`; call it without an instance. + + Building a real HybridRetriever requires an LLM and an on-disk Chroma store, + which would make this an integration test rather than a characterization of + the ranking maths. + """ + return HybridRetriever.weighted_reciprocal_rank(cast(Any, None), doc_lists) + + +def test_subdirectories_are_discovered_by_chroma_sqlite_marker(tmp_path: Path) -> None: + for name in ("reactions", "summations", "complexes"): + (tmp_path / name).mkdir() + (tmp_path / name / "chroma.sqlite3").touch() + (tmp_path / "csv_files").mkdir() # sibling data dir, not a collection + (tmp_path / "empty").mkdir() # no marker file + + assert sorted(list_chroma_subdirectories(tmp_path)) == [ + "complexes", + "reactions", + "summations", + ] + + +def test_missing_directory_yields_no_subdirectories(tmp_path: Path) -> None: + assert list_chroma_subdirectories(tmp_path / "nope") == [] + + +def test_documents_in_both_lists_outrank_documents_in_one() -> None: + """The core reason for hybrid retrieval: BM25/vector agreement should win.""" + bm25 = [_doc("agreed"), _doc("bm25 only")] + vector = [_doc("vector only"), _doc("agreed")] + + ranked = [d.page_content for d in _fuse([bm25, vector])] + + assert ranked[0] == "agreed" + assert set(ranked) == {"agreed", "bm25 only", "vector only"} + + +def test_fusion_deduplicates_on_page_content() -> None: + """De-dup key is page_content, not metadata -- identical text collapses.""" + duplicated = [[_doc("same")], [_doc("same")], [_doc("same")]] + assert len(_fuse(duplicated)) == 1 + + +def test_rank_order_within_a_single_list_is_preserved() -> None: + single = [[_doc("first"), _doc("second"), _doc("third")]] + assert [d.page_content for d in _fuse(single)] == ["first", "second", "third"] + + +def test_empty_input_is_not_an_error() -> None: + assert _fuse([[]]) == [] + + +def test_ties_are_broken_by_position_not_by_score() -> None: + """Equal RRF scores are resolved by which list came first. + + x is rank 1 in the first list and rank 2 in the second; y is the mirror + image, so both score w/61 + w/62 exactly. sorted() is stable, so the order + that survives is the order of chain.from_iterable(doc_lists) -- meaning the + caller's list order decides. + + In HybridRetriever the lists are the query variants, and within each list + BM25's results precede the vector retriever's. So on a tie, earlier query + variants win, and BM25 wins over the vector store. Deterministic, but a + consequence of iteration order rather than of relevance. + """ + a, b = [_doc("x"), _doc("y")], [_doc("y"), _doc("x")] + assert [d.page_content for d in _fuse([a, b])] == ["x", "y"] + assert [d.page_content for d in _fuse([b, a])] == ["y", "x"] + + +def test_weights_are_uniform_so_they_cannot_change_the_ordering() -> None: + """create_bm25_chroma_ensemble_retriever passes [1/n] * n. + + A constant multiplier across every list scales all scores equally, so the + weighting is currently inert. It is the obvious place to put a BM25-vs-vector + balance, but nothing uses it today. + """ + lists = [[_doc("p"), _doc("q")], [_doc("q"), _doc("r")]] + uniform_half = HybridRetriever.weighted_reciprocal_rank(cast(Any, None), lists) + assert [d.page_content for d in uniform_half] == ["q", "p", "r"] + + +@pytest.mark.requires_embeddings +def test_installed_bundle_exposes_the_expected_collections() -> None: + """Guards the bundle layout the retriever assumes: //chroma.sqlite3 + plus a sibling csv_files/.csv for BM25.""" + from util.embedding_environment import EmbeddingEnvironment + + directory = EmbeddingEnvironment.get_dir("reactome") + if directory is None or not directory.is_dir(): + pytest.skip("no reactome embeddings installed") + + collections = list_chroma_subdirectories(directory) + assert collections, "bundle contains no chroma collections" + for collection in collections: + assert ( + directory / "csv_files" / f"{collection}.csv" + ).is_file(), f"BM25 source CSV missing for collection '{collection}'" + + +def _doc_with_id(st_id: str, text: str) -> Document: + return Document(page_content=text, metadata={"st_id": st_id}) + + +def test_dedupe_keeps_the_highest_ranked_row_per_entity() -> None: + """One reaction occupies several CSV rows -- one per pathway/input/output + combination -- and those rows have different page_content, so nothing + upstream collapses them. Before this, vector search on `reactions` returned + ten results containing about five distinct reactions (issue #169).""" + docs = [ + _doc_with_id("R-HSA-1", "in pathway A"), + _doc_with_id("R-HSA-1", "in pathway B"), + _doc_with_id("R-HSA-2", "second reaction"), + _doc_with_id("R-HSA-1", "in pathway C"), + _doc_with_id("R-HSA-3", "third reaction"), + ] + kept = dedupe_by_entity(docs, limit=10) + assert [d.metadata["st_id"] for d in kept] == ["R-HSA-1", "R-HSA-2", "R-HSA-3"] + assert kept[0].page_content == "in pathway A", "keeps the highest-ranked row" + + +def test_dedupe_respects_the_limit() -> None: + docs = [_doc_with_id(f"R-HSA-{i}", f"doc {i}") for i in range(20)] + assert len(dedupe_by_entity(docs, limit=10)) == 10 + + +def test_dedupe_falls_back_to_page_content_without_st_id() -> None: + """A collection whose metadata lacks st_id degrades to the old behaviour + rather than raising.""" + docs = [_doc("same"), _doc("same"), _doc("different")] + kept = dedupe_by_entity(docs, limit=10) + assert [d.page_content for d in kept] == ["same", "different"] + + +def test_bm25_and_vector_are_fused_as_separate_lists() -> None: + """Both retrievers' top hits must score equally. + + Concatenating them into one list put every vector result at rank 11+, so the + best vector hit scored 1/71 against BM25's 1/61 (issue #170). As separate + lists both are rank 1, and a document only one of them found cannot outrank + a document they agree on. + """ + bm25 = [_doc("agreed"), _doc("bm25 only")] + vector = [_doc("agreed"), _doc("vector only")] + + ranked = [d.page_content for d in _fuse([bm25, vector])] + assert ranked[0] == "agreed", "agreement between the two retrievers wins" + # the two single-source documents are tied, so only membership is asserted + assert set(ranked[1:]) == {"bm25 only", "vector only"} + + +def test_fused_results_are_capped_per_collection() -> None: + """RRF returns every unique document it is given, not a top-N. + + Uncapped, the retriever ranked ~222 documents and sent all of them -- + roughly 32k tokens per message, which made the ranking decorative since + nothing acted on it. The cap is applied per collection so that one + collection cannot crowd out the others. + """ + many = [[_doc(f"doc {i}") for i in range(50)]] + assert len(_fuse(many)) == 50, "the fusion itself still returns everything" + assert ( + MAX_DOCUMENTS_PER_COLLECTION < 50 + ), "the cap, applied by the caller, is what bounds the prompt" diff --git a/tests/util/test_config.py b/tests/util/test_config.py new file mode 100644 index 0000000..1f011b3 --- /dev/null +++ b/tests/util/test_config.py @@ -0,0 +1,137 @@ +"""Config loading, with an emphasis on what happens when config.yml is wrong. + +`Config.from_yaml` returning None switches off every config-driven feature, +including the message quota, so the fallback path matters as much as the happy +one: a typo in config.yml must not quietly remove rate limiting. +""" + +from pathlib import Path + +import pytest + +from util.config_yml import CONFIG_DEFAULT_YML, CONFIG_YML, Config + +VALID = """ +profiles: ["React-to-Me"] +features: + postprocessing: + enabled: true + user_group: all +usage_limits: + message_rates: + - users: ["all"] + max_messages: 5 + interval: 1h +messages: + hello: + message: hi + trigger: + event: on_chat_start +""" + + +def test_the_shipped_default_config_is_valid() -> None: + """If this fails, every fallback below lands on None and features silently die.""" + config = Config.from_yaml(CONFIG_DEFAULT_YML) + assert config is not None + assert config.usage_limits.message_rates, "defaults must carry a message quota" + + +def test_valid_config_is_loaded(tmp_path: Path) -> None: + path = tmp_path / "config.yml" + path.write_text(VALID) + config = Config.from_yaml(path) + assert config is not None + assert config.usage_limits.message_rates[0].max_messages == 5 + + +def test_missing_config_falls_back_to_defaults(tmp_path: Path) -> None: + config = Config.from_yaml(tmp_path / "nope.yml") + assert config is not None + assert config.usage_limits.message_rates + + +def test_invalid_config_refuses_to_start(tmp_path: Path) -> None: + """A present-but-invalid config.yml must be fatal, not silently substituted. + + Both quiet options are wrong. Returning None disables every config-driven + feature including the quota, so one typo removed rate limiting. Falling back + to the defaults silently applies settings nobody chose -- see + test_fallback_would_have_re_enabled_a_disabled_feature for why that matters. + """ + path = tmp_path / "config.yml" + path.write_text(VALID.replace("interval: 1h", "interval: 1hr")) + with pytest.raises(SystemExit, match="Invalid config"): + Config.from_yaml(path) + + +def test_fallback_would_have_re_enabled_a_disabled_feature(tmp_path: Path) -> None: + """Why the fallback was the wrong fix, pinned so it is not reintroduced. + + An operator who turns postprocessing off and makes an unrelated typo would, + under a silent fallback, get the default config back -- which has + postprocessing enabled and a quota of 100. Enabling external web search + because of a typo elsewhere in the file is a cost and privacy change nobody + asked for. + """ + disabled = VALID.replace("enabled: true", "enabled: false") + assert "enabled: false" in disabled + + good = tmp_path / "good.yml" + good.write_text(disabled) + config = Config.from_yaml(good) + assert config is not None + assert config.features.postprocessing.enabled is False + + # the same file with a typo must now raise rather than quietly flip it back on + bad = tmp_path / "bad.yml" + bad.write_text(disabled.replace("interval: 1h", "interval: 1hr")) + with pytest.raises(SystemExit): + Config.from_yaml(bad) + + defaults = Config.from_yaml(CONFIG_DEFAULT_YML) + assert defaults is not None + assert ( + defaults.features.postprocessing.enabled is True + ), "the default this would have silently substituted" + + +def test_config_yml_as_a_directory_falls_back(tmp_path: Path) -> None: + """docker-compose bind-mounts ./config.yml; if it is absent on the host, + Docker creates a *directory* there. That means "no config supplied", not + "broken config", so it falls back rather than refusing to start.""" + path = tmp_path / "config.yml" + path.mkdir() + assert Config.from_yaml(path) is not None + + +@pytest.mark.parametrize("content", ["", "\n", "just a string", "[1, 2, 3]"]) +def test_non_mapping_config_is_fatal(tmp_path: Path, content: str) -> None: + """A file that exists but is not a config mapping is a mistake, not a default.""" + path = tmp_path / "config.yml" + path.write_text(content) + with pytest.raises(SystemExit): + Config.from_yaml(path) + + +def test_unknown_profile_is_fatal(tmp_path: Path) -> None: + """Starting with a profile the agent cannot build should not be silent.""" + path = tmp_path / "config.yml" + path.write_text(VALID.replace('["React-to-Me"]', '["Nonexistent Profile"]')) + with pytest.raises(SystemExit): + Config.from_yaml(path) + + +def test_config_paths_do_not_depend_on_the_working_directory( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The defaults must resolve from anywhere, not just the repo root. + + These used to be Path("config.yml") / Path("config_default.yml"), i.e. + relative to the process working directory, so running from a subdirectory + silently lost the config. + """ + assert CONFIG_YML.is_absolute() + assert CONFIG_DEFAULT_YML.is_absolute() + monkeypatch.chdir(tmp_path) + assert Config.from_yaml(CONFIG_DEFAULT_YML) is not None diff --git a/tests/util/test_embedding_environment.py b/tests/util/test_embedding_environment.py new file mode 100644 index 0000000..cebc5d8 --- /dev/null +++ b/tests/util/test_embedding_environment.py @@ -0,0 +1,63 @@ +"""Characterization of embeddings path resolution. + +`embeddings/current` is a single colon-separated line mapping each database to the +bundle in use. `bin/embeddings_manager use` writes it; the retriever modules read it +at import time. +""" + +from pathlib import Path + +import pytest + +import util.embedding_environment as ee +from util.embedding_environment import EmbeddingEnvironment + +REACTOME = "openai/text-embedding-3-large/reactome/Release90" +UNIPROT = "openai/text-embedding-3-large/uniprot/Release90" + + +@pytest.fixture +def archive(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + monkeypatch.setattr(ee, "EM_ARCHIVE", tmp_path) + monkeypatch.setattr(ee, "EM_CURRENT", tmp_path / "current") + return tmp_path + + +def test_no_current_file_means_no_embeddings(archive: Path) -> None: + assert EmbeddingEnvironment.get_dict() == {} + assert EmbeddingEnvironment.get_dir("reactome") is None + + +def test_database_key_is_the_parent_directory_name(archive: Path) -> None: + """The db key comes from the path, not the file -- `.../reactome/Release90` -> `reactome`.""" + (archive / "current").write_text(REACTOME) + assert EmbeddingEnvironment.get_dict() == {"reactome": Path(REACTOME)} + assert EmbeddingEnvironment.get_dir("reactome") == archive / REACTOME + + +def test_multiple_databases_are_colon_separated(archive: Path) -> None: + (archive / "current").write_text(f"{REACTOME}:{UNIPROT}") + assert set(EmbeddingEnvironment.get_dict()) == {"reactome", "uniprot"} + + +def test_get_model_strips_the_database_and_version(archive: Path) -> None: + (archive / "current").write_text(REACTOME) + assert EmbeddingEnvironment.get_model("reactome") == "openai/text-embedding-3-large" + + +def test_set_one_adds_without_disturbing_other_databases(archive: Path) -> None: + (archive / "current").write_text(REACTOME) + EmbeddingEnvironment.set_one(Path(UNIPROT)) + assert set(EmbeddingEnvironment.get_dict()) == {"reactome", "uniprot"} + + +def test_set_one_replaces_the_bundle_for_the_same_database(archive: Path) -> None: + (archive / "current").write_text(REACTOME) + newer = "openai/text-embedding-3-large/reactome/Release91" + EmbeddingEnvironment.set_one(Path(newer)) + assert EmbeddingEnvironment.get_dict() == {"reactome": Path(newer)} + + +def test_empty_current_file_is_not_an_error(archive: Path) -> None: + (archive / "current").write_text("") + assert EmbeddingEnvironment.get_dict() == {} diff --git a/tests/util/test_features.py b/tests/util/test_features.py new file mode 100644 index 0000000..f3781d4 --- /dev/null +++ b/tests/util/test_features.py @@ -0,0 +1,79 @@ +"""Feature gating. + +`postprocessing` is the external web search step, enabled for everyone in the +shipped config, so these two conditions decide whether a Tavily call happens. +""" + +import pytest +from pydantic import ValidationError + +from util.config_yml import Config +from util.config_yml.features import Feature, Features, UserGroup + +LOGGED_IN = "someone@example.org" +GUEST = None + + +def test_user_group_all_matches_everyone() -> None: + feature = Feature(enabled=True, user_group=UserGroup.all) + assert feature.matches_user_group(LOGGED_IN) is True + assert feature.matches_user_group(GUEST) is True + + +def test_user_group_logged_in_excludes_guests() -> None: + feature = Feature(enabled=True, user_group=UserGroup.logged_in) + assert feature.matches_user_group(LOGGED_IN) is True + assert feature.matches_user_group(GUEST) is False + + +def test_omitted_user_group_matches_everyone() -> None: + """user_group is optional in .config.schema.yaml; absent means unrestricted.""" + assert Feature(enabled=True).matches_user_group(GUEST) is True + + +def test_unknown_user_group_is_rejected() -> None: + with pytest.raises(ValidationError): + Feature(enabled=True, user_group="admins") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("enabled", "user_group", "user_id", "expected"), + [ + (True, "all", LOGGED_IN, True), + (True, "all", GUEST, True), + (True, "logged_in", LOGGED_IN, True), + (True, "logged_in", GUEST, False), + (False, "all", LOGGED_IN, False), + (False, "logged_in", LOGGED_IN, False), + ], +) +def test_get_feature_combines_enabled_and_group( + enabled: bool, user_group: str, user_id: str | None, expected: bool +) -> None: + config = Config( + features=Features.model_validate( + {"postprocessing": {"enabled": enabled, "user_group": user_group}} + ), + messages={}, + profiles=[], + usage_limits={"message_rates": []}, # type: ignore[arg-type] + ) + assert config.get_feature("postprocessing", user_id) is expected + + +def test_unknown_feature_id_defaults_to_enabled() -> None: + """get_feature returns True for ids it does not know about. + + That is fail-open, so adding a call for a feature that is not in the model + silently enables it for everyone rather than raising. Pinning the behaviour + rather than endorsing it. + """ + config = Config( + features=Features.model_validate( + {"postprocessing": {"enabled": False, "user_group": "all"}} + ), + messages={}, + profiles=[], + usage_limits={"message_rates": []}, # type: ignore[arg-type] + ) + assert config.get_feature("a_feature_that_does_not_exist") is True diff --git a/tests/util/test_intervals.py b/tests/util/test_intervals.py new file mode 100644 index 0000000..e448640 --- /dev/null +++ b/tests/util/test_intervals.py @@ -0,0 +1,38 @@ +from datetime import timedelta + +import pytest + +from util.config_yml.intervals import parse_interval + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ("30s", timedelta(seconds=30)), + ("1m", timedelta(minutes=1)), + ("3h", timedelta(hours=3)), # the production usage_limits interval + ("7d", timedelta(days=7)), + ("2w", timedelta(weeks=2)), + ("0s", timedelta(0)), + ("100h", timedelta(hours=100)), + ], +) +def test_parses_each_supported_unit(text: str, expected: timedelta) -> None: + assert parse_interval(text) == expected + + +@pytest.mark.parametrize( + "text", + ["", "h", "3", "3x", "3 h", "-1h", "1.5h", "3H", "3hh", "1h30m"], +) +def test_malformed_intervals_raise(text: str) -> None: + """Malformed intervals must be loud. + + This used to return timedelta(0), which silently disabled rate limiting: a + zero-length window means every queued timestamp is already outside it, so the + queue drained on every call and nobody was ever limited. Config fields now + carry INTERVAL_PATTERN, so a bad value is rejected when config.yml loads and + never reaches here -- see test_usage_limits.py. + """ + with pytest.raises(ValueError, match="malformed interval"): + parse_interval(text) diff --git a/tests/util/test_messages.py b/tests/util/test_messages.py new file mode 100644 index 0000000..b8847f9 --- /dev/null +++ b/tests/util/test_messages.py @@ -0,0 +1,75 @@ +from datetime import UTC, datetime, timedelta + +import pytest +from pydantic import ValidationError + +from util.config_yml.messages import Message, Trigger, TriggerEvent + + +def test_event_trigger_matches_only_its_own_event() -> None: + trigger = Trigger(event=TriggerEvent.on_chat_start) + assert trigger.match_trigger(TriggerEvent.on_chat_start) is True + assert trigger.match_trigger(TriggerEvent.on_message) is False + assert trigger.match_trigger(None) is False + + +def test_after_messages_fires_on_exact_equality_only() -> None: + """Note: `!=`, not `<`. An `after_messages: 3` message fires on message 3 alone.""" + trigger = Trigger(after_messages=3) + assert trigger.match_trigger(after_messages=3) is True + assert trigger.match_trigger(after_messages=2) is False + assert trigger.match_trigger(after_messages=4) is False + + +def test_start_and_end_bound_the_active_window() -> None: + past = datetime.now() - timedelta(days=1) + future = datetime.now() + timedelta(days=1) + assert Trigger(start=past, end=future).match_trigger() is True + assert Trigger(start=future).match_trigger() is False + assert Trigger(end=past).match_trigger() is False + + +def test_timezone_aware_bounds_are_converted_not_stripped() -> None: + """`Z` timestamps used to have tzinfo discarded rather than converted. + + config_default.yml writes bounds as `2025-01-01T00:00:00Z`. Stripping tzinfo + shifted the window by the host's UTC offset -- invisible at day granularity, + wrong at hour granularity. The comparison is now done in UTC on both sides. + """ + now_utc = datetime.now(UTC) + just_past = (now_utc - timedelta(minutes=5)).isoformat() + just_future = (now_utc + timedelta(minutes=5)).isoformat() + + assert Trigger.model_validate({"start": just_past}).match_trigger() is True + assert Trigger.model_validate({"start": just_future}).match_trigger() is False + assert Trigger.model_validate({"end": just_future}).match_trigger() is True + assert Trigger.model_validate({"end": just_past}).match_trigger() is False + + +def test_naive_bounds_are_treated_as_utc() -> None: + """A bound written without an offset is assumed UTC rather than local.""" + past = (datetime.now(UTC) - timedelta(days=1)).replace(tzinfo=None) + assert Trigger(end=past).match_trigger() is False + assert Trigger(start=past).match_trigger() is True + + +def test_malformed_freq_max_is_rejected() -> None: + with pytest.raises(ValidationError): + Trigger(event=TriggerEvent.on_message, freq_max="1min") + + +def test_freq_max_throttles_on_last_send_time() -> None: + trigger = Trigger(event=TriggerEvent.on_message, freq_max="1m") + just_now = datetime.now().isoformat() + long_ago = (datetime.now() - timedelta(hours=1)).isoformat() + assert ( + trigger.match_trigger(TriggerEvent.on_message, last_message=just_now) is False + ) + assert trigger.match_trigger(TriggerEvent.on_message, last_message=long_ago) is True + assert trigger.match_trigger(TriggerEvent.on_message, last_message=None) is True + + +def test_message_defaults_to_enabled_and_unrestricted() -> None: + message = Message(message="hi", trigger=Trigger(event=TriggerEvent.on_chat_start)) + assert message.enabled is True + assert message.recipients is None diff --git a/tests/util/test_usage_limits.py b/tests/util/test_usage_limits.py new file mode 100644 index 0000000..33d018f --- /dev/null +++ b/tests/util/test_usage_limits.py @@ -0,0 +1,98 @@ +"""Characterization of the production rate limiter. + +config_default.yml ships `max_messages: 100` / `interval: 3h` for `users: ["all"]`, +so this is the code path that decides whether a real user gets an answer. +""" + +from datetime import datetime, timedelta + +import pytest +from pydantic import ValidationError + +from util.config_yml.usage_limits import MessageRate, UsageLimits + + +def _ago(**kwargs: float) -> str: + return (datetime.now() - timedelta(**kwargs)).isoformat() + + +def _rate(max_messages: int = 3, interval: str = "1h") -> MessageRate: + return MessageRate(users=["all"], max_messages=max_messages, interval=interval) + + +def test_under_the_limit_is_allowed_and_records_the_message() -> None: + queue = [_ago(minutes=5)] + assert _rate().check_rate(queue) is None + assert len(queue) == 2, "the allowed message must be recorded in the queue" + + +def test_at_the_limit_is_blocked_and_does_not_record() -> None: + queue = [_ago(minutes=m) for m in (5, 10, 15)] + rate = _rate(max_messages=3) + assert rate.check_rate(queue) is rate, "returns itself to signal rate-limited" + assert len(queue) == 3, "a blocked message must not consume quota" + + +def test_check_rate_mutates_the_callers_queue() -> None: + """This in-place mutation is the contract `message_rate_limited` relies on. + + `util.chainlit_helpers.message_rate_limited` reads the queue out of user + metadata, passes it in, and writes the *same list object* back. If check_rate + ever stopped mutating, quota tracking would silently stop working. + """ + queue: list[str] = [] + original = queue + _rate().check_rate(queue) + assert queue is original + assert len(queue) == 1 + + +def test_entries_older_than_the_interval_are_evicted() -> None: + queue = [_ago(hours=5), _ago(hours=4), _ago(minutes=1)] + assert _rate(max_messages=3, interval="1h").check_rate(queue) is None + # the two stale entries are dropped, the recent one survives, the new one is added + assert len(queue) == 2 + + +def test_eviction_stops_at_the_first_in_window_entry() -> None: + """The purge loop breaks on the first fresh entry rather than scanning the rest. + + The queue is only ever appended to in chronological order, so this is sound -- + but it means an out-of-order queue would retain stale entries. + """ + queue = [_ago(minutes=1), _ago(hours=5)] + _rate(max_messages=9, interval="1h").check_rate(queue) + assert len(queue) == 3, "the stale entry behind a fresh one is not evicted" + + +@pytest.mark.parametrize("bad", ["3hr", "3", "h", "", "1h30m", "-1h"]) +def test_malformed_interval_is_rejected_at_construction(bad: str) -> None: + """A typo in config.yml must fail loudly instead of disabling the limiter. + + parse_interval() used to return timedelta(0) for anything unparseable, which + made the window zero-length: the queue drained on every call and no user was + ever limited. The field now carries the same pattern .config.schema.yaml + documents, so the config is rejected at load time instead. + """ + with pytest.raises(ValidationError): + MessageRate(users=["all"], max_messages=1, interval=bad) + + +def test_max_messages_must_be_positive() -> None: + """max_messages: 0 would block everyone; treat it as a config error.""" + with pytest.raises(ValidationError): + MessageRate(users=["all"], max_messages=0, interval="1h") + + +def test_production_interval_is_accepted() -> None: + assert _rate(max_messages=100, interval="3h").interval == "3h" + + +def test_first_matching_rule_wins() -> None: + limits = UsageLimits( + message_rates=[ + MessageRate(users=["all"], max_messages=1, interval="1h"), + MessageRate(users=["logged_in"], max_messages=100, interval="1h"), + ] + ) + assert limits.message_rates[0].max_messages == 1 diff --git a/tests/util/test_user_matching.py b/tests/util/test_user_matching.py new file mode 100644 index 0000000..32b8de7 --- /dev/null +++ b/tests/util/test_user_matching.py @@ -0,0 +1,52 @@ +import pytest + +from util.config_yml.user_matching import match_user + +LOGGED_IN = "someone@example.org" +GUEST = None + + +@pytest.mark.parametrize("spec", [None, []]) +def test_empty_spec_matches_everyone(spec: list[str] | None) -> None: + assert match_user(spec, LOGGED_IN) is True + assert match_user(spec, GUEST) is True + + +def test_all_matches_everyone() -> None: + assert match_user(["all"], LOGGED_IN) is True + assert match_user(["all"], GUEST) is True + + +def test_guests_and_logged_in_are_mutually_exclusive() -> None: + assert match_user(["guests"], GUEST) is True + assert match_user(["guests"], LOGGED_IN) is False + assert match_user(["logged_in"], LOGGED_IN) is True + assert match_user(["logged_in"], GUEST) is False + + +def test_glob_patterns_match_on_the_full_identifier() -> None: + assert match_user(["*@gmail.com"], "person@gmail.com") is True + assert match_user(["*@gmail.com"], "person@oicr.on.ca") is False + assert match_user(["*@oicr.on.ca", "*@gmail.com"], "person@gmail.com") is True + + +def test_slash_delimited_entries_are_treated_as_regex() -> None: + """Undocumented: /.../ entries are regex, which `.config.schema.yaml` disallows.""" + assert match_user(["/^admin-/"], "admin-jane") is True + assert match_user(["/^admin-/"], "jane-admin") is False + # note: re.search, not fullmatch -- the pattern is unanchored by default + assert match_user(["/oicr/"], "person@oicr.on.ca") is True + + +def test_guest_never_matches_identifier_patterns() -> None: + assert match_user(["*@gmail.com"], GUEST) is False + assert match_user(["logged_in", "*@gmail.com"], GUEST) is False + + +def test_empty_entry_is_skipped_not_a_crash() -> None: + """`entry[0]` used to index without a length check, so `users: [""]` raised.""" + assert match_user([""], LOGGED_IN) is False + assert match_user([""], GUEST) is False + assert ( + match_user(["", "all"], LOGGED_IN) is True + ), "a real entry after it still counts"