This commit is contained in:
Edward Ma 2024-12-02 08:17:22 -08:00
parent d7b159663c
commit efd842d605

View file

@ -11,7 +11,7 @@ from llama_models.datatypes import CoreModelId, SamplingStrategy
from llama_models.llama3.api.chat_format import ChatFormat from llama_models.llama3.api.chat_format import ChatFormat
from llama_models.llama3.api.datatypes import Message, ImageMedia from llama_models.llama3.api.datatypes import Message
from llama_models.llama3.api.tokenizer import Tokenizer from llama_models.llama3.api.tokenizer import Tokenizer
from openai import OpenAI from openai import OpenAI
@ -27,9 +27,7 @@ from llama_stack.providers.utils.inference.openai_compat import (
process_chat_completion_stream_response, process_chat_completion_stream_response,
) )
from llama_stack.providers.utils.inference.prompt_adapter import ( from llama_stack.providers.utils.inference.prompt_adapter import convert_message_to_dict
convert_message_to_dict,
)
from .config import SambaNovaImplConfig from .config import SambaNovaImplConfig
@ -153,10 +151,14 @@ class SambaNovaInferenceAdapter(ModelRegistryHelper, Inference):
) -> EmbeddingsResponse: ) -> EmbeddingsResponse:
raise NotImplementedError() raise NotImplementedError()
async def convert_chat_completion_request(self, request: ChatCompletionRequest) -> dict: async def convert_chat_completion_request(
self, request: ChatCompletionRequest
) -> dict:
compatible_request = self.convert_sampling_params(request.sampling_params) compatible_request = self.convert_sampling_params(request.sampling_params)
compatible_request["model"] = request.model compatible_request["model"] = request.model
compatible_request["messages"] = await self.convert_to_sambanova_message(request.messages) compatible_request["messages"] = await self.convert_to_sambanova_message(
request.messages
)
compatible_request["stream"] = request.stream compatible_request["stream"] = request.stream
compatible_request["logprobs"] = False compatible_request["logprobs"] = False
compatible_request["extra_headers"] = { compatible_request["extra_headers"] = {
@ -164,7 +166,9 @@ class SambaNovaInferenceAdapter(ModelRegistryHelper, Inference):
} }
return compatible_request return compatible_request
def convert_sampling_params(self, sampling_params: SamplingParams, legacy: bool = False) -> dict: def convert_sampling_params(
self, sampling_params: SamplingParams, legacy: bool = False
) -> dict:
params = {} params = {}
if sampling_params: if sampling_params:
@ -197,14 +201,16 @@ class SambaNovaInferenceAdapter(ModelRegistryHelper, Inference):
content["role"] = "assistant" content["role"] = "assistant"
tools = [] tools = []
for tool_call in message.tool_calls: for tool_call in message.tool_calls:
tools.append({ tools.append(
"id": tool_call.call_id, {
"function": { "id": tool_call.call_id,
"name": tool_call.name, "function": {
"arguments": json.dumps(tool_call.arguments), "name": tool_call.name,
}, "arguments": json.dumps(tool_call.arguments),
"type": "function", },
}) "type": "function",
}
)
content["tool_calls"] = tools content["tool_calls"] = tools
elif isinstance(message, ToolResponseMessage): elif isinstance(message, ToolResponseMessage):
content["role"] = "tool" content["role"] = "tool"
@ -215,5 +221,3 @@ class SambaNovaInferenceAdapter(ModelRegistryHelper, Inference):
conversation.append(content) conversation.append(content)
return conversation return conversation