[Evals API][2/n] datasets / datasetio meta-reference implementation (#288)

* skeleton dataset / datasetio

* dataset datasetio

* config

* address comments

* delete dataset_utils

* address comments

* naming fix

llama_stack/distribution/datatypes.py

@@ -14,11 +14,12 @@ from llama_stack.providers.datatypes import *  # noqa: F403
 from llama_stack.apis.models import *  # noqa: F403
 from llama_stack.apis.shields import *  # noqa: F403
 from llama_stack.apis.memory_banks import *  # noqa: F403
+from llama_stack.apis.datasets import *  # noqa: F403
+from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.inference import Inference
 from llama_stack.apis.memory import Memory
 from llama_stack.apis.safety import Safety
 
-
 LLAMA_STACK_BUILD_CONFIG_VERSION = "2"
 LLAMA_STACK_RUN_CONFIG_VERSION = "2"
 
@@ -30,18 +31,21 @@ RoutableObject = Union[
     ModelDef,
     ShieldDef,
     MemoryBankDef,
+    DatasetDef,
 ]
 
 RoutableObjectWithProvider = Union[
     ModelDefWithProvider,
     ShieldDefWithProvider,
     MemoryBankDefWithProvider,
+    DatasetDefWithProvider,
 ]
 
 RoutedProtocol = Union[
     Inference,
     Safety,
     Memory,
+    DatasetIO,
 ]
 
 

llama_stack/distribution/distribution.py

@@ -35,6 +35,10 @@ def builtin_automatically_routed_apis() -> List[AutoRoutedApiInfo]:
             routing_table_api=Api.memory_banks,
             router_api=Api.memory,
         ),
+        AutoRoutedApiInfo(
+            routing_table_api=Api.datasets,
+            router_api=Api.datasetio,
+        ),
     ]
 
 

llama_stack/distribution/resolver.py

@@ -12,6 +12,8 @@ from llama_stack.providers.datatypes import *  # noqa: F403
 from llama_stack.distribution.datatypes import *  # noqa: F403
 
 from llama_stack.apis.agents import Agents
+from llama_stack.apis.datasetio import DatasetIO
+from llama_stack.apis.datasets import Datasets
 from llama_stack.apis.inference import Inference
 from llama_stack.apis.inspect import Inspect
 from llama_stack.apis.memory import Memory
@@ -38,6 +40,8 @@ def api_protocol_map() -> Dict[Api, Any]:
         Api.safety: Safety,
         Api.shields: Shields,
         Api.telemetry: Telemetry,
+        Api.datasets: Datasets,
+        Api.datasetio: DatasetIO,
     }
 
 

llama_stack/distribution/routers/__init__.py

@@ -8,6 +8,7 @@ from typing import Any
 
 from llama_stack.distribution.datatypes import *  # noqa: F403
 from .routing_tables import (
+    DatasetsRoutingTable,
     MemoryBanksRoutingTable,
     ModelsRoutingTable,
     ShieldsRoutingTable,
@@ -23,6 +24,7 @@ async def get_routing_table_impl(
         "memory_banks": MemoryBanksRoutingTable,
         "models": ModelsRoutingTable,
         "shields": ShieldsRoutingTable,
+        "datasets": DatasetsRoutingTable,
     }
     if api.value not in api_to_tables:
         raise ValueError(f"API {api.value} not found in router map")
@@ -33,12 +35,13 @@ async def get_routing_table_impl(
 
 
 async def get_auto_router_impl(api: Api, routing_table: RoutingTable, _deps) -> Any:
-    from .routers import InferenceRouter, MemoryRouter, SafetyRouter
+    from .routers import DatasetIORouter, InferenceRouter, MemoryRouter, SafetyRouter
 
     api_to_routers = {
         "memory": MemoryRouter,
         "inference": InferenceRouter,
         "safety": SafetyRouter,
+        "datasetio": DatasetIORouter,
     }
     if api.value not in api_to_routers:
         raise ValueError(f"API {api.value} not found in router map")

llama_stack/distribution/routers/routers.py

@@ -6,11 +6,13 @@
 
 from typing import Any, AsyncGenerator, Dict, List
 
+from llama_stack.apis.datasetio.datasetio import DatasetIO
 from llama_stack.distribution.datatypes import RoutingTable
 
 from llama_stack.apis.memory import *  # noqa: F403
 from llama_stack.apis.inference import *  # noqa: F403
 from llama_stack.apis.safety import *  # noqa: F403
+from llama_stack.apis.datasetio import *  # noqa: F403
 
 
 class MemoryRouter(Memory):
@@ -160,3 +162,33 @@ class SafetyRouter(Safety):
             messages=messages,
             params=params,
         )
+
+
+class DatasetIORouter(DatasetIO):
+    def __init__(
+        self,
+        routing_table: RoutingTable,
+    ) -> None:
+        self.routing_table = routing_table
+
+    async def initialize(self) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        pass
+
+    async def get_rows_paginated(
+        self,
+        dataset_id: str,
+        rows_in_page: int,
+        page_token: Optional[str] = None,
+        filter_condition: Optional[str] = None,
+    ) -> PaginatedRowsResult:
+        return await self.routing_table.get_provider_impl(
+            dataset_id
+        ).get_rows_paginated(
+            dataset_id=dataset_id,
+            rows_in_page=rows_in_page,
+            page_token=page_token,
+            filter_condition=filter_condition,
+        )

llama_stack/distribution/routers/routing_tables.py

@@ -11,6 +11,7 @@ from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.models import *  # noqa: F403
 from llama_stack.apis.shields import *  # noqa: F403
 from llama_stack.apis.memory_banks import *  # noqa: F403
+from llama_stack.apis.datasets import *  # noqa: F403
 
 from llama_stack.distribution.datatypes import *  # noqa: F403
 
@@ -27,6 +28,10 @@ async def register_object_with_provider(obj: RoutableObject, p: Any) -> None:
         await p.register_shield(obj)
     elif api == Api.memory:
         await p.register_memory_bank(obj)
+    elif api == Api.datasetio:
+        await p.register_dataset(obj)
+    else:
+        raise ValueError(f"Unknown API {api} for registering object with provider")
 
 
 Registry = Dict[str, List[RoutableObjectWithProvider]]
@@ -80,6 +85,16 @@ class CommonRoutingTableImpl(RoutingTable):
 
                 add_objects(memory_banks)
 
+            elif api == Api.datasetio:
+                p.dataset_store = self
+                datasets = await p.list_datasets()
+
+                # do in-memory updates due to pesky Annotated unions
+                for d in datasets:
+                    d.provider_id = pid
+
+                add_objects(datasets)
+
     async def shutdown(self) -> None:
         for p in self.impls_by_provider_id.values():
             await p.shutdown()
@@ -137,6 +152,7 @@ class CommonRoutingTableImpl(RoutingTable):
             raise ValueError(f"Provider `{obj.provider_id}` not found")
 
         p = self.impls_by_provider_id[obj.provider_id]
+
         await register_object_with_provider(obj, p)
 
         if obj.identifier not in self.registry:
@@ -190,3 +206,19 @@ class MemoryBanksRoutingTable(CommonRoutingTableImpl, MemoryBanks):
         self, memory_bank: MemoryBankDefWithProvider
     ) -> None:
         await self.register_object(memory_bank)
+
+
+class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
+    async def list_datasets(self) -> List[DatasetDefWithProvider]:
+        objects = []
+        for objs in self.registry.values():
+            objects.extend(objs)
+        return objects
+
+    async def get_dataset(
+        self, dataset_identifier: str
+    ) -> Optional[DatasetDefWithProvider]:
+        return self.get_object_by_identifier(identifier)
+
+    async def register_dataset(self, dataset_def: DatasetDefWithProvider) -> None:
+        await self.register_object(dataset_def)