codex_memory/memory/
repository.rs

1use super::error::{MemoryError, Result};
2use super::models::*;
3use chrono::Utc;
4use pgvector::Vector;
5use sqlx::{PgPool, Postgres, Row, Transaction};
6use std::collections::HashMap;
7use std::time::Instant;
8use tracing::{debug, info};
9use uuid::Uuid;
10
11pub struct MemoryRepository {
12    pool: PgPool,
13}
14
15impl MemoryRepository {
16    pub fn new(pool: PgPool) -> Self {
17        Self { pool }
18    }
19
20    pub fn pool(&self) -> &PgPool {
21        &self.pool
22    }
23
24    pub async fn create_memory(&self, request: CreateMemoryRequest) -> Result<Memory> {
25        let id = Uuid::new_v4();
26        let content_hash = Memory::calculate_content_hash(&request.content);
27        let tier = request.tier.unwrap_or(MemoryTier::Working);
28
29        // Check for duplicates
30        let duplicate_exists = sqlx::query_scalar::<_, bool>(
31            "SELECT EXISTS(SELECT 1 FROM memories WHERE content_hash = $1 AND tier = $2 AND status = 'active')"
32        )
33        .bind(&content_hash)
34        .bind(tier)
35        .fetch_one(&self.pool)
36        .await?;
37
38        if duplicate_exists {
39            return Err(MemoryError::DuplicateContent {
40                tier: format!("{tier:?}"),
41            });
42        }
43
44        let embedding = request.embedding.map(Vector::from);
45
46        let memory = sqlx::query_as::<_, Memory>(
47            r#"
48            INSERT INTO memories (
49                id, content, content_hash, embedding, tier, status, 
50                importance_score, metadata, parent_id, expires_at
51            )
52            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
53            RETURNING *
54            "#,
55        )
56        .bind(id)
57        .bind(&request.content)
58        .bind(&content_hash)
59        .bind(embedding)
60        .bind(tier)
61        .bind(MemoryStatus::Active)
62        .bind(request.importance_score.unwrap_or(0.5))
63        .bind(request.metadata.unwrap_or(serde_json::json!({})))
64        .bind(request.parent_id)
65        .bind(request.expires_at)
66        .fetch_one(&self.pool)
67        .await?;
68
69        info!("Created memory {} in tier {:?}", memory.id, memory.tier);
70        Ok(memory)
71    }
72
73    pub async fn get_memory(&self, id: Uuid) -> Result<Memory> {
74        let memory = sqlx::query_as::<_, Memory>(
75            r#"
76            UPDATE memories 
77            SET access_count = access_count + 1, 
78                last_accessed_at = NOW()
79            WHERE id = $1 AND status = 'active'
80            RETURNING *
81            "#,
82        )
83        .bind(id)
84        .fetch_optional(&self.pool)
85        .await?
86        .ok_or_else(|| MemoryError::NotFound { id: id.to_string() })?;
87
88        debug!("Retrieved memory {} from tier {:?}", id, memory.tier);
89        Ok(memory)
90    }
91
92    pub async fn update_memory(&self, id: Uuid, request: UpdateMemoryRequest) -> Result<Memory> {
93        let mut tx = self.pool.begin().await?;
94
95        // Get current memory
96        let current = sqlx::query_as::<_, Memory>(
97            "SELECT * FROM memories WHERE id = $1 AND status = 'active' FOR UPDATE",
98        )
99        .bind(id)
100        .fetch_optional(&mut *tx)
101        .await?
102        .ok_or_else(|| MemoryError::NotFound { id: id.to_string() })?;
103
104        // Update fields
105        let content = request.content.as_ref().unwrap_or(&current.content);
106        let content_hash = if request.content.is_some() {
107            Memory::calculate_content_hash(content)
108        } else {
109            current.content_hash.clone()
110        };
111
112        let embedding = request.embedding.map(Vector::from).or(current.embedding);
113        let tier = request.tier.unwrap_or(current.tier);
114        let importance_score = request.importance_score.unwrap_or(current.importance_score);
115        let metadata = request.metadata.as_ref().unwrap_or(&current.metadata);
116        let expires_at = request.expires_at.or(current.expires_at);
117
118        let updated = sqlx::query_as::<_, Memory>(
119            r#"
120            UPDATE memories 
121            SET content = $2, content_hash = $3, embedding = $4, tier = $5,
122                importance_score = $6, metadata = $7, expires_at = $8,
123                updated_at = NOW()
124            WHERE id = $1
125            RETURNING *
126            "#,
127        )
128        .bind(id)
129        .bind(content)
130        .bind(&content_hash)
131        .bind(embedding)
132        .bind(tier)
133        .bind(importance_score)
134        .bind(metadata)
135        .bind(expires_at)
136        .fetch_one(&mut *tx)
137        .await?;
138
139        // Record tier migration if changed
140        if current.tier != tier {
141            self.record_migration(
142                &mut tx,
143                id,
144                current.tier,
145                tier,
146                Some("Manual update".to_string()),
147            )
148            .await?;
149        }
150
151        tx.commit().await?;
152        info!("Updated memory {}", id);
153        Ok(updated)
154    }
155
156    pub async fn delete_memory(&self, id: Uuid) -> Result<()> {
157        let result = sqlx::query(
158            "UPDATE memories SET status = 'deleted' WHERE id = $1 AND status = 'active'",
159        )
160        .bind(id)
161        .execute(&self.pool)
162        .await?;
163
164        if result.rows_affected() == 0 {
165            return Err(MemoryError::NotFound { id: id.to_string() });
166        }
167
168        info!("Soft deleted memory {}", id);
169        Ok(())
170    }
171
172    pub async fn search_memories(&self, request: SearchRequest) -> Result<SearchResponse> {
173        let start_time = Instant::now();
174
175        let search_type = request
176            .search_type
177            .as_ref()
178            .unwrap_or(&SearchType::Semantic)
179            .clone();
180        let limit = request.limit.unwrap_or(10);
181        let offset = request.offset.unwrap_or(0);
182
183        let results = match search_type {
184            SearchType::Semantic => self.semantic_search(&request).await?,
185            SearchType::Temporal => self.temporal_search(&request).await?,
186            SearchType::Hybrid => self.hybrid_search(&request).await?,
187            SearchType::FullText => self.fulltext_search(&request).await?,
188        };
189
190        let total_count = if request.include_facets.unwrap_or(false) {
191            Some(self.count_search_results(&request).await?)
192        } else {
193            None
194        };
195
196        let facets = if request.include_facets.unwrap_or(false) {
197            Some(self.generate_search_facets(&request).await?)
198        } else {
199            None
200        };
201
202        let suggestions = if request.query_text.is_some() {
203            Some(self.generate_query_suggestions(&request).await?)
204        } else {
205            None
206        };
207
208        let next_cursor = if results.len() as i32 >= limit {
209            Some(self.generate_cursor(offset + limit as i64, &request))
210        } else {
211            None
212        };
213
214        let execution_time_ms = start_time.elapsed().as_millis() as u64;
215
216        Ok(SearchResponse {
217            results,
218            total_count,
219            facets,
220            suggestions,
221            next_cursor,
222            execution_time_ms,
223        })
224    }
225
226    async fn semantic_search(&self, request: &SearchRequest) -> Result<Vec<SearchResult>> {
227        let query_embedding = if let Some(ref embedding) = request.query_embedding {
228            Vector::from(embedding.clone())
229        } else {
230            return Err(MemoryError::InvalidRequest {
231                message: "Query embedding is required for semantic search".to_string(),
232            });
233        };
234
235        let limit = request.limit.unwrap_or(10);
236        let offset = request.offset.unwrap_or(0);
237        let threshold = request.similarity_threshold.unwrap_or(0.7);
238
239        let mut query_parts = vec![
240            "SELECT m.*, 1 - (m.embedding <=> $1) as similarity_score".to_string(),
241            "FROM memories m".to_string(),
242            "WHERE m.status = 'active' AND m.embedding IS NOT NULL".to_string(),
243        ];
244
245        // Add filters
246        self.add_filters(request, &mut query_parts)?;
247
248        query_parts.push(format!("AND 1 - (m.embedding <=> $1) >= {threshold}"));
249        query_parts.push("ORDER BY similarity_score DESC".to_string());
250        query_parts.push(format!("LIMIT {limit} OFFSET {offset}"));
251
252        let query = query_parts.join(" ");
253        let rows = sqlx::query(&query)
254            .bind(&query_embedding)
255            .fetch_all(&self.pool)
256            .await?;
257
258        self.build_search_results(rows, request).await
259    }
260
261    async fn temporal_search(&self, request: &SearchRequest) -> Result<Vec<SearchResult>> {
262        let limit = request.limit.unwrap_or(10);
263        let offset = request.offset.unwrap_or(0);
264
265        let mut query_parts = vec![
266            "SELECT m.*, 0.0 as similarity_score".to_string(),
267            "FROM memories m".to_string(),
268            "WHERE m.status = 'active'".to_string(),
269        ];
270
271        // Add filters
272        self.add_filters(request, &mut query_parts)?;
273
274        query_parts.push("ORDER BY m.created_at DESC, m.updated_at DESC".to_string());
275        query_parts.push(format!("LIMIT {limit} OFFSET {offset}"));
276
277        let query = query_parts.join(" ");
278        let rows = sqlx::query(&query).fetch_all(&self.pool).await?;
279
280        self.build_search_results(rows, request).await
281    }
282
283    async fn hybrid_search(&self, request: &SearchRequest) -> Result<Vec<SearchResult>> {
284        let weights = request.hybrid_weights.as_ref().unwrap_or(&HybridWeights {
285            semantic_weight: 0.4,
286            temporal_weight: 0.3,
287            importance_weight: 0.2,
288            access_frequency_weight: 0.1,
289        });
290
291        let query_embedding = if let Some(ref embedding) = request.query_embedding {
292            Vector::from(embedding.clone())
293        } else {
294            return Err(MemoryError::InvalidRequest {
295                message: "Query embedding is required for hybrid search".to_string(),
296            });
297        };
298
299        let limit = request.limit.unwrap_or(10);
300        let offset = request.offset.unwrap_or(0);
301        let threshold = request.similarity_threshold.unwrap_or(0.5);
302
303        let query = format!(
304            r#"
305            SELECT m.*,
306                1 - (m.embedding <=> $1) as similarity_score,
307                EXTRACT(EPOCH FROM (NOW() - m.created_at))::float / 86400 as days_old,
308                m.importance_score,
309                COALESCE(m.access_count, 0) as access_count,
310                (
311                    {} * (1 - (m.embedding <=> $1)) +
312                    {} * GREATEST(0, 1 - (EXTRACT(EPOCH FROM (NOW() - m.created_at))::float / 2592000)) + -- 30 days
313                    {} * m.importance_score +
314                    {} * LEAST(1.0, COALESCE(m.access_count, 0)::float / 100.0)
315                ) as combined_score
316            FROM memories m
317            WHERE m.status = 'active'
318                AND m.embedding IS NOT NULL
319                AND 1 - (m.embedding <=> $1) >= {}
320            ORDER BY combined_score DESC
321            LIMIT {} OFFSET {}
322            "#,
323            weights.semantic_weight,
324            weights.temporal_weight,
325            weights.importance_weight,
326            weights.access_frequency_weight,
327            threshold,
328            limit,
329            offset
330        );
331
332        let rows = sqlx::query(&query)
333            .bind(&query_embedding)
334            .fetch_all(&self.pool)
335            .await?;
336
337        self.build_search_results(rows, request).await
338    }
339
340    async fn fulltext_search(&self, request: &SearchRequest) -> Result<Vec<SearchResult>> {
341        let query_text =
342            request
343                .query_text
344                .as_ref()
345                .ok_or_else(|| MemoryError::InvalidRequest {
346                    message: "Query text is required for full-text search".to_string(),
347                })?;
348
349        let limit = request.limit.unwrap_or(10);
350        let offset = request.offset.unwrap_or(0);
351
352        let query = format!(
353            r#"
354            SELECT m.*,
355                ts_rank_cd(to_tsvector('english', m.content), plainto_tsquery('english', $1)) as similarity_score
356            FROM memories m
357            WHERE m.status = 'active'
358                AND to_tsvector('english', m.content) @@ plainto_tsquery('english', $1)
359            ORDER BY similarity_score DESC
360            LIMIT {limit} OFFSET {offset}
361            "#
362        );
363
364        let rows = sqlx::query(&query)
365            .bind(query_text)
366            .fetch_all(&self.pool)
367            .await?;
368
369        self.build_search_results(rows, request).await
370    }
371
372    fn add_filters(&self, request: &SearchRequest, query_parts: &mut Vec<String>) -> Result<()> {
373        if let Some(tier) = &request.tier {
374            query_parts.push(format!("AND m.tier = '{tier:?}'"));
375        }
376
377        if let Some(date_range) = &request.date_range {
378            if let Some(start) = &date_range.start {
379                query_parts.push(format!(
380                    "AND m.created_at >= '{}'",
381                    start.format("%Y-%m-%d %H:%M:%S")
382                ));
383            }
384            if let Some(end) = &date_range.end {
385                query_parts.push(format!(
386                    "AND m.created_at <= '{}'",
387                    end.format("%Y-%m-%d %H:%M:%S")
388                ));
389            }
390        }
391
392        if let Some(importance_range) = &request.importance_range {
393            if let Some(min) = importance_range.min {
394                query_parts.push(format!("AND m.importance_score >= {min}"));
395            }
396            if let Some(max) = importance_range.max {
397                query_parts.push(format!("AND m.importance_score <= {max}"));
398            }
399        }
400
401        Ok(())
402    }
403
404    async fn build_search_results(
405        &self,
406        rows: Vec<sqlx::postgres::PgRow>,
407        request: &SearchRequest,
408    ) -> Result<Vec<SearchResult>> {
409        let mut results = Vec::new();
410        let explain_score = request.explain_score.unwrap_or(false);
411
412        for row in rows {
413            let memory = Memory {
414                id: row.try_get("id")?,
415                content: row.try_get("content")?,
416                content_hash: row.try_get("content_hash")?,
417                embedding: row.try_get("embedding")?,
418                tier: row.try_get("tier")?,
419                status: row.try_get("status")?,
420                importance_score: row.try_get("importance_score")?,
421                access_count: row.try_get("access_count")?,
422                last_accessed_at: row.try_get("last_accessed_at")?,
423                metadata: row.try_get("metadata")?,
424                parent_id: row.try_get("parent_id")?,
425                created_at: row.try_get("created_at")?,
426                updated_at: row.try_get("updated_at")?,
427                expires_at: row.try_get("expires_at")?,
428            };
429
430            let similarity_score: f32 = row.try_get("similarity_score").unwrap_or(0.0);
431            let combined_score: f32 = row.try_get("combined_score").unwrap_or(similarity_score);
432            let temporal_score: Option<f32> = row.try_get("temporal_score").ok();
433            let access_frequency_score: Option<f32> = row.try_get("access_frequency_score").ok();
434            let importance_score = memory.importance_score; // Extract before move
435
436            let score_explanation = if explain_score {
437                Some(ScoreExplanation {
438                    semantic_contribution: similarity_score * 0.4,
439                    temporal_contribution: temporal_score.unwrap_or(0.0) * 0.3,
440                    importance_contribution: (importance_score * 0.2) as f32,
441                    access_frequency_contribution: access_frequency_score.unwrap_or(0.0) * 0.1,
442                    total_score: combined_score,
443                    factors: vec![
444                        "semantic similarity".to_string(),
445                        "recency".to_string(),
446                        "importance".to_string(),
447                    ],
448                })
449            } else {
450                None
451            };
452
453            results.push(SearchResult {
454                memory,
455                similarity_score,
456                temporal_score,
457                importance_score,
458                access_frequency_score,
459                combined_score,
460                score_explanation,
461            });
462        }
463
464        debug!("Built {} search results", results.len());
465        Ok(results)
466    }
467
468    async fn count_search_results(&self, _request: &SearchRequest) -> Result<i64> {
469        // Simplified count - would implement filtering logic similar to search
470        let count: i64 =
471            sqlx::query_scalar("SELECT COUNT(*) FROM memories WHERE status = 'active'")
472                .fetch_one(&self.pool)
473                .await?;
474        Ok(count)
475    }
476
477    async fn generate_search_facets(&self, _request: &SearchRequest) -> Result<SearchFacets> {
478        // Generate tier facets
479        let tier_rows: Vec<(String, i64)> = sqlx::query_as(
480            "SELECT tier, COUNT(*) FROM memories WHERE status = 'active' GROUP BY tier",
481        )
482        .fetch_all(&self.pool)
483        .await?;
484
485        let mut tiers = HashMap::new();
486        for (tier_str, count) in tier_rows {
487            if let Ok(tier) = tier_str.parse::<MemoryTier>() {
488                tiers.insert(tier, count);
489            }
490        }
491
492        // Generate date histogram (simplified)
493        let date_histogram = vec![DateBucket {
494            date: Utc::now(),
495            count: 10,
496            interval: "day".to_string(),
497        }];
498
499        // Generate importance ranges
500        let importance_ranges = vec![
501            ImportanceRange {
502                min: 0.0,
503                max: 0.3,
504                count: 5,
505                label: "Low".to_string(),
506            },
507            ImportanceRange {
508                min: 0.3,
509                max: 0.7,
510                count: 15,
511                label: "Medium".to_string(),
512            },
513            ImportanceRange {
514                min: 0.7,
515                max: 1.0,
516                count: 8,
517                label: "High".to_string(),
518            },
519        ];
520
521        Ok(SearchFacets {
522            tiers,
523            date_histogram,
524            importance_ranges,
525            tags: HashMap::new(), // Would extract from metadata
526        })
527    }
528
529    async fn generate_query_suggestions(&self, _request: &SearchRequest) -> Result<Vec<String>> {
530        // Simplified implementation - would use ML model or query history
531        Ok(vec![
532            "recent code changes".to_string(),
533            "function definitions".to_string(),
534            "error handling patterns".to_string(),
535        ])
536    }
537
538    fn generate_cursor(&self, offset: i64, _request: &SearchRequest) -> String {
539        // Simple cursor implementation - would encode more search state in production
540        use base64::{engine::general_purpose::STANDARD, Engine};
541        STANDARD.encode(format!("offset:{offset}"))
542    }
543
544    // Legacy method for backward compatibility
545    pub async fn search_memories_simple(
546        &self,
547        request: SearchRequest,
548    ) -> Result<Vec<SearchResult>> {
549        let response = self.search_memories(request).await?;
550        Ok(response.results)
551    }
552
553    pub async fn get_memories_by_tier(
554        &self,
555        tier: MemoryTier,
556        limit: Option<i64>,
557    ) -> Result<Vec<Memory>> {
558        let limit = limit.unwrap_or(100);
559
560        let memories = sqlx::query_as::<_, Memory>(
561            r#"
562            SELECT * FROM memories 
563            WHERE tier = $1 AND status = 'active'
564            ORDER BY importance_score DESC, updated_at DESC
565            LIMIT $2
566            "#,
567        )
568        .bind(tier)
569        .bind(limit)
570        .fetch_all(&self.pool)
571        .await?;
572
573        Ok(memories)
574    }
575
576    pub async fn migrate_memory(
577        &self,
578        id: Uuid,
579        to_tier: MemoryTier,
580        reason: Option<String>,
581    ) -> Result<Memory> {
582        let mut tx = self.pool.begin().await?;
583
584        // Get current memory with lock
585        let current = sqlx::query_as::<_, Memory>(
586            "SELECT * FROM memories WHERE id = $1 AND status = 'active' FOR UPDATE",
587        )
588        .bind(id)
589        .fetch_optional(&mut *tx)
590        .await?
591        .ok_or_else(|| MemoryError::NotFound { id: id.to_string() })?;
592
593        if current.tier == to_tier {
594            return Ok(current);
595        }
596
597        // Validate tier transition
598        let valid_transition = match (current.tier, to_tier) {
599            (MemoryTier::Working, MemoryTier::Warm)
600            | (MemoryTier::Working, MemoryTier::Cold)
601            | (MemoryTier::Warm, MemoryTier::Cold)
602            | (MemoryTier::Warm, MemoryTier::Working)
603            | (MemoryTier::Cold, MemoryTier::Warm) => true,
604            _ => false,
605        };
606
607        if !valid_transition {
608            return Err(MemoryError::InvalidTierTransition {
609                from: format!("{:?}", current.tier),
610                to: format!("{to_tier:?}"),
611            });
612        }
613
614        let start = std::time::Instant::now();
615
616        // Update tier
617        let updated = sqlx::query_as::<_, Memory>(
618            r#"
619            UPDATE memories 
620            SET tier = $2, status = 'active', updated_at = NOW()
621            WHERE id = $1
622            RETURNING *
623            "#,
624        )
625        .bind(id)
626        .bind(to_tier)
627        .fetch_one(&mut *tx)
628        .await?;
629
630        let duration_ms = start.elapsed().as_millis() as i32;
631
632        // Record migration
633        self.record_migration(&mut tx, id, current.tier, to_tier, reason)
634            .await?;
635
636        tx.commit().await?;
637
638        info!(
639            "Migrated memory {} from {:?} to {:?} in {}ms",
640            id, current.tier, to_tier, duration_ms
641        );
642
643        Ok(updated)
644    }
645
646    async fn record_migration(
647        &self,
648        tx: &mut Transaction<'_, Postgres>,
649        memory_id: Uuid,
650        from_tier: MemoryTier,
651        to_tier: MemoryTier,
652        reason: Option<String>,
653    ) -> Result<()> {
654        sqlx::query(
655            r#"
656            INSERT INTO migration_history (memory_id, from_tier, to_tier, migration_reason, success)
657            VALUES ($1, $2, $3, $4, true)
658            "#,
659        )
660        .bind(memory_id)
661        .bind(from_tier)
662        .bind(to_tier)
663        .bind(reason)
664        .execute(&mut **tx)
665        .await?;
666
667        Ok(())
668    }
669
670    pub async fn get_expired_memories(&self) -> Result<Vec<Memory>> {
671        let memories = sqlx::query_as::<_, Memory>(
672            r#"
673            SELECT * FROM memories 
674            WHERE status = 'active' 
675                AND expires_at IS NOT NULL 
676                AND expires_at < NOW()
677            "#,
678        )
679        .fetch_all(&self.pool)
680        .await?;
681
682        Ok(memories)
683    }
684
685    pub async fn cleanup_expired_memories(&self) -> Result<u64> {
686        let result = sqlx::query(
687            r#"
688            UPDATE memories 
689            SET status = 'deleted' 
690            WHERE status = 'active' 
691                AND expires_at IS NOT NULL 
692                AND expires_at < NOW()
693            "#,
694        )
695        .execute(&self.pool)
696        .await?;
697
698        let count = result.rows_affected();
699        if count > 0 {
700            info!("Cleaned up {} expired memories", count);
701        }
702
703        Ok(count)
704    }
705
706    pub async fn get_migration_candidates(
707        &self,
708        tier: MemoryTier,
709        limit: i64,
710    ) -> Result<Vec<Memory>> {
711        let query = match tier {
712            MemoryTier::Working => {
713                r#"
714                SELECT * FROM memories 
715                WHERE tier = 'working' 
716                    AND status = 'active'
717                    AND (
718                        importance_score < 0.3 
719                        OR (last_accessed_at IS NOT NULL 
720                            AND last_accessed_at < NOW() - INTERVAL '24 hours')
721                    )
722                ORDER BY importance_score ASC, last_accessed_at ASC NULLS FIRST
723                LIMIT $1
724                "#
725            }
726            MemoryTier::Warm => {
727                r#"
728                SELECT * FROM memories 
729                WHERE tier = 'warm' 
730                    AND status = 'active'
731                    AND importance_score < 0.1 
732                    AND updated_at < NOW() - INTERVAL '7 days'
733                ORDER BY importance_score ASC, updated_at ASC
734                LIMIT $1
735                "#
736            }
737            MemoryTier::Cold => {
738                return Ok(Vec::new());
739            }
740        };
741
742        let memories = sqlx::query_as::<_, Memory>(query)
743            .bind(limit)
744            .fetch_all(&self.pool)
745            .await?;
746
747        Ok(memories)
748    }
749
750    pub async fn get_statistics(&self) -> Result<MemoryStatistics> {
751        let stats = sqlx::query_as::<_, MemoryStatistics>(
752            r#"
753            SELECT 
754                COUNT(*) FILTER (WHERE tier = 'working' AND status = 'active') as working_count,
755                COUNT(*) FILTER (WHERE tier = 'warm' AND status = 'active') as warm_count,
756                COUNT(*) FILTER (WHERE tier = 'cold' AND status = 'active') as cold_count,
757                COUNT(*) FILTER (WHERE status = 'active') as total_active,
758                COUNT(*) FILTER (WHERE status = 'deleted') as total_deleted,
759                AVG(importance_score) FILTER (WHERE status = 'active') as avg_importance,
760                MAX(access_count) FILTER (WHERE status = 'active') as max_access_count,
761                AVG(access_count) FILTER (WHERE status = 'active') as avg_access_count
762            FROM memories
763            "#,
764        )
765        .fetch_one(&self.pool)
766        .await?;
767
768        Ok(stats)
769    }
770}
771
772#[derive(Debug, Clone, sqlx::FromRow, serde::Serialize, serde::Deserialize)]
773pub struct MemoryStatistics {
774    pub working_count: Option<i64>,
775    pub warm_count: Option<i64>,
776    pub cold_count: Option<i64>,
777    pub total_active: Option<i64>,
778    pub total_deleted: Option<i64>,
779    pub avg_importance: Option<f64>,
780    pub max_access_count: Option<i32>,
781    pub avg_access_count: Option<f64>,
782}
783
784#[cfg(test)]
785mod tests {
786    use super::*;
787
788    #[test]
789    fn test_content_hash_generation() {
790        let content = "This is a test memory content";
791        let hash1 = Memory::calculate_content_hash(content);
792        let hash2 = Memory::calculate_content_hash(content);
793
794        assert_eq!(hash1, hash2);
795        assert_eq!(hash1.len(), 64); // SHA-256 produces 64 hex characters
796    }
797
798    #[test]
799    fn test_should_migrate() {
800        let mut memory = Memory::default();
801
802        // Working tier with low importance should migrate
803        memory.tier = MemoryTier::Working;
804        memory.importance_score = 0.2;
805        assert!(memory.should_migrate());
806
807        // Working tier with high importance should not migrate
808        memory.importance_score = 0.8;
809        assert!(!memory.should_migrate());
810
811        // Cold tier should never migrate
812        memory.tier = MemoryTier::Cold;
813        memory.importance_score = 0.0;
814        assert!(!memory.should_migrate());
815    }
816
817    #[test]
818    fn test_next_tier() {
819        let mut memory = Memory::default();
820
821        memory.tier = MemoryTier::Working;
822        assert_eq!(memory.next_tier(), Some(MemoryTier::Warm));
823
824        memory.tier = MemoryTier::Warm;
825        assert_eq!(memory.next_tier(), Some(MemoryTier::Cold));
826
827        memory.tier = MemoryTier::Cold;
828        assert_eq!(memory.next_tier(), None);
829    }
830}