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

//! Query builder for fluent API

use crate::core::{RetrievalResult, ScopeIdentifiers};
use crate::engine::NeomemxEngine;
use crate::error::NeomemxError;
use crate::error::Result;
use std::collections::HashMap;

/// Query filters
pub type QueryFilters = HashMap<String, serde_json::Value>;

/// Builder for query operations
pub struct QueryBuilder {
    engine: NeomemxEngine,
    query: Option<String>,
    scope: Option<ScopeIdentifiers>,
    limit: usize,
    filters: Option<QueryFilters>,
    rerank: bool,
    graph_expansion: bool,
}

impl QueryBuilder {
    /// Create a new query builder
    pub(crate) fn new(engine: NeomemxEngine) -> Self {
        Self {
            engine,
            query: None,
            scope: None,
            limit: 10,
            filters: None,
            rerank: false,
            graph_expansion: false,
        }
    }

    /// Set the semantic query string
    pub fn query(mut self, query: impl Into<String>) -> Self {
        self.query = Some(query.into());
        self
    }

    /// Set scoping identifiers
    pub fn with_scope(mut self, scope: ScopeIdentifiers) -> Self {
        self.scope = Some(scope);
        self
    }

    /// Set maximum number of results
    pub fn limit(mut self, limit: usize) -> Self {
        self.limit = limit;
        self
    }

    /// Add additional filters
    pub fn with_filters(mut self, filters: QueryFilters) -> Self {
        self.filters = Some(filters);
        self
    }

    /// Enable/disable reranking
    pub fn enable_rerank(mut self, enabled: bool) -> Self {
        self.rerank = enabled;
        self
    }

    /// Enable/disable graph expansion
    pub fn enable_graph_expansion(mut self, enabled: bool) -> Self {
        self.graph_expansion = enabled;
        self
    }

    /// Execute the query
    pub async fn execute(self) -> Result<RetrievalResult> {
        let scope = self.scope.ok_or_else(|| NeomemxError::required("scope"))?;

        scope.validate()?;

        if let Some(query) = self.query {
            // Semantic search
            self.engine
                .search_internal(
                    &query,
                    scope,
                    self.limit,
                    self.filters,
                    self.rerank,
                    self.graph_expansion,
                )
                .await
        } else {
            // List all
            self.engine
                .list_internal(scope, self.limit, self.filters)
                .await
        }
    }
}