pg-vector support based on Asyncpg

nano_graphrag/storage/asyncpg.py

@@ -0,0 +1,100 @@
+from nano_graphrag._storage import BaseVectorStorage
+import asyncpg
+import asyncio
+from contextlib import asynccontextmanager
+from nano_graphrag._utils import logger
+from pgvector.asyncpg import register_vector
+from nano_graphrag.graphrag import always_get_an_event_loop
+import numpy as np
+import json
+
+import nest_asyncio
+nest_asyncio.apply()
+
+class AsyncpgVectorStorage(BaseVectorStorage):
+    table_name_generator: callable = None
+    conn_fetcher: callable = None
+    cosine_better_than_threshold: float = 0.2
+    dsn = None
+    def __init__(self, dsn: str = None, conn_fetcher: callable = None, table_name_generator: callable = None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.dsn = dsn
+        self.conn_fetcher = conn_fetcher
+        assert self.dsn != None or self.conn_fetcher != None, "Must provide either dsn or conn_fetcher"
+        if self.dsn:
+            self.conn_fetcher = self.__get_conn
+        if not table_name_generator:
+            self.table_name_generator = lambda working_dir, namespace: f'{working_dir}_{namespace}_vdb'
+        self._table_name = self.table_name_generator(self.global_config["working_dir"], self.namespace)
+        self._max_batch_size = self.global_config["embedding_batch_num"]
+
+        self.cosine_better_than_threshold = self.global_config.get(
+            "query_better_than_threshold", self.cosine_better_than_threshold
+        )
+        loop = always_get_an_event_loop()
+        loop.run_until_complete(self._secure_table())
+    @asynccontextmanager
+    async def __get_conn(self):
+        try:
+            conn: asyncpg.Connection = await asyncpg.connect(self.dsn)
+            await register_vector(conn)
+            yield conn
+        finally:
+            await conn.close()
+    async def _secure_table(self):
+        async with self.conn_fetcher() as conn:
+            conn: asyncpg.Connection
+            await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
+            result = await conn.fetch(
+            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)", self._table_name)
+            table_exists = result[0]['exists']
+            if not table_exists:
+                # create the table
+                await conn.execute(f'CREATE TABLE {self._table_name} (id text PRIMARY KEY, embedding vector({self.embedding_func.embedding_dim}), data jsonb)')
+                await conn.execute(f'CREATE INDEX ON {self._table_name} USING hnsw (embedding vector_cosine_ops)')
+    async def query(self, query: str, top_k: int) -> list[dict]:
+        embedding = await self.embedding_func([query])
+        embedding = embedding[0]
+        async with self.conn_fetcher() as conn:
+
+            result = await conn.fetch(f'SELECT embedding <=> $1 as similarity, id, embedding, data FROM {self._table_name} WHERE embedding <=> $1 > $3 ORDER BY embedding <=> $1 DESC LIMIT $2', embedding, top_k, self.cosine_better_than_threshold)
+
+            rows = []
+            for row in result:
+                data = json.loads(row['data'])
+                rows.append({
+                    **data,
+                    'id': row['id'],
+                    'distance': 1 - row['similarity'],
+                    'similarity': row['similarity']
+                })
+            return rows
+    async def upsert(self, data: dict[str, dict]):
+        logger.info(f"Inserting {len(data)} vectors to {self.namespace}")
+        if not len(data):
+            logger.warning("You insert an empty data to vector DB")
+            return []
+        list_data = [
+            {
+                "__id__": k,
+                **{k1: v1 for k1, v1 in v.items() if k1 in self.meta_fields},
+            }
+            for k, v in data.items()
+        ]
+        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_list = np.concatenate(embeddings_list)
+        insert_rows = []
+        for i, d in enumerate(list_data):
+            row = [d["__id__"], embeddings_list[i], json.dumps(d)]
+            insert_rows.append(row)
+        async with self.conn_fetcher() as conn:
+            conn: asyncpg.Connection
+            stmt = f"INSERT INTO {self._table_name} (id, embedding, data) VALUES ($1, $2, $3) ON CONFLICT (id) DO UPDATE SET embedding = $2, data = $3"
+            return await conn.executemany(stmt, insert_rows)

requirements.txt

@@ -7,3 +7,6 @@ hnswlib
 xxhash
 tenacity
 dspy-ai
+pgvector==0.3.3
+asyncpg==0.29.0
+nest_asyncio==1.6.0

tests/test_asyncpg_vector_storage.py

@@ -0,0 +1,204 @@
+import numpy as np
+import pytest
+from dataclasses import asdict
+from nano_graphrag import GraphRAG
+from nano_graphrag._utils import wrap_embedding_func_with_attrs
+
+from nano_graphrag.storage.asyncpg import AsyncpgVectorStorage
+import asyncpg
+from nano_graphrag.graphrag import always_get_an_event_loop
+WORKING_DIR = "nano_graphrag_cache_asyncpg_vector_storage_test"
+dsn='postgresql://username:[email protected]:12345/db'
+
+@pytest.fixture(scope="function")
+def setup_teardown():
+
+    yield
+    loop = always_get_an_event_loop()
+    async def clean_table():
+        conn: asyncpg.Connection = await asyncpg.connect(dsn)
+        async with conn.transaction():
+            tables = await conn.fetch(
+                f"SELECT table_name FROM information_schema.tables WHERE table_name LIKE '{WORKING_DIR}%'"
+            )
+
+            for table in tables:
+                await conn.execute(f"DROP TABLE {table['table_name']} CASCADE")
+    loop.run_until_complete(clean_table())
+
+
+@wrap_embedding_func_with_attrs(embedding_dim=384, max_token_size=8192)
+async def mock_embedding(texts: list[str]) -> np.ndarray:
+    return np.random.rand(len(texts), 384)
+
+
+@pytest.fixture
+def asyncpg_storage(setup_teardown):
+    rag = GraphRAG(working_dir=WORKING_DIR, embedding_func=mock_embedding)
+    return AsyncpgVectorStorage(
+        namespace="test",
+        global_config=asdict(rag),
+        embedding_func=mock_embedding,
+        meta_fields={"entity_name"},
+        dsn=dsn
+    )
+
+
+@pytest.mark.asyncio
+async def test_upsert_and_query(asyncpg_storage):
+    test_data = {
+        "1": {"content": "Test content 1", "entity_name": "Entity 1"},
+        "2": {"content": "Test content 2", "entity_name": "Entity 2"},
+    }
+
+    await asyncpg_storage.upsert(test_data)
+
+    results = await asyncpg_storage.query("Test query", top_k=2)
+
+    assert len(results) == 2
+    assert all(isinstance(result, dict) for result in results)
+    assert all(
+        "id" in result and "distance" in result and "similarity" in result
+        for result in results
+    )
+
+
+@pytest.mark.asyncio
+async def test_persistence(setup_teardown):
+    rag = GraphRAG(working_dir=WORKING_DIR, embedding_func=mock_embedding)
+    initial_storage = AsyncpgVectorStorage(
+        namespace="test",
+        global_config=asdict(rag),
+        embedding_func=mock_embedding,
+        meta_fields={"entity_name"},
+        dsn=dsn
+    )
+
+    test_data = {
+        "1": {"content": "Test content 1", "entity_name": "Entity 1"},
+    }
+
+    await initial_storage.upsert(test_data)
+    await initial_storage.index_done_callback()
+
+    new_storage = AsyncpgVectorStorage(
+        namespace="test",
+        global_config=asdict(rag),
+        embedding_func=mock_embedding,
+        meta_fields={"entity_name"},
+        dsn=dsn
+    )
+
+    results = await new_storage.query("Test query", top_k=1)
+
+    assert len(results) == 1
+    assert results[0]["id"] == "1"
+    assert "entity_name" in results[0]
+
+
+@pytest.mark.asyncio
+async def test_persistence_large_dataset(setup_teardown):
+    rag = GraphRAG(working_dir=WORKING_DIR, embedding_func=mock_embedding)
+    initial_storage = AsyncpgVectorStorage(
+        namespace="test_large",
+        global_config=asdict(rag),
+        embedding_func=mock_embedding,
+        meta_fields={"entity_name"},
+        dsn=dsn
+    )
+
+    large_data = {
+        str(i): {"content": f"Test content {i}", "entity_name": f"Entity {i}"}
+        for i in range(1000)
+    }
+    await initial_storage.upsert(large_data)
+    await initial_storage.index_done_callback()
+
+    new_storage = AsyncpgVectorStorage(
+        namespace="test_large",
+        global_config=asdict(rag),
+        embedding_func=mock_embedding,
+        meta_fields={"entity_name"},
+        dsn=dsn
+    )
+
+    results = await new_storage.query("Test query", top_k=500)
+    assert len(results) == 500
+    assert all(result["id"] in large_data for result in results)
+
+
+@pytest.mark.asyncio
+async def test_upsert_with_existing_ids(asyncpg_storage):
+    test_data = {
+        "1": {"content": "Test content 1", "entity_name": "Entity 1"},
+        "2": {"content": "Test content 2", "entity_name": "Entity 2"},
+    }
+
+    await asyncpg_storage.upsert(test_data)
+
+    updated_data = {
+        "1": {"content": "Updated content 1", "entity_name": "Updated Entity 1"},
+        "3": {"content": "Test content 3", "entity_name": "Entity 3"},
+    }
+
+    await asyncpg_storage.upsert(updated_data)
+
+    results = await asyncpg_storage.query("Updated", top_k=3)
+
+    assert len(results) == 3
+    assert any(
+        result["id"] == "1" and result["entity_name"] == "Updated Entity 1"
+        for result in results
+    )
+    assert any(
+        result["id"] == "2" and result["entity_name"] == "Entity 2"
+        for result in results
+    )
+    assert any(
+        result["id"] == "3" and result["entity_name"] == "Entity 3"
+        for result in results
+    )
+
+
+@pytest.mark.asyncio
+async def test_large_batch_upsert(asyncpg_storage):
+    batch_size = 30
+    large_data = {
+        str(i): {"content": f"Test content {i}", "entity_name": f"Entity {i}"}
+        for i in range(batch_size)
+    }
+
+    await asyncpg_storage.upsert(large_data)
+
+    results = await asyncpg_storage.query("Test query", top_k=batch_size)
+    assert len(results) == batch_size
+    assert all(isinstance(result, dict) for result in results)
+    assert all(
+        "id" in result and "distance" in result and "similarity" in result
+        for result in results
+    )
+
+
+@pytest.mark.asyncio
+async def test_empty_data_insertion(asyncpg_storage):
+    empty_data = {}
+    await asyncpg_storage.upsert(empty_data)
+
+    results = await asyncpg_storage.query("Test query", top_k=1)
+    assert len(results) == 0
+
+
+@pytest.mark.asyncio
+async def test_query_with_no_results(asyncpg_storage):
+    results = await asyncpg_storage.query("Non-existent query", top_k=5)
+    assert len(results) == 0
+
+    test_data = {
+        "1": {"content": "Test content 1", "entity_name": "Entity 1"},
+    }
+    await asyncpg_storage.upsert(test_data)
+
+    results = await asyncpg_storage.query("Non-existent query", top_k=5)
+    assert len(results) == 1
+    assert all(0 <= result["similarity"] <= 1 for result in results)
+    assert "entity_name" in results[0]