Skip to content

Commit 9f10595

Browse files
feat(models): add SystemContentBlock support for provider-agnostic caching (#1112)
* feat(model): support prompt caching via SystemContentBlock * fix: concat text blocks for system_prompt * remove litellm and openai changes for now * integ tests * linting * linting * fix test * add test cases
1 parent 5981d36 commit 9f10595

File tree

15 files changed

+495
-50
lines changed

15 files changed

+495
-50
lines changed

src/strands/agent/agent.py

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@
5858
from ..tools.watcher import ToolWatcher
5959
from ..types._events import AgentResultEvent, InitEventLoopEvent, ModelStreamChunkEvent, ToolInterruptEvent, TypedEvent
6060
from ..types.agent import AgentInput
61-
from ..types.content import ContentBlock, Message, Messages
61+
from ..types.content import ContentBlock, Message, Messages, SystemContentBlock
6262
from ..types.exceptions import ContextWindowOverflowException
6363
from ..types.interrupt import InterruptResponseContent
6464
from ..types.tools import ToolResult, ToolUse
@@ -217,7 +217,7 @@ def __init__(
217217
model: Union[Model, str, None] = None,
218218
messages: Optional[Messages] = None,
219219
tools: Optional[list[Union[str, dict[str, str], "ToolProvider", Any]]] = None,
220-
system_prompt: Optional[str] = None,
220+
system_prompt: Optional[str | list[SystemContentBlock]] = None,
221221
structured_output_model: Optional[Type[BaseModel]] = None,
222222
callback_handler: Optional[
223223
Union[Callable[..., Any], _DefaultCallbackHandlerSentinel]
@@ -254,6 +254,7 @@ def __init__(
254254
255255
If provided, only these tools will be available. If None, all tools will be available.
256256
system_prompt: System prompt to guide model behavior.
257+
Can be a string or a list of SystemContentBlock objects for advanced features like caching.
257258
If None, the model will behave according to its default settings.
258259
structured_output_model: Pydantic model type(s) for structured output.
259260
When specified, all agent calls will attempt to return structured output of this type.
@@ -288,7 +289,8 @@ def __init__(
288289
"""
289290
self.model = BedrockModel() if not model else BedrockModel(model_id=model) if isinstance(model, str) else model
290291
self.messages = messages if messages is not None else []
291-
self.system_prompt = system_prompt
292+
# initializing self.system_prompt for backwards compatibility
293+
self.system_prompt, self._system_prompt_content = self._initialize_system_prompt(system_prompt)
292294
self._default_structured_output_model = structured_output_model
293295
self.agent_id = _identifier.validate(agent_id or _DEFAULT_AGENT_ID, _identifier.Identifier.AGENT)
294296
self.name = name or _DEFAULT_AGENT_NAME
@@ -981,6 +983,30 @@ def _filter_tool_parameters_for_recording(self, tool_name: str, input_params: di
981983
properties = tool_spec["inputSchema"]["json"]["properties"]
982984
return {k: v for k, v in input_params.items() if k in properties}
983985

986+
def _initialize_system_prompt(
987+
self, system_prompt: str | list[SystemContentBlock] | None
988+
) -> tuple[str | None, list[SystemContentBlock] | None]:
989+
"""Initialize system prompt fields from constructor input.
990+
991+
Maintains backwards compatibility by keeping system_prompt as str when string input
992+
provided, avoiding breaking existing consumers.
993+
994+
Maps system_prompt input to both string and content block representations:
995+
- If string: system_prompt=string, _system_prompt_content=[{text: string}]
996+
- If list with text elements: system_prompt=concatenated_text, _system_prompt_content=list
997+
- If list without text elements: system_prompt=None, _system_prompt_content=list
998+
- If None: system_prompt=None, _system_prompt_content=None
999+
"""
1000+
if isinstance(system_prompt, str):
1001+
return system_prompt, [{"text": system_prompt}]
1002+
elif isinstance(system_prompt, list):
1003+
# Concatenate all text elements for backwards compatibility, None if no text found
1004+
text_parts = [block["text"] for block in system_prompt if "text" in block]
1005+
system_prompt_str = "\n".join(text_parts) if text_parts else None
1006+
return system_prompt_str, system_prompt
1007+
else:
1008+
return None, None
1009+
9841010
def _append_message(self, message: Message) -> None:
9851011
"""Appends a message to the agent's list of messages and invokes the callbacks for the MessageCreatedEvent."""
9861012
self.messages.append(message)

src/strands/event_loop/event_loop.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -335,7 +335,12 @@ async def _handle_model_execution(
335335
tool_specs = agent.tool_registry.get_all_tool_specs()
336336
try:
337337
async for event in stream_messages(
338-
agent.model, agent.system_prompt, agent.messages, tool_specs, structured_output_context.tool_choice
338+
agent.model,
339+
agent.system_prompt,
340+
agent.messages,
341+
tool_specs,
342+
system_prompt_content=agent._system_prompt_content,
343+
tool_choice=structured_output_context.tool_choice,
339344
):
340345
yield event
341346

src/strands/event_loop/streaming.py

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
TypedEvent,
2323
)
2424
from ..types.citations import CitationsContentBlock
25-
from ..types.content import ContentBlock, Message, Messages
25+
from ..types.content import ContentBlock, Message, Messages, SystemContentBlock
2626
from ..types.streaming import (
2727
ContentBlockDeltaEvent,
2828
ContentBlockStart,
@@ -418,16 +418,22 @@ async def stream_messages(
418418
system_prompt: Optional[str],
419419
messages: Messages,
420420
tool_specs: list[ToolSpec],
421+
*,
421422
tool_choice: Optional[Any] = None,
423+
system_prompt_content: Optional[list[SystemContentBlock]] = None,
424+
**kwargs: Any,
422425
) -> AsyncGenerator[TypedEvent, None]:
423426
"""Streams messages to the model and processes the response.
424427
425428
Args:
426429
model: Model provider.
427-
system_prompt: The system prompt to send.
430+
system_prompt: The system prompt string, used for backwards compatibility with models that expect it.
428431
messages: List of messages to send.
429432
tool_specs: The list of tool specs.
430433
tool_choice: Optional tool choice constraint for forcing specific tool usage.
434+
system_prompt_content: The authoritative system prompt content blocks that always contains the
435+
system prompt data.
436+
**kwargs: Additional keyword arguments for future extensibility.
431437
432438
Yields:
433439
The reason for stopping, the final message, and the usage metrics
@@ -436,7 +442,14 @@ async def stream_messages(
436442

437443
messages = _normalize_messages(messages)
438444
start_time = time.time()
439-
chunks = model.stream(messages, tool_specs if tool_specs else None, system_prompt, tool_choice=tool_choice)
445+
446+
chunks = model.stream(
447+
messages,
448+
tool_specs if tool_specs else None,
449+
system_prompt,
450+
tool_choice=tool_choice,
451+
system_prompt_content=system_prompt_content,
452+
)
440453

441454
async for event in process_stream(chunks, start_time):
442455
yield event

src/strands/models/bedrock.py

Lines changed: 25 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from ..event_loop import streaming
2121
from ..tools import convert_pydantic_to_tool_spec
2222
from ..tools._tool_helpers import noop_tool
23-
from ..types.content import ContentBlock, Messages
23+
from ..types.content import ContentBlock, Messages, SystemContentBlock
2424
from ..types.exceptions import (
2525
ContextWindowOverflowException,
2626
ModelThrottledException,
@@ -187,11 +187,11 @@ def get_config(self) -> BedrockConfig:
187187
"""
188188
return self.config
189189

190-
def format_request(
190+
def _format_request(
191191
self,
192192
messages: Messages,
193193
tool_specs: Optional[list[ToolSpec]] = None,
194-
system_prompt: Optional[str] = None,
194+
system_prompt_content: Optional[list[SystemContentBlock]] = None,
195195
tool_choice: ToolChoice | None = None,
196196
) -> dict[str, Any]:
197197
"""Format a Bedrock converse stream request.
@@ -201,6 +201,7 @@ def format_request(
201201
tool_specs: List of tool specifications to make available to the model.
202202
system_prompt: System prompt to provide context to the model.
203203
tool_choice: Selection strategy for tool invocation.
204+
system_prompt_content: System prompt content blocks to provide context to the model.
204205
205206
Returns:
206207
A Bedrock converse stream request.
@@ -211,13 +212,20 @@ def format_request(
211212
)
212213
if has_tool_content:
213214
tool_specs = [noop_tool.tool_spec]
215+
216+
# Use system_prompt_content directly (copy for mutability)
217+
system_blocks: list[SystemContentBlock] = system_prompt_content.copy() if system_prompt_content else []
218+
# Add cache point if configured (backwards compatibility)
219+
if cache_prompt := self.config.get("cache_prompt"):
220+
warnings.warn(
221+
"cache_prompt is deprecated. Use SystemContentBlock with cachePoint instead.", UserWarning, stacklevel=3
222+
)
223+
system_blocks.append({"cachePoint": {"type": cache_prompt}})
224+
214225
return {
215226
"modelId": self.config["model_id"],
216227
"messages": self._format_bedrock_messages(messages),
217-
"system": [
218-
*([{"text": system_prompt}] if system_prompt else []),
219-
*([{"cachePoint": {"type": self.config["cache_prompt"]}}] if self.config.get("cache_prompt") else []),
220-
],
228+
"system": system_blocks,
221229
**(
222230
{
223231
"toolConfig": {
@@ -590,6 +598,7 @@ async def stream(
590598
system_prompt: Optional[str] = None,
591599
*,
592600
tool_choice: ToolChoice | None = None,
601+
system_prompt_content: Optional[list[SystemContentBlock]] = None,
593602
**kwargs: Any,
594603
) -> AsyncGenerator[StreamEvent, None]:
595604
"""Stream conversation with the Bedrock model.
@@ -602,6 +611,7 @@ async def stream(
602611
tool_specs: List of tool specifications to make available to the model.
603612
system_prompt: System prompt to provide context to the model.
604613
tool_choice: Selection strategy for tool invocation.
614+
system_prompt_content: System prompt content blocks to provide context to the model.
605615
**kwargs: Additional keyword arguments for future extensibility.
606616
607617
Yields:
@@ -620,7 +630,11 @@ def callback(event: Optional[StreamEvent] = None) -> None:
620630
loop = asyncio.get_event_loop()
621631
queue: asyncio.Queue[Optional[StreamEvent]] = asyncio.Queue()
622632

623-
thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt, tool_choice)
633+
# Handle backward compatibility: if system_prompt is provided but system_prompt_content is None
634+
if system_prompt and system_prompt_content is None:
635+
system_prompt_content = [{"text": system_prompt}]
636+
637+
thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt_content, tool_choice)
624638
task = asyncio.create_task(thread)
625639

626640
while True:
@@ -637,7 +651,7 @@ def _stream(
637651
callback: Callable[..., None],
638652
messages: Messages,
639653
tool_specs: Optional[list[ToolSpec]] = None,
640-
system_prompt: Optional[str] = None,
654+
system_prompt_content: Optional[list[SystemContentBlock]] = None,
641655
tool_choice: ToolChoice | None = None,
642656
) -> None:
643657
"""Stream conversation with the Bedrock model.
@@ -649,7 +663,7 @@ def _stream(
649663
callback: Function to send events to the main thread.
650664
messages: List of message objects to be processed by the model.
651665
tool_specs: List of tool specifications to make available to the model.
652-
system_prompt: System prompt to provide context to the model.
666+
system_prompt_content: System prompt content blocks to provide context to the model.
653667
tool_choice: Selection strategy for tool invocation.
654668
655669
Raises:
@@ -658,7 +672,7 @@ def _stream(
658672
"""
659673
try:
660674
logger.debug("formatting request")
661-
request = self.format_request(messages, tool_specs, system_prompt, tool_choice)
675+
request = self._format_request(messages, tool_specs, system_prompt_content, tool_choice)
662676
logger.debug("request=<%s>", request)
663677

664678
logger.debug("invoking model")

src/strands/models/model.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
from pydantic import BaseModel
88

9-
from ..types.content import Messages
9+
from ..types.content import Messages, SystemContentBlock
1010
from ..types.streaming import StreamEvent
1111
from ..types.tools import ToolChoice, ToolSpec
1212

@@ -72,6 +72,7 @@ def stream(
7272
system_prompt: Optional[str] = None,
7373
*,
7474
tool_choice: ToolChoice | None = None,
75+
system_prompt_content: list[SystemContentBlock] | None = None,
7576
**kwargs: Any,
7677
) -> AsyncIterable[StreamEvent]:
7778
"""Stream conversation with the model.
@@ -87,6 +88,7 @@ def stream(
8788
tool_specs: List of tool specifications to make available to the model.
8889
system_prompt: System prompt to provide context to the model.
8990
tool_choice: Selection strategy for tool invocation.
91+
system_prompt_content: System prompt content blocks for advanced features like caching.
9092
**kwargs: Additional keyword arguments for future extensibility.
9193
9294
Yields:

src/strands/types/content.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,11 +103,11 @@ class SystemContentBlock(TypedDict, total=False):
103103
"""Contains configurations for instructions to provide the model for how to handle input.
104104
105105
Attributes:
106-
guardContent: A content block to assess with the guardrail.
106+
cachePoint: A cache point configuration to optimize conversation history.
107107
text: A system prompt for the model.
108108
"""
109109

110-
guardContent: GuardContent
110+
cachePoint: CachePoint
111111
text: str
112112

113113

tests/fixtures/mocked_model_provider.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,9 @@ async def stream(
5858
tool_specs: Optional[list[ToolSpec]] = None,
5959
system_prompt: Optional[str] = None,
6060
tool_choice: Optional[Any] = None,
61+
*,
62+
system_prompt_content=None,
63+
**kwargs: Any,
6164
) -> AsyncGenerator[Any, None]:
6265
events = self.map_agent_message_to_events(self.agent_responses[self.index])
6366
for event in events:

0 commit comments

Comments
 (0)