mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-09-09 08:03:50 +00:00
fix: clean ai memory and cache bedrock prompts (#8847)
* fix: avoid persisting system prompts in ai memory Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * fix: keep ai memory cleanup write-side only Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * feat: add bedrock prompt caching for claude Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * test: add bedrock memory regression Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * fix: gate bedrock prompt caching by model id Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * docs: link bedrock caching allowlist source Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
centdix
co-authored by
Claude Opus 4.5
Claude Opus 4.5
parent
fc49a8fed6
commit
b1778272fc
@@ -131,6 +131,7 @@ async fn create_bedrock_client(
|
||||
fn build_tool_config_from_request(
|
||||
tools: Option<&[OpenAIToolDef]>,
|
||||
tool_choice: Option<&serde_json::Value>,
|
||||
enable_prompt_caching: bool,
|
||||
) -> Result<Option<aws_sdk_bedrockruntime::types::ToolConfiguration>> {
|
||||
if let Some(tools) = tools {
|
||||
let tool_defs: Vec<ToolDef> = tools
|
||||
@@ -163,7 +164,7 @@ fn build_tool_config_from_request(
|
||||
.map(|tc| tc == "required" || tc.as_str() == Some("required"))
|
||||
.unwrap_or(false);
|
||||
|
||||
build_tool_config(Some(&tool_defs), force_tool_use)
|
||||
build_tool_config(Some(&tool_defs), force_tool_use, enable_prompt_caching)
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
@@ -365,8 +366,13 @@ pub async fn handle_bedrock_sdk_streaming(
|
||||
.await?;
|
||||
|
||||
// Convert messages using shared conversion
|
||||
let enable_prompt_caching =
|
||||
windmill_common::ai_bedrock::bedrock_model_supports_prompt_caching(model);
|
||||
let (bedrock_messages, system_prompts) =
|
||||
windmill_common::ai_bedrock::openai_messages_to_bedrock(&openai_req.messages)?;
|
||||
windmill_common::ai_bedrock::openai_messages_to_bedrock(
|
||||
&openai_req.messages,
|
||||
enable_prompt_caching,
|
||||
)?;
|
||||
|
||||
// Build inference configuration
|
||||
let inference_config = windmill_common::ai_bedrock::create_inference_config(
|
||||
@@ -378,6 +384,7 @@ pub async fn handle_bedrock_sdk_streaming(
|
||||
let tool_config = build_tool_config_from_request(
|
||||
openai_req.tools.as_deref(),
|
||||
openai_req.tool_choice.as_ref(),
|
||||
enable_prompt_caching,
|
||||
)?;
|
||||
|
||||
// Build the SDK request
|
||||
@@ -625,8 +632,13 @@ pub async fn handle_bedrock_sdk_non_streaming(
|
||||
.await?;
|
||||
|
||||
// Convert messages using shared conversion
|
||||
let enable_prompt_caching =
|
||||
windmill_common::ai_bedrock::bedrock_model_supports_prompt_caching(model);
|
||||
let (bedrock_messages, system_prompts) =
|
||||
windmill_common::ai_bedrock::openai_messages_to_bedrock(&openai_req.messages)?;
|
||||
windmill_common::ai_bedrock::openai_messages_to_bedrock(
|
||||
&openai_req.messages,
|
||||
enable_prompt_caching,
|
||||
)?;
|
||||
|
||||
// Build inference configuration
|
||||
let inference_config = windmill_common::ai_bedrock::create_inference_config(
|
||||
@@ -638,6 +650,7 @@ pub async fn handle_bedrock_sdk_non_streaming(
|
||||
let tool_config = build_tool_config_from_request(
|
||||
openai_req.tools.as_deref(),
|
||||
openai_req.tool_choice.as_ref(),
|
||||
enable_prompt_caching,
|
||||
)?;
|
||||
|
||||
// Build the SDK request (non-streaming)
|
||||
|
||||
@@ -94,6 +94,86 @@ pub async fn check_env_credentials() -> BedrockCredentialsCheck {
|
||||
/// Constants for commonly used strings to avoid allocations
|
||||
pub const FUNCTION_TYPE: &str = "function";
|
||||
|
||||
// AWS documents Bedrock prompt-caching support as a model allowlist rather than a
|
||||
// capability exposed by the model metadata APIs:
|
||||
// https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html
|
||||
const BEDROCK_PROMPT_CACHING_SUPPORTED_MODEL_IDS: &[&str] = &[
|
||||
"anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
"anthropic.claude-opus-4-1-20250805-v1:0",
|
||||
"anthropic.claude-opus-4-20250514-v1:0",
|
||||
"anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"anthropic.claude-sonnet-4-20250514-v1:0",
|
||||
"anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"anthropic.claude-3-5-haiku-20241022-v1:0",
|
||||
"anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
];
|
||||
|
||||
fn build_default_cache_point() -> aws_sdk_bedrockruntime::types::CachePointBlock {
|
||||
aws_sdk_bedrockruntime::types::CachePointBlock::builder()
|
||||
.r#type(aws_sdk_bedrockruntime::types::CachePointType::Default)
|
||||
.build()
|
||||
.expect("cache point type is required")
|
||||
}
|
||||
|
||||
fn normalize_bedrock_model_id(model: &str) -> String {
|
||||
let model = model
|
||||
.rsplit('/')
|
||||
.next()
|
||||
.unwrap_or(model)
|
||||
.to_ascii_lowercase();
|
||||
|
||||
for prefix in ["global.", "us.", "eu.", "apac."] {
|
||||
if let Some(normalized_model) = model.strip_prefix(prefix) {
|
||||
return normalized_model.to_string();
|
||||
}
|
||||
}
|
||||
|
||||
model
|
||||
}
|
||||
|
||||
pub fn bedrock_model_supports_prompt_caching(model: &str) -> bool {
|
||||
let normalized_model = normalize_bedrock_model_id(model);
|
||||
BEDROCK_PROMPT_CACHING_SUPPORTED_MODEL_IDS.contains(&normalized_model.as_str())
|
||||
}
|
||||
|
||||
fn append_cache_point_to_system_prompts(system_prompts: &mut Vec<SystemContentBlock>) {
|
||||
if system_prompts.is_empty()
|
||||
|| matches!(
|
||||
system_prompts.last(),
|
||||
Some(SystemContentBlock::CachePoint(_))
|
||||
)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
system_prompts.push(SystemContentBlock::CachePoint(build_default_cache_point()));
|
||||
}
|
||||
|
||||
fn append_cache_point_to_last_message(messages: &mut [Message]) -> Result<(), Error> {
|
||||
let Some(last_message) = messages.last_mut() else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
if matches!(
|
||||
last_message.content().last(),
|
||||
Some(ContentBlock::CachePoint(_))
|
||||
) {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mut content = last_message.content().to_vec();
|
||||
content.push(ContentBlock::CachePoint(build_default_cache_point()));
|
||||
|
||||
*last_message = Message::builder()
|
||||
.role(last_message.role().clone())
|
||||
.set_content(Some(content))
|
||||
.build()
|
||||
.map_err(|e| Error::internal_err(format!("Failed to append cache point: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BearerTokenProvider {
|
||||
token: String,
|
||||
@@ -277,6 +357,7 @@ pub fn json_to_document(value: serde_json::Value) -> aws_smithy_types::Document
|
||||
/// Tuple of (conversation_messages, system_prompts)
|
||||
pub fn openai_messages_to_bedrock(
|
||||
messages: &[OpenAIMessage],
|
||||
enable_prompt_caching: bool,
|
||||
) -> Result<(Vec<Message>, Vec<SystemContentBlock>), Error> {
|
||||
let mut bedrock_messages = Vec::new();
|
||||
let mut system_prompts = Vec::new();
|
||||
@@ -334,6 +415,11 @@ pub fn openai_messages_to_bedrock(
|
||||
bedrock_messages.push(tool_result_message);
|
||||
}
|
||||
|
||||
if enable_prompt_caching {
|
||||
append_cache_point_to_system_prompts(&mut system_prompts);
|
||||
append_cache_point_to_last_message(&mut bedrock_messages)?;
|
||||
}
|
||||
|
||||
Ok((bedrock_messages, system_prompts))
|
||||
}
|
||||
|
||||
@@ -703,9 +789,15 @@ pub fn streaming_tool_calls_to_openai(tool_calls: Vec<StreamingToolCall>) -> Vec
|
||||
pub fn build_tool_config(
|
||||
tools: Option<&[ToolDef]>,
|
||||
force_tool_use: bool,
|
||||
enable_prompt_caching: bool,
|
||||
) -> Result<Option<aws_sdk_bedrockruntime::types::ToolConfiguration>, Error> {
|
||||
if let Some(tools) = tools {
|
||||
let bedrock_tools = openai_tools_to_bedrock(tools)?;
|
||||
let mut bedrock_tools = openai_tools_to_bedrock(tools)?;
|
||||
|
||||
if enable_prompt_caching && !matches!(bedrock_tools.last(), Some(Tool::CachePoint(_))) {
|
||||
bedrock_tools.push(Tool::CachePoint(build_default_cache_point()));
|
||||
}
|
||||
|
||||
let mut tool_config_builder = aws_sdk_bedrockruntime::types::ToolConfiguration::builder()
|
||||
.set_tools(Some(bedrock_tools));
|
||||
|
||||
@@ -724,3 +816,112 @@ pub fn build_tool_config(
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::value::RawValue;
|
||||
|
||||
fn text_message(role: &str, content: &str) -> OpenAIMessage {
|
||||
OpenAIMessage {
|
||||
role: role.to_string(),
|
||||
content: Some(OpenAIContent::Text(content.to_string())),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn test_tool() -> ToolDef {
|
||||
ToolDef {
|
||||
r#type: FUNCTION_TYPE.to_string(),
|
||||
function: crate::ai_types::ToolDefFunction {
|
||||
name: "test_tool".to_string(),
|
||||
description: Some("A test tool".to_string()),
|
||||
parameters: RawValue::from_string(
|
||||
r#"{"type":"object","properties":{},"additionalProperties":false}"#.to_string(),
|
||||
)
|
||||
.expect("valid raw json"),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_messages_to_bedrock_adds_cache_points_when_enabled() {
|
||||
let messages = vec![
|
||||
text_message("system", "Reply concisely"),
|
||||
text_message("user", "Tell me a joke"),
|
||||
];
|
||||
|
||||
let (bedrock_messages, system_prompts) =
|
||||
openai_messages_to_bedrock(&messages, true).expect("bedrock conversion succeeds");
|
||||
|
||||
assert!(matches!(
|
||||
system_prompts.last(),
|
||||
Some(SystemContentBlock::CachePoint(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
bedrock_messages
|
||||
.last()
|
||||
.and_then(|message| message.content().last()),
|
||||
Some(ContentBlock::CachePoint(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_messages_to_bedrock_skips_cache_points_when_disabled() {
|
||||
let messages = vec![
|
||||
text_message("system", "Reply concisely"),
|
||||
text_message("user", "Tell me a joke"),
|
||||
];
|
||||
|
||||
let (bedrock_messages, system_prompts) =
|
||||
openai_messages_to_bedrock(&messages, false).expect("bedrock conversion succeeds");
|
||||
|
||||
assert!(!matches!(
|
||||
system_prompts.last(),
|
||||
Some(SystemContentBlock::CachePoint(_))
|
||||
));
|
||||
assert!(!matches!(
|
||||
bedrock_messages
|
||||
.last()
|
||||
.and_then(|message| message.content().last()),
|
||||
Some(ContentBlock::CachePoint(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_tool_config_adds_cache_point_when_enabled() {
|
||||
let tools = vec![test_tool()];
|
||||
let tool_config =
|
||||
build_tool_config(Some(&tools), false, true).expect("tool config succeeds");
|
||||
|
||||
assert!(matches!(
|
||||
tool_config
|
||||
.as_ref()
|
||||
.and_then(|config| config.tools().last()),
|
||||
Some(Tool::CachePoint(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bedrock_prompt_caching_supports_documented_claude_model_ids() {
|
||||
assert!(bedrock_model_supports_prompt_caching(
|
||||
"anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
));
|
||||
assert!(bedrock_model_supports_prompt_caching(
|
||||
"global.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
));
|
||||
assert!(bedrock_model_supports_prompt_caching(
|
||||
"arn:aws:bedrock:us-east-1::inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bedrock_prompt_caching_rejects_unsupported_or_opaque_model_ids() {
|
||||
assert!(!bedrock_model_supports_prompt_caching(
|
||||
"anthropic.claude-3-haiku-20240307-v1:0"
|
||||
));
|
||||
assert!(!bedrock_model_supports_prompt_caching(
|
||||
"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/my-profile"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -18,10 +18,11 @@ use windmill_common::{client::AuthedClient, error::Error};
|
||||
|
||||
// Re-export from shared module for use by other parts of the worker
|
||||
use windmill_common::ai_bedrock::{
|
||||
bedrock_stream_event_is_block_stop, bedrock_stream_event_to_text,
|
||||
bedrock_stream_event_to_tool_delta, bedrock_stream_event_to_tool_start, build_tool_config,
|
||||
create_inference_config, format_bedrock_error, openai_messages_to_bedrock,
|
||||
streaming_tool_calls_to_openai, StreamingToolCall,
|
||||
bedrock_model_supports_prompt_caching, bedrock_stream_event_is_block_stop,
|
||||
bedrock_stream_event_to_text, bedrock_stream_event_to_tool_delta,
|
||||
bedrock_stream_event_to_tool_start, build_tool_config, create_inference_config,
|
||||
format_bedrock_error, openai_messages_to_bedrock, streaming_tool_calls_to_openai,
|
||||
StreamingToolCall,
|
||||
};
|
||||
pub use windmill_common::ai_bedrock::{check_env_credentials, BedrockClient};
|
||||
|
||||
@@ -71,13 +72,19 @@ impl BedrockQueryBuilder {
|
||||
let prepared_messages = prepare_messages_for_api(messages, client, workspace_id).await?;
|
||||
|
||||
// Convert messages to Bedrock format (separates system prompts)
|
||||
let (bedrock_messages, system_prompts) = openai_messages_to_bedrock(&prepared_messages)?;
|
||||
let enable_prompt_caching = bedrock_model_supports_prompt_caching(model);
|
||||
let (bedrock_messages, system_prompts) =
|
||||
openai_messages_to_bedrock(&prepared_messages, enable_prompt_caching)?;
|
||||
|
||||
// Build inference configuration using shared helper
|
||||
let inference_config = create_inference_config(temperature, max_tokens.map(|t| t as i32));
|
||||
|
||||
// Build tool configuration with optional ToolChoice
|
||||
let tool_config = build_tool_config(tools, structured_output_tool_name.is_some())?;
|
||||
let tool_config = build_tool_config(
|
||||
tools,
|
||||
structured_output_tool_name.is_some(),
|
||||
enable_prompt_caching,
|
||||
)?;
|
||||
|
||||
self.execute_converse_stream(
|
||||
&bedrock_client,
|
||||
|
||||
@@ -92,6 +92,38 @@ lazy_static::lazy_static! {
|
||||
const DEFAULT_MAX_AGENT_ITERATIONS: usize = 10;
|
||||
const HARD_MAX_AGENT_ITERATIONS: usize = 1000;
|
||||
|
||||
fn strip_system_messages(messages: &[OpenAIMessage]) -> Vec<OpenAIMessage> {
|
||||
messages
|
||||
.iter()
|
||||
.filter(|message| message.role != "system")
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn strip_leading_tool_messages(messages: Vec<OpenAIMessage>) -> Vec<OpenAIMessage> {
|
||||
match messages.iter().position(|message| message.role != "tool") {
|
||||
Some(first_non_tool_index) => messages.into_iter().skip(first_non_tool_index).collect(),
|
||||
None => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn prepare_auto_memory_messages_for_request(
|
||||
loaded_messages: &[OpenAIMessage],
|
||||
context_length: usize,
|
||||
) -> Vec<OpenAIMessage> {
|
||||
let start_idx = loaded_messages.len().saturating_sub(context_length);
|
||||
strip_leading_tool_messages(loaded_messages[start_idx..].to_vec())
|
||||
}
|
||||
|
||||
fn prepare_auto_memory_messages_for_persistence(
|
||||
all_messages: &[OpenAIMessage],
|
||||
context_length: usize,
|
||||
) -> Vec<OpenAIMessage> {
|
||||
let non_system_messages = strip_system_messages(all_messages);
|
||||
let start_idx = non_system_messages.len().saturating_sub(context_length);
|
||||
non_system_messages[start_idx..].to_vec()
|
||||
}
|
||||
|
||||
fn find_module_by_id(
|
||||
modules: &Vec<FlowModule>,
|
||||
target_id: &str,
|
||||
@@ -654,18 +686,10 @@ pub async fn run_agent(
|
||||
// Read messages from memory
|
||||
match read_from_memory(db, &job.workspace_id, memory_id, step_id).await {
|
||||
Ok(Some(loaded_messages)) => {
|
||||
// Take the last n messages
|
||||
let start_idx =
|
||||
loaded_messages.len().saturating_sub(*context_length);
|
||||
let mut messages_to_load = loaded_messages[start_idx..].to_vec();
|
||||
let first_non_tool_message_index =
|
||||
messages_to_load.iter().position(|m| m.role != "tool");
|
||||
|
||||
// Remove the first messages if their role is "tool" to avoid OpenAI API error
|
||||
if let Some(index) = first_non_tool_message_index {
|
||||
messages_to_load = messages_to_load[index..].to_vec();
|
||||
}
|
||||
|
||||
let messages_to_load = prepare_auto_memory_messages_for_request(
|
||||
&loaded_messages,
|
||||
*context_length,
|
||||
);
|
||||
messages.extend(messages_to_load);
|
||||
}
|
||||
Ok(None) => {}
|
||||
@@ -729,9 +753,7 @@ pub async fn run_agent(
|
||||
.unwrap_or(false);
|
||||
|
||||
if has_message && has_attachments {
|
||||
let mut parts = vec![ContentPart::Text {
|
||||
text: args.user_message.clone().unwrap(),
|
||||
}];
|
||||
let mut parts = vec![ContentPart::Text { text: args.user_message.clone().unwrap() }];
|
||||
for attachment in args.user_attachments.as_ref().unwrap() {
|
||||
if !attachment.s3.is_empty() {
|
||||
parts.push(ContentPart::S3Object { s3_object: attachment.clone() });
|
||||
@@ -1306,9 +1328,10 @@ pub async fn run_agent(
|
||||
final_messages.iter().map(|m| m.message.clone()).collect();
|
||||
|
||||
if !all_messages.is_empty() {
|
||||
// Keep only the last n messages
|
||||
let start_idx = all_messages.len().saturating_sub(*context_length);
|
||||
let messages_to_persist = all_messages[start_idx..].to_vec();
|
||||
let messages_to_persist = prepare_auto_memory_messages_for_persistence(
|
||||
&all_messages,
|
||||
*context_length,
|
||||
);
|
||||
|
||||
if let Some(memory_id) = memory_id {
|
||||
if let Err(e) = write_to_memory(
|
||||
@@ -1349,6 +1372,86 @@ pub async fn run_agent(
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn text_message(role: &str, content: &str) -> OpenAIMessage {
|
||||
OpenAIMessage {
|
||||
role: role.to_string(),
|
||||
content: Some(OpenAIContent::Text(content.to_string())),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_memory_request_preserves_messages_within_context_window() {
|
||||
let loaded_messages = vec![
|
||||
text_message("system", "instructions-a"),
|
||||
text_message("user", "first-user"),
|
||||
text_message("assistant", "first-assistant"),
|
||||
text_message("system", "instructions-b"),
|
||||
text_message("user", "second-user"),
|
||||
text_message("assistant", "second-assistant"),
|
||||
];
|
||||
|
||||
let prepared = prepare_auto_memory_messages_for_request(&loaded_messages, 3);
|
||||
let roles: Vec<&str> = prepared
|
||||
.iter()
|
||||
.map(|message| message.role.as_str())
|
||||
.collect();
|
||||
let contents: Vec<&str> = prepared
|
||||
.iter()
|
||||
.map(|message| match message.content.as_ref() {
|
||||
Some(OpenAIContent::Text(text)) => text.as_str(),
|
||||
_ => "",
|
||||
})
|
||||
.collect();
|
||||
|
||||
assert_eq!(roles, vec!["system", "user", "assistant"]);
|
||||
assert_eq!(
|
||||
contents,
|
||||
vec!["instructions-b", "second-user", "second-assistant"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_memory_request_drops_leading_tool_messages() {
|
||||
let loaded_messages = vec![
|
||||
text_message("tool", "stale-tool-result"),
|
||||
text_message("user", "hello"),
|
||||
text_message("assistant", "hi"),
|
||||
];
|
||||
|
||||
let prepared = prepare_auto_memory_messages_for_request(&loaded_messages, 10);
|
||||
let roles: Vec<&str> = prepared
|
||||
.iter()
|
||||
.map(|message| message.role.as_str())
|
||||
.collect();
|
||||
|
||||
assert_eq!(roles, vec!["user", "assistant"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_memory_persistence_excludes_system_messages() {
|
||||
let all_messages = vec![
|
||||
text_message("system", "instructions"),
|
||||
text_message("user", "hello"),
|
||||
text_message("assistant", "hi"),
|
||||
text_message("system", "duplicate-instructions"),
|
||||
text_message("user", "follow-up"),
|
||||
];
|
||||
|
||||
let persisted = prepare_auto_memory_messages_for_persistence(&all_messages, 10);
|
||||
let roles: Vec<&str> = persisted
|
||||
.iter()
|
||||
.map(|message| message.role.as_str())
|
||||
.collect();
|
||||
|
||||
assert_eq!(roles, vec!["user", "assistant", "user"]);
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle credentials check mode - check credentials without making API calls
|
||||
async fn handle_credentials_check(provider: &ProviderWithResource) -> Result<Box<RawValue>, Error> {
|
||||
let result = match &provider.kind {
|
||||
|
||||
@@ -4,17 +4,56 @@ Memory tests for AI agents.
|
||||
Tests that AI agents correctly handle conversation memory/history.
|
||||
"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from .conftest import AIAgentTestClient, create_ai_agent_flow
|
||||
from .providers import ALL_PROVIDERS, get_provider_ids
|
||||
from .providers import (
|
||||
ALL_PROVIDERS,
|
||||
BEDROCK_API_KEY,
|
||||
BEDROCK_ENV,
|
||||
BEDROCK_IAM,
|
||||
BEDROCK_IAM_SESSION,
|
||||
get_provider_ids,
|
||||
)
|
||||
|
||||
|
||||
class TestMemory:
|
||||
"""Test AI agent memory functionality."""
|
||||
|
||||
@staticmethod
|
||||
def _pick_bedrock_provider() -> dict[str, Any]:
|
||||
if os.environ.get("BEDROCK_API_KEY"):
|
||||
return BEDROCK_API_KEY
|
||||
|
||||
if (
|
||||
os.environ.get("BEDROCK_IAM_ACCESS_KEY_ID")
|
||||
and os.environ.get("BEDROCK_IAM_SECRET_ACCESS_KEY")
|
||||
):
|
||||
return BEDROCK_IAM
|
||||
|
||||
if (
|
||||
os.environ.get("BEDROCK_SESSION_ACCESS_KEY_ID")
|
||||
and os.environ.get("BEDROCK_SESSION_SECRET_ACCESS_KEY")
|
||||
and os.environ.get("BEDROCK_SESSION_TOKEN")
|
||||
):
|
||||
return BEDROCK_IAM_SESSION
|
||||
|
||||
aws_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID")
|
||||
aws_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
|
||||
aws_session_token = os.environ.get("AWS_SESSION_TOKEN")
|
||||
|
||||
if aws_access_key_id and aws_secret_access_key:
|
||||
if aws_access_key_id.startswith("ASIA") and not aws_session_token:
|
||||
pytest.skip(
|
||||
"AWS_SESSION_TOKEN required for Bedrock env fallback when AWS_ACCESS_KEY_ID is temporary (ASIA...)"
|
||||
)
|
||||
return BEDROCK_ENV
|
||||
|
||||
pytest.skip("No Bedrock credentials available for the memory regression test")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"provider_config",
|
||||
ALL_PROVIDERS,
|
||||
@@ -65,6 +104,61 @@ class TestMemory:
|
||||
|
||||
print(f"Memory test passed for {provider_config['name']}")
|
||||
|
||||
def test_memory_keeps_single_system_prompt_across_turns(
|
||||
self,
|
||||
client: AIAgentTestClient,
|
||||
setup_providers,
|
||||
):
|
||||
"""
|
||||
Test that auto memory persists conversation turns without persisting
|
||||
duplicate system prompts across repeated runs.
|
||||
"""
|
||||
provider_config = self._pick_bedrock_provider()
|
||||
|
||||
system_prompt = "You are a helpful assistant. Keep responses short."
|
||||
flow_value = create_ai_agent_flow(
|
||||
provider_input_transform=provider_config["input_transform"],
|
||||
system_prompt=system_prompt,
|
||||
context_length=6,
|
||||
)
|
||||
|
||||
memory_id = str(uuid.uuid4())
|
||||
turns = [
|
||||
"Turn one: say hi.",
|
||||
"Turn two: say hi again.",
|
||||
"Turn three: say hi one more time.",
|
||||
]
|
||||
|
||||
result = None
|
||||
for turn in turns:
|
||||
result = client.run_preview_flow(
|
||||
flow_value=flow_value,
|
||||
args={"user_message": turn},
|
||||
memory_id=memory_id,
|
||||
)
|
||||
assert result is not None
|
||||
assert "error" not in result, f"Turn failed for {provider_config['name']}: {result}"
|
||||
|
||||
assert result is not None
|
||||
messages = result.get("messages", [])
|
||||
system_messages = [msg for msg in messages if msg.get("role") == "system"]
|
||||
|
||||
assert len(system_messages) == 1, (
|
||||
f"Expected exactly one system message for {provider_config['name']}, got: {messages}"
|
||||
)
|
||||
assert system_messages[0].get("content") == system_prompt
|
||||
|
||||
user_messages = [
|
||||
msg.get("content")
|
||||
for msg in messages
|
||||
if msg.get("role") == "user"
|
||||
]
|
||||
assert user_messages == turns, (
|
||||
f"Expected all prior user turns to remain in memory for {provider_config['name']}: {messages}"
|
||||
)
|
||||
|
||||
print(f"System prompt dedupe test passed for {provider_config['name']}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
|
||||
Reference in New Issue
Block a user