Enable embedding_function in RetrieveUserProxyAgent to create initial…

… vector_db
microsoft · Nov 14, 2024 · 59088fa · 59088fa
1 parent 8a8fcd8
commit 59088fa
Showing 1 changed file with 18 additions and 5 deletions.
diff --git a/autogen/agentchat/contrib/retrieve_user_proxy_agent.py b/autogen/agentchat/contrib/retrieve_user_proxy_agent.py
@@ -379,11 +379,24 @@ def _init_db(self):
             )
             chunk_ids_set = set(chunk_ids)
             chunk_ids_set_idx = [chunk_ids.index(hash_value) for hash_value in chunk_ids_set]
-            docs = [
-                Document(id=chunk_ids[idx], content=chunks[idx], metadata=sources[idx])
-                for idx in chunk_ids_set_idx
-                if chunk_ids[idx] not in all_docs_ids
-            ]
+
+            if self._embedding_function is None:
+                docs = [
+                    Document(id=chunk_ids[idx], content=chunks[idx], metadata=sources[idx])
+                    for idx in chunk_ids_set_idx
+                    if chunk_ids[idx] not in all_docs_ids
+                ]
+            else:
+                docs = [
+                    Document(
+                        id=chunk_ids[idx],
+                        content=chunks[idx],
+                        metadata=sources[idx],
+                        embedding=self._embedding_function([chunks[idx]])[0]
+                    )
+                    for idx in chunk_ids_set_idx
+                    if chunk_ids[idx] not in all_docs_ids
+                ]
 
         self._vector_db.insert_docs(docs=docs, collection_name=self._collection_name, upsert=True)