From d6df94904952dcca671dd2b95ab580c3afee5850 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Tue, 22 Oct 2024 15:57:47 -0700
Subject: [PATCH] naming fix

---
 llama_stack/distribution/routers/routing_tables.py            | 2 +-
 .../providers/impls/meta_reference/datasetio/__init__.py      | 4 ++--
 .../providers/impls/meta_reference/datasetio/datasetio.py     | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/llama_stack/distribution/routers/routing_tables.py b/llama_stack/distribution/routers/routing_tables.py
index 66f5d24e8..db0946d81 100644
--- a/llama_stack/distribution/routers/routing_tables.py
+++ b/llama_stack/distribution/routers/routing_tables.py
@@ -217,7 +217,7 @@ class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
 
     async def get_dataset(
         self, dataset_identifier: str
-    ) -> Optional[ModelDefWithProvider]:
+    ) -> Optional[DatasetDefWithProvider]:
         return self.get_object_by_identifier(identifier)
 
     async def register_dataset(self, dataset_def: DatasetDefWithProvider) -> None:
diff --git a/llama_stack/providers/impls/meta_reference/datasetio/__init__.py b/llama_stack/providers/impls/meta_reference/datasetio/__init__.py
index 6213e575e..9a65f5c3e 100644
--- a/llama_stack/providers/impls/meta_reference/datasetio/__init__.py
+++ b/llama_stack/providers/impls/meta_reference/datasetio/__init__.py
@@ -11,8 +11,8 @@ async def get_provider_impl(
     config: MetaReferenceDatasetIOConfig,
     _deps,
 ):
-    from .datasetio import MetaReferenceDatasetioImpl
+    from .datasetio import MetaReferenceDatasetIOImpl
 
-    impl = MetaReferenceDatasetioImpl(config)
+    impl = MetaReferenceDatasetIOImpl(config)
     await impl.initialize()
     return impl
diff --git a/llama_stack/providers/impls/meta_reference/datasetio/datasetio.py b/llama_stack/providers/impls/meta_reference/datasetio/datasetio.py
index e3dbbc461..a8e648e46 100644
--- a/llama_stack/providers/impls/meta_reference/datasetio/datasetio.py
+++ b/llama_stack/providers/impls/meta_reference/datasetio/datasetio.py
@@ -90,7 +90,7 @@ class PandasDataframeDataset(BaseDataset):
         self.df = df
 
 
-class MetaReferenceDatasetioImpl(DatasetIO, DatasetsProtocolPrivate):
+class MetaReferenceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
     def __init__(self, config: MetaReferenceDatasetIOConfig) -> None:
         self.config = config
         # local registry for keeping track of datasets within the provider