diff --git a/src/modelplane/runways/responder.py b/src/modelplane/runways/responder.py index d5f39f8..adbdcdc 100644 --- a/src/modelplane/runways/responder.py +++ b/src/modelplane/runways/responder.py @@ -5,6 +5,7 @@ import mlflow from modelgauge.pipeline_runner import build_runner +from modelgauge.sut import SUTOptions from modelgauge.sut_factory import SUT_FACTORY from modelplane.runways.data import ( @@ -23,6 +24,14 @@ setup_sut_credentials, ) +# Taken from modelgauge.tests.security.BaseSecurityTest +# TODO: Figure out a way to expose the options in the CLI. +DEFAULT_SUT_OPTIONS = SUTOptions( + max_tokens=3000, + max_total_output_tokens=10000, # For reasoning SUTs. + temperature=0.01, +) + def respond( sut_id: str, @@ -34,6 +43,7 @@ def respond( num_workers: int = 1, prompt_uid_col=None, prompt_text_col=None, + sut_options: SUTOptions = DEFAULT_SUT_OPTIONS, ) -> RunArtifacts: secrets = setup_sut_credentials(sut_id) sut = SUT_FACTORY.make_instance(uid=sut_id, secrets=secrets) @@ -60,6 +70,7 @@ def respond( suts={sut.uid: sut}, prompt_uid_col=prompt_uid_col, prompt_text_col=prompt_text_col, + sut_options=sut_options, ) pipeline_runner.run(