From da01dfeae86adbc4a9643cbcdb964b26343b9021 Mon Sep 17 00:00:00 2001 From: Rubikoid Date: Fri, 10 Jul 2026 01:31:45 +0300 Subject: [PATCH 1/2] =?UTF-8?q?feat(response):=20LastMessagePromptedSchema?= =?UTF-8?q?=20=E2=80=94=20inject=20schema=20into=20last=20user=20message?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PromptedSchema previously injected its JSON-schema instruction only into the system prompt. Some models attend more strongly to the end of the conversation, and some endpoints lack a system-prompt slot entirely. Add an `inject_to: Literal["system", "last_message"]` parameter to PromptedSchema (default "system", preserving existing behaviour). When "last_message" is chosen, the instruction is stored in the new `last_message_prompt` attribute on ResponseProto instead of `system_prompt`, and the agent loop appends it as a TextInput to the final ModelRequest before each LLM call. LastMessagePromptedSchema is a convenience subclass equivalent to PromptedSchema(inner, inject_to="last_message"), re-exported from the top-level `ag2` package. Tests cover: schema construction from types/dataclasses/pydantic models, instruction landing in the last user message (and absent from the system prompt), ask-level override, callable schemas, retry re-injection, and bedrock/zai mapper no-op cases. --- ag2/__init__.py | 3 +- ag2/agent.py | 29 +++++ ag2/response/__init__.py | 3 +- ag2/response/prompted.py | 91 +++++++++++++- ag2/response/proto.py | 1 + test/agent/test_response_schema.py | 79 +++++++++++- test/config/bedrock/test_response_schema.py | 3 +- test/config/zai/test_response_schema.py | 5 +- test/response/test_prompted.py | 126 +++++++++++++++++++- 9 files changed, 330 insertions(+), 10 deletions(-) diff --git a/ag2/__init__.py b/ag2/__init__.py index 3c3f7a43c96a..59e17f0bdb23 100755 --- a/ag2/__init__.py +++ b/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", diff --git a/ag2/agent.py b/ag2/agent.py index 0b83f72a18ea..33e02aa27577 100644 --- a/ag2/agent.py +++ b/ag2/agent.py @@ -45,6 +45,7 @@ Input, ModelRequest, ModelResponse, + TextInput, ToolCallsEvent, ToolResultsEvent, UsageEvent, @@ -1311,6 +1312,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 @@ -1555,6 +1558,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], diff --git a/ag2/response/__init__.py b/ag2/response/__init__.py index f4bd14c16525..8a6f0c6956fb 100644 --- a/ag2/response/__init__.py +++ b/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", diff --git a/ag2/response/prompted.py b/ag2/response/prompted.py index 06d3d4bcfcaf..9b3d6ea00af8 100644 --- a/ag2/response/prompted.py +++ b/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,15 +101,25 @@ 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, @@ -101,3 +127,60 @@ async def validate( 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") diff --git a/ag2/response/proto.py b/ag2/response/proto.py index cb254274d322..117f0ed191c6 100644 --- a/ag2/response/proto.py +++ b/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( diff --git a/test/agent/test_response_schema.py b/test/agent/test_response_schema.py index 1a1533db76f8..72d82975e677 100644 --- a/test/agent/test_response_schema.py +++ b/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 diff --git a/test/config/bedrock/test_response_schema.py b/test/config/bedrock/test_response_schema.py index c5dfb3b4d7d2..a728a02e9eca 100644 --- a/test/config/bedrock/test_response_schema.py +++ b/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 diff --git a/test/config/zai/test_response_schema.py b/test/config/zai/test_response_schema.py index d03d2f8da228..79887b8bbe20 100644 --- a/test/config/zai/test_response_schema.py +++ b/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 diff --git a/test/response/test_prompted.py b/test/response/test_prompted.py index ce4a5c7dc87c..d08db5bc7f5e 100644 --- a/test/response/test_prompted.py +++ b/test/response/test_prompted.py @@ -9,7 +9,7 @@ import pytest from pydantic import BaseModel -from ag2 import PromptedSchema, ResponseSchema, response_schema +from ag2 import LastMessagePromptedSchema, PromptedSchema, ResponseSchema, response_schema from ag2.response import ResponseProto @@ -186,3 +186,127 @@ def parse_int(content: str) -> int: with pytest.raises(Exception): await schema.validate("not a number", context) + + +class TestInjectToLastMessage: + def test_prompted_schema_with_inject_to_last_message(self) -> None: + schema = PromptedSchema(int, inject_to="last_message") + + assert schema.system_prompt is None + assert schema.last_message_prompt is not None + assert '"type": "integer"' in schema.last_message_prompt + assert schema.json_schema is None + + def test_prompted_schema_default_injects_to_system(self) -> None: + schema = PromptedSchema(int) + + assert schema.system_prompt is not None + assert schema.last_message_prompt is None + + def test_last_message_prompted_schema_from_type(self) -> None: + schema = LastMessagePromptedSchema(int) + + assert schema.system_prompt is None + assert schema.last_message_prompt is not None + assert '"type": "integer"' in schema.last_message_prompt + assert schema.json_schema is None + + def test_last_message_prompted_schema_from_dataclass(self) -> None: + @dataclass + class User: + name: str + age: int + + schema = LastMessagePromptedSchema(User) + + assert schema.name == "User" + assert schema.system_prompt is None + assert '"name"' in schema.last_message_prompt + assert '"age"' in schema.last_message_prompt + + def test_last_message_prompted_schema_from_pydantic_model(self) -> None: + class Item(BaseModel): + title: str + price: float + + schema = LastMessagePromptedSchema(Item) + + assert schema.name == "Item" + assert schema.system_prompt is None + assert '"title"' in schema.last_message_prompt + assert '"price"' in schema.last_message_prompt + + def test_last_message_prompted_schema_wraps_response_schema(self) -> None: + inner = ResponseSchema(int, name="MyInt") + schema = LastMessagePromptedSchema(inner) + + assert schema.name == "MyInt" + assert schema.system_prompt is None + assert '"type": "integer"' in schema.last_message_prompt + + def test_last_message_prompted_schema_custom_template(self) -> None: + template = "Return JSON: ```{schema}```" + schema = LastMessagePromptedSchema(int, prompt_template=template) + + assert schema.last_message_prompt is not None + assert schema.last_message_prompt.startswith("Return JSON: ```") + assert schema.system_prompt is None + + def test_last_message_prompted_schema_no_inner_schema(self) -> None: + class NoSchemaProto(ResponseProto[str]): + def __init__(self) -> None: + self.name = "test" + self.description = None + self.json_schema = None + + async def validate(self, response, context, provider=None): + return response + + schema = LastMessagePromptedSchema(NoSchemaProto()) + assert schema.last_message_prompt is None + assert schema.system_prompt is None + + def test_is_subclass_of_prompted_schema(self) -> None: + schema = LastMessagePromptedSchema(int) + assert isinstance(schema, PromptedSchema) + + +@pytest.mark.asyncio +class TestLastMessageValidation: + async def test_validates_primitive(self) -> None: + schema = LastMessagePromptedSchema(int) + context = AsyncMock() + + result = await schema.validate('{"data": 42}', context) + assert result == 42 + + async def test_validates_dataclass(self) -> None: + @dataclass + class Point: + x: float + y: float + + schema = LastMessagePromptedSchema(Point) + context = AsyncMock() + + result = await schema.validate('{"x": 1.5, "y": 2.5}', context) + assert result == Point(x=1.5, y=2.5) + + async def test_delegates_to_inner_response_schema(self) -> None: + inner = ResponseSchema(int, name="MyInt") + schema = LastMessagePromptedSchema(inner) + context = AsyncMock() + + result = await schema.validate('{"data": 42}', context) + assert result == 42 + + async def test_delegates_to_inner_callable_schema(self) -> None: + @response_schema + def double(content: str) -> int: + return int(content) * 2 + + schema = LastMessagePromptedSchema(double) + context = AsyncMock() + + result = await schema.validate("21", context) + assert result == 42 From be1e8bd3bd7a4436bcb5f89699ea9f4e3353127a Mon Sep 17 00:00:00 2001 From: Rubikoid Date: Fri, 10 Jul 2026 01:44:23 +0300 Subject: [PATCH 2/2] docs(response): Update docs for LastMessagePromptedSchema --- website/docs/user-guide/structured_output.mdx | 31 ++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/website/docs/user-guide/structured_output.mdx b/website/docs/user-guide/structured_output.mdx index e1b961365b0f..e78a62998d41 100644 --- a/website/docs/user-guide/structured_output.mdx +++ b/website/docs/user-guide/structured_output.mdx @@ -18,7 +18,7 @@ Every turn returns an [`AgentReply`](/docs/user-guide/agents). Two surfaces matt If the model’s output cannot be parsed or fails validation, `content()` raises an error from the underlying parser (for example Pydantic’s validation errors). You can pass `retries` to automatically [re-ask the model](#validation-retries) on failure. -With the default **OpenAI** client, when the schema exposes a JSON Schema to the API, the client sends a structured `response_format` so the model is guided to emit JSON matching that schema. [`PromptedSchema`](#promptedschema-models-without-native-structured-output) is the escape hatch when the provider does not support that mechanism: the schema is injected into the system prompt instead, and `content()` still runs the same way afterward. +With the default **OpenAI** client, when the schema exposes a JSON Schema to the API, the client sends a structured `response_format` so the model is guided to emit JSON matching that schema. [`PromptedSchema`](#promptedschema-models-without-native-structured-output) is the escape hatch when the provider does not support that mechanism: the schema is injected into the prompt — the system prompt by default, or the last user message via [`LastMessagePromptedSchema`](#injecting-into-the-last-user-message) — and `content()` still runs the same way afterward. ## When to use which tool @@ -26,6 +26,7 @@ With the default **OpenAI** client, when the schema exposes a JSON Schema to the - Use **`ResponseSchema`** when you want a clear **`name`** and **`description`** in the API payload so the model knows the role of the structured payload. - Use **`@response_schema`** when you need **custom parsing**, normalization, or extra steps after JSON is read. - Use **`PromptedSchema`** when your **model or endpoint does not support** native structured output. +- Use **`LastMessagePromptedSchema`** when the model pays more attention to the end of the conversation, when the endpoint has **no system-prompt slot**, or when you want the schema instruction adjacent to the request — e.g. in chained `.ask()` calls where the last user turn carries the actual task. ## Quick start @@ -426,6 +427,34 @@ PromptedSchema( ) ``` +### Injecting into the last user message + +By default `PromptedSchema` puts the JSON-schema instruction in the **system prompt**. Some models attend more strongly to the end of the conversation. Pass `inject_to="last_message"` to append the instruction to the final user turn instead: + +```python linenums="1" hl_lines="6" +from ag2 import Agent, PromptedSchema + +agent = Agent( + "assistant", + config=config, + response_schema=PromptedSchema(int, inject_to="last_message"), +) +``` + +`LastMessagePromptedSchema` is a convenience subclass equivalent to `PromptedSchema(inner, inject_to="last_message")`: + +```python linenums="1" hl_lines="1 6" +from ag2 import Agent, LastMessagePromptedSchema + +agent = Agent( + "assistant", + config=config, + response_schema=LastMessagePromptedSchema(int), +) +``` + +This is especially useful in **chained `.ask()` calls**: each turn's schema instruction lands right next to the question the model must answer, not in a system prompt that may be far back in the conversation. The inner schema's `validate` method is unchanged — only the injection location differs. + --- ## Override schema per request