move them around

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",

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,

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"

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 ---