From 59b4a4bbe5c4b9ee908c44ba5842ff0dda0d7bb0 Mon Sep 17 00:00:00 2001 From: Yuan Tang Date: Thu, 10 Oct 2024 21:03:14 -0400 Subject: [PATCH] Fix incorrect completion() signature for Databricks provider Signed-off-by: Yuan Tang --- .../adapters/inference/databricks/databricks.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/llama_stack/providers/adapters/inference/databricks/databricks.py b/llama_stack/providers/adapters/inference/databricks/databricks.py index 2d7427253..7e8263dbf 100644 --- a/llama_stack/providers/adapters/inference/databricks/databricks.py +++ b/llama_stack/providers/adapters/inference/databricks/databricks.py @@ -48,7 +48,14 @@ class DatabricksInferenceAdapter(ModelRegistryHelper, Inference): async def shutdown(self) -> None: pass - def completion(self, request: CompletionRequest) -> AsyncGenerator: + def completion( + self, + model: str, + content: InterleavedTextMedia, + sampling_params: Optional[SamplingParams] = SamplingParams(), + stream: Optional[bool] = False, + logprobs: Optional[LogProbConfig] = None, + ) -> AsyncGenerator: raise NotImplementedError() def chat_completion(