AgentiGrid: optional retrieval grounding (RAG) for LLM proposals

agentigrid/.gitignore

@@ -114,3 +114,7 @@ desktop.ini
 AGENTS.md
 opencode.md
 /prompts/
+
+# RAG index (build locally: python -m agentigrid.rag.ingest rag/corpus)
+rag/store/
+rag/store_*/

agentigrid/README.md

@@ -286,6 +286,10 @@ When resuming, the goal and journal are loaded from the saved session, but the L
 
 The launcher sidebar includes a "Session Save/Resume" section with a Save button (available after search completes or while running) and a dropdown to resume from previously saved sessions.
 
+## Retrieval Grounding (optional RAG)
+
+AgentiGrid can ground the LLM's proposals in a curated, provenance-tagged corpus (`rag/corpus/`: ExaGO tool help, case metadata, certified response examples). It is **off by default**, touches only the generation step (never the deterministic validator), and degrades to the no-RAG baseline on any retrieval failure. Modes are `off`, `basic` and `corrective`, set with `AGENTIGRID_RAG_MODE` or the UI's **Advanced** panel. Setup, optional dependencies (`requirements-rag.txt`) and troubleshooting: [docs/RAG_USAGE.md](docs/RAG_USAGE.md).
+
 ## Interactive Steering
 
 While a search is running, you can inject steering directives from the terminal (CLI) or the GUI — without stopping and restarting the search.

agentigrid/agentigrid/exago_engine/agent_loop.py

@@ -47,6 +47,7 @@
     results_summary_for_app,
 )
 from agentigrid.exago_parsers.matpower_model import MATNetwork
+from agentigrid.rag import build_retriever
 from agentigrid.exago_parsers.opflow_results import OPFLOWResult
 from agentigrid.exago_engine import topology
 from agentigrid.exago_engine import contingency
@@ -771,6 +772,13 @@ def __init__(
         # Metered so per-run token totals (incl. prompt-cache reads/writes)
         # cover every LLM call, not only the main per-iteration call.
         self._backend: LLMBackend = UsageMeter(create_backend(config.llm))
+        # Optional retrieval grounding (off by default). AGENTIGRID_RAG_MODE is one
+        # of off | basic | corrective (AGENTIGRID_RAG=1/0 still maps to basic/off).
+        # Every mode exposes the same .enabled / .retrieve() surface, and "off"
+        # needs none of the optional RAG dependencies.
+        self._retriever = build_retriever(
+            host=os.environ.get("OLLAMA_HOST") or getattr(config.llm, "ollama_host", None) or "http://localhost:11434",
+        )
         self._executor = SimulationExecutor(config.exago, config.output)
         self._journal = SearchJournal()
         self._quiet = quiet
@@ -1457,7 +1465,22 @@ def _iteration(
         )
         self._error_feedback = None  # consumed
 
+        # Retrieval grounds the generation step only; the deterministic validator
+        # never sees it. A failed or empty retrieval leaves the prompt unchanged.
+        _retrieved = self._retriever.retrieve(goal) if self._retriever.enabled else ""
+        if _retrieved:
+            user_prompt = (
+                "=== Section B: Reference Material (retrieved) ===\n"
+                f"{_retrieved}\n\n{user_prompt}"
+            )
+
         if self._on_phase:
+            if self._retriever.enabled:
+                # "rag_retrieved:<refs>:<top score>" feeds the UI's grounding caption
+                # (0 refs under corrective = context withheld).
+                _scores = [float(x) for x in re.findall(r"score (\d+\.\d+)", _retrieved)]
+                _top = max(_scores) if _scores else 0.0
+                self._on_phase(iteration, f"rag_retrieved:{_retrieved.count('[ref ')}:{_top:.2f}")
             self._on_phase(iteration, "llm_request")
 
         response = self._backend.complete(system_prompt, user_prompt)
@@ -4999,6 +5022,19 @@ def _assemble_prompt(
 
     def _finalize(self, session: SearchSession, elapsed_seconds: float) -> None:
         """Print summary and save journal."""
+        _rag_enabled = self._retriever.enabled
+        self._journal.rag_enabled = _rag_enabled
+        try:
+            from agentigrid.rag import resolve_rag_mode, resolve_crag_grader
+            _desc = getattr(self._retriever, "describe", None)
+            self._journal.rag_config = {
+                **(_desc() if callable(_desc) else {}),
+                "mode_env": resolve_rag_mode(),
+                "grader_env": resolve_crag_grader(),
+            }
+        except Exception:
+            self._journal.rag_config = None
+        session.rag_enabled = _rag_enabled
         _totals = getattr(self._backend, "totals", None)
         if callable(_totals):
             self._journal.llm_usage = {

agentigrid/agentigrid/exago_engine/journal.py

@@ -1131,8 +1131,12 @@ def export_json(self, path: Path) -> None:
             data["session_best"] = self.session_best
         if self.load_factor is not None:
             data["load_factor"] = self.load_factor
+        if getattr(self, "rag_enabled", None) is not None:
+            data["rag_enabled"] = self.rag_enabled
         if getattr(self, "llm_usage", None) is not None:
             data["llm_usage"] = self.llm_usage
+        if getattr(self, "rag_config", None) is not None:
+            data["rag_config"] = self.rag_config
         if getattr(self, "discarded_actions", None):
             data["discarded_actions"] = self.discarded_actions
         path.write_text(json.dumps(data, indent=2), encoding="utf-8")

agentigrid/agentigrid/rag/__init__.py

@@ -0,0 +1,108 @@
+"""RAG (retrieval-augmented generation) support for AgentiGrid.
+
+Generation-stage only: retrieval grounds the LLM's spec/proposal generation.
+It is NEVER used by the deterministic validator (engine/validation.py).
+
+Retrieval is OFF by default. The mode is selected by an env var:
+
+    AGENTIGRID_RAG_MODE = off | basic | corrective
+
+For backward compatibility the legacy switch still works:
+    AGENTIGRID_RAG=1  ->  basic
+    AGENTIGRID_RAG=0/unset (and no _MODE) -> off
+
+Corrective-mode grader (AGENTIGRID_CRAG_GRADER): cosine (default) | reranker
+(local cross-encoder, needs sentence-transformers). Thresholds:
+AGENTIGRID_CRAG_TAU_LOWER / AGENTIGRID_CRAG_TAU_UPPER.
+
+Use `build_retriever(**kwargs)` to get the retriever for the active mode; every
+mode exposes the same surface (`.enabled` and `.retrieve(query) -> str`).
+"""
+from __future__ import annotations
+
+import os
+
+from .retriever import Retriever
+from .corrective import CorrectiveRetriever
+
+VALID_MODES = ("off", "basic", "corrective")
+# Graders implemented behind the corrective-RAG seam.
+VALID_CRAG_GRADERS = ("cosine", "reranker")
+
+
+def resolve_rag_mode() -> str:
+    """Return the active RAG mode, honoring AGENTIGRID_RAG_MODE and falling back
+    to the legacy AGENTIGRID_RAG switch. Unknown values degrade to 'basic'.
+    """
+    m = os.environ.get("AGENTIGRID_RAG_MODE")
+    if m:
+        m = m.strip().lower()
+        return m if m in VALID_MODES else "basic"
+    return "basic" if os.environ.get("AGENTIGRID_RAG", "0") == "1" else "off"
+
+
+def resolve_crag_grader() -> str:
+    """Return the corrective-RAG grader name from AGENTIGRID_CRAG_GRADER.
+    'cosine' (default/unset) or 'reranker'. Unknown values fall back to 'cosine'.
+    """
+    name = os.environ.get("AGENTIGRID_CRAG_GRADER", "").strip().lower()
+    return name if name in VALID_CRAG_GRADERS else "cosine"
+
+
+def _make_crag_grader():
+    """Build the optional grader for corrective mode (selected by
+    AGENTIGRID_CRAG_GRADER). 'cosine' -> None (deterministic cosine buckets, the
+    reproducible default). 'reranker' -> RerankerGrader, imported lazily so
+    sentence-transformers is only needed when it is actually requested.
+    """
+    name = resolve_crag_grader()
+    if name == "reranker":
+        from .grader_reranker import RerankerGrader
+        return RerankerGrader()  # loads lazily; failures are counted, not hidden
+    return None  # 'cosine' (and any unknown value already normalized)
+
+
+def _tau_overrides() -> dict:
+    """Corrective thresholds from AGENTIGRID_CRAG_TAU_LOWER / _UPPER, if set.
+
+    Each grader's score scale differs (cosine vs. reranker score vs. calibrated
+    probability), so its thresholds are fixed on labeled data and passed per
+    run; unset = the module defaults (0.30 / 0.50, cosine).
+    """
+    out = {}
+    for env, key in (("AGENTIGRID_CRAG_TAU_LOWER", "tau_lower"), ("AGENTIGRID_CRAG_TAU_UPPER", "tau_upper")):
+        val = os.environ.get(env)
+        if val not in (None, ""):
+            out[key] = float(val)
+    if "tau_lower" in out and "tau_upper" in out and out["tau_lower"] > out["tau_upper"]:
+        raise ValueError("AGENTIGRID_CRAG_TAU_LOWER must not exceed AGENTIGRID_CRAG_TAU_UPPER")
+    return out
+
+
+def build_retriever(**base_kwargs):
+    """Construct the retriever for the active mode.
+
+    `base_kwargs` are passed to the basic Retriever (e.g. host=, path=, model=).
+    Returns a Retriever (off/basic) or a CorrectiveRetriever (corrective); both
+    share the `.enabled` / `.retrieve()` surface.
+    """
+    mode = resolve_rag_mode()
+    base_kwargs.pop("enabled", None)
+    # Corpus selection (e.g. core vs. a larger documentation corpus): each
+    # corpus is ingested into its own store directory.
+    if os.environ.get("AGENTIGRID_RAG_STORE"):
+        base_kwargs["path"] = os.environ["AGENTIGRID_RAG_STORE"]
+    if os.environ.get("AGENTIGRID_RAG_COLLECTION"):
+        base_kwargs["collection"] = os.environ["AGENTIGRID_RAG_COLLECTION"]
+    if mode == "off":
+        return Retriever(enabled=False)
+    base = Retriever(enabled=True, **base_kwargs)
+    if mode == "corrective":
+        return CorrectiveRetriever(base, grader=_make_crag_grader(), **_tau_overrides())
+    return base  # 'basic' (and any unknown value already normalized to basic)
+
+
+__all__ = [
+    "Retriever", "CorrectiveRetriever", "build_retriever",
+    "resolve_rag_mode", "resolve_crag_grader",
+]

agentigrid/agentigrid/rag/corpus_hash.py

@@ -0,0 +1,73 @@
+"""Corpus identity: which files form a corpus and its SHA-256.
+
+Shared by rag/tools/corpus_guard.py (freeze), agentigrid.rag.ingest (records the
+hash of what it indexed) and corpus_status() (checks that a store was built
+from the frozen corpus it claims). One definition, so they can never disagree.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+
+CORPUS_SUFFIXES = (".txt", ".md")          # what ingest indexes
+MANIFEST_NAME = "corpus_manifest.json"     # written by corpus_guard --freeze
+INGEST_MANIFEST_NAME = "ingest_manifest.json"  # written by ingest into the store dir
+
+# How ingest splits the corpus into chunks. Part of the store's identity: the same
+# corpus chunked differently retrieves different text, so corpus_status refuses a
+# store built with another scheme. v2 keeps every paragraph (one worked example,
+# one schema exemplar) whole up to CHUNK_SIZE characters; v1 (800/100) cut most
+# worked examples mid-JSON.
+CHUNKING = {"scheme": "paragraph-v2", "size": 3000, "overlap": 200}
+
+
+def corpus_files(corpus_dir: Path) -> list[Path]:
+    corpus_dir = Path(corpus_dir)
+    return sorted(p for p in corpus_dir.rglob("*")
+                  if p.is_file() and p.suffix in CORPUS_SUFFIXES and p.name != MANIFEST_NAME)
+
+
+def corpus_hash(corpus_dir: Path) -> tuple[str, dict[str, str]]:
+    """Overall SHA-256 over (relative path, content hash) pairs, order-independent of disk."""
+    corpus_dir = Path(corpus_dir)
+    per_file = {}
+    for p in corpus_files(corpus_dir):
+        per_file[p.relative_to(corpus_dir).as_posix()] = hashlib.sha256(p.read_bytes()).hexdigest()
+    h = hashlib.sha256()
+    for rel in sorted(per_file):
+        h.update(f"{rel}\0{per_file[rel]}\n".encode())
+    return h.hexdigest(), per_file
+
+
+def corpus_status(corpus_dir: Path, store_dir: Path) -> dict:
+    """Is *corpus_dir* frozen, unchanged since freezing, and is *store_dir* built
+    from exactly that corpus? Returns {"ok": bool, "reason": str|None, "corpus_sha256": ...}."""
+    corpus_dir, store_dir = Path(corpus_dir), Path(store_dir)
+    out = {"ok": False, "reason": None, "corpus_dir": str(corpus_dir), "store_dir": str(store_dir),
+           "corpus_sha256": None}
+    mpath = corpus_dir / MANIFEST_NAME
+    if not mpath.exists():
+        out["reason"] = f"corpus not frozen ({mpath} missing; run corpus_guard.py --freeze)"
+        return out
+    frozen = json.loads(mpath.read_text()).get("corpus_sha256")
+    digest, _ = corpus_hash(corpus_dir)
+    out["corpus_sha256"] = digest
+    if digest != frozen:
+        out["reason"] = "corpus changed since it was frozen (re-run corpus_guard.py --freeze, then ingest)"
+        return out
+    ipath = store_dir / INGEST_MANIFEST_NAME
+    if not ipath.exists():
+        out["reason"] = f"store not built by the current ingest ({ipath} missing; re-run ingest)"
+        return out
+    ing = json.loads(ipath.read_text())
+    if ing.get("corpus_sha256") != digest:
+        out["reason"] = "store was built from a different corpus version (re-run ingest)"
+        return out
+    built = ing.get("chunking") or {}
+    if {k: built.get(k) for k in CHUNKING} != CHUNKING:
+        out["reason"] = (f"store was built with different chunking ({built or 'v1, 800/100'}; "
+                         f"current {CHUNKING}); re-run ingest")
+        return out
+    out["ok"] = True
+    return out