From ffecc17ff805575d8d176df3f76330f116dbe820 Mon Sep 17 00:00:00 2001 From: Adrian Lyjak Date: Fri, 16 Jan 2026 18:28:20 -0500 Subject: [PATCH] move them around --- .../{registry => plugins}/__init__.py | 2 +- .../{registry => plugins}/_registry.py | 0 .../src/workflows/server/server.py | 2 +- .../tests/runtime/test_identity_weak_ref.py | 97 +++++++++++++++++++ .../tests/test_registry.py | 4 +- 5 files changed, 101 insertions(+), 4 deletions(-) rename packages/llama-index-workflows/src/workflows/{registry => plugins}/__init__.py (65%) rename packages/llama-index-workflows/src/workflows/{registry => plugins}/_registry.py (100%) diff --git a/packages/llama-index-workflows/src/workflows/registry/__init__.py b/packages/llama-index-workflows/src/workflows/plugins/__init__.py similarity index 65% rename from packages/llama-index-workflows/src/workflows/registry/__init__.py rename to packages/llama-index-workflows/src/workflows/plugins/__init__.py index b3aed463..dfa225d2 100644 --- a/packages/llama-index-workflows/src/workflows/registry/__init__.py +++ b/packages/llama-index-workflows/src/workflows/plugins/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (c) 2026 LlamaIndex Inc. -from workflows.registry._registry import WorkflowRegistry +from workflows.plugins._registry import WorkflowRegistry __all__ = [ "WorkflowRegistry", diff --git a/packages/llama-index-workflows/src/workflows/registry/_registry.py b/packages/llama-index-workflows/src/workflows/plugins/_registry.py similarity index 100% rename from packages/llama-index-workflows/src/workflows/registry/_registry.py rename to packages/llama-index-workflows/src/workflows/plugins/_registry.py diff --git a/packages/llama-index-workflows/src/workflows/server/server.py b/packages/llama-index-workflows/src/workflows/server/server.py index 581eb028..89399bf0 100644 --- a/packages/llama-index-workflows/src/workflows/server/server.py +++ b/packages/llama-index-workflows/src/workflows/server/server.py @@ -37,6 +37,7 @@ from workflows.events import ( WorkflowIdleEvent, ) from workflows.handler import WorkflowHandler +from workflows.plugins import WorkflowRegistry from workflows.protocol import ( CancelHandlerResponse, HandlerData, @@ -53,7 +54,6 @@ from workflows.protocol.serializable_events import ( EventEnvelopeWithMetadata, EventValidationError, ) -from workflows.registry import WorkflowRegistry from workflows.representation import get_workflow_representation from workflows.server.abstract_workflow_store import ( AbstractWorkflowStore, diff --git a/packages/llama-index-workflows/tests/runtime/test_identity_weak_ref.py b/packages/llama-index-workflows/tests/runtime/test_identity_weak_ref.py index c5546ff3..08f88b90 100644 --- a/packages/llama-index-workflows/tests/runtime/test_identity_weak_ref.py +++ b/packages/llama-index-workflows/tests/runtime/test_identity_weak_ref.py @@ -35,3 +35,100 @@ def test_identity_weak_key_dict_removes_entry_when_object_unreferenced() -> None # Object should be collected and the dict entry removed via callback assert w() is None assert d._d == {} + + +def test_identity_weak_key_dict_keys_returns_all_keys() -> None: + d: IdentityWeakKeyDict[Unhashable, str] = IdentityWeakKeyDict() + + obj1 = Unhashable() + obj2 = Unhashable() + obj3 = Unhashable() + + d[obj1] = "one" + d[obj2] = "two" + d[obj3] = "three" + + keys = d.keys() + assert len(keys) == 3 + assert obj1 in keys + assert obj2 in keys + assert obj3 in keys + + +def test_identity_weak_key_dict_items_returns_all_pairs() -> None: + d: IdentityWeakKeyDict[Unhashable, str] = IdentityWeakKeyDict() + + obj1 = Unhashable() + obj2 = Unhashable() + + d[obj1] = "one" + d[obj2] = "two" + + items = d.items() + assert len(items) == 2 + # Find values by identity + obj1_value = next(v for k, v in items if k is obj1) + obj2_value = next(v for k, v in items if k is obj2) + assert obj1_value == "one" + assert obj2_value == "two" + + +def test_identity_weak_key_dict_keys_empty() -> None: + d: IdentityWeakKeyDict[Unhashable, str] = IdentityWeakKeyDict() + assert d.keys() == [] + + +def test_identity_weak_key_dict_items_empty() -> None: + d: IdentityWeakKeyDict[Unhashable, str] = IdentityWeakKeyDict() + assert d.items() == [] + + +def test_identity_weak_key_dict_keys_excludes_collected_objects() -> None: + d: IdentityWeakKeyDict[Unhashable, str] = IdentityWeakKeyDict() + + obj1 = Unhashable() + obj2 = Unhashable() + + d[obj1] = "one" + d[obj2] = "two" + + # Keep weak ref to verify collection + w = weakref.ref(obj2) + + # Drop strong reference to obj2 + del obj2 + gc.collect() + + # obj2 should be collected + assert w() is None + + # keys() should only return obj1 + keys = d.keys() + assert len(keys) == 1 + assert obj1 in keys + + +def test_identity_weak_key_dict_items_excludes_collected_objects() -> None: + d: IdentityWeakKeyDict[Unhashable, str] = IdentityWeakKeyDict() + + obj1 = Unhashable() + obj2 = Unhashable() + + d[obj1] = "one" + d[obj2] = "two" + + # Keep weak ref to verify collection + w = weakref.ref(obj2) + + # Drop strong reference to obj2 + del obj2 + gc.collect() + + # obj2 should be collected + assert w() is None + + # items() should only return obj1's pair + items = d.items() + assert len(items) == 1 + assert items[0][0] is obj1 + assert items[0][1] == "one" diff --git a/packages/llama-index-workflows/tests/test_registry.py b/packages/llama-index-workflows/tests/test_registry.py index 68bff80b..8f7d9ce8 100644 --- a/packages/llama-index-workflows/tests/test_registry.py +++ b/packages/llama-index-workflows/tests/test_registry.py @@ -5,8 +5,8 @@ from workflows import Workflow, step from workflows.events import StartEvent, StopEvent -from workflows.registry import WorkflowRegistry -from workflows.registry._registry import _WorkflowRegistry, compute_workflow_id +from workflows.plugins import WorkflowRegistry +from workflows.plugins._registry import _WorkflowRegistry, compute_workflow_id # --- compute_workflow_id tests ---