ai00_mem/
database.rs

1//! 数据库模�?//!
2//! 本模块提供向量数据库和图数据库的统一接口，支持SQLite、PostgreSQL和MySQL�?//! 包含向量存储、图节点和边的管理、以及高效的查询功能�?
3use crate::config::{Config, DatabaseType};
4use crate::error::{MemoryError, Result};
5use chrono::{DateTime, Utc};
6use lru::LruCache;
7use serde::{Deserialize, Serialize};
8
9use sqlx::{Pool, Row, Sqlite};
10use std::collections::HashMap;
11use std::num::NonZeroUsize;
12use std::sync::Arc;
13use tokio::sync::RwLock;
14
15/// 数据库连接池枚举
16#[derive(Debug, Clone)]
17pub enum DatabasePool {
18    Sqlite(Pool<Sqlite>),
19
20}
21
22/// 向量数据结构
23#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
24pub struct Vector {
25    pub id: String,
26    pub embedding: Vec<f32>,
27    pub metadata: HashMap<String, serde_json::Value>,
28    pub created_at: DateTime<Utc>,
29    pub updated_at: DateTime<Utc>,
30}
31
32/// 图节点数据结构
33#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
34pub struct GraphNode {
35    pub id: String,
36    pub node_type: String,
37    pub properties: HashMap<String, serde_json::Value>,
38    pub created_at: DateTime<Utc>,
39    pub updated_at: DateTime<Utc>,
40}
41
42/// 图边数据结构
43#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
44pub struct GraphEdge {
45    pub id: String,
46    pub from_node: String,
47    pub to_node: String,
48    pub edge_type: String,
49    pub weight: f32,
50    pub properties: HashMap<String, serde_json::Value>,
51    pub created_at: DateTime<Utc>,
52    pub updated_at: DateTime<Utc>,
53}
54
55/// 向量查询请求
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct VectorQueryRequest {
58    pub query_vector: Vec<f32>,
59    pub limit: Option<usize>,
60    pub threshold: Option<f32>,
61    pub filters: Option<HashMap<String, serde_json::Value>>,
62}
63
64/// 向量查询结果
65#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct VectorQueryResult {
67    pub vector: Vector,
68    pub similarity: f32,
69    pub distance: f32,
70}
71
72/// 图查询请求
73#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct GraphQueryRequest {
75    pub start_nodes: Vec<String>,
76    pub edge_types: Option<Vec<String>>,
77    pub max_depth: Option<usize>,
78    pub limit: Option<usize>,
79    pub filters: Option<HashMap<String, serde_json::Value>>,
80}
81
82/// 图查询结果
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct GraphQueryResult {
85    pub nodes: Vec<GraphNode>,
86    pub edges: Vec<GraphEdge>,
87    pub paths: Vec<Vec<String>>,
88}
89
90/// 内存缓存
91#[derive(Debug)]
92pub struct MemoryCache {
93    vectors: Arc<RwLock<LruCache<String, Vector>>>,
94    nodes: Arc<RwLock<LruCache<String, GraphNode>>>,
95    edges: Arc<RwLock<LruCache<String, GraphEdge>>>,
96    query_cache: Arc<RwLock<LruCache<String, VectorQueryResult>>>,
97}
98
99impl MemoryCache {
100    pub fn new(capacity: usize) -> Self {
101        let cache_capacity =
102            NonZeroUsize::new(capacity).unwrap_or(NonZeroUsize::new(1000).unwrap());
103
104        Self {
105            vectors: Arc::new(RwLock::new(LruCache::new(cache_capacity))),
106            nodes: Arc::new(RwLock::new(LruCache::new(cache_capacity))),
107            edges: Arc::new(RwLock::new(LruCache::new(cache_capacity))),
108            query_cache: Arc::new(RwLock::new(LruCache::new(cache_capacity))),
109        }
110    }
111
112    pub async fn get_vector(&self, id: &str) -> Option<Vector> {
113        self.vectors.read().await.peek(id).cloned()
114    }
115
116    pub async fn put_vector(&self, id: String, vector: Vector) {
117        self.vectors.write().await.put(id, vector);
118    }
119
120    pub async fn get_node(&self, id: &str) -> Option<GraphNode> {
121        self.nodes.read().await.peek(id).cloned()
122    }
123
124    pub async fn put_node(&self, id: String, node: GraphNode) {
125        self.nodes.write().await.put(id, node);
126    }
127
128    pub async fn get_edge(&self, id: &str) -> Option<GraphEdge> {
129        self.edges.read().await.peek(id).cloned()
130    }
131
132    pub async fn put_edge(&self, id: String, edge: GraphEdge) {
133        self.edges.write().await.put(id, edge);
134    }
135
136    pub async fn clear(&self) {
137        self.vectors.write().await.clear();
138        self.nodes.write().await.clear();
139        self.edges.write().await.clear();
140        self.query_cache.write().await.clear();
141    }
142
143    pub async fn get_stats(&self) -> (usize, usize, usize, usize) {
144        let vectors_len = self.vectors.read().await.len();
145        let nodes_len = self.nodes.read().await.len();
146        let edges_len = self.edges.read().await.len();
147        let query_cache_len = self.query_cache.read().await.len();
148        (vectors_len, nodes_len, edges_len, query_cache_len)
149    }
150}
151
152/// 向量图数据库
153#[derive(Debug)]
154pub struct VectorGraphDB {
155    pool: DatabasePool,
156    cache: MemoryCache,
157    config: Config,
158}
159
160impl VectorGraphDB {
161    /// 创建新的数据库实例
162    pub async fn new(config: Config) -> Result<Self> {
163        let pool = Self::create_pool(&config).await?;
164        let cache = MemoryCache::new(config.cache.lru_capacity);
165
166        let db = Self {
167            pool,
168            cache,
169            config,
170        };
171
172        // 初始化数据库表
173        db.initialize_tables().await?;
174
175        Ok(db)
176    }
177
178    /// 创建数据库连接池
179    async fn create_pool(config: &Config) -> Result<DatabasePool> {
180        match config.database.database_type {
181            DatabaseType::SQLite => {
182                let pool = sqlx::SqlitePool::connect(&config.database.url)
183                    .await
184                    .map_err(MemoryError::Database)?;
185                Ok(DatabasePool::Sqlite(pool))
186            }
187
188        }
189    }
190
191    /// 初始化数据库表
192    async fn initialize_tables(&self) -> Result<()> {
193        match &self.pool {
194            DatabasePool::Sqlite(pool) => {
195                self.create_sqlite_tables(pool).await?;
196            }
197
198        }
199        Ok(())
200    }
201
202    /// 创建SQLite表
203    async fn create_sqlite_tables(&self, pool: &Pool<Sqlite>) -> Result<()> {
204        let prefix = &self.config.database.table_prefix;
205
206        // 向量表
207        sqlx::query(&format!(
208            r#"
209            CREATE TABLE IF NOT EXISTS {}vectors (
210                id TEXT PRIMARY KEY,
211                embedding BLOB NOT NULL,
212                metadata TEXT,
213                created_at TEXT NOT NULL,
214                updated_at TEXT NOT NULL
215            )
216            "#,
217            prefix
218        ))
219        .execute(pool)
220        .await
221        .map_err(MemoryError::Database)?;
222
223        // 图节点表
224        sqlx::query(&format!(
225            r#"
226            CREATE TABLE IF NOT EXISTS {}graph_nodes (
227                id TEXT PRIMARY KEY,
228                node_type TEXT NOT NULL,
229                properties TEXT,
230                created_at TEXT NOT NULL,
231                updated_at TEXT NOT NULL
232            )
233            "#,
234            prefix
235        ))
236        .execute(pool)
237        .await
238        .map_err(MemoryError::Database)?;
239
240        // 图边表
241        sqlx::query(&format!(
242            r#"
243            CREATE TABLE IF NOT EXISTS {}graph_edges (
244                id TEXT PRIMARY KEY,
245                from_node TEXT NOT NULL,
246                to_node TEXT NOT NULL,
247                edge_type TEXT NOT NULL,
248                weight REAL NOT NULL,
249                properties TEXT,
250                created_at TEXT NOT NULL,
251                updated_at TEXT NOT NULL,
252                FOREIGN KEY (from_node) REFERENCES {}graph_nodes(id),
253                FOREIGN KEY (to_node) REFERENCES {}graph_nodes(id)
254            )
255            "#,
256            prefix, prefix, prefix
257        ))
258        .execute(pool)
259        .await
260        .map_err(MemoryError::Database)?;
261
262        // 创建索引
263        self.create_sqlite_indexes(pool).await?;
264
265        Ok(())
266    }
267
268    /// 创建SQLite索引
269    async fn create_sqlite_indexes(&self, pool: &Pool<Sqlite>) -> Result<()> {
270        let prefix = &self.config.database.table_prefix;
271
272        let indexes = vec![
273            format!("CREATE INDEX IF NOT EXISTS idx_{prefix}vectors_created_at ON {prefix}vectors(created_at)"),
274            format!("CREATE INDEX IF NOT EXISTS idx_{prefix}nodes_type ON {prefix}graph_nodes(node_type)"),
275            format!("CREATE INDEX IF NOT EXISTS idx_{prefix}nodes_created_at ON {prefix}graph_nodes(created_at)"),
276            format!("CREATE INDEX IF NOT EXISTS idx_{prefix}edges_from_node ON {prefix}graph_edges(from_node)"),
277            format!("CREATE INDEX IF NOT EXISTS idx_{prefix}edges_to_node ON {prefix}graph_edges(to_node)"),
278            format!("CREATE INDEX IF NOT EXISTS idx_{prefix}edges_type ON {prefix}graph_edges(edge_type)"),
279            format!("CREATE INDEX IF NOT EXISTS idx_{prefix}edges_weight ON {prefix}graph_edges(weight)"),
280        ];
281
282        for index_sql in indexes {
283            sqlx::query(&index_sql)
284                .execute(pool)
285                .await
286                .map_err(MemoryError::Database)?;
287        }
288
289        Ok(())
290    }
291
292
293
294    /// 插入向量
295    pub async fn insert_vector(&self, vector: &Vector) -> Result<()> {
296        let embedding_bytes = self.vector_to_bytes(&vector.embedding)?;
297        let metadata_json =
298            serde_json::to_string(&vector.metadata).map_err(MemoryError::Serialization)?;
299
300        match &self.pool {
301            DatabasePool::Sqlite(pool) => {
302                sqlx::query(&format!(
303                    "INSERT INTO {}vectors (id, embedding, metadata, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
304                    self.config.database.table_prefix
305                ))
306                .bind(&vector.id)
307                .bind(&embedding_bytes)
308                .bind(&metadata_json)
309                .bind(vector.created_at.to_rfc3339())
310                .bind(vector.updated_at.to_rfc3339())
311                .execute(pool)
312                .await
313                .map_err(MemoryError::Database)?;
314            }
315
316
317        }
318
319        // 更新缓存
320        self.cache
321            .put_vector(vector.id.clone(), vector.clone())
322            .await;
323
324        Ok(())
325    }
326
327    /// 查询向量
328    pub async fn query_vectors(
329        &self,
330        request: &VectorQueryRequest,
331    ) -> Result<Vec<VectorQueryResult>> {
332        let limit = request.limit.unwrap_or(10);
333        let threshold = request.threshold.unwrap_or(0.0);
334
335        // 从数据库获取所有向量（在实际应用中应该使用向量索引）
336        let vectors = self.get_all_vectors().await?;
337
338        let mut results = Vec::new();
339
340        for vector in vectors {
341            let similarity = self.cosine_similarity(&request.query_vector, &vector.embedding);
342
343            if similarity >= threshold {
344                // 应用过滤器
345                if let Some(filters) = &request.filters {
346                    if !self.matches_filters(&vector.metadata, filters) {
347                        continue;
348                    }
349                }
350
351                results.push(VectorQueryResult {
352                    vector,
353                    similarity,
354                    distance: 1.0 - similarity,
355                });
356            }
357        }
358
359        // 按相似度排序
360        results.sort_by(|a, b| b.similarity.partial_cmp(&a.similarity).unwrap());
361        results.truncate(limit);
362
363        Ok(results)
364    }
365
366    /// 获取所有向量
367    async fn get_all_vectors(&self) -> Result<Vec<Vector>> {
368        let mut vectors = Vec::new();
369
370        match &self.pool {
371            DatabasePool::Sqlite(pool) => {
372                let rows = sqlx::query(&format!(
373                    "SELECT id, embedding, metadata, created_at, updated_at FROM {}vectors",
374                    self.config.database.table_prefix
375                ))
376                .fetch_all(pool)
377                .await
378                .map_err(MemoryError::Database)?;
379
380                for row in rows {
381                    let id: String = row.get("id");
382                    let embedding_bytes: Vec<u8> = row.get("embedding");
383                    let metadata_str: String = row.get("metadata");
384                    let created_at_str: String = row.get("created_at");
385                    let updated_at_str: String = row.get("updated_at");
386
387                    let embedding = self.bytes_to_vector(&embedding_bytes)?;
388                    let metadata: HashMap<String, serde_json::Value> =
389                        serde_json::from_str(&metadata_str).map_err(MemoryError::Serialization)?;
390                    let created_at = DateTime::parse_from_rfc3339(&created_at_str)
391                        .map_err(|e| MemoryError::Internal {
392                            message: format!("DateTime parse failed: {}", e),
393                        })?
394                        .with_timezone(&Utc);
395                    let updated_at = DateTime::parse_from_rfc3339(&updated_at_str)
396                        .map_err(|e| MemoryError::Internal {
397                            message: format!("DateTime parse failed: {}", e),
398                        })?
399                        .with_timezone(&Utc);
400
401                    vectors.push(Vector {
402                        id,
403                        embedding,
404                        metadata,
405                        created_at,
406                        updated_at,
407                    });
408                }
409            }
410
411
412        }
413
414        Ok(vectors)
415    }
416
417    /// 插入图节点
418    pub async fn insert_node(&self, node: &GraphNode) -> Result<()> {
419        let properties_json =
420            serde_json::to_string(&node.properties).map_err(MemoryError::Serialization)?;
421
422        match &self.pool {
423            DatabasePool::Sqlite(pool) => {
424                sqlx::query(&format!(
425                    "INSERT INTO {}graph_nodes (id, node_type, properties, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
426                    self.config.database.table_prefix
427                ))
428                .bind(&node.id)
429                .bind(&node.node_type)
430                .bind(&properties_json)
431                .bind(node.created_at.to_rfc3339())
432                .bind(node.updated_at.to_rfc3339())
433                .execute(pool)
434                .await
435                .map_err(MemoryError::Database)?;
436            }
437
438        }
439
440        // 更新缓存
441        self.cache.put_node(node.id.clone(), node.clone()).await;
442
443        Ok(())
444    }
445
446    /// 插入图边
447    pub async fn insert_edge(&self, edge: &GraphEdge) -> Result<()> {
448        let properties_json =
449            serde_json::to_string(&edge.properties).map_err(MemoryError::Serialization)?;
450
451        match &self.pool {
452            DatabasePool::Sqlite(pool) => {
453                sqlx::query(&format!(
454                    "INSERT INTO {}graph_edges (id, from_node, to_node, edge_type, weight, properties, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
455                    self.config.database.table_prefix
456                ))
457                .bind(&edge.id)
458                .bind(&edge.from_node)
459                .bind(&edge.to_node)
460                .bind(&edge.edge_type)
461                .bind(edge.weight)
462                .bind(&properties_json)
463                .bind(edge.created_at.to_rfc3339())
464                .bind(edge.updated_at.to_rfc3339())
465                .execute(pool)
466                .await
467                .map_err(MemoryError::Database)?;
468            }
469
470        }
471
472        // 更新缓存
473        self.cache.put_edge(edge.id.clone(), edge.clone()).await;
474
475        Ok(())
476    }
477
478    /// 查询图
479    pub async fn query_graph(&self, request: &GraphQueryRequest) -> Result<GraphQueryResult> {
480        let mut nodes = Vec::new();
481        let mut edges = Vec::new();
482        let mut paths = Vec::new();
483
484        // 简单的图遍历实现（在实际应用中应该使用更高效的算法）
485        for start_node in &request.start_nodes {
486            let (node_results, edge_results, path_results) = self
487                .traverse_graph(
488                    start_node,
489                    request.max_depth.unwrap_or(3),
490                    &request.edge_types,
491                    &request.filters,
492                )
493                .await?;
494
495            nodes.extend(node_results);
496            edges.extend(edge_results);
497            paths.extend(path_results);
498        }
499
500        // 去重
501        nodes.sort_by(|a, b| a.id.cmp(&b.id));
502        nodes.dedup_by(|a, b| a.id == b.id);
503
504        edges.sort_by(|a, b| a.id.cmp(&b.id));
505        edges.dedup_by(|a, b| a.id == b.id);
506
507        // 应用限制
508        if let Some(limit) = request.limit {
509            nodes.truncate(limit);
510            edges.truncate(limit);
511            paths.truncate(limit);
512        }
513
514        Ok(GraphQueryResult {
515            nodes,
516            edges,
517            paths,
518        })
519    }
520
521    /// 图遍历
522    async fn traverse_graph(
523        &self,
524        start_node: &str,
525        max_depth: usize,
526        edge_types: &Option<Vec<String>>,
527        filters: &Option<HashMap<String, serde_json::Value>>,
528    ) -> Result<(Vec<GraphNode>, Vec<GraphEdge>, Vec<Vec<String>>)> {
529        let mut visited_nodes = std::collections::HashSet::new();
530        let mut result_nodes = Vec::new();
531        let mut result_edges = Vec::new();
532        let mut paths = Vec::new();
533
534        let mut queue = std::collections::VecDeque::new();
535        queue.push_back((start_node.to_string(), 0, vec![start_node.to_string()]));
536
537        while let Some((current_node, depth, path)) = queue.pop_front() {
538            if depth > max_depth || visited_nodes.contains(&current_node) {
539                continue;
540            }
541
542            visited_nodes.insert(current_node.clone());
543
544            // 获取当前节点
545            if let Ok(node) = self.get_node(&current_node).await {
546                // 应用过滤器
547                if let Some(filters) = filters {
548                    if !self.matches_filters(&node.properties, filters) {
549                        continue;
550                    }
551                }
552                result_nodes.push(node);
553            }
554
555            // 获取相邻边
556            let adjacent_edges = self.get_adjacent_edges(&current_node, edge_types).await?;
557
558            for edge in adjacent_edges {
559                result_edges.push(edge.clone());
560
561                let next_node = if edge.from_node == current_node {
562                    &edge.to_node
563                } else {
564                    &edge.from_node
565                };
566
567                if !visited_nodes.contains(next_node) {
568                    let mut new_path = path.clone();
569                    new_path.push(next_node.clone());
570                    queue.push_back((next_node.clone(), depth + 1, new_path.clone()));
571
572                    if depth < max_depth {
573                        paths.push(new_path);
574                    }
575                }
576            }
577        }
578
579        Ok((result_nodes, result_edges, paths))
580    }
581
582    /// 更新向量
583    pub async fn update_vector(&self, vector: &Vector) -> Result<()> {
584        let embedding_bytes = self.vector_to_bytes(&vector.embedding)?;
585        let metadata_json =
586            serde_json::to_string(&vector.metadata).map_err(MemoryError::Serialization)?;
587
588        match &self.pool {
589            DatabasePool::Sqlite(pool) => {
590                sqlx::query(&format!(
591                    "UPDATE {}vectors SET embedding = ?, metadata = ?, updated_at = ? WHERE id = ?",
592                    self.config.database.table_prefix
593                ))
594                .bind(&embedding_bytes)
595                .bind(&metadata_json)
596                .bind(vector.updated_at.to_rfc3339())
597                .bind(&vector.id)
598                .execute(pool)
599                .await
600                .map_err(MemoryError::Database)?;
601            }
602        }
603
604        // 更新缓存
605        self.cache.put_vector(vector.id.clone(), vector.clone()).await;
606
607        Ok(())
608    }
609
610    /// 更新节点
611    pub async fn update_node(&self, node: &GraphNode) -> Result<()> {
612        let properties_json =
613            serde_json::to_string(&node.properties).map_err(MemoryError::Serialization)?;
614
615        match &self.pool {
616            DatabasePool::Sqlite(pool) => {
617                sqlx::query(&format!(
618                    "UPDATE {}graph_nodes SET node_type = ?, properties = ?, updated_at = ? WHERE id = ?",
619                    self.config.database.table_prefix
620                ))
621                .bind(&node.node_type)
622                .bind(&properties_json)
623                .bind(node.updated_at.to_rfc3339())
624                .bind(&node.id)
625                .execute(pool)
626                .await
627                .map_err(MemoryError::Database)?;
628            }
629        }
630
631        // 更新缓存
632        self.cache.put_node(node.id.clone(), node.clone()).await;
633
634        Ok(())
635    }
636
637    /// 获取向量
638    pub async fn get_vector(&self, vector_id: &str) -> Result<Vector> {
639        // 先检查缓存
640        if let Some(vector) = self.cache.get_vector(vector_id).await {
641            return Ok(vector);
642        }
643
644        match &self.pool {
645            DatabasePool::Sqlite(pool) => {
646                let row = sqlx::query(&format!(
647                    "SELECT id, embedding, metadata, created_at, updated_at FROM {}vectors WHERE id = ?",
648                    self.config.database.table_prefix
649                ))
650                .bind(vector_id)
651                .fetch_one(pool)
652                .await
653                .map_err(MemoryError::Database)?;
654
655                let id: String = row.get("id");
656                let embedding_bytes: Vec<u8> = row.get("embedding");
657                let metadata_str: String = row.get("metadata");
658                let created_at_str: String = row.get("created_at");
659                let updated_at_str: String = row.get("updated_at");
660
661                let embedding = self.bytes_to_vector(&embedding_bytes)?;
662                let metadata: HashMap<String, serde_json::Value> =
663                    serde_json::from_str(&metadata_str).map_err(MemoryError::Serialization)?;
664                let created_at = DateTime::parse_from_rfc3339(&created_at_str)
665                    .map_err(|e| MemoryError::Internal {
666                        message: format!("DateTime parse failed: {}", e),
667                    })?
668                    .with_timezone(&Utc);
669                let updated_at = DateTime::parse_from_rfc3339(&updated_at_str)
670                    .map_err(|e| MemoryError::Internal {
671                        message: format!("DateTime parse failed: {}", e),
672                    })?
673                    .with_timezone(&Utc);
674
675                let vector = Vector {
676                    id,
677                    embedding,
678                    metadata,
679                    created_at,
680                    updated_at,
681                };
682
683                // 更新缓存
684                self.cache.put_vector(vector.id.clone(), vector.clone()).await;
685
686                Ok(vector)
687            }
688        }
689    }
690
691    /// 获取节点
692    async fn get_node(&self, node_id: &str) -> Result<GraphNode> {
693        // 先检查缓存
694        if let Some(node) = self.cache.get_node(node_id).await {
695            return Ok(node);
696        }
697
698        match &self.pool {
699            DatabasePool::Sqlite(pool) => {
700                let row = sqlx::query(&format!(
701                    "SELECT id, node_type, properties, created_at, updated_at FROM {}graph_nodes WHERE id = ?",
702                    self.config.database.table_prefix
703                ))
704                .bind(node_id)
705                .fetch_one(pool)
706                .await
707                .map_err(MemoryError::Database)?;
708
709                let id: String = row.get("id");
710                let node_type: String = row.get("node_type");
711                let properties_str: String = row.get("properties");
712                let created_at_str: String = row.get("created_at");
713                let updated_at_str: String = row.get("updated_at");
714
715                let properties: HashMap<String, serde_json::Value> =
716                    serde_json::from_str(&properties_str).map_err(MemoryError::Serialization)?;
717                let created_at = DateTime::parse_from_rfc3339(&created_at_str)
718                    .map_err(|e| MemoryError::Internal {
719                        message: format!("DateTime parse failed: {}", e),
720                    })?
721                    .with_timezone(&Utc);
722                let updated_at = DateTime::parse_from_rfc3339(&updated_at_str)
723                    .map_err(|e| MemoryError::Internal {
724                        message: format!("DateTime parse failed: {}", e),
725                    })?
726                    .with_timezone(&Utc);
727
728                let node = GraphNode {
729                    id,
730                    node_type,
731                    properties,
732                    created_at,
733                    updated_at,
734                };
735
736                // 更新缓存
737                self.cache.put_node(node.id.clone(), node.clone()).await;
738
739                Ok(node)
740            }
741        }
742    }
743
744    /// 获取相邻边
745    async fn get_adjacent_edges(
746        &self,
747        node_id: &str,
748        edge_types: &Option<Vec<String>>,
749    ) -> Result<Vec<GraphEdge>> {
750        let mut edges = Vec::new();
751
752        let type_filter = if let Some(types) = edge_types {
753            format!(
754                " AND edge_type IN ({})",
755                types.iter().map(|_| "?").collect::<Vec<_>>().join(", ")
756            )
757        } else {
758            String::new()
759        };
760
761        match &self.pool {
762            DatabasePool::Sqlite(pool) => {
763                let sql = format!(
764                    "SELECT id, from_node, to_node, edge_type, weight, properties, created_at, updated_at FROM {}graph_edges WHERE (from_node = ? OR to_node = ?){}",
765                    self.config.database.table_prefix, type_filter
766                );
767
768                let mut query = sqlx::query(&sql).bind(node_id).bind(node_id);
769
770                if let Some(types) = edge_types {
771                    for edge_type in types {
772                        query = query.bind(edge_type);
773                    }
774                }
775
776                let rows = query.fetch_all(pool).await.map_err(MemoryError::Database)?;
777
778                for row in rows {
779                    let id: String = row.get("id");
780                    let from_node: String = row.get("from_node");
781                    let to_node: String = row.get("to_node");
782                    let edge_type: String = row.get("edge_type");
783                    let weight: f32 = row.get("weight");
784                    let properties_str: String = row.get("properties");
785                    let created_at_str: String = row.get("created_at");
786                    let updated_at_str: String = row.get("updated_at");
787
788                    let properties: HashMap<String, serde_json::Value> =
789                        serde_json::from_str(&properties_str)
790                            .map_err(MemoryError::Serialization)?;
791                    let created_at = DateTime::parse_from_rfc3339(&created_at_str)
792                        .map_err(|e| MemoryError::Internal {
793                            message: format!("DateTime parse failed: {}", e),
794                        })?
795                        .with_timezone(&Utc);
796                    let updated_at = DateTime::parse_from_rfc3339(&updated_at_str)
797                        .map_err(|e| MemoryError::Internal {
798                            message: format!("DateTime parse failed: {}", e),
799                        })?
800                        .with_timezone(&Utc);
801
802                    edges.push(GraphEdge {
803                        id,
804                        from_node,
805                        to_node,
806                        edge_type,
807                        weight,
808                        properties,
809                        created_at,
810                        updated_at,
811                    });
812                }
813            }
814        }
815
816        Ok(edges)
817    }
818
819    /// 清除缓存
820    pub async fn clear_cache(&self) {
821        self.cache.clear().await;
822    }
823
824    /// 获取统计信息
825    pub async fn get_stats(&self) -> Result<(u64, u64, u64, (usize, usize, usize, usize))> {
826        let (vector_count, node_count, edge_count) = match &self.pool {
827            DatabasePool::Sqlite(pool) => {
828                let prefix = &self.config.database.table_prefix;
829
830                let vector_count: i64 =
831                    sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {}vectors", prefix))
832                        .fetch_one(pool)
833                        .await
834                        .map_err(MemoryError::Database)?;
835
836                let node_count: i64 =
837                    sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {}graph_nodes", prefix))
838                        .fetch_one(pool)
839                        .await
840                        .map_err(MemoryError::Database)?;
841
842                let edge_count: i64 =
843                    sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {}graph_edges", prefix))
844                        .fetch_one(pool)
845                        .await
846                        .map_err(MemoryError::Database)?;
847
848                (vector_count as u64, node_count as u64, edge_count as u64)
849            }
850        };
851
852        let cache_stats = self.cache.get_stats().await;
853
854        Ok((vector_count, node_count, edge_count, cache_stats))
855    }
856
857    // 辅助函数
858
859    /// 计算余弦相似度
860    fn cosine_similarity(&self, a: &[f32], b: &[f32]) -> f32 {
861        if a.len() != b.len() {
862            return 0.0;
863        }
864
865        let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
866        let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
867        let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
868
869        if norm_a == 0.0 || norm_b == 0.0 {
870            0.0
871        } else {
872            dot_product / (norm_a * norm_b)
873        }
874    }
875
876    /// 向量转字节
877    fn vector_to_bytes(&self, vector: &[f32]) -> Result<Vec<u8>> {
878        let mut bytes = Vec::with_capacity(vector.len() * 4);
879        for &value in vector {
880            bytes.extend_from_slice(&value.to_le_bytes());
881        }
882        Ok(bytes)
883    }
884
885    /// 字节转向量
886    fn bytes_to_vector(&self, bytes: &[u8]) -> Result<Vec<f32>> {
887        if bytes.len() % 4 != 0 {
888            return Err(MemoryError::InvalidVectorDimension {
889                expected: bytes.len() - (bytes.len() % 4),
890                actual: bytes.len(),
891            });
892        }
893
894        let mut vector = Vec::with_capacity(bytes.len() / 4);
895        for chunk in bytes.chunks_exact(4) {
896            let value = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
897            vector.push(value);
898        }
899
900        Ok(vector)
901    }
902
903    /// 检查过滤器匹配
904    fn matches_filters(
905        &self,
906        metadata: &HashMap<String, serde_json::Value>,
907        filters: &HashMap<String, serde_json::Value>,
908    ) -> bool {
909        for (key, expected_value) in filters {
910            if let Some(actual_value) = metadata.get(key) {
911                if actual_value != expected_value {
912                    return false;
913                }
914            } else {
915                return false;
916            }
917        }
918        true
919    }
920}
921
922#[cfg(test)]
923mod tests {
924    use super::*;
925    use tempfile::NamedTempFile;
926
927    async fn create_test_db() -> VectorGraphDB {
928        let temp_file = NamedTempFile::new().unwrap();
929        let db_url = format!("sqlite://{}", temp_file.path().display());
930
931        let config = Config {
932            database: crate::config::DatabaseConfig {
933                url: db_url,
934                ..Default::default()
935            },
936            ..Default::default()
937        };
938
939        VectorGraphDB::new(config).await.unwrap()
940    }
941
942    #[tokio::test]
943    async fn test_vector_operations() {
944        let db = create_test_db().await;
945
946        let vector = Vector {
947            id: "test_vector".to_string(),
948            embedding: vec![0.1, 0.2, 0.3, 0.4],
949            metadata: HashMap::new(),
950            created_at: Utc::now(),
951            updated_at: Utc::now(),
952        };
953
954        // 插入向量
955        db.insert_vector(&vector).await.unwrap();
956
957        // 查询向量
958        let query_request = VectorQueryRequest {
959            query_vector: vec![0.1, 0.2, 0.3, 0.4],
960            limit: Some(10),
961            threshold: Some(0.5),
962            filters: None,
963        };
964
965        let results = db.query_vectors(&query_request).await.unwrap();
966        assert!(!results.is_empty());
967        assert_eq!(results[0].vector.id, "test_vector");
968    }
969
970    #[tokio::test]
971    async fn test_graph_operations() {
972        let db = create_test_db().await;
973
974        // 创建节点
975        let node1 = GraphNode {
976            id: "node1".to_string(),
977            node_type: "test".to_string(),
978            properties: HashMap::new(),
979            created_at: Utc::now(),
980            updated_at: Utc::now(),
981        };
982
983        let node2 = GraphNode {
984            id: "node2".to_string(),
985            node_type: "test".to_string(),
986            properties: HashMap::new(),
987            created_at: Utc::now(),
988            updated_at: Utc::now(),
989        };
990
991        // 插入节点
992        db.insert_node(&node1).await.unwrap();
993        db.insert_node(&node2).await.unwrap();
994
995        // 创建边
996        let edge = GraphEdge {
997            id: "edge1".to_string(),
998            from_node: "node1".to_string(),
999            to_node: "node2".to_string(),
1000            edge_type: "connects".to_string(),
1001            weight: 0.8,
1002            properties: HashMap::new(),
1003            created_at: Utc::now(),
1004            updated_at: Utc::now(),
1005        };
1006
1007        // 插入边
1008        db.insert_edge(&edge).await.unwrap();
1009
1010        // 查询图
1011        let query_request = GraphQueryRequest {
1012            start_nodes: vec!["node1".to_string()],
1013            edge_types: None,
1014            max_depth: Some(2),
1015            limit: Some(10),
1016            filters: None,
1017        };
1018
1019        let results = db.query_graph(&query_request).await.unwrap();
1020        assert!(!results.nodes.is_empty());
1021        assert!(!results.edges.is_empty());
1022    }
1023
1024    #[tokio::test]
1025    async fn test_cosine_similarity() {
1026        let db = create_test_db().await;
1027
1028        let vec1 = vec![1.0, 0.0, 0.0];
1029        let vec2 = vec![1.0, 0.0, 0.0];
1030        let vec3 = vec![0.0, 1.0, 0.0];
1031
1032        assert!((db.cosine_similarity(&vec1, &vec2) - 1.0).abs() < 1e-6);
1033        assert!((db.cosine_similarity(&vec1, &vec3) - 0.0).abs() < 1e-6);
1034    }
1035
1036    #[tokio::test]
1037    async fn test_vector_serialization() {
1038        let db = create_test_db().await;
1039
1040        let original_vector = vec![1.5, -2.3, 0.0, 42.7];
1041        let bytes = db.vector_to_bytes(&original_vector).unwrap();
1042        let restored_vector = db.bytes_to_vector(&bytes).unwrap();
1043
1044        assert_eq!(original_vector, restored_vector);
1045    }
1046}
ai00_mem/database.rs

ai00_mem/
database.rs
