How the memory_query Tool Works in ai-memory: Hybrid Search Architecture Explained

The memory_query tool executes a hybrid lexical-entity-graph search with optional vector embedding and LLM reranking to surface ranked historical observations from indexed project sessions.

The memory_query tool provides the primary retrieval interface for the akitaonrails/ai-memory repository, enabling AI agents to query accumulated project context through an MCP (Model Context Protocol) endpoint. Implemented in crates/ai-memory-mcp/src/server.rs, this tool orchestrates a multi-stage pipeline that gracefully degrades from semantic vector search to pure lexical matching when embedding services are unavailable. Its architecture reflects a design philosophy that prioritizes result relevance through combined full-text, entity-graph, and optional neural ranking strategies.

MCP Tool Architecture and Entry Points

The memory_query tool is registered as an MCP tool that accepts query, workspace, project, and limit arguments. When an agent invokes tools/call with the name memory_query, the server implementation resolves the workspace and project identifiers into concrete WorkspaceId and ProjectId values before executing the search pipeline.

On the client side, the memory_query helper function in evals/src/retrieval/query.rs wraps the JSON-RPC transport, deserializes the response, and flattens the nested hit structures into a unified Vec<Retrieved> for downstream consumption.

Step-by-Step Execution Workflow

Argument Parsing and Scope Resolution

The tool first validates the incoming MCP request for required parameters. If the caller specifies a workspace and project, the server resolves these to internal identifiers; otherwise, it falls back to a default global marker that searches across all indexed content. This scope resolution determines which indexed segments the subsequent hybrid search will query.

Optional Vector Embedding

If a vector embedder is configured in the server instance, the tool attempts to generate a query embedding via the embed_query method. When the embedder fails—due to provider errors, model unavailability, or timeout—the system degrades gracefully to a pure lexical and entity-graph search without terminating the request.

async fn embed_query(&self, query: &str) -> Option<Vec<f32>> {
    let Some(embedder) = &self.embedder else { return None };
    match embedder.embed(query).await {
        Ok(qv) => Some(qv),
        Err(e) => {
            tracing::warn!(provider = embedder.provider(),
                           model = embedder.model(),
                           error = %e,
                           "embedder failed; degrading memory_query to FTS5 + entity + graph");
            None
        }
    }
}

This graceful degradation ensures that memory_query remains functional even when embedding infrastructure is compromised, falling back to the FTS5 and graph-based retrieval paths.

Hybrid Search Execution

The core retrieval logic resides in the search_project function, which invokes the store's hybrid_search or hybrid_search_explained method depending on whether the explain flag is enabled. This search merges three distinct ranking streams:

• FTS5 lexical match for keyword and phrase relevance
• Typed-entity match for specific identifiers like session IDs and page edges
• Graph-based relevance derived from entity-timeline walks

async fn search_project(
    &self,
    workspace_id: WorkspaceId,
    project_id: ProjectId,
    options: ProjectSearchOptions<'_>,
) -> StoreResult<Vec<(PageHit, Option<SearchExplain>)>> {
    // ...
    let fused = if options.explain {
        self.reader.hybrid_search_explained(...)
            .await?
            .into_iter()
            .map(|(hit, details)| (hit, Some(details)))
            .collect()
    } else {
        self.reader.hybrid_search(...)
            .await?
            .into_iter()
            .map(|hit| (hit, None))
            .collect()
    };
    Ok(fused)
}

The hybrid_search implementation in crates/ai-memory-store/src/fts_query.rs handles the actual index interrogation, combining these signals into an initial ranked candidate list.

Optional LLM Reranking

When an LLM reranker is configured, the tool submits the top candidate results to the reranker for relevance refinement. If the reranker times out or fails, the tool preserves the pre-rerank order rather than failing the request. After reranking (or if skipped), the list truncates to the requested limit parameter before serialization.

Result Assembly and Response Format

The tool constructs a JSON-RPC response containing a content field with a JSON object holding two arrays: hits (page paths ranked by relevance) and raw_hits (session IDs of raw observations). This structure allows agents to distinguish between structured page content and raw telemetry entries.

Client-Side Response Processing

The client implementation in evals/src/retrieval/query.rs unwraps the MCP response and deserializes the JSON payload. The memory_query async function extracts the text content from the first content array element, parses it into a QueryResponse struct, and passes it to the flatten routine.

pub async fn memory_query(
    client: &reqwest::Client,
    base_url: &str,
    workspace: &str,
    project: &str,
    query: &str,
    limit: usize,
) -> Result<Vec<Retrieved>> {
    // ...
    let text = result
        .get("content")
        .and_then(|c| c.as_array())
        .and_then(|c| c.first())
        .and_then(|c| c.get("text"))
        .and_then(|t| t.as_str())
        .ok_or_else(|| anyhow!("memory_query returned no text content: {result}"))?;
    let parsed: QueryResponse =
        serde_json::from_str(text).with_context(|| format!("parsing memory_query JSON: {text}"))?;
    Ok(flatten(parsed))
}

The flatten function converts page paths into optional session UUIDs and appends raw observation hits, ensuring page hits precede raw hits while maintaining internal ranking order:

fn flatten(resp: QueryResponse) -> Vec<Retrieved> {
    let mut out = Vec::new();
    for hit in resp.hits {
        out.push(Retrieved { session_uuid: session_uuid_from_path(&hit.path) });
    }
    for hit in resp.raw_hits {
        out.push(Retrieved { session_uuid: Some(hit.session_id) });
    }
    out
}

fn session_uuid_from_path(path: &str) -> Option<uuid::Uuid> {
    let stem = path.strip_prefix("sessions/")?.strip_suffix(".md")?;
    uuid::Uuid::parse_str(stem).ok()
}

Practical Usage Examples

To invoke memory_query from a Rust client application:

// Direct client invocation
let client = reqwest::Client::new();
let results = memory_query(
    &client,
    "http://127.0.0.1:49374",
    "my_workspace",
    "my_project",
    "how does memory_query work?",
    10,
).await?;

for hit in results {
    println!("Session UUID: {:?}", hit.session_uuid);
}

For debugging or direct MCP integration, send a raw JSON-RPC request:

curl -X POST http://localhost:49374/mcp \
  -H "Authorization: Bearer <TOKEN>" \
  -H "Accept: application/json, text/event-stream" \
  -d '{
        "jsonrpc":"2.0","id":1,"method":"tools/call",
        "params":{"name":"memory_query","arguments":{
            "query":"how does memory_query work?",
            "workspace":"my_workspace",
            "project":"my_project",
            "limit":10}}}'

Summary

• memory_query combines FTS5 lexical search, entity-graph traversal, and optional vector embedding into a unified retrieval pipeline for project sessions.
• The tool degrades gracefully when embedders or rerankers fail, ensuring high availability through fallback to lexical and graph-based search.
• Scope resolution supports both project-specific and global queries via workspace and project parameters.
• The server implementation in crates/ai-memory-mcp/src/server.rs handles embedding, hybrid search, and optional LLM reranking before returning structured JSON-RPC responses.
• Client utilities in evals/src/retrieval/query.rs flatten the response into a unified vector of Retrieved structs, parsing session UUIDs from page paths.

Frequently Asked Questions

What happens if the vector embedder fails during a memory_query call?

If the embedder returns an error or times out, the embed_query function returns None and logs a warning with the provider and model details. The tool continues execution using only the FTS5 lexical and entity-graph search components without terminating the request, ensuring the query still returns relevant results.

How does memory_query handle project scoping versus global search?

The tool accepts workspace and project arguments that resolve to internal WorkspaceId and ProjectId values. If the caller provides these identifiers, the search_project function restricts the hybrid search to that specific project index. If omitted, the server falls back to a global marker that searches across all indexed workspaces and projects.

What is the difference between hits and raw_hits in the response?

The hits array contains page paths (structured markdown files) ranked by the hybrid search algorithm, while raw_hits contains session IDs of raw observations (unstructured telemetry entries). The client-side flatten routine preserves this distinction while ensuring page hits appear before raw hits in the final ordered list consumed by agents.

Can memory_query operate without an LLM reranker?

Yes. LLM reranking is entirely optional. If no reranker is configured or if the reranker request fails, the tool returns the results ranked by the hybrid search fusion algorithm (combining FTS5, entity, and vector scores). The limit truncation occurs after the reranking step if configured, or immediately after the hybrid search if not.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}