use annotate and field discriminator

llama_stack/distribution/datatypes.py

@@ -37,12 +37,16 @@ RoutableObject = Union[
     ScoringFnDef,
 ]
 
-RoutableObjectWithProvider = Union[
-    ModelDefWithProvider,
-    ShieldDefWithProvider,
-    MemoryBankDefWithProvider,
-    DatasetDefWithProvider,
-    ScoringFnDefWithProvider,
+
+RoutableObjectWithProvider = Annotated[
+    Union[
+        ModelDefWithProvider,
+        ShieldDefWithProvider,
+        MemoryBankDefWithProvider,
+        DatasetDefWithProvider,
+        ScoringFnDefWithProvider,
+    ],
+    Field(discriminator="type"),
 ]
 
 RoutedProtocol = Union[

llama_stack/distribution/store/registry.py

@@ -2,9 +2,7 @@ import json
 
 from typing import Protocol
 
-from docs.openapi_generator.strong_typing.deserializer import create_deserializer
-
-from docs.openapi_generator.strong_typing.serialization import object_to_json
+import pydantic
 
 from llama_stack.distribution.datatypes import RoutableObjectWithProvider
 
@@ -17,7 +15,6 @@ class Registry(Protocol):
 
 
 KEY_FORMAT = "distributions:registry:{}"
-DESERIALIZER = create_deserializer(RoutableObjectWithProvider)
 
 
 class DiskRegistry(Registry):
@@ -33,20 +30,21 @@ class DiskRegistry(Registry):
         # Parse JSON string into list of objects
         objects_data = json.loads(json_str)
 
-        return [DESERIALIZER.parse(obj_str) for obj_str in objects_data]
+        return [
+            pydantic.parse_obj_as(
+                RoutableObjectWithProvider,
+                obj_str,
+            )
+            for obj_str in objects_data
+        ]
 
     # TODO: make it thread safe using CAS
     async def register(self, obj: RoutableObjectWithProvider) -> None:
-        # Get existing objects for this identifier
         existing_objects = await self.get(obj.identifier)
 
-        # Add new object to list
         existing_objects.append(obj)
 
-        # Convert all objects to JSON strings and store as JSON array
-        objects_json = [
-            object_to_json(existing_object) for existing_object in existing_objects
-        ]
+        objects_json = [obj.model_dump_json() for existing_object in existing_objects]
         await self.kvstore.set(
             KEY_FORMAT.format(obj.identifier), json.dumps(objects_json)
         )

llama_stack/distribution/store/tests/test_registry.py

@@ -0,0 +1,28 @@
+import pytest
+import pytest_asyncio
+from llama_stack.distribution.store import *
+from llama_stack.apis.memory_banks import GraphMemoryBankDef, VectorMemoryBankDef
+from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR
+from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig
+from llama_stack.distribution.datatypes import *  # noqa: F403
+
+
+@pytest.mark.asyncio
+async def test_registry():
+    registry = DiskRegistry(await kvstore_impl(SqliteKVStoreConfig()))
+    bank = VectorMemoryBankDef(
+        identifier="test_bank",
+        embedding_model="all-MiniLM-L6-v2",
+        chunk_size_in_tokens=512,
+        overlap_size_in_tokens=64,
+        provider_id="bar",
+    )
+
+    await registry.register(bank)
+    result_bank = await registry.get("test_bank")
+    # assert result_bank == bank
+    assert result_bank.identifier == bank.identifier
+    assert result_bank.embedding_model == bank.embedding_model
+    assert result_bank.chunk_size_in_tokens == bank.chunk_size_in_tokens
+    assert result_bank.overlap_size_in_tokens == bank.overlap_size_in_tokens
+    assert result_bank.provider_id == bank.provider_id