-
Notifications
You must be signed in to change notification settings - Fork 454
feat(models): add SystemContentBlock support for provider-agnostic caching #1112
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
2d4bc80
e9fe30b
d135ca5
0eb2f2b
a8e36a6
cef64d8
c2ac0b5
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -20,7 +20,7 @@ | |
| from ..event_loop import streaming | ||
| from ..tools import convert_pydantic_to_tool_spec | ||
| from ..tools._tool_helpers import noop_tool | ||
| from ..types.content import ContentBlock, Messages | ||
| from ..types.content import ContentBlock, Messages, SystemContentBlock | ||
| from ..types.exceptions import ( | ||
| ContextWindowOverflowException, | ||
| ModelThrottledException, | ||
|
|
@@ -187,11 +187,11 @@ def get_config(self) -> BedrockConfig: | |
| """ | ||
| return self.config | ||
|
|
||
| def format_request( | ||
| def _format_request( | ||
| self, | ||
| messages: Messages, | ||
| tool_specs: Optional[list[ToolSpec]] = None, | ||
| system_prompt: Optional[str] = None, | ||
| system_prompt_content: Optional[list[SystemContentBlock]] = None, | ||
| tool_choice: ToolChoice | None = None, | ||
| ) -> dict[str, Any]: | ||
| """Format a Bedrock converse stream request. | ||
|
|
@@ -201,6 +201,7 @@ def format_request( | |
| tool_specs: List of tool specifications to make available to the model. | ||
| system_prompt: System prompt to provide context to the model. | ||
| tool_choice: Selection strategy for tool invocation. | ||
| system_prompt_content: System prompt content blocks to provide context to the model. | ||
|
|
||
| Returns: | ||
| A Bedrock converse stream request. | ||
|
|
@@ -211,13 +212,20 @@ def format_request( | |
| ) | ||
| if has_tool_content: | ||
| tool_specs = [noop_tool.tool_spec] | ||
|
|
||
| # Use system_prompt_content directly (copy for mutability) | ||
| system_blocks: list[SystemContentBlock] = system_prompt_content.copy() if system_prompt_content else [] | ||
| # Add cache point if configured (backwards compatibility) | ||
| if cache_prompt := self.config.get("cache_prompt"): | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What happens if someone already specified cache points in their There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. yes, but it won't it still works. And since it does, I wanted to avoid making assumptions about ordering. Like if cache_prompt is present and the last item is CachePoint then we dedupe. It doesn't break and we emit a warning. Adding a test to demonstrate this |
||
| warnings.warn( | ||
| "cache_prompt is deprecated. Use SystemContentBlock with cachePoint instead.", UserWarning, stacklevel=3 | ||
| ) | ||
| system_blocks.append({"cachePoint": {"type": cache_prompt}}) | ||
|
|
||
| return { | ||
| "modelId": self.config["model_id"], | ||
| "messages": self._format_bedrock_messages(messages), | ||
| "system": [ | ||
| *([{"text": system_prompt}] if system_prompt else []), | ||
| *([{"cachePoint": {"type": self.config["cache_prompt"]}}] if self.config.get("cache_prompt") else []), | ||
| ], | ||
| "system": system_blocks, | ||
| **( | ||
| { | ||
| "toolConfig": { | ||
|
|
@@ -590,6 +598,7 @@ async def stream( | |
| system_prompt: Optional[str] = None, | ||
| *, | ||
| tool_choice: ToolChoice | None = None, | ||
| system_prompt_content: Optional[list[SystemContentBlock]] = None, | ||
dbschmigelski marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| **kwargs: Any, | ||
| ) -> AsyncGenerator[StreamEvent, None]: | ||
| """Stream conversation with the Bedrock model. | ||
|
|
@@ -602,6 +611,7 @@ async def stream( | |
| tool_specs: List of tool specifications to make available to the model. | ||
| system_prompt: System prompt to provide context to the model. | ||
| tool_choice: Selection strategy for tool invocation. | ||
| system_prompt_content: System prompt content blocks to provide context to the model. | ||
| **kwargs: Additional keyword arguments for future extensibility. | ||
|
|
||
| Yields: | ||
|
|
@@ -620,7 +630,7 @@ def callback(event: Optional[StreamEvent] = None) -> None: | |
| loop = asyncio.get_event_loop() | ||
| queue: asyncio.Queue[Optional[StreamEvent]] = asyncio.Queue() | ||
|
|
||
| thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt, tool_choice) | ||
| thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt_content, tool_choice) | ||
| task = asyncio.create_task(thread) | ||
|
|
||
| while True: | ||
|
|
@@ -637,7 +647,7 @@ def _stream( | |
| callback: Callable[..., None], | ||
| messages: Messages, | ||
| tool_specs: Optional[list[ToolSpec]] = None, | ||
| system_prompt: Optional[str] = None, | ||
| system_prompt_content: Optional[list[SystemContentBlock]] = None, | ||
| tool_choice: ToolChoice | None = None, | ||
| ) -> None: | ||
| """Stream conversation with the Bedrock model. | ||
|
|
@@ -649,7 +659,7 @@ def _stream( | |
| callback: Function to send events to the main thread. | ||
| messages: List of message objects to be processed by the model. | ||
| tool_specs: List of tool specifications to make available to the model. | ||
| system_prompt: System prompt to provide context to the model. | ||
| system_prompt_content: System prompt content blocks to provide context to the model. | ||
| tool_choice: Selection strategy for tool invocation. | ||
|
|
||
| Raises: | ||
|
|
@@ -658,7 +668,7 @@ def _stream( | |
| """ | ||
| try: | ||
| logger.debug("formatting request") | ||
| request = self.format_request(messages, tool_specs, system_prompt, tool_choice) | ||
| request = self._format_request(messages, tool_specs, system_prompt_content, tool_choice) | ||
| logger.debug("request=<%s>", request) | ||
|
|
||
| logger.debug("invoking model") | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -6,7 +6,7 @@ | |
|
|
||
| from pydantic import BaseModel | ||
|
|
||
| from ..types.content import Messages | ||
| from ..types.content import Messages, SystemContentBlock | ||
| from ..types.streaming import StreamEvent | ||
| from ..types.tools import ToolChoice, ToolSpec | ||
|
|
||
|
|
@@ -72,6 +72,7 @@ def stream( | |
| system_prompt: Optional[str] = None, | ||
| *, | ||
| tool_choice: ToolChoice | None = None, | ||
| system_prompt_content: list[SystemContentBlock] | None = None, | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Noting that I do not like this. But we are breaking if we did something like Even if we only passed in str if the user passed in str, which was considered, the typing is still breaking for mypy users. |
||
| **kwargs: Any, | ||
| ) -> AsyncIterable[StreamEvent]: | ||
| """Stream conversation with the model. | ||
|
|
@@ -87,6 +88,7 @@ def stream( | |
| tool_specs: List of tool specifications to make available to the model. | ||
| system_prompt: System prompt to provide context to the model. | ||
| tool_choice: Selection strategy for tool invocation. | ||
| system_prompt_content: System prompt content blocks for advanced features like caching. | ||
| **kwargs: Additional keyword arguments for future extensibility. | ||
|
|
||
| Yields: | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What if instead we just did
None, system_prompt? Is there a reason we have to force generate aself.system_promptwhen a customer passes in a list? They are opting in to a new feature and so I wouldn't expect it to break them.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
so when we call model.stream we still need to pass system_prompt. Even if a user opted in to passing this additional content, that should not break a model implementation.
So my thinking is that we will need to have the system_prompt_content -> system_prompt transformation regardless. So it is better to keep it and expose it as agent.system_prompt so it is minimally breaking in the event a user is doing "agent.system_prompt" or if a session_manager, hook, model, etc is itself making a call to agent.system_prompt
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah gotcha. So rather than requiring them to make changes in multiple places at once, they could do so gradually or be content with the transformation.