mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-11 03:12:55 +00:00
Alex
3f774d813c
Three findings from review, all on code this branch introduced.
apply_chunk was not replay-safe. commit() can report connection loss *after*
Postgres committed, and the reconnect retry then replays the write: _upsert_node
bumps doc_freq a second time and _add_edge inserts another row, since
graph_edges has no uniqueness constraint for a logical edge. The chunk's
graph_ingest_progress row is now written in the same transaction as the rows it
describes, and a replay that finds it already "done" returns (0, 0) without
touching the graph. Extraction drops its separate mark_chunk("done"): the
checkpoint and the graph can no longer disagree.
count_nodes swallows every query failure and answers 0, so extraction's
"fall back to the write count" handler could never run — a failed count after a
successful build reported an empty graph. count_nodes grows a strict mode that
re-raises; retrieval keeps the swallow, which is what routes a source to
ClassicRAG.
A zero edge weight was read as a full-strength link: `or 1.0` rewrote an
explicit 0 before the <= 0 filter. Only missing and null weights default now.
The same coercion sat in _ppr_scores, where it would have kept the ranker's
rule unreachable from the product path, so it is fixed there too.
719 lines
25 KiB
Python
719 lines
25 KiB
Python
"""Tests for the GraphRAG extraction pipeline (D28).
|
|
|
|
The LLM and the embeddings model are mocked in every test so the suite makes no
|
|
real model or network calls. A live ``GraphStore`` is exercised against the
|
|
ephemeral pytest-postgresql cluster (never the operator's dev DB) with a unique
|
|
temp ``source_id``; if pgvector is unavailable there the live tests skip.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import uuid
|
|
|
|
import pytest
|
|
|
|
import docsgpt.graphrag.extraction as extraction_module
|
|
from docsgpt.graphrag.store import GraphStore
|
|
from docsgpt.storage.db.source_config import SourceConfig
|
|
from docsgpt.vectorstore import pgconn
|
|
|
|
extract_graph_for_source = extraction_module.extract_graph_for_source
|
|
|
|
TEST_EMBEDDING_DIM = 8
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _close_pools():
|
|
"""Never leak a pool into another test; an ephemeral DSN dies with its DB."""
|
|
yield
|
|
for dsn, pool in list(pgconn._POOLS.items()):
|
|
try:
|
|
pool.close()
|
|
except Exception:
|
|
pass
|
|
pgconn._POOLS.pop(dsn, None)
|
|
|
|
|
|
def _ephemeral_dsn(info) -> str:
|
|
"""libpq DSN for the ephemeral pytest-postgresql database."""
|
|
password = f":{info.password}" if info.password else ""
|
|
return (
|
|
f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}"
|
|
)
|
|
|
|
|
|
def _live_store(monkeypatch, info):
|
|
"""Graph store on a fresh ephemeral database, schema created up front.
|
|
|
|
Construction runs no DDL any more (boot owns the schema), so the tables are
|
|
created explicitly here — what ``ensure_vector_schema`` does in production.
|
|
"""
|
|
monkeypatch.setattr(
|
|
GraphStore, "_embedding_dim", lambda self: TEST_EMBEDDING_DIM
|
|
)
|
|
dsn = _ephemeral_dsn(info)
|
|
# The pipeline builds its own GraphStore() from settings, so point those at
|
|
# the ephemeral cluster too — never at the operator's configured DB.
|
|
from docsgpt.core import settings as settings_module
|
|
|
|
monkeypatch.setattr(
|
|
settings_module.settings, "PGVECTOR_CONNECTION_STRING", dsn, raising=False
|
|
)
|
|
store = GraphStore(connection_string=dsn)
|
|
try:
|
|
store._ensure_tables()
|
|
except Exception as exc:
|
|
pytest.skip(f"pgvector extension unavailable: {exc}")
|
|
return store
|
|
|
|
|
|
class _StubLLM:
|
|
"""Stub LLM whose ``.gen`` returns crafted responses in order."""
|
|
|
|
def __init__(self, responses):
|
|
self._responses = list(responses)
|
|
self.model_id = "stub-model"
|
|
self.gen_calls = []
|
|
self._token_usage_source = None
|
|
self._request_id = None
|
|
|
|
def gen(self, model=None, messages=None, **kwargs):
|
|
self.gen_calls.append({"model": model, "messages": messages})
|
|
if not self._responses:
|
|
raise AssertionError("gen called more times than crafted responses")
|
|
response = self._responses.pop(0)
|
|
if isinstance(response, Exception):
|
|
raise response
|
|
return response
|
|
|
|
|
|
class _StubEmbedding:
|
|
"""Stub embeddings model producing deterministic fixed-dim vectors."""
|
|
|
|
def __init__(self):
|
|
self.dimension = TEST_EMBEDDING_DIM
|
|
|
|
def embed_documents(self, documents):
|
|
return [
|
|
[float(len(d) % 7)] + [0.0] * (TEST_EMBEDDING_DIM - 1)
|
|
for d in documents
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def stub_embedding(monkeypatch):
|
|
from docsgpt.core.settings import settings
|
|
|
|
# The resolver short-circuits to the remote API when this is configured,
|
|
# which would bypass the stub on a dev machine that sets it.
|
|
monkeypatch.setattr(settings, "EMBEDDINGS_BASE_URL", None)
|
|
embedding = _StubEmbedding()
|
|
monkeypatch.setattr(
|
|
extraction_module.EmbeddingsSingleton,
|
|
"get_instance",
|
|
staticmethod(lambda *a, **k: embedding),
|
|
)
|
|
return embedding
|
|
|
|
|
|
def _install_stub_llm(monkeypatch, llm):
|
|
captured = {}
|
|
|
|
def _create(*args, **kwargs):
|
|
captured["model_id"] = kwargs.get("model_id")
|
|
return llm
|
|
|
|
monkeypatch.setattr(
|
|
extraction_module.LLMCreator, "create_llm", staticmethod(_create)
|
|
)
|
|
return captured
|
|
|
|
|
|
def _chunk(doc_id, text):
|
|
return {"doc_id": doc_id, "text": text}
|
|
|
|
|
|
def _extraction_json(entities, relationships):
|
|
return json.dumps({"entities": entities, "relationships": relationships})
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestExtractionLive:
|
|
@pytest.fixture
|
|
def store(self, monkeypatch, postgresql):
|
|
store = _live_store(monkeypatch, postgresql.info)
|
|
yield store
|
|
store.close()
|
|
|
|
@pytest.fixture
|
|
def source_id(self):
|
|
return str(uuid.uuid4())
|
|
|
|
def test_entities_and_relationships_written(
|
|
self, store, source_id, monkeypatch, stub_embedding
|
|
):
|
|
try:
|
|
payload = _extraction_json(
|
|
entities=[
|
|
{"name": "Ada Lovelace", "type": "person", "description": "A mathematician."},
|
|
{"name": "Analytical Engine", "type": "machine", "description": "Early computer."},
|
|
],
|
|
relationships=[
|
|
{
|
|
"source": "Ada Lovelace",
|
|
"target": "Analytical Engine",
|
|
"type": "worked_on",
|
|
"description": "wrote algorithms for it",
|
|
"weight": 3.0,
|
|
}
|
|
],
|
|
)
|
|
llm = _StubLLM([payload])
|
|
_install_stub_llm(monkeypatch, llm)
|
|
|
|
summary = extract_graph_for_source(
|
|
source_id,
|
|
user="owner-1",
|
|
chunks=[_chunk("c1", "Ada Lovelace worked on the Analytical Engine.")],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
assert summary["nodes"] == 2
|
|
assert summary["edges"] == 1
|
|
assert summary["chunks_processed"] == 1
|
|
assert summary["failed_chunks"] == 0
|
|
assert store.count_nodes(source_id) == 2
|
|
|
|
node = store.get_node_by_normalized(source_id, "ada lovelace")
|
|
assert node is not None
|
|
mapping = store.get_chunk_ids_for_nodes(source_id, [node["id"]])
|
|
assert mapping[node["id"]] == ["c1"]
|
|
finally:
|
|
store.delete_by_source(source_id)
|
|
|
|
def test_same_entity_across_chunks_merges(
|
|
self, store, source_id, monkeypatch, stub_embedding
|
|
):
|
|
try:
|
|
payload_a = _extraction_json(
|
|
entities=[{"name": "Ada", "type": "person", "description": "first"}],
|
|
relationships=[],
|
|
)
|
|
payload_b = _extraction_json(
|
|
entities=[{"name": "Ada", "type": "person", "description": "second"}],
|
|
relationships=[],
|
|
)
|
|
llm = _StubLLM([payload_a, payload_b])
|
|
_install_stub_llm(monkeypatch, llm)
|
|
|
|
summary = extract_graph_for_source(
|
|
source_id,
|
|
user="owner-1",
|
|
chunks=[_chunk("c1", "Ada one."), _chunk("c2", "Ada two.")],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
assert summary["chunks_processed"] == 2
|
|
assert store.count_nodes(source_id) == 1
|
|
node = store.get_node_by_normalized(source_id, "ada")
|
|
assert node["doc_freq"] == 2
|
|
assert "first" in node["description"]
|
|
assert "second" in node["description"]
|
|
finally:
|
|
store.delete_by_source(source_id)
|
|
|
|
def test_checkpoint_skips_done_chunks(
|
|
self, store, source_id, monkeypatch, stub_embedding
|
|
):
|
|
try:
|
|
payload = _extraction_json(
|
|
entities=[{"name": "Ada", "type": "person", "description": "d"}],
|
|
relationships=[],
|
|
)
|
|
first_llm = _StubLLM([payload])
|
|
_install_stub_llm(monkeypatch, first_llm)
|
|
extract_graph_for_source(
|
|
source_id,
|
|
user="owner-1",
|
|
chunks=[_chunk("c1", "Ada.")],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
assert len(first_llm.gen_calls) == 1
|
|
|
|
second_llm = _StubLLM([])
|
|
_install_stub_llm(monkeypatch, second_llm)
|
|
summary = extract_graph_for_source(
|
|
source_id,
|
|
user="owner-1",
|
|
chunks=[_chunk("c1", "Ada.")],
|
|
config=SourceConfig(),
|
|
request_id="req-2",
|
|
)
|
|
assert len(second_llm.gen_calls) == 0
|
|
assert summary["chunks_processed"] == 0
|
|
finally:
|
|
store.delete_by_source(source_id)
|
|
|
|
def test_cap_limits_processing(
|
|
self, store, source_id, monkeypatch, stub_embedding
|
|
):
|
|
try:
|
|
payload = _extraction_json(
|
|
entities=[{"name": "X", "type": "t", "description": "d"}],
|
|
relationships=[],
|
|
)
|
|
llm = _StubLLM([payload, payload])
|
|
_install_stub_llm(monkeypatch, llm)
|
|
|
|
config = SourceConfig.model_validate({"graph": {"max_chunks": 2}})
|
|
summary = extract_graph_for_source(
|
|
source_id,
|
|
user="owner-1",
|
|
chunks=[_chunk(f"c{i}", f"text {i}") for i in range(5)],
|
|
config=config,
|
|
request_id="req-1",
|
|
)
|
|
|
|
assert len(llm.gen_calls) == 2
|
|
assert summary["chunks_processed"] == 2
|
|
assert summary["skipped_over_cap"] == 3
|
|
finally:
|
|
store.delete_by_source(source_id)
|
|
|
|
def test_malformed_and_error_chunks_are_skipped(
|
|
self, store, source_id, monkeypatch, stub_embedding
|
|
):
|
|
try:
|
|
good = _extraction_json(
|
|
entities=[{"name": "Ada", "type": "person", "description": "d"}],
|
|
relationships=[],
|
|
)
|
|
llm = _StubLLM([
|
|
"not json at all",
|
|
RuntimeError("model exploded"),
|
|
good,
|
|
])
|
|
_install_stub_llm(monkeypatch, llm)
|
|
|
|
summary = extract_graph_for_source(
|
|
source_id,
|
|
user="owner-1",
|
|
chunks=[
|
|
_chunk("c1", "garbage"),
|
|
_chunk("c2", "boom"),
|
|
_chunk("c3", "Ada."),
|
|
],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
assert summary["failed_chunks"] == 2
|
|
assert summary["chunks_processed"] == 1
|
|
assert store.count_nodes(source_id) == 1
|
|
progress = store.get_progress(source_id)
|
|
assert progress["c1"] == "failed"
|
|
assert progress["c2"] == "failed"
|
|
assert progress["c3"] == "done"
|
|
finally:
|
|
store.delete_by_source(source_id)
|
|
|
|
def test_exactly_one_gen_per_chunk(
|
|
self, store, source_id, monkeypatch, stub_embedding
|
|
):
|
|
try:
|
|
payload = _extraction_json(
|
|
entities=[{"name": "A", "type": "t", "description": "d"}],
|
|
relationships=[],
|
|
)
|
|
llm = _StubLLM([payload, payload, payload])
|
|
_install_stub_llm(monkeypatch, llm)
|
|
|
|
extract_graph_for_source(
|
|
source_id,
|
|
user="owner-1",
|
|
chunks=[_chunk(f"c{i}", f"text {i}") for i in range(3)],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
assert len(llm.gen_calls) == 3
|
|
finally:
|
|
store.delete_by_source(source_id)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExtractionTokenUsage:
|
|
def test_llm_tagged_for_token_usage(self, monkeypatch):
|
|
llm = _StubLLM([])
|
|
captured = _install_stub_llm(monkeypatch, llm)
|
|
|
|
built = extraction_module._build_extraction_llm(
|
|
"stub-model", user="owner-1", request_id="req-99"
|
|
)
|
|
|
|
assert built is llm
|
|
assert built._token_usage_source == "graph_extraction"
|
|
assert built._request_id == "req-99"
|
|
assert captured["model_id"] == "stub-model"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestModelResolution:
|
|
def test_per_source_override_wins(self, monkeypatch):
|
|
monkeypatch.setattr(
|
|
extraction_module.settings, "GRAPHRAG_EXTRACTION_MODEL", "setting-model"
|
|
)
|
|
monkeypatch.setattr(extraction_module.settings, "LLM_NAME", "instance-model")
|
|
config = SourceConfig.model_validate(
|
|
{"graph": {"extraction_model": "override-model"}}
|
|
)
|
|
assert (
|
|
extraction_module._resolve_extraction_model(config) == "override-model"
|
|
)
|
|
|
|
def test_setting_then_instance_default(self, monkeypatch):
|
|
monkeypatch.setattr(
|
|
extraction_module.settings, "GRAPHRAG_EXTRACTION_MODEL", "setting-model"
|
|
)
|
|
monkeypatch.setattr(extraction_module.settings, "LLM_NAME", "instance-model")
|
|
assert (
|
|
extraction_module._resolve_extraction_model(SourceConfig())
|
|
== "setting-model"
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
extraction_module.settings, "GRAPHRAG_EXTRACTION_MODEL", None
|
|
)
|
|
assert (
|
|
extraction_module._resolve_extraction_model(SourceConfig())
|
|
== "instance-model"
|
|
)
|
|
|
|
def test_max_chunks_resolution(self, monkeypatch):
|
|
monkeypatch.setattr(
|
|
extraction_module.settings,
|
|
"GRAPHRAG_MAX_CHUNKS_FOR_EXTRACTION",
|
|
2000,
|
|
)
|
|
assert extraction_module._resolve_max_chunks(SourceConfig()) == 2000
|
|
config = SourceConfig.model_validate({"graph": {"max_chunks": 5}})
|
|
assert extraction_module._resolve_max_chunks(config) == 5
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExtractionProviderResolution:
|
|
"""The extraction model decides the provider, not ``LLM_PROVIDER``.
|
|
|
|
``settings.LLM_PROVIDER`` is the deployment default (``docsgpt`` out of the
|
|
box, i.e. the hosted public endpoint). Dispatching the resolved extraction
|
|
model through it sends the call to a provider that never serves that model:
|
|
the request is rejected, the shared fallback answers instead, and the graph
|
|
is quietly built by a different model than the one configured.
|
|
"""
|
|
|
|
def _capture_create_llm(self, monkeypatch, llm=None):
|
|
captured = {}
|
|
|
|
def _create(provider, *args, **kwargs):
|
|
captured["provider"] = provider
|
|
captured["args"] = args
|
|
captured["kwargs"] = kwargs
|
|
return llm or _StubLLM([])
|
|
|
|
monkeypatch.setattr(
|
|
extraction_module.LLMCreator, "create_llm", staticmethod(_create)
|
|
)
|
|
return captured
|
|
|
|
def test_provider_comes_from_the_model_registry(self, monkeypatch):
|
|
monkeypatch.setattr(extraction_module.settings, "LLM_PROVIDER", "docsgpt")
|
|
monkeypatch.setattr(
|
|
extraction_module, "get_provider_from_model_id", lambda *a, **k: "openai"
|
|
)
|
|
monkeypatch.setattr(
|
|
extraction_module, "get_api_key_for_provider", lambda provider: "sk-openai"
|
|
)
|
|
captured = self._capture_create_llm(monkeypatch)
|
|
|
|
extraction_module._build_extraction_llm("gpt-4o-mini", "owner-1", "req-1")
|
|
|
|
assert captured["provider"] == "openai"
|
|
assert captured["kwargs"]["api_key"] == "sk-openai"
|
|
assert captured["kwargs"]["model_id"] == "gpt-4o-mini"
|
|
|
|
def test_owner_scopes_the_registry_lookup(self, monkeypatch):
|
|
"""A per-user (BYOM) model only resolves when the owner is passed."""
|
|
seen = {}
|
|
|
|
def _resolve(model_id, user_id=None):
|
|
seen["model_id"] = model_id
|
|
seen["user_id"] = user_id
|
|
return "anthropic"
|
|
|
|
monkeypatch.setattr(
|
|
extraction_module, "get_provider_from_model_id", _resolve
|
|
)
|
|
monkeypatch.setattr(
|
|
extraction_module, "get_api_key_for_provider", lambda provider: "k"
|
|
)
|
|
self._capture_create_llm(monkeypatch)
|
|
|
|
extraction_module._build_extraction_llm("byom-uuid", "owner-7", "req-1")
|
|
|
|
assert seen == {"model_id": "byom-uuid", "user_id": "owner-7"}
|
|
|
|
def test_unknown_model_falls_back_to_the_configured_provider(self, monkeypatch):
|
|
monkeypatch.setattr(extraction_module.settings, "LLM_PROVIDER", "docsgpt")
|
|
monkeypatch.setattr(
|
|
extraction_module, "get_provider_from_model_id", lambda *a, **k: None
|
|
)
|
|
monkeypatch.setattr(
|
|
extraction_module, "get_api_key_for_provider", lambda provider: "fallback-key"
|
|
)
|
|
captured = self._capture_create_llm(monkeypatch)
|
|
|
|
extraction_module._build_extraction_llm("mystery-model", "owner-1", "req-1")
|
|
|
|
assert captured["provider"] == "docsgpt"
|
|
assert captured["kwargs"]["api_key"] == "fallback-key"
|
|
|
|
def test_no_model_id_skips_the_lookup(self, monkeypatch):
|
|
monkeypatch.setattr(extraction_module.settings, "LLM_PROVIDER", "openai")
|
|
calls = []
|
|
monkeypatch.setattr(
|
|
extraction_module,
|
|
"get_provider_from_model_id",
|
|
lambda *a, **k: calls.append(a) or "anthropic",
|
|
)
|
|
monkeypatch.setattr(
|
|
extraction_module, "get_api_key_for_provider", lambda provider: "k"
|
|
)
|
|
captured = self._capture_create_llm(monkeypatch)
|
|
|
|
extraction_module._build_extraction_llm(None, "owner-1", "req-1")
|
|
|
|
assert calls == []
|
|
assert captured["provider"] == "openai"
|
|
|
|
def test_api_key_follows_the_resolved_provider(self, monkeypatch):
|
|
"""The key must match the provider actually dispatched to."""
|
|
monkeypatch.setattr(extraction_module.settings, "LLM_PROVIDER", "docsgpt")
|
|
monkeypatch.setattr(extraction_module.settings, "API_KEY", "generic-key")
|
|
monkeypatch.setattr(
|
|
extraction_module, "get_provider_from_model_id", lambda *a, **k: "anthropic"
|
|
)
|
|
keyed_for = {}
|
|
|
|
def _key(provider):
|
|
keyed_for["provider"] = provider
|
|
return "sk-anthropic"
|
|
|
|
monkeypatch.setattr(extraction_module, "get_api_key_for_provider", _key)
|
|
captured = self._capture_create_llm(monkeypatch)
|
|
|
|
extraction_module._build_extraction_llm("claude-x", "owner-1", "req-1")
|
|
|
|
assert keyed_for["provider"] == "anthropic"
|
|
assert captured["kwargs"]["api_key"] == "sk-anthropic"
|
|
assert captured["kwargs"]["api_key"] != "generic-key"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestFailedChunksAreReported:
|
|
"""Every dropped chunk has to leave a trace.
|
|
|
|
A chunk whose extraction cannot be parsed is marked ``failed`` and skipped.
|
|
That path logged nothing at all, so a graph could come back short with the
|
|
summary's ``failed_chunks`` count as the only hint and no way to tell which
|
|
chunk, or why, from the logs.
|
|
"""
|
|
|
|
def _fake_store(self, monkeypatch, chunk_ids):
|
|
from unittest.mock import MagicMock
|
|
|
|
store = MagicMock(name="GraphStore")
|
|
store.pending_chunks.return_value = list(chunk_ids)
|
|
store.apply_chunk.return_value = (1, 0)
|
|
store.count_nodes.return_value = 1
|
|
monkeypatch.setattr(
|
|
"docsgpt.graphrag.store.GraphStore", lambda *a, **k: store
|
|
)
|
|
return store
|
|
|
|
def test_unparseable_output_is_logged_with_the_chunk_id(
|
|
self, monkeypatch, caplog, stub_embedding
|
|
):
|
|
import logging
|
|
|
|
store = self._fake_store(monkeypatch, ["c1"])
|
|
_install_stub_llm(monkeypatch, _StubLLM(["not json at all"]))
|
|
|
|
with caplog.at_level(logging.WARNING, logger="docsgpt.graphrag.extraction"):
|
|
summary = extract_graph_for_source(
|
|
str(uuid.uuid4()),
|
|
user="owner-1",
|
|
chunks=[_chunk("c1", "some text")],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
assert summary["failed_chunks"] == 1
|
|
store.mark_chunk.assert_called_once()
|
|
assert store.mark_chunk.call_args.args[2] == "failed"
|
|
messages = [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING]
|
|
assert any("c1" in message for message in messages), messages
|
|
|
|
def test_llm_errors_still_name_the_chunk(
|
|
self, monkeypatch, caplog, stub_embedding
|
|
):
|
|
import logging
|
|
|
|
self._fake_store(monkeypatch, ["c7"])
|
|
_install_stub_llm(monkeypatch, _StubLLM([RuntimeError("model exploded")]))
|
|
|
|
with caplog.at_level(logging.WARNING, logger="docsgpt.graphrag.extraction"):
|
|
extract_graph_for_source(
|
|
str(uuid.uuid4()),
|
|
user="owner-1",
|
|
chunks=[_chunk("c7", "some text")],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
messages = [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING]
|
|
assert any("c7" in message for message in messages), messages
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestSummaryNodeCount:
|
|
"""``nodes`` must describe the graph, not the number of upserts."""
|
|
|
|
@pytest.fixture
|
|
def store(self, monkeypatch, postgresql):
|
|
store = _live_store(monkeypatch, postgresql.info)
|
|
yield store
|
|
store.close()
|
|
|
|
def test_repeated_entity_counts_once(
|
|
self, store, monkeypatch, stub_embedding
|
|
):
|
|
source_id = str(uuid.uuid4())
|
|
try:
|
|
payload = _extraction_json(
|
|
entities=[{"name": "Ada", "type": "person", "description": "d"}],
|
|
relationships=[],
|
|
)
|
|
_install_stub_llm(monkeypatch, _StubLLM([payload, payload]))
|
|
|
|
summary = extract_graph_for_source(
|
|
source_id,
|
|
user="owner-1",
|
|
chunks=[_chunk("c1", "Ada one."), _chunk("c2", "Ada two.")],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
# Two chunks upserted the same entity: one node in the graph.
|
|
assert store.count_nodes(source_id) == 1
|
|
assert summary["nodes"] == 1
|
|
assert summary["chunks_processed"] == 2
|
|
finally:
|
|
store.delete_by_source(source_id)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSummaryCountFailure:
|
|
"""A broken count query must not be reported as an empty graph."""
|
|
|
|
def test_a_failed_count_reports_the_write_count(
|
|
self, monkeypatch, stub_embedding
|
|
):
|
|
from unittest.mock import MagicMock
|
|
|
|
store = MagicMock(name="GraphStore")
|
|
store.pending_chunks.return_value = ["c1"]
|
|
store.apply_chunk.return_value = (2, 1)
|
|
store.count_nodes.side_effect = RuntimeError("count query failed")
|
|
monkeypatch.setattr(
|
|
"docsgpt.graphrag.store.GraphStore", lambda *a, **k: store
|
|
)
|
|
_install_stub_llm(
|
|
monkeypatch,
|
|
_StubLLM([_extraction_json([{"name": "Ada"}], [])]),
|
|
)
|
|
|
|
summary = extract_graph_for_source(
|
|
str(uuid.uuid4()),
|
|
user="owner-1",
|
|
chunks=[_chunk("c1", "Ada.")],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
# Falls back to what was actually written, not to zero.
|
|
assert summary["nodes"] == 2
|
|
# And it asked for a count that raises rather than one that returns 0,
|
|
# or the fallback above could never run.
|
|
assert store.count_nodes.call_args.kwargs.get("strict") is True
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestParsing:
|
|
def test_parses_embedded_json(self):
|
|
raw = 'sure!\n{"entities": [{"name": "A"}], "relationships": []}\nthanks'
|
|
parsed = extraction_module._parse_extraction(raw)
|
|
assert parsed["entities"] == [{"name": "A"}]
|
|
assert parsed["relationships"] == []
|
|
|
|
def test_garbage_returns_none(self):
|
|
assert extraction_module._parse_extraction("no json here") is None
|
|
assert extraction_module._parse_extraction("{bad json}") is None
|
|
assert extraction_module._parse_extraction(None) is None
|
|
|
|
def test_missing_keys_default_empty(self):
|
|
parsed = extraction_module._parse_extraction('{"foo": 1}')
|
|
assert parsed == {"entities": [], "relationships": []}
|
|
|
|
def test_chunk_id_prefers_doc_id(self):
|
|
assert extraction_module._chunk_id({"doc_id": "7"}) == "7"
|
|
assert extraction_module._chunk_id({"chunk_id": "abc"}) == "abc"
|
|
assert extraction_module._chunk_id({"id": 9}) == "9"
|
|
assert extraction_module._chunk_id({"text": "no id"}) is None
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestEmbeddingsResolution:
|
|
def test_extraction_uses_shared_resolver(self, monkeypatch):
|
|
"""Extraction must resolve embeddings through ``get_embeddings``."""
|
|
from unittest.mock import MagicMock
|
|
|
|
fake_store = MagicMock()
|
|
fake_store.pending_chunks.return_value = []
|
|
monkeypatch.setattr(
|
|
"docsgpt.graphrag.store.GraphStore", lambda *a, **k: fake_store
|
|
)
|
|
_install_stub_llm(monkeypatch, _StubLLM([]))
|
|
|
|
calls = []
|
|
fake_embedding = MagicMock()
|
|
|
|
def _resolver(*args, **kwargs):
|
|
calls.append((args, kwargs))
|
|
return fake_embedding
|
|
|
|
monkeypatch.setattr(extraction_module, "get_embeddings", _resolver)
|
|
|
|
summary = extract_graph_for_source(
|
|
str(uuid.uuid4()),
|
|
user="owner-1",
|
|
chunks=[],
|
|
config=SourceConfig(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
assert calls == [((), {})]
|
|
assert summary["chunks_processed"] == 0
|