Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Couchbase Vector DB Example Notebook and Minor Bug Fix #3804

Merged
merged 12 commits into from
Oct 17, 2024
19 changes: 8 additions & 11 deletions autogen/agentchat/contrib/vectordb/couchbase.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,16 +56,7 @@ def __init__(
wait_until_index_ready (float | None): Blocking call to wait until the database indexes are ready. None means no wait. Default is None.
wait_until_document_ready (float | None): Blocking call to wait until the database documents are ready. None means no wait. Default is None.
"""
print(
"CouchbaseVectorDB",
connection_string,
username,
password,
bucket_name,
scope_name,
collection_name,
index_name,
)

self.embedding_function = embedding_function
self.index_name = index_name

Expand Down Expand Up @@ -119,6 +110,7 @@ def create_collection(
try:
collection_mgr = self.bucket.collections()
collection_mgr.create_collection(self.scope.name, collection_name)
self.cluster.query(f"CREATE PRIMARY INDEX ON {self.bucket.name}.{self.scope.name}.{collection_name}")

except Exception:
if not get_or_create:
Expand Down Expand Up @@ -287,7 +279,12 @@ def upsert_docs(
[doc["content"]]
).tolist() # Gets new embedding even in case of document update

doc_content = {TEXT_KEY: doc["content"], "metadata": doc.get("metadata", {}), EMBEDDING_KEY: embedding}
doc_content = {
TEXT_KEY: doc["content"],
"metadata": doc.get("metadata", {}),
EMBEDDING_KEY: embedding,
"id": doc_id,
}
docs_to_upsert[doc_id] = doc_content
collection.upsert_multi(docs_to_upsert)

Expand Down
Loading
Loading