Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-03-24 20:14:28 -05:00
parent 58dcc1b33f
commit 8b6fa1f4ab
1 changed files with 164 additions and 185 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/retrieval/utils.py
+164 -185
View File
@@ -550,34 +550,30 @@ def generate_openai_batch_embeddings(
key: str = '', key: str = '',
prefix: str = None, prefix: str = None,
user: UserModel = None, user: UserModel = None,
) -> Optional[list[list[float]]]: ) -> list[list[float]]:
try: log.debug(f'generate_openai_batch_embeddings:model {model} batch size: {len(texts)}')
log.debug(f'generate_openai_batch_embeddings:model {model} batch size: {len(texts)}') json_data = {'input': texts, 'model': model}
json_data = {'input': texts, 'model': model} if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str): json_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
json_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
headers = { headers = {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'Authorization': f'Bearer {key}', 'Authorization': f'Bearer {key}',
} }
if ENABLE_FORWARD_USER_INFO_HEADERS and user: if ENABLE_FORWARD_USER_INFO_HEADERS and user:
headers = include_user_info_headers(headers, user) headers = include_user_info_headers(headers, user)
r = requests.post( r = requests.post(
f'{url}/embeddings', f'{url}/embeddings',
headers=headers, headers=headers,
json=json_data, json=json_data,
) )
r.raise_for_status() r.raise_for_status()
data = r.json() data = r.json()
if 'data' in data: if 'data' in data:
return [elem['embedding'] for elem in data['data']] return [elem['embedding'] for elem in data['data']]
else: else:
raise ValueError("Unexpected OpenAI embeddings response: missing 'data' key") raise ValueError("Unexpected OpenAI embeddings response: missing 'data' key")
except Exception as e:
log.exception(f'Error generating openai batch embeddings: {e}')
return None
async def agenerate_openai_batch_embeddings( async def agenerate_openai_batch_embeddings(
@@ -587,38 +583,34 @@ async def agenerate_openai_batch_embeddings(
key: str = '', key: str = '',
prefix: str = None, prefix: str = None,
user: UserModel = None, user: UserModel = None,
) -> Optional[list[list[float]]]: ) -> list[list[float]]:
try: log.debug(f'agenerate_openai_batch_embeddings:model {model} batch size: {len(texts)}')
log.debug(f'agenerate_openai_batch_embeddings:model {model} batch size: {len(texts)}') form_data = {'input': texts, 'model': model}
form_data = {'input': texts, 'model': model} if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str): form_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
form_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
headers = { headers = {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'Authorization': f'Bearer {key}', 'Authorization': f'Bearer {key}',
} }
if ENABLE_FORWARD_USER_INFO_HEADERS and user: if ENABLE_FORWARD_USER_INFO_HEADERS and user:
headers = include_user_info_headers(headers, user) headers = include_user_info_headers(headers, user)
async with aiohttp.ClientSession( async with aiohttp.ClientSession(
trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT) trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
) as session: ) as session:
async with session.post( async with session.post(
f'{url}/embeddings', f'{url}/embeddings',
headers=headers, headers=headers,
json=form_data, json=form_data,
ssl=AIOHTTP_CLIENT_SESSION_SSL, ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as r: ) as r:
r.raise_for_status() r.raise_for_status()
data = await r.json() data = await r.json()
if 'data' in data: if 'data' in data:
return [item['embedding'] for item in data['data']] return [item['embedding'] for item in data['data']]
else: else:
raise Exception('Something went wrong :/') raise ValueError("Unexpected OpenAI embeddings response: missing 'data' key")
except Exception as e:
log.exception(f'Error generating openai batch embeddings: {e}')
return None
def generate_azure_openai_batch_embeddings( def generate_azure_openai_batch_embeddings(
@@ -629,42 +621,38 @@ def generate_azure_openai_batch_embeddings(
version: str = '', version: str = '',
prefix: str = None, prefix: str = None,
user: UserModel = None, user: UserModel = None,
) -> Optional[list[list[float]]]: ) -> list[list[float]]:
try: log.debug(f'generate_azure_openai_batch_embeddings:deployment {model} batch size: {len(texts)}')
log.debug(f'generate_azure_openai_batch_embeddings:deployment {model} batch size: {len(texts)}') json_data = {'input': texts}
json_data = {'input': texts} if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str): json_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
json_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
url = f'{url}/openai/deployments/{model}/embeddings?api-version={version}' url = f'{url}/openai/deployments/{model}/embeddings?api-version={version}'
for _ in range(5): for _ in range(5):
headers = { headers = {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'api-key': key, 'api-key': key,
} }
if ENABLE_FORWARD_USER_INFO_HEADERS and user: if ENABLE_FORWARD_USER_INFO_HEADERS and user:
headers = include_user_info_headers(headers, user) headers = include_user_info_headers(headers, user)
r = requests.post( r = requests.post(
url, url,
headers=headers, headers=headers,
json=json_data, json=json_data,
) )
if r.status_code == 429: if r.status_code == 429:
retry = float(r.headers.get('Retry-After', '1')) retry = float(r.headers.get('Retry-After', '1'))
time.sleep(retry) time.sleep(retry)
continue continue
r.raise_for_status() r.raise_for_status()
data = r.json() data = r.json()
if 'data' in data: if 'data' in data:
return [elem['embedding'] for elem in data['data']] return [elem['embedding'] for elem in data['data']]
else: else:
raise Exception('Something went wrong :/') raise ValueError("Unexpected Azure OpenAI embeddings response: missing 'data' key")
return None raise Exception('Azure OpenAI embedding request failed: max retries (429) exceeded')
except Exception as e:
log.exception(f'Error generating azure openai batch embeddings: {e}')
return None
async def agenerate_azure_openai_batch_embeddings( async def agenerate_azure_openai_batch_embeddings(
@@ -675,40 +663,36 @@ async def agenerate_azure_openai_batch_embeddings(
version: str = '', version: str = '',
prefix: str = None, prefix: str = None,
user: UserModel = None, user: UserModel = None,
) -> Optional[list[list[float]]]: ) -> list[list[float]]:
try: log.debug(f'agenerate_azure_openai_batch_embeddings:deployment {model} batch size: {len(texts)}')
log.debug(f'agenerate_azure_openai_batch_embeddings:deployment {model} batch size: {len(texts)}') form_data = {'input': texts}
form_data = {'input': texts} if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str): form_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
form_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
full_url = f'{url}/openai/deployments/{model}/embeddings?api-version={version}' full_url = f'{url}/openai/deployments/{model}/embeddings?api-version={version}'
headers = { headers = {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'api-key': key, 'api-key': key,
} }
if ENABLE_FORWARD_USER_INFO_HEADERS and user: if ENABLE_FORWARD_USER_INFO_HEADERS and user:
headers = include_user_info_headers(headers, user) headers = include_user_info_headers(headers, user)
async with aiohttp.ClientSession( async with aiohttp.ClientSession(
trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT) trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
) as session: ) as session:
async with session.post( async with session.post(
full_url, full_url,
headers=headers, headers=headers,
json=form_data, json=form_data,
ssl=AIOHTTP_CLIENT_SESSION_SSL, ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as r: ) as r:
r.raise_for_status() r.raise_for_status()
data = await r.json() data = await r.json()
if 'data' in data: if 'data' in data:
return [item['embedding'] for item in data['data']] return [item['embedding'] for item in data['data']]
else: else:
raise Exception('Something went wrong :/') raise ValueError("Unexpected Azure OpenAI embeddings response: missing 'data' key")
except Exception as e:
log.exception(f'Error generating azure openai batch embeddings: {e}')
return None
def generate_ollama_batch_embeddings( def generate_ollama_batch_embeddings(
@@ -718,37 +702,33 @@ def generate_ollama_batch_embeddings(
key: str = '', key: str = '',
prefix: str = None, prefix: str = None,
user: UserModel = None, user: UserModel = None,
) -> Optional[list[list[float]]]: ) -> list[list[float]]:
try: log.debug(f'generate_ollama_batch_embeddings:model {model} batch size: {len(texts)}')
log.debug(f'generate_ollama_batch_embeddings:model {model} batch size: {len(texts)}') json_data = {'input': texts, 'model': model, 'truncate': True}
json_data = {'input': texts, 'model': model, 'truncate': True} if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str): json_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
json_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
headers = { headers = {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'Authorization': f'Bearer {key}', 'Authorization': f'Bearer {key}',
} }
if ENABLE_FORWARD_USER_INFO_HEADERS and user: if ENABLE_FORWARD_USER_INFO_HEADERS and user:
headers = include_user_info_headers(headers, user) headers = include_user_info_headers(headers, user)
r = requests.post( r = requests.post(
f'{url}/api/embed', f'{url}/api/embed',
headers=headers, headers=headers,
json=json_data, json=json_data,
) )
if r.status_code != 200: if r.status_code != 200:
error_detail = r.json().get('error', r.text) error_detail = r.json().get('error', r.text)
raise Exception(f'Ollama embed error ({r.status_code}): {error_detail}') raise Exception(f'Ollama embed error ({r.status_code}): {error_detail}')
data = r.json() data = r.json()
if 'embeddings' in data: if 'embeddings' in data:
return data['embeddings'] return data['embeddings']
else: else:
raise ValueError("Unexpected Ollama embeddings response: missing 'embeddings' key") raise ValueError("Unexpected Ollama embeddings response: missing 'embeddings' key")
except Exception as e:
log.exception(f'Error generating ollama batch embeddings: {e}')
return None
async def agenerate_ollama_batch_embeddings( async def agenerate_ollama_batch_embeddings(
@@ -758,41 +738,37 @@ async def agenerate_ollama_batch_embeddings(
key: str = '', key: str = '',
prefix: str = None, prefix: str = None,
user: UserModel = None, user: UserModel = None,
) -> Optional[list[list[float]]]: ) -> list[list[float]]:
try: log.debug(f'agenerate_ollama_batch_embeddings:model {model} batch size: {len(texts)}')
log.debug(f'agenerate_ollama_batch_embeddings:model {model} batch size: {len(texts)}') form_data = {'input': texts, 'model': model, 'truncate': True}
form_data = {'input': texts, 'model': model, 'truncate': True} if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str): form_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
form_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
headers = { headers = {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'Authorization': f'Bearer {key}', 'Authorization': f'Bearer {key}',
} }
if ENABLE_FORWARD_USER_INFO_HEADERS and user: if ENABLE_FORWARD_USER_INFO_HEADERS and user:
headers = include_user_info_headers(headers, user) headers = include_user_info_headers(headers, user)
async with aiohttp.ClientSession( async with aiohttp.ClientSession(
trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT) trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
) as session: ) as session:
async with session.post( async with session.post(
f'{url}/api/embed', f'{url}/api/embed',
headers=headers, headers=headers,
json=form_data, json=form_data,
ssl=AIOHTTP_CLIENT_SESSION_SSL, ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as r: ) as r:
if r.status != 200: if r.status != 200:
error_data = await r.json() error_data = await r.json()
error_detail = error_data.get('error', str(error_data)) error_detail = error_data.get('error', str(error_data))
raise Exception(f'Ollama embed error ({r.status}): {error_detail}') raise Exception(f'Ollama embed error ({r.status}): {error_detail}')
data = await r.json() data = await r.json()
if 'embeddings' in data: if 'embeddings' in data:
return data['embeddings'] return data['embeddings']
else: else:
raise Exception('Something went wrong :/') raise ValueError("Unexpected Ollama embeddings response: missing 'embeddings' key")
except Exception as e:
log.exception(f'Error generating ollama batch embeddings: {e}')
return None
def get_embedding_function( def get_embedding_function(
@@ -860,11 +836,14 @@ def get_embedding_function(
for batch in batches: for batch in batches:
batch_results.append(await embedding_function(batch, prefix=prefix, user=user)) batch_results.append(await embedding_function(batch, prefix=prefix, user=user))
# Flatten results # Flatten results — raise if any batch failed
embeddings = [] embeddings = []
for batch_embeddings in batch_results: for i, batch_embeddings in enumerate(batch_results):
if isinstance(batch_embeddings, list): if batch_embeddings is None:
embeddings.extend(batch_embeddings) raise Exception(
f'Embedding generation failed for batch {i + 1}/{len(batches)}'
)
embeddings.extend(batch_embeddings)
log.debug( log.debug(
f'generate_multiple_async: Generated {len(embeddings)} embeddings from {len(batches)} parallel batches' f'generate_multiple_async: Generated {len(embeddings)} embeddings from {len(batches)} parallel batches'