Handle caching for spaCy embeddings

Author: NathalieCharbel

src/neo4j_graphrag/experimental/components/resolver.py

@@ -334,15 +334,21 @@ def __init__(
             neo4j_database,
         )
         self.nlp = self._load_or_download_spacy_model(spacy_model)
+        self.embedding_cache: dict[str, np.ndarray] = {}
 
     def compute_similarity(self, text_a: str, text_b: str) -> float:
-        emb1 = self.nlp(text_a).vector
-        emb2 = self.nlp(text_b).vector
+        emb1 = self._get_embedding(text_a)
+        emb2 = self._get_embedding(text_b)
         sim = self._cosine_similarity(
             np.asarray(emb1, dtype=np.float64), np.asarray(emb2, dtype=np.float64)
         )
         return sim
 
+    def _get_embedding(self, text: str) -> np.ndarray:
+        if text not in self.embedding_cache:
+            self.embedding_cache[text] = self.nlp(text).vector
+        return self.embedding_cache[text]
+
     @staticmethod
     def _cosine_similarity(
         vec1: NDArray[np.float64], vec2: NDArray[np.float64]

tests/unit/experimental/components/test_resolver.py

@@ -12,7 +12,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-from unittest.mock import MagicMock, call
+from unittest.mock import MagicMock, call, patch
 
 import pytest
 from neo4j_graphrag.experimental.components.resolver import (
@@ -237,3 +237,52 @@ async def test_fuzzy_match_resolver_normalization(driver: MagicMock) -> None:
 
     sim = resolver.compute_similarity("  ALICE  ", "alice!")
     assert sim == 1
+
+
+@pytest.mark.asyncio
+async def test_spacy_resolver_caching(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 1, "name": "Alice"},
+                            {"id": 2, "name": "Alice"},
+                            {"id": 3, "name": "Bob"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 1})],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 3})],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = SpaCySemanticMatchResolver(driver=driver)
+
+    # patch spaCy NLP call to track how often embeddings are computed
+    with patch.object(resolver, "nlp", wraps=resolver.nlp) as mock_nlp:
+        await resolver.run()
+
+    # "Alice" should be embedded only once, despite being used twice.
+    # "Bob" should be embedded once.
+    assert mock_nlp.call_count == 2, (
+        f"Expected spaCy to embed each unique text once. Got {mock_nlp.call_count} "
+        f"calls."
+    )
+
+    # "Alice" and "Bob" are expected to be the only two distinct texts passed to spaCy.
+    called_texts = {call.args[0] for call in mock_nlp.call_args_list}
+    assert called_texts == {"Alice", "Bob"}