feat: Add HNSW vector storage (#18)

* Updated storage with hnswlib, unittests, benchmarking against nano vector storage, and a simple GraphRAG example

* Added kwargs for vector storage cls to pass on hyperparameters for better speed/recall tradeoffs

* Removed redundant passing of  in

---------

Co-authored-by: terence-gpt <numberchiffre@users.noreply.github.com>

Author: NumberChiffre

benchmarks/hnsw_vs_nano_vector_storage.py

@@ -0,0 +1,78 @@
+import asyncio
+import time
+import numpy as np
+from tqdm import tqdm
+from nano_graphrag import GraphRAG
+from nano_graphrag._storage import NanoVectorDBStorage, HNSWVectorStorage
+from nano_graphrag._utils import wrap_embedding_func_with_attrs
+
+
+WORKING_DIR = "./nano_graphrag_cache_benchmark_hnsw_vs_nano_vector_storage"
+DATA_LEN = 100_000
+FAKE_DIM = 1024
+BATCH_SIZE = 100000
+
+
+@wrap_embedding_func_with_attrs(embedding_dim=FAKE_DIM, max_token_size=8192)
+async def sample_embedding(texts: list[str]) -> np.ndarray:
+    return np.float32(np.random.rand(len(texts), FAKE_DIM))
+
+
+def generate_test_data():
+    return {str(i): {"content": f"Test content {i}"} for i in range(DATA_LEN)}
+
+
+async def benchmark_storage(storage_class, name):
+    rag = GraphRAG(working_dir=WORKING_DIR, embedding_func=sample_embedding)
+    storage = storage_class(
+        namespace=f"benchmark_{name}",
+        global_config=rag.__dict__,
+        embedding_func=sample_embedding,
+        meta_fields={"content"},
+    )
+
+    test_data = generate_test_data()
+    
+    print(f"Benchmarking {name}...")
+    with tqdm(total=DATA_LEN, desc=f"{name} Benchmark") as pbar:
+        start_time = time.time()
+        for i in range(0, len(test_data), BATCH_SIZE):
+            batch = {k: test_data[k] for k in list(test_data.keys())[i:i+BATCH_SIZE]}
+            await storage.upsert(batch)
+            pbar.update(min(BATCH_SIZE, DATA_LEN - i))
+        
+        insert_time = time.time() - start_time
+
+        save_start_time = time.time()
+        await storage.index_done_callback()
+        save_time = time.time() - save_start_time
+        pbar.update(1)
+
+        query_vector = np.random.rand(FAKE_DIM)
+        query_times = []
+        for _ in range(100):
+            query_start = time.time()
+            await storage.query(query_vector, top_k=10)
+            query_times.append(time.time() - query_start)
+            pbar.update(1)
+    
+    avg_query_time = sum(query_times) / len(query_times)
+    
+    print(f"{name} - Insert: {insert_time:.2f}s, Save: {save_time:.2f}s, Avg Query: {avg_query_time:.4f}s")
+    return insert_time, save_time, avg_query_time
+
+
+async def run_benchmarks():
+    print("Running NanoVectorDB benchmark...")
+    nano_insert_time, nano_save_time, nano_query_time = await benchmark_storage(NanoVectorDBStorage, "nano")
+    
+    print("\nRunning HNSWVectorStorage benchmark...")
+    hnsw_insert_time, hnsw_save_time, hnsw_query_time = await benchmark_storage(HNSWVectorStorage, "hnsw")
+    
+    print("\nBenchmark Results:")
+    print(f"NanoVectorDB - Insert: {nano_insert_time:.2f}s, Save: {nano_save_time:.2f}s, Avg Query: {nano_query_time:.4f}s")
+    print(f"HNSWVectorStorage - Insert: {hnsw_insert_time:.2f}s, Save: {hnsw_save_time:.2f}s, Avg Query: {hnsw_query_time:.4f}s")
+
+
+if __name__ == "__main__":
+    asyncio.run(run_benchmarks())

examples/using_hnsw_as_vectorDB.py

@@ -0,0 +1,61 @@
+import os
+from nano_graphrag import GraphRAG, QueryParam
+from nano_graphrag._llm import gpt_4o_mini_complete
+from nano_graphrag._storage import HNSWVectorStorage
+
+
+WORKING_DIR = "./nano_graphrag_cache_using_hnsw_as_vectorDB"
+
+
+def remove_if_exist(file):
+    if os.path.exists(file):
+        os.remove(file)
+
+
+def insert():
+    from time import time
+
+    with open("./tests/mock_data.txt", encoding="utf-8-sig") as f:
+        FAKE_TEXT = f.read()
+
+    remove_if_exist(f"{WORKING_DIR}/vdb_entities.json")
+    remove_if_exist(f"{WORKING_DIR}/kv_store_full_docs.json")
+    remove_if_exist(f"{WORKING_DIR}/kv_store_text_chunks.json")
+    remove_if_exist(f"{WORKING_DIR}/kv_store_community_reports.json")
+    remove_if_exist(f"{WORKING_DIR}/graph_chunk_entity_relation.graphml")
+    rag = GraphRAG(
+        working_dir=WORKING_DIR,
+        enable_llm_cache=True,
+        vector_db_storage_cls=HNSWVectorStorage,
+        vector_db_storage_cls_kwargs={"max_elements": 1000000, "ef_search": 100, "M": 8},
+        best_model_max_async=1,
+        cheap_model_max_async=1,
+        best_model_func=gpt_4o_mini_complete,
+        cheap_model_func=gpt_4o_mini_complete,
+    )
+    start = time()
+    rag.insert(FAKE_TEXT)
+    print("indexing time:", time() - start)
+
+
+def query():
+    rag = GraphRAG(
+        working_dir=WORKING_DIR,
+        enable_llm_cache=True,
+        vector_db_storage_cls=HNSWVectorStorage,
+        vector_db_storage_cls_kwargs={"max_elements": 1000000, "ef_search": 100, "M": 8},
+        best_model_max_async=1,
+        cheap_model_max_async=1,
+        best_model_func=gpt_4o_mini_complete,
+        cheap_model_func=gpt_4o_mini_complete,
+    )
+    print(
+        rag.query(
+            "What are the top themes in this story?", param=QueryParam(mode="global")
+        )
+    )
+
+
+if __name__ == "__main__":
+    insert()
+    query()

nano_graphrag/_storage.py

@@ -3,9 +3,10 @@
 import json
 import os
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Union, cast
-
+import pickle
+import hnswlib
 import networkx as nx
 import numpy as np
 from nano_vectordb import NanoVectorDB
@@ -115,6 +116,110 @@ async def index_done_callback(self):
         self._client.save()
 
 
+@dataclass
+class HNSWVectorStorage(BaseVectorStorage):
+    ef_construction: int = 100
+    M: int = 16
+    max_elements: int = 1000000
+    ef_search: int = 50
+    num_threads: int = -1
+    _index: Any = field(init=False)
+    _metadata: dict[str, dict] = field(default_factory=dict)
+    _current_elements: int = 0
+
+    def __post_init__(self):
+        self._index_file_name = os.path.join(
+            self.global_config["working_dir"], f"{self.namespace}_hnsw.index"
+        )
+        self._metadata_file_name = os.path.join(
+            self.global_config["working_dir"], f"{self.namespace}_hnsw_metadata.pkl"
+        )
+        self._max_batch_size = self.global_config.get("embedding_batch_num", 100)
+
+        hnsw_params = self.global_config.get("vector_db_storage_cls_kwargs", {})
+        self.ef_construction = hnsw_params.get("ef_construction", self.ef_construction)
+        self.M = hnsw_params.get("M", self.M)
+        self.max_elements = hnsw_params.get("max_elements", self.max_elements)
+        self.ef_search = hnsw_params.get("ef_search", self.ef_search)
+        self.num_threads = hnsw_params.get("num_threads", self.num_threads)
+
+        if os.path.exists(self._index_file_name) and os.path.exists(self._metadata_file_name):
+            self._index = hnswlib.Index(space='cosine', dim=self.embedding_func.embedding_dim)
+            self._index.load_index(self._index_file_name, max_elements=self.max_elements)
+            with open(self._metadata_file_name, 'rb') as f:
+                self._metadata, self._current_elements = pickle.load(f)
+            logger.info(f"Loaded existing index for {self.namespace} with {self._current_elements} elements")
+        else:
+            self._index = hnswlib.Index(space='cosine', dim=self.embedding_func.embedding_dim)
+            self._index.init_index(
+                max_elements=self.max_elements,
+                ef_construction=self.ef_construction,
+                M=self.M
+            )
+            self._index.set_ef(self.ef_search)
+            logger.info(f"Created new index for {self.namespace}")
+
+    async def upsert(self, data: dict[str, dict]):
+        logger.info(f"Inserting {len(data)} vectors to {self.namespace}")
+        if not data:
+            raise ValueError("Attempting to insert empty data to vector DB")
+
+        if self._current_elements + len(data) > self.max_elements:
+            raise ValueError(f"Cannot insert {len(data)} elements. Current: {self._current_elements}, Max: {self.max_elements}")
+
+        contents = [v["content"] for v in data.values()]
+        batches = [
+            contents[i : i + self._max_batch_size]
+            for i in range(0, len(contents), self._max_batch_size)
+        ]
+        embeddings_list = await asyncio.gather(
+            *[self.embedding_func(batch) for batch in batches]
+        )
+        embeddings = np.concatenate(embeddings_list)
+
+        ids = []
+        for id, item in data.items():
+            metadata = {k: v for k, v in item.items() if k in self.meta_fields}
+            metadata['id'] = id
+            self._metadata[id] = metadata
+            ids.append(int(id) if id.isdigit() else hash(id))
+
+        ids = np.array(ids)
+        self._index.add_items(data=embeddings, ids=ids, num_threads=self.num_threads)
+        self._current_elements += len(data)
+
+    async def query(self, query: str, top_k: int = 5) -> list[dict]:
+        if len(self._metadata) == 0:
+            return []
+        
+        if top_k >= self.ef_search:
+            raise ValueError(f"top_k must be greater than or equal to ef_search, got {top_k} and {self.ef_search}")
+
+        query_vector = await self.embedding_func([query])
+        labels, distances = self._index.knn_query(
+            data=query_vector, 
+            k=min(top_k, len(self._metadata)), 
+            num_threads=self.num_threads
+        )
+        
+        results = []
+        for label, distance in zip(labels[0], distances[0]):
+            id_str = str(label)
+            if id_str in self._metadata:
+                metadata = self._metadata[id_str]
+                results.append({
+                    **metadata,
+                    "distance": distance,
+                    "similarity": 1 - distance
+                })
+        return results
+
+    async def index_done_callback(self):
+        self._index.save_index(self._index_file_name)
+        with open(self._metadata_file_name, 'wb') as f:
+            pickle.dump((self._metadata, self._current_elements), f)
+
+
 @dataclass
 class NetworkXStorage(BaseGraphStorage):
     @staticmethod

nano_graphrag/graphrag.py

@@ -104,6 +104,7 @@ class GraphRAG:
     # storage
     key_string_value_json_storage_cls: Type[BaseKVStorage] = JsonKVStorage
     vector_db_storage_cls: Type[BaseVectorStorage] = NanoVectorDBStorage
+    vector_db_storage_cls_kwargs: dict = field(default_factory=dict)
     graph_storage_cls: Type[BaseGraphStorage] = NetworkXStorage
     enable_llm_cache: bool = True
 
@@ -150,7 +151,7 @@ def __post_init__(self):
                 namespace="entities",
                 global_config=asdict(self),
                 embedding_func=self.embedding_func,
-                meta_fields={"entity_name"},
+                meta_fields={"entity_name"}
             )
             if self.enable_local
             else None

requirements.txt

@@ -2,4 +2,5 @@ openai
 tiktoken
 networkx
 graspologic
-nano-vectordb
+nano-vectordb
+hnswlib