Merge branch 'open-webui:main' into main
This commit is contained in:
@@ -159,7 +159,7 @@ class Loader:
|
||||
elif file_ext in ["htm", "html"]:
|
||||
loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
|
||||
elif file_ext == "md":
|
||||
loader = UnstructuredMarkdownLoader(file_path)
|
||||
loader = TextLoader(file_path, autodetect_encoding=True)
|
||||
elif file_content_type == "application/epub+zip":
|
||||
loader = UnstructuredEPubLoader(file_path)
|
||||
elif (
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
from typing import Any, Dict, Generator, List, Optional, Sequence, Union
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
from langchain_core.documents import Document
|
||||
|
||||
|
||||
ALLOWED_SCHEMES = {"http", "https"}
|
||||
ALLOWED_NETLOCS = {
|
||||
"youtu.be",
|
||||
"m.youtube.com",
|
||||
"youtube.com",
|
||||
"www.youtube.com",
|
||||
"www.youtube-nocookie.com",
|
||||
"vid.plus",
|
||||
}
|
||||
|
||||
|
||||
def _parse_video_id(url: str) -> Optional[str]:
|
||||
"""Parse a YouTube URL and return the video ID if valid, otherwise None."""
|
||||
parsed_url = urlparse(url)
|
||||
|
||||
if parsed_url.scheme not in ALLOWED_SCHEMES:
|
||||
return None
|
||||
|
||||
if parsed_url.netloc not in ALLOWED_NETLOCS:
|
||||
return None
|
||||
|
||||
path = parsed_url.path
|
||||
|
||||
if path.endswith("/watch"):
|
||||
query = parsed_url.query
|
||||
parsed_query = parse_qs(query)
|
||||
if "v" in parsed_query:
|
||||
ids = parsed_query["v"]
|
||||
video_id = ids if isinstance(ids, str) else ids[0]
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
path = parsed_url.path.lstrip("/")
|
||||
video_id = path.split("/")[-1]
|
||||
|
||||
if len(video_id) != 11: # Video IDs are 11 characters long
|
||||
return None
|
||||
|
||||
return video_id
|
||||
|
||||
|
||||
class YoutubeLoader:
|
||||
"""Load `YouTube` video transcripts."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
video_id: str,
|
||||
language: Union[str, Sequence[str]] = "en",
|
||||
):
|
||||
"""Initialize with YouTube video ID."""
|
||||
_video_id = _parse_video_id(video_id)
|
||||
self.video_id = _video_id if _video_id is not None else video_id
|
||||
self._metadata = {"source": video_id}
|
||||
self.language = language
|
||||
if isinstance(language, str):
|
||||
self.language = [language]
|
||||
else:
|
||||
self.language = language
|
||||
|
||||
def load(self) -> List[Document]:
|
||||
"""Load YouTube transcripts into `Document` objects."""
|
||||
try:
|
||||
from youtube_transcript_api import (
|
||||
NoTranscriptFound,
|
||||
TranscriptsDisabled,
|
||||
YouTubeTranscriptApi,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
'Could not import "youtube_transcript_api" Python package. '
|
||||
"Please install it with `pip install youtube-transcript-api`."
|
||||
)
|
||||
|
||||
try:
|
||||
transcript_list = YouTubeTranscriptApi.list_transcripts(self.video_id)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return []
|
||||
|
||||
try:
|
||||
transcript = transcript_list.find_transcript(self.language)
|
||||
except NoTranscriptFound:
|
||||
transcript = transcript_list.find_transcript(["en"])
|
||||
|
||||
transcript_pieces: List[Dict[str, Any]] = transcript.fetch()
|
||||
|
||||
transcript = " ".join(
|
||||
map(
|
||||
lambda transcript_piece: transcript_piece["text"].strip(" "),
|
||||
transcript_pieces,
|
||||
)
|
||||
)
|
||||
return [Document(page_content=transcript, metadata=self._metadata)]
|
||||
@@ -23,6 +23,7 @@ from open_webui.apps.retrieval.vector.connector import VECTOR_DB_CLIENT
|
||||
|
||||
# Document loaders
|
||||
from open_webui.apps.retrieval.loaders.main import Loader
|
||||
from open_webui.apps.retrieval.loaders.youtube import YoutubeLoader
|
||||
|
||||
# Web search engines
|
||||
from open_webui.apps.retrieval.web.main import SearchResult
|
||||
@@ -38,6 +39,7 @@ from open_webui.apps.retrieval.web.serper import search_serper
|
||||
from open_webui.apps.retrieval.web.serply import search_serply
|
||||
from open_webui.apps.retrieval.web.serpstack import search_serpstack
|
||||
from open_webui.apps.retrieval.web.tavily import search_tavily
|
||||
from open_webui.apps.retrieval.web.bing import search_bing
|
||||
|
||||
|
||||
from open_webui.apps.retrieval.utils import (
|
||||
@@ -76,6 +78,8 @@ from open_webui.config import (
|
||||
RAG_FILE_MAX_SIZE,
|
||||
RAG_OPENAI_API_BASE_URL,
|
||||
RAG_OPENAI_API_KEY,
|
||||
RAG_OLLAMA_BASE_URL,
|
||||
RAG_OLLAMA_API_KEY,
|
||||
RAG_RELEVANCE_THRESHOLD,
|
||||
RAG_RERANKING_MODEL,
|
||||
RAG_RERANKING_MODEL_AUTO_UPDATE,
|
||||
@@ -87,6 +91,7 @@ from open_webui.config import (
|
||||
RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
|
||||
RAG_WEB_SEARCH_ENGINE,
|
||||
RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
JINA_API_KEY,
|
||||
SEARCHAPI_API_KEY,
|
||||
SEARCHAPI_ENGINE,
|
||||
SEARXNG_QUERY_URL,
|
||||
@@ -95,13 +100,20 @@ from open_webui.config import (
|
||||
SERPSTACK_API_KEY,
|
||||
SERPSTACK_HTTPS,
|
||||
TAVILY_API_KEY,
|
||||
BING_SEARCH_V7_ENDPOINT,
|
||||
BING_SEARCH_V7_SUBSCRIPTION_KEY,
|
||||
TIKA_SERVER_URL,
|
||||
UPLOAD_DIR,
|
||||
YOUTUBE_LOADER_LANGUAGE,
|
||||
DEFAULT_LOCALE,
|
||||
AppConfig,
|
||||
)
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.env import SRC_LOG_LEVELS, DEVICE_TYPE, DOCKER
|
||||
from open_webui.env import (
|
||||
SRC_LOG_LEVELS,
|
||||
DEVICE_TYPE,
|
||||
DOCKER,
|
||||
)
|
||||
from open_webui.utils.misc import (
|
||||
calculate_sha256,
|
||||
calculate_sha256_string,
|
||||
@@ -111,16 +123,17 @@ from open_webui.utils.misc import (
|
||||
from open_webui.utils.utils import get_admin_user, get_verified_user
|
||||
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter, TokenTextSplitter
|
||||
from langchain_community.document_loaders import (
|
||||
YoutubeLoader,
|
||||
)
|
||||
from langchain_core.documents import Document
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
||||
|
||||
app = FastAPI()
|
||||
app = FastAPI(
|
||||
docs_url="/docs" if ENV == "dev" else None,
|
||||
openapi_url="/openapi.json" if ENV == "dev" else None,
|
||||
redoc_url=None,
|
||||
)
|
||||
|
||||
app.state.config = AppConfig()
|
||||
|
||||
@@ -152,6 +165,9 @@ app.state.config.RAG_TEMPLATE = RAG_TEMPLATE
|
||||
app.state.config.OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
|
||||
app.state.config.OPENAI_API_KEY = RAG_OPENAI_API_KEY
|
||||
|
||||
app.state.config.OLLAMA_BASE_URL = RAG_OLLAMA_BASE_URL
|
||||
app.state.config.OLLAMA_API_KEY = RAG_OLLAMA_API_KEY
|
||||
|
||||
app.state.config.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
|
||||
|
||||
app.state.config.YOUTUBE_LOADER_LANGUAGE = YOUTUBE_LOADER_LANGUAGE
|
||||
@@ -174,6 +190,10 @@ app.state.config.SERPLY_API_KEY = SERPLY_API_KEY
|
||||
app.state.config.TAVILY_API_KEY = TAVILY_API_KEY
|
||||
app.state.config.SEARCHAPI_API_KEY = SEARCHAPI_API_KEY
|
||||
app.state.config.SEARCHAPI_ENGINE = SEARCHAPI_ENGINE
|
||||
app.state.config.JINA_API_KEY = JINA_API_KEY
|
||||
app.state.config.BING_SEARCH_V7_ENDPOINT = BING_SEARCH_V7_ENDPOINT
|
||||
app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY = BING_SEARCH_V7_SUBSCRIPTION_KEY
|
||||
|
||||
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT = RAG_WEB_SEARCH_RESULT_COUNT
|
||||
app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS = RAG_WEB_SEARCH_CONCURRENT_REQUESTS
|
||||
|
||||
@@ -185,11 +205,15 @@ def update_embedding_model(
|
||||
if embedding_model and app.state.config.RAG_EMBEDDING_ENGINE == "":
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
app.state.sentence_transformer_ef = SentenceTransformer(
|
||||
get_model_path(embedding_model, auto_update),
|
||||
device=DEVICE_TYPE,
|
||||
trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
|
||||
)
|
||||
try:
|
||||
app.state.sentence_transformer_ef = SentenceTransformer(
|
||||
get_model_path(embedding_model, auto_update),
|
||||
device=DEVICE_TYPE,
|
||||
trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
|
||||
)
|
||||
except Exception as e:
|
||||
log.debug(f"Error loading SentenceTransformer: {e}")
|
||||
app.state.sentence_transformer_ef = None
|
||||
else:
|
||||
app.state.sentence_transformer_ef = None
|
||||
|
||||
@@ -243,8 +267,16 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
|
||||
app.state.config.RAG_EMBEDDING_ENGINE,
|
||||
app.state.config.RAG_EMBEDDING_MODEL,
|
||||
app.state.sentence_transformer_ef,
|
||||
app.state.config.OPENAI_API_KEY,
|
||||
app.state.config.OPENAI_API_BASE_URL,
|
||||
(
|
||||
app.state.config.OPENAI_API_BASE_URL
|
||||
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
|
||||
else app.state.config.OLLAMA_BASE_URL
|
||||
),
|
||||
(
|
||||
app.state.config.OPENAI_API_KEY
|
||||
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
|
||||
else app.state.config.OLLAMA_API_KEY
|
||||
),
|
||||
app.state.config.RAG_EMBEDDING_BATCH_SIZE,
|
||||
)
|
||||
|
||||
@@ -294,6 +326,10 @@ async def get_embedding_config(user=Depends(get_admin_user)):
|
||||
"url": app.state.config.OPENAI_API_BASE_URL,
|
||||
"key": app.state.config.OPENAI_API_KEY,
|
||||
},
|
||||
"ollama_config": {
|
||||
"url": app.state.config.OLLAMA_BASE_URL,
|
||||
"key": app.state.config.OLLAMA_API_KEY,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -310,8 +346,14 @@ class OpenAIConfigForm(BaseModel):
|
||||
key: str
|
||||
|
||||
|
||||
class OllamaConfigForm(BaseModel):
|
||||
url: str
|
||||
key: str
|
||||
|
||||
|
||||
class EmbeddingModelUpdateForm(BaseModel):
|
||||
openai_config: Optional[OpenAIConfigForm] = None
|
||||
ollama_config: Optional[OllamaConfigForm] = None
|
||||
embedding_engine: str
|
||||
embedding_model: str
|
||||
embedding_batch_size: Optional[int] = 1
|
||||
@@ -332,6 +374,11 @@ async def update_embedding_config(
|
||||
if form_data.openai_config is not None:
|
||||
app.state.config.OPENAI_API_BASE_URL = form_data.openai_config.url
|
||||
app.state.config.OPENAI_API_KEY = form_data.openai_config.key
|
||||
|
||||
if form_data.ollama_config is not None:
|
||||
app.state.config.OLLAMA_BASE_URL = form_data.ollama_config.url
|
||||
app.state.config.OLLAMA_API_KEY = form_data.ollama_config.key
|
||||
|
||||
app.state.config.RAG_EMBEDDING_BATCH_SIZE = form_data.embedding_batch_size
|
||||
|
||||
update_embedding_model(app.state.config.RAG_EMBEDDING_MODEL)
|
||||
@@ -340,8 +387,16 @@ async def update_embedding_config(
|
||||
app.state.config.RAG_EMBEDDING_ENGINE,
|
||||
app.state.config.RAG_EMBEDDING_MODEL,
|
||||
app.state.sentence_transformer_ef,
|
||||
app.state.config.OPENAI_API_KEY,
|
||||
app.state.config.OPENAI_API_BASE_URL,
|
||||
(
|
||||
app.state.config.OPENAI_API_BASE_URL
|
||||
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
|
||||
else app.state.config.OLLAMA_BASE_URL
|
||||
),
|
||||
(
|
||||
app.state.config.OPENAI_API_KEY
|
||||
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
|
||||
else app.state.config.OLLAMA_API_KEY
|
||||
),
|
||||
app.state.config.RAG_EMBEDDING_BATCH_SIZE,
|
||||
)
|
||||
|
||||
@@ -354,6 +409,10 @@ async def update_embedding_config(
|
||||
"url": app.state.config.OPENAI_API_BASE_URL,
|
||||
"key": app.state.config.OPENAI_API_KEY,
|
||||
},
|
||||
"ollama_config": {
|
||||
"url": app.state.config.OLLAMA_BASE_URL,
|
||||
"key": app.state.config.OLLAMA_API_KEY,
|
||||
},
|
||||
}
|
||||
except Exception as e:
|
||||
log.exception(f"Problem updating embedding model: {e}")
|
||||
@@ -414,7 +473,7 @@ async def get_rag_config(user=Depends(get_admin_user)):
|
||||
"translation": app.state.YOUTUBE_LOADER_TRANSLATION,
|
||||
},
|
||||
"web": {
|
||||
"ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
|
||||
"web_loader_ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
|
||||
"search": {
|
||||
"enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
|
||||
"engine": app.state.config.RAG_WEB_SEARCH_ENGINE,
|
||||
@@ -430,6 +489,9 @@ async def get_rag_config(user=Depends(get_admin_user)):
|
||||
"tavily_api_key": app.state.config.TAVILY_API_KEY,
|
||||
"searchapi_api_key": app.state.config.SEARCHAPI_API_KEY,
|
||||
"seaarchapi_engine": app.state.config.SEARCHAPI_ENGINE,
|
||||
"jina_api_key": app.state.config.JINA_API_KEY,
|
||||
"bing_search_v7_endpoint": app.state.config.BING_SEARCH_V7_ENDPOINT,
|
||||
"bing_search_v7_subscription_key": app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
|
||||
"result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
"concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
|
||||
},
|
||||
@@ -473,6 +535,9 @@ class WebSearchConfig(BaseModel):
|
||||
tavily_api_key: Optional[str] = None
|
||||
searchapi_api_key: Optional[str] = None
|
||||
searchapi_engine: Optional[str] = None
|
||||
jina_api_key: Optional[str] = None
|
||||
bing_search_v7_endpoint: Optional[str] = None
|
||||
bing_search_v7_subscription_key: Optional[str] = None
|
||||
result_count: Optional[int] = None
|
||||
concurrent_requests: Optional[int] = None
|
||||
|
||||
@@ -519,6 +584,7 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
|
||||
|
||||
if form_data.web is not None:
|
||||
app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
|
||||
# Note: When UI "Bypass SSL verification for Websites"=True then ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION=False
|
||||
form_data.web.web_loader_ssl_verification
|
||||
)
|
||||
|
||||
@@ -540,6 +606,15 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
|
||||
app.state.config.TAVILY_API_KEY = form_data.web.search.tavily_api_key
|
||||
app.state.config.SEARCHAPI_API_KEY = form_data.web.search.searchapi_api_key
|
||||
app.state.config.SEARCHAPI_ENGINE = form_data.web.search.searchapi_engine
|
||||
|
||||
app.state.config.JINA_API_KEY = form_data.web.search.jina_api_key
|
||||
app.state.config.BING_SEARCH_V7_ENDPOINT = (
|
||||
form_data.web.search.bing_search_v7_endpoint
|
||||
)
|
||||
app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY = (
|
||||
form_data.web.search.bing_search_v7_subscription_key
|
||||
)
|
||||
|
||||
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT = form_data.web.search.result_count
|
||||
app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS = (
|
||||
form_data.web.search.concurrent_requests
|
||||
@@ -566,7 +641,7 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
|
||||
"translation": app.state.YOUTUBE_LOADER_TRANSLATION,
|
||||
},
|
||||
"web": {
|
||||
"ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
|
||||
"web_loader_ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
|
||||
"search": {
|
||||
"enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
|
||||
"engine": app.state.config.RAG_WEB_SEARCH_ENGINE,
|
||||
@@ -582,6 +657,9 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
|
||||
"serachapi_api_key": app.state.config.SEARCHAPI_API_KEY,
|
||||
"searchapi_engine": app.state.config.SEARCHAPI_ENGINE,
|
||||
"tavily_api_key": app.state.config.TAVILY_API_KEY,
|
||||
"jina_api_key": app.state.config.JINA_API_KEY,
|
||||
"bing_search_v7_endpoint": app.state.config.BING_SEARCH_V7_ENDPOINT,
|
||||
"bing_search_v7_subscription_key": app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
|
||||
"result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
"concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
|
||||
},
|
||||
@@ -643,6 +721,23 @@ async def update_query_settings(
|
||||
####################################
|
||||
|
||||
|
||||
def _get_docs_info(docs: list[Document]) -> str:
|
||||
docs_info = set()
|
||||
|
||||
# Trying to select relevant metadata identifying the document.
|
||||
for doc in docs:
|
||||
metadata = getattr(doc, "metadata", {})
|
||||
doc_name = metadata.get("name", "")
|
||||
if not doc_name:
|
||||
doc_name = metadata.get("title", "")
|
||||
if not doc_name:
|
||||
doc_name = metadata.get("source", "")
|
||||
if doc_name:
|
||||
docs_info.add(doc_name)
|
||||
|
||||
return ", ".join(docs_info)
|
||||
|
||||
|
||||
def save_docs_to_vector_db(
|
||||
docs,
|
||||
collection_name,
|
||||
@@ -651,7 +746,9 @@ def save_docs_to_vector_db(
|
||||
split: bool = True,
|
||||
add: bool = False,
|
||||
) -> bool:
|
||||
log.info(f"save_docs_to_vector_db {docs} {collection_name}")
|
||||
log.info(
|
||||
f"save_docs_to_vector_db: document {_get_docs_info(docs)} {collection_name}"
|
||||
)
|
||||
|
||||
# Check if entries with the same hash (metadata.hash) already exist
|
||||
if metadata and "hash" in metadata:
|
||||
@@ -733,8 +830,16 @@ def save_docs_to_vector_db(
|
||||
app.state.config.RAG_EMBEDDING_ENGINE,
|
||||
app.state.config.RAG_EMBEDDING_MODEL,
|
||||
app.state.sentence_transformer_ef,
|
||||
app.state.config.OPENAI_API_KEY,
|
||||
app.state.config.OPENAI_API_BASE_URL,
|
||||
(
|
||||
app.state.config.OPENAI_API_BASE_URL
|
||||
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
|
||||
else app.state.config.OLLAMA_BASE_URL
|
||||
),
|
||||
(
|
||||
app.state.config.OPENAI_API_KEY
|
||||
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
|
||||
else app.state.config.OLLAMA_API_KEY
|
||||
),
|
||||
app.state.config.RAG_EMBEDDING_BATCH_SIZE,
|
||||
)
|
||||
|
||||
@@ -959,12 +1064,10 @@ def process_youtube_video(form_data: ProcessUrlForm, user=Depends(get_verified_u
|
||||
if not collection_name:
|
||||
collection_name = calculate_sha256_string(form_data.url)[:63]
|
||||
|
||||
loader = YoutubeLoader.from_youtube_url(
|
||||
form_data.url,
|
||||
add_video_info=True,
|
||||
language=app.state.config.YOUTUBE_LOADER_LANGUAGE,
|
||||
translation=app.state.YOUTUBE_LOADER_TRANSLATION,
|
||||
loader = YoutubeLoader(
|
||||
form_data.url, language=app.state.config.YOUTUBE_LOADER_LANGUAGE
|
||||
)
|
||||
|
||||
docs = loader.load()
|
||||
content = " ".join([doc.page_content for doc in docs])
|
||||
log.debug(f"text_content: {content}")
|
||||
@@ -1150,7 +1253,20 @@ def search_web(engine: str, query: str) -> list[SearchResult]:
|
||||
else:
|
||||
raise Exception("No SEARCHAPI_API_KEY found in environment variables")
|
||||
elif engine == "jina":
|
||||
return search_jina(query, app.state.config.RAG_WEB_SEARCH_RESULT_COUNT)
|
||||
return search_jina(
|
||||
app.state.config.JINA_API_KEY,
|
||||
query,
|
||||
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
)
|
||||
elif engine == "bing":
|
||||
return search_bing(
|
||||
app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
|
||||
app.state.config.BING_SEARCH_V7_ENDPOINT,
|
||||
str(DEFAULT_LOCALE),
|
||||
query,
|
||||
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
|
||||
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
|
||||
)
|
||||
else:
|
||||
raise Exception("No search engine API key found in environment variables")
|
||||
|
||||
@@ -1180,8 +1296,12 @@ def process_web_search(form_data: SearchForm, user=Depends(get_verified_user)):
|
||||
|
||||
urls = [result.link for result in web_results]
|
||||
|
||||
loader = get_web_loader(urls)
|
||||
docs = loader.load()
|
||||
loader = get_web_loader(
|
||||
urls,
|
||||
verify_ssl=app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
|
||||
requests_per_second=app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
|
||||
)
|
||||
docs = loader.aload()
|
||||
|
||||
save_docs_to_vector_db(docs, collection_name, overwrite=True)
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ import os
|
||||
import uuid
|
||||
from typing import Optional, Union
|
||||
|
||||
import asyncio
|
||||
import requests
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
@@ -10,11 +11,6 @@ from langchain.retrievers import ContextualCompressionRetriever, EnsembleRetriev
|
||||
from langchain_community.retrievers import BM25Retriever
|
||||
from langchain_core.documents import Document
|
||||
|
||||
|
||||
from open_webui.apps.ollama.main import (
|
||||
GenerateEmbedForm,
|
||||
generate_ollama_batch_embeddings,
|
||||
)
|
||||
from open_webui.apps.retrieval.vector.connector import VECTOR_DB_CLIENT
|
||||
from open_webui.utils.misc import get_last_user_message
|
||||
|
||||
@@ -76,7 +72,7 @@ def query_doc(
|
||||
limit=k,
|
||||
)
|
||||
|
||||
log.info(f"query_doc:result {result}")
|
||||
log.info(f"query_doc:result {result.ids} {result.metadatas}")
|
||||
return result
|
||||
except Exception as e:
|
||||
print(e)
|
||||
@@ -127,7 +123,10 @@ def query_doc_with_hybrid_search(
|
||||
"metadatas": [[d.metadata for d in result]],
|
||||
}
|
||||
|
||||
log.info(f"query_doc_with_hybrid_search:result {result}")
|
||||
log.info(
|
||||
"query_doc_with_hybrid_search:result "
|
||||
+ f'{result["metadatas"]} {result["distances"]}'
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
raise e
|
||||
@@ -178,35 +177,34 @@ def merge_and_sort_query_results(
|
||||
|
||||
def query_collection(
|
||||
collection_names: list[str],
|
||||
query: str,
|
||||
queries: list[str],
|
||||
embedding_function,
|
||||
k: int,
|
||||
) -> dict:
|
||||
|
||||
results = []
|
||||
query_embedding = embedding_function(query)
|
||||
|
||||
for collection_name in collection_names:
|
||||
if collection_name:
|
||||
try:
|
||||
result = query_doc(
|
||||
collection_name=collection_name,
|
||||
k=k,
|
||||
query_embedding=query_embedding,
|
||||
)
|
||||
if result is not None:
|
||||
results.append(result.model_dump())
|
||||
except Exception as e:
|
||||
log.exception(f"Error when querying the collection: {e}")
|
||||
else:
|
||||
pass
|
||||
for query in queries:
|
||||
query_embedding = embedding_function(query)
|
||||
for collection_name in collection_names:
|
||||
if collection_name:
|
||||
try:
|
||||
result = query_doc(
|
||||
collection_name=collection_name,
|
||||
k=k,
|
||||
query_embedding=query_embedding,
|
||||
)
|
||||
if result is not None:
|
||||
results.append(result.model_dump())
|
||||
except Exception as e:
|
||||
log.exception(f"Error when querying the collection: {e}")
|
||||
else:
|
||||
pass
|
||||
|
||||
return merge_and_sort_query_results(results, k=k)
|
||||
|
||||
|
||||
def query_collection_with_hybrid_search(
|
||||
collection_names: list[str],
|
||||
query: str,
|
||||
queries: list[str],
|
||||
embedding_function,
|
||||
k: int,
|
||||
reranking_function,
|
||||
@@ -216,15 +214,16 @@ def query_collection_with_hybrid_search(
|
||||
error = False
|
||||
for collection_name in collection_names:
|
||||
try:
|
||||
result = query_doc_with_hybrid_search(
|
||||
collection_name=collection_name,
|
||||
query=query,
|
||||
embedding_function=embedding_function,
|
||||
k=k,
|
||||
reranking_function=reranking_function,
|
||||
r=r,
|
||||
)
|
||||
results.append(result)
|
||||
for query in queries:
|
||||
result = query_doc_with_hybrid_search(
|
||||
collection_name=collection_name,
|
||||
query=query,
|
||||
embedding_function=embedding_function,
|
||||
k=k,
|
||||
reranking_function=reranking_function,
|
||||
r=r,
|
||||
)
|
||||
results.append(result)
|
||||
except Exception as e:
|
||||
log.exception(
|
||||
"Error when querying the collection with " f"hybrid_search: {e}"
|
||||
@@ -281,8 +280,8 @@ def get_embedding_function(
|
||||
embedding_engine,
|
||||
embedding_model,
|
||||
embedding_function,
|
||||
openai_key,
|
||||
openai_url,
|
||||
url,
|
||||
key,
|
||||
embedding_batch_size,
|
||||
):
|
||||
if embedding_engine == "":
|
||||
@@ -292,8 +291,8 @@ def get_embedding_function(
|
||||
engine=embedding_engine,
|
||||
model=embedding_model,
|
||||
text=query,
|
||||
key=openai_key if embedding_engine == "openai" else "",
|
||||
url=openai_url if embedding_engine == "openai" else "",
|
||||
url=url,
|
||||
key=key,
|
||||
)
|
||||
|
||||
def generate_multiple(query, func):
|
||||
@@ -310,15 +309,14 @@ def get_embedding_function(
|
||||
|
||||
def get_rag_context(
|
||||
files,
|
||||
messages,
|
||||
queries,
|
||||
embedding_function,
|
||||
k,
|
||||
reranking_function,
|
||||
r,
|
||||
hybrid_search,
|
||||
):
|
||||
log.debug(f"files: {files} {messages} {embedding_function} {reranking_function}")
|
||||
query = get_last_user_message(messages)
|
||||
log.debug(f"files: {files} {queries} {embedding_function} {reranking_function}")
|
||||
|
||||
extracted_collections = []
|
||||
relevant_contexts = []
|
||||
@@ -360,7 +358,7 @@ def get_rag_context(
|
||||
try:
|
||||
context = query_collection_with_hybrid_search(
|
||||
collection_names=collection_names,
|
||||
query=query,
|
||||
queries=queries,
|
||||
embedding_function=embedding_function,
|
||||
k=k,
|
||||
reranking_function=reranking_function,
|
||||
@@ -375,7 +373,7 @@ def get_rag_context(
|
||||
if (not hybrid_search) or (context is None):
|
||||
context = query_collection(
|
||||
collection_names=collection_names,
|
||||
query=query,
|
||||
queries=queries,
|
||||
embedding_function=embedding_function,
|
||||
k=k,
|
||||
)
|
||||
@@ -467,7 +465,7 @@ def get_model_path(model: str, update_model: bool = False):
|
||||
|
||||
|
||||
def generate_openai_batch_embeddings(
|
||||
model: str, texts: list[str], key: str, url: str = "https://api.openai.com/v1"
|
||||
model: str, texts: list[str], url: str = "https://api.openai.com/v1", key: str = ""
|
||||
) -> Optional[list[list[float]]]:
|
||||
try:
|
||||
r = requests.post(
|
||||
@@ -489,29 +487,50 @@ def generate_openai_batch_embeddings(
|
||||
return None
|
||||
|
||||
|
||||
def generate_ollama_batch_embeddings(
|
||||
model: str, texts: list[str], url: str, key: str
|
||||
) -> Optional[list[list[float]]]:
|
||||
try:
|
||||
r = requests.post(
|
||||
f"{url}/api/embed",
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {key}",
|
||||
},
|
||||
json={"input": texts, "model": model},
|
||||
)
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
print(data)
|
||||
if "embeddings" in data:
|
||||
return data["embeddings"]
|
||||
else:
|
||||
raise "Something went wrong :/"
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return None
|
||||
|
||||
|
||||
def generate_embeddings(engine: str, model: str, text: Union[str, list[str]], **kwargs):
|
||||
url = kwargs.get("url", "")
|
||||
key = kwargs.get("key", "")
|
||||
|
||||
if engine == "ollama":
|
||||
if isinstance(text, list):
|
||||
embeddings = generate_ollama_batch_embeddings(
|
||||
GenerateEmbedForm(**{"model": model, "input": text})
|
||||
**{"model": model, "texts": text, "url": url, "key": key}
|
||||
)
|
||||
else:
|
||||
embeddings = generate_ollama_batch_embeddings(
|
||||
GenerateEmbedForm(**{"model": model, "input": [text]})
|
||||
**{"model": model, "texts": [text], "url": url, "key": key}
|
||||
)
|
||||
return (
|
||||
embeddings["embeddings"][0]
|
||||
if isinstance(text, str)
|
||||
else embeddings["embeddings"]
|
||||
)
|
||||
return embeddings[0] if isinstance(text, str) else embeddings
|
||||
elif engine == "openai":
|
||||
key = kwargs.get("key", "")
|
||||
url = kwargs.get("url", "https://api.openai.com/v1")
|
||||
|
||||
if isinstance(text, list):
|
||||
embeddings = generate_openai_batch_embeddings(model, text, key, url)
|
||||
embeddings = generate_openai_batch_embeddings(model, text, url, key)
|
||||
else:
|
||||
embeddings = generate_openai_batch_embeddings(model, [text], key, url)
|
||||
embeddings = generate_openai_batch_embeddings(model, [text], url, key)
|
||||
|
||||
return embeddings[0] if isinstance(text, str) else embeddings
|
||||
|
||||
|
||||
@@ -8,6 +8,14 @@ elif VECTOR_DB == "qdrant":
|
||||
from open_webui.apps.retrieval.vector.dbs.qdrant import QdrantClient
|
||||
|
||||
VECTOR_DB_CLIENT = QdrantClient()
|
||||
elif VECTOR_DB == "opensearch":
|
||||
from open_webui.apps.retrieval.vector.dbs.opensearch import OpenSearchClient
|
||||
|
||||
VECTOR_DB_CLIENT = OpenSearchClient()
|
||||
elif VECTOR_DB == "pgvector":
|
||||
from open_webui.apps.retrieval.vector.dbs.pgvector import PgvectorClient
|
||||
|
||||
VECTOR_DB_CLIENT = PgvectorClient()
|
||||
else:
|
||||
from open_webui.apps.retrieval.vector.dbs.chroma import ChromaClient
|
||||
|
||||
|
||||
@@ -13,11 +13,24 @@ from open_webui.config import (
|
||||
CHROMA_HTTP_SSL,
|
||||
CHROMA_TENANT,
|
||||
CHROMA_DATABASE,
|
||||
CHROMA_CLIENT_AUTH_PROVIDER,
|
||||
CHROMA_CLIENT_AUTH_CREDENTIALS,
|
||||
)
|
||||
|
||||
|
||||
class ChromaClient:
|
||||
def __init__(self):
|
||||
settings_dict = {
|
||||
"allow_reset": True,
|
||||
"anonymized_telemetry": False,
|
||||
}
|
||||
if CHROMA_CLIENT_AUTH_PROVIDER is not None:
|
||||
settings_dict["chroma_client_auth_provider"] = CHROMA_CLIENT_AUTH_PROVIDER
|
||||
if CHROMA_CLIENT_AUTH_CREDENTIALS is not None:
|
||||
settings_dict["chroma_client_auth_credentials"] = (
|
||||
CHROMA_CLIENT_AUTH_CREDENTIALS
|
||||
)
|
||||
|
||||
if CHROMA_HTTP_HOST != "":
|
||||
self.client = chromadb.HttpClient(
|
||||
host=CHROMA_HTTP_HOST,
|
||||
@@ -26,12 +39,12 @@ class ChromaClient:
|
||||
ssl=CHROMA_HTTP_SSL,
|
||||
tenant=CHROMA_TENANT,
|
||||
database=CHROMA_DATABASE,
|
||||
settings=Settings(allow_reset=True, anonymized_telemetry=False),
|
||||
settings=Settings(**settings_dict),
|
||||
)
|
||||
else:
|
||||
self.client = chromadb.PersistentClient(
|
||||
path=CHROMA_DATA_PATH,
|
||||
settings=Settings(allow_reset=True, anonymized_telemetry=False),
|
||||
settings=Settings(**settings_dict),
|
||||
tenant=CHROMA_TENANT,
|
||||
database=CHROMA_DATABASE,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
from opensearchpy import OpenSearch
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.apps.retrieval.vector.main import VectorItem, SearchResult, GetResult
|
||||
from open_webui.config import (
|
||||
OPENSEARCH_URI,
|
||||
OPENSEARCH_SSL,
|
||||
OPENSEARCH_CERT_VERIFY,
|
||||
OPENSEARCH_USERNAME,
|
||||
OPENSEARCH_PASSWORD,
|
||||
)
|
||||
|
||||
|
||||
class OpenSearchClient:
|
||||
def __init__(self):
|
||||
self.index_prefix = "open_webui"
|
||||
self.client = OpenSearch(
|
||||
hosts=[OPENSEARCH_URI],
|
||||
use_ssl=OPENSEARCH_SSL,
|
||||
verify_certs=OPENSEARCH_CERT_VERIFY,
|
||||
http_auth=(OPENSEARCH_USERNAME, OPENSEARCH_PASSWORD),
|
||||
)
|
||||
|
||||
def _result_to_get_result(self, result) -> GetResult:
|
||||
ids = []
|
||||
documents = []
|
||||
metadatas = []
|
||||
|
||||
for hit in result["hits"]["hits"]:
|
||||
ids.append(hit["_id"])
|
||||
documents.append(hit["_source"].get("text"))
|
||||
metadatas.append(hit["_source"].get("metadata"))
|
||||
|
||||
return GetResult(ids=ids, documents=documents, metadatas=metadatas)
|
||||
|
||||
def _result_to_search_result(self, result) -> SearchResult:
|
||||
ids = []
|
||||
distances = []
|
||||
documents = []
|
||||
metadatas = []
|
||||
|
||||
for hit in result["hits"]["hits"]:
|
||||
ids.append(hit["_id"])
|
||||
distances.append(hit["_score"])
|
||||
documents.append(hit["_source"].get("text"))
|
||||
metadatas.append(hit["_source"].get("metadata"))
|
||||
|
||||
return SearchResult(
|
||||
ids=ids, distances=distances, documents=documents, metadatas=metadatas
|
||||
)
|
||||
|
||||
def _create_index(self, index_name: str, dimension: int):
|
||||
body = {
|
||||
"mappings": {
|
||||
"properties": {
|
||||
"id": {"type": "keyword"},
|
||||
"vector": {
|
||||
"type": "dense_vector",
|
||||
"dims": dimension, # Adjust based on your vector dimensions
|
||||
"index": true,
|
||||
"similarity": "faiss",
|
||||
"method": {
|
||||
"name": "hnsw",
|
||||
"space_type": "ip", # Use inner product to approximate cosine similarity
|
||||
"engine": "faiss",
|
||||
"ef_construction": 128,
|
||||
"m": 16,
|
||||
},
|
||||
},
|
||||
"text": {"type": "text"},
|
||||
"metadata": {"type": "object"},
|
||||
}
|
||||
}
|
||||
}
|
||||
self.client.indices.create(index=f"{self.index_prefix}_{index_name}", body=body)
|
||||
|
||||
def _create_batches(self, items: list[VectorItem], batch_size=100):
|
||||
for i in range(0, len(items), batch_size):
|
||||
yield items[i : i + batch_size]
|
||||
|
||||
def has_collection(self, index_name: str) -> bool:
|
||||
# has_collection here means has index.
|
||||
# We are simply adapting to the norms of the other DBs.
|
||||
return self.client.indices.exists(index=f"{self.index_prefix}_{index_name}")
|
||||
|
||||
def delete_colleciton(self, index_name: str):
|
||||
# delete_collection here means delete index.
|
||||
# We are simply adapting to the norms of the other DBs.
|
||||
self.client.indices.delete(index=f"{self.index_prefix}_{index_name}")
|
||||
|
||||
def search(
|
||||
self, index_name: str, vectors: list[list[float]], limit: int
|
||||
) -> Optional[SearchResult]:
|
||||
query = {
|
||||
"size": limit,
|
||||
"_source": ["text", "metadata"],
|
||||
"query": {
|
||||
"script_score": {
|
||||
"query": {"match_all": {}},
|
||||
"script": {
|
||||
"source": "cosineSimilarity(params.vector, 'vector') + 1.0",
|
||||
"params": {
|
||||
"vector": vectors[0]
|
||||
}, # Assuming single query vector
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
result = self.client.search(
|
||||
index=f"{self.index_prefix}_{index_name}", body=query
|
||||
)
|
||||
|
||||
return self._result_to_search_result(result)
|
||||
|
||||
def get_or_create_index(self, index_name: str, dimension: int):
|
||||
if not self.has_index(index_name):
|
||||
self._create_index(index_name, dimension)
|
||||
|
||||
def get(self, index_name: str) -> Optional[GetResult]:
|
||||
query = {"query": {"match_all": {}}, "_source": ["text", "metadata"]}
|
||||
|
||||
result = self.client.search(
|
||||
index=f"{self.index_prefix}_{index_name}", body=query
|
||||
)
|
||||
return self._result_to_get_result(result)
|
||||
|
||||
def insert(self, index_name: str, items: list[VectorItem]):
|
||||
if not self.has_index(index_name):
|
||||
self._create_index(index_name, dimension=len(items[0]["vector"]))
|
||||
|
||||
for batch in self._create_batches(items):
|
||||
actions = [
|
||||
{
|
||||
"index": {
|
||||
"_id": item["id"],
|
||||
"_source": {
|
||||
"vector": item["vector"],
|
||||
"text": item["text"],
|
||||
"metadata": item["metadata"],
|
||||
},
|
||||
}
|
||||
}
|
||||
for item in batch
|
||||
]
|
||||
self.client.bulk(actions)
|
||||
|
||||
def upsert(self, index_name: str, items: list[VectorItem]):
|
||||
if not self.has_index(index_name):
|
||||
self._create_index(index_name, dimension=len(items[0]["vector"]))
|
||||
|
||||
for batch in self._create_batches(items):
|
||||
actions = [
|
||||
{
|
||||
"index": {
|
||||
"_id": item["id"],
|
||||
"_source": {
|
||||
"vector": item["vector"],
|
||||
"text": item["text"],
|
||||
"metadata": item["metadata"],
|
||||
},
|
||||
}
|
||||
}
|
||||
for item in batch
|
||||
]
|
||||
self.client.bulk(actions)
|
||||
|
||||
def delete(self, index_name: str, ids: list[str]):
|
||||
actions = [
|
||||
{"delete": {"_index": f"{self.index_prefix}_{index_name}", "_id": id}}
|
||||
for id in ids
|
||||
]
|
||||
self.client.bulk(body=actions)
|
||||
|
||||
def reset(self):
|
||||
indices = self.client.indices.get(index=f"{self.index_prefix}_*")
|
||||
for index in indices:
|
||||
self.client.indices.delete(index=index)
|
||||
@@ -0,0 +1,354 @@
|
||||
from typing import Optional, List, Dict, Any
|
||||
from sqlalchemy import (
|
||||
cast,
|
||||
column,
|
||||
create_engine,
|
||||
Column,
|
||||
Integer,
|
||||
select,
|
||||
text,
|
||||
Text,
|
||||
values,
|
||||
)
|
||||
from sqlalchemy.sql import true
|
||||
from sqlalchemy.pool import NullPool
|
||||
|
||||
from sqlalchemy.orm import declarative_base, scoped_session, sessionmaker
|
||||
from sqlalchemy.dialects.postgresql import JSONB, array
|
||||
from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy.ext.mutable import MutableDict
|
||||
|
||||
from open_webui.apps.retrieval.vector.main import VectorItem, SearchResult, GetResult
|
||||
from open_webui.config import PGVECTOR_DB_URL
|
||||
|
||||
VECTOR_LENGTH = 1536
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
class DocumentChunk(Base):
|
||||
__tablename__ = "document_chunk"
|
||||
|
||||
id = Column(Text, primary_key=True)
|
||||
vector = Column(Vector(dim=VECTOR_LENGTH), nullable=True)
|
||||
collection_name = Column(Text, nullable=False)
|
||||
text = Column(Text, nullable=True)
|
||||
vmetadata = Column(MutableDict.as_mutable(JSONB), nullable=True)
|
||||
|
||||
|
||||
class PgvectorClient:
|
||||
def __init__(self) -> None:
|
||||
|
||||
# if no pgvector uri, use the existing database connection
|
||||
if not PGVECTOR_DB_URL:
|
||||
from open_webui.apps.webui.internal.db import Session
|
||||
|
||||
self.session = Session
|
||||
else:
|
||||
engine = create_engine(
|
||||
PGVECTOR_DB_URL, pool_pre_ping=True, poolclass=NullPool
|
||||
)
|
||||
SessionLocal = sessionmaker(
|
||||
autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
|
||||
)
|
||||
self.session = scoped_session(SessionLocal)
|
||||
|
||||
try:
|
||||
# Ensure the pgvector extension is available
|
||||
self.session.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
|
||||
|
||||
# Create the tables if they do not exist
|
||||
# Base.metadata.create_all requires a bind (engine or connection)
|
||||
# Get the connection from the session
|
||||
connection = self.session.connection()
|
||||
Base.metadata.create_all(bind=connection)
|
||||
|
||||
# Create an index on the vector column if it doesn't exist
|
||||
self.session.execute(
|
||||
text(
|
||||
"CREATE INDEX IF NOT EXISTS idx_document_chunk_vector "
|
||||
"ON document_chunk USING ivfflat (vector vector_cosine_ops) WITH (lists = 100);"
|
||||
)
|
||||
)
|
||||
self.session.execute(
|
||||
text(
|
||||
"CREATE INDEX IF NOT EXISTS idx_document_chunk_collection_name "
|
||||
"ON document_chunk (collection_name);"
|
||||
)
|
||||
)
|
||||
self.session.commit()
|
||||
print("Initialization complete.")
|
||||
except Exception as e:
|
||||
self.session.rollback()
|
||||
print(f"Error during initialization: {e}")
|
||||
raise
|
||||
|
||||
def adjust_vector_length(self, vector: List[float]) -> List[float]:
|
||||
# Adjust vector to have length VECTOR_LENGTH
|
||||
current_length = len(vector)
|
||||
if current_length < VECTOR_LENGTH:
|
||||
# Pad the vector with zeros
|
||||
vector += [0.0] * (VECTOR_LENGTH - current_length)
|
||||
elif current_length > VECTOR_LENGTH:
|
||||
raise Exception(
|
||||
f"Vector length {current_length} not supported. Max length must be <= {VECTOR_LENGTH}"
|
||||
)
|
||||
return vector
|
||||
|
||||
def insert(self, collection_name: str, items: List[VectorItem]) -> None:
|
||||
try:
|
||||
new_items = []
|
||||
for item in items:
|
||||
vector = self.adjust_vector_length(item["vector"])
|
||||
new_chunk = DocumentChunk(
|
||||
id=item["id"],
|
||||
vector=vector,
|
||||
collection_name=collection_name,
|
||||
text=item["text"],
|
||||
vmetadata=item["metadata"],
|
||||
)
|
||||
new_items.append(new_chunk)
|
||||
self.session.bulk_save_objects(new_items)
|
||||
self.session.commit()
|
||||
print(
|
||||
f"Inserted {len(new_items)} items into collection '{collection_name}'."
|
||||
)
|
||||
except Exception as e:
|
||||
self.session.rollback()
|
||||
print(f"Error during insert: {e}")
|
||||
raise
|
||||
|
||||
def upsert(self, collection_name: str, items: List[VectorItem]) -> None:
|
||||
try:
|
||||
for item in items:
|
||||
vector = self.adjust_vector_length(item["vector"])
|
||||
existing = (
|
||||
self.session.query(DocumentChunk)
|
||||
.filter(DocumentChunk.id == item["id"])
|
||||
.first()
|
||||
)
|
||||
if existing:
|
||||
existing.vector = vector
|
||||
existing.text = item["text"]
|
||||
existing.vmetadata = item["metadata"]
|
||||
existing.collection_name = (
|
||||
collection_name # Update collection_name if necessary
|
||||
)
|
||||
else:
|
||||
new_chunk = DocumentChunk(
|
||||
id=item["id"],
|
||||
vector=vector,
|
||||
collection_name=collection_name,
|
||||
text=item["text"],
|
||||
vmetadata=item["metadata"],
|
||||
)
|
||||
self.session.add(new_chunk)
|
||||
self.session.commit()
|
||||
print(f"Upserted {len(items)} items into collection '{collection_name}'.")
|
||||
except Exception as e:
|
||||
self.session.rollback()
|
||||
print(f"Error during upsert: {e}")
|
||||
raise
|
||||
|
||||
def search(
|
||||
self,
|
||||
collection_name: str,
|
||||
vectors: List[List[float]],
|
||||
limit: Optional[int] = None,
|
||||
) -> Optional[SearchResult]:
|
||||
try:
|
||||
if not vectors:
|
||||
return None
|
||||
|
||||
# Adjust query vectors to VECTOR_LENGTH
|
||||
vectors = [self.adjust_vector_length(vector) for vector in vectors]
|
||||
num_queries = len(vectors)
|
||||
|
||||
def vector_expr(vector):
|
||||
return cast(array(vector), Vector(VECTOR_LENGTH))
|
||||
|
||||
# Create the values for query vectors
|
||||
qid_col = column("qid", Integer)
|
||||
q_vector_col = column("q_vector", Vector(VECTOR_LENGTH))
|
||||
query_vectors = (
|
||||
values(qid_col, q_vector_col)
|
||||
.data(
|
||||
[(idx, vector_expr(vector)) for idx, vector in enumerate(vectors)]
|
||||
)
|
||||
.alias("query_vectors")
|
||||
)
|
||||
|
||||
# Build the lateral subquery for each query vector
|
||||
subq = (
|
||||
select(
|
||||
DocumentChunk.id,
|
||||
DocumentChunk.text,
|
||||
DocumentChunk.vmetadata,
|
||||
(
|
||||
DocumentChunk.vector.cosine_distance(query_vectors.c.q_vector)
|
||||
).label("distance"),
|
||||
)
|
||||
.where(DocumentChunk.collection_name == collection_name)
|
||||
.order_by(
|
||||
(DocumentChunk.vector.cosine_distance(query_vectors.c.q_vector))
|
||||
)
|
||||
)
|
||||
if limit is not None:
|
||||
subq = subq.limit(limit)
|
||||
subq = subq.lateral("result")
|
||||
|
||||
# Build the main query by joining query_vectors and the lateral subquery
|
||||
stmt = (
|
||||
select(
|
||||
query_vectors.c.qid,
|
||||
subq.c.id,
|
||||
subq.c.text,
|
||||
subq.c.vmetadata,
|
||||
subq.c.distance,
|
||||
)
|
||||
.select_from(query_vectors)
|
||||
.join(subq, true())
|
||||
.order_by(query_vectors.c.qid, subq.c.distance)
|
||||
)
|
||||
|
||||
result_proxy = self.session.execute(stmt)
|
||||
results = result_proxy.all()
|
||||
|
||||
ids = [[] for _ in range(num_queries)]
|
||||
distances = [[] for _ in range(num_queries)]
|
||||
documents = [[] for _ in range(num_queries)]
|
||||
metadatas = [[] for _ in range(num_queries)]
|
||||
|
||||
if not results:
|
||||
return SearchResult(
|
||||
ids=ids,
|
||||
distances=distances,
|
||||
documents=documents,
|
||||
metadatas=metadatas,
|
||||
)
|
||||
|
||||
for row in results:
|
||||
qid = int(row.qid)
|
||||
ids[qid].append(row.id)
|
||||
distances[qid].append(row.distance)
|
||||
documents[qid].append(row.text)
|
||||
metadatas[qid].append(row.vmetadata)
|
||||
|
||||
return SearchResult(
|
||||
ids=ids, distances=distances, documents=documents, metadatas=metadatas
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error during search: {e}")
|
||||
return None
|
||||
|
||||
def query(
|
||||
self, collection_name: str, filter: Dict[str, Any], limit: Optional[int] = None
|
||||
) -> Optional[GetResult]:
|
||||
try:
|
||||
query = self.session.query(DocumentChunk).filter(
|
||||
DocumentChunk.collection_name == collection_name
|
||||
)
|
||||
|
||||
for key, value in filter.items():
|
||||
query = query.filter(DocumentChunk.vmetadata[key].astext == str(value))
|
||||
|
||||
if limit is not None:
|
||||
query = query.limit(limit)
|
||||
|
||||
results = query.all()
|
||||
|
||||
if not results:
|
||||
return None
|
||||
|
||||
ids = [[result.id for result in results]]
|
||||
documents = [[result.text for result in results]]
|
||||
metadatas = [[result.vmetadata for result in results]]
|
||||
|
||||
return GetResult(
|
||||
ids=ids,
|
||||
documents=documents,
|
||||
metadatas=metadatas,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error during query: {e}")
|
||||
return None
|
||||
|
||||
def get(
|
||||
self, collection_name: str, limit: Optional[int] = None
|
||||
) -> Optional[GetResult]:
|
||||
try:
|
||||
query = self.session.query(DocumentChunk).filter(
|
||||
DocumentChunk.collection_name == collection_name
|
||||
)
|
||||
if limit is not None:
|
||||
query = query.limit(limit)
|
||||
|
||||
results = query.all()
|
||||
|
||||
if not results:
|
||||
return None
|
||||
|
||||
ids = [[result.id for result in results]]
|
||||
documents = [[result.text for result in results]]
|
||||
metadatas = [[result.vmetadata for result in results]]
|
||||
|
||||
return GetResult(ids=ids, documents=documents, metadatas=metadatas)
|
||||
except Exception as e:
|
||||
print(f"Error during get: {e}")
|
||||
return None
|
||||
|
||||
def delete(
|
||||
self,
|
||||
collection_name: str,
|
||||
ids: Optional[List[str]] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
try:
|
||||
query = self.session.query(DocumentChunk).filter(
|
||||
DocumentChunk.collection_name == collection_name
|
||||
)
|
||||
if ids:
|
||||
query = query.filter(DocumentChunk.id.in_(ids))
|
||||
if filter:
|
||||
for key, value in filter.items():
|
||||
query = query.filter(
|
||||
DocumentChunk.vmetadata[key].astext == str(value)
|
||||
)
|
||||
deleted = query.delete(synchronize_session=False)
|
||||
self.session.commit()
|
||||
print(f"Deleted {deleted} items from collection '{collection_name}'.")
|
||||
except Exception as e:
|
||||
self.session.rollback()
|
||||
print(f"Error during delete: {e}")
|
||||
raise
|
||||
|
||||
def reset(self) -> None:
|
||||
try:
|
||||
deleted = self.session.query(DocumentChunk).delete()
|
||||
self.session.commit()
|
||||
print(
|
||||
f"Reset complete. Deleted {deleted} items from 'document_chunk' table."
|
||||
)
|
||||
except Exception as e:
|
||||
self.session.rollback()
|
||||
print(f"Error during reset: {e}")
|
||||
raise
|
||||
|
||||
def close(self) -> None:
|
||||
pass
|
||||
|
||||
def has_collection(self, collection_name: str) -> bool:
|
||||
try:
|
||||
exists = (
|
||||
self.session.query(DocumentChunk)
|
||||
.filter(DocumentChunk.collection_name == collection_name)
|
||||
.first()
|
||||
is not None
|
||||
)
|
||||
return exists
|
||||
except Exception as e:
|
||||
print(f"Error checking collection existence: {e}")
|
||||
return False
|
||||
|
||||
def delete_collection(self, collection_name: str) -> None:
|
||||
self.delete(collection_name)
|
||||
print(f"Collection '{collection_name}' deleted.")
|
||||
@@ -5,7 +5,7 @@ from qdrant_client.http.models import PointStruct
|
||||
from qdrant_client.models import models
|
||||
|
||||
from open_webui.apps.retrieval.vector.main import VectorItem, SearchResult, GetResult
|
||||
from open_webui.config import QDRANT_URI
|
||||
from open_webui.config import QDRANT_URI, QDRANT_API_KEY
|
||||
|
||||
NO_LIMIT = 999999999
|
||||
|
||||
@@ -14,7 +14,12 @@ class QdrantClient:
|
||||
def __init__(self):
|
||||
self.collection_prefix = "open-webui"
|
||||
self.QDRANT_URI = QDRANT_URI
|
||||
self.client = Qclient(url=self.QDRANT_URI) if self.QDRANT_URI else None
|
||||
self.QDRANT_API_KEY = QDRANT_API_KEY
|
||||
self.client = (
|
||||
Qclient(url=self.QDRANT_URI, api_key=self.QDRANT_API_KEY)
|
||||
if self.QDRANT_URI
|
||||
else None
|
||||
)
|
||||
|
||||
def _result_to_get_result(self, points) -> GetResult:
|
||||
ids = []
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
import logging
|
||||
import os
|
||||
from pprint import pprint
|
||||
from typing import Optional
|
||||
import requests
|
||||
from open_webui.apps.retrieval.web.main import SearchResult, get_filtered_results
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
import argparse
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
||||
"""
|
||||
Documentation: https://docs.microsoft.com/en-us/bing/search-apis/bing-web-search/overview
|
||||
"""
|
||||
|
||||
|
||||
def search_bing(
|
||||
subscription_key: str,
|
||||
endpoint: str,
|
||||
locale: str,
|
||||
query: str,
|
||||
count: int,
|
||||
filter_list: Optional[list[str]] = None,
|
||||
) -> list[SearchResult]:
|
||||
mkt = locale
|
||||
params = {"q": query, "mkt": mkt, "answerCount": count}
|
||||
headers = {"Ocp-Apim-Subscription-Key": subscription_key}
|
||||
|
||||
try:
|
||||
response = requests.get(endpoint, headers=headers, params=params)
|
||||
response.raise_for_status()
|
||||
json_response = response.json()
|
||||
results = json_response.get("webPages", {}).get("value", [])
|
||||
if filter_list:
|
||||
results = get_filtered_results(results, filter_list)
|
||||
return [
|
||||
SearchResult(
|
||||
link=result["url"],
|
||||
title=result.get("name"),
|
||||
snippet=result.get("snippet"),
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
except Exception as ex:
|
||||
log.error(f"Error: {ex}")
|
||||
raise ex
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Search Bing from the command line.")
|
||||
parser.add_argument(
|
||||
"query",
|
||||
type=str,
|
||||
default="Top 10 international news today",
|
||||
help="The search query.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--count", type=int, default=10, help="Number of search results to return."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--filter", nargs="*", help="List of filters to apply to the search results."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--locale",
|
||||
type=str,
|
||||
default="en-US",
|
||||
help="The locale to use for the search, maps to market in api",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
results = search_bing(args.locale, args.query, args.count, args.filter)
|
||||
pprint(results)
|
||||
@@ -9,7 +9,7 @@ log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
||||
|
||||
|
||||
def search_jina(query: str, count: int) -> list[SearchResult]:
|
||||
def search_jina(api_key: str, query: str, count: int) -> list[SearchResult]:
|
||||
"""
|
||||
Search using Jina's Search API and return the results as a list of SearchResult objects.
|
||||
Args:
|
||||
@@ -20,9 +20,7 @@ def search_jina(query: str, count: int) -> list[SearchResult]:
|
||||
list[SearchResult]: A list of search results
|
||||
"""
|
||||
jina_search_endpoint = "https://s.jina.ai/"
|
||||
headers = {
|
||||
"Accept": "application/json",
|
||||
}
|
||||
headers = {"Accept": "application/json", "Authorization": f"Bearer {api_key}"}
|
||||
url = str(URL(jina_search_endpoint + query))
|
||||
response = requests.get(url, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
{
|
||||
"_type": "SearchResponse",
|
||||
"queryContext": {
|
||||
"originalQuery": "Top 10 international results"
|
||||
},
|
||||
"webPages": {
|
||||
"webSearchUrl": "https://www.bing.com/search?q=Top+10+international+results",
|
||||
"totalEstimatedMatches": 687,
|
||||
"value": [
|
||||
{
|
||||
"id": "https://api.bing.microsoft.com/api/v7/#WebPages.0",
|
||||
"name": "2024 Mexican Grand Prix - F1 results and latest standings ... - PlanetF1",
|
||||
"url": "https://www.planetf1.com/news/f1-results-2024-mexican-grand-prix-race-standings",
|
||||
"datePublished": "2024-10-27T00:00:00.0000000",
|
||||
"datePublishedFreshnessText": "1 day ago",
|
||||
"isFamilyFriendly": true,
|
||||
"displayUrl": "https://www.planetf1.com/news/f1-results-2024-mexican-grand-prix-race-standings",
|
||||
"snippet": "Nico Hulkenberg and Pierre Gasly completed the top 10. A full report of the Mexican Grand Prix is available at the bottom of this article. F1 results – 2024 Mexican Grand Prix",
|
||||
"dateLastCrawled": "2024-10-28T07:15:00.0000000Z",
|
||||
"cachedPageUrl": "https://cc.bingj.com/cache.aspx?q=Top+10+international+results&d=916492551782&mkt=en-US&setlang=en-US&w=zBsfaAPyF2tUrHFHr_vFFdUm8sng4g34",
|
||||
"language": "en",
|
||||
"isNavigational": false,
|
||||
"noCache": false
|
||||
},
|
||||
{
|
||||
"id": "https://api.bing.microsoft.com/api/v7/#WebPages.1",
|
||||
"name": "F1 Results Today: HUGE Verstappen penalties cause major title change",
|
||||
"url": "https://www.gpfans.com/en/f1-news/1033512/f1-results-today-mexican-grand-prix-huge-max-verstappen-penalties-cause-major-title-change/",
|
||||
"datePublished": "2024-10-27T00:00:00.0000000",
|
||||
"datePublishedFreshnessText": "1 day ago",
|
||||
"isFamilyFriendly": true,
|
||||
"displayUrl": "https://www.gpfans.com/en/f1-news/1033512/f1-results-today-mexican-grand-prix-huge-max...",
|
||||
"snippet": "Elsewhere, Mercedes duo Lewis Hamilton and George Russell came home in P4 and P5 respectively. Meanwhile, the surprise package of the day were Haas, with both Kevin Magnussen and Nico Hulkenberg finishing inside the points.. READ MORE: RB star issues apology after red flag CRASH at Mexican GP Mexican Grand Prix 2024 results. 1. Carlos Sainz [Ferrari] 2. Lando Norris [McLaren] - +4.705",
|
||||
"dateLastCrawled": "2024-10-28T06:06:00.0000000Z",
|
||||
"cachedPageUrl": "https://cc.bingj.com/cache.aspx?q=Top+10+international+results&d=2840656522642&mkt=en-US&setlang=en-US&w=-Tbkwxnq52jZCvG7l3CtgcwT1vwAjIUD",
|
||||
"language": "en",
|
||||
"isNavigational": false,
|
||||
"noCache": false
|
||||
},
|
||||
{
|
||||
"id": "https://api.bing.microsoft.com/api/v7/#WebPages.2",
|
||||
"name": "International Power Rankings: England flying, Kangaroos cruising, Fiji rise",
|
||||
"url": "https://www.loverugbyleague.com/post/international-power-rankings-england-flying-kangaroos-cruising-fiji-rise",
|
||||
"datePublished": "2024-10-28T00:00:00.0000000",
|
||||
"datePublishedFreshnessText": "7 hours ago",
|
||||
"isFamilyFriendly": true,
|
||||
"displayUrl": "https://www.loverugbyleague.com/post/international-power-rankings-england-flying...",
|
||||
"snippet": "LRL RECOMMENDS: England player ratings from first Test against Samoa as omnificent George Williams scores perfect 10. 2. Australia (Men) – SAME. The Kangaroos remain 2nd in our Power Rankings after their 22-10 win against New Zealand in Christchurch on Sunday. As was the case in their win against Tonga last week, Mal Meninga’s side weren ...",
|
||||
"dateLastCrawled": "2024-10-28T07:09:00.0000000Z",
|
||||
"cachedPageUrl": "https://cc.bingj.com/cache.aspx?q=Top+10+international+results&d=1535008462672&mkt=en-US&setlang=en-US&w=82ujhH4Kp0iuhCS7wh1xLUFYUeetaVVm",
|
||||
"language": "en",
|
||||
"isNavigational": false,
|
||||
"noCache": false
|
||||
}
|
||||
],
|
||||
"someResultsRemoved": true
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user