some fixes and refactoring

This commit is contained in:
Vladislav 2024-12-03 18:13:51 +01:00
parent a4a7a90e42
commit 4f050be67e
2 changed files with 3 additions and 2 deletions

View file

@ -15,4 +15,4 @@ from llama_stack.providers.utils.kvstore.config import (
class LocalFSDatasetIOConfig(BaseModel): class LocalFSDatasetIOConfig(BaseModel):
kvstore: KVStoreConfig = SqliteKVStoreConfig( kvstore: KVStoreConfig = SqliteKVStoreConfig(
db_path=(RUNTIME_BASE_DIR / "localfs_datasetio.db").as_posix() db_path=(RUNTIME_BASE_DIR / "localfs_datasetio.db").as_posix()
) # Uses SQLite config specific to HF storage ) # Uses SQLite config specific to localfs storage

View file

@ -89,6 +89,7 @@ class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
self.config = config self.config = config
# local registry for keeping track of datasets within the provider # local registry for keeping track of datasets within the provider
self.dataset_infos = {} self.dataset_infos = {}
self.kvstore = None
async def initialize(self) -> None: async def initialize(self) -> None:
self.kvstore = await kvstore_impl(self.config.kvstore) self.kvstore = await kvstore_impl(self.config.kvstore)
@ -113,11 +114,11 @@ class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
) -> None: ) -> None:
# Store in kvstore # Store in kvstore
key = f"{DATASETS_PREFIX}{dataset.identifier}" key = f"{DATASETS_PREFIX}{dataset.identifier}"
dataset_impl = PandasDataframeDataset(dataset)
await self.kvstore.set( await self.kvstore.set(
key=key, key=key,
value=dataset.json(), value=dataset.json(),
) )
dataset_impl = PandasDataframeDataset(dataset)
self.dataset_infos[key] = DatasetInfo( self.dataset_infos[key] = DatasetInfo(
dataset_def=dataset, dataset_def=dataset,
dataset_impl=dataset_impl, dataset_impl=dataset_impl,