Skip to main content

Temporal-aware embedding model for AI agent memory retrieval. The first embedder that understands WHEN memories happened.

Project description

🧠 AgentRank

Temporal-aware embeddings for AI agent memory retrieval.

PyPI version License HuggingFace


The Problem

Standard embedding models (OpenAI, Cohere, MiniLM) treat "yesterday" and "6 months ago" identically. For AI agents with long-term memory, this breaks temporal reasoning completely.

AgentRank solves this with embeddings that understand:

  • When memories happened (temporal awareness)
  • 📂 What type of memory it is (episodic, semantic, procedural)
  • 21% better retrieval on agent memory benchmarks

Installation

pip install agentrank

Quick Start

from agentrank import AgentRankEmbedder

# Load model
model = AgentRankEmbedder.from_pretrained("vrushket/agentrank-base")

# Encode with temporal context
embeddings = model.encode(
    texts=["User prefers Python for backend development"],
    temporal_info=[7],        # 7 days ago
    memory_types=["semantic"] # It's a preference
)

# Use embeddings for retrieval
print(embeddings.shape)  # [1, 768]

Model Family

Model Type Params Use Case HuggingFace
AgentRank-Base Embedder 149M Best quality retrieval vrushket/agentrank-base
AgentRank-Small Embedder 33M Fast inference vrushket/agentrank-small
AgentRank-Reranker Cross-encoder 149M Accurate reranking vrushket/agentrank-reranker

Two-Stage Retrieval Pipeline

For best results, use a two-stage pipeline:

from agentrank import AgentRankEmbedder
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# Stage 1: Fast retrieval with embedder
embedder = AgentRankEmbedder.from_pretrained("vrushket/agentrank-base")
query_embedding = embedder.encode(["What's my Python preference?"])
# ... search vector DB → get top-50 candidates ...

# Stage 2: Accurate reranking with cross-encoder
reranker = AutoModelForSequenceClassification.from_pretrained("vrushket/agentrank-reranker")
tokenizer = AutoTokenizer.from_pretrained("vrushket/agentrank-reranker")

def rerank(query, candidates, top_k=10):
    scored = []
    for memory in candidates:
        inputs = tokenizer(query, memory, return_tensors="pt", truncation=True)
        with torch.no_grad():
            score = torch.sigmoid(reranker(**inputs).logits).item()
        scored.append((score, memory))
    return sorted(scored, reverse=True)[:top_k]

top_10 = rerank("What's my Python preference?", top_50_candidates)

Benchmarks

Model MRR Recall@1 Recall@5 NDCG@10
AgentRank-Base 0.6496 0.4440 99.6% 0.6786
AgentRank-Small 0.6375 0.4460 97.4% 0.6797
MPNet-base-v2 0.5351 0.3660 79.6% 0.6335
MiniLM-L6-v2 0.5297 0.3720 75.2% 0.6370

+22% MRR improvement over baseline embedding models.

Reranker Validation Accuracy Val Loss
AgentRank-Reranker 89.11% 0.2554

Key Features

Temporal Embeddings

10 learnable time buckets encode recency:

  • Today/Yesterday, This Week, This Month, Last Quarter, etc.
  • Model learns what "recent" means in context

Memory Type Embeddings

Distinguish between:

  • Episodic: Events ("We discussed Python yesterday")
  • Semantic: Facts/preferences ("User likes Python")
  • Procedural: Instructions ("To deploy, run npm build")

Architecture

  • Base: ModernBERT-base / MiniLM
  • Temporal + Type embeddings added (scaled by 0.1)
  • Trained on 500K synthetic agent memory samples
  • Hard negative mining with 7 negative types

Works Great With

CogniHive — Multi-agent memory with "who knows what" routing

pip install cognihive

Together: CogniHive routes questions to the right agent, AgentRank retrieves the right memories.


Links


Contact


License

Apache 2.0 — Free for commercial use.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

agentrank-0.1.3.tar.gz (7.6 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

agentrank-0.1.3-py3-none-any.whl (7.3 kB view details)

Uploaded Python 3

File details

Details for the file agentrank-0.1.3.tar.gz.

File metadata

  • Download URL: agentrank-0.1.3.tar.gz
  • Upload date:
  • Size: 7.6 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.3

File hashes

Hashes for agentrank-0.1.3.tar.gz
Algorithm Hash digest
SHA256 8a89b6ed6c85ee80ec46c06db721e9128b2b7f3c3d0ee8cba2535ebc3af1e045
MD5 f81ebdd48fd1824f527b6e892958c48e
BLAKE2b-256 5bd7bf14378fe169f2594482636d8ed116bbd02308d1cf670935c476ae3100ce

See more details on using hashes here.

File details

Details for the file agentrank-0.1.3-py3-none-any.whl.

File metadata

  • Download URL: agentrank-0.1.3-py3-none-any.whl
  • Upload date:
  • Size: 7.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.3

File hashes

Hashes for agentrank-0.1.3-py3-none-any.whl
Algorithm Hash digest
SHA256 76f8597d4a090b12a09f6f798a0929d72ee95476a65f13dd795eb0d7fef8c487
MD5 e57cc41366c092d1b60f4bd77049af5b
BLAKE2b-256 88208b622c5fb32b9dff4bb8aad973ad40363911d6bb77008f4483a46645fb1c

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page