adds ttl to cache control
This commit is contained in:
@@ -1004,9 +1004,9 @@ impl<W: UiWriter> Agent<W> {
|
|||||||
/// Convert cache config string to CacheControl enum
|
/// Convert cache config string to CacheControl enum
|
||||||
fn parse_cache_control(cache_config: &str) -> Option<CacheControl> {
|
fn parse_cache_control(cache_config: &str) -> Option<CacheControl> {
|
||||||
match cache_config {
|
match cache_config {
|
||||||
"ephemeral" => Some(CacheControl::Ephemeral),
|
"ephemeral" => Some(CacheControl::ephemeral()),
|
||||||
"5minute" => Some(CacheControl::FiveMinute),
|
"5minute" => Some(CacheControl::five_minute()),
|
||||||
"1hour" => Some(CacheControl::OneHour),
|
"1hour" => Some(CacheControl::one_hour()),
|
||||||
_ => {
|
_ => {
|
||||||
warn!("Invalid cache_config value: '{}'. Valid values are: ephemeral, 5minute, 1hour", cache_config);
|
warn!("Invalid cache_config value: '{}'. Valid values are: ephemeral, 5minute, 1hour", cache_config);
|
||||||
None
|
None
|
||||||
|
|||||||
@@ -172,15 +172,9 @@ impl AnthropicProvider {
|
|||||||
builder
|
builder
|
||||||
}
|
}
|
||||||
|
|
||||||
fn convert_cache_control(cache_control: &crate::CacheControl) -> AnthropicCacheControl {
|
fn convert_cache_control(cache_control: &crate::CacheControl) -> crate::CacheControl {
|
||||||
let cache_type = match cache_control {
|
// Anthropic uses the same format, so just clone it
|
||||||
crate::CacheControl::Ephemeral => "ephemeral",
|
cache_control.clone()
|
||||||
crate::CacheControl::FiveMinute => "5minute",
|
|
||||||
crate::CacheControl::OneHour => "1hour",
|
|
||||||
};
|
|
||||||
AnthropicCacheControl {
|
|
||||||
cache_type: cache_type.to_string(),
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn convert_tools(&self, tools: &[Tool]) -> Vec<AnthropicTool> {
|
fn convert_tools(&self, tools: &[Tool]) -> Vec<AnthropicTool> {
|
||||||
@@ -723,12 +717,6 @@ struct AnthropicMessage {
|
|||||||
content: Vec<AnthropicContent>,
|
content: Vec<AnthropicContent>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize)]
|
|
||||||
struct AnthropicCacheControl {
|
|
||||||
#[serde(rename = "type")]
|
|
||||||
cache_type: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize)]
|
#[derive(Debug, Serialize, Deserialize)]
|
||||||
#[serde(tag = "type")]
|
#[serde(tag = "type")]
|
||||||
enum AnthropicContent {
|
enum AnthropicContent {
|
||||||
@@ -736,7 +724,7 @@ enum AnthropicContent {
|
|||||||
Text {
|
Text {
|
||||||
text: String,
|
text: String,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
cache_control: Option<AnthropicCacheControl>,
|
cache_control: Option<crate::CacheControl>,
|
||||||
},
|
},
|
||||||
#[serde(rename = "tool_use")]
|
#[serde(rename = "tool_use")]
|
||||||
ToolUse {
|
ToolUse {
|
||||||
@@ -916,7 +904,7 @@ mod tests {
|
|||||||
let messages_with = vec![Message::with_cache_control(
|
let messages_with = vec![Message::with_cache_control(
|
||||||
MessageRole::User,
|
MessageRole::User,
|
||||||
"Hello".to_string(),
|
"Hello".to_string(),
|
||||||
crate::CacheControl::Ephemeral,
|
crate::CacheControl::ephemeral(),
|
||||||
)];
|
)];
|
||||||
let (_, anthropic_messages_with) = provider.convert_messages(&messages_with).unwrap();
|
let (_, anthropic_messages_with) = provider.convert_messages(&messages_with).unwrap();
|
||||||
let json_with = serde_json::to_string(&anthropic_messages_with).unwrap();
|
let json_with = serde_json::to_string(&anthropic_messages_with).unwrap();
|
||||||
|
|||||||
@@ -238,13 +238,9 @@ impl DatabricksProvider {
|
|||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn convert_cache_control(cache_control: &crate::CacheControl) -> DatabricksCacheControl {
|
fn convert_cache_control(cache_control: &crate::CacheControl) -> crate::CacheControl {
|
||||||
let cache_type = match cache_control {
|
// Databricks uses the same format, so just clone it
|
||||||
crate::CacheControl::Ephemeral => "ephemeral",
|
cache_control.clone()
|
||||||
crate::CacheControl::FiveMinute => "5minute",
|
|
||||||
crate::CacheControl::OneHour => "1hour",
|
|
||||||
};
|
|
||||||
DatabricksCacheControl { cache_type: cache_type.to_string() }
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn convert_messages(&self, messages: &[Message]) -> Result<Vec<DatabricksMessage>> {
|
fn convert_messages(&self, messages: &[Message]) -> Result<Vec<DatabricksMessage>> {
|
||||||
@@ -1104,12 +1100,6 @@ struct DatabricksFunction {
|
|||||||
parameters: serde_json::Value,
|
parameters: serde_json::Value,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize)]
|
|
||||||
struct DatabricksCacheControl {
|
|
||||||
#[serde(rename = "type")]
|
|
||||||
cache_type: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize)]
|
#[derive(Debug, Serialize, Deserialize)]
|
||||||
#[serde(untagged)]
|
#[serde(untagged)]
|
||||||
enum DatabricksContent {
|
enum DatabricksContent {
|
||||||
@@ -1118,7 +1108,7 @@ enum DatabricksContent {
|
|||||||
content_type: String,
|
content_type: String,
|
||||||
text: String,
|
text: String,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
cache_control: Option<DatabricksCacheControl>,
|
cache_control: Option<crate::CacheControl>,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1345,7 +1335,7 @@ mod tests {
|
|||||||
let messages_with = vec![Message::with_cache_control(
|
let messages_with = vec![Message::with_cache_control(
|
||||||
MessageRole::User,
|
MessageRole::User,
|
||||||
"Hello".to_string(),
|
"Hello".to_string(),
|
||||||
crate::CacheControl::Ephemeral,
|
crate::CacheControl::ephemeral(),
|
||||||
)];
|
)];
|
||||||
let databricks_messages_with = provider.convert_messages(&messages_with).unwrap();
|
let databricks_messages_with = provider.convert_messages(&messages_with).unwrap();
|
||||||
let json_with = serde_json::to_string(&databricks_messages_with).unwrap();
|
let json_with = serde_json::to_string(&databricks_messages_with).unwrap();
|
||||||
|
|||||||
@@ -38,13 +38,31 @@ pub struct CompletionRequest {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CacheControl {
|
||||||
|
#[serde(rename = "type")]
|
||||||
|
pub cache_type: CacheType,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub ttl: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
#[serde(rename_all = "lowercase")]
|
#[serde(rename_all = "lowercase")]
|
||||||
pub enum CacheControl {
|
pub enum CacheType {
|
||||||
Ephemeral,
|
Ephemeral,
|
||||||
#[serde(rename = "5minute")]
|
}
|
||||||
FiveMinute,
|
|
||||||
#[serde(rename = "1hour")]
|
impl CacheControl {
|
||||||
OneHour,
|
pub fn ephemeral() -> Self {
|
||||||
|
Self { cache_type: CacheType::Ephemeral, ttl: None }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn five_minute() -> Self {
|
||||||
|
Self { cache_type: CacheType::Ephemeral, ttl: Some("5m".to_string()) }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn one_hour() -> Self {
|
||||||
|
Self { cache_type: CacheType::Ephemeral, ttl: Some("1h".to_string()) }
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
@@ -220,7 +238,7 @@ mod tests {
|
|||||||
let msg = Message::with_cache_control(
|
let msg = Message::with_cache_control(
|
||||||
MessageRole::User,
|
MessageRole::User,
|
||||||
"Hello".to_string(),
|
"Hello".to_string(),
|
||||||
CacheControl::Ephemeral,
|
CacheControl::ephemeral(),
|
||||||
);
|
);
|
||||||
let json = serde_json::to_string(&msg).unwrap();
|
let json = serde_json::to_string(&msg).unwrap();
|
||||||
|
|
||||||
@@ -229,7 +247,39 @@ mod tests {
|
|||||||
"JSON should contain 'cache_control' field when configured");
|
"JSON should contain 'cache_control' field when configured");
|
||||||
assert!(json.contains("ephemeral"),
|
assert!(json.contains("ephemeral"),
|
||||||
"JSON should contain 'ephemeral' value");
|
"JSON should contain 'ephemeral' value");
|
||||||
|
assert!(json.contains("\"type\":"),
|
||||||
|
"JSON should contain 'type' field in cache_control");
|
||||||
assert!(!json.contains("null"),
|
assert!(!json.contains("null"),
|
||||||
"JSON should not contain null values");
|
"JSON should not contain null values");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_cache_control_five_minute_serialization() {
|
||||||
|
let msg = Message::with_cache_control(
|
||||||
|
MessageRole::User,
|
||||||
|
"Hello".to_string(),
|
||||||
|
CacheControl::five_minute(),
|
||||||
|
);
|
||||||
|
let json = serde_json::to_string(&msg).unwrap();
|
||||||
|
|
||||||
|
println!("Message JSON with 5-minute cache_control: {}", json);
|
||||||
|
assert!(json.contains("cache_control"), "JSON should contain 'cache_control' field");
|
||||||
|
assert!(json.contains("ephemeral"), "JSON should contain 'ephemeral' type");
|
||||||
|
assert!(json.contains("\"ttl\":\"5m\""), "JSON should contain ttl field with 5m value");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_cache_control_one_hour_serialization() {
|
||||||
|
let msg = Message::with_cache_control(
|
||||||
|
MessageRole::User,
|
||||||
|
"Hello".to_string(),
|
||||||
|
CacheControl::one_hour(),
|
||||||
|
);
|
||||||
|
let json = serde_json::to_string(&msg).unwrap();
|
||||||
|
|
||||||
|
println!("Message JSON with 1-hour cache_control: {}", json);
|
||||||
|
assert!(json.contains("cache_control"), "JSON should contain 'cache_control' field");
|
||||||
|
assert!(json.contains("ephemeral"), "JSON should contain 'ephemeral' type");
|
||||||
|
assert!(json.contains("\"ttl\":\"1h\""), "JSON should contain ttl field with 1h value");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user