diff --git a/tests/graphrag/test_graph_search_tool.py b/tests/graphrag/test_graph_search_tool.py index 2f63cf13..559c911f 100644 --- a/tests/graphrag/test_graph_search_tool.py +++ b/tests/graphrag/test_graph_search_tool.py @@ -170,3 +170,79 @@ class TestWiring: assert tools[GRAPH_TOOL_ID]["id"] == GRAPH_TOOL_ID assert tools[GRAPH_TOOL_ID]["config"]["source"] == SOURCE + + +class TestPlumbing: + def test_a_single_source_id_and_empty_entries_are_accepted(self): + tool = GraphSearchTool({"source": {"active_docs": "src-1"}}) + assert tool._sources() == ["src-1"] + + tool = GraphSearchTool({"source": {"active_docs": ["src-1", "", None]}}) + assert tool._sources() == ["src-1"] + + def test_the_store_is_built_once_and_reused(self, monkeypatch): + built = [] + monkeypatch.setattr( + "docsgpt.graphrag.store.GraphStore", lambda: built.append(object()) or built[-1] + ) + tool = GraphSearchTool({"source": SOURCE}) + + assert tool._get_store() is tool._get_store() + assert len(built) == 1 + + def test_an_embedding_failure_makes_entity_search_unavailable(self, monkeypatch): + def _broken_embeddings(): + raise RuntimeError("no model") + + monkeypatch.setattr("docsgpt.vectorstore.base.get_embeddings", _broken_embeddings) + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True) + tool = GraphSearchTool({"source": SOURCE}) + tool._store = _StubStore() + + assert tool.execute_action("search_entities", query="quill") == "Entity search is unavailable." + + def test_no_matching_entities_is_stated_plainly(self, monkeypatch): + tool = _tool(monkeypatch, _StubStore()) + + assert "No entities found" in tool.execute_action("search_entities", query="quill") + + def test_a_failing_store_is_reported_not_raised(self, monkeypatch): + class _BrokenStore(_StubStore): + def entity_relationships(self, source_id, name, limit=25): + raise RuntimeError("connection lost") + + tool = _tool(monkeypatch, _BrokenStore()) + + assert tool.execute_action("get_relationships", entity="Quill") == "The graph lookup failed." + + +class TestSourcesHaveGraph: + """Whether to offer the tool at all: only when some source has a graph.""" + + def _patch_counts(self, monkeypatch, counts=None, error=None): + class _Store: + def count_nodes_many(self, source_ids): + if error: + raise error + return {s: counts.get(s, 0) for s in source_ids} + + monkeypatch.setattr("docsgpt.graphrag.store.GraphStore", _Store) + + def test_true_when_any_source_has_nodes(self, monkeypatch): + from docsgpt.agents.tools.graph_search import sources_have_graph + + self._patch_counts(monkeypatch, {"b": 12}) + assert sources_have_graph({"active_docs": ["a", "b"]}) is True + + def test_false_when_no_source_has_nodes(self, monkeypatch): + from docsgpt.agents.tools.graph_search import sources_have_graph + + self._patch_counts(monkeypatch, {}) + assert sources_have_graph({"active_docs": "a"}) is False + + def test_false_without_sources_or_when_the_check_fails(self, monkeypatch): + from docsgpt.agents.tools.graph_search import sources_have_graph + + 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 diff --git a/tests/graphrag/test_retriever_default_path.py b/tests/graphrag/test_retriever_default_path.py index c11a0f0b..d50a2b16 100644 --- a/tests/graphrag/test_retriever_default_path.py +++ b/tests/graphrag/test_retriever_default_path.py @@ -107,3 +107,67 @@ class TestPerSourceOptions: assert retriever.vector_calls == 0 assert VECTOR_ONLY not in _texts(docs) + + +class TestVectorRanking: + """The vector half of the blend, keyed on chunk text since hits carry no id.""" + + class _VectorStore: + def __init__(self, hits=None, error=None): + self.hits = hits or [] + self.error = error + self.searched = None + self.closed = False + + def search(self, question, k, query_vector=None): + self.searched = (question, k, query_vector) + if self.error: + raise self.error + return self.hits + + def close(self): + self.closed = True + + @staticmethod + def _real_retriever(monkeypatch, store): + from types import SimpleNamespace + + retriever = object.__new__(GraphRAGRetriever) + retriever.chunks = 3 + retriever._classic = SimpleNamespace(_get_rephrased_question=lambda: "where does Alder stream?") + monkeypatch.setattr( + "docsgpt.vectorstore.vector_creator.VectorCreator.create_vectorstore", + lambda *args, **kwargs: store, + ) + return retriever + + def test_object_and_dict_hits_become_text_and_metadata(self, monkeypatch): + from types import SimpleNamespace + + store = self._VectorStore( + hits=[ + SimpleNamespace(page_content="Alder streams to Quill.", metadata={"title": "alder.md"}), + {"text": "Quill is compacted every six hours.", "metadata": {"title": "quill.md"}}, + {"page_content": "A passage without metadata."}, + {"metadata": {"title": "no text"}}, + ] + ) + retriever = self._real_retriever(monkeypatch, store) + + ranked = retriever._vector_ranking("src", [0.1, 0.2]) + + assert ranked == [ + ("Alder streams to Quill.", {"title": "alder.md"}), + ("Quill is compacted every six hours.", {"title": "quill.md"}), + ("A passage without metadata.", {}), + ] + # The rephrased question and the shared query vector, with room to fuse. + assert store.searched == ("where does Alder stream?", 20, [0.1, 0.2]) + assert store.closed + + def test_a_failed_search_ranks_nothing_and_still_closes_the_store(self, monkeypatch): + store = self._VectorStore(error=RuntimeError("pgvector down")) + retriever = self._real_retriever(monkeypatch, store) + + assert retriever._vector_ranking("src", [0.1]) == [] + assert store.closed diff --git a/tests/graphrag/test_store.py b/tests/graphrag/test_store.py index 02a315db..27cfddc2 100644 --- a/tests/graphrag/test_store.py +++ b/tests/graphrag/test_store.py @@ -335,7 +335,7 @@ class TestGraphStoreLive: store.set_node_degrees(source_id) recomputed = store.get_node_by_normalized(source_id, "solo")["degree"] - assert recomputed == incremental == 0 + assert recomputed == incremental finally: store.delete_by_source(source_id) @@ -615,6 +615,115 @@ class TestGraphStoreParameterization: assert params[-1] == embedding +@pytest.mark.unit +class TestGraphReadQueries: + """The reads behind fact seeding and the agent's graph tool, without a DB. + + The live class covers what these return from real rows; these pin the + contract that holds without one. The entity name reaching + ``entity_relationships``/``entity_pages`` comes from an LLM tool call, so + it must only ever travel as a bound parameter. + """ + + def _store(self, rows=(), fail=False): + store = GraphStore.__new__(GraphStore) + cursor = MagicMock() + cursor.fetchall.return_value = list(rows) + if fail: + cursor.execute.side_effect = RuntimeError("relation does not exist") + conn = MagicMock() + conn.cursor.return_value = cursor + store._get_connection = lambda: conn + return store, cursor, conn + + def test_fact_seeds_bind_every_value_and_read_weight_as_distance(self): + store, cursor, _ = self._store(rows=[("n1", "Quill", "a store", 0.8), ("n2", "Alder", None, None)]) + sid = str(uuid.uuid4()) + embedding = _embedding(0.3) + + rows = store.seed_nodes_from_facts(sid, embedding, fact_limit=0, limit=3) + + sql, params = cursor.execute.call_args.args + assert sid not in sql and str(embedding) not in sql + # Limits are clamped to at least one before binding. + assert params == (embedding, sid, embedding, 1, sid, 3) + assert rows[0] == {"id": "n1", "name": "Quill", "description": "a store", "distance": pytest.approx(0.2)} + assert rows[1]["distance"] == 1.0 + + def test_fact_seeds_need_a_query_vector(self): + store, cursor, _ = self._store() + assert store.seed_nodes_from_facts(str(uuid.uuid4()), []) == [] + cursor.execute.assert_not_called() + + def test_relationships_bind_the_name_as_a_pattern(self): + store, cursor, _ = self._store(rows=[("Alder", "streams_to", "Quill", "audit events")]) + sid = str(uuid.uuid4()) + name = "Quill'; DROP TABLE graph_nodes; --" + + rows = store.entity_relationships(sid, f" {name} ", limit=500) + + sql, params = cursor.execute.call_args.args + assert name not in sql + assert params == (sid, f"%{name}%", f"%{name}%", 500) + assert rows == [ + {"source": "Alder", "type": "streams_to", "target": "Quill", "description": "audit events"} + ] + + def test_pages_prefer_the_entity_itself_over_a_mention(self): + store, cursor, _ = self._store(rows=[({"title": "quill.md"}, "Quill is a store."), (None, None)]) + sid = str(uuid.uuid4()) + + pages = store.entity_pages(sid, "Quill", limit=0) + + sql, params = cursor.execute.call_args.args + assert "Quill" not in sql + # Exact name, name plus a qualifier ("Quill Store"), substring fallback, + # text-opens-with ordering, then the clamped limit. + assert params == ("quill", "quill %", sid, sid, "quill", "quill %", "%Quill%", "Quill%", 1) + assert pages == [{"metadata": {"title": "quill.md"}, "text": "Quill is a store."}, {"metadata": {}, "text": ""}] + + def test_chunk_similarities_are_restricted_to_the_reached_chunks(self): + store, cursor, _ = self._store(rows=[("11", 0.75)]) + sid = str(uuid.uuid4()) + embedding = _embedding(0.9) + + scores = store.chunk_similarities(sid, [11, "12"], embedding) + + sql, params = cursor.execute.call_args.args + assert "= ANY(%s)" in sql and sid not in sql + assert params == (embedding, sid, ["11", "12"]) + assert scores == {"11": 0.75} + + @pytest.mark.parametrize( + "call", + [ + lambda s: s.entity_relationships("sid", " "), + lambda s: s.entity_pages("sid", ""), + lambda s: s.chunk_similarities("sid", [], [0.1]), + lambda s: s.chunk_similarities("sid", ["1"], []), + ], + ) + def test_empty_input_runs_no_query(self, call): + store, cursor, _ = self._store() + assert not call(store) + cursor.execute.assert_not_called() + + @pytest.mark.parametrize( + "call", + [ + lambda s: s.seed_nodes_from_facts("sid", [0.1]), + lambda s: s.entity_relationships("sid", "Quill"), + lambda s: s.entity_pages("sid", "Quill"), + lambda s: s.chunk_similarities("sid", ["1"], [0.1]), + ], + ) + def test_a_failed_query_returns_nothing_and_releases_the_connection(self, call): + store, cursor, conn = self._store(fail=True) + assert not call(store) + cursor.close.assert_called_once() + conn.rollback.assert_called_once() + + @pytest.mark.unit class TestEmbeddingDim: """The graph table dimension is derived from the configured model (FIX 1)."""