Cache push_to_cosmos embeddings; clarify non-turn thread_id

- push_to_cosmos (sync + async): persist generated embeddings back to
  local_memory so a repeat push doesn't re-embed already-embedded
  non-turn records.
- Update test comment: _make_memory auto-generates a UUID thread_id for
  non-turn types when omitted (it does not user-scope them).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: 2 people

agent_memory_toolkit/aio/cosmos_memory_client.py

@@ -758,6 +758,9 @@ async def push_to_cosmos(self, batch_size: int = 25) -> None:
                     vectors = await self._embeddings_client.generate_batch(to_embed_text)
                     for i, vec in zip(to_embed_idx, vectors):
                         bodies[i]["embedding"] = vec
+                        # Persist the embedding back to local_memory so a
+                        # repeat push_to_cosmos() doesn't re-embed.
+                        self.local_memory[start + i]["embedding"] = vec
                 except Exception as exc:  # noqa: BLE001
                     logger.warning(
                         "push_to_cosmos: batch embedding generation failed (%s); "

agent_memory_toolkit/cosmos_memory_client.py

@@ -995,6 +995,9 @@ def push_to_cosmos(self, batch_size: int = 25) -> None:
                     vectors = self._embeddings_client.generate_batch(to_embed_text)
                     for i, vec in zip(to_embed_idx, vectors):
                         bodies[i]["embedding"] = vec
+                        # Persist the embedding back to local_memory so a
+                        # repeat push_to_cosmos() doesn't re-embed.
+                        self.local_memory[start + i]["embedding"] = vec
                 except Exception as exc:  # noqa: BLE001
                     logger.warning(
                         "push_to_cosmos: batch embedding generation failed (%s); "

tests/unit/test_cosmos_memory_client.py

@@ -138,7 +138,8 @@ def test_add_local_turn_requires_thread_id(self):
 
     def test_add_local_non_turn_thread_id_optional(self):
         mem = _make_client()
-        # Non-turn types (summary, fact, etc.) are user-scoped; thread_id is optional.
+        # Non-turn types (summary, fact, etc.) accept an omitted thread_id;
+        # _make_memory auto-generates a UUID for hierarchical-PK validity.
         mem.add_local(user_id="u1", role="user", content="profile", memory_type="user_summary")
         assert len(mem.local_memory) == 1
 
@@ -444,6 +445,26 @@ def _generate_batch(texts: list[str]) -> list[list[float]]:
         assert fact_body["embedding"] == [0.1, 0.2, 0.3]
         assert "embedding" not in turn_body
 
+    def test_push_to_cosmos_caches_embeddings_in_local_memory(self):
+        """Repeat push_to_cosmos() must not re-embed the same non-turn records."""
+        mem, container = _connected_client()
+        embed_calls: list[list[str]] = []
+
+        def _generate_batch(texts: list[str]) -> list[list[float]]:
+            embed_calls.append(list(texts))
+            return [[0.5, 0.6, 0.7] for _ in texts]
+
+        mem._embeddings_client = MagicMock()
+        mem._embeddings_client.generate_batch.side_effect = _generate_batch
+
+        mem.add_local(user_id="u1", role="user", content="fact one", memory_type="fact")
+        mem.push_to_cosmos()
+        mem.push_to_cosmos()
+
+        # Second push should not re-embed — embedding is cached on local_memory.
+        assert embed_calls == [["fact one"]]
+        assert mem.local_memory[0]["embedding"] == [0.5, 0.6, 0.7]
+
 
 class TestGetMemories:
     def test_no_filters(self):