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