From 5f85cc532dfd5db70cb9c627c3d5609e98739db5 Mon Sep 17 00:00:00 2001
From: vegu-ai-tools <152010387+vegu-ai-tools@users.noreply.github.com>
Date: Tue, 14 Apr 2026 23:57:39 +0300
Subject: [PATCH] fix: optimize model eviction logic in ChromaDBMemoryAgent

---
 src/talemate/agents/memory/__init__.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/talemate/agents/memory/__init__.py b/src/talemate/agents/memory/__init__.py
index 9a6e60ae..951b62b9 100644
--- a/src/talemate/agents/memory/__init__.py
+++ b/src/talemate/agents/memory/__init__.py
@@ -968,11 +968,17 @@ class ChromaDBMemoryAgent(MemoryAgent):
             )
 
             try:
-                # Evict cached model so device changes take effect without restart.
                 # ChromaDB caches models in a class-level dict keyed only by model name,
                 # ignoring device — so a device switch would silently reuse the old model.
+                # Evict only when the cached model's device differs from the requested one,
+                # otherwise reuse it (avoids an expensive reload on every set_db call).
                 ST = embedding_functions.SentenceTransformerEmbeddingFunction
-                if model_name in ST.models:
+                cached = ST.models.get(model_name)
+                if (
+                    cached is not None
+                    and getattr(cached, "device", None) is not None
+                    and cached.device.type != device
+                ):
                     self._release_embedding_model(ST.models.pop(model_name))
 
                 ef = ST(