rust-memory-service/src/storage/query.rs

//! Memory Query Definition
//!
//! Query structure for memory retrieval operations.

use serde::{Deserialize, Serialize};
use chrono::{DateTime, Utc};
use super::entry::MemoryType;

/// Time range for queries
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TimeRange {
    pub start: DateTime<Utc>,
    pub end: DateTime<Utc>,
}

/// Query parameters for memory retrieval
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryQuery {
    /// Filter by user ID
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user_id: Option<String>,

    /// Filter by session ID
    #[serde(skip_serializing_if = "Option::is_none")]
    pub session_id: Option<String>,

    /// Text query for semantic search
    #[serde(skip_serializing_if = "Option::is_none")]
    pub query: Option<String>,

    /// Pre-computed query embedding
    #[serde(skip_serializing_if = "Option::is_none")]
    pub query_embedding: Option<Vec<f32>>,

    /// Memory types to search (empty = all)
    #[serde(default)]
    pub memory_types: Vec<MemoryType>,

    /// Time range filter
    #[serde(skip_serializing_if = "Option::is_none")]
    pub time_range: Option<TimeRange>,

    /// Tags to filter by
    #[serde(default)]
    pub tags: Vec<String>,

    /// Maximum results to return
    #[serde(default = "default_top_k")]
    pub top_k: usize,

    /// Minimum similarity threshold (0.0 - 1.0)
    #[serde(default)]
    pub similarity_threshold: f64,

    /// Maximum age in hours
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_age_hours: Option<f64>,

    /// Sources to exclude
    #[serde(default)]
    pub exclude_sources: Vec<String>,

    /// Include expired entries
    #[serde(default)]
    pub include_expired: bool,
}

fn default_top_k() -> usize {
    10
}

impl Default for MemoryQuery {
    fn default() -> Self {
        Self {
            user_id: None,
            session_id: None,
            query: None,
            query_embedding: None,
            memory_types: Vec::new(),
            time_range: None,
            tags: Vec::new(),
            top_k: 10,
            similarity_threshold: 0.0,
            max_age_hours: None,
            exclude_sources: Vec::new(),
            include_expired: false,
        }
    }
}

impl MemoryQuery {
    /// Create a new query for a user
    pub fn for_user(user_id: String) -> Self {
        Self {
            user_id: Some(user_id),
            ..Default::default()
        }
    }

    /// Create a query for a session
    pub fn for_session(session_id: String) -> Self {
        Self {
            session_id: Some(session_id),
            ..Default::default()
        }
    }

    /// Create a semantic search query
    pub fn semantic(query: String, embedding: Vec<f32>) -> Self {
        Self {
            query: Some(query),
            query_embedding: Some(embedding),
            ..Default::default()
        }
    }

    /// Set user filter
    pub fn with_user(mut self, user_id: String) -> Self {
        self.user_id = Some(user_id);
        self
    }

    /// Set session filter
    pub fn with_session(mut self, session_id: String) -> Self {
        self.session_id = Some(session_id);
        self
    }

    /// Set memory types filter
    pub fn with_types(mut self, types: Vec<MemoryType>) -> Self {
        self.memory_types = types;
        self
    }

    /// Set top K results
    pub fn with_top_k(mut self, k: usize) -> Self {
        self.top_k = k;
        self
    }

    /// Set similarity threshold
    pub fn with_similarity_threshold(mut self, threshold: f64) -> Self {
        self.similarity_threshold = threshold;
        self
    }

    /// Set max age filter
    pub fn with_max_age_hours(mut self, hours: f64) -> Self {
        self.max_age_hours = Some(hours);
        self
    }

    /// Set time range
    pub fn with_time_range(mut self, start: DateTime<Utc>, end: DateTime<Utc>) -> Self {
        self.time_range = Some(TimeRange { start, end });
        self
    }

    /// Set tags filter
    pub fn with_tags(mut self, tags: Vec<String>) -> Self {
        self.tags = tags;
        self
    }

    /// Exclude sources
    pub fn excluding_sources(mut self, sources: Vec<String>) -> Self {
        self.exclude_sources = sources;
        self
    }

    /// Include expired entries
    pub fn including_expired(mut self) -> Self {
        self.include_expired = true;
        self
    }

    /// Generate a cache key for this query
    pub fn cache_key(&self) -> String {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();

        self.user_id.hash(&mut hasher);
        self.session_id.hash(&mut hasher);
        self.query.hash(&mut hasher);
        self.top_k.hash(&mut hasher);
        
        for mt in &self.memory_types {
            (*mt as u8).hash(&mut hasher);
        }

        format!("query:{:x}", hasher.finish())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_query_builder() {
        let query = MemoryQuery::for_user("user-1".to_string())
            .with_types(vec![MemoryType::Episodic])
            .with_top_k(5)
            .with_similarity_threshold(0.7);

        assert_eq!(query.user_id, Some("user-1".to_string()));
        assert_eq!(query.memory_types, vec![MemoryType::Episodic]);
        assert_eq!(query.top_k, 5);
        assert!((query.similarity_threshold - 0.7).abs() < 0.001);
    }

    #[test]
    fn test_cache_key() {
        let q1 = MemoryQuery::for_user("user-1".to_string()).with_top_k(5);
        let q2 = MemoryQuery::for_user("user-1".to_string()).with_top_k(5);
        let q3 = MemoryQuery::for_user("user-2".to_string()).with_top_k(5);

        assert_eq!(q1.cache_key(), q2.cache_key());
        assert_ne!(q1.cache_key(), q3.cache_key());
    }
}