mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 20:15:26 +00:00
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.
This commit is contained in:
1 parent
acb7956fe0
commit
b4c465e608
6 files changed
+130
No files matched your search
@@ -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
|
||||
|
||||
@@ -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"
|
||||
@@ -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.
|
||||
|
||||
Whitespace-only changes.
@@ -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
|
||||
@@ -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()
|
||||
Reference in new issue
Block a user