Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions 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
backend/open_webui/retrieval/utils.py
+57 -19
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
backend/open_webui/retrieval/vector/dbs/milvus.py
+5 -1
View File
@@ -8,13 +8,17 @@ from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult
from open_webui.config import (
MILVUS_URI,
MILVUS_DB,
MILVUS_TOKEN,
)
class MilvusClient:
def __init__(self):
self.collection_prefix = "open_webui"
self.client = Client(uri=MILVUS_URI, database=MILVUS_DB)
if MILVUS_TOKEN is None:
self.client = Client(uri=MILVUS_URI, database=MILVUS_DB)
else:
self.client = Client(uri=MILVUS_URI, database=MILVUS_DB, token=MILVUS_TOKEN)
def _result_to_get_result(self, result) -> GetResult:
ids = []
backend/open_webui/retrieval/web/exa.py
+76
View File
@@ -0,0 +1,76 @@
import logging
from dataclasses import dataclass
from typing import Optional
import requests
from open_webui.env import SRC_LOG_LEVELS
from open_webui.retrieval.web.main import SearchResult
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
EXA_API_BASE = "https://api.exa.ai"
@dataclass
class ExaResult:
url: str
title: str
text: str
def search_exa(
api_key: str,
query: str,
count: int,
filter_list: Optional[list[str]] = None,
) -> list[SearchResult]:
"""Search using Exa Search API and return the results as a list of SearchResult objects.
Args:
api_key (str): A Exa Search API key
query (str): The query to search for
count (int): Number of results to return
filter_list (Optional[list[str]]): List of domains to filter results by
"""
log.info(f"Searching with Exa for query: {query}")
headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
payload = {
"query": query,
"numResults": count or 5,
"includeDomains": filter_list,
"contents": {"text": True, "highlights": True},
"type": "auto", # Use the auto search type (keyword or neural)
}
try:
response = requests.post(
f"{EXA_API_BASE}/search", headers=headers, json=payload
)
response.raise_for_status()
data = response.json()
results = []
for result in data["results"]:
results.append(
ExaResult(
url=result["url"],
title=result["title"],
text=result["text"],
)
)
log.info(f"Found {len(results)} results")
return [
SearchResult(
link=result.url,
title=result.title,
snippet=result.text,
)
for result in results
]
except Exception as e:
log.error(f"Error searching Exa: {e}")
return []
backend/open_webui/retrieval/web/main.py
+4
View File
@@ -1,3 +1,5 @@
import validators
from typing import Optional
from urllib.parse import urlparse
@@ -10,6 +12,8 @@ def get_filtered_results(results, filter_list):
filtered_results = []
for result in results:
url = result.get("url") or result.get("link", "")
if not validators.url(url):
continue
domain = urlparse(url).netloc
if any(domain.endswith(filtered_domain) for filtered_domain in filter_list):
filtered_results.append(result)
backend/open_webui/retrieval/web/utils.py
+15 -4
View File
@@ -43,6 +43,17 @@ def validate_url(url: Union[str, Sequence[str]]):
return False
def safe_validate_urls(url: Sequence[str]) -> Sequence[str]:
valid_urls = []
for u in url:
try:
if validate_url(u):
valid_urls.append(u)
except ValueError:
continue
return valid_urls
def resolve_hostname(hostname):
# Get address information
addr_info = socket.getaddrinfo(hostname, None)
@@ -86,11 +97,11 @@ def get_web_loader(
verify_ssl: bool = True,
requests_per_second: int = 2,
):
# Check if the URL is valid
if not validate_url(urls):
raise ValueError(ERROR_MESSAGES.INVALID_URL)
# Check if the URLs are valid
safe_urls = safe_validate_urls([urls] if isinstance(urls, str) else urls)
return SafeWebBaseLoader(
urls,
safe_urls,
verify_ssl=verify_ssl,
requests_per_second=requests_per_second,
continue_on_failure=True,