From b4c465e60896a15862bbebc77aeb535420d80ed0 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 23 Jun 2026 00:39:27 +0100 Subject: [PATCH] feat(graphrag): config.graph sub-model + settings + pgvector availability gate GraphConfig (extraction_model?, max_chunks?, gleanings=0; extra=forbid) on SourceConfig; settings GRAPHRAG_EXTRACTION_MODEL (None=instance default) + GRAPHRAG_MAX_CHUNKS_FOR_EXTRACTION; graphrag_available() = GRAPHRAG_ENABLED and VECTOR_STORE==pgvector (D27-D32 foundation). Unit G1. --- application/core/settings.py | 5 +++ application/graphrag/__init__.py | 10 +++++ application/storage/db/source_config.py | 31 +++++++++++++ tests/graphrag/__init__.py | 0 tests/graphrag/test_graphrag_available.py | 31 +++++++++++++ tests/storage/db/test_source_config.py | 53 +++++++++++++++++++++++ 6 files changed, 130 insertions(+) create mode 100644 application/graphrag/__init__.py create mode 100644 tests/graphrag/__init__.py create mode 100644 tests/graphrag/test_graphrag_available.py diff --git a/application/core/settings.py b/application/core/settings.py index 0541571a..65b9314c 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -106,6 +106,11 @@ class Settings(BaseSettings): # Flagship GraphRAG flag. Reserved and unused for now; gates graph-aware # ingestion/retrieval when that feature lands. GRAPHRAG_ENABLED: bool = False + # Model for ingest-time graph extraction; None reuses the instance default + # model (LLM_PROVIDER/LLM_NAME). Operator-overridable (e.g. a cheaper model). + GRAPHRAG_EXTRACTION_MODEL: Optional[str] = None + # Hard cap on chunks extracted per source (cost control). + GRAPHRAG_MAX_CHUNKS_FOR_EXTRACTION: int = 2000 AGENT_NAME: str = "classic" FALLBACK_LLM_PROVIDER: Optional[str] = None # provider for fallback llm FALLBACK_LLM_NAME: Optional[str] = None # model name for fallback llm diff --git a/application/graphrag/__init__.py b/application/graphrag/__init__.py new file mode 100644 index 00000000..9931b525 --- /dev/null +++ b/application/graphrag/__init__.py @@ -0,0 +1,10 @@ +"""GraphRAG feature package (flag-gated, pgvector-only per D29).""" + +from __future__ import annotations + +from application.core.settings import settings + + +def graphrag_available() -> bool: + """Return True when GraphRAG is enabled and the store is pgvector (D29).""" + return settings.GRAPHRAG_ENABLED and settings.VECTOR_STORE == "pgvector" diff --git a/application/storage/db/source_config.py b/application/storage/db/source_config.py index 845968af..6ee354d6 100644 --- a/application/storage/db/source_config.py +++ b/application/storage/db/source_config.py @@ -49,6 +49,36 @@ class PreScreenConfig(BaseModel): return self +class GraphConfig(BaseModel): + """Ingest-time GraphRAG extraction knobs (D28; pgvector-only per D29). + + ``extraction_model`` None reuses the instance default model + (``LLM_PROVIDER``/``LLM_NAME``); ``max_chunks`` None falls back to the + ``GRAPHRAG_MAX_CHUNKS_FOR_EXTRACTION`` setting. ``gleanings`` is off by + default (cost control). + """ + + model_config = ConfigDict(extra="forbid") + + extraction_model: Optional[str] = None # None → instance default model + max_chunks: Optional[int] = None # None → GRAPHRAG_MAX_CHUNKS_FOR_EXTRACTION + gleanings: int = 0 # extra extraction passes per chunk + + @field_validator("max_chunks") + @classmethod + def _positive_max_chunks(cls, value: Optional[int]) -> Optional[int]: + if value is not None and value < 1: + raise ValueError("must be >= 1") + return value + + @field_validator("gleanings") + @classmethod + def _non_negative_gleanings(cls, value: int) -> int: + if value < 0: + raise ValueError("must be >= 0") + return value + + class ChunkingConfig(BaseModel): """Ingest-time chunking knobs (bake-time; change requires re-ingest).""" @@ -116,6 +146,7 @@ class SourceConfig(BaseModel): kind: str = "classic" # behavior selector: classic | wiki | graphrag | ... chunking: ChunkingConfig = ChunkingConfig() retrieval: RetrievalConfig = RetrievalConfig() + graph: GraphConfig = GraphConfig() def wiki_enabled(self) -> dict: """Return a config dict flipped to wiki mode + browse-as-you-go exposure. diff --git a/tests/graphrag/__init__.py b/tests/graphrag/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/graphrag/test_graphrag_available.py b/tests/graphrag/test_graphrag_available.py new file mode 100644 index 00000000..c981ea49 --- /dev/null +++ b/tests/graphrag/test_graphrag_available.py @@ -0,0 +1,31 @@ +"""Tests for the graphrag_available() pgvector-gated flag (D29).""" + +from __future__ import annotations + +import pytest + +from application.graphrag import graphrag_available +from application.core.settings import settings + + +@pytest.mark.unit +class TestGraphragAvailable: + def test_true_only_when_enabled_and_pgvector(self, monkeypatch): + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True) + monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector") + assert graphrag_available() is True + + def test_false_when_disabled(self, monkeypatch): + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False) + monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector") + assert graphrag_available() is False + + def test_false_when_store_not_pgvector(self, monkeypatch): + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True) + monkeypatch.setattr(settings, "VECTOR_STORE", "faiss") + assert graphrag_available() is False + + def test_false_when_disabled_and_not_pgvector(self, monkeypatch): + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False) + monkeypatch.setattr(settings, "VECTOR_STORE", "faiss") + assert graphrag_available() is False diff --git a/tests/storage/db/test_source_config.py b/tests/storage/db/test_source_config.py index f1883697..ab0cf28b 100644 --- a/tests/storage/db/test_source_config.py +++ b/tests/storage/db/test_source_config.py @@ -7,6 +7,7 @@ from pydantic import ValidationError from application.storage.db.source_config import ( ChunkingConfig, + GraphConfig, PreScreenConfig, RetrievalConfig, SourceConfig, @@ -21,6 +22,7 @@ class TestParseLenient: assert cfg.kind == "classic" assert cfg.chunking == ChunkingConfig() assert cfg.retrieval == RetrievalConfig() + assert cfg.graph == GraphConfig() def test_parse_none_is_all_defaults(self): cfg = SourceConfig.parse(None) @@ -148,3 +150,54 @@ class TestPreScreenConfig: # A garbage stored prescreen blob never crashes the read path. rc = RetrievalConfig.model_construct(prescreen={"candidate_k": "x"}) assert rc.prescreen_config() is None + + +@pytest.mark.unit +class TestGraphConfig: + def test_defaults(self): + g = GraphConfig() + assert g.extraction_model is None + assert g.max_chunks is None + assert g.gleanings == 0 + + def test_parse_graph_block_round_trips(self): + raw = { + "kind": "graphrag", + "graph": { + "extraction_model": "gpt-4o-mini", + "max_chunks": 500, + "gleanings": 1, + }, + } + cfg = SourceConfig.parse(raw) + assert cfg.kind == "graphrag" + assert cfg.graph.extraction_model == "gpt-4o-mini" + assert cfg.graph.max_chunks == 500 + assert cfg.graph.gleanings == 1 + dumped = cfg.model_dump() + assert SourceConfig.model_validate(dumped) == cfg + assert dumped["graph"] == { + "extraction_model": "gpt-4o-mini", + "max_chunks": 500, + "gleanings": 1, + } + + def test_forbid_unknown_graph_key(self): + with pytest.raises(ValidationError): + SourceConfig.model_validate({"graph": {"unknown": 1}}) + with pytest.raises(ValidationError): + GraphConfig(unknown=1) + + def test_negative_values_rejected(self): + with pytest.raises(ValidationError): + GraphConfig(gleanings=-1) + with pytest.raises(ValidationError): + GraphConfig(max_chunks=0) + + def test_parse_lenient_on_garbage_graph(self): + # A non-dict / invalid graph blob falls back to all-defaults on read. + assert SourceConfig.parse({"graph": "nope"}) == SourceConfig() + assert SourceConfig.parse({"graph": {"gleanings": "lots"}}) == SourceConfig() + # Missing graph block still yields the default GraphConfig. + cfg = SourceConfig.parse({"kind": "graphrag"}) + assert cfg.graph == GraphConfig()