Mirrors/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2025-12-16 11:57:51 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'dev' into main

Browse Source
This commit is contained in:
JT
2025-02-05 15:15:24 -08:00
committed by GitHub
parent 157c781b0a b8fe9497fb
commit 40dea3fbe1
148 changed files with 4380 additions and 1767 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
backend/open_webui/retrieval/utils.py
View File

@@ -15,13 +15,20 @@ from langchain_core.documents import Document
from open_webui.config import VECTOR_DB
from open_webui.retrieval.vector.connector import VECTOR_DB_CLIENT
from open_webui.utils.misc import get_last_user_message
from open_webui.models.users import UserModel
from open_webui.env import SRC_LOG_LEVELS, OFFLINE_MODE
from open_webui.env import (
SRC_LOG_LEVELS,
OFFLINE_MODE,
ENABLE_FORWARD_USER_INFO_HEADERS,
)
from open_webui.config import (
RAG_EMBEDDING_QUERY_PREFIX,
RAG_EMBEDDING_PASSAGE_PREFIX,
RAG_EMBEDDING_PREFIX_FIELD_NAME
)
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
@@ -65,9 +72,7 @@ class VectorSearchRetriever(BaseRetriever):
def query_doc(
collection_name: str,
query_embedding: list[float],
k: int,
collection_name: str, query_embedding: list[float], k: int, user: UserModel = None
):
try:
result = VECTOR_DB_CLIENT.search(
@@ -263,27 +268,31 @@ def get_embedding_function(
embedding_batch_size,
):
if embedding_engine == "":
return lambda query, prefix: embedding_function.encode(query, prompt = prefix if prefix else None).tolist()
return lambda query, prefix, user=None: embedding_function.encode(query, prompt = prefix if prefix else None).tolist()
elif embedding_engine in ["ollama", "openai"]:
func = lambda query, prefix: generate_embeddings(
func = lambda query, prefix, user=None: generate_embeddings(
engine=embedding_engine,
model=embedding_model,
text=query,
prefix=prefix,
url=url,
key=key,
user=user,
)
def generate_multiple(query, prefix, func):
def generate_multiple(query, prefix, user, func):
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))
embeddings.extend(
func(query[i : i + embedding_batch_size], prefix=prefix, user=user)
)
return embeddings
else:
return func(query, prefix)
return func(query, prefix, user)
return lambda query, prefix, user=None: generate_multiple(query, prefix, user, func)
else:
raise ValueError(f"Unknown embedding engine: {embedding_engine}")
return lambda query, prefix: generate_multiple(query, prefix, func)
def get_sources_from_files(
@@ -428,9 +437,13 @@ def get_model_path(model: str, update_model: bool = False):
def generate_openai_batch_embeddings(
model: str, texts: list[str], url: str = "https://api.openai.com/v1", key: str = "", prefix: str = None
model: str,
texts: list[str],
url: str = "https://api.openai.com/v1",
key: str = "",
prefix: str = None,
user: UserModel = None
) -> Optional[list[list[float]]]:
try:
json_data = {
"input": texts,
@@ -444,6 +457,16 @@ def generate_openai_batch_embeddings(
headers={
"Content-Type": "application/json",
"Authorization": f"Bearer {key}",
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
}
if ENABLE_FORWARD_USER_INFO_HEADERS and user
else {}
),
},
json=json_data,
)
@@ -459,7 +482,12 @@ def generate_openai_batch_embeddings(
def generate_ollama_batch_embeddings(
model: str, texts: list[str], url: str, key: str = "", prefix: str = None
model: str,
texts: list[str],
url: str,
key: str = "",
prefix: str = None,
user: UserModel = None
) -> Optional[list[list[float]]]:
try:
json_data = {
@@ -474,6 +502,16 @@ def generate_ollama_batch_embeddings(
headers={
"Content-Type": "application/json",
"Authorization": f"Bearer {key}",
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
}
if ENABLE_FORWARD_USER_INFO_HEADERS
else {}
),
},
json=json_data,
)
@@ -492,6 +530,7 @@ def generate_ollama_batch_embeddings(
def generate_embeddings(engine: str, model: str, text: Union[str, list[str]], prefix: Union[str , None] = None, **kwargs):
url = kwargs.get("url", "")
key = kwargs.get("key", "")
user = kwargs.get("user")
if prefix is not None and RAG_EMBEDDING_PREFIX_FIELD_NAME is None:
if isinstance(text, list):
@@ -502,19 +541,18 @@ def generate_embeddings(engine: str, model: str, text: Union[str, list[str]], pr
if engine == "ollama":
if isinstance(text, list):
embeddings = generate_ollama_batch_embeddings(
**{"model": model, "texts": text, "url": url, "key": key, "prefix": prefix}
**{"model": model, "texts": text, "url": url, "key": key, "prefix": prefix, "user": user}
)
else:
embeddings = generate_ollama_batch_embeddings(
**{"model": model, "texts": [text], "url": url, "key": key, "prefix": prefix}
**{"model": model, "texts": [text], "url": url, "key": key, "prefix": prefix, "user": user}
)
return embeddings[0] if isinstance(text, str) else embeddings
elif engine == "openai":
if isinstance(text, list):
embeddings = generate_openai_batch_embeddings(model, text, url, key, prefix)
embeddings = generate_openai_batch_embeddings(model, text, url, key, prefix, user)
else:
embeddings = generate_openai_batch_embeddings(model, [text], url, key, prefix)
embeddings = generate_openai_batch_embeddings(model, [text], url, key, prefix, user)
return embeddings[0] if isinstance(text, str) else embeddings