Merge remote-tracking branch 'upstream/dev' into playwright

backend/open_webui/retrieval/utils.py

@@ -15,8 +15,13 @@ 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,
+)
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
@@ -61,9 +66,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(
@@ -259,26 +262,31 @@ def get_embedding_function(
     embedding_batch_size,
 ):
     if embedding_engine == "":
-        return lambda query: embedding_function.encode(query).tolist()
+        return lambda query, user=None: embedding_function.encode(query).tolist()
     elif embedding_engine in ["ollama", "openai"]:
-        func = lambda query: generate_embeddings(
+        func = lambda query, user=None: generate_embeddings(
             engine=embedding_engine,
             model=embedding_model,
             text=query,
             url=url,
             key=key,
+            user=user,
         )
 
-        def generate_multiple(query, func):
+        def generate_multiple(query, 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]))
+                    embeddings.extend(
+                        func(query[i : i + embedding_batch_size], user=user)
+                    )
                 return embeddings
             else:
-                return func(query)
+                return func(query, user)
 
-        return lambda query: generate_multiple(query, func)
+        return lambda query, user=None: generate_multiple(query, user, func)
+    else:
+        raise ValueError(f"Unknown embedding engine: {embedding_engine}")
 
 
 def get_sources_from_files(
@@ -423,7 +431,11 @@ 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 = ""
+    model: str,
+    texts: list[str],
+    url: str = "https://api.openai.com/v1",
+    key: str = "",
+    user: UserModel = None,
 ) -> Optional[list[list[float]]]:
     try:
         r = requests.post(
@@ -431,6 +443,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={"input": texts, "model": model},
         )
@@ -446,7 +468,7 @@ def generate_openai_batch_embeddings(
 
 
 def generate_ollama_batch_embeddings(
-    model: str, texts: list[str], url: str, key: str = ""
+    model: str, texts: list[str], url: str, key: str = "", user: UserModel = None
 ) -> Optional[list[list[float]]]:
     try:
         r = requests.post(
@@ -454,6 +476,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={"input": texts, "model": model},
         )
@@ -472,22 +504,29 @@ def generate_ollama_batch_embeddings(
 def generate_embeddings(engine: str, model: str, text: Union[str, list[str]], **kwargs):
     url = kwargs.get("url", "")
     key = kwargs.get("key", "")
+    user = kwargs.get("user")
 
     if engine == "ollama":
         if isinstance(text, list):
             embeddings = generate_ollama_batch_embeddings(
-                **{"model": model, "texts": text, "url": url, "key": key}
+                **{"model": model, "texts": text, "url": url, "key": key, "user": user}
             )
         else:
             embeddings = generate_ollama_batch_embeddings(
-                **{"model": model, "texts": [text], "url": url, "key": key}
+                **{
+                    "model": model,
+                    "texts": [text],
+                    "url": url,
+                    "key": key,
+                    "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)
+            embeddings = generate_openai_batch_embeddings(model, text, url, key, user)
         else:
-            embeddings = generate_openai_batch_embeddings(model, [text], url, key)
+            embeddings = generate_openai_batch_embeddings(model, [text], url, key, user)
 
         return embeddings[0] if isinstance(text, str) else embeddings
 

backend/open_webui/retrieval/web/exa.py

@@ -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/utils.py

@@ -48,6 +48,7 @@ def validate_url(url: Union[str, Sequence[str]]):
     else:
         return False
 
+
 def safe_validate_urls(url: Sequence[str]) -> Sequence[str]:
     valid_urls = []
     for u in url:
@@ -57,6 +58,7 @@ def safe_validate_urls(url: Sequence[str]) -> Sequence[str]:
         except ValueError:
             continue
     return valid_urls
+
 def resolve_hostname(hostname):
     # Get address information
     addr_info = socket.getaddrinfo(hostname, None)