llama-stack 0.4.3__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

llama_stack/providers/inline/post_training/huggingface/post_training.py

@@ -12,17 +12,19 @@ from llama_stack.providers.inline.post_training.huggingface.config import (
 from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
 from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
 from llama_stack_api import (
-    AlgorithmConfig,
+    CancelTrainingJobRequest,
     Checkpoint,
     DatasetIO,
     Datasets,
-    DPOAlignmentConfig,
+    GetTrainingJobArtifactsRequest,
+    GetTrainingJobStatusRequest,
     JobStatus,
     ListPostTrainingJobsResponse,
     PostTrainingJob,
     PostTrainingJobArtifactsResponse,
     PostTrainingJobStatusResponse,
-    TrainingConfig,
+    PreferenceOptimizeRequest,
+    SupervisedFineTuneRequest,
 )
 
 
@@ -69,13 +71,7 @@ class HuggingFacePostTrainingImpl:
 
     async def supervised_fine_tune(
         self,
-        job_uuid: str,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
-        model: str,
-        checkpoint_dir: str | None = None,
-        algorithm_config: AlgorithmConfig | None = None,
+        request: SupervisedFineTuneRequest,
     ) -> PostTrainingJob:
         async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
             from llama_stack.providers.inline.post_training.huggingface.recipes.finetune_single_device import (
@@ -85,17 +81,17 @@ class HuggingFacePostTrainingImpl:
             on_log_message_cb("Starting HF finetuning")
 
             recipe = HFFinetuningSingleDevice(
-                job_uuid=job_uuid,
+                job_uuid=request.job_uuid,
                 datasetio_api=self.datasetio_api,
                 datasets_api=self.datasets_api,
             )
 
             resources_allocated, checkpoints = await recipe.train(
-                model=model,
-                output_dir=checkpoint_dir,
-                job_uuid=job_uuid,
-                lora_config=algorithm_config,
-                config=training_config,
+                model=request.model,
+                output_dir=request.checkpoint_dir,
+                job_uuid=request.job_uuid,
+                lora_config=request.algorithm_config,
+                config=request.training_config,
                 provider_config=self.config,
             )
 
@@ -108,17 +104,12 @@ class HuggingFacePostTrainingImpl:
             on_status_change_cb(SchedulerJobStatus.completed)
             on_log_message_cb("HF finetuning completed")
 
-        job_uuid = self._scheduler.schedule(_JOB_TYPE_SUPERVISED_FINE_TUNE, job_uuid, handler)
+        job_uuid = self._scheduler.schedule(_JOB_TYPE_SUPERVISED_FINE_TUNE, request.job_uuid, handler)
         return PostTrainingJob(job_uuid=job_uuid)
 
     async def preference_optimize(
         self,
-        job_uuid: str,
-        finetuned_model: str,
-        algorithm_config: DPOAlignmentConfig,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
+        request: PreferenceOptimizeRequest,
     ) -> PostTrainingJob:
         async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
             from llama_stack.providers.inline.post_training.huggingface.recipes.finetune_single_device_dpo import (
@@ -128,17 +119,17 @@ class HuggingFacePostTrainingImpl:
             on_log_message_cb("Starting HF DPO alignment")
 
             recipe = HFDPOAlignmentSingleDevice(
-                job_uuid=job_uuid,
+                job_uuid=request.job_uuid,
                 datasetio_api=self.datasetio_api,
                 datasets_api=self.datasets_api,
             )
 
             resources_allocated, checkpoints = await recipe.train(
-                model=finetuned_model,
-                output_dir=f"{self.config.dpo_output_dir}/{job_uuid}",
-                job_uuid=job_uuid,
-                dpo_config=algorithm_config,
-                config=training_config,
+                model=request.finetuned_model,
+                output_dir=f"{self.config.dpo_output_dir}/{request.job_uuid}",
+                job_uuid=request.job_uuid,
+                dpo_config=request.algorithm_config,
+                config=request.training_config,
                 provider_config=self.config,
             )
 
@@ -153,7 +144,7 @@ class HuggingFacePostTrainingImpl:
             on_status_change_cb(SchedulerJobStatus.completed)
             on_log_message_cb("HF DPO alignment completed")
 
-        job_uuid = self._scheduler.schedule(_JOB_TYPE_DPO_TRAINING, job_uuid, handler)
+        job_uuid = self._scheduler.schedule(_JOB_TYPE_DPO_TRAINING, request.job_uuid, handler)
         return PostTrainingJob(job_uuid=job_uuid)
 
     @staticmethod
@@ -169,8 +160,10 @@ class HuggingFacePostTrainingImpl:
         data = cls._get_artifacts_metadata_by_type(job, TrainingArtifactType.RESOURCES_STATS.value)
         return data[0] if data else None
 
-    async def get_training_job_status(self, job_uuid: str) -> PostTrainingJobStatusResponse | None:
-        job = self._scheduler.get_job(job_uuid)
+    async def get_training_job_status(
+        self, request: GetTrainingJobStatusRequest
+    ) -> PostTrainingJobStatusResponse | None:
+        job = self._scheduler.get_job(request.job_uuid)
 
         match job.status:
             # TODO: Add support for other statuses to API
@@ -186,7 +179,7 @@ class HuggingFacePostTrainingImpl:
                 raise NotImplementedError()
 
         return PostTrainingJobStatusResponse(
-            job_uuid=job_uuid,
+            job_uuid=request.job_uuid,
             status=status,
             scheduled_at=job.scheduled_at,
             started_at=job.started_at,
@@ -195,12 +188,14 @@ class HuggingFacePostTrainingImpl:
             resources_allocated=self._get_resources_allocated(job),
         )
 
-    async def cancel_training_job(self, job_uuid: str) -> None:
-        self._scheduler.cancel(job_uuid)
+    async def cancel_training_job(self, request: CancelTrainingJobRequest) -> None:
+        self._scheduler.cancel(request.job_uuid)
 
-    async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse | None:
-        job = self._scheduler.get_job(job_uuid)
-        return PostTrainingJobArtifactsResponse(job_uuid=job_uuid, checkpoints=self._get_checkpoints(job))
+    async def get_training_job_artifacts(
+        self, request: GetTrainingJobArtifactsRequest
+    ) -> PostTrainingJobArtifactsResponse | None:
+        job = self._scheduler.get_job(request.job_uuid)
+        return PostTrainingJobArtifactsResponse(job_uuid=request.job_uuid, checkpoints=self._get_checkpoints(job))
 
     async def get_training_jobs(self) -> ListPostTrainingJobsResponse:
         return ListPostTrainingJobsResponse(

llama_stack/providers/inline/post_training/huggingface/utils.py

@@ -16,7 +16,7 @@ import torch
 from datasets import Dataset
 from transformers import AutoConfig, AutoModelForCausalLM
 
-from llama_stack_api import Checkpoint, DatasetIO, TrainingConfig
+from llama_stack_api import Checkpoint, DatasetIO, IterRowsRequest, TrainingConfig
 
 if TYPE_CHECKING:
     from transformers import PretrainedConfig
@@ -135,10 +135,7 @@ def setup_torch_device(device_str: str) -> torch.device:
 async def load_rows_from_dataset(datasetio_api: DatasetIO, dataset_id: str) -> list[dict[str, Any]]:
     """Load dataset from llama stack dataset provider"""
     try:
-        all_rows = await datasetio_api.iterrows(
-            dataset_id=dataset_id,
-            limit=-1,
-        )
+        all_rows = await datasetio_api.iterrows(IterRowsRequest(dataset_id=dataset_id, limit=-1))
         if not isinstance(all_rows.data, list):
             raise RuntimeError("Expected dataset data to be a list")
         return all_rows.data

llama_stack/providers/inline/post_training/torchtune/post_training.py

@@ -12,18 +12,20 @@ from llama_stack.providers.inline.post_training.torchtune.config import (
 from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
 from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
 from llama_stack_api import (
-    AlgorithmConfig,
+    CancelTrainingJobRequest,
     Checkpoint,
     DatasetIO,
     Datasets,
-    DPOAlignmentConfig,
+    GetTrainingJobArtifactsRequest,
+    GetTrainingJobStatusRequest,
     JobStatus,
     ListPostTrainingJobsResponse,
     LoraFinetuningConfig,
     PostTrainingJob,
     PostTrainingJobArtifactsResponse,
     PostTrainingJobStatusResponse,
-    TrainingConfig,
+    PreferenceOptimizeRequest,
+    SupervisedFineTuneRequest,
 )
 
 
@@ -69,15 +71,9 @@ class TorchtunePostTrainingImpl:
 
     async def supervised_fine_tune(
         self,
-        job_uuid: str,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
-        model: str,
-        checkpoint_dir: str | None,
-        algorithm_config: AlgorithmConfig | None,
+        request: SupervisedFineTuneRequest,
     ) -> PostTrainingJob:
-        if isinstance(algorithm_config, LoraFinetuningConfig):
+        if isinstance(request.algorithm_config, LoraFinetuningConfig):
 
             async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
                 from llama_stack.providers.inline.post_training.torchtune.recipes.lora_finetuning_single_device import (
@@ -88,13 +84,13 @@ class TorchtunePostTrainingImpl:
 
                 recipe = LoraFinetuningSingleDevice(
                     self.config,
-                    job_uuid,
-                    training_config,
-                    hyperparam_search_config,
-                    logger_config,
-                    model,
-                    checkpoint_dir,
-                    algorithm_config,
+                    request.job_uuid,
+                    request.training_config,
+                    request.hyperparam_search_config,
+                    request.logger_config,
+                    request.model,
+                    request.checkpoint_dir,
+                    request.algorithm_config,
                     self.datasetio_api,
                     self.datasets_api,
                 )
@@ -112,17 +108,12 @@ class TorchtunePostTrainingImpl:
         else:
             raise NotImplementedError()
 
-        job_uuid = self._scheduler.schedule(_JOB_TYPE_SUPERVISED_FINE_TUNE, job_uuid, handler)
+        job_uuid = self._scheduler.schedule(_JOB_TYPE_SUPERVISED_FINE_TUNE, request.job_uuid, handler)
         return PostTrainingJob(job_uuid=job_uuid)
 
     async def preference_optimize(
         self,
-        job_uuid: str,
-        finetuned_model: str,
-        algorithm_config: DPOAlignmentConfig,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
+        request: PreferenceOptimizeRequest,
     ) -> PostTrainingJob:
         raise NotImplementedError()
 
@@ -144,8 +135,10 @@ class TorchtunePostTrainingImpl:
         data = cls._get_artifacts_metadata_by_type(job, TrainingArtifactType.RESOURCES_STATS.value)
         return data[0] if data else None
 
-    async def get_training_job_status(self, job_uuid: str) -> PostTrainingJobStatusResponse | None:
-        job = self._scheduler.get_job(job_uuid)
+    async def get_training_job_status(
+        self, request: GetTrainingJobStatusRequest
+    ) -> PostTrainingJobStatusResponse | None:
+        job = self._scheduler.get_job(request.job_uuid)
 
         match job.status:
             # TODO: Add support for other statuses to API
@@ -161,7 +154,7 @@ class TorchtunePostTrainingImpl:
                 raise NotImplementedError()
 
         return PostTrainingJobStatusResponse(
-            job_uuid=job_uuid,
+            job_uuid=request.job_uuid,
             status=status,
             scheduled_at=job.scheduled_at,
             started_at=job.started_at,
@@ -170,9 +163,11 @@ class TorchtunePostTrainingImpl:
             resources_allocated=self._get_resources_allocated(job),
         )
 
-    async def cancel_training_job(self, job_uuid: str) -> None:
-        self._scheduler.cancel(job_uuid)
+    async def cancel_training_job(self, request: CancelTrainingJobRequest) -> None:
+        self._scheduler.cancel(request.job_uuid)
 
-    async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse | None:
-        job = self._scheduler.get_job(job_uuid)
-        return PostTrainingJobArtifactsResponse(job_uuid=job_uuid, checkpoints=self._get_checkpoints(job))
+    async def get_training_job_artifacts(
+        self, request: GetTrainingJobArtifactsRequest
+    ) -> PostTrainingJobArtifactsResponse | None:
+        job = self._scheduler.get_job(request.job_uuid)
+        return PostTrainingJobArtifactsResponse(job_uuid=request.job_uuid, checkpoints=self._get_checkpoints(job))

llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py

@@ -50,6 +50,7 @@ from llama_stack_api import (
     DataConfig,
     DatasetIO,
     Datasets,
+    IterRowsRequest,
     LoraFinetuningConfig,
     OptimizerConfig,
     PostTrainingMetric,
@@ -334,10 +335,7 @@ class LoraFinetuningSingleDevice:
         batch_size: int,
     ) -> tuple[DistributedSampler, DataLoader]:
         async def fetch_rows(dataset_id: str):
-            return await self.datasetio_api.iterrows(
-                dataset_id=dataset_id,
-                limit=-1,
-            )
+            return await self.datasetio_api.iterrows(IterRowsRequest(dataset_id=dataset_id, limit=-1))
 
         all_rows = await fetch_rows(dataset_id)
         rows = all_rows.data

llama_stack/providers/inline/safety/code_scanner/code_scanner.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 import uuid
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from codeshield.cs import CodeShieldScanResult
@@ -15,9 +15,11 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
 )
 from llama_stack_api import (
+    GetShieldRequest,
     ModerationObject,
     ModerationObjectResults,
-    OpenAIMessageParam,
+    RunModerationRequest,
+    RunShieldRequest,
     RunShieldResponse,
     Safety,
     SafetyViolation,
@@ -51,19 +53,14 @@ class MetaReferenceCodeScannerSafetyImpl(Safety):
                 f"Unsupported Code Scanner ID: {shield.provider_resource_id}. Allowed IDs: {ALLOWED_CODE_SCANNER_MODEL_IDS}"
             )
 
-    async def run_shield(
-        self,
-        shield_id: str,
-        messages: list[OpenAIMessageParam],
-        params: dict[str, Any] = None,
-    ) -> RunShieldResponse:
-        shield = await self.shield_store.get_shield(shield_id)
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        shield = await self.shield_store.get_shield(GetShieldRequest(identifier=request.shield_id))
         if not shield:
-            raise ValueError(f"Shield {shield_id} not found")
+            raise ValueError(f"Shield {request.shield_id} not found")
 
         from codeshield.cs import CodeShield
 
-        text = "\n".join([interleaved_content_as_str(m.content) for m in messages])
+        text = "\n".join([interleaved_content_as_str(m.content) for m in request.messages])
         log.info(f"Running CodeScannerShield on {text[50:]}")
         result = await CodeShield.scan_code(text)
 
@@ -102,11 +99,11 @@ class MetaReferenceCodeScannerSafetyImpl(Safety):
             metadata=metadata,
         )
 
-    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
-        if model is None:
+    async def run_moderation(self, request: RunModerationRequest) -> ModerationObject:
+        if request.model is None:
             raise ValueError("Code scanner moderation requires a model identifier.")
 
-        inputs = input if isinstance(input, list) else [input]
+        inputs = request.input if isinstance(request.input, list) else [request.input]
         results = []
 
         from codeshield.cs import CodeShield
@@ -129,4 +126,4 @@ class MetaReferenceCodeScannerSafetyImpl(Safety):
                 )
             results.append(moderation_result)
 
-        return ModerationObject(id=str(uuid.uuid4()), model=model, results=results)
+        return ModerationObject(id=str(uuid.uuid4()), model=request.model, results=results)

llama_stack/providers/inline/safety/llama_guard/llama_guard.py

@@ -7,7 +7,6 @@
 import re
 import uuid
 from string import Template
-from typing import Any
 
 from llama_stack.core.datatypes import Api
 from llama_stack.log import get_logger
@@ -17,6 +16,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
 )
 from llama_stack_api import (
+    GetShieldRequest,
     ImageContentItem,
     Inference,
     ModerationObject,
@@ -24,6 +24,8 @@ from llama_stack_api import (
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAIMessageParam,
     OpenAIUserMessageParam,
+    RunModerationRequest,
+    RunShieldRequest,
     RunShieldResponse,
     Safety,
     SafetyViolation,
@@ -161,17 +163,12 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
         # The routing table handles the removal from the registry
         pass
 
-    async def run_shield(
-        self,
-        shield_id: str,
-        messages: list[OpenAIMessageParam],
-        params: dict[str, Any] = None,
-    ) -> RunShieldResponse:
-        shield = await self.shield_store.get_shield(shield_id)
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        shield = await self.shield_store.get_shield(GetShieldRequest(identifier=request.shield_id))
         if not shield:
-            raise ValueError(f"Unknown shield {shield_id}")
+            raise ValueError(f"Unknown shield {request.shield_id}")
 
-        messages = messages.copy()
+        messages = request.messages.copy()
         # some shields like llama-guard require the first message to be a user message
         # since this might be a tool call, first role might not be user
         if len(messages) > 0 and messages[0].role != "user":
@@ -200,30 +197,30 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
 
         return await impl.run(messages)
 
-    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
-        if model is None:
+    async def run_moderation(self, request: RunModerationRequest) -> ModerationObject:
+        if request.model is None:
             raise ValueError("Llama Guard moderation requires a model identifier.")
 
-        if isinstance(input, list):
-            messages = input.copy()
+        if isinstance(request.input, list):
+            messages = request.input.copy()
         else:
-            messages = [input]
+            messages = [request.input]
 
         # convert to user messages format with role
         messages = [OpenAIUserMessageParam(content=m) for m in messages]
 
         # Determine safety categories based on the model type
         # For known Llama Guard models, use specific categories
-        if model in LLAMA_GUARD_MODEL_IDS:
+        if request.model in LLAMA_GUARD_MODEL_IDS:
             # Use the mapped model for categories but the original model_id for inference
-            mapped_model = LLAMA_GUARD_MODEL_IDS[model]
+            mapped_model = LLAMA_GUARD_MODEL_IDS[request.model]
             safety_categories = MODEL_TO_SAFETY_CATEGORIES_MAP.get(mapped_model, DEFAULT_LG_V3_SAFETY_CATEGORIES)
         else:
             # For unknown models, use default Llama Guard 3 8B categories
             safety_categories = DEFAULT_LG_V3_SAFETY_CATEGORIES + [CAT_CODE_INTERPRETER_ABUSE]
 
         impl = LlamaGuardShield(
-            model=model,
+            model=request.model,
             inference_api=self.inference_api,
             excluded_categories=self.config.excluded_categories,
             safety_categories=safety_categories,

llama_stack/providers/inline/safety/prompt_guard/prompt_guard.py

@@ -4,17 +4,19 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Any
-
 import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 from llama_stack.core.utils.model_utils import model_local_dir
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    interleaved_content_as_str,
+)
+from llama_stack.providers.utils.safety import ShieldToModerationMixin
 from llama_stack_api import (
-    ModerationObject,
+    GetShieldRequest,
     OpenAIMessageParam,
+    RunShieldRequest,
     RunShieldResponse,
     Safety,
     SafetyViolation,
@@ -31,7 +33,7 @@ log = get_logger(name=__name__, category="safety")
 PROMPT_GUARD_MODEL = "Prompt-Guard-86M"
 
 
-class PromptGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
+class PromptGuardSafetyImpl(ShieldToModerationMixin, Safety, ShieldsProtocolPrivate):
     shield_store: ShieldStore
 
     def __init__(self, config: PromptGuardConfig, _deps) -> None:
@@ -51,20 +53,12 @@ class PromptGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
     async def unregister_shield(self, identifier: str) -> None:
         pass
 
-    async def run_shield(
-        self,
-        shield_id: str,
-        messages: list[OpenAIMessageParam],
-        params: dict[str, Any],
-    ) -> RunShieldResponse:
-        shield = await self.shield_store.get_shield(shield_id)
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        shield = await self.shield_store.get_shield(GetShieldRequest(identifier=request.shield_id))
         if not shield:
-            raise ValueError(f"Unknown shield {shield_id}")
-
-        return await self.shield.run(messages)
+            raise ValueError(f"Unknown shield {request.shield_id}")
 
-    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
-        raise NotImplementedError("run_moderation is not implemented for Prompt Guard")
+        return await self.shield.run(request.messages)
 
 
 class PromptGuardShield:

llama_stack/providers/inline/scoring/basic/scoring.py

@@ -3,16 +3,17 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from typing import Any
 
 from llama_stack_api import (
     DatasetIO,
     Datasets,
+    IterRowsRequest,
+    ScoreBatchRequest,
     ScoreBatchResponse,
+    ScoreRequest,
     ScoreResponse,
     Scoring,
     ScoringFn,
-    ScoringFnParams,
     ScoringFunctionsProtocolPrivate,
     ScoringResult,
 )
@@ -75,19 +76,15 @@ class BasicScoringImpl(
 
     async def score_batch(
         self,
-        dataset_id: str,
-        scoring_functions: dict[str, ScoringFnParams | None] = None,
-        save_results_dataset: bool = False,
+        request: ScoreBatchRequest,
     ) -> ScoreBatchResponse:
-        all_rows = await self.datasetio_api.iterrows(
-            dataset_id=dataset_id,
-            limit=-1,
-        )
-        res = await self.score(
+        all_rows = await self.datasetio_api.iterrows(IterRowsRequest(dataset_id=request.dataset_id, limit=-1))
+        score_request = ScoreRequest(
             input_rows=all_rows.data,
-            scoring_functions=scoring_functions,
+            scoring_functions=request.scoring_functions,
         )
-        if save_results_dataset:
+        res = await self.score(score_request)
+        if request.save_results_dataset:
             # TODO: persist and register dataset on to server for reading
             # self.datasets_api.register_dataset()
             raise NotImplementedError("Save results dataset not implemented yet")
@@ -98,16 +95,15 @@ class BasicScoringImpl(
 
     async def score(
         self,
-        input_rows: list[dict[str, Any]],
-        scoring_functions: dict[str, ScoringFnParams | None] = None,
+        request: ScoreRequest,
     ) -> ScoreResponse:
         res = {}
-        for scoring_fn_id in scoring_functions.keys():
+        for scoring_fn_id in request.scoring_functions.keys():
             if scoring_fn_id not in self.scoring_fn_id_impls:
                 raise ValueError(f"Scoring function {scoring_fn_id} is not supported.")
             scoring_fn = self.scoring_fn_id_impls[scoring_fn_id]
-            scoring_fn_params = scoring_functions.get(scoring_fn_id, None)
-            score_results = await scoring_fn.score(input_rows, scoring_fn_id, scoring_fn_params)
+            scoring_fn_params = request.scoring_functions.get(scoring_fn_id, None)
+            score_results = await scoring_fn.score(request.input_rows, scoring_fn_id, scoring_fn_params)
             agg_results = await scoring_fn.aggregate(score_results, scoring_fn_id, scoring_fn_params)
             res[scoring_fn_id] = ScoringResult(
                 score_rows=score_results,

llama_stack/providers/inline/scoring/braintrust/braintrust.py

@@ -29,11 +29,13 @@ from llama_stack.providers.utils.scoring.aggregation_utils import aggregate_metr
 from llama_stack_api import (
     DatasetIO,
     Datasets,
+    IterRowsRequest,
+    ScoreBatchRequest,
     ScoreBatchResponse,
+    ScoreRequest,
     ScoreResponse,
     Scoring,
     ScoringFn,
-    ScoringFnParams,
     ScoringFunctionsProtocolPrivate,
     ScoringResult,
     ScoringResultRow,
@@ -158,18 +160,17 @@ class BraintrustScoringImpl(
 
     async def score_batch(
         self,
-        dataset_id: str,
-        scoring_functions: dict[str, ScoringFnParams | None],
-        save_results_dataset: bool = False,
+        request: ScoreBatchRequest,
     ) -> ScoreBatchResponse:
         await self.set_api_key()
 
-        all_rows = await self.datasetio_api.iterrows(
-            dataset_id=dataset_id,
-            limit=-1,
+        all_rows = await self.datasetio_api.iterrows(IterRowsRequest(dataset_id=request.dataset_id, limit=-1))
+        score_request = ScoreRequest(
+            input_rows=all_rows.data,
+            scoring_functions=request.scoring_functions,
         )
-        res = await self.score(input_rows=all_rows.data, scoring_functions=scoring_functions)
-        if save_results_dataset:
+        res = await self.score(score_request)
+        if request.save_results_dataset:
             # TODO: persist and register dataset on to server for reading
             # self.datasets_api.register_dataset()
             raise NotImplementedError("Save results dataset not implemented yet")
@@ -198,21 +199,20 @@ class BraintrustScoringImpl(
 
     async def score(
         self,
-        input_rows: list[dict[str, Any]],
-        scoring_functions: dict[str, ScoringFnParams | None],
+        request: ScoreRequest,
     ) -> ScoreResponse:
         await self.set_api_key()
         res = {}
-        for scoring_fn_id in scoring_functions:
+        for scoring_fn_id in request.scoring_functions:
             if scoring_fn_id not in self.supported_fn_defs_registry:
                 raise ValueError(f"Scoring function {scoring_fn_id} is not supported.")
 
-            score_results = [await self.score_row(input_row, scoring_fn_id) for input_row in input_rows]
+            score_results = [await self.score_row(input_row, scoring_fn_id) for input_row in request.input_rows]
             aggregation_functions = self.supported_fn_defs_registry[scoring_fn_id].params.aggregation_functions
 
             # override scoring_fn params if provided
-            if scoring_functions[scoring_fn_id] is not None:
-                override_params = scoring_functions[scoring_fn_id]
+            if request.scoring_functions[scoring_fn_id] is not None:
+                override_params = request.scoring_functions[scoring_fn_id]
                 if override_params.aggregation_functions:
                     aggregation_functions = override_params.aggregation_functions