Fix embeddings attach (#1804)

Author: ahuang11

lumen/ai/vector_store.py

@@ -19,7 +19,7 @@
 import param
 import semchunk
 
-from panel import cache as pn_cache
+from panel import cache as pn_cache, state as pn_state
 from tqdm.auto import tqdm
 
 from .actor import PROMPTS_DIR, LLMUser
@@ -1116,10 +1116,22 @@ class DuckDBVectorStore(VectorStore):
         doc="Embeddings object for text processing. If None and a URI is provided, loads from the database; else NumpyEmbeddings.",
     )
 
+    # Class-level registry of connections by URI. Used to close stale
+    # connections on reload (e.g. Panel hot-reload) so the file handle is
+    # released and a fresh ATTACH can succeed.
+    _uri_connections: t.ClassVar[dict[str, duckdb.DuckDBPyConnection]] = {}
+
     def __init__(self, **params):
         super().__init__(**params)
         self._add_items_lock = asyncio.Lock()
 
+        # Reuse existing connection for this URI if still open
+        # (e.g. on Panel hot-reload within the same process).
+        if self.uri != ":memory:" and self.uri in DuckDBVectorStore._uri_connections:
+            self.connection = DuckDBVectorStore._uri_connections[self.uri]
+            self._resolve_state()
+            return
+
         connection = duckdb.connect(":memory:")
         # following the instructions from
         # https://duckdb.org/docs/stable/extensions/vss.html#persistence
@@ -1133,43 +1145,23 @@ def __init__(self, **params):
             if self.embeddings is None:
                 self.embeddings = NumpyEmbeddings()
             return
-        uri_exists = Path(self.uri).exists()
-        direct_connection = False
-        try:
-            attach_mode = "READ_ONLY" if self.read_only else "READ_WRITE"
-            connection.execute(f"ATTACH DATABASE '{self.uri}' AS embedded ({attach_mode});")
-        except (duckdb.BinderException, duckdb.IOException) as e:
-            err_msg = str(e).lower()
-            if "already attached" in err_msg or "unique file handle" in err_msg:
-                # File already held by another DuckDB connection in this process
-                # (e.g., during Panel hot-reload). Connect directly to the file instead.
-                connection.close()
-                connection = duckdb.connect(self.uri, read_only=self.read_only)
-                connection.execute("LOAD 'vss';")
-                connection.execute("SET hnsw_enable_experimental_persistence = true;")
-                direct_connection = True
-            else:
-                raise
-        except duckdb.CatalogException:
-            # handle "Failure while replaying WAL file"
-            # remove .wal uri on corruption
-            wal_path = Path(str(self.uri) + ".wal")
-            if wal_path.exists():
-                wal_path.unlink()
-            attach_mode = "READ_ONLY" if self.read_only else "READ_WRITE"
-            connection.execute(f"ATTACH DATABASE '{self.uri}' AS embedded ({attach_mode});")
-        if not direct_connection:
-            connection.execute("USE embedded;")
-        self.connection = connection
-        has_documents = (
-            connection.execute(
+
+        attach_mode = "READ_ONLY" if self.read_only else "READ_WRITE"
+        connection.execute(f"ATTACH DATABASE '{self.uri}' AS embedded ({attach_mode});")
+        connection.execute("USE embedded;")
+        DuckDBVectorStore._uri_connections[self.uri] = self.connection = connection
+        pn_state.on_session_destroyed(lambda session_context: self.close())
+        self._resolve_state()
+
+    def _resolve_state(self):
+        """Set _initialized and resolve embeddings from the active connection."""
+        self._initialized = (
+            self.connection.execute(
                 "SELECT COUNT(*) FROM information_schema.tables WHERE table_name = 'documents';"
             ).fetchone()[0]
             > 0
-        )
-        self._initialized = uri_exists and has_documents
-
-        if self.uri != ":memory:" and self._initialized:
+        ) and Path(self.uri).exists()
+        if self._initialized:
             config = self._get_embeddings_config()
             if config and self.embeddings is None:
                 module_name, class_name = config["class"].rsplit(".", 1)
@@ -1178,7 +1170,6 @@ def __init__(self, **params):
                 self.embeddings = embedding_class(**config["params"])
                 log_debug(f"Loaded embeddings {class_name} from database.")
             self._check_embeddings_consistency()
-
         if self.embeddings is None:
             self.embeddings = NumpyEmbeddings()
 
@@ -1674,6 +1665,7 @@ def close(self) -> None:
         if self.connection:
             self.connection.close()
             self.connection = None
+        DuckDBVectorStore._uri_connections.pop(self.uri, None)
 
 
 class ChromaDBVectorStore(VectorStore):

lumen/tests/ai/test_vector_store.py

@@ -644,6 +644,38 @@ async def test_api_key_not_stored_in_metadata(self, tmp_path):
         assert "api_key" not in params, "api_key should not be stored in metadata"
         store.close()
 
+    @pytest.mark.asyncio
+    async def test_reuse_connection_on_same_uri(self, tmp_path):
+        """Verifies that creating a second DuckDBVectorStore with the same URI
+        reuses the existing connection instead of raising an ATTACH error."""
+        db_path = str(tmp_path / "test_reuse.db")
+        store1 = DuckDBVectorStore(uri=db_path, embeddings=NumpyEmbeddings())
+        await store1.add([{"text": "Reuse test doc", "metadata": {"n": 1}}])
+
+        # Second store with the same URI — should reuse, not crash.
+        store2 = DuckDBVectorStore(uri=db_path, embeddings=NumpyEmbeddings())
+        assert store2.connection is store1.connection
+
+        results = await store2.query("Reuse test doc")
+        assert len(results) == 1
+        assert results[0]["text"] == "Reuse test doc"
+        store2.close()
+
+    @pytest.mark.asyncio
+    async def test_fresh_connection_after_close(self, tmp_path):
+        """Verifies that after close(), a new store gets a fresh connection
+        and can still read persisted data."""
+        db_path = str(tmp_path / "test_fresh.db")
+        store1 = DuckDBVectorStore(uri=db_path, embeddings=NumpyEmbeddings())
+        await store1.add([{"text": "Persist me"}])
+        store1.close()
+
+        store2 = DuckDBVectorStore(uri=db_path, embeddings=NumpyEmbeddings())
+        assert store2.connection is not None
+        results = await store2.query("Persist me")
+        assert len(results) == 1
+        store2.close()
+
 
 # Sample readme content for testing (mimics real-world metadata files)
 SAMPLE_README = """# Population - Data package