feat: add dataset download

Author: gusye1234

memory_benchmark/__init__.py

@@ -1,6 +1,6 @@
+from .launch import run_memory_benchmark
+from .env import CONFIG
+
 __author__ = "Memobase"
 __version__ = "0.0.1.dev0"
 __url__ = "https://github.com/memodb-io/memory-benchmark"
-
-
-print("Not ready")

memory_benchmark/datasets/base.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel
+
+
+class Message(BaseModel):
+    role: str
+    content: str
+    alias: str | None = None
+    id: str | None = None
+
+
+class Conversation(BaseModel):
+    messages: list[Message]
+    id: str | None = None
+
+
+class RemoteFile(BaseModel):
+    url: str
+    name: str
+    hash: str

memory_benchmark/datasets/download.py

@@ -0,0 +1,79 @@
+import os
+import requests
+from rich.progress import track
+from hashlib import sha256
+from ..env import HOME_PATH, console
+from .base import RemoteFile
+
+REMOTE_URLS = {
+    "locomo": [
+        RemoteFile(
+            url="https://github.com/snap-research/locomo/raw/refs/heads/main/data/locomo10.json",
+            name="locomo10.json",
+            hash="79fa87e90f04081343b8c8debecb80a9a6842b76a7aa537dc9fdf651ea698ff4",
+        )
+    ]
+}
+
+
+def exist_or_download(dataset: str):
+    if check_local_dataset_exist(dataset):
+        console.log(f"Dataset {dataset} already exists")
+    else:
+        download_from_github(dataset)
+
+
+def check_local_dataset_exist(dataset: str) -> bool:
+    assert (
+        dataset in REMOTE_URLS
+    ), f"Dataset {dataset} not found in {list(REMOTE_URLS.keys())}"
+
+    dataset_path = os.path.join(HOME_PATH, "datasets")
+    if not os.path.exists(dataset_path):
+        return False
+
+    local_path = os.path.join(dataset_path, dataset)
+    if not os.path.exists(local_path):
+        return False
+
+    for df in REMOTE_URLS[dataset]:
+        local_file = os.path.join(local_path, df.name)
+        if not os.path.exists(local_file):
+            return False
+        with open(local_file, "rb") as f:
+            file_hash = sha256(f.read()).hexdigest()
+            if file_hash != df.hash:
+                return False
+    return True
+
+
+def download_from_github(dataset: str):
+    assert (
+        dataset in REMOTE_URLS
+    ), f"Dataset {dataset} not found in {list(REMOTE_URLS.keys())}"
+
+    dataset_path = os.path.join(HOME_PATH, "datasets")
+    if not os.path.exists(dataset_path):
+        os.makedirs(dataset_path)
+
+    local_path = os.path.join(dataset_path, dataset)
+    if not os.path.exists(local_path):
+        os.makedirs(local_path)
+
+    # download the file
+    console.log(f"Downloading {dataset} from GitHub")
+    for df in track(
+        REMOTE_URLS[dataset],
+        description=f"Downloading {len(REMOTE_URLS[dataset])} files...",
+    ):
+        local_file = os.path.join(local_path, df.name)
+        response = requests.get(df.url)
+        response.raise_for_status()
+
+        file_hash = sha256(response.content).hexdigest()
+        assert (
+            file_hash == df.hash
+        ), f"Hash of {df.name}, {dataset} does not match expected hash"
+
+        with open(local_file, "wb") as f:
+            f.write(response.content)

memory_benchmark/datasets/locomo.py

@@ -0,0 +1,24 @@
+import os
+from dataclasses import dataclass, field
+from rich.console import Console
+
+
+@dataclass
+class Config:
+    llm_base_url: str | None = None
+    llm_api_key: str | None = None
+    llm_model: str = "gpt-4o-mini"
+
+    use_dataset: str = "locomo"
+    use_method: str = "memobase"
+
+    extra_kwargs: dict = field(default_factory=dict)
+
+
+CONFIG = Config()
+
+HOME_PATH = os.path.expanduser("~/.cache/memory_benchmark")
+if not os.path.exists(HOME_PATH):
+    os.makedirs(HOME_PATH)
+
+console = Console()

memory_benchmark/env.py

@@ -0,0 +1,2 @@
+async def run_memory_benchmark():
+    pass

memory_benchmark/launch.py

@@ -0,0 +1,4 @@
+from .openai import openai_complete
+
+
+llm_complete = openai_complete

memory_benchmark/llms/__init__.py

@@ -0,0 +1,35 @@
+from openai import AsyncOpenAI
+from ..env import CONFIG
+from .types import LLMResult
+
+_global_openai_async_client = None
+
+
+def get_openai_async_client_instance() -> AsyncOpenAI:
+    global _global_openai_async_client
+    if _global_openai_async_client is None:
+        _global_openai_async_client = AsyncOpenAI(
+            base_url=CONFIG.llm_base_url,
+            api_key=CONFIG.llm_api_key,
+        )
+    return _global_openai_async_client
+
+
+async def openai_complete(
+    model, prompt, system_prompt=None, history_messages=[], **kwargs
+) -> LLMResult:
+    openai_async_client = get_openai_async_client_instance()
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+
+    response = await openai_async_client.chat.completions.create(
+        model=model, messages=messages, timeout=120, **kwargs
+    )
+    return LLMResult(
+        content=response.choices[0].message.content,
+        input_tokens=response.usage.prompt_tokens,
+        output_tokens=response.usage.completion_tokens,
+    )

memory_benchmark/llms/openai.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+
+class LLMResult(BaseModel):
+    content: str
+    input_tokens: int
+    output_tokens: int

memory_benchmark/llms/types.py

@@ -0,0 +1,31 @@
+<div align="center">
+  <h1><code>memory-benchmark</code></h1>
+  <p><strong>Run benchmark results for long-term memory backends</strong></p>
+  <p>
+    <img src="https://img.shields.io/badge/python->=3.11-blue">
+    <a href="https://pypi.org/project/memory-benchmark/">
+      <img src="https://img.shields.io/pypi/v/memory-benchmark.svg">
+    </a>
+</div>
+
+
+
+## Get Started
+
+TODO
+
+
+
+## Support Datasets
+
+- [ ] [LOCOMO](https://snap-research.github.io/locomo/)
+
+
+
+## Support Methods
+
+- [ ] [Memobase](https://github.com/memodb-io/memobase)
+- [ ] [AgenticMemory](https://github.com/WujiangXu/AgenticMemory/tree/main)
+- [ ] [Mem0](https://github.com/mem0ai/mem0)
+- [ ] [MemGPT](https://github.com/letta-ai/letta)
+- [ ] [Zep](https://github.com/getzep/zep)

memory_benchmark/methods/base.py

@@ -1,3 +1,6 @@
+rich
 openai
+pydantic
+requests
 memobase
 transformers

memory_benchmark/utils.py

@@ -26,7 +26,7 @@
     url=vars2readme["__url__"],
     version=vars2readme["__version__"],
     author=vars2readme["__author__"],
-    description="A unified benchmark framework for long-term memory",
+    description="Run benchmark results for long-term memory backends",
     long_description=long_description,
     long_description_content_type="text/markdown",
     packages=find_packages(exclude=["tests"]),