litellm-rs 0.6.0

//! Cache key generation utilities
//!
//! This module provides functions for generating consistent cache keys
//! from LLM requests, including normalization and hashing.

use super::key_policy::{CACHE_KEY_SCHEMA_VERSION, stable_digest_value, versioned_key};
use super::types::CacheKey;
use crate::core::models::openai::{ChatCompletionRequest, EmbeddingRequest};
use serde::Serialize;
use serde_json::{Value, json};
use tracing::warn;

/// Prefix for chat completion cache keys
pub const CHAT_KEY_PREFIX: &str = "chat";
/// Prefix for embedding cache keys
pub const EMBEDDING_KEY_PREFIX: &str = "embed";
/// Prefix for completion cache keys
pub const COMPLETION_KEY_PREFIX: &str = "completion";

/// Generate a cache key for a chat completion request
///
/// The key is generated by hashing the normalized request parameters
/// that affect the response (model, messages, temperature, etc.)
pub fn generate_chat_key(request: &ChatCompletionRequest) -> CacheKey {
    generate_chat_key_with_user(request, None)
}

/// Generate a cache key for a chat completion request with optional user ID
pub fn generate_chat_key_with_user(
    request: &ChatCompletionRequest,
    user_id: Option<&str>,
) -> CacheKey {
    let payload = json!({
        "kind": "chat_completion",
        "schema_version": CACHE_KEY_SCHEMA_VERSION,
        "model": &request.model,
        "messages": &request.messages,
        "temperature": request.temperature,
        "max_tokens": request.max_tokens,
        "max_completion_tokens": request.max_completion_tokens,
        "top_p": request.top_p,
        "n": request.n,
        "stop": &request.stop,
        "presence_penalty": request.presence_penalty,
        "frequency_penalty": request.frequency_penalty,
        "logit_bias": &request.logit_bias,
        "functions": &request.functions,
        "function_call": &request.function_call,
        "tools": &request.tools,
        "tool_choice": &request.tool_choice,
        "parallel_tool_calls": request.parallel_tool_calls,
        "response_format": &request.response_format,
        "seed": request.seed,
        "logprobs": request.logprobs,
        "top_logprobs": request.top_logprobs,
        "modalities": &request.modalities,
        "audio": &request.audio,
        "reasoning_effort": &request.reasoning_effort,
        "service_tier": &request.service_tier,
        "prediction": &request.prediction,
        "safety_settings": &request.safety_settings,
        "cache_control": &request.cache_control,
        "extra_body": &request.extra_body,
        "user_id": user_id,
    });

    let digest = stable_digest_value(&payload);
    CacheKey::new(versioned_key(
        CHAT_KEY_PREFIX,
        Some(&request.model),
        &digest,
    ))
}

/// Generate a cache key for an embedding request
pub fn generate_embedding_key(request: &EmbeddingRequest) -> CacheKey {
    generate_embedding_key_with_user(request, None)
}

/// Generate a cache key for an embedding request with optional user ID
pub fn generate_embedding_key_with_user(
    request: &EmbeddingRequest,
    user_id: Option<&str>,
) -> CacheKey {
    let payload = json!({
        "kind": "embedding",
        "schema_version": CACHE_KEY_SCHEMA_VERSION,
        "model": &request.model,
        "input": &request.input,
        "user_id": user_id,
    });

    let digest = stable_digest_value(&payload);
    CacheKey::new(versioned_key(
        EMBEDDING_KEY_PREFIX,
        Some(&request.model),
        &digest,
    ))
}

/// Generate a cache key from arbitrary serializable request
pub fn generate_key_from_json<T: Serialize>(prefix: &str, request: &T) -> CacheKey {
    let value = serde_json::to_value(request).unwrap_or_else(|error| {
        warn!(
            "Failed to serialize request for cache key generation: {}",
            error
        );
        Value::Null
    });
    let digest = stable_digest_value(&value);

    CacheKey::new(versioned_key(prefix, None, &digest))
}

/// Generate a cache key from raw string content
pub fn generate_key_from_content(prefix: &str, content: &str) -> CacheKey {
    let digest = stable_digest_value(&json!({
        "kind": "content",
        "content": content,
    }));

    CacheKey::new(versioned_key(prefix, None, &digest))
}

/// Generate a cache key from multiple parts
pub fn generate_key_from_parts(prefix: &str, parts: &[&str]) -> CacheKey {
    let digest = stable_digest_value(&json!({
        "kind": "parts",
        "parts": parts,
    }));

    CacheKey::new(versioned_key(prefix, None, &digest))
}

// ==================== Helper Functions ====================

/// Normalize a JSON string for consistent hashing
///
/// - Sorts object keys
/// - Removes whitespace
/// - Removes timestamps and request IDs
#[cfg(test)]
fn normalize_json_string(json: &str) -> String {
    super::key_policy::canonical_json_str(json)
}

/// Check if a field should be excluded from cache key generation
#[cfg(test)]
fn is_non_deterministic_field(field: &str) -> bool {
    super::key_policy::is_non_deterministic_field(field)
}

/// Builder for constructing cache keys
pub struct CacheKeyBuilder {
    parts: Vec<String>,
    prefix: String,
}

impl CacheKeyBuilder {
    /// Create a new cache key builder with a prefix
    pub fn new(prefix: impl Into<String>) -> Self {
        Self {
            parts: Vec::new(),
            prefix: prefix.into(),
        }
    }

    /// Add a string part to the key
    pub fn with_part(mut self, part: impl Into<String>) -> Self {
        self.parts.push(part.into());
        self
    }

    /// Add an optional string part to the key
    pub fn add_optional(mut self, part: Option<impl Into<String>>) -> Self {
        if let Some(p) = part {
            self.parts.push(p.into());
        }
        self
    }

    /// Add a numeric part to the key
    pub fn add_num<N: std::fmt::Display>(mut self, num: N) -> Self {
        self.parts.push(num.to_string());
        self
    }

    /// Build the cache key
    pub fn build(self) -> CacheKey {
        let Self { parts, prefix } = self;
        let digest = stable_digest_value(&json!({
            "kind": "builder",
            "parts": parts,
        }));

        CacheKey::new(versioned_key(&prefix, None, &digest))
    }

    /// Build the cache key with explicit parts (no hashing)
    pub fn build_explicit(self) -> CacheKey {
        let key = std::iter::once(self.prefix)
            .chain(self.parts)
            .collect::<Vec<_>>()
            .join(":");

        CacheKey::new(key)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::models::openai::messages::{ChatMessage, MessageContent, MessageRole};
    use crate::core::models::openai::{
        Function, ResponseFormat, Tool, ToolChoice, ToolChoiceFunction, ToolChoiceFunctionSpec,
    };

    // Helper to create a test chat message
    fn create_user_message(content: &str) -> ChatMessage {
        ChatMessage {
            role: MessageRole::User,
            content: Some(MessageContent::Text(content.to_string())),
            name: None,
            function_call: None,
            tool_calls: None,
            tool_call_id: None,
            audio: None,
        }
    }

    // ==================== Chat Key Generation Tests ====================

    #[test]
    fn test_generate_chat_key_basic() {
        let request = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("Hello")],
            ..Default::default()
        };

        let key = generate_chat_key(&request);
        assert!(key.as_str().starts_with("chat:gpt-4:v4:"));
    }

    #[test]
    fn test_generate_chat_key_consistency() {
        let request = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("Hello")],
            temperature: Some(0.7),
            ..Default::default()
        };

        let key1 = generate_chat_key(&request);
        let key2 = generate_chat_key(&request);

        assert_eq!(key1, key2);
    }

    #[test]
    fn test_generate_chat_key_different_messages() {
        let request1 = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("Hello")],
            ..Default::default()
        };

        let request2 = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("Goodbye")],
            ..Default::default()
        };

        let key1 = generate_chat_key(&request1);
        let key2 = generate_chat_key(&request2);

        assert_ne!(key1, key2);
    }

    #[test]
    fn test_generate_chat_key_different_models() {
        let request1 = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![],
            ..Default::default()
        };

        let request2 = ChatCompletionRequest {
            model: "gpt-3.5-turbo".to_string(),
            messages: vec![],
            ..Default::default()
        };

        let key1 = generate_chat_key(&request1);
        let key2 = generate_chat_key(&request2);

        assert_ne!(key1, key2);
    }

    #[test]
    fn test_generate_chat_key_with_user() {
        let request = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![],
            ..Default::default()
        };

        let key1 = generate_chat_key_with_user(&request, Some("user-123"));
        let key2 = generate_chat_key_with_user(&request, Some("user-456"));
        let key3 = generate_chat_key_with_user(&request, None);

        assert_ne!(key1, key2);
        assert_ne!(key1, key3);
    }

    #[test]
    fn test_generate_chat_key_with_parameters() {
        let request1 = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![],
            temperature: Some(0.7),
            ..Default::default()
        };

        let request2 = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![],
            temperature: Some(0.9),
            ..Default::default()
        };

        let key1 = generate_chat_key(&request1);
        let key2 = generate_chat_key(&request2);

        assert_ne!(key1, key2);
    }

    #[test]
    fn test_generate_chat_key_includes_output_affecting_extras() {
        let base = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("hi")],
            ..Default::default()
        };

        let mut parallel = base.clone();
        parallel.parallel_tool_calls = Some(false);
        assert_ne!(generate_chat_key(&base), generate_chat_key(&parallel));

        let mut prediction = base.clone();
        prediction.prediction = Some(serde_json::json!({"type":"content","content":"42"}));
        assert_ne!(generate_chat_key(&base), generate_chat_key(&prediction));

        let mut safety = base.clone();
        safety.safety_settings =
            Some(serde_json::json!([{"category":"HARM","threshold":"BLOCK_NONE"}]));
        assert_ne!(generate_chat_key(&base), generate_chat_key(&safety));

        let mut cache_control = base.clone();
        cache_control.cache_control = Some(serde_json::json!({"type":"ephemeral"}));
        assert_ne!(generate_chat_key(&base), generate_chat_key(&cache_control));

        let mut extra = base.clone();
        extra.extra_body.insert(
            "provider_specific".to_string(),
            serde_json::json!({"k":"v"}),
        );
        assert_ne!(generate_chat_key(&base), generate_chat_key(&extra));
    }

    #[test]
    fn test_generate_chat_key_filters_volatile_extra_body_fields() {
        let base = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("hi")],
            ..Default::default()
        };

        let mut with_request_id = base.clone();
        with_request_id
            .extra_body
            .insert("request_id".to_string(), serde_json::json!("req_123"));

        assert_eq!(
            generate_chat_key(&base),
            generate_chat_key(&with_request_id)
        );
    }

    #[test]
    fn test_generate_chat_key_includes_response_format_schema() {
        let base = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("Return JSON")],
            response_format: Some(ResponseFormat {
                format_type: "json_schema".to_string(),
                json_schema: Some(serde_json::json!({
                    "name": "answer",
                    "schema": {
                        "type": "object",
                        "properties": {
                            "answer": { "type": "string" }
                        }
                    }
                })),
                response_type: None,
            }),
            ..Default::default()
        };

        let mut changed_schema = base.clone();
        changed_schema.response_format = Some(ResponseFormat {
            format_type: "json_schema".to_string(),
            json_schema: Some(serde_json::json!({
                "schema": {
                    "properties": {
                        "answer": { "type": "number" }
                    },
                    "type": "object"
                },
                "name": "answer"
            })),
            response_type: None,
        });

        assert_ne!(generate_chat_key(&base), generate_chat_key(&changed_schema));
    }

    #[test]
    fn test_generate_chat_key_preserves_response_format_schema_id_identity() {
        let base = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("Return JSON")],
            response_format: Some(ResponseFormat {
                format_type: "json_schema".to_string(),
                json_schema: Some(serde_json::json!({
                    "name": "answer",
                    "schema": {
                        "type": "object",
                        "properties": {
                            "id": { "type": "string" }
                        }
                    }
                })),
                response_type: None,
            }),
            ..Default::default()
        };

        let mut changed_id_schema = base.clone();
        changed_id_schema.response_format = Some(ResponseFormat {
            format_type: "json_schema".to_string(),
            json_schema: Some(serde_json::json!({
                "name": "answer",
                "schema": {
                    "type": "object",
                    "properties": {
                        "id": { "type": "integer" }
                    }
                }
            })),
            response_type: None,
        });

        assert_ne!(
            generate_chat_key(&base),
            generate_chat_key(&changed_id_schema)
        );
    }

    #[test]
    fn test_generate_chat_key_preserves_response_format_timestamp_identity() {
        let base = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("Return JSON")],
            response_format: Some(ResponseFormat {
                format_type: "json_schema".to_string(),
                json_schema: Some(serde_json::json!({
                    "name": "answer",
                    "schema": {
                        "type": "object",
                        "properties": {
                            "timestamp": { "type": "string" }
                        }
                    }
                })),
                response_type: None,
            }),
            ..Default::default()
        };

        let mut changed_timestamp_schema = base.clone();
        changed_timestamp_schema.response_format = Some(ResponseFormat {
            format_type: "json_schema".to_string(),
            json_schema: Some(serde_json::json!({
                "name": "answer",
                "schema": {
                    "type": "object",
                    "properties": {
                        "timestamp": { "type": "integer" }
                    }
                }
            })),
            response_type: None,
        });

        assert_ne!(
            generate_chat_key(&base),
            generate_chat_key(&changed_timestamp_schema)
        );
    }

    #[test]
    fn test_generate_chat_key_includes_tool_parameters_and_choice() {
        let base = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("Use the tool")],
            tools: Some(vec![Tool {
                tool_type: "function".to_string(),
                function: Function {
                    name: "search".to_string(),
                    description: Some("Search".to_string()),
                    parameters: Some(serde_json::json!({
                        "type": "object",
                        "properties": {
                            "query": { "type": "string" }
                        }
                    })),
                },
            }]),
            tool_choice: Some(ToolChoice::Auto("auto".to_string())),
            ..Default::default()
        };

        let mut changed_parameters = base.clone();
        changed_parameters.tools.as_mut().unwrap()[0]
            .function
            .parameters = Some(serde_json::json!({
            "type": "object",
            "properties": {
                "query": { "type": "number" }
            }
        }));

        let mut changed_choice = base.clone();
        changed_choice.tool_choice = Some(ToolChoice::Specific(ToolChoiceFunction {
            tool_type: "function".to_string(),
            function: ToolChoiceFunctionSpec {
                name: "search".to_string(),
            },
        }));

        assert_ne!(
            generate_chat_key(&base),
            generate_chat_key(&changed_parameters)
        );
        assert_ne!(generate_chat_key(&base), generate_chat_key(&changed_choice));
    }

    #[test]
    fn test_generate_chat_key_preserves_tool_parameter_id_identity() {
        let base = ChatCompletionRequest {
            model: "gpt-4".to_string(),
            messages: vec![create_user_message("Use the tool")],
            tools: Some(vec![Tool {
                tool_type: "function".to_string(),
                function: Function {
                    name: "lookup".to_string(),
                    description: Some("Lookup by ID".to_string()),
                    parameters: Some(serde_json::json!({
                        "type": "object",
                        "properties": {
                            "id": { "type": "string" }
                        }
                    })),
                },
            }]),
            ..Default::default()
        };

        let mut changed_id_schema = base.clone();
        changed_id_schema.tools.as_mut().unwrap()[0]
            .function
            .parameters = Some(serde_json::json!({
            "type": "object",
            "properties": {
                "id": { "type": "integer" }
            }
        }));

        assert_ne!(
            generate_chat_key(&base),
            generate_chat_key(&changed_id_schema)
        );
    }

    // ==================== Embedding Key Generation Tests ====================

    #[test]
    fn test_generate_embedding_key_string() {
        let request = EmbeddingRequest {
            model: "text-embedding-ada-002".to_string(),
            input: serde_json::json!("Hello world"),
            user: None,
        };

        let key = generate_embedding_key(&request);
        assert!(key.as_str().starts_with("embed:text-embedding-ada-002:v4:"));
    }

    #[test]
    fn test_generate_embedding_key_array() {
        let request = EmbeddingRequest {
            model: "text-embedding-ada-002".to_string(),
            input: serde_json::json!(["Hello", "World"]),
            user: None,
        };

        let key = generate_embedding_key(&request);
        assert!(key.as_str().starts_with("embed:text-embedding-ada-002:"));
    }

    #[test]
    fn test_generate_embedding_key_consistency() {
        let request = EmbeddingRequest {
            model: "text-embedding-3-small".to_string(),
            input: serde_json::json!("Test input"),
            user: None,
        };

        let key1 = generate_embedding_key(&request);
        let key2 = generate_embedding_key(&request);

        assert_eq!(key1, key2);
    }

    #[test]
    fn test_generate_embedding_key_array_order_is_part_of_identity() {
        let request1 = EmbeddingRequest {
            model: "text-embedding-ada-002".to_string(),
            input: serde_json::json!(["Alpha", "Beta"]),
            user: None,
        };

        let request2 = EmbeddingRequest {
            model: "text-embedding-ada-002".to_string(),
            input: serde_json::json!(["Beta", "Alpha"]),
            user: None,
        };

        let key1 = generate_embedding_key(&request1);
        let key2 = generate_embedding_key(&request2);

        assert_ne!(key1, key2);
    }

    // ==================== Generic Key Generation Tests ====================

    #[test]
    fn test_generate_key_from_content() {
        let key = generate_key_from_content("test", "some content");
        assert!(key.as_str().starts_with("test:"));
    }

    #[test]
    fn test_generate_key_from_parts() {
        let key = generate_key_from_parts("prefix", &["part1", "part2", "part3"]);
        assert!(key.as_str().starts_with("prefix:"));
    }

    #[test]
    fn test_generate_key_from_json() {
        #[derive(Serialize)]
        struct TestRequest {
            field1: String,
            field2: i32,
        }

        let request = TestRequest {
            field1: "value".to_string(),
            field2: 42,
        };

        let key = generate_key_from_json("test", &request);
        assert!(key.as_str().starts_with("test:"));
    }

    // ==================== CacheKeyBuilder Tests ====================

    #[test]
    fn test_cache_key_builder_basic() {
        let key = CacheKeyBuilder::new("chat")
            .with_part("gpt-4")
            .with_part("user-123")
            .build();

        assert!(key.as_str().starts_with("chat:"));
    }

    #[test]
    fn test_cache_key_builder_with_nums() {
        let key = CacheKeyBuilder::new("session")
            .with_part("user")
            .add_num(123)
            .add_num(456)
            .build();

        assert!(key.as_str().starts_with("session:"));
    }

    #[test]
    fn test_cache_key_builder_with_optional() {
        let key1 = CacheKeyBuilder::new("test")
            .with_part("base")
            .add_optional(Some("optional"))
            .build();

        let key2 = CacheKeyBuilder::new("test")
            .with_part("base")
            .add_optional(None::<String>)
            .build();

        assert_ne!(key1, key2);
    }

    #[test]
    fn test_cache_key_builder_explicit() {
        let key = CacheKeyBuilder::new("chat")
            .with_part("gpt-4")
            .with_part("conversation-1")
            .build_explicit();

        assert_eq!(key.as_str(), "chat:gpt-4:conversation-1");
    }

    // ==================== JSON Normalization Tests ====================

    #[test]
    fn test_normalize_json_filters_timestamp() {
        let json1 = r#"{"message": "hello", "timestamp": "2024-01-01"}"#;
        let json2 = r#"{"message": "hello", "timestamp": "2024-12-31"}"#;

        let norm1 = normalize_json_string(json1);
        let norm2 = normalize_json_string(json2);

        assert_eq!(norm1, norm2);
    }

    #[test]
    fn test_normalize_json_sorts_keys() {
        let json1 = r#"{"b": 2, "a": 1}"#;
        let json2 = r#"{"a": 1, "b": 2}"#;

        let norm1 = normalize_json_string(json1);
        let norm2 = normalize_json_string(json2);

        assert_eq!(norm1, norm2);
    }

    #[test]
    fn test_is_non_deterministic_field() {
        assert!(is_non_deterministic_field("timestamp"));
        assert!(is_non_deterministic_field("request_id"));
        assert!(is_non_deterministic_field("stream"));
        assert!(!is_non_deterministic_field("model"));
        assert!(!is_non_deterministic_field("messages"));
    }
}