skeleton dataset / datasetio

This commit is contained in:
Xi Yan 2024-10-22 11:22:39 -07:00
parent 668a495aba
commit e8de70fdbe
12 changed files with 233 additions and 2 deletions

View file

@ -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")