neomemx 0.1.2

A high-performance memory library for AI agents with semantic search
Documentation
//! Search and reranking

use std::collections::HashMap;

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

use crate::core::{ScopeIdentifiers, StoredFact};
use crate::error::Result;
use crate::vector_store::base::Filters;

pub mod semantic;

/// Trait for result reranking
#[async_trait]
pub trait ResultReranker: Send + Sync {
    /// Rerank search results
    async fn rerank(
        &self,
        query: &str,
        facts: Vec<StoredFact>,
        limit: Option<usize>,
    ) -> Result<Vec<StoredFact>>;
}

/// Kinds of memory/search that the engine can support.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MemoryKind {
    Semantic,
    Episodic,
    Procedural,
}

/// Common request payload for search strategies.
#[derive(Debug, Clone)]
pub struct SearchRequest {
    pub query: String,
    pub scope: ScopeIdentifiers,
    pub limit: usize,
    /// Optional additional filters (e.g., metadata conditions).
    pub filters: Option<Filters>,
    /// Whether to apply reranking after retrieval.
    pub rerank: bool,
    /// What memory flavor this request targets.
    pub kind: MemoryKind,
}

impl SearchRequest {
    pub fn new(query: impl Into<String>, scope: ScopeIdentifiers, limit: usize) -> Self {
        Self {
            query: query.into(),
            scope,
            limit,
            filters: None,
            rerank: false,
            kind: MemoryKind::Semantic,
        }
    }

    pub fn with_filters(mut self, filters: HashMap<String, serde_json::Value>) -> Self {
        self.filters = Some(filters);
        self
    }

    pub fn with_rerank(mut self, rerank: bool) -> Self {
        self.rerank = rerank;
        self
    }

    pub fn with_kind(mut self, kind: MemoryKind) -> Self {
        self.kind = kind;
        self
    }
}

/// Standardized search result.
#[derive(Debug, Clone)]
pub struct SearchResult {
    pub facts: Vec<StoredFact>,
}

impl SearchResult {
    pub fn new(facts: Vec<StoredFact>) -> Self {
        Self { facts }
    }

    /// Convenience to limit results.
    pub fn truncate(&mut self, limit: usize) {
        self.facts.truncate(limit);
    }
}

/// Trait for pluggable memory search strategies.
#[async_trait]
pub trait MemorySearch: Send + Sync {
    /// Execute a search given a request.
    async fn search(&self, request: SearchRequest) -> Result<SearchResult>;
}