Skip to content

Spec: Conformal Query Completeness for Graffold

Status: Draft
Author: —
Date: 2026-06-08

1. Problem

Graffold's query layer returns whatever it finds — vector similarity hits, fulltext matches, graph traversals — but provides no guarantee about what it missed. For a query like "find all proteins associated with Parkinson's disease," the system might return 12 proteins when 30 exist in the literature. The user has no signal about completeness.

This matters because: - Biomedical researchers make decisions based on what's returned (drug screening, pathway analysis) - The KG is inherently incomplete — we only ingest a fraction of published literature - Different search strategies (vector, fulltext, graph traversal) have different recall profiles - Missing a relevant protein-disease relationship can be worse than including a spurious one

Goal: Give users a declarative recall target (e.g., "I want ≥90% of relevant results") and have the system calibrate its retrieval to meet that target with statistical validity, while maximizing precision.

2. Approach

Adapt the ConRAD framework (Horchidan et al., 2026) to Graffold's architecture:

  1. Add a link prediction model that scores unobserved (entity, relation, entity) triples
  2. Introduce a conformal gate at the search strategy level — use exact graph retrieval when local evidence is sufficient, fall back to neural prediction when it isn't
  3. Offline calibration — derive per-strategy thresholds from a held-out dataset that satisfy recall targets
  4. Expose recall budgets as a query parameter — users or agents specify desired completeness

3. Architecture

┌─────────────────────────────────────────────────────────┐
│                    Query Service                         │
│                                                         │
│  query + recall_target(α)                               │
│         │                                               │
│         ▼                                               │
│  ┌─────────────────┐                                    │
│  │  Query Planner   │ ← calibration lookup (topology, α)│
│  └────────┬────────┘                                    │
│           │ threshold vector λ = [λ₁, ..., λₖ]         │
│           ▼                                             │
│  ┌─────────────────────────────────────┐                │
│  │       Conformal Gate (per hop)       │               │
│  │                                      │               │
│  │  if local_density ≥ λᵢ:             │               │
│  │    → exact retrieval only            │               │
│  │  else:                               │               │
│  │    → retrieval ∪ neural predictions  │               │
│  └─────────────┬───────────────────────┘                │
│                │                                        │
│    ┌───────────┼───────────┐                            │
│    ▼           ▼           ▼                            │
│  Graph      Vector     Link Prediction                  │
│  Traversal  Search     Model (new)                      │
│  (Neo4j)    (existing) (KG embeddings)                  │
│                                                         │
└─────────────────────────────────────────────────────────┘

New Components

Component Location Purpose
LinkPredictor src/ml/link_predictor.py Score unobserved triples (h, r, t) → [0,1]
ConformalCalibrator src/ml/conformal_calibrator.py Offline: derive threshold vectors from calibration set
ConformalGate src/agents/conformal_gate.py Runtime: route between retrieval and neural prediction per hop
CalibrationStore src/cache/calibration_store.py Persist calibrated thresholds per (KG, topology, α)

Modified Components

Component Change
QueryService Accept optional recall_target param, pass to planner
TwoPhaseAgent / DynamicQueryAgent Wrap search strategies with conformal gate
Query API models Add recall_target: float | None field

Requirements

  • Score arbitrary (head, relation, tail) triples
  • Work with Graffold's entity ID scheme (UniProt IDs, MONDO IDs, etc.)
  • Inductive: handle entities not seen during training (new proteins from recent ingestion)
  • Lightweight: inference must be fast enough for interactive queries (~100ms for batch scoring)

Options (ranked by fit)

Model Inductive? Install Latency Notes
ULTRA Yes (foundation model) git clone (no pip) ~50ms/batch Zero-shot, but requires PyG + TorchDrug
PyKEEN RotatE No (transductive) pip install pykeen ~5-20ms/batch Production-ready, needs training per KG
PyKEEN ComplEx No (transductive) pip install pykeen ~5ms/batch Simpler, often competitive
NBFNet Yes (GNN) via TorchDrug ~200ms/batch Slower but handles arbitrary topologies

Revised recommendation: Start with PyKEEN RotatE for the spike. ULTRA isn't pip-installable and has heavy dependencies (PyG, TorchDrug). PyKEEN trains fast on our data, and the transductive limitation is acceptable since our KG doesn't change between calibrations. If we need zero-shot for new entities from fresh ingestions, revisit ULTRA or use PyKEEN's re-training with warm-start.

Embedding Strategy

The link predictor needs entity and relation embeddings: - Entities: Reuse Graffold's existing graph embeddings (ingest_main.py --add-graph-embeddings) - Relations: Derive from relationship type names + properties - Scoring: score(h, r, t) = model(embed(h), embed(r), embed(t)) → [0,1]

5. Conformal Calibration

Calibration Data

For each KG, we need a calibration set of queries with known ground truth answers:

Source How
Existing query logs Historical queries where we later verified completeness
Synthetic from dense subgraphs Sample entities, remove edges, ask "find all X related to Y"
Curated benchmarks Use established protein-disease datasets (e.g., DisGeNET, OpenTargets) as ground truth
User feedback "I know these 5 proteins are relevant" → calibration signal

Minimum calibration set: ~2000 queries per topology (per ConRAD's evaluation).

Calibration Procedure

For each query topology T and risk budget α:
  1. Execute queries on calibration set using both retrieval and neural scoring
  2. Compute unified scores (retrieval facts scored high, neural predictions scored by model)
  3. Search for threshold vector λ that satisfies:
     - Expected recall ≥ (1 - α) over calibration set
     - Minimizes result set cardinality (maximizes precision)
  4. Store (KG_id, topology, α) → λ in CalibrationStore

Query Topologies in Graffold

Map Graffold's query patterns to calibratable topologies:

Graffold Pattern Topology Example
Direct lookup 1-hop projection "What diseases is BRCA1 associated with?"
Two-hop path 2-hop chain "What drugs target proteins in the MAPK pathway?"
Intersection 2ip "Proteins associated with BOTH Alzheimer's AND diabetes"
Union 2u "Proteins from study A OR study B"
Multi-hop 3p "Drugs → proteins → pathways → diseases"

6. Conformal Gate Logic

The gate decides per-operator whether to use retrieval only or retrieval + neural prediction:

class ConformalGate:
    def execute(self, input_entities, relation, threshold_λ, routing_δ=0.5):
        # Always do exact retrieval
        retrieval_results = self.graph_db.traverse(input_entities, relation)

        # Score retrieval results on unified scale [δ, 1.0]
        scored_retrieval = self.score_retrieval(retrieval_results, routing_δ)

        if threshold_λ >= routing_δ:
            # Retrieval-only mode: local graph is dense enough
            return filter(scored_retrieval, threshold_λ)
        else:
            # Hybrid mode: need neural predictions to meet recall target
            neural_candidates = self.link_predictor.predict(input_entities, relation)
            # Neural scores mapped to [0, δ)
            scored_neural = self.score_neural(neural_candidates, routing_δ)
            combined = scored_retrieval + filter(scored_neural, threshold_λ)
            return combined

Integration with Existing Agents

The conformal gate wraps individual search operations inside existing agents:

Before:  agent.search() → vector_results ∪ graph_results ∪ fulltext_results
After:   agent.search() → conformal_gate(vector, graph, fulltext, neural, threshold)

The agent's existing strategy selection (local/global/hybrid mode) remains. The conformal gate adds a completeness layer on top, filling gaps that individual strategies miss.

7. API Surface

Query Endpoint Changes

POST /v1/query
{
  "query": "What proteins are associated with Parkinson's disease?",
  "completeness": "thorough",    // NEW: "fast" | "thorough" | "exhaustive" (default: "fast")
  "recall_target": null,         // NEW: optional float override (0.0-1.0)
  "max_results": 50,
  "database": "olink3"
}

Named levels map to recall targets: | Level | α (risk budget) | Recall target | Latency budget | |-------|----------------|---------------|----------------| | fast | 1.0 (no guarantee) | best-effort | baseline | | thorough | 0.2 | ≥80% | ≤10% overhead | | exhaustive | 0.1 | ≥90% | ≤2x baseline |

Response additions: 

{
  "answer": "...",
  "sources": [...],
  "completeness": {                    // NEW
    "level": "thorough",
    "recall_target": 0.8,
    "estimated_recall": 0.84,
    "neural_augmented": true,
    "retrieval_only_hops": [1],
    "hybrid_hops": [2, 3]
  }
}

Calibration Admin Endpoint

POST /v1/platform/calibration/run
{
  "database": "olink3",
  "topologies": ["1p", "2p", "3p", "2ip"],
  "risk_budgets": [0.1, 0.2, 0.3, 0.4]
}

8. Calibration Data Generation

The hardest practical challenge: getting ground truth for calibration. Strategy:

Phase 1: Synthetic Calibration (MVP)

  • Take a dense subgraph of the existing KG
  • Randomly remove 10-40% of edges
  • Generate queries over the full graph, evaluate against the incomplete version
  • This simulates the "missing literature" problem

Phase 2: External Ground Truth

  • Import DisGeNET protein-disease associations as reference
  • Import OpenTargets drug-target-disease paths
  • Compare what's in our KG vs what these databases say → ground truth completeness

Phase 3: Live Feedback Loop

  • Track user corrections ("you missed protein X")
  • Track new ingestions that add relationships matching past queries
  • Accumulate calibration signal over time

9. Scope & Phases

  • [ ] Export triples from existing KG: uv run python scripts/export_kg_triples.py --database olink3
  • [ ] Install PyKEEN: pip install pykeen
  • [ ] Train RotatE on exported triples: uv run python scripts/spike_link_prediction.py
  • [ ] Measure MRR, Hits@10, Hits@50 on held-out test set
  • [ ] Go/no-go: MRR ≥ 0.3 and Hits@10 ≥ 0.4 → proceed. Below that → try ComplEx, increase epochs, or consider ULTRA.
  • [ ] Verify scoring latency: <200ms for batch of 1000 candidates
  • [ ] If go: save trained model for Phase A integration

Phase A — Foundation (2-3 weeks)

  • [ ] Integrate UltraQuery (or equivalent) link prediction model
  • [ ] Implement LinkPredictor with batch scoring API
  • [ ] Generate synthetic calibration datasets from existing KGs
  • [ ] Implement ConformalCalibrator (offline threshold derivation)
  • [ ] Implement CalibrationStore (Redis-backed threshold lookup)

Phase B — Query Integration (2-3 weeks)

  • [ ] Implement ConformalGate operator
  • [ ] Wire into TwoPhaseAgent as an additional search strategy
  • [ ] Add recall_target parameter to query API
  • [ ] Return completeness metadata in responses
  • [ ] Add calibration admin endpoint

Phase C — Production Hardening (2-3 weeks)

  • [ ] External ground truth import (DisGeNET, OpenTargets)
  • [ ] Per-KG calibration scheduling (recalibrate on new ingestions)
  • [ ] Monitoring: track empirical recall vs targets in production
  • [ ] Fallback behavior when calibration data is insufficient
  • [ ] Documentation and user-facing recall target guidance

10. Open Questions

  1. Model selection: ✅ Resolved — Start with UltraQuery zero-shot. Spike on actual KG data first to validate performance before committing to integration. If zero-shot MRR < 0.3 on our biomedical triples, pivot to a domain-specific model.

  2. Calibration freshness: How often does recalibration need to happen? After every ingestion job? Daily? Only when graph structure changes significantly?

  3. User experience: ✅ Resolved — Named quality levels, not raw numbers:

  4. "fast" — retrieval only, no recall guarantee (current behavior, α=1.0)
  5. "thorough" — moderate recall target (α=0.2, targeting ≥80% recall)
  6. "exhaustive" — strict recall target (α=0.1, targeting ≥90% recall)

Power users can override with explicit recall_target float if needed.

  1. Scope of neural prediction: Should the link predictor only predict missing existing relation types, or also infer novel relation types between entities?

  2. Latency budget: ✅ Resolved — Target ≤10% latency overhead for "thorough" mode. This means:

  3. Current queries: ~2-5s → "thorough" must stay under ~2.2-5.5s
  4. "Exhaustive" gets more budget: up to ~2x (4-10s) acceptable
  5. This constrains the link predictor to batch scoring in <200ms

  6. The conformal gate must bypass neural inference aggressively in dense regions

  7. Multi-tenant calibration: Does each tenant KG need its own calibration, or can we share across structurally similar KGs?

  8. Evaluation strategy: How do we validate this works before shipping? Run against a held-out subset of a known-complete reference DB?

11. Success Criteria

  • Empirical recall meets declared target within ±5% across 90% of query topologies
  • Precision ≥ 80% at recall target of 0.9 (on synthetic benchmark)
  • "Thorough" mode latency overhead ≤ 10% vs baseline retrieval-only
  • "Exhaustive" mode latency overhead ≤ 2x vs baseline
  • Zero false confidence: system never reports higher estimated recall than actual
  • Calibration runs in < 30 minutes per KG per topology
  • Link predictor batch scoring < 200ms for 1000 candidate triples