Add support to customized vectordb and embedding functions (#161)

* Add custom embedding function

* Add support to custom vector db

* Improve docstring

* Improve docstring

* Improve docstring

* Add support to customized is_termination_msg fucntion

* Add a test for customize vector db with lancedb

* Fix tests

* Add test for embedding_function

* Update docstring
This commit is contained in:
Li Jiang
2023-10-10 20:53:18 +08:00
committed by GitHub
parent 37a07a83c3
commit fa6e2a52c0
6 changed files with 192 additions and 15 deletions

View File

@@ -12,6 +12,7 @@ try:
)
from autogen.retrieve_utils import create_vector_db_from_dir, query_vector_db
import chromadb
from chromadb.utils import embedding_functions as ef
skip_test = False
except ImportError:
@@ -49,6 +50,7 @@ def test_retrievechat():
},
)
sentence_transformer_ef = ef.SentenceTransformerEmbeddingFunction()
ragproxyagent = RetrieveUserProxyAgent(
name="ragproxyagent",
human_input_mode="NEVER",
@@ -58,6 +60,7 @@ def test_retrievechat():
"chunk_token_size": 2000,
"model": config_list[0]["model"],
"client": chromadb.PersistentClient(path="/tmp/chromadb"),
"embedding_function": sentence_transformer_ef,
},
)