feat(response): inject json schema into the last message instead of system prompt

ag2/__init__.py

@@ -19,7 +19,7 @@
 from .middleware import Middleware
 from .observers import observer
 from .plugin import Plugin
-from .response import PromptedSchema, ResponseSchema, response_schema
+from .response import LastMessagePromptedSchema, PromptedSchema, ResponseSchema, response_schema
 from .spec import AgentSpec
 from .stream import MemoryStream
 from .task import Task, TaskInject, TaskSpec
@@ -41,6 +41,7 @@
     "ImageInput",
     "Inject",
     "KnowledgeConfig",
+    "LastMessagePromptedSchema",
     "MemoryStream",
     "Middleware",
     "Plugin",

ag2/agent.py

@@ -46,6 +46,7 @@
     Input,
     ModelRequest,
     ModelResponse,
+    TextInput,
     ToolCallsEvent,
     ToolResultsEvent,
     UsageEvent,
@@ -1400,6 +1401,8 @@ async def _call_client(event: BaseEvent, context: Context) -> None:
                     await context.send(DrainedModelRequest(merged.parts))
 
                 messages = await context.stream.history.get_events()
+                if final_schema and final_schema.last_message_prompt:
+                    messages = _inject_last_message_prompt(messages, final_schema.last_message_prompt)
                 result = await llm_call(messages, context)
                 # Emit usage at the point it is spent, decoupled from the
                 # response, so token accounting never depends on a response
@@ -1666,6 +1669,32 @@ def _drain_pending(context: Context) -> ModelRequest | None:
     return ModelRequest(parts)
 
 
+def _inject_last_message_prompt(
+    messages: list[BaseEvent],
+    instruction: str,
+) -> list[BaseEvent]:
+    """Return a copy of ``messages`` with ``instruction`` appended to the last user turn.
+
+    Finds the last ``ModelRequest`` and appends a ``TextInput`` carrying
+    ``instruction`` to its ``parts``. The original events are not mutated —
+    only the returned list and the replaced ``ModelRequest`` copy are new.
+    Returns the original list unchanged when there is no ``ModelRequest``.
+    """
+    last_request_idx = -1
+    original: ModelRequest | None = None
+    for i in range(len(messages) - 1, -1, -1):
+        msg = messages[i]
+        if isinstance(msg, ModelRequest):
+            last_request_idx = i
+            original = msg
+            break
+    if original is None:
+        return messages
+
+    augmented = ModelRequest([*original.parts, TextInput(instruction)])
+    return [*messages[:last_request_idx], augmented, *messages[last_request_idx + 1 :]]
+
+
 @asynccontextmanager
 async def _observer_lifecycle(
     observers: Sequence[Observer],

ag2/response/__init__.py

@@ -3,11 +3,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .callable import response_schema
-from .prompted import PromptedSchema
+from .prompted import LastMessagePromptedSchema, PromptedSchema
 from .proto import ResponseProto
 from .schema import ResponseSchema
 
 __all__ = (
+    "LastMessagePromptedSchema",
     "PromptedSchema",
     "ResponseProto",
     "ResponseSchema",

ag2/response/prompted.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import json
-from typing import overload
+from typing import Literal, overload
 
 from fast_depends import Provider
 from typing_extensions import TypeVar as TypeVar313
@@ -24,15 +24,23 @@
     "Do not include any text, markdown formatting, or explanation outside the JSON object."
 )
 
+InjectTo = Literal["system", "last_message"]
+
 
 class PromptedSchema(ResponseProto[T]):
     """Response schema that uses prompt-based instructions instead of native API structured output.
 
     Use this for models or providers that do not support built-in structured output
-    (e.g., ``response_format``). The JSON schema is injected into the system prompt,
+    (e.g., ``response_format``). The JSON schema is injected into the prompt,
     instructing the model to return valid JSON matching the schema. Validation is
     still performed using the inner schema's ``validate`` method.
 
+    By default the schema instruction is added to the system prompt. Pass
+    ``inject_to="last_message"`` (or use :class:`LastMessagePromptedSchema`) to
+    append it to the last user message instead — useful when the model pays more
+    attention to the end of the conversation, or when a system prompt is not
+    supported.
+
     Examples:
         Using with a type directly::
 
@@ -42,6 +50,10 @@ class PromptedSchema(ResponseProto[T]):
 
             schema = ResponseSchema(int | str)
             agent = Agent(..., response_schema=PromptedSchema(schema))
+
+        Injecting into the last user message::
+
+            agent = Agent(..., response_schema=LastMessagePromptedSchema(MyModel))
     """
 
     @overload
@@ -51,6 +63,7 @@ def __init__(
         /,
         *,
         prompt_template: str | None = None,
+        inject_to: InjectTo = "system",
     ) -> None: ...
 
     @overload
@@ -60,6 +73,7 @@ def __init__(
         /,
         *,
         prompt_template: str | None = None,
+        inject_to: InjectTo = "system",
     ) -> None: ...
 
     @overload
@@ -69,6 +83,7 @@ def __init__(
         /,
         *,
         prompt_template: str | None = None,
+        inject_to: InjectTo = "system",
     ) -> None: ...
 
     def __init__(
@@ -77,6 +92,7 @@ def __init__(
         /,
         *,
         prompt_template: str | None = None,
+        inject_to: InjectTo = "system",
     ) -> None:
         self._inner = ResponseSchema[T].ensure_schema(inner)
 
@@ -85,19 +101,86 @@ def __init__(
 
         self._json_schema = self._inner.json_schema
         self._prompt_template = prompt_template or _DEFAULT_PROMPT_TEMPLATE
+        self._inject_to = inject_to
+
+        instruction: str | None
         if self._json_schema:
             schema_str = json.dumps(self._json_schema, indent=2)
-            self.system_prompt = self._prompt_template.format(schema=schema_str)
+            instruction = self._prompt_template.format(schema=schema_str)
         else:
-            self.system_prompt = None
+            instruction = None
 
         # Set public property to None to avoid native JSON schema validation
         self.json_schema = None
 
+        if inject_to == "last_message":
+            self.system_prompt = None
+            self.last_message_prompt = instruction
+        else:
+            self.system_prompt = instruction
+            self.last_message_prompt = None
+
     async def validate(
         self,
         response: str,
         context: "Context",
         provider: "Provider | None" = None,
     ) -> T:
         return await self._inner.validate(response, context, provider)
+
+
+class LastMessagePromptedSchema(PromptedSchema[T]):
+    """:class:`PromptedSchema` that injects the schema instruction into the last user message.
+
+    This is a convenience subclass equivalent to
+    ``PromptedSchema(inner, inject_to="last_message")``. Some models attend more
+    strongly to the end of the conversation, and some endpoints have no system
+    prompt slot; appending the JSON-schema instruction to the final user turn
+    keeps it adjacent to the request the model must answer.
+
+    Examples:
+        Using with a type directly::
+
+            agent = Agent(..., response_schema=LastMessagePromptedSchema(MyModel))
+
+        Wrapping an existing ResponseProto::
+
+            schema = ResponseSchema(int | str)
+            agent = Agent(..., response_schema=LastMessagePromptedSchema(schema))
+    """
+
+    @overload
+    def __init__(
+        self,
+        inner: ResponseProto[T],
+        /,
+        *,
+        prompt_template: str | None = None,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        inner: type[T],
+        /,
+        *,
+        prompt_template: str | None = None,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        inner: ClassInfo,
+        /,
+        *,
+        prompt_template: str | None = None,
+    ) -> None: ...
+
+    def __init__(
+        self,
+        inner: "ResponseProto[T] | type[T] | ClassInfo",
+        /,
+        *,
+        prompt_template: str | None = None,
+    ) -> None:
+        super().__init__(inner, prompt_template=prompt_template, inject_to="last_message")

ag2/response/proto.py

@@ -18,6 +18,7 @@ class ResponseProto(ABC, Generic[T]):
     description: str | None
     json_schema: dict[str, Any] | None
     system_prompt: str | None
+    last_message_prompt: str | None = None
 
     @abstractmethod
     async def validate(

test/agent/test_response_schema.py

@@ -8,7 +8,7 @@
 import pytest
 from pydantic import BaseModel
 
-from ag2 import Agent, PromptedSchema, ResponseSchema, response_schema
+from ag2 import Agent, LastMessagePromptedSchema, PromptedSchema, ResponseSchema, response_schema
 from ag2.testing import TestConfig, TrackingConfig
 
 
@@ -248,3 +248,80 @@ async def test_ask_override_does_not_persist(self) -> None:
 
         next_reply = await reply.ask("Again!")
         assert await next_reply.content() == "42"
+
+
+@pytest.mark.asyncio()
+class TestLastMessagePromptedSchema:
+    async def test_last_message_prompted_schema_with_type(self) -> None:
+        tracking = TrackingConfig(TestConfig('{"data": 42}'))
+        agent = Agent("test", config=tracking, response_schema=LastMessagePromptedSchema(int))
+
+        reply = await agent.ask("Hi!")
+        result = await reply.content()
+
+        assert result == 42
+
+    async def test_instruction_lands_in_last_user_message(self) -> None:
+        tracking = TrackingConfig(TestConfig('{"data": 42}'))
+        agent = Agent("test", config=tracking, response_schema=LastMessagePromptedSchema(int))
+
+        await agent.ask("Hi!")
+
+        last_message = tracking.mock.call_args_list[0][0][0]
+        parts = last_message.parts
+        assert len(parts) >= 2
+        assert parts[0].content == "Hi!"
+        assert '"type": "integer"' in parts[-1].content
+
+    async def test_instruction_not_in_system_prompt(self) -> None:
+        tracking = TrackingConfig(TestConfig('{"data": 42}'))
+        agent = Agent(
+            "test",
+            config=tracking,
+            prompt="You are helpful.",
+            response_schema=LastMessagePromptedSchema(int),
+        )
+
+        await agent.ask("Hi!")
+
+        last_message = tracking.mock.call_args_list[0][0][0]
+        parts = last_message.parts
+        # The user's original question is first; the schema instruction is appended
+        assert parts[0].content == "Hi!"
+        assert "You are helpful." not in parts[-1].content
+        assert '"type": "integer"' in parts[-1].content
+
+    async def test_ask_level_override(self) -> None:
+        agent = Agent("test", config=TestConfig('{"data": 42}'))
+
+        reply = await agent.ask("Hi!", response_schema=LastMessagePromptedSchema(int))
+        result = await reply.content()
+
+        assert result == 42
+
+    async def test_with_callable(self) -> None:
+        @response_schema
+        def double(content: str) -> int:
+            return int(content) * 2
+
+        agent = Agent("test", config=TestConfig("21"), response_schema=LastMessagePromptedSchema(double))
+
+        reply = await agent.ask("Hi!")
+        result = await reply.content()
+
+        assert result == 42
+
+    async def test_retry_appends_instruction_each_call(self) -> None:
+        tracking = TrackingConfig(TestConfig("not a number", '{"data": 42}'))
+        agent = Agent("test", config=tracking, response_schema=LastMessagePromptedSchema(int))
+
+        reply = await agent.ask("Hi!")
+        result = await reply.content(retries=1)
+
+        assert result == 42
+        assert tracking.mock.call_count == 2
+
+        # Both calls should have the instruction appended to the last user message
+        for call in tracking.mock.call_args_list:
+            last_message = call[0][0]
+            assert '"type": "integer"' in last_message.parts[-1].content

test/config/bedrock/test_response_schema.py

@@ -13,7 +13,7 @@
 from ag2.config.bedrock import BedrockClient
 from ag2.config.bedrock.mappers import response_proto_to_output_config
 from ag2.events import ModelRequest, TextInput
-from ag2.response import PromptedSchema, ResponseSchema
+from ag2.response import LastMessagePromptedSchema, PromptedSchema, ResponseSchema
 from test.config.bedrock._helpers import FakeBedrockRuntime, StubSession, make_call_context
 
 
@@ -62,6 +62,7 @@ def test_output_config_nested_additional_properties() -> None:
 def test_output_config_none_without_json_schema() -> None:
     assert response_proto_to_output_config(None) is None
     assert response_proto_to_output_config(PromptedSchema(Verdict)) is None
+    assert response_proto_to_output_config(LastMessagePromptedSchema(Verdict)) is None
 
 
 @pytest.mark.asyncio

test/config/zai/test_response_schema.py

@@ -10,7 +10,7 @@
 from ag2.config.zai import ZAIClient
 from ag2.config.zai.mappers import response_proto_to_format, schema_instruction
 from ag2.events import ModelRequest, TextInput
-from ag2.response import PromptedSchema, ResponseSchema
+from ag2.response import LastMessagePromptedSchema, PromptedSchema, ResponseSchema
 from test.config.zai._helpers import FakeCompletions, FakeZAIClient, make_call_context
 
 
@@ -28,6 +28,7 @@ def test_none_response_schema_returns_none() -> None:
     assert response_proto_to_format(None) is None
     # PromptedSchema carries no native json_schema — it prompts for JSON itself.
     assert response_proto_to_format(PromptedSchema(Verdict)) is None
+    assert response_proto_to_format(LastMessagePromptedSchema(Verdict)) is None
 
 
 def test_native_schema_uses_json_mode() -> None:
@@ -47,6 +48,8 @@ def test_schema_instruction_describes_schema() -> None:
 def test_schema_instruction_skips_prompted_schema() -> None:
     # PromptedSchema supplies its own system prompt; don't double up.
     assert schema_instruction(PromptedSchema(Verdict)) is None
+    # LastMessagePromptedSchema carries no native json_schema either.
+    assert schema_instruction(LastMessagePromptedSchema(Verdict)) is None
     assert schema_instruction(None) is None