from __future__ import annotations import argparse import hashlib import json import os import sys import time import uuid from datetime import datetime, timezone from pathlib import Path from typing import Any ROOT = Path(__file__).resolve().parents[1] RUNS_DIR = ROOT / "runs" MEM0_VENV_SITE = ROOT / ".venv-mem0" / "lib" / "python3.12" / "site-packages" if MEM0_VENV_SITE.exists(): sys.path.insert(0, str(MEM0_VENV_SITE)) os.environ.setdefault("HOME", "/tmp/mem0-home") os.environ.setdefault("MEM0_DIR", "/tmp/mem0-home/.mem0") os.environ.setdefault("MEM0_TELEMETRY", "False") os.environ.setdefault("HF_HOME", "/tmp/mem0-home/huggingface") from mem0 import Memory # noqa: E402 from mem0.utils.lemmatization import lemmatize_for_bm25 # noqa: E402 from qdrant_client.models import PointStruct, SparseVector # noqa: E402 from run_sibyl_365d_500c_category_benchmark import ( # noqa: E402 MARKER_DAYS, MILESTONE_DAYS, ROLES, START_DATE, STAKEHOLDERS_PER_COMPANY, TIMELINE_DAYS, TOPICS, approx_tokens, company_index_from_case, company_key, day_date, display, marker, person_key, region, relation_key, role_index_from_case, score_context, segment, slug, spread_indices, topic, write_json, ) DEFAULT_RUN_ID = "run-2026-06-06-mem0-365d-500c-50q-category-retrieval" DEFAULT_DATA_ROOT = f"/tmp/{DEFAULT_RUN_ID}" DEFAULT_EMBEDDER_MODEL = "BAAI/bge-small-en-v1.5" DEFAULT_EMBEDDING_DIMS = 384 RUN_SCOPE = "365d-500c-50q-mem0-baseline" def iso_now() -> str: return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") def compact_json(value: Any) -> str: return json.dumps(value, sort_keys=True, separators=(",", ":"), default=str) def flatten_memory(value: Any, *, prefix: str = "") -> list[str]: if isinstance(value, dict): tokens: list[str] = [] for key, item in value.items(): nested_prefix = f"{prefix}.{key}" if prefix else str(key) tokens.extend(flatten_memory(item, prefix=nested_prefix)) return tokens if isinstance(value, list | tuple): tokens = [] for index, item in enumerate(value): tokens.extend(flatten_memory(item, prefix=f"{prefix}.{index}")) return tokens if value is None: return [] return [f"{prefix}={value}"] def make_memory(data: dict[str, Any]) -> str: return " ".join(flatten_memory(data)) def company_records(index: int, *, include_state_snapshots: bool) -> list[str]: company_slug = slug(index) company_display = display(index) records: list[str] = [] records.append( make_memory( { "tier": "entity", "category": "company", "key": company_key(index), "body": { "display_name": company_display, "segment": segment(index), "region": region(index), "timeline_start": START_DATE.isoformat(), "timeline_days": TIMELINE_DAYS, "stakeholder_count": STAKEHOLDERS_PER_COMPANY, "context_stats": { "simulated_days": TIMELINE_DAYS, "topic_cycle_size": len(TOPICS), "milestone_count": len(MILESTONE_DAYS), }, }, } ) ) for role_index, role in enumerate(ROLES, start=1): person = person_key(index, role_index) records.append( make_memory( { "tier": "entity", "category": "person", "key": person, "body": { "display_name": f"{company_display} Contact {role_index}", "role": role, "company": company_key(index), "company_display": company_display, }, } ) ) records.append( make_memory( { "tier": "entity", "category": "relationship", "key": relation_key(index, role_index), "body": { "source": person, "target": company_key(index), "type": "works_on_account", "role": role, "description": f"{person} is the {role} for {company_display}.", }, } ) ) for day_index in range(1, TIMELINE_DAYS + 1): event_marker = marker(index, day_index) event_topic = topic(day_index) event_date = day_date(day_index) records.append( make_memory( { "tier": "journal", "category": "chronology", "key": f"run365-{company_slug}-day-{day_index:03d}", "body": { "kind": "scale365_chronology_daily_update", "company": company_display, "company_slug": company_slug, "day_index": day_index, "date": event_date, "topic": event_topic, "marker": event_marker, "summary": f"{company_display} day {day_index:03d}: {event_topic} update recorded.", }, } ) ) if day_index in MILESTONE_DAYS: records.append( make_memory( { "tier": "entity", "category": "timeline_milestone", "key": f"run365-{company_slug}-day-{day_index:03d}", "body": { "company": company_key(index), "company_display": company_display, "day_index": day_index, "date": event_date, "topic": event_topic, "marker": event_marker, "decision": f"{company_display} checkpoint day {day_index:03d} accepted for {event_topic}.", }, } ) ) if include_state_snapshots and (day_index % 10 == 0 or day_index == TIMELINE_DAYS): records.append( make_memory( { "tier": "state", "key": f"run365-current-status-{company_slug}", "snapshot_day_index": day_index, "body": { "company": company_display, "last_day_index": day_index, "last_date": event_date, "last_marker": event_marker, "status": "active", }, } ) ) if not include_state_snapshots: records.append( make_memory( { "tier": "state", "key": f"run365-current-status-{company_slug}", "body": { "company": company_display, "last_day_index": TIMELINE_DAYS, "last_date": day_date(TIMELINE_DAYS), "last_marker": marker(index, TIMELINE_DAYS), "status": "active", }, } ) ) return records def build_questions(company_count: int, questions_per_category: int) -> list[dict[str, Any]]: cases: list[dict[str, Any]] = [] for index in spread_indices(questions_per_category, company_count, offset=0): cases.append( { "category": "status", "id": f"status_company_{index:03d}", "query": f"What is the current status for {display(index)} after the 365 day chronology?", "retrieval_query": f"{display(index)} current status day 365", "expected_contains": [display(index), str(TIMELINE_DAYS), day_date(TIMELINE_DAYS), "active"], } ) for position, index in enumerate(spread_indices(questions_per_category, company_count, offset=7)): day_index = MILESTONE_DAYS[position % len(MILESTONE_DAYS)] cases.append( { "category": "milestone", "id": f"milestone_company_{index:03d}_day_{day_index:03d}", "query": f"What was the {display(index)} day {day_index} milestone accepted for?", "retrieval_query": f"{display(index)} day {day_index} milestone", "expected_contains": [display(index), str(day_index), day_date(day_index), topic(day_index)], } ) for index in spread_indices(questions_per_category, company_count, offset=13): cases.append( { "category": "context_stat", "id": f"context_company_{index:03d}", "query": f"What segment and region context is stored for {display(index)}?", "retrieval_query": display(index), "expected_contains": [display(index), segment(index), region(index), str(TIMELINE_DAYS)], } ) for position, index in enumerate(spread_indices(questions_per_category, company_count, offset=19)): role_index = (position % len(ROLES)) + 1 cases.append( { "category": "role", "id": f"role_company_{index:03d}_{role_index}", "query": f"Who is the {ROLES[role_index - 1]} for {display(index)}?", "retrieval_query": f"{display(index)} {ROLES[role_index - 1]}", "expected_contains": [f"{display(index)} Contact {role_index}", ROLES[role_index - 1]], } ) for position, index in enumerate(spread_indices(questions_per_category, company_count, offset=29)): day_index = MARKER_DAYS[position % len(MARKER_DAYS)] cases.append( { "category": "marker", "id": f"marker_company_{index:03d}_day_{day_index:03d}", "query": f"What marker is associated with {display(index)} day {day_index}?", "retrieval_query": marker(index, day_index), "expected_contains": [marker(index, day_index)], } ) for position, index in enumerate(spread_indices(questions_per_category, company_count, offset=37)): day_index = 1 + ((position * 17) % TIMELINE_DAYS) cases.append( { "category": "temporal_topic", "id": f"topic_company_{index:03d}_day_{day_index:03d}", "query": f"What topic was recorded for {display(index)} on day {day_index}?", "retrieval_query": f"{display(index)} day {day_index} topic", "expected_contains": [display(index), str(day_index), day_date(day_index), topic(day_index)], } ) for index in range(1, questions_per_category + 1): name = f"Phantom Prospect {index:03d}" safe = name.lower().replace(" ", "_") cases.append( { "category": "negative_trap", "id": f"negative_{safe}", "query": f"What is the current status for {name} after the 365 day chronology?", "retrieval_query": f"{name} current status day 365", "expected_unknown": True, "forbidden_contains": ["active", str(TIMELINE_DAYS), day_date(TIMELINE_DAYS)], } ) return cases def mem0_config(data_root: Path, collection_name: str, embedder_model: str, embedding_dims: int) -> dict[str, Any]: return { "version": "v1.1", "history_db_path": str(data_root / "history.db"), "llm": {"provider": "openai", "config": {"api_key": "not-used-in-infer-false", "model": "gpt-5-mini"}}, "embedder": {"provider": "fastembed", "config": {"model": embedder_model, "embedding_dims": embedding_dims}}, "vector_store": { "provider": "qdrant", "config": { "collection_name": collection_name, "path": str(data_root / "qdrant"), "embedding_model_dims": embedding_dims, "on_disk": True, }, }, } def ingest_public_api(memory: Memory, *, company_count: int, batch_size: int, include_state_snapshots: bool, progress: bool) -> dict[str, Any]: started = time.perf_counter() added = 0 for index in range(1, company_count + 1): records = company_records(index, include_state_snapshots=include_state_snapshots) for offset in range(0, len(records), batch_size): messages = [{"role": "user", "content": content} for content in records[offset : offset + batch_size]] result = memory.add( messages, run_id=RUN_SCOPE, infer=False, metadata={"benchmark": "sibyl-365d-500c-category", "company": display(index), "company_index": index}, ) added += len(result.get("results", [])) if progress and (index == 1 or index % 10 == 0): print(f"[ingest] company {index}/{company_count} records={added} elapsed={time.perf_counter() - started:.1f}s", file=sys.stderr, flush=True) return {"mode": "mem0_public_add_infer_false", "records_added": added, "elapsed_seconds": round(time.perf_counter() - started, 3)} def ingest_vector_store_batch(memory: Memory, *, company_count: int, batch_size: int, include_state_snapshots: bool, progress: bool) -> dict[str, Any]: started = time.perf_counter() added = 0 contents: list[str] = [] ids: list[str] = [] payloads: list[dict[str, Any]] = [] def flush() -> None: nonlocal contents, ids, payloads, added if not contents: return if hasattr(memory.embedding_model, "dense_model"): vectors = [embedding.tolist() for embedding in memory.embedding_model.dense_model.embed(contents)] else: vectors = [memory.embedding_model.embed(content, "add") for content in contents] points: list[PointStruct] = [] sparse_vectors: list[Any | None] = [None] * len(contents) if getattr(memory.vector_store, "_has_bm25_slot", False): encoder = memory.vector_store._get_bm25_encoder() if encoder is not None: texts = [payload.get("text_lemmatized") or payload.get("data", "") for payload in payloads] sparse_vectors = list(encoder.embed(texts)) for position, vector in enumerate(vectors): named_vectors: dict[str, Any] = {"": vector} sparse = sparse_vectors[position] if sparse is not None: named_vectors["bm25"] = SparseVector(indices=sparse.indices.tolist(), values=sparse.values.tolist()) points.append(PointStruct(id=ids[position], vector=named_vectors, payload=payloads[position])) memory.vector_store.client.upsert(collection_name=memory.vector_store.collection_name, points=points) added += len(contents) contents = [] ids = [] payloads = [] for index in range(1, company_count + 1): for content in company_records(index, include_state_snapshots=include_state_snapshots): contents.append(content) ids.append(str(uuid.uuid4())) created_at = iso_now() payloads.append( { "data": content, "hash": hashlib.md5(content.encode()).hexdigest(), "created_at": created_at, "updated_at": created_at, "text_lemmatized": lemmatize_for_bm25(content), "run_id": RUN_SCOPE, "role": "user", "benchmark": "sibyl-365d-500c-category", "company": display(index), "company_index": index, } ) if len(contents) >= batch_size: flush() if progress and (index == 1 or index % 10 == 0): print(f"[ingest] company {index}/{company_count} records={added + len(contents)} elapsed={time.perf_counter() - started:.1f}s", file=sys.stderr, flush=True) flush() return {"mode": "mem0_vector_store_batch_infer_false_equivalent", "records_added": added, "elapsed_seconds": round(time.perf_counter() - started, 3)} def result_row(item: dict[str, Any]) -> dict[str, Any]: metadata = item.get("metadata") or {} memory = item.get("memory", "") return { "tier": "mem0", "key": item.get("id"), "score": item.get("score"), "category": metadata.get("company"), "body": memory, "snippet": memory[:500], } def search_context(memory: Memory, case: dict[str, Any], top_k: int) -> list[dict[str, Any]]: query = str(case.get("retrieval_query") or case.get("query") or "") result = memory.search(query=query, top_k=top_k, filters={"run_id": RUN_SCOPE}, threshold=0.0, rerank=False) return [result_row(item) for item in result.get("results", [])] def summarize_questions(questions: list[dict[str, Any]]) -> dict[str, Any]: summary: dict[str, Any] = {} for item in questions: category = item["category"] bucket = summary.setdefault(category, {"passed": 0, "total": 0}) bucket["total"] += 1 bucket["passed"] += int(bool(item["retrieval_score"]["passed"])) return summary def build_markdown(raw: dict[str, Any]) -> str: rows = "\n".join( f"| {category} | {row['passed']} / {row['total']} |" for category, row in raw["summary"]["retrieval_by_category"].items() ) failures = [ { "id": item["id"], "category": item["category"], "query": item["query"], "expected_contains": item.get("expected_contains", []), "forbidden_contains": item.get("forbidden_contains", []), "retrieval_score": item["retrieval_score"], "context_sample": item["context"][:3], } for item in raw["questions"] if not item["retrieval_score"]["passed"] ] return f"""# Mem0 365d 500c Category Baseline Run: `{raw["run_id"]}` ## Scope - System: Mem0 open-source local baseline. - Corpus: {raw["dataset"]["company_count"]} companies, {raw["dataset"]["stakeholder_count"]} stakeholders, {raw["dataset"]["timeline_days"]} simulated days. - Stored memories: {raw["ingest"]["records_added"]}. - Questions: {raw["summary"]["question_total"]}, {raw["dataset"]["questions_per_category"]} per category. - Mode: `{raw["ingest"]["mode"]}`. - Inference: `infer=False`, raw structured memory records, no LLM extraction. - Vector store: Qdrant local. - Embedder: fastembed `{raw["mem0"]["embedder_model"]}`. - Sparse BM25: `{raw["mem0"]["sparse_bm25"]}`. - Answer mode: skipped. ## Summary | Metric | Value | | --- | ---: | | Retrieval passed | {raw["summary"]["retrieval_passed"]} / {raw["summary"]["question_total"]} | | Avg context rows | {raw["summary"]["avg_context_rows"]:.2f} | | Avg context tokens | {raw["summary"]["avg_context_tokens"]:.2f} | | Ingest elapsed seconds | {raw["ingest"]["elapsed_seconds"]} | | Retrieval elapsed seconds | {raw["summary"]["retrieval_elapsed_seconds"]} | | Model/API cost USD | 0.000000 | ## Retrieval By Category | Category | Retrieval passed | | --- | ---: | {rows} ## Failures ```json {json.dumps(failures[:25], indent=2, sort_keys=True, default=str)} ``` """ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run Mem0 local baseline on the Sibyl 365d/500-company category benchmark.") parser.add_argument("--run-id", default=DEFAULT_RUN_ID) parser.add_argument("--data-root", default=DEFAULT_DATA_ROOT) parser.add_argument("--company-limit", type=int, default=500) parser.add_argument("--questions-per-category", type=int, default=50) parser.add_argument("--question-limit", type=int, default=None) parser.add_argument("--top-k", type=int, default=8) parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--embedder-model", default=DEFAULT_EMBEDDER_MODEL) parser.add_argument("--embedding-dims", type=int, default=DEFAULT_EMBEDDING_DIMS) parser.add_argument("--collection-name", default=None) parser.add_argument("--dense-only", action="store_true") parser.add_argument("--skip-ingest", action="store_true") parser.add_argument("--include-state-snapshots", action="store_true") parser.add_argument("--public-add", action="store_true") parser.add_argument("--progress", action="store_true") return parser.parse_args() def main() -> None: args = parse_args() if args.company_limit < 1 or args.company_limit > 500: raise SystemExit("--company-limit must be between 1 and 500") if args.questions_per_category < 1: raise SystemExit("--questions-per-category must be positive") data_root = Path(args.data_root) data_root.mkdir(parents=True, exist_ok=True) collection_name = args.collection_name or args.run_id.replace("-", "_") memory = Memory.from_config(mem0_config(data_root, collection_name, args.embedder_model, args.embedding_dims)) if args.dense_only: memory.vector_store._has_bm25_slot = False raw: dict[str, Any] = { "run_id": args.run_id, "ts_started": iso_now(), "source": str(data_root), "mem0": { "version": "2.0.4", "run_scope": RUN_SCOPE, "vector_store": "qdrant_local", "embedder": "fastembed", "embedder_model": args.embedder_model, "embedding_dims": args.embedding_dims, "sparse_bm25": not args.dense_only, }, } if args.skip_ingest: ingest = {"mode": "skipped", "records_added": None, "elapsed_seconds": 0} elif args.public_add: ingest = ingest_public_api( memory, company_count=args.company_limit, batch_size=args.batch_size, include_state_snapshots=args.include_state_snapshots, progress=args.progress, ) else: ingest = ingest_vector_store_batch( memory, company_count=args.company_limit, batch_size=args.batch_size, include_state_snapshots=args.include_state_snapshots, progress=args.progress, ) questions = build_questions(args.company_limit, args.questions_per_category) if args.question_limit is not None: questions = questions[: args.question_limit] started_questions = time.perf_counter() question_rows: list[dict[str, Any]] = [] for position, case in enumerate(questions, start=1): context = search_context(memory, case, args.top_k) row = { **case, "context_count": len(context), "context_tokens_estimate": approx_tokens(context), "context": context, "retrieval_score": score_context(case, context), } question_rows.append(row) if args.progress and (position == 1 or position % 25 == 0): print(f"[questions] {position}/{len(questions)} elapsed={time.perf_counter() - started_questions:.1f}s", file=sys.stderr, flush=True) retrieval_passed = sum(1 for item in question_rows if item["retrieval_score"]["passed"]) raw["ts_completed"] = iso_now() raw["ingest"] = ingest raw["dataset"] = { "company_count": args.company_limit, "stakeholder_count": args.company_limit * STAKEHOLDERS_PER_COMPANY, "timeline_days": TIMELINE_DAYS, "questions_per_category": args.questions_per_category, "include_state_snapshots": args.include_state_snapshots, "top_k": args.top_k, } raw["questions"] = question_rows raw["summary"] = { "question_total": len(question_rows), "retrieval_passed": retrieval_passed, "retrieval_failed": len(question_rows) - retrieval_passed, "retrieval_by_category": summarize_questions(question_rows), "avg_context_rows": sum(item["context_count"] for item in question_rows) / len(question_rows), "avg_context_tokens": sum(item["context_tokens_estimate"] for item in question_rows) / len(question_rows), "retrieval_elapsed_seconds": round(time.perf_counter() - started_questions, 3), "estimated_cost": {"estimated_total_usd": 0.0}, } raw["status"] = "PASS" if retrieval_passed == len(question_rows) else "FAIL" RUNS_DIR.mkdir(parents=True, exist_ok=True) raw_path = RUNS_DIR / f"{args.run_id}.raw_result.json" md_path = RUNS_DIR / f"{args.run_id}.md" write_json(raw_path, raw) md_path.write_text(build_markdown(raw), encoding="utf8") print(json.dumps({"status": raw["status"], "summary": raw["summary"], "raw": str(raw_path), "md": str(md_path)}, indent=2)) if __name__ == "__main__": main()