Repository navigation
feat(agent): mark Bedrock Converse prompts for caching behind a flag [PC-5190] - #1142
Merged
Merged
Conversation
tudormatei1
force-pushed
the
feat/bedrock-prompt-caching
branch
4 times, most recently
from
October 14, 2026 13:32
50c8a4d to
8c31d27
Compare
tudormatei1
changed the title
tudormatei1
marked this pull request as ready for review
October 15, 2026 10:56
There was a problem hiding this comment.
🟡 Changes recommended
Legacy Bedrock clients bypass the BYO exclusion and receive cache markers when caching is enabled.
1 open finding
What changed in this PR
Adds opt-in Bedrock prompt caching to agent LLM calls, with provider eligibility checks and exclusions for non-reusable calls.
Changes:
- Threads
prompt_cachingfrom agent configuration into LLM nodes. - Adds Claude/Nova cache markers and BYO eligibility checks.
- Adds regression tests and bumps the version to
0.18.40.
| File | Description |
|---|---|
uv.lock |
Updates package version. |
tests/chat/test_bedrock_payload_handler.py |
Tests provider and ownership eligibility. |
tests/agent/react/test_llm_node.py |
Tests cache markers and excluded calls. |
tests/agent/react/test_create_agent.py |
Tests configuration propagation. |
src/uipath_langchain/chat/handlers/bedrock.py |
Adds Bedrock caching eligibility. |
src/uipath_langchain/chat/handlers/base.py |
Adds default caching hook. |
src/uipath_langchain/agent/react/types.py |
Adds opt-in configuration field. |
src/uipath_langchain/agent/react/llm_node.py |
Applies caching to eligible calls. |
src/uipath_langchain/agent/react/agent.py |
Forwards caching configuration. |
pyproject.toml |
Bumps version to 0.18.40. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+79
to
+81
| settings = getattr(model, "client_settings", None) | ||
| if settings is None: | ||
| return False |
| } | ||
|
|
||
|
|
||
| def _is_byo_model(model: Any) -> bool: |
Contributor
There was a problem hiding this comment.
Is checking the subscription not sufficient here? Why are we checking multiple fields in order?
Contributor
Author
There was a problem hiding this comment.
It is. Simplified in 153edc0: the lookup now passes byo_connection_id, and BYO means modelSubscriptionType != "UiPathOwned". The gateway always sets byomDetails together with the subscription type, so that check never changed the result.
tudormatei1
force-pushed
the
feat/bedrock-prompt-caching
branch
2 times, most recently
from
November 4, 2026 10:03
153edc0 to
e53ed5b
Compare
Low-code and conversational agents never send prompt cache markers, so their Claude calls on Bedrock almost never hit the cache (the LLM Gateway reports 1.3% for agentsruntimefallback and 0% for agentsplaygroundfallback). With EnableBedrockPromptCaching on, each LLM turn passes cache_control to bind_tools, and langchain-aws places cachePoint blocks after the tools, the system prompt and the last message. LLMCall carries a required cache_prompt: the normal turn of every node caches, while the forced-extraction retry and the AutoToolChoice finish call don't, since their prefix differs from the run's turns and a write there is never read. The Bedrock Converse handler only returns cache_control for Claude and Nova, the models that accept checkpoints (the same check as langchain-aws's own caching middleware), and never for BYO connections. Other providers return nothing.
Discovery can resolve a BYO-only model name without a connection id. The client then sets base_model_id from the BYO deployment's model, and since the gate read base_model_id as the model to check, a BYO Claude still got cache markers. UiPath-owned models never set base_model_id, so it now counts as a BYO signal next to byo_connection_id, and the model check reads model_id. The handler gate gets its own tests (Claude, us. inference profiles, Nova, non-caching families, BYO on the real UiPathChatBedrockConverse). The node tests now count checkpoints across the whole request and check their content, and pass a mocked bedrock_client so building the model no longer probes AWS for credentials.
The gate inferred BYO from byo_connection_id or base_model_id, but the client only sets base_model_id for BYO additions whose customerModel parses as provider.model. A BYO replacement is listed under the UiPath model name with no customerModel, and agents select it by name, so it still got cache markers; so did additions backed by an inference profile ARN. The handler now reads the model's discovery entry through its client settings (cached, no extra I/O) and treats it as BYO when byomDetails is set or modelSubscriptionType isn't UiPathOwned, the same signals the client uses to prefer UiPath-owned entries. A failed lookup turns caching off. The tests build models through get_chat_model with only discovery patched. Version bumped to 0.18.36; main already released 0.18.35.
The LLM node now only says whether a call's prefix is reused (cache_prompt). The Bedrock Converse handler reads EnableBedrockPromptCaching and decides whether markers apply, so a second provider is one handler override with its own flag. The BYO check moves to the handlers' base module since it only reads discovery.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The caching switch now arrives the same way as strict_mode and parallel_tool_calls: AgentGraphConfig.prompt_caching, create_agent, create_llm_node, the node. The caller reads the flag; uipath_langchain no longer touches the registry for it, and the Bedrock handler only decides whether the model takes cache markers. The per-call field is renamed to LLMCall.cacheable and defaults to True, so only the extraction retry and the finish call mention it.
The handler's BYO check goes through get_model_info, which logs an INFO line under uipath run whenever no resource overwrite matches. The model never changes for a node, so the kwargs are now resolved when the node is built instead of on every turn. prompt_caching moves after tool_choice in ConversationalLLMNode so existing positional calls keep binding the same. Adds the tests the review found missing: an explicit BYO connection next to a UiPath-owned entry of the same name, an entry without modelSubscriptionType, and a non-Bedrock model with caching on.
Its fail-closed rule (a failed discovery lookup counts as BYO) is a caching decision, not a general BYO test, and Bedrock Converse is the only caller.
LLM Gateway discovery always sets modelSubscriptionType and byomDetails together, so byomDetails never changed the answer. The explicit connection check now comes from the lookup itself: get_model_info with the model's byo_connection_id returns the BYO entry instead of a UiPath-owned entry with the same name. An entry without a subscription type now counts as BYO, so it gets no cache markers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
tudormatei1
force-pushed
the
feat/bedrock-prompt-caching
branch
from
November 5, 2026 13:25
e53ed5b to
84f1246
Compare
andreiancuta-uipath
approved these changes
Nov 5, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Low-code and conversational agents never send prompt cache markers, so their Claude calls on Bedrock almost never hit the cache. The LLM Gateway team measured 1.3% for agentsruntimefallback and 0% for agentsplaygroundfallback in September, on roughly $120k of spend, and estimates $25-56k a month in savings. Advanced agents aren't affected: deepagents already adds its own caching middleware.
Callers turn it on with
AgentGraphConfig.prompt_caching, whichcreate_agentpasses to the LLM node the same way asstrict_modeandparallel_tool_calls. uipath-agents sets it from theEnableBedrockPromptCachingflag (UiPath/uipath-agents-python#804); this repo doesn't read the flag itself. When it's on, the node resolveshandler.get_prompt_caching_kwargs()once when it's built and_invoke_modelmerges the result into thebind_toolskwargs, and for Claude and Nova the Bedrock Converse handler returnscache_control, which langchain-aws turns intocachePointblocks after the tools, the system prompt and the last message (3 of Bedrock's 4).LLMCall.cacheable(defaultTrue) isFalsefor the forced-extraction retry and the AutoToolChoice finish call, since their prefix differs from the run's turns (thinking stripped, or a reduced tool list) and a cache write there is never read.The Bedrock Converse handler only returns
cache_controlfor Claude and Nova, the model families langchain-aws'sBedrockPromptCachingMiddlewaretreats as cache-capable, and never for BYO deployments. The check readsmodel_id; langchain-aws also looks atbase_model_id, but the client only sets that for BYO entries, which are skipped anyway. It looks the model up in discovery with itsbyo_connection_id(the list is already cached, so no extra calls) and skips caching unlessmodelSubscriptionTypeisUiPathOwned. An entry without a subscription type gets no markers. That covers BYO replacements selected by name, which are listed under the UiPath model name, BYO additions and explicit connections. A failed lookup turns caching off. Custom BYO gateways can rewrite the Converse body, and the rollout targets UiPath-billed traffic. For UiPath-owned agents traffic this changes nothing today, since every Bedrock model the agents Model Hub serves is Claude 4.x/5.x (older ids are routed to Sonnet 4.5), but it keeps a 400 off the table for BYO Llama/Mistral or a future non-Claude model. Other handlers return{}. The node holds no provider logic: it knows whether caching is on and which calls reuse their prefix, and the handler decides how to mark it. Adding another provider later (Claude on Vertex throughAnthropicPayloadHandlertakes the samecache_controlkwarg) is one handler override. The BYO check stays private to the Bedrock handler until a second handler needs it.Things to know:
prompt_cachingdefaults to off, so coded agents and other callers are unchanged. The flag is created for agents on alpha in UiPath/gitops-centralized-cluster#600765. UiPath/uipath-agents-python#804 reads it and sets the field, and has to ship with the agents pin bump to this release, since the field doesn't exist before it.Testing: the handler tests cover the model gate (Claude via a
us.profile, Nova, Llama) and the BYO gate, building the model throughget_chat_modelwith only discovery patched (UiPath-owned, an entry withoutmodelSubscriptionType, a BYO replacement, a BYO addition by ARN, an explicit BYO connection next to a UiPath-owned entry of the same name, a failed lookup). TheLLMNodetests drive a realChatBedrockConversewith a mocked boto client and countcachePointblocks in the request: 3 on each of two turns withprompt_cachingon and 0 with it off, 3 on a conversational turn, 0 on the extraction and finish calls, and nocache_controlfor a non-Bedrock model. Acreate_agenttest checks thatAgentGraphConfig.prompt_cachingreaches the bound model. Each check was confirmed by breaking the code it guards (setting dropped at each step from the config to every node class,cacheableignored, extraction or finish call cached, BYO check removed, lookup failure cached, Nova dropped, subscription ignored, connection id left out of the lookup). Full suite: 4136 passed, 3 skipped. mypy and ruff are clean.Live check on alpha (this branch plus UiPath/uipath-agents-python#804 run locally): a low-code agent on Claude Sonnet 4.5 (Bedrock) with a 3.3k-token prompt, 5 runs per row, starting from a cold cache. Cost is input tokens weighted by Bedrock's cache pricing (reads 0.1x, writes 1.25x), relative to the same runs with the flag off.
Every warm run hit the cache, and the cached token counts were identical from run to run. A new question still re-writes about 400 tokens beyond the question itself, which looks like the tool definitions, so a new input settles at 25% rather than ~11%. LLM latency didn't change at this prompt size. Not covered live: caching across several LLM calls in one run, and BYO models through the gateway.
Version bumped to 0.19.2 (main is on 0.19.1).