introduce openai_compat with the completions (not chat-completions) API

This keeps the prompt encoding layer in our control (see
`chat_completion_request_to_prompt()` method)

llama_stack/providers/utils/inference/augment_messages.py

@@ -3,6 +3,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+from llama_models.llama3.api.chat_format import ChatFormat
 from termcolor import cprint
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.inference import *  # noqa: F403
@@ -19,6 +20,14 @@ from llama_models.sku_list import resolve_model
 from llama_stack.providers.utils.inference import supported_inference_models
 
 
+def chat_completion_request_to_prompt(
+    request: ChatCompletionRequest, formatter: ChatFormat
+) -> str:
+    messages = augment_messages_for_tools(request)
+    model_input = formatter.encode_dialog_prompt(messages)
+    return formatter.tokenizer.decode(model_input.tokens)
+
+
 def augment_messages_for_tools(request: ChatCompletionRequest) -> List[Message]:
     """Reads chat completion request and augments the messages to handle tools.
     For eg. for llama_3_1, add system message with the appropriate tools or
@@ -48,7 +57,6 @@ def augment_messages_for_tools(request: ChatCompletionRequest) -> List[Message]:
 def augment_messages_for_tools_llama_3_1(
     request: ChatCompletionRequest,
 ) -> List[Message]:
-
     assert request.tool_choice == ToolChoice.auto, "Only `ToolChoice.auto` supported"
 
     existing_messages = request.messages

llama_stack/providers/utils/inference/openai_compat.py

@@ -0,0 +1,187 @@
+# 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 AsyncGenerator, Optional
+
+from llama_models.llama3.api.chat_format import ChatFormat
+
+from llama_models.llama3.api.datatypes import StopReason
+
+from llama_stack.apis.inference import *  # noqa: F403
+
+from pydantic import BaseModel
+
+
+class OpenAICompatCompletionChoiceDelta(BaseModel):
+    content: str
+
+
+class OpenAICompatCompletionChoice(BaseModel):
+    finish_reason: Optional[str] = None
+    text: Optional[str] = None
+    delta: Optional[OpenAICompatCompletionChoiceDelta] = None
+
+
+class OpenAICompatCompletionResponse(BaseModel):
+    choices: List[OpenAICompatCompletionChoice]
+
+
+def get_sampling_options(request: ChatCompletionRequest) -> dict:
+    options = {}
+    if params := request.sampling_params:
+        for attr in {"temperature", "top_p", "top_k", "max_tokens"}:
+            if getattr(params, attr):
+                options[attr] = getattr(params, attr)
+
+        if params.repetition_penalty is not None and params.repetition_penalty != 1.0:
+            options["repeat_penalty"] = params.repetition_penalty
+
+    return options
+
+
+def text_from_choice(choice) -> str:
+    if hasattr(choice, "delta") and choice.delta:
+        return choice.delta.content
+
+    return choice.text
+
+
+def process_chat_completion_response(
+    request: ChatCompletionRequest,
+    response: OpenAICompatCompletionResponse,
+    formatter: ChatFormat,
+) -> ChatCompletionResponse:
+    choice = response.choices[0]
+
+    stop_reason = None
+    if reason := choice.finish_reason:
+        if reason in ["stop", "eos"]:
+            stop_reason = StopReason.end_of_turn
+        elif reason == "length":
+            stop_reason = StopReason.out_of_tokens
+
+    if stop_reason is None:
+        stop_reason = StopReason.out_of_tokens
+
+    completion_message = formatter.decode_assistant_message_from_content(
+        text_from_choice(choice), stop_reason
+    )
+    return ChatCompletionResponse(
+        completion_message=completion_message,
+        logprobs=None,
+    )
+
+
+async def process_chat_completion_stream_response(
+    request: ChatCompletionRequest,
+    stream: AsyncGenerator[OpenAICompatCompletionResponse, None],
+    formatter: ChatFormat,
+) -> AsyncGenerator:
+    yield ChatCompletionResponseStreamChunk(
+        event=ChatCompletionResponseEvent(
+            event_type=ChatCompletionResponseEventType.start,
+            delta="",
+        )
+    )
+
+    buffer = ""
+    ipython = False
+    stop_reason = None
+
+    async for chunk in stream:
+        choice = chunk.choices[0]
+        finish_reason = choice.finish_reason
+
+        if finish_reason:
+            if stop_reason is None and finish_reason in ["stop", "eos"]:
+                stop_reason = StopReason.end_of_turn
+            elif stop_reason is None and finish_reason == "length":
+                stop_reason = StopReason.out_of_tokens
+            break
+
+        text = text_from_choice(choice)
+        # check if its a tool call ( aka starts with <|python_tag|> )
+        if not ipython and text.startswith("<|python_tag|>"):
+            ipython = True
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=ToolCallDelta(
+                        content="",
+                        parse_status=ToolCallParseStatus.started,
+                    ),
+                )
+            )
+            buffer += text
+            continue
+
+        if ipython:
+            if text == "<|eot_id|>":
+                stop_reason = StopReason.end_of_turn
+                text = ""
+                continue
+            elif text == "<|eom_id|>":
+                stop_reason = StopReason.end_of_message
+                text = ""
+                continue
+
+            buffer += text
+            delta = ToolCallDelta(
+                content=text,
+                parse_status=ToolCallParseStatus.in_progress,
+            )
+
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=delta,
+                    stop_reason=stop_reason,
+                )
+            )
+        else:
+            buffer += text
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=text,
+                    stop_reason=stop_reason,
+                )
+            )
+
+    # parse tool calls and report errors
+    message = formatter.decode_assistant_message_from_content(buffer, stop_reason)
+    parsed_tool_calls = len(message.tool_calls) > 0
+    if ipython and not parsed_tool_calls:
+        yield ChatCompletionResponseStreamChunk(
+            event=ChatCompletionResponseEvent(
+                event_type=ChatCompletionResponseEventType.progress,
+                delta=ToolCallDelta(
+                    content="",
+                    parse_status=ToolCallParseStatus.failure,
+                ),
+                stop_reason=stop_reason,
+            )
+        )
+
+    for tool_call in message.tool_calls:
+        yield ChatCompletionResponseStreamChunk(
+            event=ChatCompletionResponseEvent(
+                event_type=ChatCompletionResponseEventType.progress,
+                delta=ToolCallDelta(
+                    content=tool_call,
+                    parse_status=ToolCallParseStatus.success,
+                ),
+                stop_reason=stop_reason,
+            )
+        )
+
+    yield ChatCompletionResponseStreamChunk(
+        event=ChatCompletionResponseEvent(
+            event_type=ChatCompletionResponseEventType.complete,
+            delta="",
+            stop_reason=stop_reason,
+        )
+    )