Embedding function should always return a list of a list of vectors (#…

…3570) Co-authored-by: patrickaubin-abbott <[email protected]> Co-authored-by: Li Jiang <[email protected]>
microsoft · Sep 26, 2024 · 94064da · 94064da
1 parent 0d5163b
commit 94064da
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/autogen/agentchat/contrib/vectordb/pgvectordb.py b/autogen/agentchat/contrib/vectordb/pgvectordb.py
@@ -415,7 +415,8 @@ def query(
         cursor = self.client.cursor()
         results = []
         for query_text in query_texts:
-            vector = self.embedding_function(query_text, convert_to_tensor=False).tolist()
+            vector = self.embedding_function(query_text)
+
             if distance_type.lower() == "cosine":
                 index_function = "<=>"
             elif distance_type.lower() == "euclidean":
@@ -619,7 +620,7 @@ def __init__(
         if embedding_function:
             self.embedding_function = embedding_function
         else:
-            self.embedding_function = SentenceTransformer("all-MiniLM-L6-v2").encode
+            self.embedding_function = lambda s: SentenceTransformer("all-MiniLM-L6-v2").encode(s).tolist()
         self.metadata = metadata
         register_vector(self.client)
         self.active_collection = None