Support sys_prompt behavior in inference

Author: ehhuang

llama_stack/apis/inference/inference.py

@@ -308,14 +308,46 @@ class CompletionResponseStreamChunk(BaseModel):
     logprobs: Optional[List[TokenLogProbs]] = None
 
 
+@json_schema_type
+class SystemMessageBehavior(Enum):
+    """Config for how to override the default system prompt.
+
+    :cvar append: Appends the provided system message to the default system prompt:
+        https://www.llama.com/docs/model-cards-and-prompt-formats/llama3_2/#-function-definitions-in-the-system-prompt-
+    :cvar replace: Replaces the default system prompt with the provided system message. The system message can include the string
+        '{{function_definitions}}' to indicate where the function definitions should be inserted.
+    """
+
+    append = "append"
+    replace = "replace"
+
+
+@json_schema_type
+class ToolConfig(BaseModel):
+    """Configuration for tool use.
+
+    :param tool_choice: (Optional) Whether tool use is required or automatic. Defaults to ToolChoice.auto.
+    :param tool_prompt_format: (Optional) Instructs the model how to format tool calls. By default, Llama Stack will attempt to use a format that is best adapted to the model.
+        - `ToolPromptFormat.json`: The tool calls are formatted as a JSON object.
+        - `ToolPromptFormat.function_tag`: The tool calls are enclosed in a <function=function_name> tag.
+        - `ToolPromptFormat.python_list`: The tool calls are output as Python syntax -- a list of function calls.
+    """
+
+    tool_choice: Optional[ToolChoice] = Field(default=ToolChoice.auto)
+    tool_prompt_format: Optional[ToolPromptFormat] = Field(default=None)
+    system_message_behavior: SystemMessageBehavior = Field(default=SystemMessageBehavior.append)
+
+
 # This is an internally used class
+@json_schema_type
 class ChatCompletionRequest(BaseModel):
     model: str
     messages: List[Message]
     sampling_params: Optional[SamplingParams] = SamplingParams()
+
     tools: Optional[List[ToolDefinition]] = Field(default_factory=list)
-    tool_choice: Optional[ToolChoice] = Field(default=ToolChoice.auto)
-    tool_prompt_format: Optional[ToolPromptFormat] = Field(default=None)
+    tool_config: Optional[ToolConfig] = Field(default_factory=ToolConfig)
+
     response_format: Optional[ResponseFormat] = None
     stream: Optional[bool] = False
     logprobs: Optional[LogProbConfig] = None
@@ -404,6 +436,7 @@ async def chat_completion(
         response_format: Optional[ResponseFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionResponseStreamChunk]]:
         """Generate a chat completion for the given messages using the specified model.
 
@@ -412,15 +445,20 @@ async def chat_completion(
         :param sampling_params: Parameters to control the sampling strategy
         :param tools: (Optional) List of tool definitions available to the model
         :param tool_choice: (Optional) Whether tool use is required or automatic. Defaults to ToolChoice.auto.
+            .. deprecated::
+               Use tool_config instead.
         :param tool_prompt_format: (Optional) Instructs the model how to format tool calls. By default, Llama Stack will attempt to use a format that is best adapted to the model.
             - `ToolPromptFormat.json`: The tool calls are formatted as a JSON object.
             - `ToolPromptFormat.function_tag`: The tool calls are enclosed in a <function=function_name> tag.
             - `ToolPromptFormat.python_list`: The tool calls are output as Python syntax -- a list of function calls.
+            .. deprecated::
+               Use tool_config instead.
         :param response_format: (Optional) Grammar specification for guided (structured) decoding. There are two options:
             - `ResponseFormat.json_schema`: The grammar is a JSON schema. Most providers support this format.
             - `ResponseFormat.grammar`: The grammar is a BNF grammar. This format is more flexible, but not all providers support it.
         :param stream: (Optional) If True, generate an SSE event stream of the response. Defaults to False.
         :param logprobs: (Optional) If specified, log probabilities for each token position will be returned.
+        :param tool_config: (Optional) Configuration for tool use.
         :returns: If stream=False, returns a ChatCompletionResponse with the full completion.
                  If stream=True, returns an SSE event stream of ChatCompletionResponseStreamChunk
         """

llama_stack/distribution/routers/routers.py

@@ -24,6 +24,7 @@
     ResponseFormat,
     SamplingParams,
     ToolChoice,
+    ToolConfig,
     ToolDefinition,
     ToolPromptFormat,
 )
@@ -132,12 +133,23 @@ async def chat_completion(
         tool_prompt_format: Optional[ToolPromptFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> AsyncGenerator:
         model = await self.routing_table.get_model(model_id)
         if model is None:
             raise ValueError(f"Model '{model_id}' not found")
         if model.model_type == ModelType.embedding:
             raise ValueError(f"Model '{model_id}' is an embedding model and does not support chat completions")
+        if tool_config:
+            if tool_choice != tool_config.tool_choice:
+                raise ValueError("tool_choice and tool_config.tool_choice must match")
+            if tool_prompt_format != tool_config.tool_prompt_format:
+                raise ValueError("tool_prompt_format and tool_config.tool_prompt_format must match")
+        else:
+            tool_config = ToolConfig(
+                tool_choice=tool_choice,
+                tool_prompt_format=tool_prompt_format,
+            )
         params = dict(
             model_id=model_id,
             messages=messages,
@@ -148,6 +160,7 @@ async def chat_completion(
             response_format=response_format,
             stream=stream,
             logprobs=logprobs,
+            tool_config=tool_config,
         )
         provider = self.routing_table.get_provider_impl(model_id)
         if stream:

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

@@ -400,7 +400,7 @@ def chat_completion(
         yield from self.generate(
             model_input=self.formatter.encode_dialog_prompt(
                 request.messages,
-                request.tool_prompt_format,
+                request.tool_config.tool_prompt_format,
             ),
             max_gen_len=max_gen_len,
             temperature=temperature,

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

@@ -38,6 +38,7 @@
     ResponseFormat,
     TokenLogProbs,
     ToolChoice,
+    ToolConfig,
 )
 from llama_stack.apis.models import Model, ModelType
 from llama_stack.providers.datatypes import ModelsProtocolPrivate
@@ -252,6 +253,7 @@ async def chat_completion(
         tool_prompt_format: Optional[ToolPromptFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> AsyncGenerator:
         if logprobs:
             assert logprobs.top_k == 1, f"Unexpected top_k={logprobs.top_k}"
@@ -262,11 +264,10 @@ async def chat_completion(
             messages=messages,
             sampling_params=sampling_params,
             tools=tools or [],
-            tool_choice=tool_choice,
-            tool_prompt_format=tool_prompt_format,
             response_format=response_format,
             stream=stream,
             logprobs=logprobs,
+            tool_config=tool_config,
         )
         self.check_model(request)
 

llama_stack/providers/inline/inference/sentence_transformers/sentence_transformers.py

@@ -17,6 +17,7 @@
     ToolChoice,
     ToolDefinition,
     ToolPromptFormat,
+    ToolConfig,
 )
 from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.embedding_mixin import (
@@ -71,5 +72,6 @@ async def chat_completion(
         tool_prompt_format: Optional[ToolPromptFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> AsyncGenerator:
         raise ValueError("Sentence transformers don't support chat completion")

llama_stack/providers/inline/inference/vllm/vllm.py

@@ -30,6 +30,7 @@
     ResponseFormat,
     SamplingParams,
     ToolChoice,
+    ToolConfig,
     ToolDefinition,
     ToolPromptFormat,
 )
@@ -159,6 +160,7 @@ async def chat_completion(
         response_format: Optional[ResponseFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> ChatCompletionResponse | ChatCompletionResponseStreamChunk:
         assert self.engine is not None
 
@@ -167,10 +169,9 @@ async def chat_completion(
             messages=messages,
             sampling_params=sampling_params,
             tools=tools or [],
-            tool_choice=tool_choice,
-            tool_prompt_format=tool_prompt_format,
             stream=stream,
             logprobs=logprobs,
+            tool_config=tool_config,
         )
 
         log.info("Sampling params: %s", sampling_params)

llama_stack/providers/remote/inference/bedrock/bedrock.py

@@ -24,6 +24,7 @@
     ResponseFormat,
     SamplingParams,
     ToolChoice,
+    ToolConfig,
     ToolDefinition,
     ToolPromptFormat,
 )
@@ -102,18 +103,18 @@ async def chat_completion(
         tool_prompt_format: Optional[ToolPromptFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionResponseStreamChunk]]:
         model = await self.model_store.get_model(model_id)
         request = ChatCompletionRequest(
             model=model.provider_resource_id,
             messages=messages,
             sampling_params=sampling_params,
             tools=tools or [],
-            tool_choice=tool_choice,
-            tool_prompt_format=tool_prompt_format,
             response_format=response_format,
             stream=stream,
             logprobs=logprobs,
+            tool_config=tool_config,
         )
 
         if stream:

llama_stack/providers/remote/inference/cerebras/cerebras.py

@@ -24,6 +24,7 @@
     ResponseFormat,
     SamplingParams,
     ToolChoice,
+    ToolConfig,
     ToolDefinition,
     ToolPromptFormat,
 )
@@ -128,6 +129,7 @@ async def chat_completion(
         response_format: Optional[ResponseFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> AsyncGenerator:
         model = await self.model_store.get_model(model_id)
         request = ChatCompletionRequest(
@@ -140,6 +142,7 @@ async def chat_completion(
             response_format=response_format,
             stream=stream,
             logprobs=logprobs,
+            tool_config=tool_config,
         )
 
         if stream:

llama_stack/providers/remote/inference/databricks/databricks.py

@@ -89,16 +89,16 @@ async def chat_completion(
         tool_prompt_format: Optional[ToolPromptFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> AsyncGenerator:
         request = ChatCompletionRequest(
             model=model,
             messages=messages,
             sampling_params=sampling_params,
             tools=tools or [],
-            tool_choice=tool_choice,
-            tool_prompt_format=tool_prompt_format,
             stream=stream,
             logprobs=logprobs,
+            tool_config=tool_config,
         )
 
         client = OpenAI(base_url=self.config.url, api_key=self.config.api_token)

llama_stack/providers/remote/inference/fireworks/fireworks.py

@@ -25,6 +25,7 @@
     ResponseFormatType,
     SamplingParams,
     ToolChoice,
+    ToolConfig,
     ToolDefinition,
     ToolPromptFormat,
 )
@@ -204,18 +205,18 @@ async def chat_completion(
         response_format: Optional[ResponseFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> AsyncGenerator:
         model = await self.model_store.get_model(model_id)
         request = ChatCompletionRequest(
             model=model.provider_resource_id,
             messages=messages,
             sampling_params=sampling_params,
             tools=tools or [],
-            tool_choice=tool_choice,
-            tool_prompt_format=tool_prompt_format,
             response_format=response_format,
             stream=stream,
             logprobs=logprobs,
+            tool_config=tool_config,
         )
 
         if stream:

llama_stack/providers/remote/inference/groq/groq.py

@@ -99,6 +99,7 @@ async def chat_completion(
         tool_prompt_format: Optional[ToolPromptFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionResponseStreamChunk]]:
         model_id = self.get_provider_model_id(model_id)
         if model_id == "llama-3.2-3b-preview":
@@ -115,10 +116,9 @@ async def chat_completion(
                 sampling_params=sampling_params,
                 response_format=response_format,
                 tools=tools,
-                tool_choice=tool_choice,
-                tool_prompt_format=tool_prompt_format,
                 stream=stream,
                 logprobs=logprobs,
+                tool_config=tool_config,
             )
         )
 

llama_stack/providers/remote/inference/groq/groq_utils.py

@@ -79,7 +79,7 @@ def convert_chat_completion_request(
         # so we exclude it for now
         warnings.warn("repetition_penalty is not supported")
 
-    if request.tool_prompt_format != ToolPromptFormat.json:
+    if request.tool_config.tool_prompt_format != ToolPromptFormat.json:
         warnings.warn("tool_prompt_format is not used by Groq. Ignoring.")
 
     sampling_options = get_sampling_strategy_options(request.sampling_params)
@@ -93,7 +93,7 @@ def convert_chat_completion_request(
         temperature=sampling_options.get("temperature", 1.0),
         top_p=sampling_options.get("top_p", 1.0),
         tools=[_convert_groq_tool_definition(tool) for tool in request.tools or []],
-        tool_choice=request.tool_choice.value if request.tool_choice else None,
+        tool_choice=(request.tool_config.tool_choice.value if request.tool_config.tool_choice else None),
     )
 
 

llama_stack/providers/remote/inference/nvidia/nvidia.py

@@ -171,6 +171,7 @@ async def chat_completion(
         tool_prompt_format: Optional[ToolPromptFormat] = None,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
     ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionResponseStreamChunk]]:
         if tool_prompt_format:
             warnings.warn("tool_prompt_format is not supported by NVIDIA NIM, ignoring")
@@ -184,10 +185,9 @@ async def chat_completion(
                 sampling_params=sampling_params,
                 response_format=response_format,
                 tools=tools,
-                tool_choice=tool_choice,
-                tool_prompt_format=tool_prompt_format,
                 stream=stream,
                 logprobs=logprobs,
+                tool_config=tool_config,
             ),
             n=1,
         )

llama_stack/providers/remote/inference/nvidia/openai_utils.py

@@ -282,9 +282,9 @@ async def convert_chat_completion_request(
 
     if request.tools:
         payload.update(tools=[_convert_tooldef_to_openai_tool(tool) for tool in request.tools])
-        if request.tool_choice:
+        if request.tool_config.tool_choice:
             payload.update(
-                tool_choice=request.tool_choice.value
+                tool_choice=request.tool_config.tool_choice.value
             )  # we cannot include tool_choice w/o tools, server will complain
 
     if request.logprobs: