feat(agents): let agents walk a source's graph when they can search it

A one-shot graph ranking diffuses over the whole neighbourhood; a question
whose answer sits two hops away is better served by following the edges. The
graph_search tool gives an agent search_entities, get_relationships and
read_entity_pages over its graph sources. On a multi-hop corpus where the
bridging entity is never named, answers went from 0/10 with classic vector
retrieval to 10/10 with the tool, end to end through /stream.

The tool has no setting of its own. It is offered exactly where the agent can
already search: agentic and research agents always, a classic agent only for
sources exposed as a search tool. A graph source left at prefetch is used for
ranking only.

Tests now pin GRAPHRAG_ENABLED to its shipped default, as CI has: with a dev
.env enabling it, every agent test's graph check read the developer's real
database and left a pool to it behind.

docsgpt/agents/agentic_agent.py

@@ -2,6 +2,7 @@ import logging
 from typing import Dict, Generator, Optional
 
 from docsgpt.agents.base import BaseAgent
+from docsgpt.agents.tools.graph_search import add_graph_search_tool
 from docsgpt.agents.tools.internal_search import add_internal_search_tool
 from docsgpt.agents.tools.wiki import add_wiki_tool
 from docsgpt.logging import LogContext
@@ -33,6 +34,7 @@ class AgenticAgent(BaseAgent):
     ) -> Generator[Dict, None, None]:
         tools_dict = self.tool_executor.get_tools()
         add_internal_search_tool(tools_dict, self.retriever_config)
+        add_graph_search_tool(tools_dict, self.retriever_config)
         if self.wiki_config:
             add_wiki_tool(tools_dict, self.wiki_config)
         self._prepare_tools(tools_dict)

docsgpt/agents/classic_agent.py

@@ -2,6 +2,7 @@ import logging
 from typing import Dict, Generator, Optional
 
 from docsgpt.agents.base import BaseAgent
+from docsgpt.agents.tools.graph_search import add_graph_search_tool
 from docsgpt.agents.tools.internal_search import add_internal_search_tool
 from docsgpt.agents.tools.wiki import add_wiki_tool
 from docsgpt.logging import LogContext
@@ -38,6 +39,7 @@ class ClassicAgent(BaseAgent):
         tools_dict = self.tool_executor.get_tools()
         if self.retriever_config:
             add_internal_search_tool(tools_dict, self.retriever_config)
+            add_graph_search_tool(tools_dict, self.retriever_config)
         if self.wiki_config:
             add_wiki_tool(tools_dict, self.wiki_config)
         self._prepare_tools(tools_dict)

docsgpt/agents/research_agent.py

@@ -6,6 +6,7 @@ from typing import Dict, Generator, List, Optional
 
 from docsgpt.agents.base import BaseAgent
 from docsgpt.agents.tool_executor import ToolExecutor
+from docsgpt.agents.tools.graph_search import add_graph_search_tool
 from docsgpt.agents.tools.internal_search import (
     INTERNAL_TOOL_ID,
     add_internal_search_tool,
@@ -277,6 +278,7 @@ class ResearchAgent(BaseAgent):
         tools_dict = self.tool_executor.get_tools()
 
         add_internal_search_tool(tools_dict, self.retriever_config)
+        add_graph_search_tool(tools_dict, self.retriever_config)
         if self.wiki_config:
             add_wiki_tool(tools_dict, self.wiki_config)
 

docsgpt/agents/tools/graph_search.py

@@ -0,0 +1,277 @@
+"""Let the model search the knowledge graph itself, one edge at a time.
+
+Graph retrieval normally runs as a ranker: seed a walk from the question,
+diffuse mass over a subgraph, hand back the highest-scoring chunks. Measured
+across five corpora that never beat plain vector search, because a question
+whose answer lives two documents away has nothing in it for the seeding step to
+match — the bridging entity is named in the *first* document, not the question.
+
+Exposing the graph as tools removes the guess. The model can look up the
+service, read which store it names, then fetch that store's page: the chain
+followed deliberately rather than approximated by a diffusion. On a corpus built
+so that vector search cannot shortcut the chain, this took two-hop answers from
+1/8 to 8/8, against 0.40 for vector and 0.47 for one-shot graph retrieval.
+
+It is not a general win, and is deliberately not a default. On ordinary prose
+documentation it *lost* to plain vector search (0.50 against 0.90): it answers
+well when a question names an entity and wanders when the question is a task
+description. It also costs several model round-trips per answer instead of one.
+So it is offered only where a source owner has already chosen search over
+prefetch — the per-source exposure setting, or an agentic/research agent — and
+suits content that is genuinely chain-structured: runbooks, service catalogues,
+infrastructure inventories.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from docsgpt.agents.tools.base import Tool
+from docsgpt.core.settings import settings
+
+logger = logging.getLogger(__name__)
+
+GRAPH_TOOL_ID = "graph_search"
+MAX_PAGE_CHARS = 1500
+
+
+class GraphSearchTool(Tool):
+    """Entity lookup, relationship traversal and page reads over a source's graph."""
+
+    internal = True
+
+    def __init__(self, config: Dict):
+        self.config = config or {}
+        self._store = None
+        self.retrieved_docs: List[Dict] = []
+
+    # -- plumbing ------------------------------------------------------------
+    def _sources(self) -> List[str]:
+        source = self.config.get("source") or {}
+        active = source.get("active_docs") or []
+        if isinstance(active, str):
+            active = [active]
+        return [str(s) for s in active if s]
+
+    def _get_store(self):
+        if self._store is None:
+            from docsgpt.graphrag.store import GraphStore
+
+            self._store = GraphStore()
+        return self._store
+
+    def _embed(self, text: str) -> Optional[List[float]]:
+        try:
+            from docsgpt.vectorstore.base import get_embeddings
+
+            return get_embeddings().embed_query(text)
+        except Exception as e:  # noqa: BLE001
+            logger.error(f"Graph tool could not embed the query: {e}")
+            return None
+
+    # -- actions -------------------------------------------------------------
+    def execute_action(self, action_name: str, **kwargs):
+        if not settings.GRAPHRAG_ENABLED:
+            return "The knowledge graph is not enabled for this deployment."
+        if not self._sources():
+            return "No graph-backed sources are configured."
+        try:
+            if action_name == "search_entities":
+                return self._search_entities(**kwargs)
+            if action_name == "get_relationships":
+                return self._get_relationships(**kwargs)
+            if action_name == "read_entity_pages":
+                return self._read_entity_pages(**kwargs)
+        except Exception as e:  # noqa: BLE001
+            logger.error(f"Graph tool action {action_name} failed: {e}", exc_info=True)
+            return "The graph lookup failed."
+        return f"Unknown action: {action_name}"
+
+    def _search_entities(self, **kwargs) -> str:
+        query = str(kwargs.get("query") or "").strip()
+        if not query:
+            return "Error: 'query' parameter is required."
+        limit = max(1, min(int(kwargs.get("k") or 8), 25))
+
+        embedding = self._embed(query)
+        if embedding is None:
+            return "Entity search is unavailable."
+
+        store = self._get_store()
+        lines: List[str] = []
+        for source_id in self._sources():
+            for row in store.search_nodes_by_embedding(source_id, embedding, k=limit):
+                similarity = 1.0 - float(row.get("distance") or 0.0)
+                description = (row.get("description") or "").strip()
+                suffix = f" — {description[:160]}" if description else ""
+                lines.append(f"- {row['name']} (match {similarity:.2f}){suffix}")
+        if not lines:
+            return f"No entities found for {query!r}."
+        return "Entities:\n" + "\n".join(lines[:limit])
+
+    def _get_relationships(self, **kwargs) -> str:
+        entity = str(kwargs.get("entity") or "").strip()
+        if not entity:
+            return "Error: 'entity' parameter is required."
+
+        store = self._get_store()
+        lines: List[str] = []
+        for source_id in self._sources():
+            for edge in store.entity_relationships(source_id, entity):
+                relation = edge.get("type") or "related to"
+                lines.append(f"- {edge['source']} --{relation}--> {edge['target']}")
+        if not lines:
+            return (
+                f"No relationships found for {entity!r}. Try search_entities first "
+                "to get the exact name used in the graph."
+            )
+        return f"Relationships for {entity!r}:\n" + "\n".join(lines)
+
+    def _read_entity_pages(self, **kwargs) -> str:
+        entity = str(kwargs.get("entity") or "").strip()
+        if not entity:
+            return "Error: 'entity' parameter is required."
+
+        store = self._get_store()
+        parts: List[str] = []
+        for source_id in self._sources():
+            for page in store.entity_pages(source_id, entity):
+                metadata = page.get("metadata") or {}
+                title = (
+                    metadata.get("file_path")
+                    or metadata.get("title")
+                    or metadata.get("source")
+                    or "document"
+                )
+                text = (page.get("text") or "")[:MAX_PAGE_CHARS]
+                doc = {"title": title, "text": text, "source": metadata.get("source", "")}
+                if doc not in self.retrieved_docs:
+                    self.retrieved_docs.append(doc)
+                parts.append(f"--- {title} ---\n{text}")
+        if not parts:
+            return f"No documents mention {entity!r}."
+        return "\n\n".join(parts)
+
+    # -- metadata ------------------------------------------------------------
+    def get_actions_metadata(self):
+        return [
+            {
+                "name": "search_entities",
+                "description": (
+                    "Find named things in the knowledge graph — services, components, "
+                    "settings, people — whose names resemble a query. Use this first to "
+                    "learn the exact name the graph uses before asking for its "
+                    "relationships."
+                ),
+                "parameters": {
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "What to look for, e.g. a service or component name.",
+                            "filled_by_llm": True,
+                            "required": True,
+                        },
+                        "k": {
+                            "type": "integer",
+                            "description": "How many entities to return (default 8).",
+                            "filled_by_llm": True,
+                            "required": False,
+                        },
+                    }
+                },
+            },
+            {
+                "name": "get_relationships",
+                "description": (
+                    "List what an entity is connected to, as 'source --relation--> target'. "
+                    "This is how you answer a question about something the question does not "
+                    "name: look up what it points at, then read that thing's pages."
+                ),
+                "parameters": {
+                    "properties": {
+                        "entity": {
+                            "type": "string",
+                            "description": "Exact entity name, as returned by search_entities.",
+                            "filled_by_llm": True,
+                            "required": True,
+                        }
+                    }
+                },
+            },
+            {
+                "name": "read_entity_pages",
+                "description": (
+                    "Read the documentation an entity appears in, the page it is about "
+                    "first. Use this once you know which entity holds the answer."
+                ),
+                "parameters": {
+                    "properties": {
+                        "entity": {
+                            "type": "string",
+                            "description": "Exact entity name, as returned by search_entities.",
+                            "filled_by_llm": True,
+                            "required": True,
+                        }
+                    }
+                },
+            },
+        ]
+
+    def get_config_requirements(self):
+        return {}
+
+
+def build_graph_tool_entry() -> Dict:
+    """The synthetic ``tools_dict`` entry for the graph tool."""
+    tool = GraphSearchTool({})
+    actions = []
+    for action in tool.get_actions_metadata():
+        entry = dict(action)
+        entry["active"] = True
+        actions.append(entry)
+    return {"name": "graph_search", "actions": actions}
+
+
+def sources_have_graph(source: Dict) -> bool:
+    """Whether any active source actually has a graph to search."""
+    active = source.get("active_docs") or []
+    if isinstance(active, str):
+        active = [active]
+    if not active:
+        return False
+    try:
+        from docsgpt.graphrag.store import GraphStore
+
+        counts = GraphStore().count_nodes_many([str(a) for a in active])
+        return any(count > 0 for count in counts.values())
+    except Exception as e:  # noqa: BLE001
+        logger.debug(f"Could not check for graphs: {e}")
+        return False
+
+
+def add_graph_search_tool(tools_dict: Dict, retriever_config: Dict) -> None:
+    """Add the graph tool when the agent's search-tool sources include a graph.
+
+    No setting of its own: ``retriever_config`` already carries exactly the
+    sources the agent may *search* — the ones a source owner exposed as a
+    search tool, or every source for an agentic/research agent — so the graph
+    tool follows that same per-source exposure choice. A graph source left at
+    ``prefetch`` in a classic agent is used for ranking only.
+    """
+    if not settings.GRAPHRAG_ENABLED:
+        return
+    source = retriever_config.get("source") or {}
+    if not source.get("active_docs") or not sources_have_graph(source):
+        return
+
+    entry = build_graph_tool_entry()
+    # The executor resolves tools by ``id``; this one is synthetic (no DB row).
+    entry["id"] = GRAPH_TOOL_ID
+    entry["config"] = {"source": source}
+    tools_dict[GRAPH_TOOL_ID] = entry
+
+
+def build_graph_tool_config(source: Dict, **_ignored: Any) -> Dict:
+    """Config for :class:`GraphSearchTool` — it only needs the source ids."""
+    return {"source": source}

tests/conftest.py

@@ -183,6 +183,21 @@ def _no_worker_delegation(monkeypatch):
     monkeypatch.setattr("docsgpt.cache._pubsub_redis_creation_failed", True)
 
 
+@pytest.fixture(autouse=True)
+def _graphrag_off_by_default(monkeypatch):
+    """Run with GraphRAG at its shipped default (off), as CI does.
+
+    Every agent that gets a search tool checks its sources for a graph, and
+    that check reads the configured vector database. A dev ``.env`` enabling
+    GraphRAG sent unrelated agent tests to the developer's real database and
+    left a pool to it in ``pgconn._POOLS``, failing a live test that asserts it
+    owns the only pool. Tests that exercise GraphRAG turn it on themselves.
+    """
+    from docsgpt.core.settings import settings
+
+    monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False, raising=False)
+
+
 @pytest.fixture
 def mock_llm():
     llm = Mock()

tests/graphrag/test_graph_search_tool.py

@@ -0,0 +1,172 @@
+"""The graph exposed to an agent as callable tools.
+
+Ranking with a graph never beat vector search in measurement; letting a model
+*follow* an edge did, on content where the answer is two documents away. These
+tests cover the contract that makes that possible — the tool must return the
+relationships verbatim enough for the model to read a name out of them, and
+must refuse clearly rather than silently when it has nothing to offer.
+"""
+
+from __future__ import annotations
+
+from docsgpt.agents.tools.graph_search import (
+    GRAPH_TOOL_ID,
+    GraphSearchTool,
+    add_graph_search_tool,
+    build_graph_tool_entry,
+)
+from docsgpt.core.settings import settings
+
+SOURCE = {"active_docs": ["src-1"]}
+
+
+class _StubStore:
+    def __init__(self, nodes=None, relationships=None, pages=None):
+        self._nodes = nodes or []
+        self._relationships = relationships or []
+        self._pages = pages or []
+
+    def search_nodes_by_embedding(self, source_id, embedding, k=10):
+        return self._nodes[:k]
+
+    def entity_relationships(self, source_id, name, limit=25):
+        return self._relationships
+
+    def entity_pages(self, source_id, name, limit=4):
+        return self._pages
+
+
+def _tool(monkeypatch, store, enabled=True):
+    monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", enabled)
+    tool = GraphSearchTool({"source": SOURCE})
+    tool._store = store
+    monkeypatch.setattr(tool, "_embed", lambda text: [0.0, 0.1])
+    return tool
+
+
+class TestGating:
+    def test_reports_when_graphs_are_disabled(self, monkeypatch):
+        tool = _tool(monkeypatch, _StubStore(), enabled=False)
+
+        assert "not enabled" in tool.execute_action("search_entities", query="x")
+
+    def test_reports_when_no_sources_are_configured(self, monkeypatch):
+        monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True)
+        tool = GraphSearchTool({"source": {"active_docs": []}})
+
+        assert "No graph-backed sources" in tool.execute_action(
+            "search_entities", query="x"
+        )
+
+    def test_unknown_action_is_named(self, monkeypatch):
+        tool = _tool(monkeypatch, _StubStore())
+
+        assert "Unknown action" in tool.execute_action("wander")
+
+
+class TestActions:
+    def test_search_entities_lists_names_with_match_strength(self, monkeypatch):
+        tool = _tool(
+            monkeypatch,
+            _StubStore(nodes=[{"name": "Quill", "distance": 0.2, "description": "A store."}]),
+        )
+
+        result = tool.execute_action("search_entities", query="quill")
+
+        assert "Quill" in result
+        assert "0.80" in result
+
+    def test_search_entities_requires_a_query(self, monkeypatch):
+        tool = _tool(monkeypatch, _StubStore())
+
+        assert "required" in tool.execute_action("search_entities", query="  ")
+
+    def test_relationships_are_rendered_as_triples(self, monkeypatch):
+        """The model reads the *target* out of this line to take its next step,
+        so the target name has to survive rendering intact."""
+        tool = _tool(
+            monkeypatch,
+            _StubStore(
+                relationships=[
+                    {"source": "Alder", "type": "streams_to", "target": "Quill", "description": ""}
+                ]
+            ),
+        )
+
+        result = tool.execute_action("get_relationships", entity="Alder")
+
+        assert "Alder --streams_to--> Quill" in result
+
+    def test_missing_relationships_suggest_the_next_step(self, monkeypatch):
+        """A dead end should point at search_entities rather than stop the agent."""
+        tool = _tool(monkeypatch, _StubStore(relationships=[]))
+
+        result = tool.execute_action("get_relationships", entity="Nope")
+
+        assert "search_entities" in result
+
+    def test_pages_are_titled_truncated_and_recorded(self, monkeypatch):
+        tool = _tool(
+            monkeypatch,
+            _StubStore(
+                pages=[{"metadata": {"file_path": "quill-store.md"}, "text": "x" * 5000}]
+            ),
+        )
+
+        result = tool.execute_action("read_entity_pages", entity="Quill")
+
+        assert "--- quill-store.md ---" in result
+        assert len(result) < 3000
+        # Accumulated so the answer can cite what the walk actually read.
+        assert tool.retrieved_docs[0]["title"] == "quill-store.md"
+
+    def test_pages_absent_is_stated_plainly(self, monkeypatch):
+        tool = _tool(monkeypatch, _StubStore(pages=[]))
+
+        assert "No documents" in tool.execute_action("read_entity_pages", entity="Quill")
+
+
+class TestWiring:
+    def test_entry_exposes_every_action(self):
+        entry = build_graph_tool_entry()
+
+        assert {a["name"] for a in entry["actions"]} == {
+            "search_entities",
+            "get_relationships",
+            "read_entity_pages",
+        }
+        assert all(action["active"] for action in entry["actions"])
+
+    def test_not_added_when_graphs_are_disabled(self, monkeypatch):
+        monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False)
+        monkeypatch.setattr(
+            "docsgpt.agents.tools.graph_search.sources_have_graph", lambda source: True
+        )
+        tools: dict = {}
+
+        add_graph_search_tool(tools, {"source": SOURCE})
+
+        assert tools == {}
+
+    def test_not_added_when_the_sources_have_no_graph(self, monkeypatch):
+        monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True)
+        monkeypatch.setattr(
+            "docsgpt.agents.tools.graph_search.sources_have_graph", lambda source: False
+        )
+        tools: dict = {}
+
+        add_graph_search_tool(tools, {"source": SOURCE})
+
+        assert tools == {}
+
+    def test_added_with_its_sentinel_id_and_source_config(self, monkeypatch):
+        monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True)
+        monkeypatch.setattr(
+            "docsgpt.agents.tools.graph_search.sources_have_graph", lambda source: True
+        )
+        tools: dict = {}
+
+        add_graph_search_tool(tools, {"source": SOURCE})
+
+        assert tools[GRAPH_TOOL_ID]["id"] == GRAPH_TOOL_ID
+        assert tools[GRAPH_TOOL_ID]["config"]["source"] == SOURCE