persist registered objects with distribution (#354)

* persist registered objects with distribution

* linter fixes

* comment

* use annotate and field discriminator

* workign tests

* donot use global state

* precommit failures fixed

* add back Any

* fix imports

* remove unnecessary changes in ollama

* precommit failures fixed

* make kvstore configurable for dist and rename registry

* add comment about registry list return

* fix linter errors

* use registry to hydrate

* remove debug print

* linter fixes

* remove kvstore.db

* rename distribution_registry_store

---------

Co-authored-by: Dinesh Yeduguru <dineshyv@fb.com>

llama_stack/apis/datasets/datasets.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Any, Dict, List, Optional, Protocol
+from typing import Any, Dict, List, Literal, Optional, Protocol
 
 from llama_models.llama3.api.datatypes import URL
 
@@ -32,6 +32,7 @@ class DatasetDef(BaseModel):
 
 @json_schema_type
 class DatasetDefWithProvider(DatasetDef):
+    type: Literal["dataset"] = "dataset"
     provider_id: str = Field(
         description="ID of the provider which serves this dataset",
     )

llama_stack/apis/models/models.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+from typing import Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
@@ -25,6 +25,7 @@ class ModelDef(BaseModel):
 
 @json_schema_type
 class ModelDefWithProvider(ModelDef):
+    type: Literal["model"] = "model"
     provider_id: str = Field(
         description="The provider ID for this model",
     )

llama_stack/apis/scoring_functions/scoring_functions.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+from typing import Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
@@ -53,6 +53,7 @@ class ScoringFnDef(BaseModel):
 
 @json_schema_type
 class ScoringFnDefWithProvider(ScoringFnDef):
+    type: Literal["scoring_fn"] = "scoring_fn"
     provider_id: str = Field(
         description="ID of the provider which serves this dataset",
     )

llama_stack/apis/shields/shields.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 from enum import Enum
-from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+from typing import Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
@@ -34,6 +34,7 @@ class ShieldDef(BaseModel):
 
 @json_schema_type
 class ShieldDefWithProvider(ShieldDef):
+    type: Literal["shield"] = "shield"
     provider_id: str = Field(
         description="The provider ID for this shield type",
     )