This commit is contained in:
Timothy Jaeryang Baek 2025-09-11 21:29:56 +04:00
parent bbe6564bfd
commit 2185fc61c0
2 changed files with 15 additions and 7 deletions

View file

@ -432,13 +432,14 @@ def get_embedding_function(
if isinstance(query, list):
embeddings = []
for i in range(0, len(query), embedding_batch_size):
embeddings.extend(
func(
query[i : i + embedding_batch_size],
prefix=prefix,
user=user,
)
batch_embeddings = func(
query[i : i + embedding_batch_size],
prefix=prefix,
user=user,
)
if isinstance(batch_embeddings, list):
embeddings.extend(batch_embeddings)
return embeddings
else:
return func(query, prefix, user)

View file

@ -1334,7 +1334,7 @@ def save_docs_to_vector_db(
)
return True
log.info(f"adding to collection {collection_name}")
log.info(f"generating embeddings for {collection_name}")
embedding_function = get_embedding_function(
request.app.state.config.RAG_EMBEDDING_ENGINE,
request.app.state.config.RAG_EMBEDDING_MODEL,
@ -1381,11 +1381,18 @@ def save_docs_to_vector_db(
for idx, text in enumerate(texts)
]
log.info(f"adding to collection {collection_name}")
VECTOR_DB_CLIENT.insert(
collection_name=collection_name,
items=items,
)
# Validate the number of items inserted
result = VECTOR_DB_CLIENT.query(
collection_name=collection_name,
filter={"metadata": metadata} if metadata else None,
)
return True
except Exception as e:
log.exception(e)