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.
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
- GitHub: github.com/vmore2/AgentRank-base
- HuggingFace: huggingface.co/vrushket
- CogniHive: pypi.org/project/cognihive
Contact
- Author: Vrushket More
- Email: vrushket2604@gmail.com
- LinkedIn: linkedin.com/in/vrushket-more-07b38b25b
- Issues: GitHub Issues
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
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
8a89b6ed6c85ee80ec46c06db721e9128b2b7f3c3d0ee8cba2535ebc3af1e045
|
|
| MD5 |
f81ebdd48fd1824f527b6e892958c48e
|
|
| BLAKE2b-256 |
5bd7bf14378fe169f2594482636d8ed116bbd02308d1cf670935c476ae3100ce
|
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
76f8597d4a090b12a09f6f798a0929d72ee95476a65f13dd795eb0d7fef8c487
|
|
| MD5 |
e57cc41366c092d1b60f4bd77049af5b
|
|
| BLAKE2b-256 |
88208b622c5fb32b9dff4bb8aad973ad40363911d6bb77008f4483a46645fb1c
|