Implement lazy initialization for embedding function and add API key setter in ChromaDB class.

gnosis/chroma_client.py

@@ -11,13 +11,29 @@ class ChromaDB:
         """Initialize the client."""
         self.client = chromadb.PersistentClient(path=path)
         self.api_key = api_key
-        self.embedding_function = OpenAIEmbeddingFunction(
-            api_key=self.api_key, model_name="text-embedding-3-small"
-        )
+        self._embedding_function = None  # ← lazy, not created yet
         self.client.heartbeat()
 
+    @property
+    def embedding_function(self):
+        """Create embedding function only when API key is available."""
+        if not self.api_key:
+            return None
+        if self._embedding_function is None:
+            self._embedding_function = OpenAIEmbeddingFunction(
+                api_key=self.api_key, model_name="text-embedding-3-small"
+            )
+        return self._embedding_function
+
+    def set_api_key(self, api_key):
+        """Update the API key and reset the embedding function."""
+        self.api_key = api_key
+        self._embedding_function = None  # force recreation with new key
+
     def get_collection(self, name):
         """Get a Chroma collection."""
+        if not self.embedding_function:
+            return st.error("Add your OpenAI API key first.")
         try:
             collection = self.client.get_collection(
                 name=name,
@@ -29,6 +45,8 @@ class ChromaDB:
 
     def create_collection(self, name):
         """Create a Chroma collection."""
+        if not self.embedding_function:
+            return st.error("Add your OpenAI API key first.")
         try:
             collection = self.client.get_or_create_collection(
                 name=name,

gnosis/components/handlers.py

@@ -8,6 +8,8 @@ def set_api_key():
     """Set the OpenAI API key."""
     openai.api_key = st.session_state.api_key
     st.session_state.api_message = gm.api_message(openai.api_key)
+    if "chroma_db" in st.session_state:
+        st.session_state.chroma_db.set_api_key(openai.api_key)
 
 
 def click_wk_button():