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.datatypes import Message, ImageMedia
from llama_models.llama3.api.datatypes import Message
from llama_models.llama3.api.tokenizer import Tokenizer
from openai import OpenAI
@ -27,9 +27,7 @@ from llama_stack.providers.utils.inference.openai_compat import (
process_chat_completion_stream_response,
)
from llama_stack.providers.utils.inference.prompt_adapter import (
convert_message_to_dict,
)
from llama_stack.providers.utils.inference.prompt_adapter import convert_message_to_dict
from .config import SambaNovaImplConfig
@ -153,10 +151,14 @@ class SambaNovaInferenceAdapter(ModelRegistryHelper, Inference):
) -> EmbeddingsResponse:
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["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["logprobs"] = False
compatible_request["extra_headers"] = {
@ -164,7 +166,9 @@ class SambaNovaInferenceAdapter(ModelRegistryHelper, Inference):
}
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 = {}
if sampling_params:
@ -197,14 +201,16 @@ class SambaNovaInferenceAdapter(ModelRegistryHelper, Inference):
content["role"] = "assistant"
tools = []
for tool_call in message.tool_calls:
tools.append({
tools.append(
{
"id": tool_call.call_id,
"function": {
"name": tool_call.name,
"arguments": json.dumps(tool_call.arguments),
},
"type": "function",
})
}
)
content["tool_calls"] = tools
elif isinstance(message, ToolResponseMessage):
content["role"] = "tool"
@ -215,5 +221,3 @@ class SambaNovaInferenceAdapter(ModelRegistryHelper, Inference):
conversation.append(content)
return conversation