diff --git a/docsgpt/agents/tools/graph_search.py b/docsgpt/agents/tools/graph_search.py index c2c15a7b..b370949e 100644 --- a/docsgpt/agents/tools/graph_search.py +++ b/docsgpt/agents/tools/graph_search.py @@ -28,7 +28,8 @@ import logging from typing import Any, Dict, List, Optional from docsgpt.agents.tools.base import Tool -from docsgpt.core.settings import settings +from docsgpt.graphrag import graphrag_available +from docsgpt.retriever.labels import labels_from_metadata logger = logging.getLogger(__name__) @@ -61,6 +62,24 @@ class GraphSearchTool(Tool): self._store = GraphStore() return self._store + def _release_store(self) -> None: + """Hand the pooled connection back at the end of an action. + + The executor caches this tool for the whole agent run, so a store kept + between actions pins one connection of the shared pgvector pool across + every LLM round trip of that run -- minutes at a time, and enough + concurrent runs exhaust the pool. ``GraphRAGRetriever`` releases its + store before falling back for the same reason. Checking one back out + costs a pool acquire. + """ + store, self._store = self._store, None + if store is None: + return + try: + store.close() + except Exception as exc: # noqa: BLE001 -- releasing must not fail an action + logger.debug(f"Graph tool could not release its store: {exc}") + def _embed(self, text: str) -> Optional[List[float]]: try: from docsgpt.vectorstore.base import get_embeddings @@ -72,7 +91,9 @@ class GraphSearchTool(Tool): # -- actions ------------------------------------------------------------- def execute_action(self, action_name: str, **kwargs): - if not settings.GRAPHRAG_ENABLED: + # The graph lives in the pgvector store, so the flag alone is not + # enough: under another vector store there is no graph to read. + if not graphrag_available(): return "The knowledge graph is not enabled for this deployment." if not self._sources(): return "No graph-backed sources are configured." @@ -86,6 +107,8 @@ class GraphSearchTool(Tool): except Exception as e: # noqa: BLE001 logger.error(f"Graph tool action {action_name} failed: {e}", exc_info=True) return "The graph lookup failed." + finally: + self._release_store() return f"Unknown action: {action_name}" def _search_entities(self, **kwargs) -> str: @@ -137,18 +160,17 @@ class GraphSearchTool(Tool): 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", "")} + text = page.get("text") or "" + # The retrievers' own labelling: a page read here and the same + # chunk retrieved by internal_search are one document, and + # citations key on (source, title). Labelling it differently + # gives that document two citation numbers. + labels = labels_from_metadata(page.get("metadata"), text, source_id) + doc = {**labels, "text": text} if doc not in self.retrieved_docs: self.retrieved_docs.append(doc) - parts.append(f"--- {title} ---\n{text}") + header = labels["filename"] or labels["title"] + parts.append(f"--- {header} ---\n{text[:MAX_PAGE_CHARS]}") if not parts: return f"No documents mention {entity!r}." return "\n\n".join(parts) @@ -259,7 +281,7 @@ def add_graph_search_tool(tools_dict: Dict, retriever_config: Dict) -> None: 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: + if not graphrag_available(): return source = retriever_config.get("source") or {} if not source.get("active_docs") or not sources_have_graph(source): diff --git a/tests/graphrag/test_graph_search_tool.py b/tests/graphrag/test_graph_search_tool.py index 559c911f..465a1c25 100644 --- a/tests/graphrag/test_graph_search_tool.py +++ b/tests/graphrag/test_graph_search_tool.py @@ -9,6 +9,8 @@ must refuse clearly rather than silently when it has nothing to offer. from __future__ import annotations +import pytest + from docsgpt.agents.tools.graph_search import ( GRAPH_TOOL_ID, GraphSearchTool, @@ -38,6 +40,7 @@ class _StubStore: def _tool(monkeypatch, store, enabled=True): monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", enabled) + monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector") tool = GraphSearchTool({"source": SOURCE}) tool._store = store monkeypatch.setattr(tool, "_embed", lambda text: [0.0, 0.1]) @@ -52,6 +55,7 @@ class TestGating: def test_reports_when_no_sources_are_configured(self, monkeypatch): monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True) + monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector") tool = GraphSearchTool({"source": {"active_docs": []}}) assert "No graph-backed sources" in tool.execute_action( @@ -109,16 +113,49 @@ class TestActions: tool = _tool( monkeypatch, _StubStore( - pages=[{"metadata": {"file_path": "quill-store.md"}, "text": "x" * 5000}] + pages=[ + { + "metadata": {"title": "quill-store.md", "source": "quill-store.md"}, + "text": "x" * 5000, + } + ] ), ) result = tool.execute_action("read_entity_pages", entity="Quill") assert "--- quill-store.md ---" in result + # Only what the model reads is truncated. assert len(result) < 3000 # Accumulated so the answer can cite what the walk actually read. assert tool.retrieved_docs[0]["title"] == "quill-store.md" + assert len(tool.retrieved_docs[0]["text"]) == 5000 + + def test_page_labels_match_what_the_retrievers_record(self, monkeypatch): + """A page read here and the same chunk retrieved by internal_search are + one document. The citation manager keys on (source, title), so labels + derived differently give the same document two citation numbers.""" + from docsgpt.retriever.labels import labels_from_metadata + + metadata = {"title": "Quill Store", "source": "quill-store.md"} + text = "Quill is a write-ahead store." + tool = _tool(monkeypatch, _StubStore(pages=[{"metadata": metadata, "text": text}])) + + tool.execute_action("read_entity_pages", entity="Quill") + + expected = labels_from_metadata(metadata, text, "src-1") + doc = tool.retrieved_docs[0] + assert {k: doc[k] for k in ("title", "source", "filename")} == expected + # The full chunk text, so the doc dedupes against the retriever's copy; + # only what the model reads is truncated. + assert doc["text"] == text + + def test_a_page_with_no_metadata_falls_back_to_its_source_id(self, monkeypatch): + tool = _tool(monkeypatch, _StubStore(pages=[{"metadata": {}, "text": "body"}])) + + tool.execute_action("read_entity_pages", entity="Quill") + + assert tool.retrieved_docs[0]["source"] == "src-1" def test_pages_absent_is_stated_plainly(self, monkeypatch): tool = _tool(monkeypatch, _StubStore(pages=[])) @@ -150,6 +187,7 @@ class TestWiring: def test_not_added_when_the_sources_have_no_graph(self, monkeypatch): monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True) + monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector") monkeypatch.setattr( "docsgpt.agents.tools.graph_search.sources_have_graph", lambda source: False ) @@ -161,6 +199,7 @@ class TestWiring: def test_added_with_its_sentinel_id_and_source_config(self, monkeypatch): monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True) + monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector") monkeypatch.setattr( "docsgpt.agents.tools.graph_search.sources_have_graph", lambda source: True ) @@ -246,3 +285,66 @@ class TestSourcesHaveGraph: assert sources_have_graph({"active_docs": []}) is False self._patch_counts(monkeypatch, error=RuntimeError("no pgvector")) assert sources_have_graph({"active_docs": ["a"]}) is False + + +class TestGraphsMustBeAvailable: + """The graph lives in the pgvector store, so the flag alone is not enough. + + With another vector store configured the graph tables are not the ones the + sources were ingested into; everything else in the app asks + ``graphrag_available()``, which requires both. + """ + + def test_the_tool_is_not_offered_without_pgvector(self, monkeypatch): + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True) + monkeypatch.setattr(settings, "VECTOR_STORE", "faiss") + monkeypatch.setattr( + "docsgpt.agents.tools.graph_search.sources_have_graph", + lambda source: pytest.fail("must not reach the database"), + ) + tools = {} + add_graph_search_tool(tools, {"source": SOURCE}) + assert tools == {} + + def test_actions_report_it_rather_than_querying(self, monkeypatch): + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True) + monkeypatch.setattr(settings, "VECTOR_STORE", "faiss") + tool = GraphSearchTool({"source": SOURCE}) + tool._store = _StubStore(nodes=[{"name": "Quill", "distance": 0.1}]) + + assert "not enabled" in tool.execute_action("search_entities", query="quill") + + +class TestPooledConnection: + """The tool is cached for the whole agent run; its connection must not be.""" + + class _ClosingStore(_StubStore): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.closed = 0 + + def close(self): + self.closed += 1 + + def test_the_connection_goes_back_after_each_action(self, monkeypatch): + store = self._ClosingStore(relationships=[{"source": "A", "target": "B", "type": "r"}]) + tool = _tool(monkeypatch, store) + + tool.execute_action("get_relationships", entity="A") + + # Held open, one pooled connection would be pinned across every LLM + # round trip of the run. + assert store.closed == 1 + assert tool._store is None + + def test_a_failing_action_still_releases_it(self, monkeypatch): + class _Broken(self._ClosingStore): + def entity_relationships(self, source_id, name, limit=25): + raise RuntimeError("connection lost") + + store = _Broken() + tool = _tool(monkeypatch, store) + + assert tool.execute_action("get_relationships", entity="A") == "The graph lookup failed." + assert store.closed == 1 + assert tool._store is None