feat(response): inject json schema into the last message instead of system prompt - #3070
feat(response): inject json schema into the last message instead of system prompt#3070Rubikoid wants to merge 2 commits into
Conversation
…ser message 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.
|
|
github-actions
Bot
added
area:docs
genisis0x
left a comment
genisis0x
left a comment
There was a problem hiding this comment.
The motivation makes sense — for providers without native structured output, and especially without a usable system prompt, putting the schema instruction near the generation boundary is a real improvement, and the inject_to / LastMessagePromptedSchema opt-in keeps the default behavior intact. _inject_last_message_prompt is also nicely done: it rebuilds the list and the ModelRequest copy instead of mutating, and no-ops safely. Multi-provider tests (bedrock/zai/agent/prompted) are a good touch.
Two things I'd look at, both around the "last message" semantics:
-
It targets the last
ModelRequest, not literally the last event. In a plain.askwith no tools that's the same thing, so it works. But once there are tool round-trips, the most recentModelRequest(the user turn) can sit several events back — behindToolResultsEvents and priorModelResponses — so the schema instruction is no longer adjacent to where the model generates. Since the whole point is "the model pays more attention to the end of the conversation," it might be worth appending a fresh trailing turn at generation time rather than augmenting the last user request, so the instruction is always last. At minimum, a note in the docstring that "last message" means "last user request" would avoid surprise. -
Silent no-op when there's no
ModelRequest._inject_last_message_promptreturnsmessagesunchanged if it finds noModelRequest, which means the schema guidance is dropped entirely for that call (rather than, say, falling back to the system prompt). That's an easy corner to hit with a history that starts from a non-request event. Consider falling back to system injection (or emitting a warning) in that branch, and adding a test for the no-ModelRequestcase to lock whichever behavior you choose.
Minor: appending the instruction as a TextInput onto the existing user ModelRequest.parts merges it into the user's own message. That's probably fine, but a one-line comment on why it's merged into the last request vs. sent as a distinct message would help the next reader.
3 participants
Why are these changes needed?
I ran into the following scenario: I had a warmed-up session with a model that was supposed to respond with a structured response on .ask.
My provider didn't support "full, via API" schema passing, and
PromptedSchemainserted the schema into the system prompt - as a result, the model couldn't make sense of its own prior history and produced JSON poorly (only from the second or third attempt).After switching to LastMessagePromptedSchema, the model started producing JSON on the first or second try.
I'm not sure this is the most elegant or correct implementation — if it's not, let's refine it together.
Related issue number
Did't find corresponding.
Checks
AI assistance