Fix the round trips

integrations/lightrag-memgraph/src/lightrag_memgraph/docstatus_impl.py

@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 import os
 from dataclasses import dataclass
@@ -276,26 +277,35 @@ async def get_docs_paginated(
         params.update(filter_params)
 
         driver = await get_driver()
-        async with driver.session(database=get_database(), default_access_mode="READ") as session:
-            count_result = await session.run(
-                f"MATCH (n:`{self._label}`) {where} RETURN count(n) AS total",
-                **filter_params,
-            )
-            count_record = await count_result.single()
-            await count_result.consume()
-            total = count_record["total"] if count_record else 0
 
-            page_result = await session.run(
-                f"""
-                MATCH (n:`{self._label}`) {where}
-                RETURN n.id AS id, n.data AS data
-                ORDER BY {sort_expr} {order}
-                SKIP $skip LIMIT $limit
-                """,
-                **params,
-            )
-            records = [record async for record in page_result]
-            await page_result.consume()
+        async def _count() -> int:
+            async with driver.session(database=get_database(), default_access_mode="READ") as session:
+                result = await session.run(
+                    f"MATCH (n:`{self._label}`) {where} RETURN count(n) AS total",
+                    **filter_params,
+                )
+                record = await result.single()
+                await result.consume()
+                return record["total"] if record else 0
+
+        async def _page() -> list:
+            async with driver.session(database=get_database(), default_access_mode="READ") as session:
+                result = await session.run(
+                    f"""
+                    MATCH (n:`{self._label}`) {where}
+                    RETURN n.id AS id, n.data AS data
+                    ORDER BY {sort_expr} {order}
+                    SKIP $skip LIMIT $limit
+                    """,
+                    **params,
+                )
+                records = [record async for record in result]
+                await result.consume()
+                return records
+
+        # No data dependency between the count and the page -- run them concurrently
+        # instead of two sequential round trips.
+        total, records = await asyncio.gather(_count(), _page())
 
         docs: list[tuple[str, DocProcessingStatus]] = []
         for record in records:

integrations/lightrag-memgraph/src/lightrag_memgraph/kv_impl.py

@@ -121,20 +121,29 @@ async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
         if not data:
             return
         current_time = int(time.time())
-        # create_time only set for genuinely new records (matches JsonKVStorage semantics).
-        existing = set(data.keys()) - await self.filter_keys(set(data.keys()))
+        # MERGE picks new vs. existing payload via ON CREATE/ON MATCH -- no separate filter_keys() round trip.
         entries = []
         for k, v in data.items():
             value = dict(v)
             if self.namespace.endswith("text_chunks") and "llm_cache_list" not in value:
                 value["llm_cache_list"] = []
-            if k in existing:
-                value["update_time"] = current_time
-            else:
-                value["create_time"] = current_time
-                value["update_time"] = current_time
             value["_id"] = k
-            entries.append({"id": k, "data": json.dumps(value, ensure_ascii=False), "ts": current_time})
+
+            new_value = dict(value)
+            new_value["create_time"] = current_time
+            new_value["update_time"] = current_time
+
+            update_value = dict(value)
+            update_value["update_time"] = current_time
+
+            entries.append(
+                {
+                    "id": k,
+                    "new_data": json.dumps(new_value, ensure_ascii=False),
+                    "update_data": json.dumps(update_value, ensure_ascii=False),
+                    "ts": current_time,
+                }
+            )
 
         driver = await get_driver()
         async with driver.session(database=get_database()) as session:
@@ -143,8 +152,9 @@ async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
                     f"""
                     UNWIND $entries AS e
                     MERGE (n:`{self._label}` {{id: e.id}})
-                    ON CREATE SET n.created_at = e.ts
-                    SET n.data = e.data, n.updated_at = e.ts
+                    ON CREATE SET n.data = e.new_data, n.created_at = e.ts
+                    ON MATCH SET n.data = e.update_data
+                    SET n.updated_at = e.ts
                     """,
                     entries=entries,
                 )