scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-11-15 09:31:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

refac

Browse source
This commit is contained in:
Timothy Jaeryang Baek 2026-10-29 05:27:31 +09:00
parent cb83239a8a
commit 58536b6503
3 changed files with 75 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
backend/open_webui/env.py
View file

@ -758,6 +758,15 @@ try:
except ValueError:
AIOHTTP_POOL_DNS_TTL = 300
# Target chunks buffered before insertion, rounded up to a whole embedding batch.
try:
RAG_EMBEDDING_BUFFER_SIZE = int(os.getenv('RAG_EMBEDDING_BUFFER_SIZE', '1000'))
if RAG_EMBEDDING_BUFFER_SIZE <= 0:
raise ValueError
except ValueError:
log.warning('RAG_EMBEDDING_BUFFER_SIZE must be a positive integer. Defaulting to 1000.')
RAG_EMBEDDING_BUFFER_SIZE = 1000
RAG_EMBEDDING_TIMEOUT = os.getenv('RAG_EMBEDDING_TIMEOUT', '')
if RAG_EMBEDDING_TIMEOUT == '':

18
backend/open_webui/retrieval/vector/dbs/s3vector.py
View file

@ -592,11 +592,13 @@ class S3VectorClient(VectorDBBase):
if ids:
# Delete by specific vector IDs/keys
log.info("Deleting %s vectors by IDs from collection '%s'", len(ids), collection_name)
self.client.delete_vectors(
vectorBucketName=self.bucket_name,
indexName=collection_name,
keys=ids,
)
# S3 Vectors accepts at most 500 keys per delete request.
for start in range(0, len(ids), 500):
self.client.delete_vectors(
vectorBucketName=self.bucket_name,
indexName=collection_name,
keys=ids[start : start + 500],
)
log.info("Deleted %s vectors from index '%s'", len(ids), collection_name)
elif filter:
@ -623,11 +625,7 @@ class S3VectorClient(VectorDBBase):
log.info('Found %s vectors matching filter, deleting them', len(matching_ids))
# Delete the matching vectors by ID
self.client.delete_vectors(
vectorBucketName=self.bucket_name,
indexName=collection_name,
keys=matching_ids,
)
self.delete(collection_name, ids=matching_ids)
log.info("Deleted %s vectors from index '%s' using filter", len(matching_ids), collection_name)
else:
log.warning('No vectors found matching the filter criteria')

84
backend/open_webui/routers/retrieval.py
View file

@ -7,6 +7,7 @@ import mimetypes
import os
import re
import shutil
import time
import uuid
from datetime import datetime
from pathlib import Path
@ -49,6 +50,7 @@ from open_webui.env import (
AIOHTTP_CLIENT_SESSION_SSL,
DEVICE_TYPE,
DOCKER,
RAG_EMBEDDING_BUFFER_SIZE,
RAG_EMBEDDING_TIMEOUT,
SENTENCE_TRANSFORMERS_BACKEND,
SENTENCE_TRANSFORMERS_CROSS_ENCODER_BACKEND,
@ -1533,38 +1535,68 @@ def save_docs_to_vector_db(
concurrent_requests=config['rag.embedding_concurrent_requests'],
)
batch_size = config['rag.embedding_batch_size']
buffer_size = RAG_EMBEDDING_BUFFER_SIZE
if isinstance(batch_size, int) and batch_size > 0:
# Round up to whole provider batches to preserve request boundaries.
buffer_size = batch_size * ((buffer_size + batch_size - 1) // batch_size)
# Run async embedding in sync context using the main event loop
# This allows the main loop to stay responsive to health checks during long operations
embedding_timeout = RAG_EMBEDDING_TIMEOUT
future = asyncio.run_coroutine_threadsafe(
embedding_function(
list(map(lambda x: x.replace('\n', ' '), texts)),
prefix=config['rag.embedding_content_prefix'],
user=user,
),
request.app.state.main_loop,
)
embeddings = future.result(timeout=embedding_timeout)
log.info('embeddings generated %s for %s items', len(embeddings), len(texts))
items = [
{
'id': str(uuid.uuid4()),
'text': text,
'vector': embeddings[idx],
'metadata': metadatas[idx],
}
for idx, text in enumerate(texts)
]
log.info('adding to collection %s', collection_name)
get_vector_db_client().insert(
collection_name=collection_name,
items=items,
)
inserted_ids = []
try:
for start in range(0, len(texts), buffer_size):
buffer_texts = texts[start : start + buffer_size]
future = asyncio.run_coroutine_threadsafe(
embedding_function(
[text.replace('\n', ' ') for text in buffer_texts],
prefix=config['rag.embedding_content_prefix'],
user=user,
),
request.app.state.main_loop,
)
started_at = time.monotonic()
try:
embeddings = future.result(timeout=embedding_timeout)
except TimeoutError:
future.cancel()
raise
if embedding_timeout is not None:
embedding_timeout = max(embedding_timeout - (time.monotonic() - started_at), 0)
log.info('embeddings generated %s for %s items', len(embeddings), len(buffer_texts))
log.info('added %s items to collection %s', len(items), collection_name)
items = [
{
'id': str(uuid.uuid4()),
'text': text,
'vector': embeddings[idx],
'metadata': metadatas[start + idx],
}
for idx, text in enumerate(buffer_texts)
]
inserted_ids.extend(item['id'] for item in items)
get_vector_db_client().insert(
collection_name=collection_name,
items=items,
)
del items, embeddings, future
except Exception:
if inserted_ids:
try:
if add:
get_vector_db_client().delete(collection_name=collection_name, ids=inserted_ids)
else:
# This call created/replaced the collection; an empty one would make retries skip indexing.
get_vector_db_client().delete_collection(collection_name=collection_name)
except Exception:
log.exception('failed to remove partially added items from collection %s', collection_name)
raise
log.info('added %s items to collection %s', len(texts), collection_name)
return True
except Exception as e:
log.exception(e)