test_scoring

llama_stack/distribution/routers/__init__.py

@@ -11,6 +11,7 @@ from .routing_tables import (
     DatasetsRoutingTable,
     MemoryBanksRoutingTable,
     ModelsRoutingTable,
+    ScoringFunctionsRoutingTable,
     ShieldsRoutingTable,
 )
 
@@ -25,7 +26,9 @@ async def get_routing_table_impl(
         "models": ModelsRoutingTable,
         "shields": ShieldsRoutingTable,
         "datasets": DatasetsRoutingTable,
+        "scoring_functions": ScoringFunctionsRoutingTable,
     }
+
     if api.value not in api_to_tables:
         raise ValueError(f"API {api.value} not found in router map")
 
@@ -35,17 +38,30 @@ async def get_routing_table_impl(
 
 
 async def get_auto_router_impl(api: Api, routing_table: RoutingTable, _deps) -> Any:
-    from .routers import DatasetIORouter, InferenceRouter, MemoryRouter, SafetyRouter
+    from .routers import (
+        DatasetIORouter,
+        InferenceRouter,
+        MemoryRouter,
+        SafetyRouter,
+        ScoringRouter,
+    )
 
     api_to_routers = {
         "memory": MemoryRouter,
         "inference": InferenceRouter,
         "safety": SafetyRouter,
         "datasetio": DatasetIORouter,
+        "scoring": ScoringRouter,
     }
     if api.value not in api_to_routers:
         raise ValueError(f"API {api.value} not found in router map")
 
+    # api_with_deps = {"scoring"}
+    # if api.value in api_with_deps:
+    #     impl = api_to_routers[api.value](routing_table, _deps)
+    # else:
+    #     impl = api_to_routers[api.value](routing_table)
+
     impl = api_to_routers[api.value](routing_table)
     await impl.initialize()
     return impl

llama_stack/distribution/routers/routers.py

@@ -13,6 +13,7 @@ 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
+from llama_stack.apis.scoring import *  # noqa: F403
 
 
 class MemoryRouter(Memory):
@@ -192,3 +193,28 @@ class DatasetIORouter(DatasetIO):
             page_token=page_token,
             filter_condition=filter_condition,
         )
+
+
+class ScoringRouter(Scoring):
+    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 score_batch(
+        self, dataset_id: str, scoring_functions: List[str]
+    ) -> ScoreBatchResponse:
+        # TODO
+        pass
+
+    async def score(
+        self, input_rows: List[Dict[str, Any]], scoring_functions: List[str]
+    ) -> ScoreResponse:
+        pass

llama_stack/distribution/routers/routing_tables.py

@@ -218,7 +218,25 @@ class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
     async def get_dataset(
         self, dataset_identifier: str
     ) -> Optional[DatasetDefWithProvider]:
-        return self.get_object_by_identifier(identifier)
+        return self.get_object_by_identifier(dataset_identifier)
 
     async def register_dataset(self, dataset_def: DatasetDefWithProvider) -> None:
         await self.register_object(dataset_def)
+
+
+class ScoringFunctionsRoutingTable(CommonRoutingTableImpl, Scoring):
+    async def list_scoring_functions(self) -> List[ScoringFunctionDefWithProvider]:
+        objects = []
+        for objs in self.registry.values():
+            objects.extend(objs)
+        return objects
+
+    async def get_scoring_function(
+        self, name: str
+    ) -> Optional[ScoringFunctionDefWithProvider]:
+        return self.get_object_by_identifier(name)
+
+    async def register_scoring_function(
+        self, function_def: ScoringFunctionDefWithProvider
+    ) -> None:
+        await self.register_object(function_def)