chore: import upstream snapshot with attribution
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,421 @@
|
||||
"""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
|
||||
pgvector DB with a unique temp ``source_id`` and torn down via
|
||||
``delete_by_source``; if no DB is reachable the live tests skip.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
import application.graphrag.extraction as extraction_module
|
||||
from application.graphrag.store import GraphStore
|
||||
from application.storage.db.source_config import SourceConfig
|
||||
|
||||
extract_graph_for_source = extraction_module.extract_graph_for_source
|
||||
|
||||
TEST_EMBEDDING_DIM = 8
|
||||
|
||||
|
||||
def _resolve_connection_string():
|
||||
from application.core.settings import settings
|
||||
|
||||
conn = getattr(settings, "PGVECTOR_CONNECTION_STRING", None)
|
||||
if not conn and getattr(settings, "POSTGRES_URI", None):
|
||||
from application.core.db_uri import normalize_pgvector_connection_string
|
||||
|
||||
conn = normalize_pgvector_connection_string(settings.POSTGRES_URI)
|
||||
return conn
|
||||
|
||||
|
||||
_GRAPH_TABLES = (
|
||||
"graph_node_chunks",
|
||||
"graph_edges",
|
||||
"graph_nodes",
|
||||
"graph_ingest_progress",
|
||||
)
|
||||
|
||||
|
||||
def _drop_graph_tables(conn_string):
|
||||
import psycopg
|
||||
|
||||
with psycopg.connect(conn_string) as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(
|
||||
f"DROP TABLE IF EXISTS {', '.join(_GRAPH_TABLES)} CASCADE;"
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _live_store(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
GraphStore, "_embedding_dim", lambda self: TEST_EMBEDDING_DIM
|
||||
)
|
||||
conn = _resolve_connection_string()
|
||||
if not conn:
|
||||
pytest.skip("No pgvector connection string configured")
|
||||
try:
|
||||
_drop_graph_tables(conn)
|
||||
store = GraphStore(connection_string=conn)
|
||||
except Exception as exc:
|
||||
pytest.skip(f"pgvector DB not reachable: {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):
|
||||
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):
|
||||
store = _live_store(monkeypatch)
|
||||
yield store
|
||||
_drop_graph_tables(store._connection_string)
|
||||
|
||||
@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 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
|
||||
@@ -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
|
||||
@@ -0,0 +1,475 @@
|
||||
"""Tests for the GraphRAG GraphStore (on-demand tables in the pgvector DB).
|
||||
|
||||
Two layers:
|
||||
|
||||
* A live-pg integration test that exercises the real DDL + SQL against the
|
||||
pgvector store DB (same connection-string source as ``PGVectorStore``). It
|
||||
uses a unique temp ``source_id`` and tears down every row it creates.
|
||||
* A mock-cursor test that asserts the parameterized SQL shapes — ``source_id``
|
||||
and embeddings are bound params, never interpolated.
|
||||
|
||||
The embedding dimension is mocked everywhere so the suite never loads the real
|
||||
SentenceTransformer model: the live store creates ``TEST_EMBEDDING_DIM`` vectors
|
||||
and the helpers build matching ones.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
import application.graphrag.store as store_module
|
||||
|
||||
GraphStore = store_module.GraphStore
|
||||
|
||||
TEST_EMBEDDING_DIM = 8
|
||||
|
||||
_REAL_EMBEDDING_DIM = GraphStore._embedding_dim
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _mock_embedding_dim(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
GraphStore, "_embedding_dim", lambda self: TEST_EMBEDDING_DIM
|
||||
)
|
||||
|
||||
|
||||
def _resolve_connection_string():
|
||||
from application.core.settings import settings
|
||||
|
||||
conn = getattr(settings, "PGVECTOR_CONNECTION_STRING", None)
|
||||
if not conn and getattr(settings, "POSTGRES_URI", None):
|
||||
from application.core.db_uri import normalize_pgvector_connection_string
|
||||
|
||||
conn = normalize_pgvector_connection_string(settings.POSTGRES_URI)
|
||||
return conn
|
||||
|
||||
|
||||
_GRAPH_TABLES = (
|
||||
"graph_node_chunks",
|
||||
"graph_edges",
|
||||
"graph_nodes",
|
||||
"graph_ingest_progress",
|
||||
)
|
||||
|
||||
|
||||
def _drop_graph_tables(conn_string):
|
||||
import psycopg
|
||||
|
||||
with psycopg.connect(conn_string) as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(
|
||||
f"DROP TABLE IF EXISTS {', '.join(_GRAPH_TABLES)} CASCADE;"
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _live_store():
|
||||
conn = _resolve_connection_string()
|
||||
if not conn:
|
||||
pytest.skip("No pgvector connection string configured")
|
||||
try:
|
||||
_drop_graph_tables(conn)
|
||||
store = GraphStore(connection_string=conn)
|
||||
except Exception as exc:
|
||||
pytest.skip(f"pgvector DB not reachable: {exc}")
|
||||
return store
|
||||
|
||||
|
||||
def _embedding(seed: float) -> list:
|
||||
vec = [0.0] * TEST_EMBEDDING_DIM
|
||||
vec[0] = seed
|
||||
return vec
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestGraphStoreLive:
|
||||
@pytest.fixture
|
||||
def store(self):
|
||||
store = _live_store()
|
||||
yield store
|
||||
_drop_graph_tables(store._connection_string)
|
||||
|
||||
@pytest.fixture
|
||||
def source_id(self):
|
||||
return str(uuid.uuid4())
|
||||
|
||||
def test_ensure_tables_idempotent(self, store):
|
||||
store._ensure_tables()
|
||||
store._ensure_tables()
|
||||
|
||||
def test_upsert_node_merges_by_normalized_name(self, store, source_id):
|
||||
try:
|
||||
first = store.upsert_node(
|
||||
source_id=source_id,
|
||||
name="Ada Lovelace",
|
||||
normalized_name="ada lovelace",
|
||||
type="person",
|
||||
description="A mathematician.",
|
||||
name_embedding=_embedding(1.0),
|
||||
)
|
||||
second = store.upsert_node(
|
||||
source_id=source_id,
|
||||
name="Ada Lovelace",
|
||||
normalized_name="ada lovelace",
|
||||
type="person",
|
||||
description="Wrote the first algorithm.",
|
||||
)
|
||||
assert first == second
|
||||
|
||||
node = store.get_node_by_normalized(source_id, "ada lovelace")
|
||||
assert node is not None
|
||||
assert node["id"] == first
|
||||
assert node["doc_freq"] == 2
|
||||
assert "mathematician" in node["description"]
|
||||
assert "first algorithm" in node["description"]
|
||||
|
||||
duplicate = store.upsert_node(
|
||||
source_id=source_id,
|
||||
name="Ada Lovelace",
|
||||
normalized_name="ada lovelace",
|
||||
description="Wrote the first algorithm.",
|
||||
)
|
||||
assert duplicate == first
|
||||
node = store.get_node_by_normalized(source_id, "ada lovelace")
|
||||
assert node["description"].count("first algorithm") == 1
|
||||
|
||||
assert store.count_nodes(source_id) == 1
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_add_edge_and_link_chunk(self, store, source_id):
|
||||
try:
|
||||
a = store.upsert_node(source_id, "A", "a", "thing", "desc a")
|
||||
b = store.upsert_node(source_id, "B", "b", "thing", "desc b")
|
||||
store.add_edge(
|
||||
source_id, a, b, "related", "a relates to b", 2.0, ["chunk-1"]
|
||||
)
|
||||
store.link_node_chunk(source_id, a, "chunk-1")
|
||||
store.link_node_chunk(source_id, a, "chunk-1")
|
||||
store.link_node_chunk(source_id, b, "chunk-1")
|
||||
|
||||
mapping = store.get_chunk_ids_for_nodes(source_id, [a, b])
|
||||
assert mapping[a] == ["chunk-1"]
|
||||
assert mapping[b] == ["chunk-1"]
|
||||
|
||||
store.set_node_degrees(source_id)
|
||||
node_a = store.get_node_by_normalized(source_id, "a")
|
||||
assert node_a["degree"] == 1
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_apply_chunk_writes_nodes_links_and_edges(self, store, source_id):
|
||||
"""One transactional write: entities linked to the chunk, edges added,
|
||||
and a bare relationship endpoint upserted but not chunk-linked."""
|
||||
try:
|
||||
entities = [
|
||||
{"name": "Ada", "normalized_name": "ada", "type": "person",
|
||||
"description": "mathematician"},
|
||||
{"name": "Engine", "normalized_name": "engine", "type": "machine",
|
||||
"description": None},
|
||||
]
|
||||
relationships = [
|
||||
{"source": "Ada", "target": "Engine", "type": "designed",
|
||||
"description": "Ada designed the Engine", "weight": 2.0},
|
||||
# 'Babbage' is only an endpoint — upserted edge-only.
|
||||
{"source": "Babbage", "target": "Engine", "type": "built",
|
||||
"description": None, "weight": 1.0},
|
||||
]
|
||||
name_embeddings = {
|
||||
"ada": [0.1] * store._embedding_dim(),
|
||||
"engine": [0.2] * store._embedding_dim(),
|
||||
"babbage": [0.3] * store._embedding_dim(),
|
||||
}
|
||||
|
||||
nodes, edges = store.apply_chunk(
|
||||
source_id, "c1", entities, relationships, name_embeddings
|
||||
)
|
||||
assert nodes == 2 # only entities are counted
|
||||
assert edges == 2
|
||||
|
||||
ada = store.get_node_by_normalized(source_id, "ada")
|
||||
engine = store.get_node_by_normalized(source_id, "engine")
|
||||
babbage = store.get_node_by_normalized(source_id, "babbage")
|
||||
assert ada is not None and engine is not None
|
||||
assert babbage is not None # endpoint upserted
|
||||
|
||||
mapping = store.get_chunk_ids_for_nodes(
|
||||
source_id, [ada["id"], engine["id"], babbage["id"]]
|
||||
)
|
||||
assert mapping[ada["id"]] == ["c1"]
|
||||
assert mapping[engine["id"]] == ["c1"]
|
||||
# Bare endpoint is not linked to the chunk.
|
||||
assert babbage["id"] not in mapping
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_self_loop_degree_agrees_across_paths(self, store, source_id):
|
||||
"""``add_edge``'s incremental +1 and ``set_node_degrees`` recompute must
|
||||
agree on a self-loop (count it once)."""
|
||||
try:
|
||||
node = store.upsert_node(source_id, "Solo", "solo")
|
||||
store.add_edge(source_id, node, node, "self")
|
||||
|
||||
incremental = store.get_node_by_normalized(source_id, "solo")["degree"]
|
||||
assert incremental == 1
|
||||
|
||||
store.set_node_degrees(source_id)
|
||||
recomputed = store.get_node_by_normalized(source_id, "solo")["degree"]
|
||||
assert recomputed == 1
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_search_nodes_by_embedding(self, store, source_id):
|
||||
try:
|
||||
near = store.upsert_node(
|
||||
source_id, "Near", "near", "thing", "d", _embedding(1.0)
|
||||
)
|
||||
store.upsert_node(
|
||||
source_id, "Far", "far", "thing", "d", _embedding(-1.0)
|
||||
)
|
||||
results = store.search_nodes_by_embedding(source_id, _embedding(1.0), k=2)
|
||||
assert len(results) == 2
|
||||
assert results[0]["id"] == near
|
||||
assert results[0]["distance"] <= results[1]["distance"]
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_get_subgraph_bounded(self, store, source_id):
|
||||
try:
|
||||
a = store.upsert_node(source_id, "A", "a")
|
||||
b = store.upsert_node(source_id, "B", "b")
|
||||
c = store.upsert_node(source_id, "C", "c")
|
||||
store.add_edge(source_id, a, b, "rel")
|
||||
store.add_edge(source_id, b, c, "rel")
|
||||
|
||||
one_hop = store.get_subgraph(source_id, [a], hops=1)
|
||||
node_ids = {n["id"] for n in one_hop["nodes"]}
|
||||
assert a in node_ids and b in node_ids
|
||||
assert c not in node_ids
|
||||
|
||||
two_hop = store.get_subgraph(source_id, [a], hops=2)
|
||||
node_ids = {n["id"] for n in two_hop["nodes"]}
|
||||
assert {a, b, c} <= node_ids
|
||||
assert len(two_hop["edges"]) >= 2
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_get_subgraph_frontier_truncation_is_deterministic(
|
||||
self, store, source_id, monkeypatch
|
||||
):
|
||||
"""Bounded expansion must pick the same neighbors run-to-run so PPR (G5)
|
||||
is reproducible."""
|
||||
try:
|
||||
hub = store.upsert_node(source_id, "Hub", "hub")
|
||||
leaves = []
|
||||
for i in range(6):
|
||||
leaf = store.upsert_node(source_id, f"L{i}", f"l{i}")
|
||||
store.add_edge(source_id, hub, leaf, "rel")
|
||||
leaves.append(leaf)
|
||||
|
||||
monkeypatch.setattr(store_module, "MAX_SUBGRAPH_NODES", 4)
|
||||
|
||||
first = {n["id"] for n in store.get_subgraph(source_id, [hub])["nodes"]}
|
||||
second = {n["id"] for n in store.get_subgraph(source_id, [hub])["nodes"]}
|
||||
assert first == second
|
||||
assert len(first) == 4
|
||||
|
||||
kept_leaves = sorted(leaves)[:3]
|
||||
assert first == {hub, *kept_leaves}
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_get_graph_overview_bounded_by_degree(self, store, source_id):
|
||||
try:
|
||||
hub = store.upsert_node(source_id, "Hub", "hub")
|
||||
leaves = [
|
||||
store.upsert_node(source_id, f"L{i}", f"l{i}") for i in range(4)
|
||||
]
|
||||
for leaf in leaves:
|
||||
store.add_edge(source_id, hub, leaf, "rel")
|
||||
store.set_node_degrees(source_id)
|
||||
|
||||
overview = store.get_graph_overview(source_id, limit=3)
|
||||
node_ids = [n["id"] for n in overview["nodes"]]
|
||||
assert len(node_ids) == 3
|
||||
# The hub has the highest degree, so it must lead the bounded set.
|
||||
assert node_ids[0] == hub
|
||||
# Edges only connect nodes that survived the limit.
|
||||
for edge in overview["edges"]:
|
||||
assert edge["source"] in node_ids
|
||||
assert edge["target"] in node_ids
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_get_graph_overview_empty_source(self, store, source_id):
|
||||
overview = store.get_graph_overview(source_id)
|
||||
assert overview == {"nodes": [], "edges": []}
|
||||
|
||||
def test_get_node_detail_with_linked_chunks(self, store, source_id):
|
||||
try:
|
||||
node = store.upsert_node(
|
||||
source_id, "Ada", "ada", "person", "A mathematician."
|
||||
)
|
||||
store.link_node_chunk(source_id, node, "chunk-1")
|
||||
|
||||
detail = store.get_node_detail(source_id, node)
|
||||
assert detail is not None
|
||||
assert detail["name"] == "Ada"
|
||||
assert detail["description"] == "A mathematician."
|
||||
chunk_ids = [c["chunk_id"] for c in detail["chunks"]]
|
||||
assert "chunk-1" in chunk_ids
|
||||
|
||||
assert store.get_node_detail(source_id, str(uuid.uuid4())) is None
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_checkpoint_pending_and_mark(self, store, source_id):
|
||||
try:
|
||||
all_chunks = ["c1", "c2", "c3"]
|
||||
assert store.pending_chunks(source_id, all_chunks) == all_chunks
|
||||
|
||||
store.mark_chunk(source_id, "c1", "done")
|
||||
store.mark_chunk(source_id, "c2", "pending")
|
||||
assert store.pending_chunks(source_id, all_chunks) == ["c2", "c3"]
|
||||
|
||||
store.mark_chunk(source_id, "c2", "done")
|
||||
assert store.pending_chunks(source_id, all_chunks) == ["c3"]
|
||||
|
||||
progress = store.get_progress(source_id)
|
||||
assert progress["c1"] == "done"
|
||||
assert progress["c2"] == "done"
|
||||
finally:
|
||||
store.delete_by_source(source_id)
|
||||
|
||||
def test_delete_by_source_isolation(self, store):
|
||||
keep = str(uuid.uuid4())
|
||||
drop = str(uuid.uuid4())
|
||||
try:
|
||||
k = store.upsert_node(keep, "K", "k")
|
||||
d = store.upsert_node(drop, "D", "d")
|
||||
store.add_edge(keep, k, k, "self")
|
||||
store.add_edge(drop, d, d, "self")
|
||||
store.link_node_chunk(keep, k, "kc")
|
||||
store.link_node_chunk(drop, d, "dc")
|
||||
store.mark_chunk(keep, "kc", "done")
|
||||
store.mark_chunk(drop, "dc", "done")
|
||||
|
||||
store.delete_by_source(drop)
|
||||
|
||||
assert store.count_nodes(drop) == 0
|
||||
assert store.get_progress(drop) == {}
|
||||
assert store.count_nodes(keep) == 1
|
||||
assert store.get_progress(keep) == {"kc": "done"}
|
||||
finally:
|
||||
store.delete_by_source(keep)
|
||||
store.delete_by_source(drop)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGraphStoreParameterization:
|
||||
"""Asserts SQL is parameterized without touching a real DB."""
|
||||
|
||||
def _store_with_mock_conn(self):
|
||||
store = GraphStore.__new__(GraphStore)
|
||||
cursor = MagicMock()
|
||||
cursor.fetchone.return_value = [str(uuid.uuid4())]
|
||||
cursor.fetchall.return_value = []
|
||||
conn = MagicMock()
|
||||
conn.cursor.return_value = cursor
|
||||
store._connection = conn
|
||||
store._get_connection = lambda: conn
|
||||
return store, cursor
|
||||
|
||||
def test_delete_by_source_binds_source_id(self):
|
||||
store, cursor = self._store_with_mock_conn()
|
||||
sid = str(uuid.uuid4())
|
||||
store.delete_by_source(sid)
|
||||
|
||||
for call in cursor.execute.call_args_list:
|
||||
sql = call.args[0]
|
||||
params = call.args[1] if len(call.args) > 1 else None
|
||||
assert "WHERE source_id = %s" in sql
|
||||
assert sid not in sql
|
||||
assert params == (sid,)
|
||||
|
||||
def test_search_binds_embedding_and_source(self):
|
||||
store, cursor = self._store_with_mock_conn()
|
||||
sid = str(uuid.uuid4())
|
||||
embedding = _embedding(0.5)
|
||||
store.search_nodes_by_embedding(sid, embedding, k=5)
|
||||
|
||||
sql, params = cursor.execute.call_args.args[0], cursor.execute.call_args.args[1]
|
||||
assert "%s::vector" in sql
|
||||
assert "source_id = %s" in sql
|
||||
assert sid not in sql
|
||||
assert str(embedding) not in sql
|
||||
assert params == (embedding, sid, embedding, 5)
|
||||
|
||||
def test_graph_overview_binds_source_and_clamps_limit(self):
|
||||
from application.graphrag.store import GRAPH_OVERVIEW_MAX_LIMIT
|
||||
|
||||
store, cursor = self._store_with_mock_conn()
|
||||
cursor.fetchall.return_value = []
|
||||
sid = str(uuid.uuid4())
|
||||
|
||||
store.get_graph_overview(sid, limit=10_000)
|
||||
|
||||
sql, params = (
|
||||
cursor.execute.call_args.args[0],
|
||||
cursor.execute.call_args.args[1],
|
||||
)
|
||||
assert "source_id = %s" in sql
|
||||
assert sid not in sql
|
||||
# An empty node fetch short-circuits; only the node query ran, and the
|
||||
# limit is clamped to the hard cap before binding.
|
||||
assert params == (sid, GRAPH_OVERVIEW_MAX_LIMIT)
|
||||
|
||||
def test_upsert_node_binds_all_values(self):
|
||||
store, cursor = self._store_with_mock_conn()
|
||||
sid = str(uuid.uuid4())
|
||||
embedding = _embedding(0.1)
|
||||
store.upsert_node(sid, "Name", "name", "type", "desc", embedding)
|
||||
|
||||
sql, params = cursor.execute.call_args.args[0], cursor.execute.call_args.args[1]
|
||||
assert "ON CONFLICT (source_id, normalized_name) DO UPDATE" in sql
|
||||
assert sid not in sql
|
||||
assert "name" not in [t for t in sql.split() if t == sid]
|
||||
assert params[1] == sid
|
||||
assert params[-1] == embedding
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestEmbeddingDim:
|
||||
"""The graph table dimension is derived from the configured model (FIX 1)."""
|
||||
|
||||
def test_uses_configured_model_dimension(self, monkeypatch):
|
||||
from application.vectorstore import base as base_module
|
||||
|
||||
fake_embedding = MagicMock()
|
||||
fake_embedding.dimension = 1536
|
||||
monkeypatch.setattr(
|
||||
base_module.EmbeddingsSingleton,
|
||||
"get_instance",
|
||||
staticmethod(lambda *a, **k: fake_embedding),
|
||||
)
|
||||
monkeypatch.setattr(GraphStore, "_embedding_dim", _REAL_EMBEDDING_DIM)
|
||||
|
||||
store = GraphStore.__new__(GraphStore)
|
||||
assert store._embedding_dim() == 1536
|
||||
|
||||
def test_falls_back_to_default_dimension(self, monkeypatch):
|
||||
from application.vectorstore import base as base_module
|
||||
|
||||
fake_embedding = object()
|
||||
monkeypatch.setattr(
|
||||
base_module.EmbeddingsSingleton,
|
||||
"get_instance",
|
||||
staticmethod(lambda *a, **k: fake_embedding),
|
||||
)
|
||||
monkeypatch.setattr(GraphStore, "_embedding_dim", _REAL_EMBEDDING_DIM)
|
||||
|
||||
store = GraphStore.__new__(GraphStore)
|
||||
assert store._embedding_dim() == store_module.DEFAULT_NAME_EMBEDDING_DIM
|
||||
Reference in New Issue
Block a user