Fix conversion to RawMessage everywhere

llama_stack/providers/inline/inference/meta_reference/generation.py

@@ -25,7 +25,6 @@ from fairscale.nn.model_parallel.initialize import (
 )
 from llama_models.llama3.api.args import ModelArgs
 from llama_models.llama3.api.chat_format import ChatFormat, LLMInput
-from llama_models.llama3.api.datatypes import RawContent, RawMessage
 from llama_models.llama3.api.tokenizer import Tokenizer
 from llama_models.llama3.reference_impl.model import Transformer
 from llama_models.llama3.reference_impl.multimodal.model import (
@@ -39,6 +38,10 @@ from llama_stack.apis.inference import *  # noqa: F403
 from lmformatenforcer import JsonSchemaParser, TokenEnforcer, TokenEnforcerTokenizerData
 
 from llama_stack.distribution.utils.model_utils import model_local_dir
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    ChatCompletionRequestWithRawContent,
+    CompletionRequestWithRawContent,
+)
 
 from .config import (
     Fp8QuantizationConfig,
@@ -50,14 +53,6 @@ from .config import (
 log = logging.getLogger(__name__)
 
 
-class ChatCompletionRequestWithRawContent(ChatCompletionRequest):
-    messages: List[RawMessage]
-
-
-class CompletionRequestWithRawContent(CompletionRequest):
-    content: RawContent
-
-
 def model_checkpoint_dir(model) -> str:
     checkpoint_dir = Path(model_local_dir(model.descriptor()))
 

llama_stack/providers/inline/inference/meta_reference/inference.py

@@ -12,7 +12,6 @@ from typing import AsyncGenerator, List, Optional, Union
 from llama_models.datatypes import Model
 
 from llama_models.llama3.api.datatypes import (
-    RawMessage,
     SamplingParams,
     StopReason,
     ToolDefinition,
@@ -53,14 +52,10 @@ from llama_stack.providers.utils.inference.model_registry import (
 from llama_stack.providers.utils.inference.prompt_adapter import (
     augment_content_with_response_format_prompt,
     chat_completion_request_to_messages,
-    interleaved_content_convert_to_raw,
+    convert_request_to_raw,
 )
 from .config import MetaReferenceInferenceConfig
-from .generation import (
-    ChatCompletionRequestWithRawContent,
-    CompletionRequestWithRawContent,
-    Llama,
-)
+from .generation import Llama
 from .model_parallel import LlamaModelParallelGenerator
 
 log = logging.getLogger(__name__)
@@ -450,20 +445,3 @@ class MetaReferenceInferenceImpl(
         else:
             for x in impl():
                 yield x
-
-
-async def convert_request_to_raw(
-    request: Union[ChatCompletionRequest, CompletionRequest],
-) -> Union[ChatCompletionRequestWithRawContent, CompletionRequestWithRawContent]:
-    if isinstance(request, ChatCompletionRequest):
-        messages = []
-        for m in request.messages:
-            content = await interleaved_content_convert_to_raw(m.content)
-            d = m.model_dump()
-            d["content"] = content
-            messages.append(RawMessage(**d))
-        request.messages = messages
-    else:
-        request.content = await interleaved_content_convert_to_raw(request.content)
-
-    return request