Spaces:
Paused
Paused
File size: 2,594 Bytes
8c1b9fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | from __future__ import annotations
import pytest
from auralynq.embeddings import get_embedder
from auralynq.ingest.models import Chunk, SourceType
from auralynq.retrieval.pathrag.builder import build_from_chunks
from auralynq.retrieval.pathrag.retriever import PathRAGRetriever
from auralynq.vectorstore.memory_store import MemoryStore
@pytest.fixture
def graph_and_store(tmp_path):
chunks = [
Chunk(
id="c0",
doc_id="d",
ordinal=0,
source="geo.txt",
source_type=SourceType.text,
text="Paris is the capital of France. France is located in Europe.",
),
Chunk(
id="c1",
doc_id="d",
ordinal=1,
source="geo2.txt",
source_type=SourceType.text,
text="France is located in Europe. Europe contains many countries.",
),
Chunk(
id="c2",
doc_id="d",
ordinal=2,
source="geo3.txt",
source_type=SourceType.text,
text="The Seine flows through Paris.",
),
]
kg = build_from_chunks(chunks)
emb = get_embedder()
store = MemoryStore(path=tmp_path / "ms")
store.upsert(chunks, emb.embed([c.text for c in chunks]))
return kg, store
def test_seed_entities_found(graph_and_store):
kg, store = graph_and_store
r = PathRAGRetriever(kg, store=store)
seeds = r.seed_entities("Tell me about Paris and France")
assert seeds
def test_pathrag_returns_paths_and_chunks(graph_and_store):
kg, store = graph_and_store
r = PathRAGRetriever(kg, store=store, max_hops=3)
res = r.retrieve("How is Paris connected to Europe through France?", k=5)
assert res.method == "pathrag"
assert res.metadata["seeds"]
assert res.metadata["paths"], "expected at least one relational path"
# Paths should be reliability-scored and rendered to text.
first = res.metadata["paths"][0]
assert "reliability" in first and first["text"]
def test_pathrag_flow_pruning_bounds_paths(graph_and_store):
kg, store = graph_and_store
r = PathRAGRetriever(kg, store=store, max_paths=2)
res = r.retrieve("Paris France Europe relationship", k=5)
assert len(res.metadata["paths"]) <= 2
def test_pathrag_empty_graph_is_safe(tmp_path):
from auralynq.retrieval.pathrag.graph import KnowledgeGraph
r = PathRAGRetriever(KnowledgeGraph(), store=MemoryStore(path=tmp_path / "ms"))
res = r.retrieve("anything at all", k=5)
assert res.chunks == []
assert res.metadata["paths"] == []
|