Rework InterleavedContentMedia datatype so URL downloading is in llama-stack

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

@@ -24,7 +24,8 @@ from fairscale.nn.model_parallel.initialize import (
     model_parallel_is_initialized,
 )
 from llama_models.llama3.api.args import ModelArgs
-from llama_models.llama3.api.chat_format import ChatFormat, ModelInput
+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 (
@@ -38,10 +39,6 @@ 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 (
-    augment_content_with_response_format_prompt,
-    chat_completion_request_to_messages,
-)
 
 from .config import (
     Fp8QuantizationConfig,
@@ -53,6 +50,14 @@ 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()))
 
@@ -206,7 +211,7 @@ class Llama:
     @torch.inference_mode()
     def generate(
         self,
-        model_input: ModelInput,
+        model_input: LLMInput,
         max_gen_len: int,
         temperature: float = 0.6,
         top_p: float = 0.9,
@@ -343,7 +348,7 @@ class Llama:
 
     def completion(
         self,
-        request: CompletionRequest,
+        request: CompletionRequestWithRawContent,
     ) -> Generator:
         sampling_params = request.sampling_params
         max_gen_len = sampling_params.max_tokens
@@ -354,10 +359,7 @@ class Llama:
         ):
             max_gen_len = self.model.params.max_seq_len - 1
 
-        content = augment_content_with_response_format_prompt(
-            request.response_format, request.content
-        )
-        model_input = self.formatter.encode_content(content)
+        model_input = self.formatter.encode_content(request.content)
         yield from self.generate(
             model_input=model_input,
             max_gen_len=max_gen_len,
@@ -374,10 +376,8 @@ class Llama:
 
     def chat_completion(
         self,
-        request: ChatCompletionRequest,
+        request: ChatCompletionRequestWithRawContent,
     ) -> Generator:
-        messages = chat_completion_request_to_messages(request, self.llama_model)
-
         sampling_params = request.sampling_params
         max_gen_len = sampling_params.max_tokens
         if (
@@ -389,7 +389,7 @@ class Llama:
 
         yield from self.generate(
             model_input=self.formatter.encode_dialog_prompt(
-                messages,
+                request.messages,
                 request.tool_prompt_format,
             ),
             max_gen_len=max_gen_len,

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

@@ -7,25 +7,59 @@
 import asyncio
 import logging
 
-from typing import AsyncGenerator, List
+from typing import AsyncGenerator, List, Optional, Union
 
+from llama_models.datatypes import Model
+
+from llama_models.llama3.api.datatypes import (
+    RawMessage,
+    SamplingParams,
+    StopReason,
+    ToolDefinition,
+    ToolPromptFormat,
+)
 from llama_models.sku_list import resolve_model
 
-from llama_models.llama3.api.datatypes import *  # noqa: F403
+from llama_stack.apis.inference import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseEvent,
+    ChatCompletionResponseEventType,
+    ChatCompletionResponseStreamChunk,
+    CompletionRequest,
+    CompletionResponse,
+    CompletionResponseStreamChunk,
+    Inference,
+    InterleavedContent,
+    LogProbConfig,
+    Message,
+    ResponseFormat,
+    TokenLogProbs,
+    ToolCallDelta,
+    ToolCallParseStatus,
+    ToolChoice,
+)
 
-from llama_stack.providers.utils.inference.model_registry import build_model_alias
-from llama_stack.apis.inference import *  # noqa: F403
+from llama_stack.apis.models import ModelType
 from llama_stack.providers.datatypes import ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.embedding_mixin import (
     SentenceTransformerEmbeddingMixin,
 )
-from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
+from llama_stack.providers.utils.inference.model_registry import (
+    build_model_alias,
+    ModelRegistryHelper,
+)
 from llama_stack.providers.utils.inference.prompt_adapter import (
-    convert_image_media_to_url,
-    request_has_media,
+    augment_content_with_response_format_prompt,
+    chat_completion_request_to_messages,
+    interleaved_content_convert_to_raw,
 )
 from .config import MetaReferenceInferenceConfig
-from .generation import Llama
+from .generation import (
+    ChatCompletionRequestWithRawContent,
+    CompletionRequestWithRawContent,
+    Llama,
+)
 from .model_parallel import LlamaModelParallelGenerator
 
 log = logging.getLogger(__name__)
@@ -90,7 +124,7 @@ class MetaReferenceInferenceImpl(
     async def completion(
         self,
         model_id: str,
-        content: InterleavedTextMedia,
+        content: InterleavedContent,
         sampling_params: Optional[SamplingParams] = SamplingParams(),
         response_format: Optional[ResponseFormat] = None,
         stream: Optional[bool] = False,
@@ -99,6 +133,7 @@ class MetaReferenceInferenceImpl(
         if logprobs:
             assert logprobs.top_k == 1, f"Unexpected top_k={logprobs.top_k}"
 
+        content = augment_content_with_response_format_prompt(response_format, content)
         request = CompletionRequest(
             model=model_id,
             content=content,
@@ -108,7 +143,7 @@ class MetaReferenceInferenceImpl(
             logprobs=logprobs,
         )
         self.check_model(request)
-        request = await request_with_localized_media(request)
+        request = await convert_request_to_raw(request)
 
         if request.stream:
             return self._stream_completion(request)
@@ -233,7 +268,13 @@ class MetaReferenceInferenceImpl(
             logprobs=logprobs,
         )
         self.check_model(request)
-        request = await request_with_localized_media(request)
+
+        # augment and rewrite messages depending on the model
+        request.messages = chat_completion_request_to_messages(
+            request, self.model.core_model_id.value
+        )
+        # download media and convert to raw content so we can send it to the model
+        request = await convert_request_to_raw(request)
 
         if self.config.create_distributed_process_group:
             if SEMAPHORE.locked():
@@ -406,29 +447,16 @@ class MetaReferenceInferenceImpl(
                 yield x
 
 
-async def request_with_localized_media(
+async def convert_request_to_raw(
     request: Union[ChatCompletionRequest, CompletionRequest],
-) -> Union[ChatCompletionRequest, CompletionRequest]:
-    if not request_has_media(request):
-        return request
-
-    async def _convert_single_content(content):
-        if isinstance(content, ImageMedia):
-            url = await convert_image_media_to_url(content, download=True)
-            return ImageMedia(image=URL(uri=url))
-        else:
-            return content
-
-    async def _convert_content(content):
-        if isinstance(content, list):
-            return [await _convert_single_content(c) for c in content]
-        else:
-            return await _convert_single_content(content)
-
+) -> Union[ChatCompletionRequestWithRawContent, CompletionRequestWithRawContent]:
     if isinstance(request, ChatCompletionRequest):
+        messages = []
         for m in request.messages:
-            m.content = await _convert_content(m.content)
+            content = await interleaved_content_convert_to_raw(m.content)
+            messages.append(RawMessage(**m.model_dump(), content=content))
+        request.messages = messages
     else:
-        request.content = await _convert_content(request.content)
+        request.content = await interleaved_content_convert_to_raw(request.content)
 
     return request