Fix incorrect completion() signature for Databricks provider (#236)

docs/getting_started.md

@@ -73,7 +73,7 @@ docker run -it -p 5000:5000 -v ~/.llama:/root/.llama --gpus=all llamastack-local
 ```
 
 > [!NOTE]
-> `~/.llama` should be the path containing downloaded weights of Llama models. 
+> `~/.llama` should be the path containing downloaded weights of Llama models.
 
 
 #### Via conda

llama_stack/providers/adapters/inference/databricks/__init__.py

@@ -7,10 +7,11 @@
 from .config import DatabricksImplConfig
 from .databricks import DatabricksInferenceAdapter
 
+
 async def get_adapter_impl(config: DatabricksImplConfig, _deps):
     assert isinstance(
         config, DatabricksImplConfig
     ), f"Unexpected config type: {type(config)}"
     impl = DatabricksInferenceAdapter(config)
     await impl.initialize()
-    return impl
+    return impl

llama_stack/providers/adapters/inference/databricks/config.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Optional
 
 from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
@@ -19,4 +18,4 @@ class DatabricksImplConfig(BaseModel):
     api_token: str = Field(
         default=None,
         description="The Databricks API token",
-    )
+    )

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(

llama_stack/providers/impls/meta_reference/safety/llama_guard.py

@@ -170,7 +170,7 @@ class LlamaGuardShield(ShieldBase):
         for i in range(1, len(messages)):
             if messages[i].role == messages[i - 1].role:
                 raise ValueError(
-                    f"Messages must alternate between user and assistant. Message {i} has the same role as message {i-1}"
+                    f"Messages must alternate between user and assistant. Message {i} has the same role as message {i - 1}"
                 )
         return messages
 

llama_stack/providers/impls/vllm/__init__.py

@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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 .config import VLLMConfig