wehub-resource-sync
131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
import asyncio
|
|
import types
|
|
import typing as t
|
|
|
|
import pytest
|
|
|
|
from ragas.testset.graph import KnowledgeGraph, Node, NodeType
|
|
from ragas.testset.transforms.base import BaseGraphTransformation
|
|
from ragas.testset.transforms.engine import Parallel, apply_transforms, get_desc
|
|
|
|
|
|
class DummyTransformation(BaseGraphTransformation):
|
|
def __init__(self, name="Dummy"):
|
|
self.name = name
|
|
|
|
def generate_execution_plan(self, kg):
|
|
return [self.double(node) for node in kg.nodes]
|
|
|
|
async def transform(
|
|
self, kg: KnowledgeGraph
|
|
) -> t.List[t.Tuple[Node, t.Tuple[str, t.Any]]]:
|
|
filtered = self.filter(kg)
|
|
nodes = sorted(
|
|
filtered.nodes, key=lambda n: n.get_property("page_content") or ""
|
|
)
|
|
return [(node, await self.double(node)) for node in nodes]
|
|
|
|
async def double(self, node):
|
|
# Repeat the text in a single node's 'page_content' property
|
|
content = node.get_property("page_content")
|
|
if content is not None:
|
|
node.properties["page_content"] = content * 2
|
|
return node
|
|
|
|
|
|
@pytest.fixture
|
|
def kg():
|
|
import string
|
|
|
|
kg = KnowledgeGraph()
|
|
for letter in string.ascii_uppercase[:10]:
|
|
node = Node(
|
|
properties={"page_content": letter},
|
|
type=NodeType.DOCUMENT,
|
|
)
|
|
kg.add(node)
|
|
return kg
|
|
|
|
|
|
def test_parallel_stores_transformations():
|
|
t1 = DummyTransformation("A")
|
|
t2 = DummyTransformation("B")
|
|
p = Parallel(t1, t2)
|
|
assert p.transformations == [t1, t2]
|
|
|
|
|
|
def test_parallel_generate_execution_plan_aggregates(kg):
|
|
t1 = DummyTransformation("A")
|
|
t2 = DummyTransformation("B")
|
|
p = Parallel(t1, t2)
|
|
coros = p.generate_execution_plan(kg)
|
|
assert len(coros) == len(kg.nodes) * 2 # Each transformation runs on each node
|
|
assert all(isinstance(c, types.CoroutineType) for c in coros)
|
|
|
|
# Await all coroutines to avoid RuntimeWarning
|
|
async def run_all():
|
|
await asyncio.gather(*coros)
|
|
|
|
asyncio.run(run_all())
|
|
|
|
|
|
def test_parallel_nested(kg):
|
|
t1 = DummyTransformation("A")
|
|
t2 = DummyTransformation("B")
|
|
p_inner = Parallel(t1)
|
|
p_outer = Parallel(p_inner, t2)
|
|
coros = p_outer.generate_execution_plan(kg)
|
|
assert len(coros) == len(kg.nodes) * 2 # Each transformation runs on each node
|
|
assert all(isinstance(c, types.CoroutineType) for c in coros)
|
|
|
|
# Await all coroutines to avoid RuntimeWarning
|
|
async def run_all():
|
|
await asyncio.gather(*coros)
|
|
|
|
asyncio.run(run_all())
|
|
|
|
|
|
def test_get_desc_parallel_and_single():
|
|
t1 = DummyTransformation("A")
|
|
p = Parallel(t1)
|
|
desc_p = get_desc(p)
|
|
desc_t = get_desc(t1)
|
|
assert "Parallel" not in desc_t
|
|
assert "DummyTransformation" in desc_p or "DummyTransformation" in desc_t
|
|
|
|
|
|
def test_apply_transforms_single(kg):
|
|
t1 = DummyTransformation()
|
|
apply_transforms(kg, t1)
|
|
# All nodes' page_content should be doubled
|
|
for node in kg.nodes:
|
|
content = node.get_property("page_content")
|
|
assert content == (content[0] * 2)
|
|
|
|
|
|
def test_apply_transforms_list(kg):
|
|
t1 = DummyTransformation()
|
|
t2 = DummyTransformation()
|
|
apply_transforms(kg, [t1, t2])
|
|
# Each transformation doubles the content, so after two: x -> xxxx
|
|
for node in kg.nodes:
|
|
content = node.get_property("page_content")
|
|
assert content == (content[0] * 2 * 2)
|
|
|
|
|
|
def test_apply_transforms_parallel(kg):
|
|
t1 = DummyTransformation()
|
|
t2 = DummyTransformation()
|
|
p = Parallel(t1, t2)
|
|
apply_transforms(kg, p)
|
|
# Each transformation in parallel doubles the content, but both operate on the same initial state, so after both: x -> xx (not xxxx)
|
|
for node in kg.nodes:
|
|
content = node.get_property("page_content")
|
|
assert content == (content[0] * 2 * 2)
|
|
|
|
|
|
def test_apply_transforms_invalid():
|
|
kg = KnowledgeGraph()
|
|
with pytest.raises(ValueError):
|
|
apply_transforms(kg, 123) # type: ignore
|