Merge branch 'main' into fix_DPOAlignmentConfig_schema

.pre-commit-config.yaml

@@ -129,6 +129,22 @@ repos:
         require_serial: true
         always_run: true
         files: ^llama_stack/.*$
+      - id: forbid-pytest-asyncio
+        name: Block @pytest.mark.asyncio and @pytest_asyncio.fixture
+        entry: bash
+        language: system
+        types: [python]
+        pass_filenames: true
+        args:
+          - -c
+          - |
+            grep -EnH '^[^#]*@pytest\.mark\.asyncio|@pytest_asyncio\.fixture' "$@" && {
+              echo;
+              echo "❌ Do not use @pytest.mark.asyncio or @pytest_asyncio.fixture."
+              echo "   pytest is already configured with async-mode=auto."
+              echo;
+              exit 1;
+            } || true
 
 ci:
     autofix_commit_msg: 🎨 [pre-commit.ci] Auto format from pre-commit.com hooks

llama_stack/cli/stack/run.py

@@ -47,8 +47,7 @@ class StackRun(Subcommand):
         self.parser.add_argument(
             "--image-name",
             type=str,
-            default=os.environ.get("CONDA_DEFAULT_ENV"),
+            help="Name of the image to run.",
-            help="Name of the image to run. Defaults to the current environment",
         )
         self.parser.add_argument(
             "--env",

tests/integration/telemetry/test_telemetry.py

@@ -5,41 +5,183 @@
 # the root directory of this source tree.
 
 import time
+from datetime import UTC, datetime
 from uuid import uuid4
 
 import pytest
 from llama_stack_client import Agent
 
 
-@pytest.mark.skip(reason="telemetry is not stable")
+@pytest.fixture(scope="module", autouse=True)
-def test_agent_query_spans(llama_stack_client, text_model_id):
+def setup_telemetry_data(llama_stack_client, text_model_id):
+    """Setup fixture that creates telemetry data before tests run."""
     agent = Agent(llama_stack_client, model=text_model_id, instructions="You are a helpful assistant")
-    session_id = agent.create_session(f"test-session-{uuid4()}")
+
-    agent.create_turn(
+    session_id = agent.create_session(f"test-setup-session-{uuid4()}")
-        messages=[
+
-            {
+    messages = [
-                "role": "user",
+        "What is 2 + 2?",
-                "content": "Give me a sentence that contains the word: hello",
+        "Tell me a short joke",
-            }
+    ]
-        ],
+
-        session_id=session_id,
+    for msg in messages:
-        stream=False,
+        agent.create_turn(
+            messages=[{"role": "user", "content": msg}],
+            session_id=session_id,
+            stream=False,
+        )
+
+    for i in range(2):
+        llama_stack_client.inference.chat_completion(
+            model_id=text_model_id, messages=[{"role": "user", "content": f"Test trace {i}"}]
+        )
+
+    start_time = time.time()
+
+    while time.time() - start_time < 30:
+        traces = llama_stack_client.telemetry.query_traces(limit=10)
+        if len(traces) >= 4:
+            break
+        time.sleep(1)
+
+    if len(traces) < 4:
+        pytest.fail(f"Failed to create sufficient telemetry data after 30s. Got {len(traces)} traces.")
+
+    yield
+
+
+def test_query_traces_basic(llama_stack_client):
+    """Test basic trace querying functionality with proper data validation."""
+    all_traces = llama_stack_client.telemetry.query_traces(limit=5)
+
+    assert isinstance(all_traces, list), "Should return a list of traces"
+    assert len(all_traces) >= 4, "Should have at least 4 traces from setup"
+
+    # Verify trace structure and data quality
+    first_trace = all_traces[0]
+    assert hasattr(first_trace, "trace_id"), "Trace should have trace_id"
+    assert hasattr(first_trace, "start_time"), "Trace should have start_time"
+    assert hasattr(first_trace, "root_span_id"), "Trace should have root_span_id"
+
+    # Validate trace_id is a valid UUID format
+    assert isinstance(first_trace.trace_id, str) and len(first_trace.trace_id) > 0, (
+        "trace_id should be non-empty string"
     )
 
-    # Wait for the span to be logged
+    # Validate start_time format and not in the future
-    time.sleep(2)
+    now = datetime.now(UTC)
+    if isinstance(first_trace.start_time, str):
+        trace_time = datetime.fromisoformat(first_trace.start_time.replace("Z", "+00:00"))
+    else:
+        # start_time is already a datetime object
+        trace_time = first_trace.start_time
+        if trace_time.tzinfo is None:
+            trace_time = trace_time.replace(tzinfo=UTC)
 
-    agent_logs = []
+    # Ensure trace time is not in the future (but allow any age in the past for persistent test data)
+    time_diff = (now - trace_time).total_seconds()
+    assert time_diff >= 0, f"Trace start_time should not be in the future, got {time_diff}s"
 
-    for span in llama_stack_client.telemetry.query_spans(
+    # Validate root_span_id exists and is non-empty
-        attribute_filters=[
+    assert isinstance(first_trace.root_span_id, str) and len(first_trace.root_span_id) > 0, (
-            {"key": "session_id", "op": "eq", "value": session_id},
+        "root_span_id should be non-empty string"
-        ],
+    )
-        attributes_to_return=["input", "output"],
-    ):
-        if span.attributes["output"] != "no shields":
-            agent_logs.append(span.attributes)
 
-    assert len(agent_logs) == 1
+    # Test querying specific trace by ID
-    assert "Give me a sentence that contains the word: hello" in agent_logs[0]["input"]
+    specific_trace = llama_stack_client.telemetry.get_trace(trace_id=first_trace.trace_id)
-    assert "hello" in agent_logs[0]["output"].lower()
+    assert specific_trace.trace_id == first_trace.trace_id, "Retrieved trace should match requested ID"
+    assert specific_trace.start_time == first_trace.start_time, "Retrieved trace should have same start_time"
+    assert specific_trace.root_span_id == first_trace.root_span_id, "Retrieved trace should have same root_span_id"
+
+    # Test pagination with proper validation
+    recent_traces = llama_stack_client.telemetry.query_traces(limit=3, offset=0)
+    assert len(recent_traces) <= 3, "Should return at most 3 traces when limit=3"
+    assert len(recent_traces) >= 1, "Should return at least 1 trace"
+
+    # Verify all traces have required fields
+    for trace in recent_traces:
+        assert hasattr(trace, "trace_id") and trace.trace_id, "Each trace should have non-empty trace_id"
+        assert hasattr(trace, "start_time") and trace.start_time, "Each trace should have non-empty start_time"
+        assert hasattr(trace, "root_span_id") and trace.root_span_id, "Each trace should have non-empty root_span_id"
+
+
+def test_query_spans_basic(llama_stack_client):
+    """Test basic span querying functionality with proper validation."""
+    spans = llama_stack_client.telemetry.query_spans(attribute_filters=[], attributes_to_return=[])
+
+    assert isinstance(spans, list), "Should return a list of spans"
+    assert len(spans) >= 1, "Should have at least one span from setup"
+
+    # Verify span structure and data quality
+    first_span = spans[0]
+    required_attrs = ["span_id", "name", "trace_id"]
+    for attr in required_attrs:
+        assert hasattr(first_span, attr), f"Span should have {attr} attribute"
+        assert getattr(first_span, attr), f"Span {attr} should not be empty"
+
+    # Validate span data types and content
+    assert isinstance(first_span.span_id, str) and len(first_span.span_id) > 0, "span_id should be non-empty string"
+    assert isinstance(first_span.name, str) and len(first_span.name) > 0, "span name should be non-empty string"
+    assert isinstance(first_span.trace_id, str) and len(first_span.trace_id) > 0, "trace_id should be non-empty string"
+
+    # Verify span belongs to a valid trace (test with traces we know exist)
+    all_traces = llama_stack_client.telemetry.query_traces(limit=10)
+    trace_ids = {t.trace_id for t in all_traces}
+    if first_span.trace_id in trace_ids:
+        trace = llama_stack_client.telemetry.get_trace(trace_id=first_span.trace_id)
+        assert trace is not None, "Should be able to retrieve trace for valid trace_id"
+        assert trace.trace_id == first_span.trace_id, "Trace ID should match span's trace_id"
+
+    # Test with span filtering and validate results
+    filtered_spans = llama_stack_client.telemetry.query_spans(
+        attribute_filters=[{"key": "name", "op": "eq", "value": first_span.name}],
+        attributes_to_return=["name", "span_id"],
+    )
+    assert isinstance(filtered_spans, list), "Should return a list with span name filter"
+
+    # Validate filtered spans if filtering works
+    if len(filtered_spans) > 0:
+        for span in filtered_spans:
+            assert hasattr(span, "name"), "Filtered spans should have name attribute"
+            assert hasattr(span, "span_id"), "Filtered spans should have span_id attribute"
+            assert span.name == first_span.name, "Filtered spans should match the filter criteria"
+            assert isinstance(span.span_id, str) and len(span.span_id) > 0, "Filtered span_id should be valid"
+
+    # Test that all spans have consistent structure
+    for span in spans:
+        for attr in required_attrs:
+            assert hasattr(span, attr) and getattr(span, attr), f"All spans should have non-empty {attr}"
+
+
+def test_telemetry_pagination(llama_stack_client):
+    """Test pagination in telemetry queries."""
+    # Get total count of traces
+    all_traces = llama_stack_client.telemetry.query_traces(limit=20)
+    total_count = len(all_traces)
+    assert total_count >= 4, "Should have at least 4 traces from setup"
+
+    # Test trace pagination
+    page1 = llama_stack_client.telemetry.query_traces(limit=2, offset=0)
+    page2 = llama_stack_client.telemetry.query_traces(limit=2, offset=2)
+
+    assert len(page1) == 2, "First page should have exactly 2 traces"
+    assert len(page2) >= 1, "Second page should have at least 1 trace"
+
+    # Verify no overlap between pages
+    page1_ids = {t.trace_id for t in page1}
+    page2_ids = {t.trace_id for t in page2}
+    assert len(page1_ids.intersection(page2_ids)) == 0, "Pages should contain different traces"
+
+    # Test ordering
+    ordered_traces = llama_stack_client.telemetry.query_traces(limit=5, order_by=["start_time"])
+    assert len(ordered_traces) >= 4, "Should have at least 4 traces for ordering test"
+
+    # Verify ordering by start_time
+    for i in range(len(ordered_traces) - 1):
+        current_time = ordered_traces[i].start_time
+        next_time = ordered_traces[i + 1].start_time
+        assert current_time <= next_time, f"Traces should be ordered by start_time: {current_time} > {next_time}"
+
+    # Test limit behavior
+    limited = llama_stack_client.telemetry.query_traces(limit=3)
+    assert len(limited) == 3, "Should return exactly 3 traces when limit=3"

tests/unit/models/test_prompt_adapter.py

@@ -4,14 +4,14 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import asyncio
+import pytest
-import unittest
 
 from llama_stack.apis.inference import (
     ChatCompletionRequest,
     CompletionMessage,
     StopReason,
     SystemMessage,
+    SystemMessageBehavior,
     ToolCall,
     ToolConfig,
     UserMessage,
@@ -25,264 +25,275 @@ from llama_stack.models.llama.datatypes import (
 from llama_stack.providers.utils.inference.prompt_adapter import (
     chat_completion_request_to_messages,
     chat_completion_request_to_prompt,
+    interleaved_content_as_str,
 )
 
 MODEL = "Llama3.1-8B-Instruct"
 MODEL3_2 = "Llama3.2-3B-Instruct"
 
 
-class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
+@pytest.mark.asyncio
-    async def asyncSetUp(self):
+async def test_system_default():
-        asyncio.get_running_loop().set_debug(False)
+    content = "Hello !"
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[
+            UserMessage(content=content),
+        ],
+    )
+    messages = chat_completion_request_to_messages(request, MODEL)
+    assert len(messages) == 2
+    assert messages[-1].content == content
+    assert "Cutting Knowledge Date: December 2023" in interleaved_content_as_str(messages[0].content)
 
-    async def test_system_default(self):
-        content = "Hello !"
-        request = ChatCompletionRequest(
-            model=MODEL,
-            messages=[
-                UserMessage(content=content),
-            ],
-        )
-        messages = chat_completion_request_to_messages(request, MODEL)
-        self.assertEqual(len(messages), 2)
-        self.assertEqual(messages[-1].content, content)
-        self.assertTrue("Cutting Knowledge Date: December 2023" in messages[0].content)
 
-    async def test_system_builtin_only(self):
+@pytest.mark.asyncio
-        content = "Hello !"
+async def test_system_builtin_only():
-        request = ChatCompletionRequest(
+    content = "Hello !"
-            model=MODEL,
+    request = ChatCompletionRequest(
-            messages=[
+        model=MODEL,
-                UserMessage(content=content),
+        messages=[
-            ],
+            UserMessage(content=content),
-            tools=[
+        ],
-                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+        tools=[
-                ToolDefinition(tool_name=BuiltinTool.brave_search),
+            ToolDefinition(tool_name=BuiltinTool.code_interpreter),
-            ],
+            ToolDefinition(tool_name=BuiltinTool.brave_search),
-        )
+        ],
-        messages = chat_completion_request_to_messages(request, MODEL)
+    )
-        self.assertEqual(len(messages), 2)
+    messages = chat_completion_request_to_messages(request, MODEL)
-        self.assertEqual(messages[-1].content, content)
+    assert len(messages) == 2
-        self.assertTrue("Cutting Knowledge Date: December 2023" in messages[0].content)
+    assert messages[-1].content == content
-        self.assertTrue("Tools: brave_search" in messages[0].content)
+    assert "Cutting Knowledge Date: December 2023" in interleaved_content_as_str(messages[0].content)
+    assert "Tools: brave_search" in interleaved_content_as_str(messages[0].content)
 
-    async def test_system_custom_only(self):
-        content = "Hello !"
-        request = ChatCompletionRequest(
-            model=MODEL,
-            messages=[
-                UserMessage(content=content),
-            ],
-            tools=[
-                ToolDefinition(
-                    tool_name="custom1",
-                    description="custom1 tool",
-                    parameters={
-                        "param1": ToolParamDefinition(
-                            param_type="str",
-                            description="param1 description",
-                            required=True,
-                        ),
-                    },
-                )
-            ],
-            tool_config=ToolConfig(tool_prompt_format=ToolPromptFormat.json),
-        )
-        messages = chat_completion_request_to_messages(request, MODEL)
-        self.assertEqual(len(messages), 3)
-        self.assertTrue("Environment: ipython" in messages[0].content)
 
-        self.assertTrue("Return function calls in JSON format" in messages[1].content)
+@pytest.mark.asyncio
-        self.assertEqual(messages[-1].content, content)
+async def test_system_custom_only():
+    content = "Hello !"
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[
+            UserMessage(content=content),
+        ],
+        tools=[
+            ToolDefinition(
+                tool_name="custom1",
+                description="custom1 tool",
+                parameters={
+                    "param1": ToolParamDefinition(
+                        param_type="str",
+                        description="param1 description",
+                        required=True,
+                    ),
+                },
+            )
+        ],
+        tool_config=ToolConfig(tool_prompt_format=ToolPromptFormat.json),
+    )
+    messages = chat_completion_request_to_messages(request, MODEL)
+    assert len(messages) == 3
+    assert "Environment: ipython" in interleaved_content_as_str(messages[0].content)
 
-    async def test_system_custom_and_builtin(self):
+    assert "Return function calls in JSON format" in interleaved_content_as_str(messages[1].content)
-        content = "Hello !"
+    assert messages[-1].content == content
-        request = ChatCompletionRequest(
-            model=MODEL,
-            messages=[
-                UserMessage(content=content),
-            ],
-            tools=[
-                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
-                ToolDefinition(tool_name=BuiltinTool.brave_search),
-                ToolDefinition(
-                    tool_name="custom1",
-                    description="custom1 tool",
-                    parameters={
-                        "param1": ToolParamDefinition(
-                            param_type="str",
-                            description="param1 description",
-                            required=True,
-                        ),
-                    },
-                ),
-            ],
-        )
-        messages = chat_completion_request_to_messages(request, MODEL)
-        self.assertEqual(len(messages), 3)
 
-        self.assertTrue("Environment: ipython" in messages[0].content)
-        self.assertTrue("Tools: brave_search" in messages[0].content)
 
-        self.assertTrue("Return function calls in JSON format" in messages[1].content)
+@pytest.mark.asyncio
-        self.assertEqual(messages[-1].content, content)
+async def test_system_custom_and_builtin():
-
+    content = "Hello !"
-    async def test_completion_message_encoding(self):
+    request = ChatCompletionRequest(
-        request = ChatCompletionRequest(
+        model=MODEL,
-            model=MODEL3_2,
+        messages=[
-            messages=[
+            UserMessage(content=content),
-                UserMessage(content="hello"),
+        ],
-                CompletionMessage(
+        tools=[
-                    content="",
+            ToolDefinition(tool_name=BuiltinTool.code_interpreter),
-                    stop_reason=StopReason.end_of_turn,
+            ToolDefinition(tool_name=BuiltinTool.brave_search),
-                    tool_calls=[
+            ToolDefinition(
-                        ToolCall(
+                tool_name="custom1",
-                            tool_name="custom1",
+                description="custom1 tool",
-                            arguments={"param1": "value1"},
+                parameters={
-                            call_id="123",
+                    "param1": ToolParamDefinition(
-                        )
+                        param_type="str",
-                    ],
+                        description="param1 description",
-                ),
+                        required=True,
-            ],
+                    ),
-            tools=[
+                },
-                ToolDefinition(
-                    tool_name="custom1",
-                    description="custom1 tool",
-                    parameters={
-                        "param1": ToolParamDefinition(
-                            param_type="str",
-                            description="param1 description",
-                            required=True,
-                        ),
-                    },
-                ),
-            ],
-            tool_config=ToolConfig(tool_prompt_format=ToolPromptFormat.python_list),
-        )
-        prompt = await chat_completion_request_to_prompt(request, request.model)
-        self.assertIn('[custom1(param1="value1")]', prompt)
-
-        request.model = MODEL
-        request.tool_config.tool_prompt_format = ToolPromptFormat.json
-        prompt = await chat_completion_request_to_prompt(request, request.model)
-        self.assertIn(
-            '{"type": "function", "name": "custom1", "parameters": {"param1": "value1"}}',
-            prompt,
-        )
-
-    async def test_user_provided_system_message(self):
-        content = "Hello !"
-        system_prompt = "You are a pirate"
-        request = ChatCompletionRequest(
-            model=MODEL,
-            messages=[
-                SystemMessage(content=system_prompt),
-                UserMessage(content=content),
-            ],
-            tools=[
-                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
-            ],
-        )
-        messages = chat_completion_request_to_messages(request, MODEL)
-        self.assertEqual(len(messages), 2, messages)
-        self.assertTrue(messages[0].content.endswith(system_prompt))
-
-        self.assertEqual(messages[-1].content, content)
-
-    async def test_repalce_system_message_behavior_builtin_tools(self):
-        content = "Hello !"
-        system_prompt = "You are a pirate"
-        request = ChatCompletionRequest(
-            model=MODEL,
-            messages=[
-                SystemMessage(content=system_prompt),
-                UserMessage(content=content),
-            ],
-            tools=[
-                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
-            ],
-            tool_config=ToolConfig(
-                tool_choice="auto",
-                tool_prompt_format="python_list",
-                system_message_behavior="replace",
             ),
-        )
+        ],
-        messages = chat_completion_request_to_messages(request, MODEL3_2)
+    )
-        self.assertEqual(len(messages), 2, messages)
+    messages = chat_completion_request_to_messages(request, MODEL)
-        self.assertTrue(messages[0].content.endswith(system_prompt))
+    assert len(messages) == 3
-        self.assertIn("Environment: ipython", messages[0].content)
-        self.assertEqual(messages[-1].content, content)
 
-    async def test_repalce_system_message_behavior_custom_tools(self):
+    assert "Environment: ipython" in interleaved_content_as_str(messages[0].content)
-        content = "Hello !"
+    assert "Tools: brave_search" in interleaved_content_as_str(messages[0].content)
-        system_prompt = "You are a pirate"
+
-        request = ChatCompletionRequest(
+    assert "Return function calls in JSON format" in interleaved_content_as_str(messages[1].content)
-            model=MODEL,
+    assert messages[-1].content == content
-            messages=[
+
-                SystemMessage(content=system_prompt),
+
-                UserMessage(content=content),
+@pytest.mark.asyncio
-            ],
+async def test_completion_message_encoding():
-            tools=[
+    request = ChatCompletionRequest(
-                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+        model=MODEL3_2,
-                ToolDefinition(
+        messages=[
-                    tool_name="custom1",
+            UserMessage(content="hello"),
-                    description="custom1 tool",
+            CompletionMessage(
-                    parameters={
+                content="",
-                        "param1": ToolParamDefinition(
+                stop_reason=StopReason.end_of_turn,
-                            param_type="str",
+                tool_calls=[
-                            description="param1 description",
+                    ToolCall(
-                            required=True,
+                        tool_name="custom1",
-                        ),
+                        arguments={"param1": "value1"},
-                    },
+                        call_id="123",
-                ),
+                    )
-            ],
+                ],
-            tool_config=ToolConfig(
-                tool_choice="auto",
-                tool_prompt_format="python_list",
-                system_message_behavior="replace",
             ),
-        )
+        ],
-        messages = chat_completion_request_to_messages(request, MODEL3_2)
+        tools=[
-
+            ToolDefinition(
-        self.assertEqual(len(messages), 2, messages)
+                tool_name="custom1",
-        self.assertTrue(messages[0].content.endswith(system_prompt))
+                description="custom1 tool",
-        self.assertIn("Environment: ipython", messages[0].content)
+                parameters={
-        self.assertEqual(messages[-1].content, content)
+                    "param1": ToolParamDefinition(
-
+                        param_type="str",
-    async def test_replace_system_message_behavior_custom_tools_with_template(self):
+                        description="param1 description",
-        content = "Hello !"
+                        required=True,
-        system_prompt = "You are a pirate {{ function_description }}"
+                    ),
-        request = ChatCompletionRequest(
+                },
-            model=MODEL,
-            messages=[
-                SystemMessage(content=system_prompt),
-                UserMessage(content=content),
-            ],
-            tools=[
-                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
-                ToolDefinition(
-                    tool_name="custom1",
-                    description="custom1 tool",
-                    parameters={
-                        "param1": ToolParamDefinition(
-                            param_type="str",
-                            description="param1 description",
-                            required=True,
-                        ),
-                    },
-                ),
-            ],
-            tool_config=ToolConfig(
-                tool_choice="auto",
-                tool_prompt_format="python_list",
-                system_message_behavior="replace",
             ),
-        )
+        ],
-        messages = chat_completion_request_to_messages(request, MODEL3_2)
+        tool_config=ToolConfig(tool_prompt_format=ToolPromptFormat.python_list),
+    )
+    prompt = await chat_completion_request_to_prompt(request, request.model)
+    assert '[custom1(param1="value1")]' in prompt
 
-        self.assertEqual(len(messages), 2, messages)
+    request.model = MODEL
-        self.assertIn("Environment: ipython", messages[0].content)
+    request.tool_config = ToolConfig(tool_prompt_format=ToolPromptFormat.json)
-        self.assertIn("You are a pirate", messages[0].content)
+    prompt = await chat_completion_request_to_prompt(request, request.model)
-        # function description is present in the system prompt
+    assert '{"type": "function", "name": "custom1", "parameters": {"param1": "value1"}}' in prompt
-        self.assertIn('"name": "custom1"', messages[0].content)
+
-        self.assertEqual(messages[-1].content, content)
+
+@pytest.mark.asyncio
+async def test_user_provided_system_message():
+    content = "Hello !"
+    system_prompt = "You are a pirate"
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[
+            SystemMessage(content=system_prompt),
+            UserMessage(content=content),
+        ],
+        tools=[
+            ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+        ],
+    )
+    messages = chat_completion_request_to_messages(request, MODEL)
+    assert len(messages) == 2
+    assert interleaved_content_as_str(messages[0].content).endswith(system_prompt)
+
+    assert messages[-1].content == content
+
+
+@pytest.mark.asyncio
+async def test_replace_system_message_behavior_builtin_tools():
+    content = "Hello !"
+    system_prompt = "You are a pirate"
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[
+            SystemMessage(content=system_prompt),
+            UserMessage(content=content),
+        ],
+        tools=[
+            ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+        ],
+        tool_config=ToolConfig(
+            tool_choice="auto",
+            tool_prompt_format=ToolPromptFormat.python_list,
+            system_message_behavior=SystemMessageBehavior.replace,
+        ),
+    )
+    messages = chat_completion_request_to_messages(request, MODEL3_2)
+    assert len(messages) == 2
+    assert interleaved_content_as_str(messages[0].content).endswith(system_prompt)
+    assert "Environment: ipython" in interleaved_content_as_str(messages[0].content)
+    assert messages[-1].content == content
+
+
+@pytest.mark.asyncio
+async def test_replace_system_message_behavior_custom_tools():
+    content = "Hello !"
+    system_prompt = "You are a pirate"
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[
+            SystemMessage(content=system_prompt),
+            UserMessage(content=content),
+        ],
+        tools=[
+            ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+            ToolDefinition(
+                tool_name="custom1",
+                description="custom1 tool",
+                parameters={
+                    "param1": ToolParamDefinition(
+                        param_type="str",
+                        description="param1 description",
+                        required=True,
+                    ),
+                },
+            ),
+        ],
+        tool_config=ToolConfig(
+            tool_choice="auto",
+            tool_prompt_format=ToolPromptFormat.python_list,
+            system_message_behavior=SystemMessageBehavior.replace,
+        ),
+    )
+    messages = chat_completion_request_to_messages(request, MODEL3_2)
+
+    assert len(messages) == 2
+    assert interleaved_content_as_str(messages[0].content).endswith(system_prompt)
+    assert "Environment: ipython" in interleaved_content_as_str(messages[0].content)
+    assert messages[-1].content == content
+
+
+@pytest.mark.asyncio
+async def test_replace_system_message_behavior_custom_tools_with_template():
+    content = "Hello !"
+    system_prompt = "You are a pirate {{ function_description }}"
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[
+            SystemMessage(content=system_prompt),
+            UserMessage(content=content),
+        ],
+        tools=[
+            ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+            ToolDefinition(
+                tool_name="custom1",
+                description="custom1 tool",
+                parameters={
+                    "param1": ToolParamDefinition(
+                        param_type="str",
+                        description="param1 description",
+                        required=True,
+                    ),
+                },
+            ),
+        ],
+        tool_config=ToolConfig(
+            tool_choice="auto",
+            tool_prompt_format=ToolPromptFormat.python_list,
+            system_message_behavior=SystemMessageBehavior.replace,
+        ),
+    )
+    messages = chat_completion_request_to_messages(request, MODEL3_2)
+
+    assert len(messages) == 2
+    assert "Environment: ipython" in interleaved_content_as_str(messages[0].content)
+    assert "You are a pirate" in interleaved_content_as_str(messages[0].content)
+    # function description is present in the system prompt
+    assert '"name": "custom1"' in interleaved_content_as_str(messages[0].content)
+    assert messages[-1].content == content

tests/unit/server/test_replace_env_vars.py

@@ -5,73 +5,86 @@
 # the root directory of this source tree.
 
 import os
-import unittest
+
+import pytest
 
 from llama_stack.distribution.stack import replace_env_vars
 
 
-class TestReplaceEnvVars(unittest.TestCase):
+@pytest.fixture
-    def setUp(self):
+def setup_env_vars():
-        # Clear any existing environment variables we'll use in tests
+    # Clear any existing environment variables we'll use in tests
-        for var in ["TEST_VAR", "EMPTY_VAR", "ZERO_VAR"]:
+    for var in ["TEST_VAR", "EMPTY_VAR", "ZERO_VAR"]:
-            if var in os.environ:
+        if var in os.environ:
-                del os.environ[var]
+            del os.environ[var]
 
-        # Set up test environment variables
+    # Set up test environment variables
-        os.environ["TEST_VAR"] = "test_value"
+    os.environ["TEST_VAR"] = "test_value"
-        os.environ["EMPTY_VAR"] = ""
+    os.environ["EMPTY_VAR"] = ""
-        os.environ["ZERO_VAR"] = "0"
+    os.environ["ZERO_VAR"] = "0"
 
-    def test_simple_replacement(self):
+    yield
-        self.assertEqual(replace_env_vars("${env.TEST_VAR}"), "test_value")
 
-    def test_default_value_when_not_set(self):
+    # Cleanup after test
-        self.assertEqual(replace_env_vars("${env.NOT_SET:=default}"), "default")
+    for var in ["TEST_VAR", "EMPTY_VAR", "ZERO_VAR"]:
-
+        if var in os.environ:
-    def test_default_value_when_set(self):
+            del os.environ[var]
-        self.assertEqual(replace_env_vars("${env.TEST_VAR:=default}"), "test_value")
-
-    def test_default_value_when_empty(self):
-        self.assertEqual(replace_env_vars("${env.EMPTY_VAR:=default}"), "default")
-
-    def test_none_value_when_empty(self):
-        self.assertEqual(replace_env_vars("${env.EMPTY_VAR:=}"), None)
-
-    def test_value_when_set(self):
-        self.assertEqual(replace_env_vars("${env.TEST_VAR:=}"), "test_value")
-
-    def test_empty_var_no_default(self):
-        self.assertEqual(replace_env_vars("${env.EMPTY_VAR_NO_DEFAULT:+}"), None)
-
-    def test_conditional_value_when_set(self):
-        self.assertEqual(replace_env_vars("${env.TEST_VAR:+conditional}"), "conditional")
-
-    def test_conditional_value_when_not_set(self):
-        self.assertEqual(replace_env_vars("${env.NOT_SET:+conditional}"), None)
-
-    def test_conditional_value_when_empty(self):
-        self.assertEqual(replace_env_vars("${env.EMPTY_VAR:+conditional}"), None)
-
-    def test_conditional_value_with_zero(self):
-        self.assertEqual(replace_env_vars("${env.ZERO_VAR:+conditional}"), "conditional")
-
-    def test_mixed_syntax(self):
-        self.assertEqual(
-            replace_env_vars("${env.TEST_VAR:=default} and ${env.NOT_SET:+conditional}"), "test_value and "
-        )
-        self.assertEqual(
-            replace_env_vars("${env.NOT_SET:=default} and ${env.TEST_VAR:+conditional}"), "default and conditional"
-        )
-
-    def test_nested_structures(self):
-        data = {
-            "key1": "${env.TEST_VAR:=default}",
-            "key2": ["${env.NOT_SET:=default}", "${env.TEST_VAR:+conditional}"],
-            "key3": {"nested": "${env.NOT_SET:+conditional}"},
-        }
-        expected = {"key1": "test_value", "key2": ["default", "conditional"], "key3": {"nested": None}}
-        self.assertEqual(replace_env_vars(data), expected)
 
 
-if __name__ == "__main__":
+def test_simple_replacement(setup_env_vars):
-    unittest.main()
+    assert replace_env_vars("${env.TEST_VAR}") == "test_value"
+
+
+def test_default_value_when_not_set(setup_env_vars):
+    assert replace_env_vars("${env.NOT_SET:=default}") == "default"
+
+
+def test_default_value_when_set(setup_env_vars):
+    assert replace_env_vars("${env.TEST_VAR:=default}") == "test_value"
+
+
+def test_default_value_when_empty(setup_env_vars):
+    assert replace_env_vars("${env.EMPTY_VAR:=default}") == "default"
+
+
+def test_none_value_when_empty(setup_env_vars):
+    assert replace_env_vars("${env.EMPTY_VAR:=}") is None
+
+
+def test_value_when_set(setup_env_vars):
+    assert replace_env_vars("${env.TEST_VAR:=}") == "test_value"
+
+
+def test_empty_var_no_default(setup_env_vars):
+    assert replace_env_vars("${env.EMPTY_VAR_NO_DEFAULT:+}") is None
+
+
+def test_conditional_value_when_set(setup_env_vars):
+    assert replace_env_vars("${env.TEST_VAR:+conditional}") == "conditional"
+
+
+def test_conditional_value_when_not_set(setup_env_vars):
+    assert replace_env_vars("${env.NOT_SET:+conditional}") is None
+
+
+def test_conditional_value_when_empty(setup_env_vars):
+    assert replace_env_vars("${env.EMPTY_VAR:+conditional}") is None
+
+
+def test_conditional_value_with_zero(setup_env_vars):
+    assert replace_env_vars("${env.ZERO_VAR:+conditional}") == "conditional"
+
+
+def test_mixed_syntax(setup_env_vars):
+    assert replace_env_vars("${env.TEST_VAR:=default} and ${env.NOT_SET:+conditional}") == "test_value and "
+    assert replace_env_vars("${env.NOT_SET:=default} and ${env.TEST_VAR:+conditional}") == "default and conditional"
+
+
+def test_nested_structures(setup_env_vars):
+    data = {
+        "key1": "${env.TEST_VAR:=default}",
+        "key2": ["${env.NOT_SET:=default}", "${env.TEST_VAR:+conditional}"],
+        "key3": {"nested": "${env.NOT_SET:+conditional}"},
+    }
+    expected = {"key1": "test_value", "key2": ["default", "conditional"], "key3": {"nested": None}}
+    assert replace_env_vars(data) == expected