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): if isinstance(query, list):
embeddings = [] embeddings = []
for i in range(0, len(query), embedding_batch_size): for i in range(0, len(query), embedding_batch_size):
embeddings.extend( batch_embeddings = func(
func( query[i : i + embedding_batch_size],
query[i : i + embedding_batch_size], prefix=prefix,
prefix=prefix, user=user,
user=user,
)
) )
if isinstance(batch_embeddings, list):
embeddings.extend(batch_embeddings)
return embeddings return embeddings
else: else:
return func(query, prefix, user) return func(query, prefix, user)

View File

@@ -1334,7 +1334,7 @@ def save_docs_to_vector_db(
) )
return True return True
log.info(f"adding to collection {collection_name}") log.info(f"generating embeddings for {collection_name}")
embedding_function = get_embedding_function( embedding_function = get_embedding_function(
request.app.state.config.RAG_EMBEDDING_ENGINE, request.app.state.config.RAG_EMBEDDING_ENGINE,
request.app.state.config.RAG_EMBEDDING_MODEL, request.app.state.config.RAG_EMBEDDING_MODEL,
@@ -1381,11 +1381,18 @@ def save_docs_to_vector_db(
for idx, text in enumerate(texts) for idx, text in enumerate(texts)
] ]
log.info(f"adding to collection {collection_name}")
VECTOR_DB_CLIENT.insert( VECTOR_DB_CLIENT.insert(
collection_name=collection_name, collection_name=collection_name,
items=items, 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 return True
except Exception as e: except Exception as e:
log.exception(e) log.exception(e)