Compare commits

..

1 Commits

Author SHA1 Message Date
kartik-mem0 5b9bf44272 chore: update changelog, bump SDK and package versions to 3.0.8 and 2.0.6 2026-06-13 18:49:06 +05:30
29 changed files with 2033 additions and 447 deletions
-4
View File
@@ -1,4 +0,0 @@
[submodule "evaluation"]
path = evaluation
url = https://github.com/mem0ai/memory-benchmarks
branch = main
+12 -13
View File
@@ -12,7 +12,7 @@ This file provides context for AI coding assistants (Claude Code, Cursor, GitHub
## Repository Structure
This is a **polyglot monorepo** containing Python and TypeScript packages, CLIs, servers, plugins, and documentation.
This is a **polyglot monorepo** containing Python and TypeScript packages, CLIs, servers, plugins, documentation, and evaluation tooling.
### Key Directories
@@ -32,7 +32,7 @@ This is a **polyglot monorepo** containing Python and TypeScript packages, CLIs,
| `skills/` | Claude Code skill definitions. Reference skills (SDK knowledge, always-on): `mem0/`, `mem0-cli/`, `mem0-vercel-ai-sdk/`. Pipeline skills (run on demand): `mem0-integrate/`, `mem0-test-integration/`, `mem0-oss-to-platform/` |
| `docs/` | Documentation site (Mintlify) |
| `tests/` | Python SDK tests (pytest) |
| `evaluation/` | Submodule → [`mem0ai/memory-benchmarks`](https://github.com/mem0ai/memory-benchmarks) — benchmarking (LOCOMO, LongMemEval, BEAM) lives in that repo |
| `evaluation/` | Benchmarking framework — LOCOMO evals, experiment runner, score generation |
| `examples/` | Sample projects & runnable demos — apps, Chrome extension, multi-agent patterns, and Jupyter notebooks (`notebooks/`) |
| `pr-reviews/` | Pull request review materials |
| `scripts/` | Repo-wide utility scripts (e.g., `check-llms-txt-coverage.py` for docs/llms.txt sync) |
@@ -246,19 +246,18 @@ make docs # or: cd docs && mintlify dev
- **API spec:** `docs/openapi.json`
- **Structure:** `api-reference/`, `open-source/`, `platform/`, `integrations/`, `cookbooks/`, `core-concepts/`
### Evaluation / Benchmarking
Benchmarking lives in the external [`mem0ai/memory-benchmarks`](https://github.com/mem0ai/memory-benchmarks) repo (LOCOMO + LongMemEval + BEAM). The in-repo `evaluation/` path is a **git submodule** pinned to that repo's `main` — populate it with `git submodule update --init evaluation` (or clone mem0 with `--recurse-submodules`), or clone the benchmarks repo standalone:
### Evaluation (`evaluation/`)
```bash
git clone https://github.com/mem0ai/memory-benchmarks.git
cd memory-benchmarks
pip install -r requirements.txt
# Run a benchmark (Mem0 Cloud; use docker compose for OSS)
python -m benchmarks.locomo.run --project-name my-test --backend cloud --mem0-api-key $MEM0_API_KEY
python -m benchmarks.longmemeval.run --project-name my-test --backend cloud --mem0-api-key $MEM0_API_KEY --all-questions
python -m benchmarks.beam.run --project-name my-test --backend cloud --mem0-api-key $MEM0_API_KEY --chat-sizes 100K --conversations 0-9
cd evaluation
make run-mem0-add # Run mem0 add experiments
make run-mem0-search # Run mem0 search experiments
make run-mem0-plus-add # With graph memory
make run-mem0-plus-search # With graph memory
make run-rag # RAG baseline
make run-full-context # Full context baseline
make run-langmem # LangMem comparison
make run-openai # OpenAI comparison
```
## Core APIs
Submodule evaluation deleted from 4b61c5d31b
+31
View File
@@ -0,0 +1,31 @@
# Run the experiments
run-mem0-add:
python run_experiments.py --technique_type mem0 --method add
run-mem0-search:
python run_experiments.py --technique_type mem0 --method search --output_folder results/ --top_k 30
run-mem0-plus-add:
python run_experiments.py --technique_type mem0 --method add --is_graph
run-mem0-plus-search:
python run_experiments.py --technique_type mem0 --method search --is_graph --output_folder results/ --top_k 30
run-rag:
python run_experiments.py --technique_type rag --chunk_size 500 --num_chunks 1 --output_folder results/
run-full-context:
python run_experiments.py --technique_type rag --chunk_size -1 --num_chunks 1 --output_folder results/
run-langmem:
python run_experiments.py --technique_type langmem --output_folder results/
run-zep-add:
python run_experiments.py --technique_type zep --method add --output_folder results/
run-zep-search:
python run_experiments.py --technique_type zep --method search --output_folder results/
run-openai:
python run_experiments.py --technique_type openai --output_folder results/
+198
View File
@@ -0,0 +1,198 @@
# Mem0: Building Production‑Ready AI Agents with Scalable Long‑Term Memory
[![arXiv](https://img.shields.io/badge/arXiv-Paper-b31b1b.svg)](https://arxiv.org/abs/2504.19413)
[![Website](https://img.shields.io/badge/Website-Project-blue)](https://mem0.ai/research)
This repository contains the code and dataset for our paper: **Mem0: Building Production‑Ready AI Agents with Scalable Long‑Term Memory**.
## 📋 Overview
This project evaluates Mem0 and compares it with different memory and retrieval techniques for AI systems:
1. **Established LOCOMO Benchmarks**: We evaluate against five established approaches from the literature: LoCoMo, ReadAgent, MemoryBank, MemGPT, and A-Mem.
2. **Open-Source Memory Solutions**: We test promising open-source memory architectures including LangMem, which provides flexible memory management capabilities.
3. **RAG Systems**: We implement Retrieval-Augmented Generation with various configurations, testing different chunk sizes and retrieval counts to optimize performance.
4. **Full-Context Processing**: We examine the effectiveness of passing the entire conversation history within the context window of the LLM as a baseline approach.
5. **Proprietary Memory Systems**: We evaluate OpenAI's built-in memory feature available in their ChatGPT interface to compare against commercial solutions.
6. **Third-Party Memory Providers**: We incorporate Zep, a specialized memory management platform designed for AI agents, to assess the performance of dedicated memory infrastructure.
We test these techniques on the LOCOMO dataset, which contains conversational data with various question types to evaluate memory recall and understanding.
## 🔍 Dataset
The LOCOMO dataset used in our experiments can be downloaded from our Google Drive repository:
[Download LOCOMO Dataset](https://drive.google.com/drive/folders/1L-cTjTm0ohMsitsHg4dijSPJtqNflwX-?usp=drive_link)
The dataset contains conversational data specifically designed to test memory recall and understanding across various question types and complexity levels.
Place the dataset files in the `dataset/` directory:
- `locomo10.json`: Original dataset
- `locomo10_rag.json`: Dataset formatted for RAG experiments
## 📁 Project Structure
```
.
├── src/ # Source code for different memory techniques
│ ├── mem0/ # Implementation of the Mem0 technique
│ ├── openai/ # Implementation of the OpenAI memory
│ ├── zep/ # Implementation of the Zep memory
│ ├── rag.py # Implementation of the RAG technique
│ └── langmem.py # Implementation of the Language-based memory
├── metrics/ # Code for evaluation metrics
├── results/ # Results of experiments
├── dataset/ # Dataset files
├── evals.py # Evaluation script
├── run_experiments.py # Script to run experiments
├── generate_scores.py # Script to generate scores from results
└── prompts.py # Prompts used for the models
```
## 🚀 Getting Started
### Prerequisites
Create a `.env` file with your API keys and configurations. The following keys are required:
```
# OpenAI API key for GPT models and embeddings
OPENAI_API_KEY="your-openai-api-key"
# Mem0 API keys (for Mem0 and Mem0+ techniques)
MEM0_API_KEY="your-mem0-api-key"
MEM0_PROJECT_ID="your-mem0-project-id"
MEM0_ORGANIZATION_ID="your-mem0-organization-id"
# Model configuration
MODEL="gpt-4o-mini" # or your preferred model
EMBEDDING_MODEL="text-embedding-3-small" # or your preferred embedding model
ZEP_API_KEY="api-key-from-zep"
```
### Running Experiments
You can run experiments using the provided Makefile commands:
#### Memory Techniques
```bash
# Run Mem0 experiments
make run-mem0-add # Add memories using Mem0
make run-mem0-search # Search memories using Mem0
# Run Mem0+ experiments (with graph-based search)
make run-mem0-plus-add # Add memories using Mem0+
make run-mem0-plus-search # Search memories using Mem0+
# Run RAG experiments
make run-rag # Run RAG with chunk size 500
make run-full-context # Run RAG with full context
# Run LangMem experiments
make run-langmem # Run LangMem
# Run Zep experiments
make run-zep-add # Add memories using Zep
make run-zep-search # Search memories using Zep
# Run OpenAI experiments
make run-openai # Run OpenAI experiments
```
Alternatively, you can run experiments directly with custom parameters:
```bash
python run_experiments.py --technique_type [mem0|rag|langmem] [additional parameters]
```
#### Command-line Parameters:
| Parameter | Description | Default |
|-----------|-------------|---------|
| `--technique_type` | Memory technique to use (mem0, rag, langmem) | mem0 |
| `--method` | Method to use (add, search) | add |
| `--chunk_size` | Chunk size for processing | 1000 |
| `--top_k` | Number of top memories to retrieve | 30 |
| `--filter_memories` | Whether to filter memories | False |
| `--is_graph` | Whether to use graph-based search | False |
| `--num_chunks` | Number of chunks to process for RAG | 1 |
### 📊 Evaluation
To evaluate results, run:
```bash
python evals.py --input_file [path_to_results] --output_file [output_path]
```
This script:
1. Processes each question-answer pair
2. Calculates BLEU and F1 scores automatically
3. Uses an LLM judge to evaluate answer correctness
4. Saves the combined results to the output file
### 📈 Generating Scores
Generate final scores with:
```bash
python generate_scores.py
```
This script:
1. Loads the evaluation metrics data
2. Calculates mean scores for each category (BLEU, F1, LLM)
3. Reports the number of questions per category
4. Calculates overall mean scores across all categories
Example output:
```
Mean Scores Per Category:
bleu_score f1_score llm_score count
category
1 0.xxxx 0.xxxx 0.xxxx xx
2 0.xxxx 0.xxxx 0.xxxx xx
3 0.xxxx 0.xxxx 0.xxxx xx
Overall Mean Scores:
bleu_score 0.xxxx
f1_score 0.xxxx
llm_score 0.xxxx
```
## 📏 Evaluation Metrics
We use several metrics to evaluate the performance of different memory techniques:
1. **BLEU Score**: Measures the similarity between the model's response and the ground truth
2. **F1 Score**: Measures the harmonic mean of precision and recall
3. **LLM Score**: A binary score (0 or 1) determined by an LLM judge evaluating the correctness of responses
4. **Token Consumption**: Number of tokens required to generate final answer.
5. **Latency**: Time required during search and to generate response.
## 📚 Citation
If you use this code or dataset in your research, please cite our paper:
```bibtex
@article{mem0,
title={Mem0: Building Production-Ready AI Agents with Scalable Long-Term Memory},
author={Chhikara, Prateek and Khant, Dev and Aryan, Saket and Singh, Taranjeet and Yadav, Deshraj},
journal={arXiv preprint arXiv:2504.19413},
year={2025}
}
```
## 📄 License
[MIT License](LICENSE)
## 👥 Contributors
- [Prateek Chhikara](https://github.com/prateekchhikara)
- [Dev Khant](https://github.com/Dev-Khant)
- [Saket Aryan](https://github.com/whysosaket)
- [Taranjeet Singh](https://github.com/taranjeet)
- [Deshraj Yadav](https://github.com/deshraj)
+81
View File
@@ -0,0 +1,81 @@
import argparse
import concurrent.futures
import json
import threading
from collections import defaultdict
from metrics.llm_judge import evaluate_llm_judge
from metrics.utils import calculate_bleu_scores, calculate_metrics
from tqdm import tqdm
def process_item(item_data):
k, v = item_data
local_results = defaultdict(list)
for item in v:
gt_answer = str(item["answer"])
pred_answer = str(item["response"])
category = str(item["category"])
question = str(item["question"])
# Skip category 5
if category == "5":
continue
metrics = calculate_metrics(pred_answer, gt_answer)
bleu_scores = calculate_bleu_scores(pred_answer, gt_answer)
llm_score = evaluate_llm_judge(question, gt_answer, pred_answer)
local_results[k].append(
{
"question": question,
"answer": gt_answer,
"response": pred_answer,
"category": category,
"bleu_score": bleu_scores["bleu1"],
"f1_score": metrics["f1"],
"llm_score": llm_score,
}
)
return local_results
def main():
parser = argparse.ArgumentParser(description="Evaluate RAG results")
parser.add_argument(
"--input_file", type=str, default="results/rag_results_500_k1.json", help="Path to the input dataset file"
)
parser.add_argument(
"--output_file", type=str, default="evaluation_metrics.json", help="Path to save the evaluation results"
)
parser.add_argument("--max_workers", type=int, default=10, help="Maximum number of worker threads")
args = parser.parse_args()
with open(args.input_file, "r") as f:
data = json.load(f)
results = defaultdict(list)
results_lock = threading.Lock()
# Use ThreadPoolExecutor with specified workers
with concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor:
futures = [executor.submit(process_item, item_data) for item_data in data.items()]
for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures)):
local_results = future.result()
with results_lock:
for k, items in local_results.items():
results[k].extend(items)
# Save results to JSON file
with open(args.output_file, "w") as f:
json.dump(results, f, indent=4)
print(f"Results saved to {args.output_file}")
if __name__ == "__main__":
main()
+34
View File
@@ -0,0 +1,34 @@
import json
import pandas as pd
# Load the evaluation metrics data
with open("evaluation_metrics.json", "r") as f:
data = json.load(f)
# Flatten the data into a list of question items
all_items = []
for key in data:
all_items.extend(data[key])
# Convert to DataFrame
df = pd.DataFrame(all_items)
# Convert category to numeric type
df["category"] = pd.to_numeric(df["category"])
# Calculate mean scores by category
result = df.groupby("category").agg({"bleu_score": "mean", "f1_score": "mean", "llm_score": "mean"}).round(4)
# Add count of questions per category
result["count"] = df.groupby("category").size()
# Print the results
print("Mean Scores Per Category:")
print(result)
# Calculate overall means
overall_means = df.agg({"bleu_score": "mean", "f1_score": "mean", "llm_score": "mean"}).round(4)
print("\nOverall Mean Scores:")
print(overall_means)
+130
View File
@@ -0,0 +1,130 @@
import argparse
import json
from collections import defaultdict
import numpy as np
from openai import OpenAI
from mem0.memory.utils import extract_json
client = OpenAI()
ACCURACY_PROMPT = """
Your task is to label an answer to a question as ’CORRECT’ or ’WRONG’. You will be given the following data:
(1) a question (posed by one user to another user),
(2) a ’gold’ (ground truth) answer,
(3) a generated answer
which you will score as CORRECT/WRONG.
The point of the question is to ask about something one user should know about the other user based on their prior conversations.
The gold answer will usually be a concise and short answer that includes the referenced topic, for example:
Question: Do you remember what I got the last time I went to Hawaii?
Gold answer: A shell necklace
The generated answer might be much longer, but you should be generous with your grading - as long as it touches on the same topic as the gold answer, it should be counted as CORRECT.
For time related questions, the gold answer will be a specific date, month, year, etc. The generated answer might be much longer or use relative time references (like "last Tuesday" or "next month"), but you should be generous with your grading - as long as it refers to the same date or time period as the gold answer, it should be counted as CORRECT. Even if the format differs (e.g., "May 7th" vs "7 May"), consider it CORRECT if it's the same date.
Now it's time for the real question:
Question: {question}
Gold answer: {gold_answer}
Generated answer: {generated_answer}
First, provide a short (one sentence) explanation of your reasoning, then finish with CORRECT or WRONG.
Do NOT include both CORRECT and WRONG in your response, or it will break the evaluation script.
Just return the label CORRECT or WRONG in a json format with the key as "label".
"""
def evaluate_llm_judge(question, gold_answer, generated_answer):
"""Evaluate the generated answer against the gold answer using an LLM judge."""
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "user",
"content": ACCURACY_PROMPT.format(
question=question, gold_answer=gold_answer, generated_answer=generated_answer
),
}
],
response_format={"type": "json_object"},
temperature=0.0,
)
label = json.loads(extract_json(response.choices[0].message.content))["label"]
return 1 if label == "CORRECT" else 0
def main():
"""Main function to evaluate RAG results using LLM judge."""
parser = argparse.ArgumentParser(description="Evaluate RAG results using LLM judge")
parser.add_argument(
"--input_file",
type=str,
default="results/default_run_v4_k30_new_graph.json",
help="Path to the input dataset file",
)
args = parser.parse_args()
dataset_path = args.input_file
output_path = f"results/llm_judge_{dataset_path.split('/')[-1]}"
with open(dataset_path, "r") as f:
data = json.load(f)
LLM_JUDGE = defaultdict(list)
RESULTS = defaultdict(list)
index = 0
for k, v in data.items():
for x in v:
question = x["question"]
gold_answer = x["answer"]
generated_answer = x["response"]
category = x["category"]
# Skip category 5
if int(category) == 5:
continue
# Evaluate the answer
label = evaluate_llm_judge(question, gold_answer, generated_answer)
LLM_JUDGE[category].append(label)
# Store the results
RESULTS[index].append(
{
"question": question,
"gt_answer": gold_answer,
"response": generated_answer,
"category": category,
"llm_label": label,
}
)
# Save intermediate results
with open(output_path, "w") as f:
json.dump(RESULTS, f, indent=4)
# Print current accuracy for all categories
print("All categories accuracy:")
for cat, results in LLM_JUDGE.items():
if results: # Only print if there are results for this category
print(f" Category {cat}: {np.mean(results):.4f} ({sum(results)}/{len(results)})")
print("------------------------------------------")
index += 1
# Save final results
with open(output_path, "w") as f:
json.dump(RESULTS, f, indent=4)
# Print final summary
print("PATH: ", dataset_path)
print("------------------------------------------")
for k, v in LLM_JUDGE.items():
print(k, np.mean(v))
if __name__ == "__main__":
main()
+211
View File
@@ -0,0 +1,211 @@
"""
Borrowed from https://github.com/WujiangXu/AgenticMemory/blob/main/utils.py
@article{xu2025mem,
title={A-mem: Agentic memory for llm agents},
author={Xu, Wujiang and Liang, Zujie and Mei, Kai and Gao, Hang and Tan, Juntao
and Zhang, Yongfeng},
journal={arXiv preprint arXiv:2502.12110},
year={2025}
}
"""
import statistics
from collections import defaultdict
from typing import Dict, List, Union
import nltk
from bert_score import score as bert_score
from nltk.translate.bleu_score import SmoothingFunction, sentence_bleu
from nltk.translate.meteor_score import meteor_score
from rouge_score import rouge_scorer
from sentence_transformers import SentenceTransformer
# from load_dataset import load_locomo_dataset, QA, Turn, Session, Conversation
from sentence_transformers.util import pytorch_cos_sim
# Download required NLTK data
try:
nltk.download("punkt", quiet=True)
nltk.download("wordnet", quiet=True)
except Exception as e:
print(f"Error downloading NLTK data: {e}")
# Initialize SentenceTransformer model (this will be reused)
try:
sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
except Exception as e:
print(f"Warning: Could not load SentenceTransformer model: {e}")
sentence_model = None
def simple_tokenize(text):
"""Simple tokenization function."""
# Convert to string if not already
text = str(text)
return text.lower().replace(".", " ").replace(",", " ").replace("!", " ").replace("?", " ").split()
def calculate_rouge_scores(prediction: str, reference: str) -> Dict[str, float]:
"""Calculate ROUGE scores for prediction against reference."""
scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
scores = scorer.score(reference, prediction)
return {
"rouge1_f": scores["rouge1"].fmeasure,
"rouge2_f": scores["rouge2"].fmeasure,
"rougeL_f": scores["rougeL"].fmeasure,
}
def calculate_bleu_scores(prediction: str, reference: str) -> Dict[str, float]:
"""Calculate BLEU scores with different n-gram settings."""
pred_tokens = nltk.word_tokenize(prediction.lower())
ref_tokens = [nltk.word_tokenize(reference.lower())]
weights_list = [(1, 0, 0, 0), (0.5, 0.5, 0, 0), (0.33, 0.33, 0.33, 0), (0.25, 0.25, 0.25, 0.25)]
smooth = SmoothingFunction().method1
scores = {}
for n, weights in enumerate(weights_list, start=1):
try:
score = sentence_bleu(ref_tokens, pred_tokens, weights=weights, smoothing_function=smooth)
except Exception as e:
print(f"Error calculating BLEU score: {e}")
score = 0.0
scores[f"bleu{n}"] = score
return scores
def calculate_bert_scores(prediction: str, reference: str) -> Dict[str, float]:
"""Calculate BERTScore for semantic similarity."""
try:
P, R, F1 = bert_score([prediction], [reference], lang="en", verbose=False)
return {"bert_precision": P.item(), "bert_recall": R.item(), "bert_f1": F1.item()}
except Exception as e:
print(f"Error calculating BERTScore: {e}")
return {"bert_precision": 0.0, "bert_recall": 0.0, "bert_f1": 0.0}
def calculate_meteor_score(prediction: str, reference: str) -> float:
"""Calculate METEOR score for the prediction."""
try:
return meteor_score([reference.split()], prediction.split())
except Exception as e:
print(f"Error calculating METEOR score: {e}")
return 0.0
def calculate_sentence_similarity(prediction: str, reference: str) -> float:
"""Calculate sentence embedding similarity using SentenceBERT."""
if sentence_model is None:
return 0.0
try:
# Encode sentences
embedding1 = sentence_model.encode([prediction], convert_to_tensor=True)
embedding2 = sentence_model.encode([reference], convert_to_tensor=True)
# Calculate cosine similarity
similarity = pytorch_cos_sim(embedding1, embedding2).item()
return float(similarity)
except Exception as e:
print(f"Error calculating sentence similarity: {e}")
return 0.0
def calculate_metrics(prediction: str, reference: str) -> Dict[str, float]:
"""Calculate comprehensive evaluation metrics for a prediction."""
# Handle empty or None values
if not prediction or not reference:
return {
"exact_match": 0,
"f1": 0.0,
"rouge1_f": 0.0,
"rouge2_f": 0.0,
"rougeL_f": 0.0,
"bleu1": 0.0,
"bleu2": 0.0,
"bleu3": 0.0,
"bleu4": 0.0,
"bert_f1": 0.0,
"meteor": 0.0,
"sbert_similarity": 0.0,
}
# Convert to strings if they're not already
prediction = str(prediction).strip()
reference = str(reference).strip()
# Calculate exact match
exact_match = int(prediction.lower() == reference.lower())
# Calculate token-based F1 score
pred_tokens = set(simple_tokenize(prediction))
ref_tokens = set(simple_tokenize(reference))
common_tokens = pred_tokens & ref_tokens
if not pred_tokens or not ref_tokens:
f1 = 0.0
else:
precision = len(common_tokens) / len(pred_tokens)
recall = len(common_tokens) / len(ref_tokens)
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
# Calculate all scores
bleu_scores = calculate_bleu_scores(prediction, reference)
# Combine all metrics
metrics = {
"exact_match": exact_match,
"f1": f1,
**bleu_scores,
}
return metrics
def aggregate_metrics(
all_metrics: List[Dict[str, float]], all_categories: List[int]
) -> Dict[str, Dict[str, Union[float, Dict[str, float]]]]:
"""Calculate aggregate statistics for all metrics, split by category."""
if not all_metrics:
return {}
# Initialize aggregates for overall and per-category metrics
aggregates = defaultdict(list)
category_aggregates = defaultdict(lambda: defaultdict(list))
# Collect all values for each metric, both overall and per category
for metrics, category in zip(all_metrics, all_categories):
for metric_name, value in metrics.items():
aggregates[metric_name].append(value)
category_aggregates[category][metric_name].append(value)
# Calculate statistics for overall metrics
results = {"overall": {}}
for metric_name, values in aggregates.items():
results["overall"][metric_name] = {
"mean": statistics.mean(values),
"std": statistics.stdev(values) if len(values) > 1 else 0.0,
"median": statistics.median(values),
"min": min(values),
"max": max(values),
"count": len(values),
}
# Calculate statistics for each category
for category in sorted(category_aggregates.keys()):
results[f"category_{category}"] = {}
for metric_name, values in category_aggregates[category].items():
if values: # Only calculate if we have values for this category
results[f"category_{category}"][metric_name] = {
"mean": statistics.mean(values),
"std": statistics.stdev(values) if len(values) > 1 else 0.0,
"median": statistics.median(values),
"min": min(values),
"max": max(values),
"count": len(values),
}
return results
+147
View File
@@ -0,0 +1,147 @@
ANSWER_PROMPT_GRAPH = """
You are an intelligent memory assistant tasked with retrieving accurate information from
conversation memories.
# CONTEXT:
You have access to memories from two speakers in a conversation. These memories contain
timestamped information that may be relevant to answering the question. You also have
access to knowledge graph relations for each user, showing connections between entities,
concepts, and events relevant to that user.
# INSTRUCTIONS:
1. Carefully analyze all provided memories from both speakers
2. Pay special attention to the timestamps to determine the answer
3. If the question asks about a specific event or fact, look for direct evidence in the
memories
4. If the memories contain contradictory information, prioritize the most recent memory
5. If there is a question about time references (like "last year", "two months ago",
etc.), calculate the actual date based on the memory timestamp. For example, if a
memory from 4 May 2022 mentions "went to India last year," then the trip occurred
in 2021.
6. Always convert relative time references to specific dates, months, or years. For
example, convert "last year" to "2022" or "two months ago" to "March 2023" based
on the memory timestamp. Ignore the reference while answering the question.
7. Focus only on the content of the memories from both speakers. Do not confuse
character names mentioned in memories with the actual users who created those
memories.
8. The answer should be less than 5-6 words.
9. Use the knowledge graph relations to understand the user's knowledge network and
identify important relationships between entities in the user's world.
# APPROACH (Think step by step):
1. First, examine all memories that contain information related to the question
2. Examine the timestamps and content of these memories carefully
3. Look for explicit mentions of dates, times, locations, or events that answer the
question
4. If the answer requires calculation (e.g., converting relative time references),
show your work
5. Analyze the knowledge graph relations to understand the user's knowledge context
6. Formulate a precise, concise answer based solely on the evidence in the memories
7. Double-check that your answer directly addresses the question asked
8. Ensure your final answer is specific and avoids vague time references
Memories for user {{speaker_1_user_id}}:
{{speaker_1_memories}}
Relations for user {{speaker_1_user_id}}:
{{speaker_1_graph_memories}}
Memories for user {{speaker_2_user_id}}:
{{speaker_2_memories}}
Relations for user {{speaker_2_user_id}}:
{{speaker_2_graph_memories}}
Question: {{question}}
Answer:
"""
ANSWER_PROMPT = """
You are an intelligent memory assistant tasked with retrieving accurate information from conversation memories.
# CONTEXT:
You have access to memories from two speakers in a conversation. These memories contain
timestamped information that may be relevant to answering the question.
# INSTRUCTIONS:
1. Carefully analyze all provided memories from both speakers
2. Pay special attention to the timestamps to determine the answer
3. If the question asks about a specific event or fact, look for direct evidence in the memories
4. If the memories contain contradictory information, prioritize the most recent memory
5. If there is a question about time references (like "last year", "two months ago", etc.),
calculate the actual date based on the memory timestamp. For example, if a memory from
4 May 2022 mentions "went to India last year," then the trip occurred in 2021.
6. Always convert relative time references to specific dates, months, or years. For example,
convert "last year" to "2022" or "two months ago" to "March 2023" based on the memory
timestamp. Ignore the reference while answering the question.
7. Focus only on the content of the memories from both speakers. Do not confuse character
names mentioned in memories with the actual users who created those memories.
8. The answer should be less than 5-6 words.
# APPROACH (Think step by step):
1. First, examine all memories that contain information related to the question
2. Examine the timestamps and content of these memories carefully
3. Look for explicit mentions of dates, times, locations, or events that answer the question
4. If the answer requires calculation (e.g., converting relative time references), show your work
5. Formulate a precise, concise answer based solely on the evidence in the memories
6. Double-check that your answer directly addresses the question asked
7. Ensure your final answer is specific and avoids vague time references
Memories for user {{speaker_1_user_id}}:
{{speaker_1_memories}}
Memories for user {{speaker_2_user_id}}:
{{speaker_2_memories}}
Question: {{question}}
Answer:
"""
ANSWER_PROMPT_ZEP = """
You are an intelligent memory assistant tasked with retrieving accurate information from conversation memories.
# CONTEXT:
You have access to memories from a conversation. These memories contain
timestamped information that may be relevant to answering the question.
# INSTRUCTIONS:
1. Carefully analyze all provided memories
2. Pay special attention to the timestamps to determine the answer
3. If the question asks about a specific event or fact, look for direct evidence in the memories
4. If the memories contain contradictory information, prioritize the most recent memory
5. If there is a question about time references (like "last year", "two months ago", etc.),
calculate the actual date based on the memory timestamp. For example, if a memory from
4 May 2022 mentions "went to India last year," then the trip occurred in 2021.
6. Always convert relative time references to specific dates, months, or years. For example,
convert "last year" to "2022" or "two months ago" to "March 2023" based on the memory
timestamp. Ignore the reference while answering the question.
7. Focus only on the content of the memories. Do not confuse character
names mentioned in memories with the actual users who created those memories.
8. The answer should be less than 5-6 words.
# APPROACH (Think step by step):
1. First, examine all memories that contain information related to the question
2. Examine the timestamps and content of these memories carefully
3. Look for explicit mentions of dates, times, locations, or events that answer the question
4. If the answer requires calculation (e.g., converting relative time references), show your work
5. Formulate a precise, concise answer based solely on the evidence in the memories
6. Double-check that your answer directly addresses the question asked
7. Ensure your final answer is specific and avoids vague time references
Memories:
{{memories}}
Question: {{question}}
Answer:
"""
+75
View File
@@ -0,0 +1,75 @@
import argparse
import os
from src.langmem import LangMemManager
from src.memzero.add import MemoryADD
from src.memzero.search import MemorySearch
from src.openai.predict import OpenAIPredict
from src.rag import RAGManager
from src.utils import METHODS, TECHNIQUES
from src.zep.add import ZepAdd
from src.zep.search import ZepSearch
class Experiment:
def __init__(self, technique_type, chunk_size):
self.technique_type = technique_type
self.chunk_size = chunk_size
def run(self):
print(f"Running experiment with technique: {self.technique_type}, chunk size: {self.chunk_size}")
def main():
parser = argparse.ArgumentParser(description="Run memory experiments")
parser.add_argument("--technique_type", choices=TECHNIQUES, default="mem0", help="Memory technique to use")
parser.add_argument("--method", choices=METHODS, default="add", help="Method to use")
parser.add_argument("--chunk_size", type=int, default=1000, help="Chunk size for processing")
parser.add_argument("--output_folder", type=str, default="results/", help="Output path for results")
parser.add_argument("--top_k", type=int, default=30, help="Number of top memories to retrieve")
parser.add_argument("--filter_memories", action="store_true", default=False, help="Whether to filter memories")
parser.add_argument("--is_graph", action="store_true", default=False, help="Whether to use graph-based search")
parser.add_argument("--num_chunks", type=int, default=1, help="Number of chunks to process")
args = parser.parse_args()
# Add your experiment logic here
print(f"Running experiments with technique: {args.technique_type}, chunk size: {args.chunk_size}")
if args.technique_type == "mem0":
if args.method == "add":
memory_manager = MemoryADD(data_path="dataset/locomo10.json", is_graph=args.is_graph)
memory_manager.process_all_conversations()
elif args.method == "search":
output_file_path = os.path.join(
args.output_folder,
f"mem0_results_top_{args.top_k}_filter_{args.filter_memories}_graph_{args.is_graph}.json",
)
memory_searcher = MemorySearch(output_file_path, args.top_k, args.filter_memories, args.is_graph)
memory_searcher.process_data_file("dataset/locomo10.json")
elif args.technique_type == "rag":
output_file_path = os.path.join(args.output_folder, f"rag_results_{args.chunk_size}_k{args.num_chunks}.json")
rag_manager = RAGManager(data_path="dataset/locomo10_rag.json", chunk_size=args.chunk_size, k=args.num_chunks)
rag_manager.process_all_conversations(output_file_path)
elif args.technique_type == "langmem":
output_file_path = os.path.join(args.output_folder, "langmem_results.json")
langmem_manager = LangMemManager(dataset_path="dataset/locomo10_rag.json")
langmem_manager.process_all_conversations(output_file_path)
elif args.technique_type == "zep":
if args.method == "add":
zep_manager = ZepAdd(data_path="dataset/locomo10.json")
zep_manager.process_all_conversations("1")
elif args.method == "search":
output_file_path = os.path.join(args.output_folder, "zep_search_results.json")
zep_manager = ZepSearch()
zep_manager.process_data_file("dataset/locomo10.json", "1", output_file_path)
elif args.technique_type == "openai":
output_file_path = os.path.join(args.output_folder, "openai_results.json")
openai_manager = OpenAIPredict()
openai_manager.process_data_file("dataset/locomo10.json", output_file_path)
else:
raise ValueError(f"Invalid technique type: {args.technique_type}")
if __name__ == "__main__":
main()
+185
View File
@@ -0,0 +1,185 @@
import json
import multiprocessing as mp
import os
import time
from collections import defaultdict
from dotenv import load_dotenv
from jinja2 import Template
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
from langgraph.store.memory import InMemoryStore
from langgraph.utils.config import get_store
from langmem import create_manage_memory_tool, create_search_memory_tool
from openai import OpenAI
from prompts import ANSWER_PROMPT
from tqdm import tqdm
load_dotenv()
client = OpenAI()
ANSWER_PROMPT_TEMPLATE = Template(ANSWER_PROMPT)
def get_answer(question, speaker_1_user_id, speaker_1_memories, speaker_2_user_id, speaker_2_memories):
prompt = ANSWER_PROMPT_TEMPLATE.render(
question=question,
speaker_1_user_id=speaker_1_user_id,
speaker_1_memories=speaker_1_memories,
speaker_2_user_id=speaker_2_user_id,
speaker_2_memories=speaker_2_memories,
)
t1 = time.time()
response = client.chat.completions.create(
model=os.getenv("MODEL"), messages=[{"role": "system", "content": prompt}], temperature=0.0
)
t2 = time.time()
return response.choices[0].message.content, t2 - t1
def prompt(state):
"""Prepare the messages for the LLM."""
store = get_store()
memories = store.search(
("memories",),
query=state["messages"][-1].content,
)
system_msg = f"""You are a helpful assistant.
## Memories
<memories>
{memories}
</memories>
"""
return [{"role": "system", "content": system_msg}, *state["messages"]]
class LangMem:
def __init__(
self,
):
self.store = InMemoryStore(
index={
"dims": 1536,
"embed": f"openai:{os.getenv('EMBEDDING_MODEL')}",
}
)
self.checkpointer = MemorySaver() # Checkpoint graph state
self.agent = create_react_agent(
f"openai:{os.getenv('MODEL')}",
prompt=prompt,
tools=[
create_manage_memory_tool(namespace=("memories",)),
create_search_memory_tool(namespace=("memories",)),
],
store=self.store,
checkpointer=self.checkpointer,
)
def add_memory(self, message, config):
return self.agent.invoke({"messages": [{"role": "user", "content": message}]}, config=config)
def search_memory(self, query, config):
try:
t1 = time.time()
response = self.agent.invoke({"messages": [{"role": "user", "content": query}]}, config=config)
t2 = time.time()
return response["messages"][-1].content, t2 - t1
except Exception as e:
print(f"Error in search_memory: {e}")
return "", t2 - t1
class LangMemManager:
def __init__(self, dataset_path):
self.dataset_path = dataset_path
with open(self.dataset_path, "r") as f:
self.data = json.load(f)
def process_all_conversations(self, output_file_path):
OUTPUT = defaultdict(list)
# Process conversations in parallel with multiple workers
def process_conversation(key_value_pair):
key, value = key_value_pair
result = defaultdict(list)
chat_history = value["conversation"]
questions = value["question"]
agent1 = LangMem()
agent2 = LangMem()
config = {"configurable": {"thread_id": f"thread-{key}"}}
speakers = set()
# Identify speakers
for conv in chat_history:
speakers.add(conv["speaker"])
if len(speakers) != 2:
raise ValueError(f"Expected 2 speakers, got {len(speakers)}")
speaker1 = list(speakers)[0]
speaker2 = list(speakers)[1]
# Add memories for each message
for conv in tqdm(chat_history, desc=f"Processing messages {key}", leave=False):
message = f"{conv['timestamp']} | {conv['speaker']}: {conv['text']}"
if conv["speaker"] == speaker1:
agent1.add_memory(message, config)
elif conv["speaker"] == speaker2:
agent2.add_memory(message, config)
else:
raise ValueError(f"Expected speaker1 or speaker2, got {conv['speaker']}")
# Process questions
for q in tqdm(questions, desc=f"Processing questions {key}", leave=False):
category = q["category"]
if int(category) == 5:
continue
answer = q["answer"]
question = q["question"]
response1, speaker1_memory_time = agent1.search_memory(question, config)
response2, speaker2_memory_time = agent2.search_memory(question, config)
generated_answer, response_time = get_answer(question, speaker1, response1, speaker2, response2)
result[key].append(
{
"question": question,
"answer": answer,
"response1": response1,
"response2": response2,
"category": category,
"speaker1_memory_time": speaker1_memory_time,
"speaker2_memory_time": speaker2_memory_time,
"response_time": response_time,
"response": generated_answer,
}
)
return result
# Use multiprocessing to process conversations in parallel
with mp.Pool(processes=10) as pool:
results = list(
tqdm(
pool.imap(process_conversation, list(self.data.items())),
total=len(self.data),
desc="Processing conversations",
)
)
# Combine results from all workers
for result in results:
for key, items in result.items():
OUTPUT[key].extend(items)
# Save final results
with open(output_file_path, "w") as f:
json.dump(OUTPUT, f, indent=4)
+141
View File
@@ -0,0 +1,141 @@
import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dotenv import load_dotenv
from tqdm import tqdm
from mem0 import MemoryClient
load_dotenv()
# Update custom instructions
custom_instructions = """
Generate personal memories that follow these guidelines:
1. Each memory should be self-contained with complete context, including:
- The person's name, do not use "user" while creating memories
- Personal details (career aspirations, hobbies, life circumstances)
- Emotional states and reactions
- Ongoing journeys or future plans
- Specific dates when events occurred
2. Include meaningful personal narratives focusing on:
- Identity and self-acceptance journeys
- Family planning and parenting
- Creative outlets and hobbies
- Mental health and self-care activities
- Career aspirations and education goals
- Important life events and milestones
3. Make each memory rich with specific details rather than general statements
- Include timeframes (exact dates when possible)
- Name specific activities (e.g., "charity race for mental health" rather than just "exercise")
- Include emotional context and personal growth elements
4. Extract memories only from user messages, not incorporating assistant responses
5. Format each memory as a paragraph with a clear narrative structure that captures the person's experience, challenges, and aspirations
"""
class MemoryADD:
def __init__(self, data_path=None, batch_size=2, is_graph=False):
self.mem0_client = MemoryClient(
api_key=os.getenv("MEM0_API_KEY"),
org_id=os.getenv("MEM0_ORGANIZATION_ID"),
project_id=os.getenv("MEM0_PROJECT_ID"),
)
self.mem0_client.update_project(custom_instructions=custom_instructions)
self.batch_size = batch_size
self.data_path = data_path
self.data = None
self.is_graph = is_graph
if data_path:
self.load_data()
def load_data(self):
with open(self.data_path, "r") as f:
self.data = json.load(f)
return self.data
def add_memory(self, user_id, message, metadata, retries=3):
for attempt in range(retries):
try:
_ = self.mem0_client.add(
message, user_id=user_id, version="v2", metadata=metadata, enable_graph=self.is_graph
)
return
except Exception as e:
if attempt < retries - 1:
time.sleep(1) # Wait before retrying
continue
else:
raise e
def add_memories_for_speaker(self, speaker, messages, timestamp, desc):
for i in tqdm(range(0, len(messages), self.batch_size), desc=desc):
batch_messages = messages[i : i + self.batch_size]
self.add_memory(speaker, batch_messages, metadata={"timestamp": timestamp})
def process_conversation(self, item, idx):
conversation = item["conversation"]
speaker_a = conversation["speaker_a"]
speaker_b = conversation["speaker_b"]
speaker_a_user_id = f"{speaker_a}_{idx}"
speaker_b_user_id = f"{speaker_b}_{idx}"
# delete all memories for the two users
self.mem0_client.delete_all(user_id=speaker_a_user_id)
self.mem0_client.delete_all(user_id=speaker_b_user_id)
for key in conversation.keys():
if key in ["speaker_a", "speaker_b"] or "date" in key or "timestamp" in key:
continue
date_time_key = key + "_date_time"
timestamp = conversation[date_time_key]
chats = conversation[key]
messages = []
messages_reverse = []
for chat in chats:
if chat["speaker"] == speaker_a:
messages.append({"role": "user", "content": f"{speaker_a}: {chat['text']}"})
messages_reverse.append({"role": "assistant", "content": f"{speaker_a}: {chat['text']}"})
elif chat["speaker"] == speaker_b:
messages.append({"role": "assistant", "content": f"{speaker_b}: {chat['text']}"})
messages_reverse.append({"role": "user", "content": f"{speaker_b}: {chat['text']}"})
else:
raise ValueError(f"Unknown speaker: {chat['speaker']}")
# add memories for the two users on different threads
thread_a = threading.Thread(
target=self.add_memories_for_speaker,
args=(speaker_a_user_id, messages, timestamp, "Adding Memories for Speaker A"),
)
thread_b = threading.Thread(
target=self.add_memories_for_speaker,
args=(speaker_b_user_id, messages_reverse, timestamp, "Adding Memories for Speaker B"),
)
thread_a.start()
thread_b.start()
thread_a.join()
thread_b.join()
print("Messages added successfully")
def process_all_conversations(self, max_workers=10):
if not self.data:
raise ValueError("No data loaded. Please set data_path and call load_data() first.")
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [executor.submit(self.process_conversation, item, idx) for idx, item in enumerate(self.data)]
for future in futures:
future.result()
+215
View File
@@ -0,0 +1,215 @@
import json
import os
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from dotenv import load_dotenv
from jinja2 import Template
from openai import OpenAI
from prompts import ANSWER_PROMPT, ANSWER_PROMPT_GRAPH
from tqdm import tqdm
from mem0 import MemoryClient
load_dotenv()
class MemorySearch:
def __init__(self, output_path="results.json", top_k=10, filter_memories=False, is_graph=False):
self.mem0_client = MemoryClient(
api_key=os.getenv("MEM0_API_KEY"),
org_id=os.getenv("MEM0_ORGANIZATION_ID"),
project_id=os.getenv("MEM0_PROJECT_ID"),
)
self.top_k = top_k
self.openai_client = OpenAI()
self.results = defaultdict(list)
self.output_path = output_path
self.filter_memories = filter_memories
self.is_graph = is_graph
if self.is_graph:
self.ANSWER_PROMPT = ANSWER_PROMPT_GRAPH
else:
self.ANSWER_PROMPT = ANSWER_PROMPT
def search_memory(self, user_id, query, max_retries=3, retry_delay=1):
start_time = time.time()
retries = 0
while retries < max_retries:
try:
if self.is_graph:
print("Searching with graph")
memories = self.mem0_client.search(
query,
user_id=user_id,
top_k=self.top_k,
filter_memories=self.filter_memories,
enable_graph=True,
output_format="v1.1",
)
else:
memories = self.mem0_client.search(
query, user_id=user_id, top_k=self.top_k, filter_memories=self.filter_memories
)
break
except Exception as e:
print("Retrying...")
retries += 1
if retries >= max_retries:
raise e
time.sleep(retry_delay)
end_time = time.time()
if not self.is_graph:
semantic_memories = [
{
"memory": memory["memory"],
"timestamp": memory["metadata"]["timestamp"],
"score": round(memory["score"], 2),
}
for memory in memories
]
graph_memories = None
else:
semantic_memories = [
{
"memory": memory["memory"],
"timestamp": memory["metadata"]["timestamp"],
"score": round(memory["score"], 2),
}
for memory in memories["results"]
]
graph_memories = [
{"source": relation["source"], "relationship": relation["relationship"], "target": relation["target"]}
for relation in memories["relations"]
]
return semantic_memories, graph_memories, end_time - start_time
def answer_question(self, speaker_1_user_id, speaker_2_user_id, question, answer, category):
speaker_1_memories, speaker_1_graph_memories, speaker_1_memory_time = self.search_memory(
speaker_1_user_id, question
)
speaker_2_memories, speaker_2_graph_memories, speaker_2_memory_time = self.search_memory(
speaker_2_user_id, question
)
search_1_memory = [f"{item['timestamp']}: {item['memory']}" for item in speaker_1_memories]
search_2_memory = [f"{item['timestamp']}: {item['memory']}" for item in speaker_2_memories]
template = Template(self.ANSWER_PROMPT)
answer_prompt = template.render(
speaker_1_user_id=speaker_1_user_id.split("_")[0],
speaker_2_user_id=speaker_2_user_id.split("_")[0],
speaker_1_memories=json.dumps(search_1_memory, indent=4),
speaker_2_memories=json.dumps(search_2_memory, indent=4),
speaker_1_graph_memories=json.dumps(speaker_1_graph_memories, indent=4),
speaker_2_graph_memories=json.dumps(speaker_2_graph_memories, indent=4),
question=question,
)
t1 = time.time()
response = self.openai_client.chat.completions.create(
model=os.getenv("MODEL"), messages=[{"role": "system", "content": answer_prompt}], temperature=0.0
)
t2 = time.time()
response_time = t2 - t1
return (
response.choices[0].message.content,
speaker_1_memories,
speaker_2_memories,
speaker_1_memory_time,
speaker_2_memory_time,
speaker_1_graph_memories,
speaker_2_graph_memories,
response_time,
)
def process_question(self, val, speaker_a_user_id, speaker_b_user_id):
question = val.get("question", "")
answer = val.get("answer", "")
category = val.get("category", -1)
evidence = val.get("evidence", [])
adversarial_answer = val.get("adversarial_answer", "")
(
response,
speaker_1_memories,
speaker_2_memories,
speaker_1_memory_time,
speaker_2_memory_time,
speaker_1_graph_memories,
speaker_2_graph_memories,
response_time,
) = self.answer_question(speaker_a_user_id, speaker_b_user_id, question, answer, category)
result = {
"question": question,
"answer": answer,
"category": category,
"evidence": evidence,
"response": response,
"adversarial_answer": adversarial_answer,
"speaker_1_memories": speaker_1_memories,
"speaker_2_memories": speaker_2_memories,
"num_speaker_1_memories": len(speaker_1_memories),
"num_speaker_2_memories": len(speaker_2_memories),
"speaker_1_memory_time": speaker_1_memory_time,
"speaker_2_memory_time": speaker_2_memory_time,
"speaker_1_graph_memories": speaker_1_graph_memories,
"speaker_2_graph_memories": speaker_2_graph_memories,
"response_time": response_time,
}
# Save results after each question is processed
with open(self.output_path, "w") as f:
json.dump(self.results, f, indent=4)
return result
def process_data_file(self, file_path):
with open(file_path, "r") as f:
data = json.load(f)
for idx, item in tqdm(enumerate(data), total=len(data), desc="Processing conversations"):
qa = item["qa"]
conversation = item["conversation"]
speaker_a = conversation["speaker_a"]
speaker_b = conversation["speaker_b"]
speaker_a_user_id = f"{speaker_a}_{idx}"
speaker_b_user_id = f"{speaker_b}_{idx}"
for question_item in tqdm(
qa, total=len(qa), desc=f"Processing questions for conversation {idx}", leave=False
):
result = self.process_question(question_item, speaker_a_user_id, speaker_b_user_id)
self.results[idx].append(result)
# Save results after each question is processed
with open(self.output_path, "w") as f:
json.dump(self.results, f, indent=4)
# Final save at the end
with open(self.output_path, "w") as f:
json.dump(self.results, f, indent=4)
def process_questions_parallel(self, qa_list, speaker_a_user_id, speaker_b_user_id, max_workers=1):
def process_single_question(val):
result = self.process_question(val, speaker_a_user_id, speaker_b_user_id)
# Save results after each question is processed
with open(self.output_path, "w") as f:
json.dump(self.results, f, indent=4)
return result
with ThreadPoolExecutor(max_workers=max_workers) as executor:
results = list(
tqdm(executor.map(process_single_question, qa_list), total=len(qa_list), desc="Answering Questions")
)
# Final save at the end
with open(self.output_path, "w") as f:
json.dump(self.results, f, indent=4)
return results
+131
View File
@@ -0,0 +1,131 @@
import argparse
import json
import os
import time
from collections import defaultdict
from dotenv import load_dotenv
from jinja2 import Template
from openai import OpenAI
from tqdm import tqdm
load_dotenv()
ANSWER_PROMPT = """
You are an intelligent memory assistant tasked with retrieving accurate information from conversation memories.
# CONTEXT:
You have access to memories from a conversation. These memories contain
timestamped information that may be relevant to answering the question.
# INSTRUCTIONS:
1. Carefully analyze all provided memories
2. Pay special attention to the timestamps to determine the answer
3. If the question asks about a specific event or fact, look for direct evidence in the memories
4. If the memories contain contradictory information, prioritize the most recent memory
5. If there is a question about time references (like "last year", "two months ago", etc.),
calculate the actual date based on the memory timestamp. For example, if a memory from
4 May 2022 mentions "went to India last year," then the trip occurred in 2021.
6. Always convert relative time references to specific dates, months, or years. For example,
convert "last year" to "2022" or "two months ago" to "March 2023" based on the memory
timestamp. Ignore the reference while answering the question.
7. Focus only on the content of the memories. Do not confuse character
names mentioned in memories with the actual users who created those memories.
8. The answer should be less than 5-6 words.
# APPROACH (Think step by step):
1. First, examine all memories that contain information related to the question
2. Examine the timestamps and content of these memories carefully
3. Look for explicit mentions of dates, times, locations, or events that answer the question
4. If the answer requires calculation (e.g., converting relative time references), show your work
5. Formulate a precise, concise answer based solely on the evidence in the memories
6. Double-check that your answer directly addresses the question asked
7. Ensure your final answer is specific and avoids vague time references
Memories:
{{memories}}
Question: {{question}}
Answer:
"""
class OpenAIPredict:
def __init__(self, model="gpt-4o-mini"):
self.model = model
self.openai_client = OpenAI()
self.results = defaultdict(list)
def search_memory(self, idx):
with open(f"memories/{idx}.txt", "r") as file:
memories = file.read()
return memories, 0
def process_question(self, val, idx):
question = val.get("question", "")
answer = val.get("answer", "")
category = val.get("category", -1)
evidence = val.get("evidence", [])
adversarial_answer = val.get("adversarial_answer", "")
response, search_memory_time, response_time, context = self.answer_question(idx, question)
result = {
"question": question,
"answer": answer,
"category": category,
"evidence": evidence,
"response": response,
"adversarial_answer": adversarial_answer,
"search_memory_time": search_memory_time,
"response_time": response_time,
"context": context,
}
return result
def answer_question(self, idx, question):
memories, search_memory_time = self.search_memory(idx)
template = Template(ANSWER_PROMPT)
answer_prompt = template.render(memories=memories, question=question)
t1 = time.time()
response = self.openai_client.chat.completions.create(
model=os.getenv("MODEL"), messages=[{"role": "system", "content": answer_prompt}], temperature=0.0
)
t2 = time.time()
response_time = t2 - t1
return response.choices[0].message.content, search_memory_time, response_time, memories
def process_data_file(self, file_path, output_file_path):
with open(file_path, "r") as f:
data = json.load(f)
for idx, item in tqdm(enumerate(data), total=len(data), desc="Processing conversations"):
qa = item["qa"]
for question_item in tqdm(
qa, total=len(qa), desc=f"Processing questions for conversation {idx}", leave=False
):
result = self.process_question(question_item, idx)
self.results[idx].append(result)
# Save results after each question is processed
with open(output_file_path, "w") as f:
json.dump(self.results, f, indent=4)
# Final save at the end
with open(output_file_path, "w") as f:
json.dump(self.results, f, indent=4)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--output_file_path", type=str, required=True)
args = parser.parse_args()
openai_predict = OpenAIPredict()
openai_predict.process_data_file("../../dataset/locomo10.json", args.output_file_path)
+183
View File
@@ -0,0 +1,183 @@
import json
import os
import time
from collections import defaultdict
import numpy as np
import tiktoken
from dotenv import load_dotenv
from jinja2 import Template
from openai import OpenAI
from tqdm import tqdm
load_dotenv()
PROMPT = """
# Question:
{{QUESTION}}
# Context:
{{CONTEXT}}
# Short answer:
"""
class RAGManager:
def __init__(self, data_path="dataset/locomo10_rag.json", chunk_size=500, k=1):
self.model = os.getenv("MODEL")
self.client = OpenAI()
self.data_path = data_path
self.chunk_size = chunk_size
self.k = k
def generate_response(self, question, context):
template = Template(PROMPT)
prompt = template.render(CONTEXT=context, QUESTION=question)
max_retries = 3
retries = 0
while retries <= max_retries:
try:
t1 = time.time()
response = self.client.chat.completions.create(
model=self.model,
messages=[
{
"role": "system",
"content": "You are a helpful assistant that can answer "
"questions based on the provided context."
"If the question involves timing, use the conversation date for reference."
"Provide the shortest possible answer."
"Use words directly from the conversation when possible."
"Avoid using subjects in your answer.",
},
{"role": "user", "content": prompt},
],
temperature=0,
)
t2 = time.time()
return response.choices[0].message.content.strip(), t2 - t1
except Exception as e:
retries += 1
if retries > max_retries:
raise e
time.sleep(1) # Wait before retrying
def clean_chat_history(self, chat_history):
cleaned_chat_history = ""
for c in chat_history:
cleaned_chat_history += f"{c['timestamp']} | {c['speaker']}: {c['text']}\n"
return cleaned_chat_history
def calculate_embedding(self, document):
response = self.client.embeddings.create(model=os.getenv("EMBEDDING_MODEL"), input=document)
return response.data[0].embedding
def calculate_similarity(self, embedding1, embedding2):
return np.dot(embedding1, embedding2) / (np.linalg.norm(embedding1) * np.linalg.norm(embedding2))
def search(self, query, chunks, embeddings, k=1):
"""
Search for the top-k most similar chunks to the query.
Args:
query: The query string
chunks: List of text chunks
embeddings: List of embeddings for each chunk
k: Number of top chunks to return (default: 1)
Returns:
combined_chunks: The combined text of the top-k chunks
search_time: Time taken for the search
"""
t1 = time.time()
query_embedding = self.calculate_embedding(query)
similarities = [self.calculate_similarity(query_embedding, embedding) for embedding in embeddings]
# Get indices of top-k most similar chunks
if k == 1:
# Original behavior - just get the most similar chunk
top_indices = [np.argmax(similarities)]
else:
# Get indices of top-k chunks
top_indices = np.argsort(similarities)[-k:][::-1]
# Combine the top-k chunks
combined_chunks = "\n<->\n".join([chunks[i] for i in top_indices])
t2 = time.time()
return combined_chunks, t2 - t1
def create_chunks(self, chat_history, chunk_size=500):
"""
Create chunks using tiktoken for more accurate token counting
"""
# Get the encoding for the model
encoding = tiktoken.encoding_for_model(os.getenv("EMBEDDING_MODEL"))
documents = self.clean_chat_history(chat_history)
if chunk_size == -1:
return [documents], []
chunks = []
# Encode the document
tokens = encoding.encode(documents)
# Split into chunks based on token count
for i in range(0, len(tokens), chunk_size):
chunk_tokens = tokens[i : i + chunk_size]
chunk = encoding.decode(chunk_tokens)
chunks.append(chunk)
embeddings = []
for chunk in chunks:
embedding = self.calculate_embedding(chunk)
embeddings.append(embedding)
return chunks, embeddings
def process_all_conversations(self, output_file_path):
with open(self.data_path, "r") as f:
data = json.load(f)
FINAL_RESULTS = defaultdict(list)
for key, value in tqdm(data.items(), desc="Processing conversations"):
chat_history = value["conversation"]
questions = value["question"]
chunks, embeddings = self.create_chunks(chat_history, self.chunk_size)
for item in tqdm(questions, desc="Answering questions", leave=False):
question = item["question"]
answer = item.get("answer", "")
category = item["category"]
if self.chunk_size == -1:
context = chunks[0]
search_time = 0
else:
context, search_time = self.search(question, chunks, embeddings, k=self.k)
response, response_time = self.generate_response(question, context)
FINAL_RESULTS[key].append(
{
"question": question,
"answer": answer,
"category": category,
"context": context,
"response": response,
"search_time": search_time,
"response_time": response_time,
}
)
with open(output_file_path, "w+") as f:
json.dump(FINAL_RESULTS, f, indent=4)
# Save results
with open(output_file_path, "w+") as f:
json.dump(FINAL_RESULTS, f, indent=4)
+3
View File
@@ -0,0 +1,3 @@
TECHNIQUES = ["mem0", "rag", "langmem", "zep", "openai"]
METHODS = ["add", "search"]
+76
View File
@@ -0,0 +1,76 @@
import argparse
import json
import os
from dotenv import load_dotenv
from tqdm import tqdm
from zep_cloud import Message
from zep_cloud.client import Zep
load_dotenv()
class ZepAdd:
def __init__(self, data_path=None):
self.zep_client = Zep(api_key=os.getenv("ZEP_API_KEY"))
self.data_path = data_path
self.data = None
if data_path:
self.load_data()
def load_data(self):
with open(self.data_path, "r") as f:
self.data = json.load(f)
return self.data
def process_conversation(self, run_id, item, idx):
conversation = item["conversation"]
user_id = f"run_id_{run_id}_experiment_user_{idx}"
session_id = f"run_id_{run_id}_experiment_session_{idx}"
# # delete all memories for the two users
# self.zep_client.user.delete(user_id=user_id)
# self.zep_client.memory.delete(session_id=session_id)
self.zep_client.user.add(user_id=user_id)
self.zep_client.memory.add_session(
user_id=user_id,
session_id=session_id,
)
print("Starting to add memories... for user", user_id)
for key in tqdm(conversation.keys(), desc=f"Processing user {user_id}"):
if key in ["speaker_a", "speaker_b"] or "date" in key:
continue
date_time_key = key + "_date_time"
timestamp = conversation[date_time_key]
chats = conversation[key]
for chat in tqdm(chats, desc=f"Adding chats for {key}", leave=False):
self.zep_client.memory.add(
session_id=session_id,
messages=[
Message(
role=chat["speaker"],
role_type="user",
content=f"{timestamp}: {chat['text']}",
)
],
)
def process_all_conversations(self, run_id):
if not self.data:
raise ValueError("No data loaded. Please set data_path and call load_data() first.")
for idx, item in tqdm(enumerate(self.data)):
if idx == 0:
self.process_conversation(run_id, item, idx)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--run_id", type=str, required=True)
args = parser.parse_args()
zep_add = ZepAdd(data_path="../../dataset/locomo10.json")
zep_add.process_all_conversations(args.run_id)
+140
View File
@@ -0,0 +1,140 @@
import argparse
import json
import os
import time
from collections import defaultdict
from dotenv import load_dotenv
from jinja2 import Template
from openai import OpenAI
from prompts import ANSWER_PROMPT_ZEP
from tqdm import tqdm
from zep_cloud import EntityEdge, EntityNode
from zep_cloud.client import Zep
load_dotenv()
TEMPLATE = """
FACTS and ENTITIES represent relevant context to the current conversation.
# These are the most relevant facts and their valid date ranges
# format: FACT (Date range: from - to)
{facts}
# These are the most relevant entities
# ENTITY_NAME: entity summary
{entities}
"""
class ZepSearch:
def __init__(self):
self.zep_client = Zep(api_key=os.getenv("ZEP_API_KEY"))
self.results = defaultdict(list)
self.openai_client = OpenAI()
def format_edge_date_range(self, edge: EntityEdge) -> str:
# return f"{datetime(edge.valid_at).strftime('%Y-%m-%d %H:%M:%S') if edge.valid_at else 'date unknown'} - {(edge.invalid_at.strftime('%Y-%m-%d %H:%M:%S') if edge.invalid_at else 'present')}"
return f"{edge.valid_at if edge.valid_at else 'date unknown'} - {(edge.invalid_at if edge.invalid_at else 'present')}"
def compose_search_context(self, edges: list[EntityEdge], nodes: list[EntityNode]) -> str:
facts = [f" - {edge.fact} ({self.format_edge_date_range(edge)})" for edge in edges]
entities = [f" - {node.name}: {node.summary}" for node in nodes]
return TEMPLATE.format(facts="\n".join(facts), entities="\n".join(entities))
def search_memory(self, run_id, idx, query, max_retries=3, retry_delay=1):
start_time = time.time()
retries = 0
while retries < max_retries:
try:
user_id = f"run_id_{run_id}_experiment_user_{idx}"
edges_results = (
self.zep_client.graph.search(
user_id=user_id, reranker="cross_encoder", query=query, scope="edges", limit=20
)
).edges
node_results = (
self.zep_client.graph.search(user_id=user_id, reranker="rrf", query=query, scope="nodes", limit=20)
).nodes
context = self.compose_search_context(edges_results, node_results)
break
except Exception as e:
print("Retrying...")
retries += 1
if retries >= max_retries:
raise e
time.sleep(retry_delay)
end_time = time.time()
return context, end_time - start_time
def process_question(self, run_id, val, idx):
question = val.get("question", "")
answer = val.get("answer", "")
category = val.get("category", -1)
evidence = val.get("evidence", [])
adversarial_answer = val.get("adversarial_answer", "")
response, search_memory_time, response_time, context = self.answer_question(run_id, idx, question)
result = {
"question": question,
"answer": answer,
"category": category,
"evidence": evidence,
"response": response,
"adversarial_answer": adversarial_answer,
"search_memory_time": search_memory_time,
"response_time": response_time,
"context": context,
}
return result
def answer_question(self, run_id, idx, question):
context, search_memory_time = self.search_memory(run_id, idx, question)
template = Template(ANSWER_PROMPT_ZEP)
answer_prompt = template.render(memories=context, question=question)
t1 = time.time()
response = self.openai_client.chat.completions.create(
model=os.getenv("MODEL"), messages=[{"role": "system", "content": answer_prompt}], temperature=0.0
)
t2 = time.time()
response_time = t2 - t1
return response.choices[0].message.content, search_memory_time, response_time, context
def process_data_file(self, file_path, run_id, output_file_path):
with open(file_path, "r") as f:
data = json.load(f)
for idx, item in tqdm(enumerate(data), total=len(data), desc="Processing conversations"):
qa = item["qa"]
for question_item in tqdm(
qa, total=len(qa), desc=f"Processing questions for conversation {idx}", leave=False
):
result = self.process_question(run_id, question_item, idx)
self.results[idx].append(result)
# Save results after each question is processed
with open(output_file_path, "w") as f:
json.dump(self.results, f, indent=4)
# Final save at the end
with open(output_file_path, "w") as f:
json.dump(self.results, f, indent=4)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--run_id", type=str, required=True)
args = parser.parse_args()
zep_search = ZepSearch()
zep_search.process_data_file("../../dataset/locomo10.json", args.run_id, "results/zep_search_results.json")
+1 -28
View File
@@ -3,9 +3,7 @@
*
* Used to stitch PostHog identities: SDKs and CLIs persist anonymous
* distinct_id values here, and the TS MemoryClient reads those on init to
* fire $identify and merge them into the email identity. The OSS user_id is
* only a stitch candidate when telemetry.oss_used_at proves OSS Memory
* actually ran here (the platform client never mints ids itself).
* fire $identify and merge them into the email identity.
*
* Node-only. Browsers (no `process.versions.node`) no-op.
*/
@@ -13,7 +11,6 @@
export interface Mem0AnonIds {
oss?: string;
cli?: string;
ossUsedAt?: string;
aliasedPairs: string[];
}
@@ -113,10 +110,6 @@ export async function readMem0AnonIds(): Promise<Mem0AnonIds | null> {
typeof telemetry.anonymous_id === "string"
? telemetry.anonymous_id
: undefined,
ossUsedAt:
typeof telemetry.oss_used_at === "string" && telemetry.oss_used_at
? telemetry.oss_used_at
: undefined,
aliasedPairs: Array.isArray(telemetry.aliased_pairs)
? telemetry.aliased_pairs.filter(
(item: unknown) => typeof item === "string",
@@ -144,26 +137,6 @@ export async function isMem0Aliased(
return aliasedPairs.includes(aliasPairMarker(node, anonId, email));
}
export async function markMem0OssUsed(): Promise<void> {
const node = await getNodeFs();
if (!node) return;
try {
const config = loadConfig(node) ?? {};
const telemetry =
config.telemetry && typeof config.telemetry === "object"
? config.telemetry
: {};
if (typeof telemetry.oss_used_at === "string" && telemetry.oss_used_at) {
return;
}
telemetry.oss_used_at = new Date().toISOString();
config.telemetry = telemetry;
writeConfig(node, config);
} catch {
// Best-effort: read-only filesystems and unwritable paths just skip.
}
}
export async function markMem0Aliased(
anonId: string,
email: string,
+11 -10
View File
@@ -27,7 +27,12 @@ import {
isTelemetryEnabled,
telemetry,
} from "./telemetry";
import { isMem0Aliased, markMem0Aliased, readMem0AnonIds } from "./config";
import {
getOrCreateMem0UserId,
isMem0Aliased,
markMem0Aliased,
readMem0AnonIds,
} from "./config";
import { camelToSnake, camelToSnakeKeys, snakeToCamelKeys } from "./utils";
import { createExceptionFromResponse, MemoryError } from "../common/exceptions";
@@ -146,16 +151,12 @@ export default class MemoryClient {
try {
const email = this.telemetryId;
if (!email || !email.includes("@")) return;
const sharedAnonId = await getOrCreateMem0UserId();
const anonIds = await readMem0AnonIds();
if (!anonIds) return;
// The OSS user_id is only a stitch candidate when telemetry.oss_used_at
// proves OSS Memory actually ran here; never mint an id from the
// platform client path. The CLI anon id is presence-based (only the
// CLIs write it).
const candidates = [
anonIds.ossUsedAt ? anonIds.oss : undefined,
anonIds.cli,
].filter((id): id is string => !!id && id !== email);
if (!anonIds && !sharedAnonId) return;
const candidates = [anonIds?.oss || sharedAnonId, anonIds?.cli].filter(
(id): id is string => !!id && id !== email,
);
const seen = new Set<string>();
for (const anonId of candidates) {
if (seen.has(anonId) || (await isMem0Aliased(anonId, email))) continue;
-1
View File
@@ -85,7 +85,6 @@ class UnifiedTelemetry implements TelemetryClient {
properties: {
$anon_distinct_id: anonId,
client_source: "typescript",
client_version: version,
$lib: "posthog-node",
},
};
@@ -13,7 +13,6 @@ import {
getOrCreateMem0UserId,
isMem0Aliased,
markMem0Aliased,
markMem0OssUsed,
readMem0AnonIds,
} from "../config";
import { TEST_API_KEY } from "./helpers";
@@ -146,55 +145,6 @@ describe("config.ts — readMem0AnonIds / markMem0Aliased", () => {
markMem0Aliased("oss-uuid", "user@example.com"),
).resolves.toBeUndefined();
});
test("markMem0OssUsed writes marker preserving other fields", async () => {
fs.writeFileSync(
path.join(tmpHome, "config.json"),
JSON.stringify({
user_id: "oss-uuid",
telemetry: { anonymous_id: "cli-anon", aliased_pairs: ["pair-marker"] },
}),
);
await markMem0OssUsed();
const written = JSON.parse(
fs.readFileSync(path.join(tmpHome, "config.json"), "utf8"),
);
expect(written.user_id).toBe("oss-uuid");
expect(written.telemetry.anonymous_id).toBe("cli-anon");
expect(written.telemetry.aliased_pairs).toEqual(["pair-marker"]);
expect(written.telemetry.oss_used_at).toBeTruthy();
});
test("markMem0OssUsed is idempotent", async () => {
await markMem0OssUsed();
const first = JSON.parse(
fs.readFileSync(path.join(tmpHome, "config.json"), "utf8"),
).telemetry.oss_used_at;
await markMem0OssUsed();
const second = JSON.parse(
fs.readFileSync(path.join(tmpHome, "config.json"), "utf8"),
).telemetry.oss_used_at;
expect(second).toBe(first);
});
test("readMem0AnonIds surfaces ossUsedAt", async () => {
fs.writeFileSync(
path.join(tmpHome, "config.json"),
JSON.stringify({
user_id: "oss-uuid",
telemetry: { oss_used_at: "2026-01-01T00:00:00.000Z" },
}),
);
const ids = await readMem0AnonIds();
expect(ids?.oss).toBe("oss-uuid");
expect(ids?.ossUsedAt).toBe("2026-01-01T00:00:00.000Z");
});
test("markMem0OssUsed does not throw when target dir is unwritable", async () => {
fs.writeFileSync(path.join(tmpHome, "blocker"), "x");
process.env.MEM0_DIR = path.join(tmpHome, "blocker"); // file used as dir
await expect(markMem0OssUsed()).resolves.toBeUndefined();
});
});
// ─── telemetry.captureIdentify ───────────────────────────────
@@ -216,7 +166,6 @@ describe("telemetry.captureIdentify", () => {
expect(payload.event).toBe("$identify");
expect(payload.distinct_id).toBe("user@example.com");
expect(payload.properties.$anon_distinct_id).toBe("anon-uuid");
expect(payload.properties.client_version).toBeTruthy();
expect(payload.properties.$process_person_profile).toBeUndefined();
});
@@ -271,10 +220,7 @@ describe("MemoryClient — _maybeAliasAnonToEmail", () => {
test("fires $identify on first init and persists pair marker", async () => {
fs.writeFileSync(
path.join(tmpHome, "config.json"),
JSON.stringify({
user_id: "oss-uuid",
telemetry: { oss_used_at: "2026-01-01T00:00:00.000Z" },
}),
JSON.stringify({ user_id: "oss-uuid" }),
);
const fetchMock = setupMockFetchWithPostHog();
@@ -298,38 +244,17 @@ describe("MemoryClient — _maybeAliasAnonToEmail", () => {
expect(written.telemetry.aliased_pairs).toHaveLength(1);
});
test("platform-only init mints nothing and fires no $identify", async () => {
// Regression: the client used to create the shared anon id itself and
// immediately stitch it, merging an event-less anon person into every
// platform user. The platform path must now be read-only.
test("platform-first init creates shared anon ID and identifies it", async () => {
const fetchMock = setupMockFetchWithPostHog();
const client = makeStubClient("test@example.com");
await (client as any)._maybeAliasAnonToEmail();
expect(fs.existsSync(path.join(tmpHome, "config.json"))).toBe(false);
const identifyCalls = (fetchMock.mock.calls as any[]).filter(
([, init]: [string, RequestInit]) => {
if (!init?.body) return false;
return JSON.parse(init.body as string).event === "$identify";
},
const written = JSON.parse(
fs.readFileSync(path.join(tmpHome, "config.json"), "utf8"),
);
expect(identifyCalls.length).toBe(0);
});
test("oss id without oss_used_at marker is not stitched; cli id still is", async () => {
fs.writeFileSync(
path.join(tmpHome, "config.json"),
JSON.stringify({
user_id: "oss-uuid",
telemetry: { anonymous_id: "cli-anon" },
}),
);
const fetchMock = setupMockFetchWithPostHog();
const client = makeStubClient("test@example.com");
await (client as any)._maybeAliasAnonToEmail();
expect(written.user_id).toBeTruthy();
expect(written.telemetry.aliased_pairs).toHaveLength(1);
const identifyCalls = (fetchMock.mock.calls as any[]).filter(
([, init]: [string, RequestInit]) => {
@@ -339,32 +264,8 @@ describe("MemoryClient — _maybeAliasAnonToEmail", () => {
);
expect(identifyCalls.length).toBe(1);
const body = JSON.parse(identifyCalls[0][1].body);
expect(body.properties.$anon_distinct_id).toBe("cli-anon");
});
test("stitches config written by the Python SDK (cross-SDK key compat)", async () => {
// Python writes datetime.now(timezone.utc).isoformat(): +00:00 offset.
fs.writeFileSync(
path.join(tmpHome, "config.json"),
JSON.stringify({
user_id: "oss-uuid",
telemetry: { oss_used_at: "2026-01-01T00:00:00+00:00" },
}),
);
const fetchMock = setupMockFetchWithPostHog();
const client = makeStubClient("test@example.com");
await (client as any)._maybeAliasAnonToEmail();
const identifyCalls = (fetchMock.mock.calls as any[]).filter(
([, init]: [string, RequestInit]) => {
if (!init?.body) return false;
return JSON.parse(init.body as string).event === "$identify";
},
);
expect(identifyCalls.length).toBe(1);
const body = JSON.parse(identifyCalls[0][1].body);
expect(body.properties.$anon_distinct_id).toBe("oss-uuid");
expect(body.distinct_id).toBe("test@example.com");
expect(body.properties.$anon_distinct_id).toBe(written.user_id);
});
test("second init does not refire $identify", async () => {
@@ -372,7 +273,7 @@ describe("MemoryClient — _maybeAliasAnonToEmail", () => {
path.join(tmpHome, "config.json"),
JSON.stringify({
user_id: "oss-uuid",
telemetry: { oss_used_at: "2026-01-01T00:00:00.000Z" },
telemetry: {},
}),
);
await markMem0Aliased("oss-uuid", "test@example.com");
@@ -395,10 +296,7 @@ describe("MemoryClient — _maybeAliasAnonToEmail", () => {
path.join(tmpHome, "config.json"),
JSON.stringify({
user_id: "oss-uuid",
telemetry: {
anonymous_id: "cli-anon",
oss_used_at: "2026-01-01T00:00:00.000Z",
},
telemetry: { anonymous_id: "cli-anon" },
}),
);
const fetchMock = setupMockFetchWithPostHog();
@@ -509,18 +407,4 @@ describe("config.ts in browser-like environment", () => {
jest.resetModules();
}
});
test("markMem0OssUsed is a no-op when not Node", async () => {
const originalProcess = global.process;
// @ts-expect-error force-undefining global to simulate a browser
delete global.process;
try {
jest.resetModules();
const { markMem0OssUsed: browserMark } = await import("../config");
await expect(browserMark()).resolves.toBeUndefined();
} finally {
global.process = originalProcess;
jest.resetModules();
}
});
});
+2 -8
View File
@@ -40,7 +40,7 @@ import {
} from "./memory.types";
import { parse_vision_messages } from "../utils/memory";
import { HistoryManager } from "../storage/base";
import { captureClientEvent, isTelemetryEnabled } from "../utils/telemetry";
import { captureClientEvent } from "../utils/telemetry";
import {
detectScaleThresholdFromAddResult,
detectScaleThresholdFromTopK,
@@ -72,7 +72,7 @@ import {
ScoredResult,
} from "../utils/scoring";
import { getDefaultVectorStoreDbPath } from "../utils/sqlite";
import { getOrCreateMem0UserId, markMem0OssUsed } from "../../../client/config";
import { getOrCreateMem0UserId } from "../../../client/config";
// Entity params that must be passed via filters - check both snake_case and camelCase
const ENTITY_PARAMS = [
@@ -470,12 +470,6 @@ export class Memory {
try {
await this._getTelemetryId();
// Record genuine OSS usage so the platform client may stitch this
// machine's anon id to a platform email later.
if (isTelemetryEnabled()) {
await markMem0OssUsed();
}
// Capture initialization event
await captureClientEvent("init", this, {
api_version: this.apiVersion,
+5 -7
View File
@@ -19,13 +19,16 @@ from mem0.client.types import (
from mem0.client.utils import api_error_handler
# Exception classes are referenced in docstrings only
from mem0.memory.setup import get_user_id, is_aliased, is_oss_used, mark_aliased, read_anon_ids
from mem0.memory.setup import get_user_id, is_aliased, mark_aliased, read_anon_ids, setup_config
from mem0.memory.telemetry import capture_client_event, client_telemetry
logger = logging.getLogger(__name__)
warnings.filterwarnings("default", category=DeprecationWarning)
# Setup user config
setup_config()
# Entity parameters that must be passed via filters, not top-level
ENTITY_PARAMS = frozenset({"user_id", "agent_id", "app_id", "run_id"})
@@ -42,10 +45,6 @@ def _validate_and_trim_search_query(query: str) -> str:
def _maybe_alias_anon_to_email(user_email):
"""Fire $identify per prior anon ID so PostHog merges them into email.
The OSS user_id is only a stitch candidate when telemetry.oss_used_at
proves OSS Memory actually ran here; bare user_id values were historically
minted at import time by platform-only clients and must not be stitched.
The CLI anon id keeps presence-based stitching (only the CLIs write it).
Idempotent via telemetry.aliased_pairs: only writes markers when
telemetry is actually enabled, so disabling/re-enabling MEM0_TELEMETRY still works.
Best-effort: never raises.
@@ -56,9 +55,8 @@ def _maybe_alias_anon_to_email(user_email):
return
try:
anon_ids = read_anon_ids()
oss_id = anon_ids.get("oss") if is_oss_used() else None
seen = set()
for anon_id in (oss_id, anon_ids.get("cli")):
for anon_id in (anon_ids.get("oss"), anon_ids.get("cli")):
if not anon_id or anon_id == user_email or anon_id in seen:
continue
seen.add(anon_id)
+2 -11
View File
@@ -24,7 +24,7 @@ from mem0.configs.prompts import (
)
from mem0.exceptions import ValidationError as Mem0ValidationError
from mem0.memory.base import MemoryBase
from mem0.memory.setup import mark_oss_used, mem0_dir, setup_config
from mem0.memory.setup import mem0_dir, setup_config
from mem0.memory.storage import SQLiteManager
from mem0.memory.telemetry import MEM0_TELEMETRY, capture_event
from mem0.memory.notices import (
@@ -370,6 +370,7 @@ def _entity_collection_name(provider: str, collection_name: str) -> str:
return f"{collection_name}{separator}entities"
setup_config()
logger = logging.getLogger(__name__)
_PROJECT_UPDATE_UNSUPPORTED_ERROR = "Project updates are not supported by the OSS Memory SDK."
@@ -468,13 +469,6 @@ class Memory(MemoryBase):
self.config.vector_store.provider,
)
# Mint the OSS anon id (and record genuine OSS usage) only on real OSS
# init, never at import time: platform-only clients must not create
# stitch candidates. Must run before the first capture_event so the
# telemetry singleton picks up the minted user_id.
if MEM0_TELEMETRY:
setup_config()
mark_oss_used()
capture_event("mem0.init", self, {"sync_type": "sync"})
@property
@@ -2015,9 +2009,6 @@ class AsyncMemory(MemoryBase):
self.config.vector_store.provider,
)
if MEM0_TELEMETRY:
setup_config()
mark_oss_used()
capture_event("mem0.init", self, {"sync_type": "async"})
@property
+4 -32
View File
@@ -3,7 +3,6 @@ import logging
import os
import tempfile
import uuid
from datetime import datetime, timezone
from hashlib import sha256
# Set up the directory path
@@ -57,12 +56,10 @@ def _write_config(config):
def setup_config():
"""Ensure ~/.mem0/config.json exists with a top-level user_id.
Called lazily from OSS Memory/AsyncMemory init (never at import time, so
platform-only MemoryClient usage mints no anon id). Idempotent: backfills
user_id for users whose config was written by the CLI (which writes
telemetry.anonymous_id but no top-level user_id). Without this, OSS Python
telemetry is silently dropped because get_user_id() returns None when
user_id is missing.
Idempotent: backfills user_id for users whose config was written by the
CLI (which writes telemetry.anonymous_id but no top-level user_id).
Without this, OSS Python telemetry is silently dropped because
get_user_id() returns None when user_id is missing.
"""
config = _load_config()
if config.get("user_id"):
@@ -134,31 +131,6 @@ def mark_aliased(anon_id, email):
_write_config(config)
def is_oss_used():
"""Return whether OSS Memory has genuinely initialized on this machine.
Gates identity stitching: the top-level user_id is only a stitch candidate
when telemetry.oss_used_at is present, because bare user_id values were
historically minted at import time by platform-only clients.
"""
config = _load_config()
telemetry = config.get("telemetry") if isinstance(config.get("telemetry"), dict) else {}
return bool(telemetry.get("oss_used_at"))
def mark_oss_used():
"""Persist telemetry.oss_used_at on first genuine OSS Memory init. Never raises."""
config = _load_config()
telemetry = config.get("telemetry")
if not isinstance(telemetry, dict):
telemetry = {}
if telemetry.get("oss_used_at"):
return
telemetry["oss_used_at"] = datetime.now(timezone.utc).isoformat()
config["telemetry"] = telemetry
_write_config(config)
def get_or_create_user_id(vector_store=None):
"""Store user_id in vector store and return it.
+1 -5
View File
@@ -129,11 +129,7 @@ class AnonymousTelemetry:
self.posthog.capture(
distinct_id=email,
event="$identify",
properties={
"$anon_distinct_id": anon_id,
"client_source": "python",
"client_version": mem0.__version__,
},
properties={"$anon_distinct_id": anon_id, "client_source": "python"},
)
return True
except Exception as e:
+4 -201
View File
@@ -164,68 +164,11 @@ class TestMarkAliased:
assert not setup_module.is_aliased("oss-uuid", "other@example.com")
# ─── oss_used marker ─────────────────────────────────────────────────────────
class TestOssUsedMarker:
def test_mark_writes_marker_preserving_fields(self, tmp_mem0_dir):
import mem0.memory.setup as setup_module
_write_config(
tmp_mem0_dir,
{
"user_id": "oss-uuid",
"telemetry": {"anonymous_id": "cli-anon", "aliased_pairs": ["pair-marker"]},
},
)
setup_module.mark_oss_used()
config = json.loads((tmp_mem0_dir / "config.json").read_text())
assert config["user_id"] == "oss-uuid"
assert config["telemetry"]["anonymous_id"] == "cli-anon"
assert config["telemetry"]["aliased_pairs"] == ["pair-marker"]
assert config["telemetry"]["oss_used_at"]
assert setup_module.is_oss_used()
def test_mark_is_idempotent(self, tmp_mem0_dir):
import mem0.memory.setup as setup_module
setup_module.mark_oss_used()
first = json.loads((tmp_mem0_dir / "config.json").read_text())["telemetry"]["oss_used_at"]
setup_module.mark_oss_used()
second = json.loads((tmp_mem0_dir / "config.json").read_text())["telemetry"]["oss_used_at"]
assert first == second
def test_is_oss_used_false_when_no_config(self, tmp_mem0_dir):
import mem0.memory.setup as setup_module
assert not setup_module.is_oss_used()
def test_is_oss_used_false_without_marker(self, tmp_mem0_dir):
import mem0.memory.setup as setup_module
_write_config(tmp_mem0_dir, {"user_id": "oss-uuid", "telemetry": {"anonymous_id": "cli-anon"}})
assert not setup_module.is_oss_used()
def test_is_oss_used_false_on_malformed_json(self, tmp_mem0_dir):
import mem0.memory.setup as setup_module
(tmp_mem0_dir / "config.json").write_text("{not json")
assert not setup_module.is_oss_used()
def test_mark_does_not_raise_when_config_path_unwritable(self, tmp_mem0_dir):
import mem0.memory.setup as setup_module
(tmp_mem0_dir / "config.json").mkdir() # open() on a directory fails
setup_module.mark_oss_used() # must not raise
assert not setup_module.is_oss_used()
# ─── capture_identify ────────────────────────────────────────────────────────
class TestCaptureIdentify:
def test_fires_identify_with_anon_distinct_id(self):
import mem0
import mem0.memory.telemetry as telemetry_module
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
@@ -238,7 +181,6 @@ class TestCaptureIdentify:
assert kwargs["distinct_id"] == "user@example.com"
assert kwargs["event"] == "$identify"
assert kwargs["properties"]["$anon_distinct_id"] == "anon-123"
assert kwargs["properties"]["client_version"] == mem0.__version__
def test_skips_when_anon_equals_email(self):
import mem0.memory.telemetry as telemetry_module
@@ -288,7 +230,8 @@ class TestCaptureIdentify:
class TestMaybeAliasAnonToEmail:
"""Test the alias helper in isolation by mocking out the config readers
and the telemetry client, so no real ~/.mem0/config.json is touched."""
and the telemetry client, since module-level setup_config() side effects
make end-to-end fixturing awkward."""
def test_fires_identify_for_oss_uuid(self):
from mem0.client import main as client_main
@@ -299,7 +242,6 @@ class TestMaybeAliasAnonToEmail:
"read_anon_ids",
return_value={"oss": "oss-uuid", "cli": None, "aliased_pairs": []},
),
patch.object(client_main, "is_oss_used", return_value=True),
patch.object(client_main, "is_aliased", return_value=False),
patch.object(client_main, "mark_aliased") as mark,
patch.object(client_main, "client_telemetry") as telemetry,
@@ -310,7 +252,6 @@ class TestMaybeAliasAnonToEmail:
mark.assert_called_once_with("oss-uuid", "user@example.com")
def test_fires_identify_for_cli_anon(self):
"""The CLI anon id is presence-based: it stitches even without the oss_used marker."""
from mem0.client import main as client_main
with (
@@ -319,7 +260,6 @@ class TestMaybeAliasAnonToEmail:
"read_anon_ids",
return_value={"oss": None, "cli": "cli-anon-xyz", "aliased_pairs": []},
),
patch.object(client_main, "is_oss_used", return_value=False),
patch.object(client_main, "is_aliased", return_value=False),
patch.object(client_main, "mark_aliased"),
patch.object(client_main, "client_telemetry") as telemetry,
@@ -337,7 +277,6 @@ class TestMaybeAliasAnonToEmail:
"read_anon_ids",
return_value={"oss": "oss-uuid", "cli": "cli-anon", "aliased_pairs": []},
),
patch.object(client_main, "is_oss_used", return_value=True),
patch.object(client_main, "is_aliased", return_value=False),
patch.object(client_main, "mark_aliased"),
patch.object(client_main, "client_telemetry") as telemetry,
@@ -349,26 +288,6 @@ class TestMaybeAliasAnonToEmail:
assert ("oss-uuid", "user@example.com") in calls
assert ("cli-anon", "user@example.com") in calls
def test_oss_id_requires_oss_used_marker(self):
"""Without proof of real OSS usage, the oss user_id must never be
stitched (it may have been minted by an old platform-only client)."""
from mem0.client import main as client_main
with (
patch.object(
client_main,
"read_anon_ids",
return_value={"oss": "oss-uuid", "cli": "cli-anon", "aliased_pairs": []},
),
patch.object(client_main, "is_oss_used", return_value=False),
patch.object(client_main, "is_aliased", return_value=False),
patch.object(client_main, "mark_aliased"),
patch.object(client_main, "client_telemetry") as telemetry,
):
telemetry.capture_identify.return_value = True
client_main._maybe_alias_anon_to_email("user@example.com")
telemetry.capture_identify.assert_called_once_with("cli-anon", "user@example.com")
def test_skips_when_pair_already_aliased(self):
from mem0.client import main as client_main
@@ -378,7 +297,6 @@ class TestMaybeAliasAnonToEmail:
"read_anon_ids",
return_value={"oss": "oss-uuid", "cli": None, "aliased_pairs": ["pair-marker"]},
),
patch.object(client_main, "is_oss_used", return_value=True),
patch.object(client_main, "is_aliased", return_value=True),
patch.object(client_main, "mark_aliased") as mark,
patch.object(client_main, "client_telemetry") as telemetry,
@@ -424,7 +342,6 @@ class TestMaybeAliasAnonToEmail:
"read_anon_ids",
return_value={"oss": "oss-uuid", "cli": None, "aliased_pairs": []},
),
patch.object(client_main, "is_oss_used", return_value=True),
patch.object(client_main, "is_aliased", return_value=False),
patch.object(client_main, "mark_aliased") as mark,
patch.object(client_main, "client_telemetry", mock_telemetry),
@@ -442,7 +359,6 @@ class TestMaybeAliasAnonToEmail:
"read_anon_ids",
return_value={"oss": "user@example.com", "cli": None, "aliased_pairs": []},
),
patch.object(client_main, "is_oss_used", return_value=True),
patch.object(client_main, "is_aliased", return_value=False),
patch.object(client_main, "mark_aliased"),
patch.object(client_main, "client_telemetry") as telemetry,
@@ -470,12 +386,8 @@ class TestEndToEndIdempotency:
calls fire $identify exactly once thanks to the persisted pair marker."""
def test_second_call_is_noop_after_pair_marker_persisted(self, tmp_mem0_dir):
# Pre-populate config with an OSS user_id plus the oss_used_at marker
# (without the marker the oss id is no longer a stitch candidate).
_write_config(
tmp_mem0_dir,
{"user_id": "oss-uuid", "telemetry": {"oss_used_at": "2026-01-01T00:00:00+00:00"}},
)
# Pre-populate config with an OSS user_id only.
_write_config(tmp_mem0_dir, {"user_id": "oss-uuid"})
# Reload setup so it uses the tempdir, then reload client.main so it
# picks up the freshly-loaded read_anon_ids/mark_aliased bindings.
import mem0.memory.setup as setup_module
@@ -497,112 +409,3 @@ class TestEndToEndIdempotency:
config = json.loads((tmp_mem0_dir / "config.json").read_text())
assert len(config["telemetry"]["aliased_pairs"]) == 1
def test_oss_then_platform_stitches_exactly_once(self, tmp_mem0_dir):
"""Genuine OSS usage (setup_config + mark_oss_used, as Memory.__init__
does) followed by platform inits stitches the minted id exactly once."""
import mem0.memory.setup as setup_module
setup_module.setup_config()
setup_module.mark_oss_used()
oss_id = json.loads((tmp_mem0_dir / "config.json").read_text())["user_id"]
from mem0.client import main as client_main
importlib.reload(client_main)
with patch.object(client_main, "client_telemetry") as telemetry:
telemetry.capture_identify.return_value = True
client_main._maybe_alias_anon_to_email("user@example.com")
telemetry.capture_identify.assert_called_once_with(oss_id, "user@example.com")
client_main._maybe_alias_anon_to_email("user@example.com")
assert telemetry.capture_identify.call_count == 1
config = json.loads((tmp_mem0_dir / "config.json").read_text())
assert len(config["telemetry"]["aliased_pairs"]) == 1
assert config["telemetry"]["oss_used_at"]
# ─── Platform-only init must not mint or stitch ──────────────────────────────
class TestPlatformOnlyDoesNotMint:
"""Regression for the bug where importing the client minted an anon id at
import time and then stitched it on MemoryClient init, merging an
event-less anon person into every platform user."""
def test_no_config_created_and_no_identify(self, tmp_mem0_dir):
from mem0.client import main as client_main
# Re-execute the client module: it must not create config.json.
importlib.reload(client_main)
assert not (tmp_mem0_dir / "config.json").exists()
with patch.object(client_main, "client_telemetry") as telemetry:
telemetry.capture_identify.return_value = True
client_main._maybe_alias_anon_to_email("user@example.com")
telemetry.capture_identify.assert_not_called()
assert not (tmp_mem0_dir / "config.json").exists()
# ─── Lazy minting from OSS Memory init ───────────────────────────────────────
class TestLazyOssMinting:
"""OSS Memory/AsyncMemory init mints the anon id and records genuine OSS
usage; nothing happens at import time or with telemetry disabled."""
def test_memory_init_mints_and_marks(self, tmp_path):
import mem0.memory.main as memory_main
with (
patch.object(memory_main, "EmbedderFactory"),
patch.object(memory_main, "VectorStoreFactory"),
patch.object(memory_main, "LlmFactory"),
patch.object(memory_main, "SQLiteManager"),
patch.object(memory_main, "capture_event"),
patch.object(memory_main, "mem0_dir", str(tmp_path)),
patch.object(memory_main, "setup_config") as setup,
patch.object(memory_main, "mark_oss_used") as mark,
patch.object(memory_main, "MEM0_TELEMETRY", True),
):
memory_main.Memory()
setup.assert_called_once()
mark.assert_called_once()
def test_memory_init_skips_minting_when_telemetry_disabled(self, tmp_path):
import mem0.memory.main as memory_main
with (
patch.object(memory_main, "EmbedderFactory"),
patch.object(memory_main, "VectorStoreFactory"),
patch.object(memory_main, "LlmFactory"),
patch.object(memory_main, "SQLiteManager"),
patch.object(memory_main, "capture_event"),
patch.object(memory_main, "mem0_dir", str(tmp_path)),
patch.object(memory_main, "setup_config") as setup,
patch.object(memory_main, "mark_oss_used") as mark,
patch.object(memory_main, "MEM0_TELEMETRY", False),
):
memory_main.Memory()
setup.assert_not_called()
mark.assert_not_called()
def test_async_memory_init_mints_and_marks(self, tmp_path):
import mem0.memory.main as memory_main
with (
patch.object(memory_main, "EmbedderFactory"),
patch.object(memory_main, "VectorStoreFactory"),
patch.object(memory_main, "LlmFactory"),
patch.object(memory_main, "SQLiteManager"),
patch.object(memory_main, "capture_event"),
patch.object(memory_main, "mem0_dir", str(tmp_path)),
patch.object(memory_main, "setup_config") as setup,
patch.object(memory_main, "mark_oss_used") as mark,
patch.object(memory_main, "MEM0_TELEMETRY", True),
):
memory_main.AsyncMemory()
setup.assert_called_once()
mark.assert_called_once()