mirror of
https://github.com/microsoft/autogen.git
synced 2026-02-15 07:55:17 -05:00
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:
@@ -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,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user