neomemx 0.1.2

A high-performance memory library for AI agents with semantic search
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

//! Base trait for vector store implementations

use std::collections::HashMap;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};

use crate::error::Result;

/// Output data from vector store operations.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OutputData {
    /// The unique identifier of the vector.
    pub id: String,
    /// The similarity score.
    pub score: Option<f32>,
    /// The associated payload/metadata.
    pub payload: HashMap<String, serde_json::Value>,
}

impl OutputData {
    /// Creates a new OutputData instance.
    pub fn new(
        id: String,
        score: Option<f32>,
        payload: HashMap<String, serde_json::Value>,
    ) -> Self {
        Self { id, score, payload }
    }

    /// Retrieves a string value from the payload by key.
    pub fn get_string(&self, key: &str) -> Option<String> {
        self.payload
            .get(key)
            .and_then(|v| v.as_str())
            .map(String::from)
    }

    /// Retrieves the "data" field from the payload.
    pub fn get_data(&self) -> Option<String> {
        self.get_string("data")
    }
}

/// Filters for vector store queries
pub type Filters = HashMap<String, serde_json::Value>;

/// Base trait for vector store implementations.
#[async_trait]
pub trait VectorStoreBase: Send + Sync {
    /// Creates a new collection (if supported).
    async fn create_collection(&self, name: &str) -> Result<()>;
    /// Inserts vectors into the store.
    async fn insert(
        &self,
        vectors: Vec<Vec<f32>>,
        payloads: Option<Vec<HashMap<String, serde_json::Value>>>,
        ids: Option<Vec<String>>,
    ) -> Result<()>;
    /// Searches for similar vectors.
    async fn search(
        &self,
        query: &str,
        vectors: &[f32],
        limit: usize,
        filters: Option<Filters>,
    ) -> Result<Vec<OutputData>>;
    /// Deletes a vector by ID.
    async fn delete(&self, vector_id: &str) -> Result<()>;

    /// Deletes multiple vectors by ID.
    async fn delete_batch(&self, vector_ids: &[String]) -> Result<()> {
        use futures::future;
        let futures: Vec<_> = vector_ids.iter().map(|id| self.delete(id)).collect();
        future::try_join_all(futures).await?;
        Ok(())
    }

    /// Updates a vector and/or its payload.
    async fn update(
        &self,
        vector_id: &str,
        vector: Option<Vec<f32>>,
        payload: Option<HashMap<String, serde_json::Value>>,
    ) -> Result<()>;
    /// Retrieves a vector by ID.
    async fn get(&self, vector_id: &str) -> Result<Option<OutputData>>;

    /// Retrieves multiple vectors by ID.
    async fn get_batch(&self, vector_ids: &[String]) -> Result<Vec<OutputData>> {
        use futures::future;
        let futures: Vec<_> = vector_ids.iter().map(|id| self.get(id)).collect();
        let results = future::try_join_all(futures).await?;
        Ok(results.into_iter().flatten().collect())
    }
    /// Lists all collections.
    async fn list_collections(&self) -> Result<Vec<String>>;
    /// Deletes the current collection.
    async fn delete_collection(&self) -> Result<()>;
    /// Retrieves collection information.
    async fn collection_info(&self) -> Result<serde_json::Value>;
    /// Lists vectors in the store.
    async fn list(&self, filters: Option<Filters>, limit: usize) -> Result<Vec<OutputData>>;
    /// Resets the vector store.
    async fn reset(&self) -> Result<()>;
}