graphrag_cli/
query_history.rs1use chrono::{DateTime, Utc};
4use color_eyre::eyre::Result;
5use serde::{Deserialize, Serialize};
6use std::path::PathBuf;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct QueryEntry {
11 pub query: String,
13 pub timestamp: DateTime<Utc>,
15 pub duration_ms: u128,
17 pub results_count: usize,
19 pub results_preview: Vec<String>,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct QueryHistory {
26 entries: Vec<QueryEntry>,
28 #[serde(default = "default_max_entries")]
30 max_entries: usize,
31}
32
33fn default_max_entries() -> usize {
34 1000
35}
36
37impl QueryHistory {
38 pub fn new() -> Self {
40 Self {
41 entries: Vec::new(),
42 max_entries: default_max_entries(),
43 }
44 }
45
46 pub fn add_entry(&mut self, entry: QueryEntry) {
48 self.entries.insert(0, entry);
49
50 if self.entries.len() > self.max_entries {
52 self.entries.truncate(self.max_entries);
53 }
54 }
55
56 #[allow(dead_code)]
58 pub fn entries(&self) -> &[QueryEntry] {
59 &self.entries
60 }
61
62 #[allow(dead_code)]
64 pub fn last_n(&self, n: usize) -> &[QueryEntry] {
65 let end = n.min(self.entries.len());
66 &self.entries[..end]
67 }
68
69 #[allow(dead_code)]
71 pub fn clear(&mut self) {
72 self.entries.clear()
73 }
74
75 pub fn total_queries(&self) -> usize {
77 self.entries.len()
78 }
79
80 #[allow(dead_code)]
82 pub async fn save(&self, path: &PathBuf) -> Result<()> {
83 let json = serde_json::to_string_pretty(self)?;
84 tokio::fs::write(path, json).await?;
85 Ok(())
86 }
87
88 pub async fn load(path: &PathBuf) -> Result<Self> {
90 let content = tokio::fs::read_to_string(path).await?;
91 let history: Self = serde_json::from_str(&content)?;
92 Ok(history)
93 }
94}
95
96impl Default for QueryHistory {
97 fn default() -> Self {
98 Self::new()
99 }
100}
101
102#[cfg(test)]
103mod tests {
104 use super::*;
105
106 #[test]
107 fn test_add_entry() {
108 let mut history = QueryHistory::new();
109 let entry = QueryEntry {
110 query: "test".to_string(),
111 timestamp: Utc::now(),
112 duration_ms: 100,
113 results_count: 5,
114 results_preview: vec![],
115 };
116
117 history.add_entry(entry.clone());
118 assert_eq!(history.total_queries(), 1);
119 assert_eq!(history.entries()[0].query, "test");
120 }
121
122 #[test]
123 fn test_max_entries() {
124 let mut history = QueryHistory::new();
125 history.max_entries = 5;
126
127 for i in 0..10 {
128 history.add_entry(QueryEntry {
129 query: format!("query {}", i),
130 timestamp: Utc::now(),
131 duration_ms: 100,
132 results_count: 1,
133 results_preview: vec![],
134 });
135 }
136
137 assert_eq!(history.total_queries(), 5);
138 }
139}