Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
60ac9ce149 | ||
|
|
9315e3036f | ||
|
|
b3ede5b7c0 | ||
|
|
3553fc79dd | ||
|
|
f322cf82b9 | ||
|
|
73c975ba68 | ||
|
|
931d579ba5 |
@@ -0,0 +1,4 @@
|
||||
[submodule "evaluation"]
|
||||
path = evaluation
|
||||
url = https://github.com/mem0ai/memory-benchmarks
|
||||
branch = main
|
||||
@@ -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, documentation, and evaluation tooling.
|
||||
This is a **polyglot monorepo** containing Python and TypeScript packages, CLIs, servers, plugins, and documentation.
|
||||
|
||||
### Key Directories
|
||||
|
||||
@@ -32,9 +32,8 @@ 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/` | Benchmarking framework — LOCOMO evals, experiment runner, score generation |
|
||||
| `examples/` | Sample projects — demo apps, Chrome extension, multi-agent patterns |
|
||||
| `cookbooks/` | Jupyter notebooks — customer support chatbot, AutoGen integration |
|
||||
| `evaluation/` | Submodule → [`mem0ai/memory-benchmarks`](https://github.com/mem0ai/memory-benchmarks) — benchmarking (LOCOMO, LongMemEval, BEAM) lives in that repo |
|
||||
| `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) |
|
||||
|
||||
@@ -247,18 +246,19 @@ make docs # or: cd docs && mintlify dev
|
||||
- **API spec:** `docs/openapi.json`
|
||||
- **Structure:** `api-reference/`, `open-source/`, `platform/`, `integrations/`, `cookbooks/`, `core-concepts/`
|
||||
|
||||
### Evaluation (`evaluation/`)
|
||||
### 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:
|
||||
|
||||
```bash
|
||||
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
|
||||
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
|
||||
```
|
||||
|
||||
## Core APIs
|
||||
|
||||
@@ -1026,7 +1026,8 @@ def get_user_preferences(user_id: str):
|
||||
### AutoGen Integration
|
||||
|
||||
```python
|
||||
from cookbooks.helper.mem0_teachability import Mem0Teachability
|
||||
# Mem0Teachability lives in examples/notebooks/helper/ — see examples/notebooks/mem0-autogen.ipynb
|
||||
from helper.mem0_teachability import Mem0Teachability
|
||||
from mem0 import Memory
|
||||
|
||||
# Add memory capability to AutoGen agents
|
||||
|
||||
@@ -7,6 +7,22 @@ mode: "wide"
|
||||
<Tabs>
|
||||
<Tab title="Python">
|
||||
|
||||
<Update label="2026-06-13" description="v2.0.6">
|
||||
|
||||
**New Features:**
|
||||
- **Memory:** Add a contextual OSS-to-Platform notices system that surfaces occasional, situation-aware messages (first run, scale/performance thresholds, slow queries, and when temporal/decay features are relevant) pointing to the corresponding Mem0 Platform capabilities; disable via `MEM0_TELEMETRY=false` ([#5494](https://github.com/mem0ai/mem0/pull/5494))
|
||||
|
||||
**Bug Fixes:**
|
||||
- **Memory:** Prevent a crash in `parse_vision_messages` when vision support is disabled ([#5487](https://github.com/mem0ai/mem0/pull/5487))
|
||||
- **Vector Stores:** Expose the `https` option on the Qdrant vector store configuration so TLS endpoints can be targeted explicitly ([#5380](https://github.com/mem0ai/mem0/pull/5380))
|
||||
- **Vector Stores:** Use valid S3 Vectors entity index names, fixing index operations that failed on invalid names ([#5416](https://github.com/mem0ai/mem0/pull/5416))
|
||||
- **Vector Stores:** Fix `search()` crashing with a `TypeError` in the LangChain vector store when a result score is `None` ([#5072](https://github.com/mem0ai/mem0/pull/5072))
|
||||
- **Vector Stores:** Use `is not None` instead of a truthiness check for vector/payload in the PGVector `update()` path, so valid empty/zero values are no longer skipped ([#5488](https://github.com/mem0ai/mem0/pull/5488))
|
||||
- **Vector Stores:** Index the Valkey `memory` field as `TEXT` rather than `TAG` so full-text search behaves correctly ([#5443](https://github.com/mem0ai/mem0/pull/5443))
|
||||
- **Vector Stores:** Implement `$not` filter support in the ChromaDB vector store ([#5485](https://github.com/mem0ai/mem0/pull/5485))
|
||||
|
||||
</Update>
|
||||
|
||||
<Update label="2026-06-10" description="v2.0.5">
|
||||
|
||||
**New Features:**
|
||||
@@ -960,6 +976,17 @@ See the [OSS v1 to v2 migration guide](https://docs.mem0.ai/migration/oss-v1-to-
|
||||
|
||||
<Tab title="TypeScript">
|
||||
|
||||
<Update label="2026-06-13" description="v3.0.8">
|
||||
|
||||
**New Features:**
|
||||
- **Memory:** Add a contextual OSS-to-Platform notices system that surfaces occasional, situation-aware messages (first run, scale/performance thresholds, slow queries, and when temporal/decay features are relevant) pointing to the corresponding Mem0 Platform capabilities; disable via `MEM0_TELEMETRY=false` ([#5494](https://github.com/mem0ai/mem0/pull/5494))
|
||||
|
||||
**Security:**
|
||||
- **Dependencies:** Upgrade `@langchain/community` to `^1.1.18` to remediate CVE-2026-27795 and CVE-2026-26019 ([#5510](https://github.com/mem0ai/mem0/pull/5510))
|
||||
- **Dependencies:** Resolve all open MEDIUM Dependabot alerts via pnpm overrides ([#5489](https://github.com/mem0ai/mem0/pull/5489))
|
||||
|
||||
</Update>
|
||||
|
||||
<Update label="2026-06-10" description="v3.0.7">
|
||||
|
||||
**New Features:**
|
||||
|
||||
Submodule
+1
Submodule evaluation added at 4b61c5d31b
@@ -1,31 +0,0 @@
|
||||
|
||||
# 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/
|
||||
@@ -1,198 +0,0 @@
|
||||
# Mem0: Building Production‑Ready AI Agents with Scalable Long‑Term Memory
|
||||
|
||||
[](https://arxiv.org/abs/2504.19413)
|
||||
[](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)
|
||||
|
||||
@@ -1,81 +0,0 @@
|
||||
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()
|
||||
@@ -1,34 +0,0 @@
|
||||
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)
|
||||
@@ -1,130 +0,0 @@
|
||||
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()
|
||||
@@ -1,211 +0,0 @@
|
||||
"""
|
||||
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
|
||||
@@ -1,147 +0,0 @@
|
||||
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:
|
||||
"""
|
||||
@@ -1,75 +0,0 @@
|
||||
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()
|
||||
@@ -1,185 +0,0 @@
|
||||
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)
|
||||
@@ -1,141 +0,0 @@
|
||||
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()
|
||||
@@ -1,215 +0,0 @@
|
||||
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
|
||||
@@ -1,131 +0,0 @@
|
||||
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)
|
||||
@@ -1,183 +0,0 @@
|
||||
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)
|
||||
@@ -1,3 +0,0 @@
|
||||
TECHNIQUES = ["mem0", "rag", "langmem", "zep", "openai"]
|
||||
|
||||
METHODS = ["add", "search"]
|
||||
@@ -1,76 +0,0 @@
|
||||
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)
|
||||
@@ -1,140 +0,0 @@
|
||||
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")
|
||||
@@ -555,7 +555,7 @@
|
||||
"# - Enables creation of AI agents with long-term memory and learning abilities.\n",
|
||||
"# - Improves consistency and reduces repetition in user-agent interactions.\n",
|
||||
"\n",
|
||||
"from cookbooks.helper.mem0_teachability import Mem0Teachability\n",
|
||||
"from helper.mem0_teachability import Mem0Teachability\n",
|
||||
"\n",
|
||||
"teachability = Mem0Teachability(\n",
|
||||
" verbosity=2, # for visibility of what's happening\n",
|
||||
@@ -2,19 +2,14 @@
|
||||
|
||||
All notable changes to the `@mem0/opencode-plugin` will be documented in this file.
|
||||
|
||||
## 0.1.4 — Anonymous usage telemetry
|
||||
|
||||
### Added
|
||||
|
||||
- **PostHog telemetry (`telemetry.ts`):** Anonymous, fire-and-forget usage events. Opt out with `MEM0_TELEMETRY=false`. Only fires when an API key is present; never sends memory content, prompts, or the API key — only an anonymized `sha256(apiKey)[:32]` identity plus event type, platform, and plugin version. Emits the same schema as the Mem0 editor plugin (`plugin.*` events, `source: "plugin"`, `platform: "opencode"`) so OpenCode appears as a `platform` in the shared plugin dashboard. Events: `plugin.session_start` (with memory count) and `plugin.tool_use` (`add` / `search` / `update` / `delete`).
|
||||
|
||||
## 0.1.3 — File-context injection, session summaries & activity timeline
|
||||
## 0.1.3 — File-context injection, session summaries & activity timeline, anonymous telemetry
|
||||
|
||||
### Added
|
||||
|
||||
- **File-context injection (`tool.execute.before` / Read):** Before the agent reads a file, the plugin searches mem0 for memories referencing that file path and injects prior work as system context. Gates on file size (>= 1,500 bytes). Gives the agent "I've worked on this file before" awareness automatically.
|
||||
- **Stop hook session summary (`experimental.session.compacting`):** Enhanced session compaction to store a structured `session_summary` memory with `infer=True`, letting the mem0 backend AI extract key facts (request, decisions, learnings, next steps). Previously only stored a raw stats string.
|
||||
- **SessionStart activity timeline:** The initial memory loading now formats recent memories with type icons (⚖️ decision, 🔴 bug_fix, 🔵 task_learning, etc.) and relative age indicators (2h ago, 1d ago) instead of bare text. Provides a visual "Recent Activity" timeline on first message.
|
||||
- **PostHog telemetry (`telemetry.ts`):** Anonymous, fire-and-forget usage events. Opt out with `MEM0_TELEMETRY=false`. Only fires when an API key is present; never sends memory content, prompts, or the API key — only an anonymized `sha256(apiKey)[:32]` identity plus event type, platform, and plugin version. Emits the same schema as the Mem0 editor plugin (`plugin.*` events, `source: "plugin"`, `platform: "opencode"`) so OpenCode appears as a `platform` in the shared plugin dashboard. Events: `plugin.session_start` (with memory count) and `plugin.tool_use` (`add` / `search` / `update` / `delete`).
|
||||
|
||||
### Changed
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
"name": "@mem0/opencode-plugin",
|
||||
"dependencies": {
|
||||
"@opencode-ai/plugin": "^1.0.162",
|
||||
"mem0ai": "^3.0.5",
|
||||
"mem0ai": "^3.0.7",
|
||||
},
|
||||
"devDependencies": {
|
||||
"bun-types": ">=1.3.14",
|
||||
@@ -462,7 +462,7 @@
|
||||
|
||||
"md5": ["md5@2.3.0", "", { "dependencies": { "charenc": "0.0.2", "crypt": "0.0.2", "is-buffer": "~1.1.6" } }, "sha512-T1GITYmFaKuO91vxyoQMFETst+O71VUPEU3ze5GNzDm0OWdP8v1ziTaAEPUr/3kLsY3Sftgz242A1SetQiDL7g=="],
|
||||
|
||||
"mem0ai": ["mem0ai@3.0.5", "", { "dependencies": { "axios": "^1.15.2", "openai": "^4.93.0", "uuid": "9.0.1", "zod": "^3.24.1" }, "peerDependencies": { "@anthropic-ai/sdk": "^0.40.1", "@azure/identity": "^4.0.0", "@azure/search-documents": "^12.0.0", "@cloudflare/workers-types": "^4.20250504.0", "@google/genai": "^1.2.0", "@langchain/core": "^1.1.47", "@mistralai/mistralai": "^1.5.2", "@qdrant/js-client-rest": "1.13.0", "@supabase/supabase-js": "^2.49.1", "@types/jest": "29.5.14", "@types/pg": "8.11.0", "better-sqlite3": "^12.6.2", "cloudflare": "^4.2.0", "compromise": "^14.0.0", "groq-sdk": "0.3.0", "natural": "^8.0.1", "ollama": "^0.5.14", "pg": "8.11.3", "redis": "^4.6.13" } }, "sha512-W/R59d5fMpUGHhPEnyoo36GSz5NFJbAs+vS4BxoIvE+t19mIJfoz/2FJSKIw80mT8AkeNYpDfzc/DYRu2IzYIw=="],
|
||||
"mem0ai": ["mem0ai@3.0.7", "", { "dependencies": { "axios": "^1.16.0", "openai": "^4.93.0", "uuid": "9.0.1", "zod": "^3.24.1" }, "peerDependencies": { "@anthropic-ai/sdk": "^0.40.1", "@azure/identity": "^4.0.0", "@azure/search-documents": "^12.0.0", "@cloudflare/workers-types": "^4.20250504.0", "@google/genai": "^1.40.0", "@langchain/core": "^1.1.47", "@mistralai/mistralai": "^1.5.2", "@qdrant/js-client-rest": "^1.18.0", "@supabase/supabase-js": "^2.49.1", "@types/jest": "29.5.14", "@types/pg": "8.11.0", "better-sqlite3": "^12.6.2", "cloudflare": "^4.2.0", "compromise": "^14.0.0", "groq-sdk": "0.3.0", "natural": "^8.0.1", "ollama": "^0.5.14", "pg": "8.11.3", "redis": "^4.6.13" } }, "sha512-CUHzX7DyeKTHcI3aDsSqY9LXTD7GcFxf988796TuOa4yJgGuF2Xd2NROcBhVFRo3r9y8fVmbo3c5TF9jv1KlKw=="],
|
||||
|
||||
"memjs": ["memjs@1.3.2", "", {}, "sha512-qUEg2g8vxPe+zPn09KidjIStHPtoBO8Cttm8bgJFWWabbsjQ9Av9Ky+6UcvKx6ue0LLb/LEhtcyQpRyKfzeXcg=="],
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@mem0/opencode-plugin",
|
||||
"version": "0.1.4",
|
||||
"version": "0.1.3",
|
||||
"type": "module",
|
||||
"description": "Mem0 persistent memory plugin for OpenCode — add, search, and manage memories across sessions",
|
||||
"main": "dist/index.js",
|
||||
@@ -59,7 +59,7 @@
|
||||
},
|
||||
"dependencies": {
|
||||
"@opencode-ai/plugin": "^1.0.162",
|
||||
"mem0ai": "^3.0.5"
|
||||
"mem0ai": "^3.0.7"
|
||||
},
|
||||
"devDependencies": {
|
||||
"bun-types": ">=1.3.14",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "mem0ai",
|
||||
"version": "3.0.7",
|
||||
"version": "3.0.9",
|
||||
"description": "The Memory Layer For Your AI Apps",
|
||||
"main": "./dist/index.js",
|
||||
"module": "./dist/index.mjs",
|
||||
|
||||
@@ -36,10 +36,29 @@ import {
|
||||
SearchMemoryOptions,
|
||||
DeleteAllMemoryOptions,
|
||||
GetAllMemoryOptions,
|
||||
UpdateProjectOptions,
|
||||
} from "./memory.types";
|
||||
import { parse_vision_messages } from "../utils/memory";
|
||||
import { HistoryManager } from "../storage/base";
|
||||
import { captureClientEvent } from "../utils/telemetry";
|
||||
import {
|
||||
detectScaleThresholdFromAddResult,
|
||||
detectScaleThresholdFromTopK,
|
||||
detectPerformanceSlowQuery,
|
||||
detectTemporalUsageFromMetadata,
|
||||
detectTemporalUsageFromSearch,
|
||||
displayDecayUsageNotice,
|
||||
displayFirstRunNotice,
|
||||
displayPerformanceSlowQueryNotice,
|
||||
displayScaleThresholdNotice,
|
||||
displayTemporalUsageNotice,
|
||||
getDecayFeatureErrorMessage,
|
||||
getDecayUsageDeleteCountAfterSuccess,
|
||||
getTemporalFeatureErrorMessage,
|
||||
isDecayUsageDeleteEligible,
|
||||
PerformanceSlowQueryTrigger,
|
||||
ScaleThresholdTrigger,
|
||||
} from "../utils/notices";
|
||||
import { lemmatizeForBm25 } from "../utils/lemmatization";
|
||||
import {
|
||||
extractEntities,
|
||||
@@ -494,6 +513,73 @@ export class Memory {
|
||||
}
|
||||
}
|
||||
|
||||
private async _displayFirstRunNotice(triggerFunction: string) {
|
||||
try {
|
||||
await this._getTelemetryId();
|
||||
await displayFirstRunNotice(this, triggerFunction);
|
||||
} catch {}
|
||||
}
|
||||
|
||||
private async _displayDecayUsageNotice(trigger: {
|
||||
triggerFunction: "delete" | "delete_all";
|
||||
triggerSource: "delete_count" | "delete_all";
|
||||
triggerReason: "repeated_deletes" | "bulk_delete";
|
||||
deleteCount?: number;
|
||||
deletedCount?: number;
|
||||
}) {
|
||||
try {
|
||||
await this._getTelemetryId();
|
||||
await displayDecayUsageNotice(this, trigger);
|
||||
} catch {}
|
||||
}
|
||||
|
||||
private async _displayTemporalUsageNotice(trigger: {
|
||||
triggerFunction: "add" | "search";
|
||||
triggerSource: "metadata" | "query" | "filter";
|
||||
triggerReason:
|
||||
| "date_like_metadata"
|
||||
| "relative_phrase"
|
||||
| "date_like_query"
|
||||
| "date_range_filter";
|
||||
}) {
|
||||
try {
|
||||
await this._getTelemetryId();
|
||||
await displayTemporalUsageNotice(this, trigger);
|
||||
} catch {}
|
||||
}
|
||||
|
||||
private async _displayScaleThresholdNotice(trigger: ScaleThresholdTrigger) {
|
||||
try {
|
||||
await this._getTelemetryId();
|
||||
await displayScaleThresholdNotice(this, trigger);
|
||||
} catch {}
|
||||
}
|
||||
|
||||
private async _displayPerformanceSlowQueryNotice(
|
||||
trigger: PerformanceSlowQueryTrigger,
|
||||
) {
|
||||
try {
|
||||
await this._getTelemetryId();
|
||||
await displayPerformanceSlowQueryNotice(this, trigger);
|
||||
} catch {}
|
||||
}
|
||||
|
||||
private async _getNoticeTelemetryId() {
|
||||
try {
|
||||
if (
|
||||
!this.telemetryId ||
|
||||
this.telemetryId === "anonymous" ||
|
||||
this.telemetryId === "anonymous-supabase"
|
||||
) {
|
||||
this.telemetryId = (await getOrCreateMem0UserId()) || "anonymous";
|
||||
}
|
||||
return this.telemetryId;
|
||||
} catch {
|
||||
this.telemetryId = "anonymous";
|
||||
return this.telemetryId;
|
||||
}
|
||||
}
|
||||
|
||||
static fromConfig(configDict: Record<string, any>): Memory {
|
||||
try {
|
||||
const config = MemoryConfigSchema.parse(configDict);
|
||||
@@ -504,10 +590,29 @@ export class Memory {
|
||||
}
|
||||
}
|
||||
|
||||
async updateProject(options: UpdateProjectOptions = {}): Promise<never> {
|
||||
if (options?.decay === true) {
|
||||
await this._getNoticeTelemetryId();
|
||||
throw new Error(await getDecayFeatureErrorMessage(this));
|
||||
}
|
||||
|
||||
throw new Error("Project updates are not supported by the OSS Memory SDK.");
|
||||
}
|
||||
|
||||
async add(
|
||||
messages: string | Message[],
|
||||
config: AddMemoryOptions,
|
||||
): Promise<SearchResult> {
|
||||
if (config?.timestamp !== undefined) {
|
||||
await this._getNoticeTelemetryId();
|
||||
throw new Error(
|
||||
await getTemporalFeatureErrorMessage(this, {
|
||||
triggerFunction: "add",
|
||||
triggerParameter: "timestamp",
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// Validate messages input
|
||||
if (messages === undefined || messages === null) {
|
||||
throw new Error(
|
||||
@@ -515,6 +620,10 @@ export class Memory {
|
||||
);
|
||||
}
|
||||
|
||||
const temporalUsageNotice = detectTemporalUsageFromMetadata(
|
||||
config?.metadata,
|
||||
);
|
||||
|
||||
await this._ensureInitialized();
|
||||
await this._captureEvent("add", {
|
||||
message_count: Array.isArray(messages) ? messages.length : 1,
|
||||
@@ -554,6 +663,27 @@ export class Memory {
|
||||
infer,
|
||||
);
|
||||
|
||||
if (temporalUsageNotice) {
|
||||
await this._displayTemporalUsageNotice({
|
||||
triggerFunction: "add",
|
||||
triggerSource: temporalUsageNotice.triggerSource,
|
||||
triggerReason: temporalUsageNotice.triggerReason,
|
||||
});
|
||||
} else {
|
||||
const scaleThresholdNotice = await detectScaleThresholdFromAddResult(
|
||||
this,
|
||||
vectorStoreResult,
|
||||
);
|
||||
if (scaleThresholdNotice) {
|
||||
await this._displayScaleThresholdNotice({
|
||||
triggerFunction: "add",
|
||||
...scaleThresholdNotice,
|
||||
});
|
||||
} else {
|
||||
await this._displayFirstRunNotice("add");
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
results: vectorStoreResult,
|
||||
};
|
||||
@@ -996,7 +1126,10 @@ export class Memory {
|
||||
async get(memoryId: string): Promise<MemoryItem | null> {
|
||||
await this._ensureInitialized();
|
||||
const memory = await this.vectorStore.get(memoryId);
|
||||
if (!memory) return null;
|
||||
if (!memory) {
|
||||
await this._displayFirstRunNotice("get");
|
||||
return null;
|
||||
}
|
||||
|
||||
const filters = {
|
||||
...(memory.payload.user_id && { user_id: memory.payload.user_id }),
|
||||
@@ -1031,13 +1164,30 @@ export class Memory {
|
||||
}
|
||||
}
|
||||
|
||||
return { ...memoryItem, ...filters };
|
||||
const result = { ...memoryItem, ...filters };
|
||||
await this._displayFirstRunNotice("get");
|
||||
return result;
|
||||
}
|
||||
|
||||
async search(
|
||||
query: string,
|
||||
config: SearchMemoryOptions,
|
||||
): Promise<SearchResult> {
|
||||
if (config?.referenceDate !== undefined) {
|
||||
await this._getNoticeTelemetryId();
|
||||
throw new Error(
|
||||
await getTemporalFeatureErrorMessage(this, {
|
||||
triggerFunction: "search",
|
||||
triggerParameter: "referenceDate",
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
const temporalUsageNotice = detectTemporalUsageFromSearch(
|
||||
query,
|
||||
config?.filters,
|
||||
);
|
||||
|
||||
// Reject top-level entity params - must use filters instead
|
||||
rejectTopLevelEntityParams(config as Record<string, any>, "search");
|
||||
|
||||
@@ -1105,6 +1255,8 @@ export class Memory {
|
||||
);
|
||||
}
|
||||
|
||||
const searchStartMs = Date.now();
|
||||
|
||||
// Step 1: Preprocess query
|
||||
const queryLemmatized = lemmatizeForBm25(query);
|
||||
const queryEntities = extractEntities(query);
|
||||
@@ -1280,9 +1432,41 @@ export class Memory {
|
||||
};
|
||||
});
|
||||
|
||||
return {
|
||||
const result = {
|
||||
results,
|
||||
};
|
||||
const searchElapsedMs = Date.now() - searchStartMs;
|
||||
if (temporalUsageNotice) {
|
||||
await this._displayTemporalUsageNotice({
|
||||
triggerFunction: "search",
|
||||
triggerSource: temporalUsageNotice.triggerSource,
|
||||
triggerReason: temporalUsageNotice.triggerReason,
|
||||
});
|
||||
} else {
|
||||
const scaleThresholdNotice = detectScaleThresholdFromTopK(topK);
|
||||
if (scaleThresholdNotice) {
|
||||
await this._displayScaleThresholdNotice({
|
||||
triggerFunction: "search",
|
||||
...scaleThresholdNotice,
|
||||
});
|
||||
} else {
|
||||
const performanceSlowQueryNotice = detectPerformanceSlowQuery(
|
||||
searchElapsedMs,
|
||||
topK,
|
||||
results.length,
|
||||
);
|
||||
if (performanceSlowQueryNotice) {
|
||||
await this._displayPerformanceSlowQueryNotice({
|
||||
triggerFunction: "search",
|
||||
triggerReason: "slow_query",
|
||||
...performanceSlowQueryNotice,
|
||||
});
|
||||
} else {
|
||||
await this._displayFirstRunNotice("search");
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
async update(memoryId: string, data: string): Promise<{ message: string }> {
|
||||
@@ -1290,14 +1474,28 @@ export class Memory {
|
||||
await this._captureEvent("update", { memory_id: memoryId });
|
||||
const embedding = await this.embedder.embed(data);
|
||||
await this.updateMemory(memoryId, data, { [data]: embedding });
|
||||
return { message: "Memory updated successfully!" };
|
||||
const result = { message: "Memory updated successfully!" };
|
||||
await this._displayFirstRunNotice("update");
|
||||
return result;
|
||||
}
|
||||
|
||||
async delete(memoryId: string): Promise<{ message: string }> {
|
||||
await this._ensureInitialized();
|
||||
await this._captureEvent("delete", { memory_id: memoryId });
|
||||
await this.deleteMemory(memoryId);
|
||||
return { message: "Memory deleted successfully!" };
|
||||
const result = { message: "Memory deleted successfully!" };
|
||||
const deleteCount = getDecayUsageDeleteCountAfterSuccess();
|
||||
if (isDecayUsageDeleteEligible(deleteCount)) {
|
||||
await this._displayDecayUsageNotice({
|
||||
triggerFunction: "delete",
|
||||
triggerSource: "delete_count",
|
||||
triggerReason: "repeated_deletes",
|
||||
deleteCount,
|
||||
});
|
||||
} else {
|
||||
await this._displayFirstRunNotice("delete");
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
async deleteAll(
|
||||
@@ -1328,12 +1526,25 @@ export class Memory {
|
||||
await this.deleteMemory(memory.id);
|
||||
}
|
||||
|
||||
return { message: "Memories deleted successfully!" };
|
||||
const result = { message: "Memories deleted successfully!" };
|
||||
if (memories.length > 0) {
|
||||
await this._displayDecayUsageNotice({
|
||||
triggerFunction: "delete_all",
|
||||
triggerSource: "delete_all",
|
||||
triggerReason: "bulk_delete",
|
||||
deletedCount: memories.length,
|
||||
});
|
||||
} else {
|
||||
await this._displayFirstRunNotice("delete_all");
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
async history(memoryId: string): Promise<any[]> {
|
||||
await this._ensureInitialized();
|
||||
return this.db.getHistory(memoryId);
|
||||
const result = await this.db.getHistory(memoryId);
|
||||
await this._displayFirstRunNotice("history");
|
||||
return result;
|
||||
}
|
||||
|
||||
async reset(): Promise<void> {
|
||||
@@ -1385,6 +1596,7 @@ export class Memory {
|
||||
console.error(this._initError);
|
||||
});
|
||||
await this._initPromise;
|
||||
await this._displayFirstRunNotice("reset");
|
||||
}
|
||||
|
||||
async getAll(config: GetAllMemoryOptions): Promise<SearchResult> {
|
||||
@@ -1452,7 +1664,17 @@ export class Memory {
|
||||
...(mem.payload.run_id && { run_id: mem.payload.run_id }),
|
||||
}));
|
||||
|
||||
return { results };
|
||||
const result = { results };
|
||||
const scaleThresholdNotice = detectScaleThresholdFromTopK(topK);
|
||||
if (scaleThresholdNotice) {
|
||||
await this._displayScaleThresholdNotice({
|
||||
triggerFunction: "get_all",
|
||||
...scaleThresholdNotice,
|
||||
});
|
||||
} else {
|
||||
await this._displayFirstRunNotice("get_all");
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private async createMemory(
|
||||
|
||||
@@ -11,6 +11,7 @@ export interface AddMemoryOptions extends Entity {
|
||||
metadata?: Record<string, any>;
|
||||
filters?: SearchFilters;
|
||||
infer?: boolean;
|
||||
timestamp?: number | string | Date | null;
|
||||
}
|
||||
|
||||
export interface SearchMemoryOptions {
|
||||
@@ -18,6 +19,7 @@ export interface SearchMemoryOptions {
|
||||
filters?: SearchFilters;
|
||||
threshold?: number;
|
||||
explain?: boolean;
|
||||
referenceDate?: number | string | Date | null;
|
||||
}
|
||||
|
||||
export interface GetAllMemoryOptions {
|
||||
@@ -26,3 +28,8 @@ export interface GetAllMemoryOptions {
|
||||
}
|
||||
|
||||
export interface DeleteAllMemoryOptions extends Entity {}
|
||||
|
||||
export interface UpdateProjectOptions {
|
||||
decay?: boolean;
|
||||
[key: string]: any;
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -13,10 +13,12 @@ let version =
|
||||
// Safely check for process.env in different environments
|
||||
let MEM0_TELEMETRY = true;
|
||||
try {
|
||||
MEM0_TELEMETRY = process?.env?.MEM0_TELEMETRY === "false" ? false : true;
|
||||
MEM0_TELEMETRY =
|
||||
process?.env?.MEM0_TELEMETRY?.toLowerCase() === "false" ? false : true;
|
||||
} catch (error) {}
|
||||
const POSTHOG_API_KEY = "phc_hgJkUVJFYtmaJqrvf6CYN67TIQ8yhXAkWzUn9AMU4yX";
|
||||
const POSTHOG_HOST = "https://us.i.posthog.com/i/v0/e/";
|
||||
const NOTICE_EVENT_NAME = "mem0.notice_displayed";
|
||||
|
||||
// Default sampling rate for hot-path OSS events. Lifecycle events always fire at 100%.
|
||||
// Override via MEM0_TELEMETRY_SAMPLE_RATE env var. Mirrors mem0/memory/telemetry.py.
|
||||
@@ -35,7 +37,11 @@ const MEM0_TELEMETRY_SAMPLE_RATE: number = ((): number => {
|
||||
})();
|
||||
|
||||
// Events that bypass sampling. Keep in sync with _captureEvent call sites in memory/index.ts.
|
||||
const LIFECYCLE_EVENTS: ReadonlySet<string> = new Set(["init", "reset"]);
|
||||
const ALWAYS_SEND_EVENTS: ReadonlySet<string> = new Set([
|
||||
"init",
|
||||
"reset",
|
||||
"notice_displayed",
|
||||
]);
|
||||
|
||||
class UnifiedTelemetry implements TelemetryClient {
|
||||
private apiKey: string;
|
||||
@@ -91,6 +97,10 @@ class UnifiedTelemetry implements TelemetryClient {
|
||||
|
||||
const telemetry = new UnifiedTelemetry(POSTHOG_API_KEY, POSTHOG_HOST);
|
||||
|
||||
function isTelemetryEnabled(): boolean {
|
||||
return MEM0_TELEMETRY;
|
||||
}
|
||||
|
||||
async function captureClientEvent(
|
||||
eventName: string,
|
||||
instance: TelemetryInstance,
|
||||
@@ -102,8 +112,8 @@ async function captureClientEvent(
|
||||
}
|
||||
|
||||
// >= so that rate=0 drops everything and rate=1 keeps everything (Math.random() ∈ [0, 1)).
|
||||
const isLifecycle = LIFECYCLE_EVENTS.has(eventName);
|
||||
if (!isLifecycle && Math.random() >= MEM0_TELEMETRY_SAMPLE_RATE) {
|
||||
const alwaysSend = ALWAYS_SEND_EVENTS.has(eventName);
|
||||
if (!alwaysSend && Math.random() >= MEM0_TELEMETRY_SAMPLE_RATE) {
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -116,7 +126,7 @@ async function captureClientEvent(
|
||||
client_source: "nodejs",
|
||||
...additionalData,
|
||||
// sample_rate set AFTER the spread so callers can never override it
|
||||
sample_rate: isLifecycle ? 1.0 : MEM0_TELEMETRY_SAMPLE_RATE,
|
||||
sample_rate: alwaysSend ? 1.0 : MEM0_TELEMETRY_SAMPLE_RATE,
|
||||
};
|
||||
|
||||
await telemetry.captureEvent(
|
||||
@@ -126,4 +136,36 @@ async function captureClientEvent(
|
||||
);
|
||||
}
|
||||
|
||||
export { telemetry, captureClientEvent };
|
||||
async function captureNoticeEvent(
|
||||
instance: TelemetryInstance,
|
||||
properties: Record<string, any> = {},
|
||||
) {
|
||||
if (!instance.telemetryId) return;
|
||||
|
||||
const eventData: TelemetryEventData = {
|
||||
function: `${instance.constructor.name}`,
|
||||
method: "notice_displayed",
|
||||
api_host: instance.host,
|
||||
timestamp: new Date().toISOString(),
|
||||
client_version: version,
|
||||
client_source: "nodejs",
|
||||
...properties,
|
||||
sample_rate: 1.0,
|
||||
};
|
||||
|
||||
await telemetry.captureEvent(
|
||||
instance.telemetryId,
|
||||
NOTICE_EVENT_NAME,
|
||||
eventData,
|
||||
);
|
||||
}
|
||||
|
||||
export {
|
||||
POSTHOG_API_KEY,
|
||||
POSTHOG_HOST,
|
||||
NOTICE_EVENT_NAME,
|
||||
telemetry,
|
||||
captureClientEvent,
|
||||
captureNoticeEvent,
|
||||
isTelemetryEnabled,
|
||||
};
|
||||
|
||||
@@ -418,6 +418,7 @@ describe("Memory – LM Studio end-to-end flow", () => {
|
||||
}));
|
||||
jest.doMock("../src/utils/telemetry", () => ({
|
||||
captureClientEvent: jest.fn().mockResolvedValue(undefined),
|
||||
isTelemetryEnabled: jest.fn(() => false),
|
||||
}));
|
||||
|
||||
MemoryClass = require("../src/memory").Memory;
|
||||
|
||||
@@ -290,6 +290,7 @@ describe("Memory – auto-initialization", () => {
|
||||
|
||||
jest.doMock("../src/utils/telemetry", () => ({
|
||||
captureClientEvent: jest.fn().mockResolvedValue(undefined),
|
||||
isTelemetryEnabled: jest.fn(() => false),
|
||||
}));
|
||||
|
||||
MemoryClass = require("../src/memory").Memory;
|
||||
|
||||
@@ -0,0 +1,387 @@
|
||||
/// <reference types="jest" />
|
||||
import * as fs from "fs";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
|
||||
jest.setTimeout(15000);
|
||||
|
||||
jest.mock("../src/embeddings/google", () => ({
|
||||
GoogleEmbedder: jest.fn(),
|
||||
}));
|
||||
jest.mock("../src/llms/google", () => ({
|
||||
GoogleLLM: jest.fn(),
|
||||
}));
|
||||
|
||||
const mockEmbedding = new Array(1536).fill(0.1);
|
||||
jest.mock("../src/embeddings/openai", () => ({
|
||||
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
|
||||
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||
embedBatch: jest
|
||||
.fn()
|
||||
.mockImplementation((texts: string[]) =>
|
||||
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||
),
|
||||
embeddingDims: 1536,
|
||||
})),
|
||||
}));
|
||||
|
||||
jest.mock("../src/llms/openai", () => ({
|
||||
OpenAILLM: jest.fn().mockImplementation(() => ({
|
||||
generateResponse: jest.fn(),
|
||||
})),
|
||||
}));
|
||||
|
||||
const DECAY_COPY =
|
||||
'Memory decay requires Mem0 Platform. Get a free API key at https://app.mem0.ai?utm_source=oss_sdk&utm_medium=in_product&utm_campaign=decay_stub&utm_content=node_error and use `import MemoryClient from "mem0ai"`.';
|
||||
const PLAIN_DECAY_ERROR =
|
||||
"The decay parameter is not supported by the OSS Memory SDK.";
|
||||
const PROJECT_UPDATE_ERROR =
|
||||
"Project updates are not supported by the OSS Memory SDK.";
|
||||
|
||||
function makeTempMem0Dir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "mem0-node-decay-feature-"));
|
||||
}
|
||||
|
||||
function decayPayload(overrides: Record<string, any> = {}) {
|
||||
return {
|
||||
notices: {
|
||||
decay_stub: {
|
||||
enabled: true,
|
||||
notice_type: "error",
|
||||
copy: DECAY_COPY,
|
||||
...overrides,
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function createFetchMock(options: {
|
||||
variant?: string;
|
||||
payload?: unknown;
|
||||
failFlags?: boolean;
|
||||
flagEnabled?: boolean;
|
||||
}) {
|
||||
const calls: any[] = [];
|
||||
const fetchMock = jest.fn(async (url: string | URL, init?: RequestInit) => {
|
||||
const target = String(url);
|
||||
|
||||
if (target.includes("/flags")) {
|
||||
if (options.failFlags) {
|
||||
throw new Error("flag evaluation failed");
|
||||
}
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({
|
||||
flags: {
|
||||
"mem0-oss-notices": {
|
||||
key: "mem0-oss-notices",
|
||||
enabled: options.flagEnabled ?? true,
|
||||
variant: options.variant ?? "displayed",
|
||||
metadata: {
|
||||
payload:
|
||||
options.payload === undefined
|
||||
? JSON.stringify(decayPayload())
|
||||
: options.payload,
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
if (target.includes("/i/v0/e/")) {
|
||||
calls.push(JSON.parse(String(init?.body)));
|
||||
return {
|
||||
ok: true,
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({}),
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
});
|
||||
|
||||
return { fetchMock, calls };
|
||||
}
|
||||
|
||||
function noticeEvents(calls: any[]) {
|
||||
return calls.filter((call) => call.event === "mem0.notice_displayed");
|
||||
}
|
||||
|
||||
async function createMemory() {
|
||||
const { Memory } = await import("../src/memory");
|
||||
return new Memory({
|
||||
version: "v1.1",
|
||||
embedder: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "text-embedding-3-small" },
|
||||
},
|
||||
vectorStore: {
|
||||
provider: "memory",
|
||||
config: {
|
||||
collectionName: `test-decay-feature-${Date.now()}-${Math.random()}`,
|
||||
dimension: 1536,
|
||||
dbPath: ":memory:",
|
||||
},
|
||||
},
|
||||
llm: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "gpt-5-mini" },
|
||||
},
|
||||
historyDbPath: ":memory:",
|
||||
});
|
||||
}
|
||||
|
||||
describe("Node OSS decay feature error notice", () => {
|
||||
let originalMem0Dir: string | undefined;
|
||||
let originalTelemetry: string | undefined;
|
||||
let originalSampleRate: string | undefined;
|
||||
let originalFetch: typeof global.fetch;
|
||||
|
||||
beforeEach(() => {
|
||||
originalMem0Dir = process.env.MEM0_DIR;
|
||||
originalTelemetry = process.env.MEM0_TELEMETRY;
|
||||
originalSampleRate = process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
originalFetch = global.fetch;
|
||||
|
||||
process.env.MEM0_DIR = makeTempMem0Dir();
|
||||
process.env.MEM0_TELEMETRY = "true";
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = "1";
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalMem0Dir === undefined) delete process.env.MEM0_DIR;
|
||||
else process.env.MEM0_DIR = originalMem0Dir;
|
||||
|
||||
if (originalTelemetry === undefined) delete process.env.MEM0_TELEMETRY;
|
||||
else process.env.MEM0_TELEMETRY = originalTelemetry;
|
||||
|
||||
if (originalSampleRate === undefined) {
|
||||
delete process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
} else {
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = originalSampleRate;
|
||||
}
|
||||
|
||||
global.fetch = originalFetch;
|
||||
jest.restoreAllMocks();
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
it("raises CTA copy for displayed and emits displayed=true", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject({ decay: true })).rejects.toThrow(
|
||||
DECAY_COPY,
|
||||
);
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "decay_stub",
|
||||
notice_type: "error",
|
||||
flag_key: "mem0-oss-notices",
|
||||
variant: "displayed",
|
||||
displayed: true,
|
||||
payload: DECAY_COPY,
|
||||
notice_config_found: true,
|
||||
sync_type: "async",
|
||||
trigger_function: "update_project",
|
||||
trigger_parameter: "decay",
|
||||
sample_rate: 1,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("raises CTA copy for holdout and emits displayed=true", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "holdout" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject({ decay: true })).rejects.toThrow(
|
||||
DECAY_COPY,
|
||||
);
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "decay_stub",
|
||||
variant: "holdout",
|
||||
displayed: true,
|
||||
trigger_function: "update_project",
|
||||
trigger_parameter: "decay",
|
||||
}),
|
||||
);
|
||||
expect(notices[0].properties.bypass_reason).toBeUndefined();
|
||||
});
|
||||
|
||||
it("uses plain error for unknown future variants and emits not_displayed", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "silent" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject({ decay: true })).rejects.toThrow(
|
||||
PLAIN_DECAY_ERROR,
|
||||
);
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "decay_stub",
|
||||
variant: "silent",
|
||||
displayed: false,
|
||||
bypass_reason: "not_displayed",
|
||||
payload: DECAY_COPY,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("uses plain error for disabled payload and emits payload_disabled", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "holdout",
|
||||
payload: JSON.stringify(decayPayload({ enabled: false })),
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject({ decay: true })).rejects.toThrow(
|
||||
PLAIN_DECAY_ERROR,
|
||||
);
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "decay_stub",
|
||||
displayed: false,
|
||||
bypass_reason: "payload_disabled",
|
||||
disabled_reason: "payload_disabled",
|
||||
notice_config_found: true,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it.each([
|
||||
[
|
||||
"missing config",
|
||||
JSON.stringify({ notices: {} }),
|
||||
"missing_notice_config",
|
||||
],
|
||||
[
|
||||
"missing copy",
|
||||
JSON.stringify(decayPayload({ copy: "" })),
|
||||
"missing_copy",
|
||||
],
|
||||
["malformed payload", "{not-json", "missing_notice_config"],
|
||||
])("uses plain error for %s", async (_label, payload, bypassReason) => {
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject({ decay: true })).rejects.toThrow(
|
||||
PLAIN_DECAY_ERROR,
|
||||
);
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "decay_stub",
|
||||
displayed: false,
|
||||
bypass_reason: bypassReason,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("uses plain error and emits no event when the blunt flag is disabled", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
flagEnabled: false,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject({ decay: true })).rejects.toThrow(
|
||||
PLAIN_DECAY_ERROR,
|
||||
);
|
||||
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("uses plain error and emits no event when PostHog fails", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ failFlags: true });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject({ decay: true })).rejects.toThrow(
|
||||
PLAIN_DECAY_ERROR,
|
||||
);
|
||||
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("uses plain error and skips flag evaluation when telemetry is off", async () => {
|
||||
process.env.MEM0_TELEMETRY = "False";
|
||||
jest.resetModules();
|
||||
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject({ decay: true })).rejects.toThrow(
|
||||
PLAIN_DECAY_ERROR,
|
||||
);
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["empty options", {}],
|
||||
["decay false", { decay: false }],
|
||||
["non-decay options", { customInstructions: "Updated" }],
|
||||
])("does not emit notice telemetry for %s", async (_label, options) => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject(options)).rejects.toThrow(
|
||||
PROJECT_UPDATE_ERROR,
|
||||
);
|
||||
|
||||
expect(
|
||||
fetchMock.mock.calls.filter(([url]) => String(url).includes("/flags")),
|
||||
).toHaveLength(0);
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("does not trigger first-run notice or update_project telemetry", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.updateProject({ decay: true })).rejects.toThrow(
|
||||
DECAY_COPY,
|
||||
);
|
||||
|
||||
const eventNames = calls.map((call) => call.event);
|
||||
expect(eventNames).toContain("mem0.notice_displayed");
|
||||
expect(eventNames).not.toContain("mem0.update_project");
|
||||
expect(
|
||||
noticeEvents(calls).some(
|
||||
(call) => call.properties.notice_id === "first_run",
|
||||
),
|
||||
).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,526 @@
|
||||
/// <reference types="jest" />
|
||||
import * as fs from "fs";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
|
||||
jest.setTimeout(15000);
|
||||
|
||||
jest.mock("../src/embeddings/google", () => ({
|
||||
GoogleEmbedder: jest.fn(),
|
||||
}));
|
||||
jest.mock("../src/llms/google", () => ({
|
||||
GoogleLLM: jest.fn(),
|
||||
}));
|
||||
|
||||
const mockEmbedding = new Array(1536).fill(0.1);
|
||||
jest.mock("../src/embeddings/openai", () => ({
|
||||
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
|
||||
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||
embedBatch: jest
|
||||
.fn()
|
||||
.mockImplementation((texts: string[]) =>
|
||||
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||
),
|
||||
embeddingDims: 1536,
|
||||
})),
|
||||
}));
|
||||
|
||||
jest.mock("../src/llms/openai", () => ({
|
||||
OpenAILLM: jest.fn().mockImplementation(() => ({
|
||||
generateResponse: jest.fn(),
|
||||
})),
|
||||
}));
|
||||
|
||||
const DECAY_USAGE_COPY =
|
||||
"Tip: Python fallback copy with memory.project.update(decay=True).";
|
||||
|
||||
function makeTempMem0Dir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "mem0-node-decay-usage-"));
|
||||
}
|
||||
|
||||
function configPath(): string {
|
||||
return path.join(process.env.MEM0_DIR as string, "config.json");
|
||||
}
|
||||
|
||||
function readConfig(): Record<string, any> {
|
||||
const file = configPath();
|
||||
if (!fs.existsSync(file)) return {};
|
||||
return JSON.parse(fs.readFileSync(file, "utf8"));
|
||||
}
|
||||
|
||||
function writeConfig(config: Record<string, any>) {
|
||||
fs.mkdirSync(path.dirname(configPath()), { recursive: true });
|
||||
fs.writeFileSync(configPath(), JSON.stringify(config, null, 4));
|
||||
}
|
||||
|
||||
function consumeFirstRun() {
|
||||
writeConfig({
|
||||
user_id: "node-decay-usage-test-user",
|
||||
notice_state: {
|
||||
first_run: {
|
||||
consumed: true,
|
||||
trigger_function: "test_setup",
|
||||
variant: "test",
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function decayUsagePayload(overrides: Record<string, any> = {}) {
|
||||
return {
|
||||
notices: {
|
||||
decay_usage: {
|
||||
enabled: true,
|
||||
notice_type: "log_line",
|
||||
copy: DECAY_USAGE_COPY,
|
||||
...overrides,
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function createFetchMock(options: {
|
||||
variant?: string;
|
||||
payload?: unknown;
|
||||
failFlags?: boolean;
|
||||
flagEnabled?: boolean;
|
||||
}) {
|
||||
const calls: any[] = [];
|
||||
const fetchMock = jest.fn(async (url: string | URL, init?: RequestInit) => {
|
||||
const target = String(url);
|
||||
|
||||
if (target.includes("/flags")) {
|
||||
if (options.failFlags) {
|
||||
throw new Error("flag evaluation failed");
|
||||
}
|
||||
const payload =
|
||||
options.payload === undefined
|
||||
? JSON.stringify(decayUsagePayload())
|
||||
: options.payload;
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({
|
||||
flags: {
|
||||
"mem0-oss-notices": {
|
||||
key: "mem0-oss-notices",
|
||||
enabled: options.flagEnabled ?? true,
|
||||
variant: options.variant ?? "displayed",
|
||||
metadata: { payload },
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
if (target.includes("/i/v0/e/")) {
|
||||
calls.push(JSON.parse(String(init?.body)));
|
||||
return {
|
||||
ok: true,
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({}),
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
});
|
||||
|
||||
return { fetchMock, calls };
|
||||
}
|
||||
|
||||
function noticeEvents(calls: any[]) {
|
||||
return calls.filter((call) => call.event === "mem0.notice_displayed");
|
||||
}
|
||||
|
||||
function flagRequestCount(fetchMock: jest.Mock) {
|
||||
return fetchMock.mock.calls.filter(([url]) => String(url).includes("/flags"))
|
||||
.length;
|
||||
}
|
||||
|
||||
async function createMemory() {
|
||||
const { Memory } = await import("../src/memory");
|
||||
return new Memory({
|
||||
version: "v1.1",
|
||||
embedder: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "text-embedding-3-small" },
|
||||
},
|
||||
vectorStore: {
|
||||
provider: "memory",
|
||||
config: {
|
||||
collectionName: `test-decay-usage-${Date.now()}-${Math.random()}`,
|
||||
dimension: 1536,
|
||||
dbPath: ":memory:",
|
||||
},
|
||||
},
|
||||
llm: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "gpt-5-mini" },
|
||||
},
|
||||
historyDbPath: ":memory:",
|
||||
});
|
||||
}
|
||||
|
||||
async function addMemories(memory: any, userId: string, count: number) {
|
||||
const ids: string[] = [];
|
||||
for (let index = 0; index < count; index++) {
|
||||
const result = await memory.add(`Decay usage memory ${index}`, {
|
||||
userId,
|
||||
infer: false,
|
||||
});
|
||||
ids.push(result.results[0].id);
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
describe("Node OSS decay usage notice", () => {
|
||||
let originalMem0Dir: string | undefined;
|
||||
let originalTelemetry: string | undefined;
|
||||
let originalSampleRate: string | undefined;
|
||||
let originalFetch: typeof global.fetch;
|
||||
let stderrSpy: jest.SpyInstance;
|
||||
|
||||
beforeEach(() => {
|
||||
originalMem0Dir = process.env.MEM0_DIR;
|
||||
originalTelemetry = process.env.MEM0_TELEMETRY;
|
||||
originalSampleRate = process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
originalFetch = global.fetch;
|
||||
|
||||
process.env.MEM0_DIR = makeTempMem0Dir();
|
||||
process.env.MEM0_TELEMETRY = "true";
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = "1";
|
||||
stderrSpy = jest
|
||||
.spyOn(process.stderr, "write")
|
||||
.mockImplementation(() => true);
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalMem0Dir === undefined) delete process.env.MEM0_DIR;
|
||||
else process.env.MEM0_DIR = originalMem0Dir;
|
||||
|
||||
if (originalTelemetry === undefined) delete process.env.MEM0_TELEMETRY;
|
||||
else process.env.MEM0_TELEMETRY = originalTelemetry;
|
||||
|
||||
if (originalSampleRate === undefined) {
|
||||
delete process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
} else {
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = originalSampleRate;
|
||||
}
|
||||
|
||||
global.fetch = originalFetch;
|
||||
jest.restoreAllMocks();
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
it("does not evaluate or write decay state before the 5th successful delete", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
const ids = await addMemories(memory, "decay-delete-user", 4);
|
||||
|
||||
for (const id of ids) {
|
||||
await memory.delete(id);
|
||||
}
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.decay_usage).toBeUndefined();
|
||||
expect(stderrSpy).not.toHaveBeenCalledWith(
|
||||
expect.stringContaining(DECAY_USAGE_COPY),
|
||||
);
|
||||
});
|
||||
|
||||
it("evaluates on the 5th successful delete and records delete_count", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
const ids = await addMemories(memory, "decay-delete-user", 5);
|
||||
|
||||
for (const id of ids) {
|
||||
await memory.delete(id);
|
||||
}
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "decay_usage",
|
||||
notice_type: "log_line",
|
||||
flag_key: "mem0-oss-notices",
|
||||
variant: "displayed",
|
||||
displayed: true,
|
||||
payload: DECAY_USAGE_COPY,
|
||||
notice_config_found: true,
|
||||
sync_type: "async",
|
||||
trigger_function: "delete",
|
||||
trigger_source: "delete_count",
|
||||
trigger_reason: "repeated_deletes",
|
||||
delete_count: 5,
|
||||
sample_rate: 1,
|
||||
}),
|
||||
);
|
||||
expect(stderrSpy.mock.calls.flat().join("")).toContain(DECAY_USAGE_COPY);
|
||||
expect(readConfig().notice_state.decay_usage.events).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("does not count or emit when delete fails", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
jest
|
||||
.spyOn(memory as any, "deleteMemory")
|
||||
.mockRejectedValue(new Error("delete failed"));
|
||||
|
||||
await expect(memory.delete("memory-id")).rejects.toThrow("delete failed");
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.decay_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("evaluates deleteAll after deleting at least one memory", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addMemories(memory, "decay-delete-all-user", 3);
|
||||
|
||||
await memory.deleteAll({ userId: "decay-delete-all-user" });
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "decay_usage",
|
||||
variant: "displayed",
|
||||
displayed: true,
|
||||
payload: DECAY_USAGE_COPY,
|
||||
trigger_function: "delete_all",
|
||||
trigger_source: "delete_all",
|
||||
trigger_reason: "bulk_delete",
|
||||
deleted_count: 3,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("does not evaluate deleteAll when no memories are deleted", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.deleteAll({ userId: "empty-delete-all-user" });
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.decay_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not evaluate when deleteAll fails", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.deleteAll({} as any)).rejects.toThrow(
|
||||
"At least one filter is required",
|
||||
);
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.decay_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("is silent for holdout and emits displayed=false", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "holdout" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addMemories(memory, "decay-holdout-user", 3);
|
||||
|
||||
await memory.deleteAll({ userId: "decay-holdout-user" });
|
||||
|
||||
expect(stderrSpy).not.toHaveBeenCalledWith(
|
||||
expect.stringContaining(DECAY_USAGE_COPY),
|
||||
);
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "decay_usage",
|
||||
variant: "holdout",
|
||||
displayed: false,
|
||||
bypass_reason: "holdout",
|
||||
trigger_function: "delete_all",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it.each([
|
||||
[
|
||||
"disabled payload",
|
||||
JSON.stringify(decayUsagePayload({ enabled: false })),
|
||||
"payload_disabled",
|
||||
"payload_disabled",
|
||||
],
|
||||
[
|
||||
"missing config",
|
||||
JSON.stringify({ notices: {} }),
|
||||
"missing_notice_config",
|
||||
undefined,
|
||||
],
|
||||
[
|
||||
"missing copy",
|
||||
JSON.stringify({
|
||||
notices: {
|
||||
decay_usage: { enabled: true, notice_type: "log_line" },
|
||||
},
|
||||
}),
|
||||
"missing_copy",
|
||||
undefined,
|
||||
],
|
||||
["malformed payload", "{not-json", "missing_notice_config", undefined],
|
||||
])(
|
||||
"stays silent and emits safe bypass for %s",
|
||||
async (_label, payload, bypassReason, disabledReason) => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addMemories(memory, "decay-bypass-user", 3);
|
||||
|
||||
await memory.deleteAll({ userId: "decay-bypass-user" });
|
||||
|
||||
expect(stderrSpy).not.toHaveBeenCalledWith(
|
||||
expect.stringContaining(DECAY_USAGE_COPY),
|
||||
);
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "decay_usage",
|
||||
displayed: false,
|
||||
bypass_reason: bypassReason,
|
||||
...(disabledReason && { disabled_reason: disabledReason }),
|
||||
}),
|
||||
);
|
||||
},
|
||||
);
|
||||
|
||||
it("does not emit or consume cap when the blunt flag is disabled", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
flagEnabled: false,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addMemories(memory, "decay-flag-disabled-user", 3);
|
||||
|
||||
await memory.deleteAll({ userId: "decay-flag-disabled-user" });
|
||||
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.decay_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not emit or consume cap when PostHog fails", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ failFlags: true });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addMemories(memory, "decay-posthog-failure-user", 3);
|
||||
|
||||
await memory.deleteAll({ userId: "decay-posthog-failure-user" });
|
||||
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.decay_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("skips flag evaluation, event emission, and state writes when telemetry is off", async () => {
|
||||
process.env.MEM0_TELEMETRY = "False";
|
||||
jest.resetModules();
|
||||
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addMemories(memory, "decay-telemetry-off-user", 3);
|
||||
|
||||
await memory.deleteAll({ userId: "decay-telemetry-off-user" });
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.decay_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("blocks the 11th evaluated opportunity before flag evaluation", async () => {
|
||||
consumeFirstRun();
|
||||
const now = new Date();
|
||||
writeConfig({
|
||||
...readConfig(),
|
||||
notice_state: {
|
||||
...readConfig().notice_state,
|
||||
first_run: readConfig().notice_state.first_run,
|
||||
decay_usage: {
|
||||
events: Array.from({ length: 10 }, (_, index) => ({
|
||||
evaluated_at: new Date(now.getTime() - index * 1000).toISOString(),
|
||||
variant: "displayed",
|
||||
})),
|
||||
},
|
||||
},
|
||||
});
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addMemories(memory, "decay-cap-user", 1);
|
||||
|
||||
await memory.deleteAll({ userId: "decay-cap-user" });
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state.decay_usage.events).toHaveLength(10);
|
||||
});
|
||||
|
||||
it("does not consume first-run on a qualifying decay usage call", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload: JSON.stringify({
|
||||
notices: {
|
||||
first_run: {
|
||||
enabled: true,
|
||||
notice_type: "log_line",
|
||||
copy: "First-run should not appear",
|
||||
},
|
||||
decay_usage: {
|
||||
enabled: true,
|
||||
notice_type: "log_line",
|
||||
copy: DECAY_USAGE_COPY,
|
||||
},
|
||||
},
|
||||
}),
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addMemories(memory, "decay-priority-user", 3);
|
||||
const config = readConfig();
|
||||
delete config.notice_state.first_run;
|
||||
writeConfig(config);
|
||||
|
||||
await memory.deleteAll({ userId: "decay-priority-user" });
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties.notice_id).toBe("decay_usage");
|
||||
expect(readConfig().notice_state.first_run).toBeUndefined();
|
||||
expect(readConfig().notice_state.decay_usage.events).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,409 @@
|
||||
/// <reference types="jest" />
|
||||
import * as fs from "fs";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
|
||||
jest.setTimeout(15000);
|
||||
|
||||
jest.mock("../src/embeddings/google", () => ({
|
||||
GoogleEmbedder: jest.fn(),
|
||||
}));
|
||||
jest.mock("../src/llms/google", () => ({
|
||||
GoogleLLM: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock("../src/llms/openai", () => ({
|
||||
OpenAILLM: jest.fn().mockImplementation(() => ({
|
||||
generateResponse: jest.fn().mockResolvedValue(
|
||||
JSON.stringify({
|
||||
memory: [{ id: "0", text: "stored fact", attributed_to: "user" }],
|
||||
}),
|
||||
),
|
||||
})),
|
||||
}));
|
||||
|
||||
const mockEmbedding = new Array(1536).fill(0.1);
|
||||
jest.mock("../src/embeddings/openai", () => ({
|
||||
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
|
||||
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||
embedBatch: jest
|
||||
.fn()
|
||||
.mockImplementation((texts: string[]) =>
|
||||
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||
),
|
||||
embeddingDims: 1536,
|
||||
})),
|
||||
}));
|
||||
|
||||
const FIRST_RUN_COPY = "First-run CTA from PostHog";
|
||||
|
||||
function makeTempMem0Dir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "mem0-node-first-run-"));
|
||||
}
|
||||
|
||||
function firstRunPayload(overrides: Record<string, any> = {}) {
|
||||
return {
|
||||
notices: {
|
||||
first_run: {
|
||||
enabled: true,
|
||||
notice_type: "log_line",
|
||||
copy: FIRST_RUN_COPY,
|
||||
...overrides,
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function createFetchMock(options: {
|
||||
variant?: string;
|
||||
payload?: Record<string, any>;
|
||||
failFlags?: boolean;
|
||||
}) {
|
||||
const calls: any[] = [];
|
||||
const fetchMock = jest.fn(async (url: string | URL, init?: RequestInit) => {
|
||||
const target = String(url);
|
||||
|
||||
if (target.includes("/flags")) {
|
||||
if (options.failFlags) {
|
||||
throw new Error("flag evaluation failed");
|
||||
}
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({
|
||||
flags: {
|
||||
"mem0-oss-notices": {
|
||||
key: "mem0-oss-notices",
|
||||
enabled: true,
|
||||
variant: options.variant ?? "displayed",
|
||||
metadata: {
|
||||
payload: JSON.stringify(options.payload ?? firstRunPayload()),
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
if (target.includes("/i/v0/e/")) {
|
||||
calls.push(JSON.parse(String(init?.body)));
|
||||
return {
|
||||
ok: true,
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({}),
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
});
|
||||
|
||||
return { fetchMock, calls };
|
||||
}
|
||||
|
||||
function noticeEvents(calls: any[]) {
|
||||
return calls.filter((call) => call.event === "mem0.notice_displayed");
|
||||
}
|
||||
|
||||
async function createMemory() {
|
||||
const { Memory } = await import("../src/memory");
|
||||
return new Memory({
|
||||
version: "v1.1",
|
||||
embedder: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "text-embedding-3-small" },
|
||||
},
|
||||
vectorStore: {
|
||||
provider: "memory",
|
||||
config: {
|
||||
collectionName: `test-first-run-${Date.now()}-${Math.random()}`,
|
||||
dimension: 1536,
|
||||
dbPath: ":memory:",
|
||||
},
|
||||
},
|
||||
llm: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "gpt-5-mini" },
|
||||
},
|
||||
historyDbPath: ":memory:",
|
||||
});
|
||||
}
|
||||
|
||||
describe("Node OSS first-run notice", () => {
|
||||
let originalMem0Dir: string | undefined;
|
||||
let originalTelemetry: string | undefined;
|
||||
let originalSampleRate: string | undefined;
|
||||
let originalFetch: typeof global.fetch;
|
||||
let stderrSpy: jest.SpyInstance;
|
||||
|
||||
beforeEach(() => {
|
||||
originalMem0Dir = process.env.MEM0_DIR;
|
||||
originalTelemetry = process.env.MEM0_TELEMETRY;
|
||||
originalSampleRate = process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
originalFetch = global.fetch;
|
||||
|
||||
process.env.MEM0_DIR = makeTempMem0Dir();
|
||||
process.env.MEM0_TELEMETRY = "true";
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = "1";
|
||||
stderrSpy = jest
|
||||
.spyOn(process.stderr, "write")
|
||||
.mockImplementation(() => true);
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalMem0Dir === undefined) delete process.env.MEM0_DIR;
|
||||
else process.env.MEM0_DIR = originalMem0Dir;
|
||||
|
||||
if (originalTelemetry === undefined) delete process.env.MEM0_TELEMETRY;
|
||||
else process.env.MEM0_TELEMETRY = originalTelemetry;
|
||||
|
||||
if (originalSampleRate === undefined) {
|
||||
delete process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
} else {
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = originalSampleRate;
|
||||
}
|
||||
|
||||
global.fetch = originalFetch;
|
||||
jest.restoreAllMocks();
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
it("shows the displayed first-run copy once after a successful public call", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Direct storage content", {
|
||||
userId: "first-run-user",
|
||||
infer: false,
|
||||
});
|
||||
await memory.add("Second direct storage content", {
|
||||
userId: "first-run-user",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
const stderrOutput = stderrSpy.mock.calls.flat().join("");
|
||||
expect(stderrOutput.match(new RegExp(FIRST_RUN_COPY, "g"))).toHaveLength(1);
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "first_run",
|
||||
notice_type: "log_line",
|
||||
flag_key: "mem0-oss-notices",
|
||||
variant: "displayed",
|
||||
displayed: true,
|
||||
payload: FIRST_RUN_COPY,
|
||||
notice_config_found: true,
|
||||
sync_type: "async",
|
||||
trigger_function: "add",
|
||||
sample_rate: 1,
|
||||
}),
|
||||
);
|
||||
|
||||
const { __noticeTestHooks } = await import("../src/utils/notices");
|
||||
const config = __noticeTestHooks.loadMem0Config();
|
||||
expect(config.notice_state.first_run).toEqual(
|
||||
expect.objectContaining({
|
||||
consumed: true,
|
||||
trigger_function: "add",
|
||||
variant: "displayed",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("treats missing enabled as enabled, matching Python payload semantics", async () => {
|
||||
const payload = firstRunPayload();
|
||||
delete (payload.notices.first_run as Record<string, any>).enabled;
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Missing enabled payload content", {
|
||||
userId: "first-run-missing-enabled",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
const stderrOutput = stderrSpy.mock.calls.flat().join("");
|
||||
expect(stderrOutput).toContain(FIRST_RUN_COPY);
|
||||
expect(noticeEvents(calls)[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "first_run",
|
||||
displayed: true,
|
||||
payload: FIRST_RUN_COPY,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("is silent for holdout but still emits and consumes the opportunity", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "holdout" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Holdout content", {
|
||||
userId: "first-run-holdout",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
expect(stderrSpy).not.toHaveBeenCalledWith(
|
||||
expect.stringContaining(FIRST_RUN_COPY),
|
||||
);
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "first_run",
|
||||
variant: "holdout",
|
||||
displayed: false,
|
||||
bypass_reason: "holdout",
|
||||
notice_config_found: true,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("is silent for disabled payloads and records a safe bypass", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload: firstRunPayload({ enabled: false }),
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Disabled payload content", {
|
||||
userId: "first-run-disabled",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
const stderrOutput = stderrSpy.mock.calls.flat().join("");
|
||||
expect(stderrOutput).not.toContain(FIRST_RUN_COPY);
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "first_run",
|
||||
displayed: false,
|
||||
bypass_reason: "payload_disabled",
|
||||
disabled_reason: "payload_disabled",
|
||||
notice_config_found: true,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("lets the payload disabled kill switch take precedence over holdout", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "holdout",
|
||||
payload: firstRunPayload({ enabled: false }),
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Disabled holdout payload content", {
|
||||
userId: "first-run-disabled-holdout",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "first_run",
|
||||
variant: "holdout",
|
||||
displayed: false,
|
||||
bypass_reason: "payload_disabled",
|
||||
disabled_reason: "payload_disabled",
|
||||
notice_config_found: true,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("does not evaluate flags, emit, or write state when telemetry is off", async () => {
|
||||
process.env.MEM0_TELEMETRY = "false";
|
||||
jest.resetModules();
|
||||
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Telemetry off content", {
|
||||
userId: "first-run-off",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
const { __noticeTestHooks } = await import("../src/utils/notices");
|
||||
expect(__noticeTestHooks.loadMem0Config().notice_state).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not export internal notice test hooks from the public OSS entrypoint", async () => {
|
||||
const publicOssEntry = await import("../src");
|
||||
|
||||
expect(publicOssEntry).not.toHaveProperty("__noticeTestHooks");
|
||||
});
|
||||
|
||||
it("treats uppercase False as telemetry off", async () => {
|
||||
process.env.MEM0_TELEMETRY = "False";
|
||||
jest.resetModules();
|
||||
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Uppercase telemetry off content", {
|
||||
userId: "first-run-off-uppercase",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
const { __noticeTestHooks } = await import("../src/utils/notices");
|
||||
expect(__noticeTestHooks.loadMem0Config().notice_state).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not run after a failed public call", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.add("Missing owner", {} as any)).rejects.toThrow(
|
||||
"One of the filters: userId, agentId or runId is required!",
|
||||
);
|
||||
|
||||
expect(
|
||||
fetchMock.mock.calls.filter(([url]) => String(url).includes("/flags")),
|
||||
).toHaveLength(0);
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("does not break the successful operation when flag evaluation fails", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ failFlags: true });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
const result = await memory.add("Flag failure content", {
|
||||
userId: "first-run-failure",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
expect(result.results).toHaveLength(1);
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("can trigger from get() when get is the first successful public call", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(memory.get("missing-memory-id")).resolves.toBeNull();
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties.trigger_function).toBe("get");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,227 @@
|
||||
/// <reference types="jest" />
|
||||
import * as fs from "fs";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
import type { TelemetryInstance } from "../src/utils/telemetry.types";
|
||||
|
||||
function makeTempMem0Dir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "mem0-node-notices-"));
|
||||
}
|
||||
|
||||
function makeInstance(
|
||||
overrides: Partial<TelemetryInstance> = {},
|
||||
): TelemetryInstance {
|
||||
return {
|
||||
telemetryId: "notice-test-user",
|
||||
constructor: { name: "Memory" },
|
||||
host: "https://test.example.com",
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe("Node OSS notice foundation", () => {
|
||||
let originalMem0Dir: string | undefined;
|
||||
let originalTelemetry: string | undefined;
|
||||
let originalFetch: typeof global.fetch;
|
||||
|
||||
beforeEach(() => {
|
||||
originalMem0Dir = process.env.MEM0_DIR;
|
||||
originalTelemetry = process.env.MEM0_TELEMETRY;
|
||||
originalFetch = global.fetch;
|
||||
process.env.MEM0_DIR = makeTempMem0Dir();
|
||||
process.env.MEM0_TELEMETRY = "true";
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalMem0Dir === undefined) delete process.env.MEM0_DIR;
|
||||
else process.env.MEM0_DIR = originalMem0Dir;
|
||||
|
||||
if (originalTelemetry === undefined) delete process.env.MEM0_TELEMETRY;
|
||||
else process.env.MEM0_TELEMETRY = originalTelemetry;
|
||||
|
||||
global.fetch = originalFetch;
|
||||
jest.restoreAllMocks();
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
it("writes notice state into isolated MEM0_DIR config and preserves user_id", async () => {
|
||||
const notices = await import("../src/utils/notices");
|
||||
const configPath = notices.getMem0ConfigPath();
|
||||
|
||||
fs.mkdirSync(path.dirname(configPath), { recursive: true });
|
||||
fs.writeFileSync(configPath, JSON.stringify({ user_id: "existing-user" }));
|
||||
|
||||
expect(
|
||||
notices.recordNoticeOpportunity("foundation_notice", {
|
||||
variant: "displayed",
|
||||
}),
|
||||
).toBe(true);
|
||||
|
||||
const config = JSON.parse(fs.readFileSync(configPath, "utf8"));
|
||||
expect(config.user_id).toBe("existing-user");
|
||||
expect(config.notice_state.foundation_notice.events).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("writes config through a temp-file path and leaves no temp file behind", async () => {
|
||||
const notices = await import("../src/utils/notices");
|
||||
|
||||
expect(notices.writeMem0ConfigAtomic({ user_id: "atomic-user" })).toBe(
|
||||
true,
|
||||
);
|
||||
|
||||
const configDir = path.dirname(notices.getMem0ConfigPath());
|
||||
const tempFiles = fs
|
||||
.readdirSync(configDir)
|
||||
.filter((name) => name.endsWith(".tmp"));
|
||||
expect(tempFiles).toHaveLength(0);
|
||||
expect(notices.loadMem0Config().user_id).toBe("atomic-user");
|
||||
});
|
||||
|
||||
it("allows 10 evaluated opportunities in a rolling window and blocks the 11th", async () => {
|
||||
const notices = await import("../src/utils/notices");
|
||||
const now = new Date("2026-06-11T12:00:00.000Z");
|
||||
let state: Record<string, any> = {};
|
||||
|
||||
for (let i = 0; i < 10; i++) {
|
||||
const nextState = notices.appendNoticeCapEvent(
|
||||
state,
|
||||
{ variant: "displayed", index: i },
|
||||
{ now: new Date(now.getTime() + i) },
|
||||
);
|
||||
expect(nextState).not.toBeNull();
|
||||
state = nextState!;
|
||||
}
|
||||
|
||||
expect(notices.hasNoticeCapRoom(state, { now })).toBe(false);
|
||||
expect(
|
||||
notices.appendNoticeCapEvent(state, { variant: "displayed" }, { now }),
|
||||
).toBeNull();
|
||||
});
|
||||
|
||||
it("drops old cap events outside the rolling window", async () => {
|
||||
const notices = await import("../src/utils/notices");
|
||||
const oldEvent = {
|
||||
evaluated_at: "2026-06-01T00:00:00.000Z",
|
||||
variant: "displayed",
|
||||
};
|
||||
const now = new Date("2026-06-11T12:00:00.000Z");
|
||||
|
||||
expect(notices.hasNoticeCapRoom({ events: [oldEvent] }, { now })).toBe(
|
||||
true,
|
||||
);
|
||||
const nextState = notices.appendNoticeCapEvent(
|
||||
{ events: [oldEvent] },
|
||||
{ variant: "displayed" },
|
||||
{ now },
|
||||
);
|
||||
expect(nextState?.events).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("does not evaluate flags or write notice state when telemetry is off", async () => {
|
||||
process.env.MEM0_TELEMETRY = "false";
|
||||
jest.resetModules();
|
||||
|
||||
const fetchMock = jest.fn();
|
||||
global.fetch = fetchMock as any;
|
||||
const notices = await import("../src/utils/notices");
|
||||
|
||||
await expect(
|
||||
notices.evaluateNoticeFlag("notice-test-user", { fetchImpl: fetchMock }),
|
||||
).resolves.toBeNull();
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(
|
||||
notices.recordNoticeOpportunity("foundation_notice", {
|
||||
variant: "displayed",
|
||||
}),
|
||||
).toBe(false);
|
||||
expect(fs.existsSync(notices.getMem0ConfigPath())).toBe(false);
|
||||
});
|
||||
|
||||
it("returns null when PostHog flag evaluation fails", async () => {
|
||||
const notices = await import("../src/utils/notices");
|
||||
const fetchMock = jest.fn().mockRejectedValue(new Error("network down"));
|
||||
|
||||
await expect(
|
||||
notices.evaluateNoticeFlag("notice-test-user", { fetchImpl: fetchMock }),
|
||||
).resolves.toBeNull();
|
||||
});
|
||||
|
||||
it("returns null when PostHog flag evaluation times out", async () => {
|
||||
const notices = await import("../src/utils/notices");
|
||||
const fetchMock = jest.fn(
|
||||
(_url: string | URL | Request, init?: RequestInit) =>
|
||||
new Promise((_resolve, reject) => {
|
||||
init?.signal?.addEventListener("abort", () => {
|
||||
reject(new Error("aborted"));
|
||||
});
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(
|
||||
notices.evaluateNoticeFlag("notice-test-user", {
|
||||
fetchImpl: fetchMock as any,
|
||||
timeoutMs: 1,
|
||||
}),
|
||||
).resolves.toBeNull();
|
||||
});
|
||||
|
||||
it("parses displayed variant and JSON payload from PostHog flags response", async () => {
|
||||
const notices = await import("../src/utils/notices");
|
||||
const payload = {
|
||||
notices: {
|
||||
foundation_notice: {
|
||||
enabled: true,
|
||||
notice_type: "log_line",
|
||||
copy: "Foundation notice",
|
||||
},
|
||||
},
|
||||
};
|
||||
const fetchMock = jest.fn().mockResolvedValue({
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({
|
||||
flags: {
|
||||
"mem0-oss-notices": {
|
||||
key: "mem0-oss-notices",
|
||||
enabled: true,
|
||||
variant: "displayed",
|
||||
metadata: { payload: JSON.stringify(payload) },
|
||||
},
|
||||
},
|
||||
}),
|
||||
});
|
||||
|
||||
const result = await notices.evaluateNoticeFlag("notice-test-user", {
|
||||
fetchImpl: fetchMock,
|
||||
});
|
||||
expect(result?.variant).toBe("displayed");
|
||||
const parsed = notices.getNoticeConfigFromPayload(
|
||||
result?.payload,
|
||||
"foundation_notice",
|
||||
);
|
||||
expect(parsed.found).toBe(true);
|
||||
expect(parsed.config?.copy).toBe("Foundation notice");
|
||||
});
|
||||
|
||||
it("emits mem0.notice_displayed with sample_rate=1", async () => {
|
||||
const fetchMock = jest.fn().mockResolvedValue({
|
||||
ok: true,
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const notices = await import("../src/utils/notices");
|
||||
|
||||
await notices.emitNoticeDisplayed(makeInstance(), {
|
||||
notice_id: "foundation_notice",
|
||||
notice_type: "log_line",
|
||||
flag_key: "mem0-oss-notices",
|
||||
variant: "displayed",
|
||||
displayed: true,
|
||||
});
|
||||
|
||||
const body = JSON.parse(fetchMock.mock.calls[0][1].body);
|
||||
expect(body.event).toBe("mem0.notice_displayed");
|
||||
expect(body.properties.sample_rate).toBe(1);
|
||||
expect(body.properties.notice_id).toBe("foundation_notice");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,492 @@
|
||||
/// <reference types="jest" />
|
||||
import * as fs from "fs";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
|
||||
jest.setTimeout(15000);
|
||||
|
||||
jest.mock("../src/embeddings/google", () => ({
|
||||
GoogleEmbedder: jest.fn(),
|
||||
}));
|
||||
jest.mock("../src/llms/google", () => ({
|
||||
GoogleLLM: jest.fn(),
|
||||
}));
|
||||
|
||||
const mockEmbedding = new Array(1536).fill(0.1);
|
||||
jest.mock("../src/embeddings/openai", () => ({
|
||||
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
|
||||
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||
embedBatch: jest
|
||||
.fn()
|
||||
.mockImplementation((texts: string[]) =>
|
||||
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||
),
|
||||
embeddingDims: 1536,
|
||||
})),
|
||||
}));
|
||||
|
||||
jest.mock("../src/llms/openai", () => ({
|
||||
OpenAILLM: jest.fn().mockImplementation(() => ({
|
||||
generateResponse: jest.fn(),
|
||||
})),
|
||||
}));
|
||||
|
||||
const PERFORMANCE_COPY =
|
||||
"Mem0 Platform is optimized for this type of workload. Its retrieval benchmarked at ~0.8-1.09s p50 across LoCoMo, LongMemEval, and BEAM 1M/10M workloads; you can use it for free by getting an API key at: https://app.mem0.ai";
|
||||
|
||||
function makeTempMem0Dir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "mem0-node-performance-"));
|
||||
}
|
||||
|
||||
function configPath(): string {
|
||||
return path.join(process.env.MEM0_DIR as string, "config.json");
|
||||
}
|
||||
|
||||
function readConfig(): Record<string, any> {
|
||||
const file = configPath();
|
||||
if (!fs.existsSync(file)) return {};
|
||||
return JSON.parse(fs.readFileSync(file, "utf8"));
|
||||
}
|
||||
|
||||
function writeConfig(config: Record<string, any>) {
|
||||
fs.mkdirSync(path.dirname(configPath()), { recursive: true });
|
||||
fs.writeFileSync(configPath(), JSON.stringify(config, null, 4));
|
||||
}
|
||||
|
||||
function consumeFirstRun() {
|
||||
writeConfig({
|
||||
user_id: "node-performance-test-user",
|
||||
notice_state: {
|
||||
first_run: {
|
||||
consumed: true,
|
||||
trigger_function: "test_setup",
|
||||
variant: "test",
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function performancePayload(overrides: Record<string, any> = {}) {
|
||||
return {
|
||||
notices: {
|
||||
performance_slow_query: {
|
||||
enabled: true,
|
||||
notice_type: "log_line",
|
||||
copy: PERFORMANCE_COPY,
|
||||
...overrides,
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function createFetchMock(options: {
|
||||
variant?: string;
|
||||
payload?: unknown;
|
||||
failFlags?: boolean;
|
||||
flagEnabled?: boolean;
|
||||
}) {
|
||||
const calls: any[] = [];
|
||||
const fetchMock = jest.fn(async (url: string | URL, init?: RequestInit) => {
|
||||
const target = String(url);
|
||||
|
||||
if (target.includes("/flags")) {
|
||||
if (options.failFlags) {
|
||||
throw new Error("flag evaluation failed");
|
||||
}
|
||||
const payload =
|
||||
options.payload === undefined
|
||||
? JSON.stringify(performancePayload())
|
||||
: options.payload;
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({
|
||||
flags: {
|
||||
"mem0-oss-notices": {
|
||||
key: "mem0-oss-notices",
|
||||
enabled: options.flagEnabled ?? true,
|
||||
variant: options.variant ?? "displayed",
|
||||
metadata: { payload },
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
if (target.includes("/i/v0/e/")) {
|
||||
calls.push(JSON.parse(String(init?.body)));
|
||||
return {
|
||||
ok: true,
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({}),
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
});
|
||||
|
||||
return { fetchMock, calls };
|
||||
}
|
||||
|
||||
function noticeEvents(calls: any[]) {
|
||||
return calls.filter((call) => call.event === "mem0.notice_displayed");
|
||||
}
|
||||
|
||||
function performanceEvents(calls: any[]) {
|
||||
return noticeEvents(calls).filter(
|
||||
(call) => call.properties.notice_id === "performance_slow_query",
|
||||
);
|
||||
}
|
||||
|
||||
function flagRequestCount(fetchMock: jest.Mock) {
|
||||
return fetchMock.mock.calls.filter(([url]) => String(url).includes("/flags"))
|
||||
.length;
|
||||
}
|
||||
|
||||
function mockSearchElapsed(elapsedMs: number) {
|
||||
let now = 1000;
|
||||
jest.spyOn(Date, "now").mockImplementation(() => {
|
||||
const current = now;
|
||||
now += elapsedMs;
|
||||
return current;
|
||||
});
|
||||
}
|
||||
|
||||
async function createMemory() {
|
||||
const { Memory } = await import("../src/memory");
|
||||
return new Memory({
|
||||
version: "v1.1",
|
||||
embedder: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "text-embedding-3-small" },
|
||||
},
|
||||
vectorStore: {
|
||||
provider: "memory",
|
||||
config: {
|
||||
collectionName: `test-performance-${Date.now()}-${Math.random()}`,
|
||||
dimension: 1536,
|
||||
dbPath: ":memory:",
|
||||
},
|
||||
},
|
||||
llm: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "gpt-5-mini" },
|
||||
},
|
||||
historyDbPath: ":memory:",
|
||||
});
|
||||
}
|
||||
|
||||
async function addSeed(memory: any, userId = "performance-user") {
|
||||
await memory.add("The user's favorite drink is green tea.", {
|
||||
userId,
|
||||
infer: false,
|
||||
});
|
||||
}
|
||||
|
||||
describe("Node OSS performance slow query notice", () => {
|
||||
let originalMem0Dir: string | undefined;
|
||||
let originalTelemetry: string | undefined;
|
||||
let originalSampleRate: string | undefined;
|
||||
let originalFetch: typeof global.fetch;
|
||||
let stderrSpy: jest.SpyInstance;
|
||||
|
||||
beforeEach(() => {
|
||||
originalMem0Dir = process.env.MEM0_DIR;
|
||||
originalTelemetry = process.env.MEM0_TELEMETRY;
|
||||
originalSampleRate = process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
originalFetch = global.fetch;
|
||||
|
||||
process.env.MEM0_DIR = makeTempMem0Dir();
|
||||
process.env.MEM0_TELEMETRY = "true";
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = "1";
|
||||
stderrSpy = jest
|
||||
.spyOn(process.stderr, "write")
|
||||
.mockImplementation(() => true);
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalMem0Dir === undefined) delete process.env.MEM0_DIR;
|
||||
else process.env.MEM0_DIR = originalMem0Dir;
|
||||
|
||||
if (originalTelemetry === undefined) delete process.env.MEM0_TELEMETRY;
|
||||
else process.env.MEM0_TELEMETRY = originalTelemetry;
|
||||
|
||||
if (originalSampleRate === undefined) {
|
||||
delete process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
} else {
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = originalSampleRate;
|
||||
}
|
||||
|
||||
global.fetch = originalFetch;
|
||||
jest.restoreAllMocks();
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
it("displays and emits safe fields after a slow successful search", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSeed(memory);
|
||||
stderrSpy.mockClear();
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
const result = await memory.search("favorite drink private text", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
expect(result.results.length).toBeGreaterThan(0);
|
||||
expect(stderrSpy).toHaveBeenCalledWith(`${PERFORMANCE_COPY}\n`);
|
||||
const event = performanceEvents(calls)[0];
|
||||
expect(event.properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "performance_slow_query",
|
||||
notice_type: "log_line",
|
||||
flag_key: "mem0-oss-notices",
|
||||
variant: "displayed",
|
||||
displayed: true,
|
||||
payload: PERFORMANCE_COPY,
|
||||
notice_config_found: true,
|
||||
sync_type: "async",
|
||||
trigger_function: "search",
|
||||
trigger_reason: "slow_query",
|
||||
elapsed_ms: 2345,
|
||||
threshold_ms: 2000,
|
||||
top_k: 3,
|
||||
result_count: result.results.length,
|
||||
sample_rate: 1,
|
||||
}),
|
||||
);
|
||||
const serialized = JSON.stringify(event.properties);
|
||||
expect(serialized).not.toContain("private text");
|
||||
expect(serialized).not.toContain("performance-user");
|
||||
});
|
||||
|
||||
it("is silent for holdout and emits displayed=false", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "holdout" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSeed(memory);
|
||||
stderrSpy.mockClear();
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
expect(stderrSpy).not.toHaveBeenCalled();
|
||||
expect(performanceEvents(calls)[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "performance_slow_query",
|
||||
variant: "holdout",
|
||||
displayed: false,
|
||||
bypass_reason: "holdout",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it.each([
|
||||
[
|
||||
"disabled payload",
|
||||
JSON.stringify(performancePayload({ enabled: false })),
|
||||
"payload_disabled",
|
||||
],
|
||||
[
|
||||
"missing config",
|
||||
JSON.stringify({ notices: {} }),
|
||||
"missing_notice_config",
|
||||
],
|
||||
[
|
||||
"missing copy",
|
||||
JSON.stringify(performancePayload({ copy: "" })),
|
||||
"missing_copy",
|
||||
],
|
||||
["malformed payload", "{not-json", "missing_notice_config"],
|
||||
])("is silent and safe for %s", async (_label, payload, bypassReason) => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSeed(memory);
|
||||
stderrSpy.mockClear();
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
expect(stderrSpy).not.toHaveBeenCalled();
|
||||
expect(performanceEvents(calls)[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
displayed: false,
|
||||
bypass_reason: bypassReason,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("does not evaluate fast searches", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSeed(memory);
|
||||
mockSearchElapsed(100);
|
||||
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(performanceEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.performance_slow_query).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not evaluate failed searches", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
await expect(
|
||||
memory.search("private slow query", {
|
||||
filters: {},
|
||||
topK: 3,
|
||||
}),
|
||||
).rejects.toThrow("filters must contain");
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(performanceEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.performance_slow_query).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not consume cap or emit when the blunt flag is disabled", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
flagEnabled: false,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSeed(memory);
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
expect(performanceEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.performance_slow_query).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not consume cap or emit when PostHog fails", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ failFlags: true });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSeed(memory);
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
expect(performanceEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.performance_slow_query).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does nothing when telemetry is off", async () => {
|
||||
process.env.MEM0_TELEMETRY = "false";
|
||||
jest.resetModules();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSeed(memory);
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(performanceEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state).toBeUndefined();
|
||||
});
|
||||
|
||||
it("caps evaluated opportunities at 10 per rolling week", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSeed(memory);
|
||||
stderrSpy.mockClear();
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
for (let index = 1; index <= 11; index++) {
|
||||
await memory.search(`favorite drink ${index}`, {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
}
|
||||
|
||||
expect(performanceEvents(calls)).toHaveLength(10);
|
||||
expect(flagRequestCount(fetchMock)).toBe(10);
|
||||
expect(stderrSpy).toHaveBeenCalledTimes(10);
|
||||
expect(
|
||||
readConfig().notice_state.performance_slow_query.events,
|
||||
).toHaveLength(10);
|
||||
});
|
||||
|
||||
it("lets temporal usage and scale threshold take priority", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSeed(memory);
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
await memory.search("what happened last week?", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 50,
|
||||
});
|
||||
|
||||
expect(performanceEvents(calls)).toHaveLength(0);
|
||||
const noticeIds = noticeEvents(calls).map(
|
||||
(call) => call.properties.notice_id,
|
||||
);
|
||||
expect(noticeIds).toEqual(["temporal_usage", "scale_threshold"]);
|
||||
});
|
||||
|
||||
it("does not consume first-run on the same qualifying performance call", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
mockSearchElapsed(2345);
|
||||
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "performance-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
expect(performanceEvents(calls)).toHaveLength(1);
|
||||
expect(readConfig().notice_state.first_run).toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,463 @@
|
||||
/// <reference types="jest" />
|
||||
import * as fs from "fs";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
|
||||
jest.setTimeout(15000);
|
||||
|
||||
jest.mock("../src/embeddings/google", () => ({
|
||||
GoogleEmbedder: jest.fn(),
|
||||
}));
|
||||
jest.mock("../src/llms/google", () => ({
|
||||
GoogleLLM: jest.fn(),
|
||||
}));
|
||||
|
||||
const mockEmbedding = new Array(1536).fill(0.1);
|
||||
jest.mock("../src/embeddings/openai", () => ({
|
||||
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
|
||||
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||
embedBatch: jest
|
||||
.fn()
|
||||
.mockImplementation((texts: string[]) =>
|
||||
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||
),
|
||||
embeddingDims: 1536,
|
||||
})),
|
||||
}));
|
||||
|
||||
jest.mock("../src/llms/openai", () => ({
|
||||
OpenAILLM: jest.fn().mockImplementation(() => ({
|
||||
generateResponse: jest.fn(),
|
||||
})),
|
||||
}));
|
||||
|
||||
const SCALE_TOP_K_COPY = "Scale top {top_k}";
|
||||
const SCALE_MEMORY_COUNT_COPY = "Scale count {memory_count}";
|
||||
|
||||
function makeTempMem0Dir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "mem0-node-scale-"));
|
||||
}
|
||||
|
||||
function configPath(): string {
|
||||
return path.join(process.env.MEM0_DIR as string, "config.json");
|
||||
}
|
||||
|
||||
function readConfig(): Record<string, any> {
|
||||
const file = configPath();
|
||||
if (!fs.existsSync(file)) return {};
|
||||
return JSON.parse(fs.readFileSync(file, "utf8"));
|
||||
}
|
||||
|
||||
function writeConfig(config: Record<string, any>) {
|
||||
fs.mkdirSync(path.dirname(configPath()), { recursive: true });
|
||||
fs.writeFileSync(configPath(), JSON.stringify(config, null, 4));
|
||||
}
|
||||
|
||||
function consumeFirstRun(scaleState: Record<string, any> = {}) {
|
||||
writeConfig({
|
||||
user_id: "node-scale-test-user",
|
||||
notice_state: {
|
||||
first_run: {
|
||||
consumed: true,
|
||||
trigger_function: "test_setup",
|
||||
variant: "test",
|
||||
},
|
||||
...(Object.keys(scaleState).length > 0 && {
|
||||
scale_threshold: scaleState,
|
||||
}),
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function scalePayload(overrides: Record<string, any> = {}) {
|
||||
return {
|
||||
notices: {
|
||||
scale_threshold: {
|
||||
enabled: true,
|
||||
notice_type: "log_line",
|
||||
copies: {
|
||||
top_k: SCALE_TOP_K_COPY,
|
||||
memory_count: SCALE_MEMORY_COUNT_COPY,
|
||||
},
|
||||
...overrides,
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function createFetchMock(options: {
|
||||
variant?: string;
|
||||
payload?: unknown;
|
||||
failFlags?: boolean;
|
||||
flagEnabled?: boolean;
|
||||
}) {
|
||||
const calls: any[] = [];
|
||||
const fetchMock = jest.fn(async (url: string | URL, init?: RequestInit) => {
|
||||
const target = String(url);
|
||||
|
||||
if (target.includes("/flags")) {
|
||||
if (options.failFlags) {
|
||||
throw new Error("flag evaluation failed");
|
||||
}
|
||||
const payload =
|
||||
options.payload === undefined
|
||||
? JSON.stringify(scalePayload())
|
||||
: options.payload;
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({
|
||||
flags: {
|
||||
"mem0-oss-notices": {
|
||||
key: "mem0-oss-notices",
|
||||
enabled: options.flagEnabled ?? true,
|
||||
variant: options.variant ?? "displayed",
|
||||
metadata: { payload },
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
if (target.includes("/i/v0/e/")) {
|
||||
calls.push(JSON.parse(String(init?.body)));
|
||||
return {
|
||||
ok: true,
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({}),
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
});
|
||||
|
||||
return { fetchMock, calls };
|
||||
}
|
||||
|
||||
function noticeEvents(calls: any[]) {
|
||||
return calls.filter((call) => call.event === "mem0.notice_displayed");
|
||||
}
|
||||
|
||||
function scaleEvents(calls: any[]) {
|
||||
return noticeEvents(calls).filter(
|
||||
(call) => call.properties.notice_id === "scale_threshold",
|
||||
);
|
||||
}
|
||||
|
||||
function flagRequestCount(fetchMock: jest.Mock) {
|
||||
return fetchMock.mock.calls.filter(([url]) => String(url).includes("/flags"))
|
||||
.length;
|
||||
}
|
||||
|
||||
async function createMemory() {
|
||||
const { Memory } = await import("../src/memory");
|
||||
return new Memory({
|
||||
version: "v1.1",
|
||||
embedder: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "text-embedding-3-small" },
|
||||
},
|
||||
vectorStore: {
|
||||
provider: "memory",
|
||||
config: {
|
||||
collectionName: `test-scale-${Date.now()}-${Math.random()}`,
|
||||
dimension: 1536,
|
||||
dbPath: ":memory:",
|
||||
},
|
||||
},
|
||||
llm: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "gpt-5-mini" },
|
||||
},
|
||||
historyDbPath: ":memory:",
|
||||
});
|
||||
}
|
||||
|
||||
async function addSeed(memory: any, userId = "scale-user") {
|
||||
await memory.add("The user's favorite drink is green tea.", {
|
||||
userId,
|
||||
infer: false,
|
||||
});
|
||||
}
|
||||
|
||||
describe("Node OSS scale threshold notice", () => {
|
||||
let originalMem0Dir: string | undefined;
|
||||
let originalTelemetry: string | undefined;
|
||||
let originalSampleRate: string | undefined;
|
||||
let originalFetch: typeof global.fetch;
|
||||
let stderrSpy: jest.SpyInstance;
|
||||
|
||||
beforeEach(() => {
|
||||
originalMem0Dir = process.env.MEM0_DIR;
|
||||
originalTelemetry = process.env.MEM0_TELEMETRY;
|
||||
originalSampleRate = process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
originalFetch = global.fetch;
|
||||
|
||||
process.env.MEM0_DIR = makeTempMem0Dir();
|
||||
process.env.MEM0_TELEMETRY = "true";
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = "1";
|
||||
stderrSpy = jest
|
||||
.spyOn(process.stderr, "write")
|
||||
.mockImplementation(() => true);
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalMem0Dir === undefined) delete process.env.MEM0_DIR;
|
||||
else process.env.MEM0_DIR = originalMem0Dir;
|
||||
|
||||
if (originalTelemetry === undefined) delete process.env.MEM0_TELEMETRY;
|
||||
else process.env.MEM0_TELEMETRY = originalTelemetry;
|
||||
|
||||
if (originalSampleRate === undefined) {
|
||||
delete process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
} else {
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = originalSampleRate;
|
||||
}
|
||||
|
||||
global.fetch = originalFetch;
|
||||
jest.restoreAllMocks();
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
it("detects high topK after successful search", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await addSeed(memory);
|
||||
stderrSpy.mockClear();
|
||||
const result = await memory.search("favorite drink private text", {
|
||||
filters: { user_id: "scale-user" },
|
||||
topK: 50,
|
||||
});
|
||||
|
||||
expect(result.results.length).toBeGreaterThan(0);
|
||||
expect(stderrSpy).toHaveBeenCalledWith("Scale top 50\n");
|
||||
const event = scaleEvents(calls)[0];
|
||||
expect(event.properties).toMatchObject({
|
||||
notice_id: "scale_threshold",
|
||||
displayed: true,
|
||||
trigger_function: "search",
|
||||
trigger_source: "top_k",
|
||||
trigger_reason: "high_top_k",
|
||||
top_k: 50,
|
||||
threshold: 50,
|
||||
sample_rate: 1,
|
||||
});
|
||||
expect(JSON.stringify(event.properties)).not.toContain("private text");
|
||||
});
|
||||
|
||||
it("detects high topK after successful getAll", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await addSeed(memory);
|
||||
stderrSpy.mockClear();
|
||||
const result = await memory.getAll({
|
||||
filters: { user_id: "scale-user" },
|
||||
topK: 50,
|
||||
});
|
||||
|
||||
expect(result.results.length).toBeGreaterThan(0);
|
||||
expect(stderrSpy).toHaveBeenCalledWith("Scale top 50\n");
|
||||
const event = scaleEvents(calls)[0];
|
||||
expect(event.properties).toMatchObject({
|
||||
notice_id: "scale_threshold",
|
||||
displayed: true,
|
||||
trigger_function: "get_all",
|
||||
trigger_source: "top_k",
|
||||
trigger_reason: "high_top_k",
|
||||
top_k: 50,
|
||||
threshold: 50,
|
||||
});
|
||||
});
|
||||
|
||||
it("does not evaluate topK below the threshold", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await addSeed(memory);
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "scale-user" },
|
||||
topK: 49,
|
||||
});
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(scaleEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.scale_threshold).toBeUndefined();
|
||||
});
|
||||
|
||||
it("marks memory-count threshold evaluated before PostHog display", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ failFlags: true });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
(memory as any).vectorStore.count = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ count: 2000 });
|
||||
|
||||
await memory.add("Scale count threshold fixture.", {
|
||||
userId: "scale-user",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
expect((memory as any).vectorStore.count).toHaveBeenCalledTimes(1);
|
||||
expect(
|
||||
readConfig().notice_state.scale_threshold
|
||||
.memory_count_threshold_evaluated,
|
||||
).toBe(true);
|
||||
expect(scaleEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("does not count provider memories once threshold was evaluated", async () => {
|
||||
consumeFirstRun({ memory_count_threshold_evaluated: true });
|
||||
const { fetchMock } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
(memory as any).vectorStore.count = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ count: 2000 });
|
||||
|
||||
await memory.add("Scale count already evaluated fixture.", {
|
||||
userId: "scale-user",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
expect((memory as any).vectorStore.count).not.toHaveBeenCalled();
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
});
|
||||
|
||||
it("throttles under-threshold provider counts", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
(memory as any).vectorStore.count = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ count: 1999 });
|
||||
|
||||
await memory.add("Scale count under threshold one.", {
|
||||
userId: "scale-user",
|
||||
infer: false,
|
||||
});
|
||||
await memory.add("Scale count under threshold two.", {
|
||||
userId: "scale-user",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
expect((memory as any).vectorStore.count).toHaveBeenCalledTimes(1);
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
});
|
||||
|
||||
it("records holdout and disabled variants silently", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "holdout",
|
||||
payload: JSON.stringify(scalePayload({ enabled: false })),
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await addSeed(memory);
|
||||
stderrSpy.mockClear();
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "scale-user" },
|
||||
topK: 50,
|
||||
});
|
||||
|
||||
expect(stderrSpy).not.toHaveBeenCalled();
|
||||
const event = scaleEvents(calls)[0];
|
||||
expect(event.properties).toMatchObject({
|
||||
displayed: false,
|
||||
bypass_reason: "payload_disabled",
|
||||
disabled_reason: "payload_disabled",
|
||||
variant: "holdout",
|
||||
});
|
||||
});
|
||||
|
||||
it("treats missing enabled as enabled for scale payloads", async () => {
|
||||
consumeFirstRun();
|
||||
const payload = scalePayload();
|
||||
delete (payload.notices.scale_threshold as Record<string, any>).enabled;
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload: JSON.stringify(payload),
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await addSeed(memory);
|
||||
stderrSpy.mockClear();
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "scale-user" },
|
||||
topK: 50,
|
||||
});
|
||||
|
||||
expect(stderrSpy.mock.calls.flat().join("")).toContain("Scale top 50");
|
||||
expect(scaleEvents(calls)[0].properties).toMatchObject({
|
||||
notice_id: "scale_threshold",
|
||||
displayed: true,
|
||||
payload: "Scale top 50",
|
||||
});
|
||||
});
|
||||
|
||||
it("caps evaluated scale opportunities at 10 per week", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await addSeed(memory);
|
||||
stderrSpy.mockClear();
|
||||
for (let index = 0; index < 11; index++) {
|
||||
await memory.search(`favorite drink ${index}`, {
|
||||
filters: { user_id: "scale-user" },
|
||||
topK: 50,
|
||||
});
|
||||
}
|
||||
|
||||
expect(scaleEvents(calls)).toHaveLength(10);
|
||||
expect(flagRequestCount(fetchMock)).toBe(10);
|
||||
expect(stderrSpy).toHaveBeenCalledTimes(10);
|
||||
expect(readConfig().notice_state.scale_threshold.events).toHaveLength(10);
|
||||
});
|
||||
|
||||
it("does not consume first-run on the same qualifying scale call", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.getAll({
|
||||
filters: { user_id: "scale-user" },
|
||||
topK: 50,
|
||||
});
|
||||
|
||||
expect(scaleEvents(calls)).toHaveLength(1);
|
||||
expect(readConfig().notice_state?.first_run).toBeUndefined();
|
||||
});
|
||||
|
||||
it("telemetry off skips flag evaluation, event emission, and state writes", async () => {
|
||||
process.env.MEM0_TELEMETRY = "false";
|
||||
jest.resetModules();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.getAll({
|
||||
filters: { user_id: "scale-user" },
|
||||
topK: 50,
|
||||
});
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(scaleEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state).toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,516 @@
|
||||
/// <reference types="jest" />
|
||||
import * as fs from "fs";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
|
||||
jest.setTimeout(15000);
|
||||
|
||||
jest.mock("../src/embeddings/google", () => ({
|
||||
GoogleEmbedder: jest.fn(),
|
||||
}));
|
||||
jest.mock("../src/llms/google", () => ({
|
||||
GoogleLLM: jest.fn(),
|
||||
}));
|
||||
|
||||
const mockEmbedding = new Array(1536).fill(0.1);
|
||||
const mockEmbed = jest.fn().mockResolvedValue(mockEmbedding);
|
||||
const mockEmbedBatch = jest
|
||||
.fn()
|
||||
.mockImplementation((texts: string[]) =>
|
||||
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||
);
|
||||
|
||||
jest.mock("../src/embeddings/openai", () => ({
|
||||
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
|
||||
embed: mockEmbed,
|
||||
embedBatch: mockEmbedBatch,
|
||||
embeddingDims: 1536,
|
||||
})),
|
||||
}));
|
||||
|
||||
const mockGenerateResponse = jest.fn();
|
||||
jest.mock("../src/llms/openai", () => ({
|
||||
OpenAILLM: jest.fn().mockImplementation(() => ({
|
||||
generateResponse: mockGenerateResponse,
|
||||
})),
|
||||
}));
|
||||
|
||||
const TEMPORAL_COPY =
|
||||
"Temporal reasoning requires a Mem0 API key. Get one for free at https://app.mem0.ai";
|
||||
const PLAIN_TIMESTAMP_ERROR =
|
||||
"The timestamp parameter is not supported by the OSS Memory SDK.";
|
||||
const PLAIN_REFERENCE_DATE_ERROR =
|
||||
"The referenceDate parameter is not supported by the OSS Memory SDK.";
|
||||
|
||||
function makeTempMem0Dir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "mem0-node-temporal-feature-"));
|
||||
}
|
||||
|
||||
function temporalPayload(overrides: Record<string, any> = {}) {
|
||||
return {
|
||||
notices: {
|
||||
temporal_stub: {
|
||||
enabled: true,
|
||||
notice_type: "error",
|
||||
copy: TEMPORAL_COPY,
|
||||
...overrides,
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function createFetchMock(options: {
|
||||
variant?: string;
|
||||
payload?: unknown;
|
||||
failFlags?: boolean;
|
||||
flagEnabled?: boolean;
|
||||
}) {
|
||||
const calls: any[] = [];
|
||||
const fetchMock = jest.fn(async (url: string | URL, init?: RequestInit) => {
|
||||
const target = String(url);
|
||||
|
||||
if (target.includes("/flags")) {
|
||||
if (options.failFlags) {
|
||||
throw new Error("flag evaluation failed");
|
||||
}
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({
|
||||
flags: {
|
||||
"mem0-oss-notices": {
|
||||
key: "mem0-oss-notices",
|
||||
enabled: options.flagEnabled ?? true,
|
||||
variant: options.variant ?? "displayed",
|
||||
metadata: {
|
||||
payload:
|
||||
options.payload === undefined
|
||||
? JSON.stringify(temporalPayload())
|
||||
: options.payload,
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
if (target.includes("/i/v0/e/")) {
|
||||
calls.push(JSON.parse(String(init?.body)));
|
||||
return {
|
||||
ok: true,
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({}),
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
});
|
||||
|
||||
return { fetchMock, calls };
|
||||
}
|
||||
|
||||
function noticeEvents(calls: any[]) {
|
||||
return calls.filter((call) => call.event === "mem0.notice_displayed");
|
||||
}
|
||||
|
||||
async function createMemory() {
|
||||
const { Memory } = await import("../src/memory");
|
||||
return new Memory({
|
||||
version: "v1.1",
|
||||
embedder: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "text-embedding-3-small" },
|
||||
},
|
||||
vectorStore: {
|
||||
provider: "memory",
|
||||
config: {
|
||||
collectionName: `test-temporal-feature-${Date.now()}-${Math.random()}`,
|
||||
dimension: 1536,
|
||||
dbPath: ":memory:",
|
||||
},
|
||||
},
|
||||
llm: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "gpt-5-mini" },
|
||||
},
|
||||
historyDbPath: ":memory:",
|
||||
});
|
||||
}
|
||||
|
||||
describe("Node OSS temporal feature error notice", () => {
|
||||
let originalMem0Dir: string | undefined;
|
||||
let originalTelemetry: string | undefined;
|
||||
let originalSampleRate: string | undefined;
|
||||
let originalFetch: typeof global.fetch;
|
||||
|
||||
beforeEach(() => {
|
||||
originalMem0Dir = process.env.MEM0_DIR;
|
||||
originalTelemetry = process.env.MEM0_TELEMETRY;
|
||||
originalSampleRate = process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
originalFetch = global.fetch;
|
||||
|
||||
process.env.MEM0_DIR = makeTempMem0Dir();
|
||||
process.env.MEM0_TELEMETRY = "true";
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = "1";
|
||||
mockEmbed.mockClear();
|
||||
mockEmbedBatch.mockClear();
|
||||
mockGenerateResponse.mockClear();
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalMem0Dir === undefined) delete process.env.MEM0_DIR;
|
||||
else process.env.MEM0_DIR = originalMem0Dir;
|
||||
|
||||
if (originalTelemetry === undefined) delete process.env.MEM0_TELEMETRY;
|
||||
else process.env.MEM0_TELEMETRY = originalTelemetry;
|
||||
|
||||
if (originalSampleRate === undefined) {
|
||||
delete process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
} else {
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = originalSampleRate;
|
||||
}
|
||||
|
||||
global.fetch = originalFetch;
|
||||
jest.restoreAllMocks();
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
it.each([
|
||||
["add", "timestamp"],
|
||||
["search", "referenceDate"],
|
||||
])(
|
||||
"raises CTA copy for displayed %s(%s) and emits displayed=true",
|
||||
async (triggerFunction, triggerParameter) => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
if (triggerFunction === "add") {
|
||||
await expect(
|
||||
memory.add("Temporal add", {
|
||||
userId: "temporal-user",
|
||||
timestamp: 1778112000,
|
||||
}),
|
||||
).rejects.toThrow(TEMPORAL_COPY);
|
||||
} else {
|
||||
await expect(
|
||||
memory.search("Temporal search", {
|
||||
filters: { user_id: "temporal-user" },
|
||||
referenceDate: "2026-05-06",
|
||||
}),
|
||||
).rejects.toThrow(TEMPORAL_COPY);
|
||||
}
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_stub",
|
||||
notice_type: "error",
|
||||
flag_key: "mem0-oss-notices",
|
||||
variant: "displayed",
|
||||
displayed: true,
|
||||
payload: TEMPORAL_COPY,
|
||||
notice_config_found: true,
|
||||
sync_type: "async",
|
||||
trigger_function: triggerFunction,
|
||||
trigger_parameter: triggerParameter,
|
||||
sample_rate: 1,
|
||||
}),
|
||||
);
|
||||
},
|
||||
);
|
||||
|
||||
it.each([
|
||||
["add", "timestamp"],
|
||||
["search", "referenceDate"],
|
||||
])(
|
||||
"raises CTA copy for holdout %s(%s) and emits displayed=true",
|
||||
async (triggerFunction, triggerParameter) => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "holdout" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
if (triggerFunction === "add") {
|
||||
await expect(
|
||||
memory.add("Temporal add", {
|
||||
userId: "temporal-user",
|
||||
timestamp: null,
|
||||
}),
|
||||
).rejects.toThrow(TEMPORAL_COPY);
|
||||
} else {
|
||||
await expect(
|
||||
memory.search("Temporal search", {
|
||||
filters: { user_id: "temporal-user" },
|
||||
referenceDate: null,
|
||||
}),
|
||||
).rejects.toThrow(TEMPORAL_COPY);
|
||||
}
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_stub",
|
||||
variant: "holdout",
|
||||
displayed: true,
|
||||
trigger_function: triggerFunction,
|
||||
trigger_parameter: triggerParameter,
|
||||
}),
|
||||
);
|
||||
expect(notices[0].properties.bypass_reason).toBeUndefined();
|
||||
},
|
||||
);
|
||||
|
||||
it("uses plain error for unknown future variants and emits not_displayed", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "silent" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(
|
||||
memory.add("Temporal add", {
|
||||
userId: "temporal-user",
|
||||
timestamp: 1778112000,
|
||||
}),
|
||||
).rejects.toThrow(PLAIN_TIMESTAMP_ERROR);
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_stub",
|
||||
variant: "silent",
|
||||
displayed: false,
|
||||
bypass_reason: "not_displayed",
|
||||
payload: TEMPORAL_COPY,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("treats missing enabled as enabled for feature-error payloads", async () => {
|
||||
const payload = temporalPayload();
|
||||
delete (payload.notices.temporal_stub as Record<string, any>).enabled;
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload: JSON.stringify(payload),
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(
|
||||
memory.add("Temporal add", {
|
||||
userId: "temporal-user",
|
||||
timestamp: 1778112000,
|
||||
}),
|
||||
).rejects.toThrow(TEMPORAL_COPY);
|
||||
|
||||
expect(noticeEvents(calls)[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_stub",
|
||||
displayed: true,
|
||||
payload: TEMPORAL_COPY,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("uses timestamp plain error for disabled payload and emits payload_disabled", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload: JSON.stringify(temporalPayload({ enabled: false })),
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(
|
||||
memory.add("Temporal add", {
|
||||
userId: "temporal-user",
|
||||
timestamp: 1778112000,
|
||||
}),
|
||||
).rejects.toThrow(PLAIN_TIMESTAMP_ERROR);
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_stub",
|
||||
displayed: false,
|
||||
bypass_reason: "payload_disabled",
|
||||
disabled_reason: "payload_disabled",
|
||||
notice_config_found: true,
|
||||
trigger_function: "add",
|
||||
trigger_parameter: "timestamp",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it.each([
|
||||
[
|
||||
"missing config",
|
||||
JSON.stringify({ notices: {} }),
|
||||
"missing_notice_config",
|
||||
],
|
||||
[
|
||||
"missing copy",
|
||||
JSON.stringify(temporalPayload({ copy: "" })),
|
||||
"missing_copy",
|
||||
],
|
||||
["malformed payload", "{not-json", "missing_notice_config"],
|
||||
])(
|
||||
"uses referenceDate plain error for %s",
|
||||
async (_label, payload, bypassReason) => {
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(
|
||||
memory.search("Temporal search", {
|
||||
filters: { user_id: "temporal-user" },
|
||||
referenceDate: "2026-05-06",
|
||||
}),
|
||||
).rejects.toThrow(PLAIN_REFERENCE_DATE_ERROR);
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_stub",
|
||||
displayed: false,
|
||||
bypass_reason: bypassReason,
|
||||
trigger_function: "search",
|
||||
trigger_parameter: "referenceDate",
|
||||
}),
|
||||
);
|
||||
},
|
||||
);
|
||||
|
||||
it("uses plain error and emits no event when the blunt flag is disabled", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
flagEnabled: false,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(
|
||||
memory.add("Temporal add", {
|
||||
userId: "temporal-user",
|
||||
timestamp: 1778112000,
|
||||
}),
|
||||
).rejects.toThrow(PLAIN_TIMESTAMP_ERROR);
|
||||
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("uses plain error and emits no event when PostHog fails", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ failFlags: true });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(
|
||||
memory.search("Temporal search", {
|
||||
filters: { user_id: "temporal-user" },
|
||||
referenceDate: "2026-05-06",
|
||||
}),
|
||||
).rejects.toThrow(PLAIN_REFERENCE_DATE_ERROR);
|
||||
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("uses plain error and skips flag evaluation when telemetry is off", async () => {
|
||||
process.env.MEM0_TELEMETRY = "False";
|
||||
jest.resetModules();
|
||||
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(
|
||||
memory.add("Temporal add", {
|
||||
userId: "temporal-user",
|
||||
timestamp: 1778112000,
|
||||
}),
|
||||
).rejects.toThrow(PLAIN_TIMESTAMP_ERROR);
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(noticeEvents(calls)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("throws before add validation, normal telemetry, embeddings, and first-run", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
const addToVectorStoreSpy = jest.spyOn(memory as any, "addToVectorStore");
|
||||
|
||||
await expect(
|
||||
memory.add(
|
||||
undefined as any,
|
||||
{
|
||||
timestamp: 1778112000,
|
||||
} as any,
|
||||
),
|
||||
).rejects.toThrow(TEMPORAL_COPY);
|
||||
|
||||
expect(mockEmbed).not.toHaveBeenCalled();
|
||||
expect(addToVectorStoreSpy).not.toHaveBeenCalled();
|
||||
expect(calls.map((call) => call.event)).not.toContain("mem0.add");
|
||||
expect(
|
||||
noticeEvents(calls).some(
|
||||
(call) => call.properties.notice_id === "first_run",
|
||||
),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it("throws before search validation, normal telemetry, embeddings, and first-run", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(
|
||||
memory.search("Temporal search", {
|
||||
userId: "invalid-top-level-user",
|
||||
referenceDate: "2026-05-06",
|
||||
} as any),
|
||||
).rejects.toThrow(TEMPORAL_COPY);
|
||||
|
||||
expect(mockEmbed).not.toHaveBeenCalled();
|
||||
expect(calls.map((call) => call.event)).not.toContain("mem0.search");
|
||||
expect(
|
||||
noticeEvents(calls).some(
|
||||
(call) => call.properties.notice_id === "first_run",
|
||||
),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it("leaves normal add/search calls unchanged when temporal options are omitted", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "holdout" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
const addResult = await memory.add("Normal add", {
|
||||
userId: "normal-user",
|
||||
infer: false,
|
||||
});
|
||||
expect(addResult.results).toHaveLength(1);
|
||||
|
||||
const searchResult = await memory.search("Normal add", {
|
||||
filters: { user_id: "normal-user" },
|
||||
topK: 3,
|
||||
});
|
||||
expect(searchResult.results.length).toBeGreaterThanOrEqual(0);
|
||||
|
||||
const eventNames = calls.map((call) => call.event);
|
||||
expect(eventNames).toContain("mem0.add");
|
||||
expect(eventNames).toContain("mem0.search");
|
||||
expect(
|
||||
noticeEvents(calls).some(
|
||||
(call) => call.properties.notice_id === "temporal_stub",
|
||||
),
|
||||
).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,594 @@
|
||||
/// <reference types="jest" />
|
||||
import * as fs from "fs";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
|
||||
jest.setTimeout(15000);
|
||||
|
||||
jest.mock("../src/embeddings/google", () => ({
|
||||
GoogleEmbedder: jest.fn(),
|
||||
}));
|
||||
jest.mock("../src/llms/google", () => ({
|
||||
GoogleLLM: jest.fn(),
|
||||
}));
|
||||
|
||||
const mockEmbedding = new Array(1536).fill(0.1);
|
||||
jest.mock("../src/embeddings/openai", () => ({
|
||||
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
|
||||
embed: jest.fn().mockResolvedValue(mockEmbedding),
|
||||
embedBatch: jest
|
||||
.fn()
|
||||
.mockImplementation((texts: string[]) =>
|
||||
Promise.resolve(texts.map(() => mockEmbedding)),
|
||||
),
|
||||
embeddingDims: 1536,
|
||||
})),
|
||||
}));
|
||||
|
||||
jest.mock("../src/llms/openai", () => ({
|
||||
OpenAILLM: jest.fn().mockImplementation(() => ({
|
||||
generateResponse: jest.fn(),
|
||||
})),
|
||||
}));
|
||||
|
||||
const TEMPORAL_USAGE_COPY =
|
||||
"This looks like a time-aware memory workflow. Mem0 Platform has temporal reasoning built in. Use `timestamp` when adding memories and `reference_date` when searching.";
|
||||
|
||||
function makeTempMem0Dir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "mem0-node-temporal-usage-"));
|
||||
}
|
||||
|
||||
function configPath(): string {
|
||||
return path.join(process.env.MEM0_DIR as string, "config.json");
|
||||
}
|
||||
|
||||
function readConfig(): Record<string, any> {
|
||||
const file = configPath();
|
||||
if (!fs.existsSync(file)) return {};
|
||||
return JSON.parse(fs.readFileSync(file, "utf8"));
|
||||
}
|
||||
|
||||
function writeConfig(config: Record<string, any>) {
|
||||
fs.mkdirSync(path.dirname(configPath()), { recursive: true });
|
||||
fs.writeFileSync(configPath(), JSON.stringify(config, null, 4));
|
||||
}
|
||||
|
||||
function consumeFirstRun() {
|
||||
writeConfig({
|
||||
user_id: "node-temporal-usage-test-user",
|
||||
notice_state: {
|
||||
first_run: {
|
||||
consumed: true,
|
||||
trigger_function: "test_setup",
|
||||
variant: "test",
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function temporalUsagePayload(overrides: Record<string, any> = {}) {
|
||||
return {
|
||||
notices: {
|
||||
temporal_usage: {
|
||||
enabled: true,
|
||||
notice_type: "log_line",
|
||||
copy: TEMPORAL_USAGE_COPY,
|
||||
...overrides,
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function createFetchMock(options: {
|
||||
variant?: string;
|
||||
payload?: unknown;
|
||||
failFlags?: boolean;
|
||||
flagEnabled?: boolean;
|
||||
}) {
|
||||
const calls: any[] = [];
|
||||
const fetchMock = jest.fn(async (url: string | URL, init?: RequestInit) => {
|
||||
const target = String(url);
|
||||
|
||||
if (target.includes("/flags")) {
|
||||
if (options.failFlags) {
|
||||
throw new Error("flag evaluation failed");
|
||||
}
|
||||
const payload =
|
||||
options.payload === undefined
|
||||
? JSON.stringify(temporalUsagePayload())
|
||||
: options.payload;
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({
|
||||
flags: {
|
||||
"mem0-oss-notices": {
|
||||
key: "mem0-oss-notices",
|
||||
enabled: options.flagEnabled ?? true,
|
||||
variant: options.variant ?? "displayed",
|
||||
metadata: { payload },
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
if (target.includes("/i/v0/e/")) {
|
||||
calls.push(JSON.parse(String(init?.body)));
|
||||
return {
|
||||
ok: true,
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({}),
|
||||
text: jest.fn().mockResolvedValue(""),
|
||||
};
|
||||
});
|
||||
|
||||
return { fetchMock, calls };
|
||||
}
|
||||
|
||||
function noticeEvents(calls: any[]) {
|
||||
return calls.filter((call) => call.event === "mem0.notice_displayed");
|
||||
}
|
||||
|
||||
function temporalUsageEvents(calls: any[]) {
|
||||
return noticeEvents(calls).filter(
|
||||
(call) => call.properties.notice_id === "temporal_usage",
|
||||
);
|
||||
}
|
||||
|
||||
function flagRequestCount(fetchMock: jest.Mock) {
|
||||
return fetchMock.mock.calls.filter(([url]) => String(url).includes("/flags"))
|
||||
.length;
|
||||
}
|
||||
|
||||
async function createMemory() {
|
||||
const { Memory } = await import("../src/memory");
|
||||
return new Memory({
|
||||
version: "v1.1",
|
||||
embedder: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "text-embedding-3-small" },
|
||||
},
|
||||
vectorStore: {
|
||||
provider: "memory",
|
||||
config: {
|
||||
collectionName: `test-temporal-usage-${Date.now()}-${Math.random()}`,
|
||||
dimension: 1536,
|
||||
dbPath: ":memory:",
|
||||
},
|
||||
},
|
||||
llm: {
|
||||
provider: "openai",
|
||||
config: { apiKey: "test-key", model: "gpt-5-mini" },
|
||||
},
|
||||
historyDbPath: ":memory:",
|
||||
});
|
||||
}
|
||||
|
||||
async function addSearchSeed(memory: any, userId = "temporal-user") {
|
||||
await memory.add("The user's favorite drink is green tea.", {
|
||||
userId,
|
||||
infer: false,
|
||||
});
|
||||
}
|
||||
|
||||
describe("Node OSS temporal usage notice", () => {
|
||||
let originalMem0Dir: string | undefined;
|
||||
let originalTelemetry: string | undefined;
|
||||
let originalSampleRate: string | undefined;
|
||||
let originalFetch: typeof global.fetch;
|
||||
let stderrSpy: jest.SpyInstance;
|
||||
|
||||
beforeEach(() => {
|
||||
originalMem0Dir = process.env.MEM0_DIR;
|
||||
originalTelemetry = process.env.MEM0_TELEMETRY;
|
||||
originalSampleRate = process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
originalFetch = global.fetch;
|
||||
|
||||
process.env.MEM0_DIR = makeTempMem0Dir();
|
||||
process.env.MEM0_TELEMETRY = "true";
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = "1";
|
||||
stderrSpy = jest
|
||||
.spyOn(process.stderr, "write")
|
||||
.mockImplementation(() => true);
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalMem0Dir === undefined) delete process.env.MEM0_DIR;
|
||||
else process.env.MEM0_DIR = originalMem0Dir;
|
||||
|
||||
if (originalTelemetry === undefined) delete process.env.MEM0_TELEMETRY;
|
||||
else process.env.MEM0_TELEMETRY = originalTelemetry;
|
||||
|
||||
if (originalSampleRate === undefined) {
|
||||
delete process.env.MEM0_TELEMETRY_SAMPLE_RATE;
|
||||
} else {
|
||||
process.env.MEM0_TELEMETRY_SAMPLE_RATE = originalSampleRate;
|
||||
}
|
||||
|
||||
global.fetch = originalFetch;
|
||||
jest.restoreAllMocks();
|
||||
jest.resetModules();
|
||||
});
|
||||
|
||||
it("detects timestamp-like metadata after successful add", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
const result = await memory.add("Temporal metadata memory", {
|
||||
userId: "temporal-user",
|
||||
infer: false,
|
||||
metadata: { event_date: "2025-04-09" },
|
||||
});
|
||||
|
||||
expect(result.results).toHaveLength(1);
|
||||
expect(stderrSpy).toHaveBeenCalledWith(`${TEMPORAL_USAGE_COPY}\n`);
|
||||
const notices = temporalUsageEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_usage",
|
||||
notice_type: "log_line",
|
||||
flag_key: "mem0-oss-notices",
|
||||
variant: "displayed",
|
||||
displayed: true,
|
||||
payload: TEMPORAL_USAGE_COPY,
|
||||
notice_config_found: true,
|
||||
sync_type: "async",
|
||||
trigger_function: "add",
|
||||
trigger_source: "metadata",
|
||||
trigger_reason: "date_like_metadata",
|
||||
sample_rate: 1,
|
||||
}),
|
||||
);
|
||||
expect(readConfig().notice_state.temporal_usage.events).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("emits query trigger fields after successful search", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSearchSeed(memory);
|
||||
|
||||
await memory.search("what happened last week?", {
|
||||
filters: { user_id: "temporal-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
const notices = temporalUsageEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_usage",
|
||||
displayed: true,
|
||||
trigger_function: "search",
|
||||
trigger_source: "query",
|
||||
trigger_reason: "relative_phrase",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("detects temporal range filters after successful search", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSearchSeed(memory);
|
||||
|
||||
await memory.search("favorite drink", {
|
||||
filters: {
|
||||
user_id: "temporal-user",
|
||||
AND: [{ created_at: { gte: "2025-04-01" } }],
|
||||
},
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
const notices = temporalUsageEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_usage",
|
||||
displayed: true,
|
||||
trigger_function: "search",
|
||||
trigger_source: "filter",
|
||||
trigger_reason: "date_range_filter",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("does not evaluate for normal non-temporal add/search calls", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Normal memory", {
|
||||
userId: "normal-user",
|
||||
infer: false,
|
||||
metadata: { category: "planning" },
|
||||
});
|
||||
await memory.search("favorite drink", {
|
||||
filters: { user_id: "normal-user" },
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
expect(flagRequestCount(fetchMock)).toBe(0);
|
||||
expect(temporalUsageEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.temporal_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not evaluate when add/search fail", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await expect(
|
||||
memory.add("Invalid temporal metadata add", {
|
||||
infer: false,
|
||||
metadata: { event_date: "2025-04-09" },
|
||||
} as any),
|
||||
).rejects.toThrow("One of the filters");
|
||||
await expect(
|
||||
memory.search("what happened last week?", {
|
||||
filters: {},
|
||||
}),
|
||||
).rejects.toThrow("filters must contain");
|
||||
|
||||
expect(temporalUsageEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.temporal_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("is silent for holdout and emits displayed=false", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "holdout" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Temporal metadata memory", {
|
||||
userId: "temporal-user",
|
||||
infer: false,
|
||||
metadata: { event_date: "2025-04-09" },
|
||||
});
|
||||
|
||||
expect(stderrSpy).not.toHaveBeenCalled();
|
||||
const notices = temporalUsageEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_usage",
|
||||
variant: "holdout",
|
||||
displayed: false,
|
||||
bypass_reason: "holdout",
|
||||
trigger_source: "metadata",
|
||||
trigger_reason: "date_like_metadata",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it.each([
|
||||
[
|
||||
"disabled payload",
|
||||
JSON.stringify(temporalUsagePayload({ enabled: false })),
|
||||
"payload_disabled",
|
||||
],
|
||||
[
|
||||
"missing config",
|
||||
JSON.stringify({ notices: {} }),
|
||||
"missing_notice_config",
|
||||
],
|
||||
[
|
||||
"missing copy",
|
||||
JSON.stringify(temporalUsagePayload({ copy: "" })),
|
||||
"missing_copy",
|
||||
],
|
||||
["malformed payload", "{not-json", "missing_notice_config"],
|
||||
])("is silent and safe for %s", async (_label, payload, bypassReason) => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
payload,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Temporal metadata memory", {
|
||||
userId: "temporal-user",
|
||||
infer: false,
|
||||
metadata: { event_date: "2025-04-09" },
|
||||
});
|
||||
|
||||
expect(stderrSpy).not.toHaveBeenCalled();
|
||||
const notices = temporalUsageEvents(calls);
|
||||
expect(notices).toHaveLength(1);
|
||||
expect(notices[0].properties).toEqual(
|
||||
expect.objectContaining({
|
||||
notice_id: "temporal_usage",
|
||||
displayed: false,
|
||||
bypass_reason: bypassReason,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("does not consume cap or emit when the blunt flag is disabled", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({
|
||||
variant: "displayed",
|
||||
flagEnabled: false,
|
||||
});
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Temporal metadata memory", {
|
||||
userId: "temporal-user",
|
||||
infer: false,
|
||||
metadata: { event_date: "2025-04-09" },
|
||||
});
|
||||
|
||||
expect(temporalUsageEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.temporal_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not consume cap or emit when PostHog fails", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ failFlags: true });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Temporal metadata memory", {
|
||||
userId: "temporal-user",
|
||||
infer: false,
|
||||
metadata: { event_date: "2025-04-09" },
|
||||
});
|
||||
|
||||
expect(temporalUsageEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state?.temporal_usage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does nothing when telemetry is off", async () => {
|
||||
process.env.MEM0_TELEMETRY = "False";
|
||||
jest.resetModules();
|
||||
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Temporal metadata memory", {
|
||||
userId: "temporal-user",
|
||||
infer: false,
|
||||
metadata: { event_date: "2025-04-09" },
|
||||
});
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(temporalUsageEvents(calls)).toHaveLength(0);
|
||||
expect(readConfig().notice_state).toBeUndefined();
|
||||
});
|
||||
|
||||
it("caps evaluated opportunities at 10 per rolling week", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSearchSeed(memory);
|
||||
|
||||
for (let index = 1; index <= 11; index++) {
|
||||
await memory.search(`what happened last week ${index}?`, {
|
||||
filters: { user_id: "temporal-user" },
|
||||
topK: 3,
|
||||
});
|
||||
}
|
||||
|
||||
expect(temporalUsageEvents(calls)).toHaveLength(10);
|
||||
expect(flagRequestCount(fetchMock)).toBe(10);
|
||||
expect(readConfig().notice_state.temporal_usage.events).toHaveLength(10);
|
||||
expect(stderrSpy).toHaveBeenCalledTimes(10);
|
||||
});
|
||||
|
||||
it("does not consume first-run on the same qualifying temporal usage call", async () => {
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
|
||||
await memory.add("Temporal metadata memory", {
|
||||
userId: "temporal-user",
|
||||
infer: false,
|
||||
metadata: { event_date: "2025-04-09" },
|
||||
});
|
||||
|
||||
expect(readConfig().notice_state.first_run).toBeUndefined();
|
||||
await memory.add("Normal follow-up memory", {
|
||||
userId: "temporal-user",
|
||||
infer: false,
|
||||
});
|
||||
|
||||
const notices = noticeEvents(calls);
|
||||
expect(notices.map((call) => call.properties.notice_id)).toEqual([
|
||||
"temporal_usage",
|
||||
"first_run",
|
||||
]);
|
||||
expect(readConfig().notice_state.first_run.consumed).toBe(true);
|
||||
});
|
||||
|
||||
it("does not include raw query, metadata, or filter values in telemetry props", async () => {
|
||||
consumeFirstRun();
|
||||
const { fetchMock, calls } = createFetchMock({ variant: "displayed" });
|
||||
global.fetch = fetchMock as any;
|
||||
const memory = await createMemory();
|
||||
await addSearchSeed(memory);
|
||||
|
||||
await memory.search("private trip since 2025-04-09", {
|
||||
filters: {
|
||||
user_id: "temporal-user",
|
||||
created_at: { gte: "2025-04-09" },
|
||||
},
|
||||
topK: 3,
|
||||
});
|
||||
|
||||
const props = temporalUsageEvents(calls)[0].properties;
|
||||
const serialized = JSON.stringify(props);
|
||||
expect(serialized).not.toContain("private trip");
|
||||
expect(serialized).not.toContain("2025-04-09");
|
||||
expect(serialized).not.toContain("temporal-user");
|
||||
});
|
||||
|
||||
it("detects expected query, metadata, and filter cases conservatively", async () => {
|
||||
const { __noticeTestHooks } = await import("../src/utils/notices");
|
||||
|
||||
expect(
|
||||
__noticeTestHooks.detectTemporalUsageFromSearch("notes from today", null),
|
||||
).toEqual({ triggerSource: "query", triggerReason: "relative_phrase" });
|
||||
expect(
|
||||
__noticeTestHooks.detectTemporalUsageFromSearch(
|
||||
"notes from 2025-04-09",
|
||||
null,
|
||||
),
|
||||
).toEqual({ triggerSource: "query", triggerReason: "date_like_query" });
|
||||
expect(
|
||||
__noticeTestHooks.detectTemporalUsageFromMetadata({
|
||||
event_date: "2025-04-09",
|
||||
}),
|
||||
).toEqual({
|
||||
triggerSource: "metadata",
|
||||
triggerReason: "date_like_metadata",
|
||||
});
|
||||
expect(
|
||||
__noticeTestHooks.detectTemporalUsageFromSearch("favorite drink", {
|
||||
AND: [{ created_at: { gte: "2025-04-01" } }],
|
||||
}),
|
||||
).toEqual({ triggerSource: "filter", triggerReason: "date_range_filter" });
|
||||
expect(
|
||||
__noticeTestHooks.detectTemporalUsageFromSearch("favorite drink", null),
|
||||
).toBeNull();
|
||||
expect(
|
||||
__noticeTestHooks.detectTemporalUsageFromMetadata({
|
||||
category: "planning",
|
||||
}),
|
||||
).toBeNull();
|
||||
});
|
||||
|
||||
it("detectors never raise for cyclic metadata or filters", async () => {
|
||||
const { __noticeTestHooks } = await import("../src/utils/notices");
|
||||
const metadata: Record<string, any> = {};
|
||||
metadata.self = metadata;
|
||||
const filters: Record<string, any> = {};
|
||||
filters.AND = [filters];
|
||||
|
||||
expect(
|
||||
__noticeTestHooks.detectTemporalUsageFromMetadata(metadata),
|
||||
).toBeNull();
|
||||
expect(
|
||||
__noticeTestHooks.detectTemporalUsageFromSearch(
|
||||
"favorite drink",
|
||||
filters,
|
||||
),
|
||||
).toBeNull();
|
||||
});
|
||||
});
|
||||
@@ -61,6 +61,13 @@ describe("telemetry sampling", () => {
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("notice_displayed event fires even at the highest random value", async () => {
|
||||
randomSpy.mockReturnValue(0.999);
|
||||
const { captureClientEvent } = await import("../src/utils/telemetry");
|
||||
await captureClientEvent("notice_displayed", makeInstance());
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("init event payload has sample_rate: 1.0", async () => {
|
||||
randomSpy.mockReturnValue(0.999);
|
||||
const { captureClientEvent } = await import("../src/utils/telemetry");
|
||||
|
||||
@@ -923,6 +923,7 @@ describe("Memory class – backward compat with all providers", () => {
|
||||
|
||||
jest.doMock("../src/utils/telemetry", () => ({
|
||||
captureClientEvent: jest.fn().mockResolvedValue(undefined),
|
||||
isTelemetryEnabled: jest.fn(() => false),
|
||||
}));
|
||||
|
||||
MemoryClass = require("../src/memory").Memory;
|
||||
|
||||
+195
-2
@@ -5,6 +5,7 @@ import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
import warnings
|
||||
from copy import deepcopy
|
||||
@@ -26,6 +27,29 @@ from mem0.memory.base import MemoryBase
|
||||
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 (
|
||||
PERFORMANCE_SLOW_QUERY_THRESHOLD_SECONDS,
|
||||
detect_scale_threshold_from_add_result,
|
||||
detect_scale_threshold_from_top_k,
|
||||
detect_decay_usage_from_delete,
|
||||
detect_decay_usage_from_delete_all,
|
||||
detect_temporal_usage_from_metadata,
|
||||
detect_temporal_usage_from_search,
|
||||
display_decay_usage_notice,
|
||||
display_decay_usage_notice_async,
|
||||
display_first_run_notice,
|
||||
display_first_run_notice_async,
|
||||
display_performance_slow_query_notice,
|
||||
display_performance_slow_query_notice_async,
|
||||
display_scale_threshold_notice,
|
||||
display_scale_threshold_notice_async,
|
||||
display_temporal_usage_notice,
|
||||
display_temporal_usage_notice_async,
|
||||
get_decay_feature_error_message,
|
||||
get_decay_feature_error_message_async,
|
||||
get_temporal_feature_error_message,
|
||||
get_temporal_feature_error_message_async,
|
||||
)
|
||||
from mem0.memory.utils import (
|
||||
extract_json,
|
||||
parse_messages,
|
||||
@@ -349,6 +373,36 @@ def _entity_collection_name(provider: str, collection_name: str) -> str:
|
||||
setup_config()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_PROJECT_UPDATE_UNSUPPORTED_ERROR = "Project updates are not supported by the OSS Memory SDK."
|
||||
|
||||
|
||||
class _OSSProject:
|
||||
def update(
|
||||
self,
|
||||
custom_instructions: Optional[str] = None,
|
||||
custom_categories: Optional[list] = None,
|
||||
retrieval_criteria: Optional[list] = None,
|
||||
multilingual: Optional[bool] = None,
|
||||
decay: Optional[bool] = None,
|
||||
):
|
||||
if decay is True:
|
||||
raise ValueError(get_decay_feature_error_message("sync", "project.update", "decay"))
|
||||
raise ValueError(_PROJECT_UPDATE_UNSUPPORTED_ERROR)
|
||||
|
||||
|
||||
class _AsyncOSSProject:
|
||||
async def update(
|
||||
self,
|
||||
custom_instructions: Optional[str] = None,
|
||||
custom_categories: Optional[list] = None,
|
||||
retrieval_criteria: Optional[list] = None,
|
||||
multilingual: Optional[bool] = None,
|
||||
decay: Optional[bool] = None,
|
||||
):
|
||||
if decay is True:
|
||||
raise ValueError(await get_decay_feature_error_message_async("async", "project.update", "decay"))
|
||||
raise ValueError(_PROJECT_UPDATE_UNSUPPORTED_ERROR)
|
||||
|
||||
|
||||
class Memory(MemoryBase):
|
||||
def __init__(self, config: MemoryConfig = MemoryConfig()):
|
||||
@@ -417,6 +471,10 @@ class Memory(MemoryBase):
|
||||
|
||||
capture_event("mem0.init", self, {"sync_type": "sync"})
|
||||
|
||||
@property
|
||||
def project(self):
|
||||
return _OSSProject()
|
||||
|
||||
@property
|
||||
def entity_store(self):
|
||||
"""Lazily initialize entity store on first use."""
|
||||
@@ -600,6 +658,7 @@ class Memory(MemoryBase):
|
||||
agent_id: Optional[str] = None,
|
||||
run_id: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
timestamp: Optional[Any] = None,
|
||||
infer: bool = True,
|
||||
memory_type: Optional[str] = None,
|
||||
prompt: Optional[str] = None,
|
||||
@@ -617,6 +676,7 @@ class Memory(MemoryBase):
|
||||
agent_id (str, optional): ID of the agent creating the memory. Defaults to None.
|
||||
run_id (str, optional): ID of the run creating the memory. Defaults to None.
|
||||
metadata (dict, optional): Metadata to store with the memory. Defaults to None.
|
||||
timestamp (Any, optional): Platform-only temporal parameter. Not supported in OSS.
|
||||
infer (bool, optional): If True (default), an LLM is used to extract key facts from
|
||||
'messages' and decide whether to add, update, or delete related memories.
|
||||
If False, 'messages' are added as raw memories directly.
|
||||
@@ -639,7 +699,10 @@ class Memory(MemoryBase):
|
||||
LLMError: If LLM operations fail.
|
||||
DatabaseError: If database operations fail.
|
||||
"""
|
||||
if timestamp is not None:
|
||||
raise ValueError(get_temporal_feature_error_message("sync", "add", "timestamp"))
|
||||
|
||||
temporal_usage_notice = detect_temporal_usage_from_metadata(metadata)
|
||||
processed_metadata, effective_filters = _build_filters_and_metadata(
|
||||
user_id=user_id,
|
||||
agent_id=agent_id,
|
||||
@@ -671,6 +734,13 @@ class Memory(MemoryBase):
|
||||
|
||||
if agent_id is not None and memory_type == MemoryType.PROCEDURAL.value:
|
||||
results = self._create_procedural_memory(messages, metadata=processed_metadata, prompt=prompt)
|
||||
scale_threshold_notice = detect_scale_threshold_from_add_result(self, results)
|
||||
if temporal_usage_notice:
|
||||
display_temporal_usage_notice(self, "sync", "add", *temporal_usage_notice)
|
||||
elif scale_threshold_notice:
|
||||
display_scale_threshold_notice(self, "sync", "add", *scale_threshold_notice)
|
||||
else:
|
||||
display_first_run_notice(self, "sync", "add")
|
||||
return results
|
||||
|
||||
if self.config.llm.config.get("enable_vision"):
|
||||
@@ -679,6 +749,13 @@ class Memory(MemoryBase):
|
||||
messages = parse_vision_messages(messages)
|
||||
|
||||
vector_store_result = self._add_to_vector_store(messages, processed_metadata, effective_filters, infer, prompt=prompt)
|
||||
scale_threshold_notice = detect_scale_threshold_from_add_result(self, vector_store_result)
|
||||
if temporal_usage_notice:
|
||||
display_temporal_usage_notice(self, "sync", "add", *temporal_usage_notice)
|
||||
elif scale_threshold_notice:
|
||||
display_scale_threshold_notice(self, "sync", "add", *scale_threshold_notice)
|
||||
else:
|
||||
display_first_run_notice(self, "sync", "add")
|
||||
return {"results": vector_store_result}
|
||||
|
||||
def _add_to_vector_store(self, messages, metadata, filters, infer, prompt=None):
|
||||
@@ -1005,6 +1082,7 @@ class Memory(MemoryBase):
|
||||
capture_event("mem0.get", self, {"memory_id": memory_id, "sync_type": "sync"})
|
||||
memory = self.vector_store.get(vector_id=memory_id)
|
||||
if not memory:
|
||||
display_first_run_notice(self, "sync", "get")
|
||||
return None
|
||||
|
||||
promoted_payload_keys = [
|
||||
@@ -1033,6 +1111,7 @@ class Memory(MemoryBase):
|
||||
if additional_metadata:
|
||||
result_item["metadata"] = additional_metadata
|
||||
|
||||
display_first_run_notice(self, "sync", "get")
|
||||
return result_item
|
||||
|
||||
def get_all(
|
||||
@@ -1088,6 +1167,7 @@ class Memory(MemoryBase):
|
||||
)
|
||||
|
||||
limit = top_k
|
||||
scale_threshold_notice = detect_scale_threshold_from_top_k(top_k)
|
||||
|
||||
keys, encoded_ids = process_telemetry_filters(effective_filters)
|
||||
capture_event(
|
||||
@@ -1096,6 +1176,10 @@ class Memory(MemoryBase):
|
||||
|
||||
all_memories_result = self._get_all_from_vector_store(effective_filters, limit)
|
||||
|
||||
if scale_threshold_notice:
|
||||
display_scale_threshold_notice(self, "sync", "get_all", *scale_threshold_notice)
|
||||
else:
|
||||
display_first_run_notice(self, "sync", "get_all")
|
||||
return {"results": all_memories_result}
|
||||
|
||||
def _get_all_from_vector_store(self, filters, limit):
|
||||
@@ -1154,6 +1238,7 @@ class Memory(MemoryBase):
|
||||
threshold: float = 0.1,
|
||||
rerank: bool = False,
|
||||
explain: bool = False,
|
||||
reference_date: Optional[Any] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
@@ -1185,6 +1270,7 @@ class Memory(MemoryBase):
|
||||
threshold (float, optional): Minimum score for a memory to be included. Defaults to 0.1.
|
||||
rerank (bool, optional): Whether to rerank results. Defaults to False.
|
||||
explain (bool, optional): Whether to include score_details for each result. Defaults to False.
|
||||
reference_date (Any, optional): Platform-only temporal parameter. Not supported in OSS.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the search results under a "results" key.
|
||||
@@ -1194,12 +1280,16 @@ class Memory(MemoryBase):
|
||||
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id,
|
||||
or if threshold/top_k values are invalid.
|
||||
"""
|
||||
if reference_date is not None:
|
||||
raise ValueError(get_temporal_feature_error_message("sync", "search", "reference_date"))
|
||||
|
||||
# Reject top-level entity params - must use filters instead
|
||||
_reject_top_level_entity_params(kwargs, "search")
|
||||
|
||||
# Validate search parameters (before applying defaults)
|
||||
_validate_search_params(threshold=threshold, top_k=top_k)
|
||||
query = _validate_and_trim_search_query(query)
|
||||
temporal_usage_notice = detect_temporal_usage_from_search(query, filters)
|
||||
|
||||
# Validate and trim entity IDs in filters
|
||||
effective_filters = filters.copy() if filters else {}
|
||||
@@ -1222,6 +1312,7 @@ class Memory(MemoryBase):
|
||||
)
|
||||
|
||||
limit = top_k
|
||||
scale_threshold_notice = detect_scale_threshold_from_top_k(top_k)
|
||||
|
||||
# Apply enhanced metadata filtering if advanced operators are detected
|
||||
if self._has_advanced_operators(effective_filters):
|
||||
@@ -1250,7 +1341,9 @@ class Memory(MemoryBase):
|
||||
},
|
||||
)
|
||||
|
||||
search_start = time.perf_counter()
|
||||
original_memories = self._search_vector_store(query, effective_filters, limit, threshold, explain=explain)
|
||||
search_elapsed_seconds = time.perf_counter() - search_start
|
||||
|
||||
# Apply reranking if enabled and reranker is available
|
||||
if rerank and self.reranker and original_memories:
|
||||
@@ -1260,6 +1353,21 @@ class Memory(MemoryBase):
|
||||
except Exception as e:
|
||||
logger.warning(f"Reranking failed, using original results: {e}")
|
||||
|
||||
if temporal_usage_notice:
|
||||
display_temporal_usage_notice(self, "sync", "search", *temporal_usage_notice)
|
||||
elif scale_threshold_notice:
|
||||
display_scale_threshold_notice(self, "sync", "search", *scale_threshold_notice)
|
||||
elif search_elapsed_seconds > PERFORMANCE_SLOW_QUERY_THRESHOLD_SECONDS:
|
||||
display_performance_slow_query_notice(
|
||||
self,
|
||||
"sync",
|
||||
"search",
|
||||
search_elapsed_seconds,
|
||||
top_k,
|
||||
len(original_memories),
|
||||
)
|
||||
else:
|
||||
display_first_run_notice(self, "sync", "search")
|
||||
return {"results": original_memories}
|
||||
|
||||
def _process_metadata_filters(self, metadata_filters: Dict[str, Any]) -> Dict[str, Any]:
|
||||
@@ -1569,6 +1677,7 @@ class Memory(MemoryBase):
|
||||
existing_embeddings = {data: self.embedding_model.embed(data, "update")}
|
||||
|
||||
self._update_memory(memory_id, data, existing_embeddings, metadata)
|
||||
display_first_run_notice(self, "sync", "update")
|
||||
return {"message": "Memory updated successfully!"}
|
||||
|
||||
def delete(self, memory_id):
|
||||
@@ -1585,6 +1694,11 @@ class Memory(MemoryBase):
|
||||
raise ValueError(f"Memory with id {memory_id} not found")
|
||||
|
||||
self._delete_memory(memory_id, existing_memory)
|
||||
decay_usage_notice = detect_decay_usage_from_delete()
|
||||
if decay_usage_notice:
|
||||
display_decay_usage_notice(self, "sync", "delete", *decay_usage_notice)
|
||||
else:
|
||||
display_first_run_notice(self, "sync", "delete")
|
||||
return {"message": "Memory deleted successfully!"}
|
||||
|
||||
def delete_all(self, user_id: Optional[str] = None, agent_id: Optional[str] = None, run_id: Optional[str] = None):
|
||||
@@ -1618,6 +1732,11 @@ class Memory(MemoryBase):
|
||||
|
||||
logger.info(f"Deleted {len(memories)} memories")
|
||||
|
||||
decay_usage_notice = detect_decay_usage_from_delete_all(len(memories))
|
||||
if decay_usage_notice:
|
||||
display_decay_usage_notice(self, "sync", "delete_all", *decay_usage_notice)
|
||||
else:
|
||||
display_first_run_notice(self, "sync", "delete_all")
|
||||
return {"message": "Memories deleted successfully!"}
|
||||
|
||||
def history(self, memory_id):
|
||||
@@ -1631,7 +1750,9 @@ class Memory(MemoryBase):
|
||||
list: List of changes for the memory.
|
||||
"""
|
||||
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "sync"})
|
||||
return self.db.get_history(memory_id)
|
||||
history = self.db.get_history(memory_id)
|
||||
display_first_run_notice(self, "sync", "history")
|
||||
return history
|
||||
|
||||
def _create_memory(self, data, existing_embeddings, metadata=None):
|
||||
logger.debug(f"Creating memory with {data=}")
|
||||
@@ -1831,6 +1952,7 @@ class Memory(MemoryBase):
|
||||
self._entity_store = None
|
||||
|
||||
capture_event("mem0.reset", self, {"sync_type": "sync"})
|
||||
display_first_run_notice(self, "sync", "reset")
|
||||
|
||||
def close(self):
|
||||
"""Release resources held by this Memory instance (SQLite connections, etc.)."""
|
||||
@@ -1889,6 +2011,10 @@ class AsyncMemory(MemoryBase):
|
||||
|
||||
capture_event("mem0.init", self, {"sync_type": "async"})
|
||||
|
||||
@property
|
||||
def project(self):
|
||||
return _AsyncOSSProject()
|
||||
|
||||
@property
|
||||
def entity_store(self):
|
||||
"""Lazily initialize entity store on first use."""
|
||||
@@ -2058,6 +2184,7 @@ class AsyncMemory(MemoryBase):
|
||||
agent_id: Optional[str] = None,
|
||||
run_id: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
timestamp: Optional[Any] = None,
|
||||
infer: bool = True,
|
||||
memory_type: Optional[str] = None,
|
||||
prompt: Optional[str] = None,
|
||||
@@ -2072,6 +2199,7 @@ class AsyncMemory(MemoryBase):
|
||||
agent_id (str, optional): ID of the agent creating the memory. Defaults to None.
|
||||
run_id (str, optional): ID of the run creating the memory. Defaults to None.
|
||||
metadata (dict, optional): Metadata to store with the memory. Defaults to None.
|
||||
timestamp (Any, optional): Platform-only temporal parameter. Not supported in OSS.
|
||||
infer (bool, optional): Whether to infer the memories. Defaults to True.
|
||||
memory_type (str, optional): Type of memory to create. Defaults to None.
|
||||
Pass "procedural_memory" to create procedural memories.
|
||||
@@ -2080,6 +2208,10 @@ class AsyncMemory(MemoryBase):
|
||||
Returns:
|
||||
dict: A dictionary containing the result of the memory addition operation.
|
||||
"""
|
||||
if timestamp is not None:
|
||||
raise ValueError(await get_temporal_feature_error_message_async("async", "add", "timestamp"))
|
||||
|
||||
temporal_usage_notice = detect_temporal_usage_from_metadata(metadata)
|
||||
processed_metadata, effective_filters = _build_filters_and_metadata(
|
||||
user_id=user_id, agent_id=agent_id, run_id=run_id, input_metadata=metadata
|
||||
)
|
||||
@@ -2107,6 +2239,13 @@ class AsyncMemory(MemoryBase):
|
||||
results = await self._create_procedural_memory(
|
||||
messages, metadata=processed_metadata, prompt=prompt, llm=llm
|
||||
)
|
||||
scale_threshold_notice = await asyncio.to_thread(detect_scale_threshold_from_add_result, self, results)
|
||||
if temporal_usage_notice:
|
||||
await display_temporal_usage_notice_async(self, "async", "add", *temporal_usage_notice)
|
||||
elif scale_threshold_notice:
|
||||
await display_scale_threshold_notice_async(self, "async", "add", *scale_threshold_notice)
|
||||
else:
|
||||
await display_first_run_notice_async(self, "async", "add")
|
||||
return results
|
||||
|
||||
if self.config.llm.config.get("enable_vision"):
|
||||
@@ -2115,6 +2254,13 @@ class AsyncMemory(MemoryBase):
|
||||
messages = parse_vision_messages(messages)
|
||||
|
||||
vector_store_result = await self._add_to_vector_store(messages, processed_metadata, effective_filters, infer, prompt=prompt)
|
||||
scale_threshold_notice = await asyncio.to_thread(detect_scale_threshold_from_add_result, self, vector_store_result)
|
||||
if temporal_usage_notice:
|
||||
await display_temporal_usage_notice_async(self, "async", "add", *temporal_usage_notice)
|
||||
elif scale_threshold_notice:
|
||||
await display_scale_threshold_notice_async(self, "async", "add", *scale_threshold_notice)
|
||||
else:
|
||||
await display_first_run_notice_async(self, "async", "add")
|
||||
return {"results": vector_store_result}
|
||||
|
||||
async def _add_to_vector_store(
|
||||
@@ -2448,6 +2594,7 @@ class AsyncMemory(MemoryBase):
|
||||
capture_event("mem0.get", self, {"memory_id": memory_id, "sync_type": "async"})
|
||||
memory = await asyncio.to_thread(self.vector_store.get, vector_id=memory_id)
|
||||
if not memory:
|
||||
await display_first_run_notice_async(self, "async", "get")
|
||||
return None
|
||||
|
||||
promoted_payload_keys = [
|
||||
@@ -2476,6 +2623,7 @@ class AsyncMemory(MemoryBase):
|
||||
if additional_metadata:
|
||||
result_item["metadata"] = additional_metadata
|
||||
|
||||
await display_first_run_notice_async(self, "async", "get")
|
||||
return result_item
|
||||
|
||||
async def get_all(
|
||||
@@ -2531,6 +2679,7 @@ class AsyncMemory(MemoryBase):
|
||||
)
|
||||
|
||||
limit = top_k
|
||||
scale_threshold_notice = detect_scale_threshold_from_top_k(top_k)
|
||||
|
||||
keys, encoded_ids = process_telemetry_filters(effective_filters)
|
||||
capture_event(
|
||||
@@ -2539,6 +2688,10 @@ class AsyncMemory(MemoryBase):
|
||||
|
||||
all_memories_result = await self._get_all_from_vector_store(effective_filters, limit)
|
||||
|
||||
if scale_threshold_notice:
|
||||
await display_scale_threshold_notice_async(self, "async", "get_all", *scale_threshold_notice)
|
||||
else:
|
||||
await display_first_run_notice_async(self, "async", "get_all")
|
||||
return {"results": all_memories_result}
|
||||
|
||||
async def _get_all_from_vector_store(self, filters, limit):
|
||||
@@ -2597,6 +2750,7 @@ class AsyncMemory(MemoryBase):
|
||||
threshold: float = 0.1,
|
||||
rerank: bool = False,
|
||||
explain: bool = False,
|
||||
reference_date: Optional[Any] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
@@ -2628,6 +2782,7 @@ class AsyncMemory(MemoryBase):
|
||||
threshold (float, optional): Minimum score for a memory to be included. Defaults to 0.1.
|
||||
rerank (bool, optional): Whether to rerank results. Defaults to False.
|
||||
explain (bool, optional): Whether to include score_details for each result. Defaults to False.
|
||||
reference_date (Any, optional): Platform-only temporal parameter. Not supported in OSS.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the search results under a "results" key.
|
||||
@@ -2637,12 +2792,18 @@ class AsyncMemory(MemoryBase):
|
||||
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id,
|
||||
or if threshold/top_k values are invalid.
|
||||
"""
|
||||
if reference_date is not None:
|
||||
raise ValueError(
|
||||
await get_temporal_feature_error_message_async("async", "search", "reference_date")
|
||||
)
|
||||
|
||||
# Reject top-level entity params - must use filters instead
|
||||
_reject_top_level_entity_params(kwargs, "search")
|
||||
|
||||
# Validate search parameters (before applying defaults)
|
||||
_validate_search_params(threshold=threshold, top_k=top_k)
|
||||
query = _validate_and_trim_search_query(query)
|
||||
temporal_usage_notice = detect_temporal_usage_from_search(query, filters)
|
||||
|
||||
# Validate and trim entity IDs in filters
|
||||
effective_filters = filters.copy() if filters else {}
|
||||
@@ -2667,6 +2828,7 @@ class AsyncMemory(MemoryBase):
|
||||
)
|
||||
|
||||
limit = top_k
|
||||
scale_threshold_notice = detect_scale_threshold_from_top_k(top_k)
|
||||
|
||||
# Apply enhanced metadata filtering if advanced operators are detected
|
||||
if self._has_advanced_operators(effective_filters):
|
||||
@@ -2695,7 +2857,9 @@ class AsyncMemory(MemoryBase):
|
||||
},
|
||||
)
|
||||
|
||||
search_start = time.perf_counter()
|
||||
original_memories = await self._search_vector_store(query, effective_filters, limit, threshold, explain=explain)
|
||||
search_elapsed_seconds = time.perf_counter() - search_start
|
||||
|
||||
# Apply reranking if enabled and reranker is available
|
||||
if rerank and self.reranker and original_memories:
|
||||
@@ -2708,6 +2872,21 @@ class AsyncMemory(MemoryBase):
|
||||
except Exception as e:
|
||||
logger.warning(f"Reranking failed, using original results: {e}")
|
||||
|
||||
if temporal_usage_notice:
|
||||
await display_temporal_usage_notice_async(self, "async", "search", *temporal_usage_notice)
|
||||
elif scale_threshold_notice:
|
||||
await display_scale_threshold_notice_async(self, "async", "search", *scale_threshold_notice)
|
||||
elif search_elapsed_seconds > PERFORMANCE_SLOW_QUERY_THRESHOLD_SECONDS:
|
||||
await display_performance_slow_query_notice_async(
|
||||
self,
|
||||
"async",
|
||||
"search",
|
||||
search_elapsed_seconds,
|
||||
top_k,
|
||||
len(original_memories),
|
||||
)
|
||||
else:
|
||||
await display_first_run_notice_async(self, "async", "search")
|
||||
return {"results": original_memories}
|
||||
|
||||
def _process_metadata_filters(self, metadata_filters: Dict[str, Any]) -> Dict[str, Any]:
|
||||
@@ -3008,6 +3187,7 @@ class AsyncMemory(MemoryBase):
|
||||
existing_embeddings = {data: embeddings}
|
||||
|
||||
await self._update_memory(memory_id, data, existing_embeddings, metadata)
|
||||
await display_first_run_notice_async(self, "async", "update")
|
||||
return {"message": "Memory updated successfully!"}
|
||||
|
||||
async def delete(self, memory_id):
|
||||
@@ -3024,6 +3204,11 @@ class AsyncMemory(MemoryBase):
|
||||
raise ValueError(f"Memory with id {memory_id} not found")
|
||||
|
||||
await self._delete_memory(memory_id, existing_memory)
|
||||
decay_usage_notice = detect_decay_usage_from_delete()
|
||||
if decay_usage_notice:
|
||||
await display_decay_usage_notice_async(self, "async", "delete", *decay_usage_notice)
|
||||
else:
|
||||
await display_first_run_notice_async(self, "async", "delete")
|
||||
return {"message": "Memory deleted successfully!"}
|
||||
|
||||
async def delete_all(self, user_id=None, agent_id=None, run_id=None):
|
||||
@@ -3060,6 +3245,11 @@ class AsyncMemory(MemoryBase):
|
||||
|
||||
logger.info(f"Deleted {len(memories[0])} memories")
|
||||
|
||||
decay_usage_notice = detect_decay_usage_from_delete_all(len(memories[0]))
|
||||
if decay_usage_notice:
|
||||
await display_decay_usage_notice_async(self, "async", "delete_all", *decay_usage_notice)
|
||||
else:
|
||||
await display_first_run_notice_async(self, "async", "delete_all")
|
||||
return {"message": "Memories deleted successfully!"}
|
||||
|
||||
async def history(self, memory_id):
|
||||
@@ -3073,7 +3263,9 @@ class AsyncMemory(MemoryBase):
|
||||
list: List of changes for the memory.
|
||||
"""
|
||||
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "async"})
|
||||
return await asyncio.to_thread(self.db.get_history, memory_id)
|
||||
history = await asyncio.to_thread(self.db.get_history, memory_id)
|
||||
await display_first_run_notice_async(self, "async", "history")
|
||||
return history
|
||||
|
||||
async def _create_memory(self, data, existing_embeddings, metadata=None):
|
||||
logger.debug(f"Creating memory with {data=}")
|
||||
@@ -3291,6 +3483,7 @@ class AsyncMemory(MemoryBase):
|
||||
)
|
||||
|
||||
capture_event("mem0.reset", self, {"sync_type": "async"})
|
||||
await display_first_run_notice_async(self, "async", "reset")
|
||||
|
||||
def close(self):
|
||||
"""Release resources held by this AsyncMemory instance."""
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+13
-1
@@ -1,6 +1,7 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import uuid
|
||||
from hashlib import sha256
|
||||
|
||||
@@ -34,10 +35,21 @@ def _load_config():
|
||||
def _write_config(config):
|
||||
"""Best-effort write of ~/.mem0/config.json. Never raises."""
|
||||
path = _config_path()
|
||||
temp_path = None
|
||||
try:
|
||||
with open(path, "w") as f:
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with tempfile.NamedTemporaryFile("w", dir=os.path.dirname(path), delete=False) as f:
|
||||
temp_path = f.name
|
||||
json.dump(config, f, indent=4)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(temp_path, path)
|
||||
except Exception as e:
|
||||
if temp_path:
|
||||
try:
|
||||
os.unlink(temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
_logger.debug("Failed to write mem0 config %s: %s", path, e)
|
||||
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ from mem0.memory.setup import get_or_create_user_id
|
||||
MEM0_TELEMETRY = os.environ.get("MEM0_TELEMETRY", "True")
|
||||
PROJECT_API_KEY = "phc_hgJkUVJFYtmaJqrvf6CYN67TIQ8yhXAkWzUn9AMU4yX"
|
||||
HOST = "https://us.i.posthog.com"
|
||||
FEATURE_FLAGS_REQUEST_TIMEOUT_SECONDS = 0.5
|
||||
|
||||
if isinstance(MEM0_TELEMETRY, str):
|
||||
MEM0_TELEMETRY = MEM0_TELEMETRY.lower() in ("true", "1", "yes")
|
||||
@@ -49,7 +50,9 @@ MEM0_TELEMETRY_SAMPLE_RATE = _parse_sample_rate(os.environ.get("MEM0_TELEMETRY_S
|
||||
# Events that bypass sampling and always fire. Keep this set in sync with the
|
||||
# event names passed to capture_event() in mem0/memory/main.py.
|
||||
# $identify is included so PostHog person-merging is never lost to sampling.
|
||||
_LIFECYCLE_EVENTS = frozenset({"mem0.init", "mem0.reset", "mem0._create_procedural_memory", "$identify"})
|
||||
_LIFECYCLE_EVENTS = frozenset(
|
||||
{"mem0.init", "mem0.reset", "mem0._create_procedural_memory", "mem0.notice_displayed", "$identify"}
|
||||
)
|
||||
|
||||
|
||||
def _sampling_before_send(msg):
|
||||
@@ -77,15 +80,15 @@ class AnonymousTelemetry:
|
||||
self.user_id = None
|
||||
return
|
||||
|
||||
try:
|
||||
self.posthog = Posthog(project_api_key=PROJECT_API_KEY, host=HOST, before_send=before_send)
|
||||
except TypeError:
|
||||
# posthog <4.5.0 does not accept before_send; fall back without sampling.
|
||||
_logger.debug("posthog.Posthog does not accept before_send; upgrade to >=4.5.0 for sampling")
|
||||
self.posthog = Posthog(project_api_key=PROJECT_API_KEY, host=HOST)
|
||||
self.posthog = Posthog(
|
||||
project_api_key=PROJECT_API_KEY,
|
||||
host=HOST,
|
||||
before_send=before_send,
|
||||
feature_flags_request_timeout_seconds=FEATURE_FLAGS_REQUEST_TIMEOUT_SECONDS,
|
||||
)
|
||||
self.user_id = get_or_create_user_id(vector_store)
|
||||
|
||||
def capture_event(self, event_name, properties=None, user_email=None):
|
||||
def capture_event(self, event_name, properties=None, user_email=None, flags=None):
|
||||
if self.posthog is None:
|
||||
return
|
||||
|
||||
@@ -109,7 +112,10 @@ class AnonymousTelemetry:
|
||||
**properties,
|
||||
}
|
||||
try:
|
||||
self.posthog.capture(distinct_id=distinct_id, event=event_name, properties=properties)
|
||||
capture_kwargs = {"distinct_id": distinct_id, "properties": properties}
|
||||
if flags is not None:
|
||||
capture_kwargs["flags"] = flags
|
||||
self.posthog.capture(event_name, **capture_kwargs)
|
||||
except Exception as e:
|
||||
_logger.debug("Failed to capture telemetry event %r: %s", event_name, e)
|
||||
|
||||
|
||||
+2
-2
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "mem0ai"
|
||||
version = "2.0.5"
|
||||
version = "2.0.6"
|
||||
description = "Long-term memory for AI Agents"
|
||||
authors = [
|
||||
{ name = "Mem0", email = "support@mem0.ai" }
|
||||
@@ -17,7 +17,7 @@ dependencies = [
|
||||
"qdrant-client>=1.12.0",
|
||||
"pydantic>=2.7.3",
|
||||
"openai>=1.90.0",
|
||||
"posthog>=4.5.0",
|
||||
"posthog>=7.14.0",
|
||||
"pytz>=2024.1",
|
||||
"sqlalchemy>=2.0.31",
|
||||
"protobuf>=5.29.6,<7.0.0",
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
import pytest
|
||||
|
||||
from mem0.memory import main as memory_main
|
||||
from mem0.memory.main import AsyncMemory, Memory
|
||||
|
||||
|
||||
def test_sync_project_update_decay_true_raises_with_notice_message(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def get_error(sync_type, trigger_function, trigger_parameter):
|
||||
calls.append((sync_type, trigger_function, trigger_parameter))
|
||||
return "blocked decay"
|
||||
|
||||
monkeypatch.setattr(memory_main, "get_decay_feature_error_message", get_error)
|
||||
|
||||
with pytest.raises(ValueError, match="blocked decay"):
|
||||
Memory.__new__(Memory).project.update(decay=True)
|
||||
|
||||
assert calls == [("sync", "project.update", "decay")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_project_update_decay_true_raises_with_notice_message(monkeypatch):
|
||||
calls = []
|
||||
|
||||
async def get_error(sync_type, trigger_function, trigger_parameter):
|
||||
calls.append((sync_type, trigger_function, trigger_parameter))
|
||||
return "blocked async decay"
|
||||
|
||||
monkeypatch.setattr(memory_main, "get_decay_feature_error_message_async", get_error)
|
||||
|
||||
with pytest.raises(ValueError, match="blocked async decay"):
|
||||
await AsyncMemory.__new__(AsyncMemory).project.update(decay=True)
|
||||
|
||||
assert calls == [("async", "project.update", "decay")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kwargs", [{}, {"decay": False}])
|
||||
def test_sync_project_update_non_trigger_raises_plain_error_without_notice(monkeypatch, kwargs):
|
||||
monkeypatch.setattr(
|
||||
memory_main,
|
||||
"get_decay_feature_error_message",
|
||||
lambda *args: pytest.fail("decay feature notice should not run"),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Project updates are not supported by the OSS Memory SDK."):
|
||||
Memory.__new__(Memory).project.update(**kwargs)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kwargs", [{}, {"decay": False}])
|
||||
async def test_async_project_update_non_trigger_raises_plain_error_without_notice(monkeypatch, kwargs):
|
||||
monkeypatch.setattr(
|
||||
memory_main,
|
||||
"get_decay_feature_error_message_async",
|
||||
lambda *args: pytest.fail("decay feature notice should not run"),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Project updates are not supported by the OSS Memory SDK."):
|
||||
await AsyncMemory.__new__(AsyncMemory).project.update(**kwargs)
|
||||
@@ -0,0 +1,214 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from mem0.memory import main as memory_main
|
||||
from mem0.memory.main import AsyncMemory, Memory
|
||||
|
||||
|
||||
def make_sync_memory():
|
||||
memory = Memory.__new__(Memory)
|
||||
memory.vector_store = MagicMock()
|
||||
memory._delete_memory = MagicMock()
|
||||
return memory
|
||||
|
||||
|
||||
def make_async_memory():
|
||||
memory = AsyncMemory.__new__(AsyncMemory)
|
||||
memory.vector_store = MagicMock()
|
||||
memory._delete_memory = AsyncMock()
|
||||
return memory
|
||||
|
||||
|
||||
def test_sync_delete_decay_usage_runs_after_success(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
existing_memory = SimpleNamespace(id="memory-1")
|
||||
memory.vector_store.get.return_value = existing_memory
|
||||
decay_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(
|
||||
memory_main,
|
||||
"detect_decay_usage_from_delete",
|
||||
MagicMock(return_value=("delete_count", "repeated_deletes", 5, None)),
|
||||
)
|
||||
monkeypatch.setattr(memory_main, "display_decay_usage_notice", decay_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
result = Memory.delete(memory, "memory-1")
|
||||
|
||||
assert result == {"message": "Memory deleted successfully!"}
|
||||
memory._delete_memory.assert_called_once_with("memory-1", existing_memory)
|
||||
decay_notice.assert_called_once_with(
|
||||
memory,
|
||||
"sync",
|
||||
"delete",
|
||||
"delete_count",
|
||||
"repeated_deletes",
|
||||
5,
|
||||
None,
|
||||
)
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_delete_below_threshold_uses_first_run_notice(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
memory.vector_store.get.return_value = SimpleNamespace(id="memory-1")
|
||||
decay_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "detect_decay_usage_from_delete", MagicMock(return_value=None))
|
||||
monkeypatch.setattr(memory_main, "display_decay_usage_notice", decay_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
Memory.delete(memory, "memory-1")
|
||||
|
||||
decay_notice.assert_not_called()
|
||||
first_run_notice.assert_called_once_with(memory, "sync", "delete")
|
||||
|
||||
|
||||
def test_sync_delete_failure_does_not_trigger_decay_usage_notice(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
memory.vector_store.get.return_value = None
|
||||
detect_decay = MagicMock()
|
||||
decay_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "detect_decay_usage_from_delete", detect_decay)
|
||||
monkeypatch.setattr(memory_main, "display_decay_usage_notice", decay_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
with pytest.raises(ValueError, match="Memory with id memory-1 not found"):
|
||||
Memory.delete(memory, "memory-1")
|
||||
|
||||
detect_decay.assert_not_called()
|
||||
decay_notice.assert_not_called()
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_delete_all_decay_usage_runs_after_success(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
memories = [SimpleNamespace(id="memory-1"), SimpleNamespace(id="memory-2")]
|
||||
memory.vector_store.list.return_value = (memories, None)
|
||||
decay_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
detect_decay = MagicMock(return_value=("delete_all", "bulk_delete", None, 2))
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "detect_decay_usage_from_delete_all", detect_decay)
|
||||
monkeypatch.setattr(memory_main, "display_decay_usage_notice", decay_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
result = Memory.delete_all(memory, user_id="u1")
|
||||
|
||||
assert result == {"message": "Memories deleted successfully!"}
|
||||
assert memory._delete_memory.call_count == 2
|
||||
detect_decay.assert_called_once_with(2)
|
||||
decay_notice.assert_called_once_with(
|
||||
memory,
|
||||
"sync",
|
||||
"delete_all",
|
||||
"delete_all",
|
||||
"bulk_delete",
|
||||
None,
|
||||
2,
|
||||
)
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_delete_all_zero_deletes_uses_first_run_notice(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
memory.vector_store.list.return_value = ([], None)
|
||||
decay_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "detect_decay_usage_from_delete_all", MagicMock(return_value=None))
|
||||
monkeypatch.setattr(memory_main, "display_decay_usage_notice", decay_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
Memory.delete_all(memory, user_id="u1")
|
||||
|
||||
decay_notice.assert_not_called()
|
||||
first_run_notice.assert_called_once_with(memory, "sync", "delete_all")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_delete_decay_usage_runs_after_success(monkeypatch):
|
||||
memory = make_async_memory()
|
||||
existing_memory = SimpleNamespace(id="memory-1")
|
||||
memory.vector_store.get.return_value = existing_memory
|
||||
decay_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(
|
||||
memory_main,
|
||||
"detect_decay_usage_from_delete",
|
||||
MagicMock(return_value=("delete_count", "repeated_deletes", 5, None)),
|
||||
)
|
||||
monkeypatch.setattr(memory_main, "display_decay_usage_notice_async", decay_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
result = await AsyncMemory.delete(memory, "memory-1")
|
||||
|
||||
assert result == {"message": "Memory deleted successfully!"}
|
||||
memory._delete_memory.assert_awaited_once_with("memory-1", existing_memory)
|
||||
decay_notice.assert_awaited_once_with(
|
||||
memory,
|
||||
"async",
|
||||
"delete",
|
||||
"delete_count",
|
||||
"repeated_deletes",
|
||||
5,
|
||||
None,
|
||||
)
|
||||
first_run_notice.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_delete_failure_does_not_trigger_decay_usage_notice(monkeypatch):
|
||||
memory = make_async_memory()
|
||||
memory.vector_store.get.return_value = None
|
||||
detect_decay = MagicMock()
|
||||
decay_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "detect_decay_usage_from_delete", detect_decay)
|
||||
monkeypatch.setattr(memory_main, "display_decay_usage_notice_async", decay_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
with pytest.raises(ValueError, match="Memory with id memory-1 not found"):
|
||||
await AsyncMemory.delete(memory, "memory-1")
|
||||
|
||||
detect_decay.assert_not_called()
|
||||
decay_notice.assert_not_awaited()
|
||||
first_run_notice.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_delete_all_decay_usage_runs_after_success(monkeypatch):
|
||||
memory = make_async_memory()
|
||||
memories = [SimpleNamespace(id="memory-1"), SimpleNamespace(id="memory-2")]
|
||||
memory.vector_store.list.return_value = (memories, None)
|
||||
decay_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
detect_decay = MagicMock(return_value=("delete_all", "bulk_delete", None, 2))
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "detect_decay_usage_from_delete_all", detect_decay)
|
||||
monkeypatch.setattr(memory_main, "display_decay_usage_notice_async", decay_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
result = await AsyncMemory.delete_all(memory, user_id="u1")
|
||||
|
||||
assert result == {"message": "Memories deleted successfully!"}
|
||||
assert memory._delete_memory.await_count == 2
|
||||
detect_decay.assert_called_once_with(2)
|
||||
decay_notice.assert_awaited_once_with(
|
||||
memory,
|
||||
"async",
|
||||
"delete_all",
|
||||
"delete_all",
|
||||
"bulk_delete",
|
||||
None,
|
||||
2,
|
||||
)
|
||||
first_run_notice.assert_not_awaited()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,445 @@
|
||||
from copy import deepcopy
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from mem0.memory import notices
|
||||
from mem0.memory import main as memory_main
|
||||
from mem0.memory.main import AsyncMemory, Memory
|
||||
|
||||
|
||||
class FakeFlags:
|
||||
def __init__(self, variant, payload):
|
||||
self.variant = variant
|
||||
self.payload = payload
|
||||
|
||||
def get_flag(self, key):
|
||||
assert key == notices.FLAG_KEY
|
||||
return self.variant
|
||||
|
||||
def get_flag_payload(self, key):
|
||||
assert key == notices.FLAG_KEY
|
||||
return self.payload
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_notice_process_state():
|
||||
notices._performance_slow_query_capacity_reached_in_process = False
|
||||
yield
|
||||
notices._performance_slow_query_capacity_reached_in_process = False
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def notice_harness(monkeypatch):
|
||||
config = {}
|
||||
telemetry = MagicMock()
|
||||
telemetry.user_id = "oss-user"
|
||||
|
||||
def write_config(updated):
|
||||
saved = deepcopy(updated)
|
||||
config.clear()
|
||||
config.update(saved)
|
||||
|
||||
monkeypatch.setattr(notices, "_load_config", lambda: config)
|
||||
monkeypatch.setattr(notices, "_write_config", write_config)
|
||||
monkeypatch.setattr(notices.telemetry_module, "MEM0_TELEMETRY", True)
|
||||
monkeypatch.setattr(notices.telemetry_module, "_get_oss_telemetry", lambda: telemetry)
|
||||
return config, telemetry
|
||||
|
||||
|
||||
def configure_flag(telemetry, variant, payload):
|
||||
flags = FakeFlags(variant, payload)
|
||||
telemetry.posthog.evaluate_flags.return_value = flags
|
||||
return flags
|
||||
|
||||
|
||||
def performance_payload(copy="Performance CTA", enabled=True, notice_type="log_line"):
|
||||
payload = {
|
||||
"notices": {
|
||||
"performance_slow_query": {
|
||||
"enabled": enabled,
|
||||
"notice_type": notice_type,
|
||||
}
|
||||
}
|
||||
}
|
||||
if copy is not None:
|
||||
payload["notices"]["performance_slow_query"]["copy"] = copy
|
||||
return payload
|
||||
|
||||
|
||||
def make_sync_memory(search_results=None):
|
||||
memory = Memory.__new__(Memory)
|
||||
memory.api_version = "v1.1"
|
||||
memory.reranker = None
|
||||
memory._search_vector_store = MagicMock(return_value=search_results or [])
|
||||
return memory
|
||||
|
||||
|
||||
def make_async_memory(search_results=None):
|
||||
memory = AsyncMemory.__new__(AsyncMemory)
|
||||
memory.api_version = "v1.1"
|
||||
memory.reranker = None
|
||||
memory._search_vector_store = AsyncMock(return_value=search_results or [])
|
||||
return memory
|
||||
|
||||
|
||||
def test_sync_slow_search_triggers_performance_notice_after_success(monkeypatch):
|
||||
results = [{"id": "m1"}, {"id": "m2"}]
|
||||
memory = make_sync_memory(search_results=results)
|
||||
performance_notice = MagicMock()
|
||||
temporal_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main.time, "perf_counter", MagicMock(side_effect=[100.0, 102.1]))
|
||||
monkeypatch.setattr(memory_main, "display_performance_slow_query_notice", performance_notice)
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
result = Memory.search(memory, "favorite drink", filters={"user_id": "u1"}, top_k=3)
|
||||
|
||||
assert result == {"results": results}
|
||||
memory._search_vector_store.assert_called_once()
|
||||
performance_notice.assert_called_once_with(memory, "sync", "search", pytest.approx(2.1), 3, 2)
|
||||
temporal_notice.assert_not_called()
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_fast_search_uses_first_run_notice(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
performance_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main.time, "perf_counter", MagicMock(side_effect=[100.0, 101.0]))
|
||||
monkeypatch.setattr(memory_main, "display_performance_slow_query_notice", performance_notice)
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
Memory.search(memory, "favorite drink", filters={"user_id": "u1"})
|
||||
|
||||
performance_notice.assert_not_called()
|
||||
first_run_notice.assert_called_once_with(memory, "sync", "search")
|
||||
|
||||
|
||||
def test_sync_failed_search_does_not_trigger_performance_notice(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
memory._search_vector_store.side_effect = RuntimeError("search failure")
|
||||
performance_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main.time, "perf_counter", MagicMock(return_value=100.0))
|
||||
monkeypatch.setattr(memory_main, "display_performance_slow_query_notice", performance_notice)
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
with pytest.raises(RuntimeError, match="search failure"):
|
||||
Memory.search(memory, "favorite drink", filters={"user_id": "u1"})
|
||||
|
||||
performance_notice.assert_not_called()
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_temporal_usage_takes_precedence_over_slow_search(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
performance_notice = MagicMock()
|
||||
temporal_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main.time, "perf_counter", MagicMock(side_effect=[100.0, 102.1]))
|
||||
monkeypatch.setattr(memory_main, "display_performance_slow_query_notice", performance_notice)
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
Memory.search(memory, "what happened last week?", filters={"user_id": "u1"})
|
||||
|
||||
temporal_notice.assert_called_once_with(memory, "sync", "search", "query", "relative_phrase")
|
||||
performance_notice.assert_not_called()
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_scale_takes_precedence_over_slow_search(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
performance_notice = MagicMock()
|
||||
scale_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main.time, "perf_counter", MagicMock(side_effect=[100.0, 102.1]))
|
||||
monkeypatch.setattr(memory_main, "display_performance_slow_query_notice", performance_notice)
|
||||
monkeypatch.setattr(memory_main, "display_scale_threshold_notice", scale_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
Memory.search(memory, "favorite drink", filters={"user_id": "u1"}, top_k=50)
|
||||
|
||||
scale_notice.assert_called_once_with(
|
||||
memory,
|
||||
"sync",
|
||||
"search",
|
||||
"top_k",
|
||||
"high_top_k",
|
||||
50,
|
||||
None,
|
||||
notices.SCALE_TOP_K_THRESHOLD,
|
||||
)
|
||||
performance_notice.assert_not_called()
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_slow_search_triggers_performance_notice_after_success(monkeypatch):
|
||||
results = [{"id": "m1"}]
|
||||
memory = make_async_memory(search_results=results)
|
||||
performance_notice = AsyncMock()
|
||||
temporal_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main.time, "perf_counter", MagicMock(side_effect=[100.0, 102.1]))
|
||||
monkeypatch.setattr(memory_main, "display_performance_slow_query_notice_async", performance_notice)
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice_async", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
result = await AsyncMemory.search(memory, "favorite drink", filters={"user_id": "u1"}, top_k=4)
|
||||
|
||||
assert result == {"results": results}
|
||||
memory._search_vector_store.assert_awaited_once()
|
||||
performance_notice.assert_awaited_once_with(memory, "async", "search", pytest.approx(2.1), 4, 1)
|
||||
temporal_notice.assert_not_awaited()
|
||||
first_run_notice.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_fast_search_uses_first_run_notice(monkeypatch):
|
||||
memory = make_async_memory()
|
||||
performance_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main.time, "perf_counter", MagicMock(side_effect=[100.0, 101.0]))
|
||||
monkeypatch.setattr(memory_main, "display_performance_slow_query_notice_async", performance_notice)
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice_async", AsyncMock())
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
await AsyncMemory.search(memory, "favorite drink", filters={"user_id": "u1"})
|
||||
|
||||
performance_notice.assert_not_awaited()
|
||||
first_run_notice.assert_awaited_once_with(memory, "async", "search")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_failed_search_does_not_trigger_performance_notice(monkeypatch):
|
||||
memory = make_async_memory()
|
||||
memory._search_vector_store.side_effect = RuntimeError("search failure")
|
||||
performance_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main.time, "perf_counter", MagicMock(return_value=100.0))
|
||||
monkeypatch.setattr(memory_main, "display_performance_slow_query_notice_async", performance_notice)
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice_async", AsyncMock())
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
with pytest.raises(RuntimeError, match="search failure"):
|
||||
await AsyncMemory.search(memory, "favorite drink", filters={"user_id": "u1"})
|
||||
|
||||
performance_notice.assert_not_awaited()
|
||||
first_run_notice.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_temporal_usage_takes_precedence_over_slow_search(monkeypatch):
|
||||
memory = make_async_memory()
|
||||
performance_notice = AsyncMock()
|
||||
temporal_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main.time, "perf_counter", MagicMock(side_effect=[100.0, 102.1]))
|
||||
monkeypatch.setattr(memory_main, "display_performance_slow_query_notice_async", performance_notice)
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice_async", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
await AsyncMemory.search(memory, "what happened last week?", filters={"user_id": "u1"})
|
||||
|
||||
temporal_notice.assert_awaited_once_with(memory, "async", "search", "query", "relative_phrase")
|
||||
performance_notice.assert_not_awaited()
|
||||
first_run_notice.assert_not_awaited()
|
||||
|
||||
|
||||
def test_performance_slow_query_displayed_logs_and_captures_event(notice_harness, capsys):
|
||||
config, telemetry = notice_harness
|
||||
flags = configure_flag(telemetry, "displayed", performance_payload())
|
||||
|
||||
notices.display_performance_slow_query_notice(
|
||||
MagicMock(),
|
||||
"sync",
|
||||
"search",
|
||||
elapsed_seconds=2.345,
|
||||
top_k=20,
|
||||
result_count=7,
|
||||
)
|
||||
|
||||
assert capsys.readouterr().err == "Performance CTA\n"
|
||||
telemetry.posthog.evaluate_flags.assert_called_once_with("oss-user", flag_keys=[notices.FLAG_KEY])
|
||||
telemetry.capture_event.assert_called_once()
|
||||
event_name, props = telemetry.capture_event.call_args.args
|
||||
assert event_name == notices.NOTICE_EVENT
|
||||
assert props["notice_id"] == "performance_slow_query"
|
||||
assert props["notice_type"] == "log_line"
|
||||
assert props["variant"] == "displayed"
|
||||
assert props["displayed"] is True
|
||||
assert props["payload"] == "Performance CTA"
|
||||
assert props["bypass_reason"] is None
|
||||
assert props["disabled_reason"] is None
|
||||
assert props["notice_config_found"] is True
|
||||
assert props["sync_type"] == "sync"
|
||||
assert props["trigger_function"] == "search"
|
||||
assert props["trigger_reason"] == "slow_query"
|
||||
assert props["elapsed_ms"] == 2345
|
||||
assert props["threshold_ms"] == 2000
|
||||
assert props["top_k"] == 20
|
||||
assert props["result_count"] == 7
|
||||
assert telemetry.capture_event.call_args.kwargs["flags"] is flags
|
||||
assert len(config["notice_state"]["performance_slow_query"]["events"]) == 1
|
||||
|
||||
|
||||
def test_performance_slow_query_holdout_is_silent_but_captures_event(notice_harness, capsys):
|
||||
_, telemetry = notice_harness
|
||||
configure_flag(telemetry, "holdout", performance_payload())
|
||||
|
||||
notices.display_performance_slow_query_notice(
|
||||
MagicMock(),
|
||||
"sync",
|
||||
"search",
|
||||
elapsed_seconds=2.1,
|
||||
top_k=10,
|
||||
result_count=2,
|
||||
)
|
||||
|
||||
assert capsys.readouterr().err == ""
|
||||
props = telemetry.capture_event.call_args.args[1]
|
||||
assert props["displayed"] is False
|
||||
assert props["bypass_reason"] == "holdout"
|
||||
assert props["trigger_reason"] == "slow_query"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("payload", "expected_reason", "expected_found"),
|
||||
[
|
||||
({}, "missing_notice_config", False),
|
||||
({"notices": {}}, "missing_notice_config", False),
|
||||
({"notices": "not-an-object"}, "missing_notice_config", False),
|
||||
(performance_payload(copy=None), "missing_copy", True),
|
||||
(performance_payload(enabled=False, copy="hidden"), "payload_disabled", True),
|
||||
],
|
||||
)
|
||||
def test_performance_slow_query_bad_or_disabled_payload_is_silent_and_safe(
|
||||
notice_harness, payload, expected_reason, expected_found, capsys
|
||||
):
|
||||
_, telemetry = notice_harness
|
||||
configure_flag(telemetry, "displayed", payload)
|
||||
|
||||
notices.display_performance_slow_query_notice(
|
||||
MagicMock(),
|
||||
"sync",
|
||||
"search",
|
||||
elapsed_seconds=2.1,
|
||||
top_k=10,
|
||||
result_count=2,
|
||||
)
|
||||
|
||||
assert capsys.readouterr().err == ""
|
||||
props = telemetry.capture_event.call_args.args[1]
|
||||
assert props["displayed"] is False
|
||||
assert props["bypass_reason"] == expected_reason
|
||||
assert props["notice_config_found"] is expected_found
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", [None, False])
|
||||
def test_performance_slow_query_blunt_flag_disable_does_not_capture_or_consume(
|
||||
notice_harness, variant, capsys
|
||||
):
|
||||
config, telemetry = notice_harness
|
||||
configure_flag(telemetry, variant, performance_payload())
|
||||
|
||||
notices.display_performance_slow_query_notice(
|
||||
MagicMock(),
|
||||
"sync",
|
||||
"search",
|
||||
elapsed_seconds=2.1,
|
||||
top_k=10,
|
||||
result_count=2,
|
||||
)
|
||||
|
||||
assert capsys.readouterr().err == ""
|
||||
telemetry.capture_event.assert_not_called()
|
||||
assert config.get("notice_state") is None
|
||||
|
||||
|
||||
def test_performance_slow_query_telemetry_disabled_does_not_touch_posthog_or_state(
|
||||
monkeypatch, capsys
|
||||
):
|
||||
load_config = MagicMock(return_value={})
|
||||
write_config = MagicMock()
|
||||
get_telemetry = MagicMock()
|
||||
|
||||
monkeypatch.setattr(notices, "_load_config", load_config)
|
||||
monkeypatch.setattr(notices, "_write_config", write_config)
|
||||
monkeypatch.setattr(notices.telemetry_module, "MEM0_TELEMETRY", False)
|
||||
monkeypatch.setattr(notices.telemetry_module, "_get_oss_telemetry", get_telemetry)
|
||||
|
||||
notices.display_performance_slow_query_notice(
|
||||
MagicMock(),
|
||||
"sync",
|
||||
"search",
|
||||
elapsed_seconds=2.1,
|
||||
top_k=10,
|
||||
result_count=2,
|
||||
)
|
||||
|
||||
load_config.assert_not_called()
|
||||
write_config.assert_not_called()
|
||||
get_telemetry.assert_not_called()
|
||||
assert capsys.readouterr().err == ""
|
||||
|
||||
|
||||
def test_performance_slow_query_cap_blocks_before_posthog_eval(notice_harness, capsys):
|
||||
config, telemetry = notice_harness
|
||||
configure_flag(telemetry, "displayed", performance_payload())
|
||||
|
||||
for _ in range(notices.PERFORMANCE_SLOW_QUERY_CAP):
|
||||
notices.display_performance_slow_query_notice(
|
||||
MagicMock(),
|
||||
"sync",
|
||||
"search",
|
||||
elapsed_seconds=2.1,
|
||||
top_k=10,
|
||||
result_count=2,
|
||||
)
|
||||
|
||||
notices.display_performance_slow_query_notice(
|
||||
MagicMock(),
|
||||
"sync",
|
||||
"search",
|
||||
elapsed_seconds=2.1,
|
||||
top_k=10,
|
||||
result_count=2,
|
||||
)
|
||||
|
||||
assert capsys.readouterr().err == "Performance CTA\n" * notices.PERFORMANCE_SLOW_QUERY_CAP
|
||||
assert telemetry.posthog.evaluate_flags.call_count == notices.PERFORMANCE_SLOW_QUERY_CAP
|
||||
assert telemetry.capture_event.call_count == notices.PERFORMANCE_SLOW_QUERY_CAP
|
||||
assert len(config["notice_state"]["performance_slow_query"]["events"]) == notices.PERFORMANCE_SLOW_QUERY_CAP
|
||||
|
||||
|
||||
def test_performance_slow_query_props_do_not_include_raw_user_inputs(notice_harness):
|
||||
_, telemetry = notice_harness
|
||||
configure_flag(telemetry, "displayed", performance_payload(copy="safe copy"))
|
||||
|
||||
notices.display_performance_slow_query_notice(
|
||||
MagicMock(),
|
||||
"sync",
|
||||
"search",
|
||||
elapsed_seconds=2.1,
|
||||
top_k=10,
|
||||
result_count=2,
|
||||
)
|
||||
|
||||
props = telemetry.capture_event.call_args.args[1]
|
||||
assert "favorite drink" not in str(props)
|
||||
assert "user_id" not in str(props)
|
||||
assert "green tea" not in str(props)
|
||||
@@ -0,0 +1,96 @@
|
||||
import pytest
|
||||
|
||||
from mem0.memory import main as memory_main
|
||||
from mem0.memory.main import AsyncMemory, Memory
|
||||
|
||||
|
||||
def test_sync_add_timestamp_raises_before_validation(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def get_error(sync_type, trigger_function, trigger_parameter):
|
||||
calls.append((sync_type, trigger_function, trigger_parameter))
|
||||
return "blocked timestamp"
|
||||
|
||||
monkeypatch.setattr(memory_main, "get_temporal_feature_error_message", get_error)
|
||||
|
||||
with pytest.raises(ValueError, match="blocked timestamp"):
|
||||
Memory.add(Memory.__new__(Memory), "hello", timestamp=123)
|
||||
|
||||
assert calls == [("sync", "add", "timestamp")]
|
||||
|
||||
|
||||
def test_sync_search_reference_date_raises_before_validation(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def get_error(sync_type, trigger_function, trigger_parameter):
|
||||
calls.append((sync_type, trigger_function, trigger_parameter))
|
||||
return "blocked reference date"
|
||||
|
||||
monkeypatch.setattr(memory_main, "get_temporal_feature_error_message", get_error)
|
||||
|
||||
with pytest.raises(ValueError, match="blocked reference date"):
|
||||
Memory.search(Memory.__new__(Memory), "what happened last week?", reference_date="2025-03-21")
|
||||
|
||||
assert calls == [("sync", "search", "reference_date")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_add_timestamp_raises_before_validation(monkeypatch):
|
||||
calls = []
|
||||
|
||||
async def get_error(sync_type, trigger_function, trigger_parameter):
|
||||
calls.append((sync_type, trigger_function, trigger_parameter))
|
||||
return "blocked async timestamp"
|
||||
|
||||
monkeypatch.setattr(memory_main, "get_temporal_feature_error_message_async", get_error)
|
||||
|
||||
with pytest.raises(ValueError, match="blocked async timestamp"):
|
||||
await AsyncMemory.add(AsyncMemory.__new__(AsyncMemory), "hello", timestamp=123)
|
||||
|
||||
assert calls == [("async", "add", "timestamp")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_search_reference_date_raises_before_validation(monkeypatch):
|
||||
calls = []
|
||||
|
||||
async def get_error(sync_type, trigger_function, trigger_parameter):
|
||||
calls.append((sync_type, trigger_function, trigger_parameter))
|
||||
return "blocked async reference date"
|
||||
|
||||
monkeypatch.setattr(memory_main, "get_temporal_feature_error_message_async", get_error)
|
||||
|
||||
with pytest.raises(ValueError, match="blocked async reference date"):
|
||||
await AsyncMemory.search(
|
||||
AsyncMemory.__new__(AsyncMemory),
|
||||
"what happened last week?",
|
||||
reference_date="2025-03-21",
|
||||
)
|
||||
|
||||
assert calls == [("async", "search", "reference_date")]
|
||||
|
||||
|
||||
def test_sync_add_without_timestamp_does_not_call_temporal_feature_notice(monkeypatch):
|
||||
get_error = monkeypatch.setattr(
|
||||
memory_main,
|
||||
"get_temporal_feature_error_message",
|
||||
lambda *args: pytest.fail("temporal feature notice should not run"),
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match="At least one of 'user_id', 'agent_id', or 'run_id'"):
|
||||
Memory.add(Memory.__new__(Memory), "hello")
|
||||
|
||||
assert get_error is None
|
||||
|
||||
|
||||
def test_sync_search_without_reference_date_does_not_call_temporal_feature_notice(monkeypatch):
|
||||
get_error = monkeypatch.setattr(
|
||||
memory_main,
|
||||
"get_temporal_feature_error_message",
|
||||
lambda *args: pytest.fail("temporal feature notice should not run"),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="filters must contain"):
|
||||
Memory.search(Memory.__new__(Memory), "hello")
|
||||
|
||||
assert get_error is None
|
||||
@@ -0,0 +1,202 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from mem0.memory import main as memory_main
|
||||
from mem0.memory.main import AsyncMemory, Memory
|
||||
|
||||
|
||||
def make_sync_memory():
|
||||
memory = Memory.__new__(Memory)
|
||||
memory.config = SimpleNamespace(llm=SimpleNamespace(config={}))
|
||||
memory.api_version = "v1.1"
|
||||
memory.reranker = None
|
||||
memory._add_to_vector_store = MagicMock(return_value=[])
|
||||
memory._search_vector_store = MagicMock(return_value=[])
|
||||
return memory
|
||||
|
||||
|
||||
def make_async_memory():
|
||||
memory = AsyncMemory.__new__(AsyncMemory)
|
||||
memory.config = SimpleNamespace(llm=SimpleNamespace(config={}))
|
||||
memory.api_version = "v1.1"
|
||||
memory.reranker = None
|
||||
memory._add_to_vector_store = AsyncMock(return_value=[])
|
||||
memory._search_vector_store = AsyncMock(return_value=[])
|
||||
return memory
|
||||
|
||||
|
||||
def test_sync_add_temporal_metadata_triggers_notice_after_success(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
temporal_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
result = Memory.add(
|
||||
memory,
|
||||
"The user visited Paris.",
|
||||
user_id="u1",
|
||||
metadata={"event_date": "2025-04-09"},
|
||||
infer=False,
|
||||
)
|
||||
|
||||
assert result == {"results": []}
|
||||
memory._add_to_vector_store.assert_called_once()
|
||||
temporal_notice.assert_called_once_with(memory, "sync", "add", "metadata", "date_like_metadata")
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_add_non_temporal_metadata_uses_first_run_notice(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
temporal_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
Memory.add(memory, "The user likes tea.", user_id="u1", metadata={"topic": "drink"}, infer=False)
|
||||
|
||||
temporal_notice.assert_not_called()
|
||||
first_run_notice.assert_called_once_with(memory, "sync", "add")
|
||||
|
||||
|
||||
def test_sync_add_failure_does_not_trigger_temporal_usage_notice(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
memory._add_to_vector_store.side_effect = RuntimeError("vector failure")
|
||||
temporal_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
with pytest.raises(RuntimeError, match="vector failure"):
|
||||
Memory.add(
|
||||
memory,
|
||||
"The user visited Paris.",
|
||||
user_id="u1",
|
||||
metadata={"event_date": "2025-04-09"},
|
||||
infer=False,
|
||||
)
|
||||
|
||||
temporal_notice.assert_not_called()
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_search_temporal_query_triggers_notice_after_success(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
temporal_notice = MagicMock()
|
||||
first_run_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", first_run_notice)
|
||||
|
||||
result = Memory.search(memory, "what happened last week?", filters={"user_id": "u1"})
|
||||
|
||||
assert result == {"results": []}
|
||||
memory._search_vector_store.assert_called_once()
|
||||
temporal_notice.assert_called_once_with(memory, "sync", "search", "query", "relative_phrase")
|
||||
first_run_notice.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_search_temporal_filter_triggers_notice_after_success(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
temporal_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", MagicMock())
|
||||
|
||||
Memory.search(
|
||||
memory,
|
||||
"favorite drink",
|
||||
filters={"user_id": "u1", "created_at": {"gte": "2025-04-01"}},
|
||||
)
|
||||
|
||||
temporal_notice.assert_called_once_with(memory, "sync", "search", "filter", "date_range_filter")
|
||||
|
||||
|
||||
def test_sync_search_failure_does_not_trigger_temporal_usage_notice(monkeypatch):
|
||||
memory = make_sync_memory()
|
||||
memory._search_vector_store.side_effect = RuntimeError("search failure")
|
||||
temporal_notice = MagicMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice", MagicMock())
|
||||
|
||||
with pytest.raises(RuntimeError, match="search failure"):
|
||||
Memory.search(memory, "what happened last week?", filters={"user_id": "u1"})
|
||||
|
||||
temporal_notice.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_add_temporal_metadata_triggers_notice_after_success(monkeypatch):
|
||||
memory = make_async_memory()
|
||||
temporal_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice_async", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
result = await AsyncMemory.add(
|
||||
memory,
|
||||
"The user visited Paris.",
|
||||
user_id="u1",
|
||||
metadata={"event_date": "2025-04-09"},
|
||||
infer=False,
|
||||
)
|
||||
|
||||
assert result == {"results": []}
|
||||
memory._add_to_vector_store.assert_awaited_once()
|
||||
temporal_notice.assert_awaited_once_with(memory, "async", "add", "metadata", "date_like_metadata")
|
||||
first_run_notice.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_add_runs_scale_detection_in_thread(monkeypatch):
|
||||
memory = make_async_memory()
|
||||
scale_detector = MagicMock(return_value=("memory_count", "memory_count_threshold", None, 2000, 2000))
|
||||
scale_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
to_thread_calls = []
|
||||
|
||||
async def to_thread(fn, *args, **kwargs):
|
||||
to_thread_calls.append((fn, args, kwargs))
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(memory_main, "detect_scale_threshold_from_add_result", scale_detector)
|
||||
monkeypatch.setattr(memory_main.asyncio, "to_thread", to_thread)
|
||||
monkeypatch.setattr(memory_main, "display_scale_threshold_notice_async", scale_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
result = await AsyncMemory.add(memory, "The user likes tea.", user_id="u1", infer=False)
|
||||
|
||||
assert result == {"results": []}
|
||||
assert to_thread_calls == [(scale_detector, (memory, []), {})]
|
||||
scale_detector.assert_called_once_with(memory, [])
|
||||
scale_notice.assert_awaited_once_with(
|
||||
memory,
|
||||
"async",
|
||||
"add",
|
||||
"memory_count",
|
||||
"memory_count_threshold",
|
||||
None,
|
||||
2000,
|
||||
2000,
|
||||
)
|
||||
first_run_notice.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_search_temporal_query_triggers_notice_after_success(monkeypatch):
|
||||
memory = make_async_memory()
|
||||
temporal_notice = AsyncMock()
|
||||
first_run_notice = AsyncMock()
|
||||
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||
monkeypatch.setattr(memory_main, "display_temporal_usage_notice_async", temporal_notice)
|
||||
monkeypatch.setattr(memory_main, "display_first_run_notice_async", first_run_notice)
|
||||
|
||||
result = await AsyncMemory.search(memory, "what happened last week?", filters={"user_id": "u1"})
|
||||
|
||||
assert result == {"results": []}
|
||||
memory._search_vector_store.assert_awaited_once()
|
||||
temporal_notice.assert_awaited_once_with(memory, "async", "search", "query", "relative_phrase")
|
||||
first_run_notice.assert_not_awaited()
|
||||
@@ -70,6 +70,10 @@ class TestTelemetryEnabled:
|
||||
with patch("mem0.memory.telemetry.get_or_create_user_id", return_value="test-user"):
|
||||
at = telemetry_module.AnonymousTelemetry()
|
||||
mock_posthog.assert_called_once()
|
||||
assert (
|
||||
mock_posthog.call_args.kwargs["feature_flags_request_timeout_seconds"]
|
||||
== telemetry_module.FEATURE_FLAGS_REQUEST_TIMEOUT_SECONDS
|
||||
)
|
||||
assert at.posthog is not None
|
||||
assert at.user_id == "test-user"
|
||||
|
||||
@@ -84,6 +88,22 @@ class TestTelemetryEnabled:
|
||||
telemetry_module.capture_event("test.event", mock_memory)
|
||||
mock_at.capture_event.assert_called_once()
|
||||
|
||||
def test_anonymous_capture_event_passes_flags_to_posthog(self):
|
||||
"""capture_event() should use PostHog's event-first API and preserve flag snapshots."""
|
||||
flags = MagicMock()
|
||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||
with patch("mem0.memory.telemetry.Posthog") as mock_posthog_cls:
|
||||
with patch("mem0.memory.telemetry.get_or_create_user_id", return_value="test-user"):
|
||||
at = telemetry_module.AnonymousTelemetry()
|
||||
at.capture_event("test.event", {"key": "value"}, flags=flags)
|
||||
|
||||
mock_posthog_cls.return_value.capture.assert_called_once()
|
||||
args, kwargs = mock_posthog_cls.return_value.capture.call_args
|
||||
assert args == ("test.event",)
|
||||
assert kwargs["distinct_id"] == "test-user"
|
||||
assert kwargs["flags"] is flags
|
||||
assert kwargs["properties"]["key"] == "value"
|
||||
|
||||
def test_capture_client_event_sends_when_enabled(self):
|
||||
"""capture_client_event() should call client_telemetry.capture_event when enabled."""
|
||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||
|
||||
@@ -162,16 +162,22 @@ class TestBeforeSendWiring:
|
||||
# before_send is None (default), not _sampling_before_send
|
||||
assert kwargs.get("before_send") is None
|
||||
|
||||
def test_anonymous_telemetry_falls_back_when_posthog_rejects_before_send(self):
|
||||
"""If posthog (older version) rejects before_send, construction still succeeds."""
|
||||
def test_anonymous_telemetry_constructs_posthog_with_current_kwargs(self):
|
||||
"""AnonymousTelemetry uses the supported PostHog constructor shape."""
|
||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||
with patch("mem0.memory.telemetry.Posthog") as mock_posthog_cls:
|
||||
with patch("mem0.memory.telemetry.get_or_create_user_id", return_value="u"):
|
||||
# First call (with before_send) raises TypeError; second call succeeds.
|
||||
mock_posthog_cls.side_effect = [TypeError("unexpected kwarg before_send"), object()]
|
||||
at = telemetry_module.AnonymousTelemetry(before_send=telemetry_module._sampling_before_send)
|
||||
# Constructor was called twice: once with before_send, once without
|
||||
assert mock_posthog_cls.call_count == 2
|
||||
mock_posthog_cls.assert_called_once()
|
||||
_, kwargs = mock_posthog_cls.call_args
|
||||
assert kwargs["project_api_key"] == telemetry_module.PROJECT_API_KEY
|
||||
assert kwargs["host"] == telemetry_module.HOST
|
||||
assert kwargs["before_send"] is telemetry_module._sampling_before_send
|
||||
assert (
|
||||
kwargs["feature_flags_request_timeout_seconds"]
|
||||
== telemetry_module.FEATURE_FLAGS_REQUEST_TIMEOUT_SECONDS
|
||||
)
|
||||
assert at.user_id == "u"
|
||||
assert at.posthog is not None
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user