API Reference

Core

class core.AgenticCollection(name: str = 'default', dimension: int | None = None, indexer_type: str = 'faiss', sparse_indexer_type: str | None = None, truncate_dim: int | None = None, **indexer_kwargs)

Bases: Collection

Specialized collection for autonomous agent memory. Supports search loops, feedback, and self-healing ranking.

agentic_search(query_vector: ndarray, top_k: int = 5, feedback_weight: float = 1.0) → List[Tuple[Dict[str, Any], float]]

Search with additive self-healing logic using stored feedback scores.

feedback(doc_id: str, label: str = 'good')

Provide feedback on a retrieval result. Adjusts metadata to influence future rankings.

class core.CacheCollection(name: str = 'default', dimension: int | None = None, indexer_type: str = 'faiss', sparse_indexer_type: str | None = None, truncate_dim: int | None = None, **indexer_kwargs)

Bases: Collection

Specialized collection for Retrieval-Augmented KV Caching (RA-KVC). Supports storing high-dimensional activation tensors.

add_cache(vectors: ndarray | List[List[float]], activations: Dict[str, ndarray], metadata: List[Dict[str, Any]] | None = None, quantize: bool = False)

Add embeddings and their associated KV cache activations.

get_cache(metadata: Dict[str, Any]) → Dict[str, ndarray]

Retrieve activations for a given metadata result.

class core.ClusterCollection(n_clusters: int = 10, **kwargs)

Bases: Collection

Specialized collection for ClusterKV-style optimizations. Implements semantic clustering of vectors for improved retrieval throughput.

cluster_data()

Perform semantic clustering on the current collection data.

search_clustered(query: ndarray, top_k: int = 5) → List[Tuple[Dict[str, Any], float]]

Search for vectors by first identifying the most relevant cluster.

class core.Collection(name: str = 'default', dimension: int | None = None, indexer_type: str = 'faiss', sparse_indexer_type: str | None = None, truncate_dim: int | None = None, **indexer_kwargs)

Bases: object

Primary interface for managing embeddings and metadata. Provides a high-level API for indexing, search, and I/O.

add(vectors: ndarray | List[List[float]], metadata: List[Dict[str, Any]] | None = None)

Add vectors and metadata to the collection.

add_images(image_paths: List[str], model: str = 'openai/clip-vit-base-patch32', metadata: List[Dict[str, Any]] | None = None)

Embed and add images to the collection.

benchmark(indexers: List[str] | None = None, top_k: int = 5)

Benchmark multiple indexers on the current collection data.

Parameters:

• indexers – List of indexer names to compare (e.g. [“faiss”, “hnswlib”]). If None, benchmarks all available indexers.

• top_k – Number of neighbors to search for during benchmark.

evaluate(indexer_type: str = 'faiss-hnsw', top_k: int = 10, **kwargs) → Dict[str, Any]

Evaluate an indexer’s recall and latency against an exact search baseline.

Returns:

Dictionary with ‘recall’ and ‘latency_ms’ metrics.

export_to_production(backend: str, connection_url: str, collection_name: str | None = None)

One-click export from local Embenx collection to production clusters. Supported backends: ‘qdrant’, ‘milvus’.

classmethod from_numpy(path: str, **kwargs)

Load a collection from a .npy or .npz file.

classmethod from_parquet(path: str, vector_col: str = 'vector', **kwargs)

Load a collection from a Parquet file.

generate_synthetic_queries(text_key: str = 'text', n_queries_per_doc: int = 1, num_docs: int = 100, model: str = 'gpt-4o-mini', custom_prompt: str | None = None, output_path: str | None = None, api_base: str | None = None, **llm_kwargs) → List[Dict[str, Any]]

Generate synthetic search queries for documents in the collection using an LLM. Supports local Ollama/vLLM via api_base and llm_kwargs.

hybrid_search(query_vector: ndarray | List[float], query_text: str, top_k: int = 5, dense_weight: float = 0.5, sparse_weight: float = 0.5, where: Dict[str, Any] | None = None) → List[Tuple[Dict[str, Any], float]]

Perform hybrid search combining dense and sparse results using Reciprocal Rank Fusion (RRF).

search(query: ndarray | List[float], top_k: int = 5, where: Dict[str, Any] | None = None, reranker: callable | RerankHandler | None = None, query_text: str | None = None) → List[Tuple[Dict[str, Any], float]]

Search the collection for the nearest neighbors.

Parameters:

• query – Vector to search for.

• top_k – Number of results to return.

• where – Metadata filter dictionary.

• reranker – A callable or RerankHandler for re-scoring.

• query_text – Original text for reranking context.

search_image(image_path: str, model: str = 'openai/clip-vit-base-patch32', top_k: int = 5) → List[Tuple[Dict[str, Any], float]]

Search for similar items using an image query.

search_trajectory(trajectory: ndarray | List[List[float]], top_k: int = 5, pooling: str = 'mean', where: Dict[str, Any] | None = None) → List[Tuple[Dict[str, Any], float]]

Search for similar trajectories (sequences of vectors).

Parameters:

• trajectory – Sequence of vectors representing a state/action trajectory.

• top_k – Number of results to return.

• pooling – Method to pool the trajectory into a single search vector (‘mean’ or ‘max’).

• where – Metadata filter dictionary.

to_parquet(path: str, vector_col: str = 'vector')

Save the collection to a Parquet file.

class core.Session(session_id: str, dimension: int, storage_dir: str = '.embenx_sessions')

Bases: object

Managed agentic session with automatic temporal decay and persistence.

add_interaction(vector: ndarray | List[float], text: str, **metadata)

Add a new interaction to the session memory.

cleanup()

Delete session data.

retrieve_context(query_vector: ndarray, top_k: int = 5, recency_weight: float = 0.4)

Retrieve relevant context from the session with recency bias.

class core.SpatialCollection(name: str = 'default', dimension: int | None = None, indexer_type: str = 'faiss', sparse_indexer_type: str | None = None, truncate_dim: int | None = None, **indexer_kwargs)

Bases: Collection

Specialized collection for ESWM (Episodic Spatial World Memory). Supports navigation trajectories and spatial-aware retrieval.

add_spatial(vectors: ndarray | List[List[float]], coords: ndarray, metadata: List[Dict[str, Any]] | None = None)

Add semantic embeddings and their associated spatial coordinates (x, y, z).

search_spatial(query_vector: ndarray, current_coords: ndarray, top_k: int = 5, spatial_radius: float = 10.0) → List[Tuple[Dict[str, Any], float]]

Spatial-aware search that favors episodic memories near the current location.

class core.StateCollection(name: str = 'default', dimension: int | None = None, indexer_type: str = 'faiss', sparse_indexer_type: str | None = None, truncate_dim: int | None = None, **indexer_kwargs)

Bases: Collection

Specialized collection for State Space Model (SSM) hydration. Supports storing hidden states (h0).

add_states(vectors: ndarray | List[List[float]], states: ndarray, metadata: List[Dict[str, Any]] | None = None)

Add embeddings and their associated SSM hidden states.

get_state(metadata: Dict[str, Any]) → ndarray

Retrieve hidden state for a given metadata result.

class core.TemporalCollection(name: str = 'default', dimension: int | None = None, indexer_type: str = 'faiss', sparse_indexer_type: str | None = None, truncate_dim: int | None = None, **indexer_kwargs)

Bases: Collection

Specialized collection for Echo-style temporal episodic memory. Supports time-stamped embeddings and recency-biased retrieval.

add_temporal(vectors: ndarray | List[List[float]], timestamps: List[float] | None = None, metadata: List[Dict[str, Any]] | None = None)

Add embeddings with Unix timestamps.

search_temporal(query_vector: ndarray, top_k: int = 5, recency_weight: float = 0.5, time_window: Tuple[float, float] | None = None, where: Dict[str, Any] | None = None) → List[Tuple[Dict[str, Any], float]]

Temporal-aware search that ranks results by similarity and recency.

Benchmark

benchmark.benchmark_single_indexer(name, indexer_cls, dimension, embeddings, metadata, console, cleanup=True)

benchmark.display_results(results, console)

benchmark.generate_report(results: List[Dict[str, Any]], dataset_name: str, output_path: str = 'benchmark_report.md')

Generate a formatted Markdown technical report from benchmark results.

benchmark.get_memory_usage()

benchmark.load_custom_indexer(script_path: str, console: Console)

Dynamically load a class inheriting from BaseIndexer from a given script.

benchmark.run_benchmark(dataset_name: str, split: str, text_column: str, max_docs: int, indexer_names: List[str], model_name: str, console: Console, data_files: str = None, cleanup: bool = True, custom_indexer_script: str = None, subset: str = 'default')

Run Embenx benchmarks. Matches original signature for test compatibility.

Reranking

class rerank.RerankHandler(model_name: str = 'flashrank', model_type: str = 'flashrank', **kwargs)

Bases: object

Unified handler for reranking retrieved results.

rerank(query: str, results: List[Tuple[Dict[str, Any], float]], top_k: int = 5) → List[Tuple[Dict[str, Any], float]]

Rerank a list of (metadata, distance) tuples. Expects a ‘text’ field in metadata for reranking.

Data & Zoo

data.list_zoo() → list

List all available pre-built collections in the zoo.

data.load_documents(dataset_name: str, subset: str = 'default', split: str = 'train', max_docs: int = 100) → List[Dict[str, Any]]

Load documents from Hugging Face or local files.

data.load_from_zoo(dataset_name: str, cache_dir: str = '.embenx_cache') → Collection

Download and load a pre-built collection from the Embenx Retrieval Zoo.

data.save_collection(collection: Collection, path: str)

Save a collection’s vectors and metadata to disk.

LLM & Embedding

class llm.Embedder(model_name: str, batch_size: int = 32, truncate_dim: int | None = None)

Bases: object

embed_query(query: str) → List[float]

embed_texts(texts: List[str]) → List[List[float]]

Embed a list of texts using LiteLLM with batching.

class llm.Generator(model_name: str = 'gpt-4o-mini', api_base: str | None = None, **kwargs)

Bases: object

generate(prompt: str) → str

Generate text using LiteLLM.

Agentic & MCP

async mcp_server.call_tool_impl(name: str, arguments: Dict[str, Any]) → List[TextContent | ImageContent | EmbeddedResource]

Handle tool calls from the AI agent.

async mcp_server.list_tools_impl() → List[Tool]

List available tools for agentic memory.

async mcp_server.run()

Indexers

class indexers.BaseIndexer(name: str, dimension: int)

Bases: ABC

abstractmethod build_index(embeddings: List[List[float]], metadata: List[Dict[str, Any]]) → None

Build or insert embeddings into the index.

Parameters:

• embeddings – List of embedding vectors.

• metadata – List of metadata dictionaries corresponding to each embedding.

cleanup() → None

Optional method to clean up temporary resources or files.

abstractmethod get_size() → int

Return the approximate memory footprint or disk size in bytes.

abstractmethod search(query_embedding: List[float], top_k: int = 5) → List[Tuple[Dict[str, Any], float]]

Search the index and return a list of (metadata, distance/score) tuples.

Parameters:

• query_embedding – The embedding vector to search for.

• top_k – Number of nearest neighbors to return.

Returns:

List of tuples containing (metadata, distance).

indexers.get_indexer_map()

class indexers.base.BaseIndexer(name: str, dimension: int)

Bases: ABC

abstractmethod build_index(embeddings: List[List[float]], metadata: List[Dict[str, Any]]) → None

Build or insert embeddings into the index.

Parameters:

• embeddings – List of embedding vectors.

• metadata – List of metadata dictionaries corresponding to each embedding.

cleanup() → None

Optional method to clean up temporary resources or files.

abstractmethod get_size() → int

Return the approximate memory footprint or disk size in bytes.

abstractmethod search(query_embedding: List[float], top_k: int = 5) → List[Tuple[Dict[str, Any], float]]

Search the index and return a list of (metadata, distance/score) tuples.

Parameters:

• query_embedding – The embedding vector to search for.

• top_k – Number of nearest neighbors to return.

Returns:

List of tuples containing (metadata, distance).