open-webui/backend/open_webui/apps/retrieval/main.py

1401 lines
48 KiB
Python
Raw Normal View History

2024-10-03 04:14:58 +00:00
# TODO: Merge this with the webui_app and make it a single app
2024-08-27 22:10:27 +00:00
import json
import logging
import mimetypes
import os
import shutil
2024-09-28 00:23:09 +00:00
2024-08-27 22:10:27 +00:00
import uuid
2024-06-08 04:18:04 +00:00
from datetime import datetime
2024-02-18 05:06:08 +00:00
from pathlib import Path
2024-08-27 22:10:27 +00:00
from typing import Iterator, Optional, Sequence, Union
2024-01-07 06:59:22 +00:00
2024-09-10 01:27:50 +00:00
from fastapi import Depends, FastAPI, File, Form, HTTPException, UploadFile, status
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
2024-10-26 04:46:14 +00:00
import tiktoken
2024-09-10 01:27:50 +00:00
2024-10-13 10:02:02 +00:00
2024-10-21 06:45:15 +00:00
from open_webui.storage.provider import Storage
2024-10-13 10:02:02 +00:00
from open_webui.apps.webui.models.knowledge import Knowledges
2024-09-28 00:23:09 +00:00
from open_webui.apps.retrieval.vector.connector import VECTOR_DB_CLIENT
# Document loaders
2024-09-28 17:17:21 +00:00
from open_webui.apps.retrieval.loaders.main import Loader
2024-09-28 00:23:09 +00:00
# Web search engines
from open_webui.apps.retrieval.web.main import SearchResult
from open_webui.apps.retrieval.web.utils import get_web_loader
from open_webui.apps.retrieval.web.brave import search_brave
from open_webui.apps.retrieval.web.duckduckgo import search_duckduckgo
from open_webui.apps.retrieval.web.google_pse import search_google_pse
from open_webui.apps.retrieval.web.jina_search import search_jina
from open_webui.apps.retrieval.web.searchapi import search_searchapi
from open_webui.apps.retrieval.web.searxng import search_searxng
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
2024-10-28 09:33:52 +00:00
from open_webui.apps.retrieval.web.bing import search_bing
2024-09-28 00:23:09 +00:00
2024-09-27 23:28:45 +00:00
from open_webui.apps.retrieval.utils import (
2024-08-27 22:10:27 +00:00
get_embedding_function,
get_model_path,
query_collection,
query_collection_with_hybrid_search,
query_doc,
query_doc_with_hybrid_search,
2024-02-18 05:06:08 +00:00
)
2024-09-28 00:23:09 +00:00
from open_webui.apps.webui.models.files import Files
from open_webui.config import (
2024-08-27 22:10:27 +00:00
BRAVE_SEARCH_API_KEY,
2024-10-26 04:46:14 +00:00
TIKTOKEN_ENCODING_NAME,
2024-10-13 09:07:50 +00:00
RAG_TEXT_SPLITTER,
2024-08-27 22:10:27 +00:00
CHUNK_OVERLAP,
CHUNK_SIZE,
2024-07-02 00:11:09 +00:00
CONTENT_EXTRACTION_ENGINE,
2024-08-27 22:10:27 +00:00
CORS_ALLOW_ORIGIN,
ENABLE_RAG_HYBRID_SEARCH,
ENABLE_RAG_LOCAL_WEB_FETCH,
ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
ENABLE_RAG_WEB_SEARCH,
ENV,
GOOGLE_PSE_API_KEY,
GOOGLE_PSE_ENGINE_ID,
PDF_EXTRACT_IMAGES,
2024-04-14 21:55:00 +00:00
RAG_EMBEDDING_ENGINE,
2024-02-18 19:16:10 +00:00
RAG_EMBEDDING_MODEL,
2024-04-25 12:49:59 +00:00
RAG_EMBEDDING_MODEL_AUTO_UPDATE,
RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
RAG_EMBEDDING_BATCH_SIZE,
2024-08-27 22:10:27 +00:00
RAG_FILE_MAX_COUNT,
RAG_FILE_MAX_SIZE,
RAG_OPENAI_API_BASE_URL,
RAG_OPENAI_API_KEY,
RAG_RELEVANCE_THRESHOLD,
2024-04-22 20:49:58 +00:00
RAG_RERANKING_MODEL,
2024-04-25 12:49:59 +00:00
RAG_RERANKING_MODEL_AUTO_UPDATE,
2024-04-22 20:49:58 +00:00
RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
2024-09-16 10:01:04 +00:00
DEFAULT_RAG_TEMPLATE,
2024-02-18 06:41:03 +00:00
RAG_TEMPLATE,
2024-08-27 22:10:27 +00:00
RAG_TOP_K,
RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
2024-08-27 22:10:27 +00:00
RAG_WEB_SEARCH_ENGINE,
RAG_WEB_SEARCH_RESULT_COUNT,
2024-11-04 01:07:24 +00:00
JINA_API_KEY,
2024-08-27 22:10:27 +00:00
SEARCHAPI_API_KEY,
SEARCHAPI_ENGINE,
2024-06-02 02:03:56 +00:00
SEARXNG_QUERY_URL,
SERPER_API_KEY,
SERPLY_API_KEY,
2024-08-27 22:10:27 +00:00
SERPSTACK_API_KEY,
SERPSTACK_HTTPS,
TAVILY_API_KEY,
2024-11-04 01:07:24 +00:00
BING_SEARCH_V7_ENDPOINT,
BING_SEARCH_V7_SUBSCRIPTION_KEY,
2024-08-27 22:10:27 +00:00
TIKA_SERVER_URL,
UPLOAD_DIR,
YOUTUBE_LOADER_LANGUAGE,
2024-10-28 09:33:52 +00:00
DEFAULT_LOCALE,
2024-08-27 22:10:27 +00:00
AppConfig,
2024-02-18 05:06:08 +00:00
)
from open_webui.constants import ERROR_MESSAGES
from open_webui.env import (
SRC_LOG_LEVELS,
DEVICE_TYPE,
DOCKER,
)
2024-09-10 01:27:50 +00:00
from open_webui.utils.misc import (
calculate_sha256,
calculate_sha256_string,
extract_folders_after_data_docs,
sanitize_filename,
)
from open_webui.utils.utils import get_admin_user, get_verified_user
2024-10-13 09:07:50 +00:00
from langchain.text_splitter import RecursiveCharacterTextSplitter, TokenTextSplitter
2024-08-27 22:10:27 +00:00
from langchain_community.document_loaders import (
YoutubeLoader,
)
from langchain_core.documents import Document
2024-09-28 00:23:09 +00:00
2024-01-07 06:59:22 +00:00
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
2024-01-07 06:07:20 +00:00
app = FastAPI()
app.state.config = AppConfig()
2024-04-26 18:41:39 +00:00
app.state.config.TOP_K = RAG_TOP_K
app.state.config.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD
2024-08-27 13:51:40 +00:00
app.state.config.FILE_MAX_SIZE = RAG_FILE_MAX_SIZE
app.state.config.FILE_MAX_COUNT = RAG_FILE_MAX_COUNT
app.state.config.ENABLE_RAG_HYBRID_SEARCH = ENABLE_RAG_HYBRID_SEARCH
app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
)
2024-04-25 22:31:21 +00:00
2024-07-02 00:11:09 +00:00
app.state.config.CONTENT_EXTRACTION_ENGINE = CONTENT_EXTRACTION_ENGINE
app.state.config.TIKA_SERVER_URL = TIKA_SERVER_URL
2024-10-13 09:07:50 +00:00
app.state.config.TEXT_SPLITTER = RAG_TEXT_SPLITTER
2024-10-26 04:46:14 +00:00
app.state.config.TIKTOKEN_ENCODING_NAME = TIKTOKEN_ENCODING_NAME
2024-10-13 09:07:50 +00:00
app.state.config.CHUNK_SIZE = CHUNK_SIZE
app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP
2024-04-10 07:33:45 +00:00
app.state.config.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
app.state.config.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
app.state.config.RAG_EMBEDDING_BATCH_SIZE = RAG_EMBEDDING_BATCH_SIZE
app.state.config.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
app.state.config.RAG_TEMPLATE = RAG_TEMPLATE
2024-04-10 07:33:45 +00:00
app.state.config.OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
app.state.config.OPENAI_API_KEY = RAG_OPENAI_API_KEY
2024-04-10 07:33:45 +00:00
app.state.config.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
2024-04-14 21:55:00 +00:00
app.state.config.YOUTUBE_LOADER_LANGUAGE = YOUTUBE_LOADER_LANGUAGE
2024-05-08 17:47:05 +00:00
app.state.YOUTUBE_LOADER_TRANSLATION = None
2024-06-02 02:03:56 +00:00
app.state.config.ENABLE_RAG_WEB_SEARCH = ENABLE_RAG_WEB_SEARCH
2024-06-02 02:40:48 +00:00
app.state.config.RAG_WEB_SEARCH_ENGINE = RAG_WEB_SEARCH_ENGINE
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST = RAG_WEB_SEARCH_DOMAIN_FILTER_LIST
2024-06-02 02:40:48 +00:00
2024-06-02 02:03:56 +00:00
app.state.config.SEARXNG_QUERY_URL = SEARXNG_QUERY_URL
app.state.config.GOOGLE_PSE_API_KEY = GOOGLE_PSE_API_KEY
app.state.config.GOOGLE_PSE_ENGINE_ID = GOOGLE_PSE_ENGINE_ID
2024-06-02 02:40:48 +00:00
app.state.config.BRAVE_SEARCH_API_KEY = BRAVE_SEARCH_API_KEY
2024-06-02 02:03:56 +00:00
app.state.config.SERPSTACK_API_KEY = SERPSTACK_API_KEY
app.state.config.SERPSTACK_HTTPS = SERPSTACK_HTTPS
app.state.config.SERPER_API_KEY = SERPER_API_KEY
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
2024-11-04 01:07:24 +00:00
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
2024-06-02 02:03:56 +00:00
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
2024-04-25 12:49:59 +00:00
def update_embedding_model(
embedding_model: str,
2024-09-17 20:58:06 +00:00
auto_update: bool = False,
2024-04-25 12:49:59 +00:00
):
if embedding_model and app.state.config.RAG_EMBEDDING_ENGINE == "":
2024-10-13 07:21:06 +00:00
from sentence_transformers import SentenceTransformer
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
2024-04-25 12:49:59 +00:00
else:
app.state.sentence_transformer_ef = None
def update_reranking_model(
reranking_model: str,
2024-09-17 20:58:06 +00:00
auto_update: bool = False,
2024-04-25 12:49:59 +00:00
):
if reranking_model:
2024-09-16 10:36:43 +00:00
if any(model in reranking_model for model in ["jinaai/jina-colbert-v2"]):
2024-09-17 21:07:04 +00:00
try:
from open_webui.apps.retrieval.models.colbert import ColBERT
2024-09-19 16:40:23 +00:00
app.state.sentence_transformer_rf = ColBERT(
2024-09-29 21:20:37 +00:00
get_model_path(reranking_model, auto_update),
env="docker" if DOCKER else None,
2024-09-19 16:40:23 +00:00
)
2024-09-17 21:13:51 +00:00
except Exception as e:
log.error(f"ColBERT: {e}")
2024-09-17 21:07:04 +00:00
app.state.sentence_transformer_rf = None
app.state.config.ENABLE_RAG_HYBRID_SEARCH = False
2024-09-16 09:46:39 +00:00
else:
import sentence_transformers
2024-09-16 09:46:39 +00:00
try:
app.state.sentence_transformer_rf = sentence_transformers.CrossEncoder(
2024-09-17 20:58:06 +00:00
get_model_path(reranking_model, auto_update),
2024-09-16 09:46:39 +00:00
device=DEVICE_TYPE,
trust_remote_code=RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
)
except:
log.error("CrossEncoder error")
app.state.sentence_transformer_rf = None
app.state.config.ENABLE_RAG_HYBRID_SEARCH = False
2024-04-25 12:49:59 +00:00
else:
app.state.sentence_transformer_rf = None
update_embedding_model(
app.state.config.RAG_EMBEDDING_MODEL,
2024-04-25 12:49:59 +00:00
RAG_EMBEDDING_MODEL_AUTO_UPDATE,
)
update_reranking_model(
app.state.config.RAG_RERANKING_MODEL,
2024-04-25 12:49:59 +00:00
RAG_RERANKING_MODEL_AUTO_UPDATE,
)
2024-02-18 06:29:52 +00:00
2024-04-27 19:38:50 +00:00
app.state.EMBEDDING_FUNCTION = get_embedding_function(
app.state.config.RAG_EMBEDDING_ENGINE,
app.state.config.RAG_EMBEDDING_MODEL,
2024-04-27 19:38:50 +00:00
app.state.sentence_transformer_ef,
app.state.config.OPENAI_API_KEY,
app.state.config.OPENAI_API_BASE_URL,
app.state.config.RAG_EMBEDDING_BATCH_SIZE,
2024-04-27 19:38:50 +00:00
)
2024-01-07 06:07:20 +00:00
app.add_middleware(
CORSMiddleware,
2024-08-18 21:17:26 +00:00
allow_origins=CORS_ALLOW_ORIGIN,
2024-01-07 06:07:20 +00:00
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
2024-01-07 07:40:51 +00:00
class CollectionNameForm(BaseModel):
2024-09-28 00:29:08 +00:00
collection_name: Optional[str] = None
2024-01-07 06:59:22 +00:00
2024-09-28 00:29:08 +00:00
class ProcessUrlForm(CollectionNameForm):
2024-01-07 07:40:51 +00:00
url: str
2024-03-26 06:47:08 +00:00
class SearchForm(CollectionNameForm):
query: str
2024-01-07 06:07:20 +00:00
@app.get("/")
async def get_status():
2024-02-18 06:29:52 +00:00
return {
"status": True,
"chunk_size": app.state.config.CHUNK_SIZE,
"chunk_overlap": app.state.config.CHUNK_OVERLAP,
"template": app.state.config.RAG_TEMPLATE,
"embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
"embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
"reranking_model": app.state.config.RAG_RERANKING_MODEL,
"embedding_batch_size": app.state.config.RAG_EMBEDDING_BATCH_SIZE,
2024-02-19 19:05:45 +00:00
}
2024-04-14 22:31:40 +00:00
@app.get("/embedding")
async def get_embedding_config(user=Depends(get_admin_user)):
2024-02-19 19:05:45 +00:00
return {
"status": True,
"embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
"embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
"embedding_batch_size": app.state.config.RAG_EMBEDDING_BATCH_SIZE,
2024-04-14 23:15:39 +00:00
"openai_config": {
"url": app.state.config.OPENAI_API_BASE_URL,
"key": app.state.config.OPENAI_API_KEY,
2024-04-14 23:15:39 +00:00
},
2024-02-19 19:05:45 +00:00
}
2024-04-22 20:49:58 +00:00
@app.get("/reranking")
async def get_reraanking_config(user=Depends(get_admin_user)):
return {
"status": True,
"reranking_model": app.state.config.RAG_RERANKING_MODEL,
}
2024-04-22 20:49:58 +00:00
2024-04-14 23:15:39 +00:00
class OpenAIConfigForm(BaseModel):
url: str
key: str
2024-02-19 19:05:45 +00:00
class EmbeddingModelUpdateForm(BaseModel):
2024-04-14 23:15:39 +00:00
openai_config: Optional[OpenAIConfigForm] = None
2024-04-14 22:31:40 +00:00
embedding_engine: str
2024-02-19 19:05:45 +00:00
embedding_model: str
embedding_batch_size: Optional[int] = 1
2024-02-19 19:05:45 +00:00
2024-04-14 22:31:40 +00:00
@app.post("/embedding/update")
async def update_embedding_config(
2024-02-19 19:05:45 +00:00
form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
):
2024-04-04 18:07:42 +00:00
log.info(
f"Updating embedding model: {app.state.config.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
2024-02-19 19:05:45 +00:00
)
try:
app.state.config.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
app.state.config.RAG_EMBEDDING_MODEL = form_data.embedding_model
2024-04-14 22:31:40 +00:00
if app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
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
app.state.config.RAG_EMBEDDING_BATCH_SIZE = form_data.embedding_batch_size
2024-05-19 13:51:32 +00:00
update_embedding_model(app.state.config.RAG_EMBEDDING_MODEL)
2024-04-27 19:38:50 +00:00
app.state.EMBEDDING_FUNCTION = get_embedding_function(
app.state.config.RAG_EMBEDDING_ENGINE,
app.state.config.RAG_EMBEDDING_MODEL,
2024-04-27 19:38:50 +00:00
app.state.sentence_transformer_ef,
app.state.config.OPENAI_API_KEY,
app.state.config.OPENAI_API_BASE_URL,
app.state.config.RAG_EMBEDDING_BATCH_SIZE,
2024-04-27 19:38:50 +00:00
)
return {
"status": True,
"embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
"embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
"embedding_batch_size": app.state.config.RAG_EMBEDDING_BATCH_SIZE,
2024-04-14 23:15:39 +00:00
"openai_config": {
"url": app.state.config.OPENAI_API_BASE_URL,
"key": app.state.config.OPENAI_API_KEY,
2024-04-14 23:15:39 +00:00
},
}
except Exception as e:
log.exception(f"Problem updating embedding model: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=ERROR_MESSAGES.DEFAULT(e),
)
2024-02-18 06:29:52 +00:00
2024-04-22 20:49:58 +00:00
class RerankingModelUpdateForm(BaseModel):
reranking_model: str
2024-04-22 20:49:58 +00:00
@app.post("/reranking/update")
async def update_reranking_config(
form_data: RerankingModelUpdateForm, user=Depends(get_admin_user)
):
log.info(
f"Updating reranking model: {app.state.config.RAG_RERANKING_MODEL} to {form_data.reranking_model}"
2024-04-22 20:49:58 +00:00
)
try:
app.state.config.RAG_RERANKING_MODEL = form_data.reranking_model
update_reranking_model(app.state.config.RAG_RERANKING_MODEL, True)
2024-04-22 20:49:58 +00:00
return {
"status": True,
"reranking_model": app.state.config.RAG_RERANKING_MODEL,
2024-04-22 20:49:58 +00:00
}
except Exception as e:
log.exception(f"Problem updating reranking model: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=ERROR_MESSAGES.DEFAULT(e),
)
2024-03-10 20:32:34 +00:00
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
2024-02-18 06:29:52 +00:00
return {
"status": True,
"pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
2024-07-02 00:11:09 +00:00
"content_extraction": {
"engine": app.state.config.CONTENT_EXTRACTION_ENGINE,
"tika_server_url": app.state.config.TIKA_SERVER_URL,
},
2024-03-10 20:32:34 +00:00
"chunk": {
2024-10-13 11:24:13 +00:00
"text_splitter": app.state.config.TEXT_SPLITTER,
"chunk_size": app.state.config.CHUNK_SIZE,
"chunk_overlap": app.state.config.CHUNK_OVERLAP,
2024-03-10 20:32:34 +00:00
},
2024-10-13 11:24:13 +00:00
"file": {
"max_size": app.state.config.FILE_MAX_SIZE,
"max_count": app.state.config.FILE_MAX_COUNT,
},
2024-05-08 17:47:05 +00:00
"youtube": {
"language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
2024-05-08 17:47:05 +00:00
"translation": app.state.YOUTUBE_LOADER_TRANSLATION,
},
2024-06-02 02:03:56 +00:00
"web": {
"web_loader_ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
2024-06-02 02:03:56 +00:00
"search": {
2024-06-02 03:08:08 +00:00
"enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
2024-06-02 02:40:48 +00:00
"engine": app.state.config.RAG_WEB_SEARCH_ENGINE,
2024-06-02 02:03:56 +00:00
"searxng_query_url": app.state.config.SEARXNG_QUERY_URL,
"google_pse_api_key": app.state.config.GOOGLE_PSE_API_KEY,
"google_pse_engine_id": app.state.config.GOOGLE_PSE_ENGINE_ID,
2024-06-02 02:40:48 +00:00
"brave_search_api_key": app.state.config.BRAVE_SEARCH_API_KEY,
2024-06-02 02:03:56 +00:00
"serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
"serpstack_https": app.state.config.SERPSTACK_HTTPS,
"serper_api_key": app.state.config.SERPER_API_KEY,
"serply_api_key": app.state.config.SERPLY_API_KEY,
"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,
2024-11-04 01:07:24 +00:00
"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,
2024-06-02 02:03:56 +00:00
"result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
"concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
2024-06-02 02:40:48 +00:00
},
2024-06-02 02:03:56 +00:00
},
2024-02-18 06:29:52 +00:00
}
2024-08-27 15:05:24 +00:00
class FileConfig(BaseModel):
max_size: Optional[int] = None
max_count: Optional[int] = None
2024-07-02 00:11:09 +00:00
class ContentExtractionConfig(BaseModel):
engine: str = ""
tika_server_url: Optional[str] = None
2024-02-18 06:29:52 +00:00
class ChunkParamUpdateForm(BaseModel):
2024-10-13 11:24:13 +00:00
text_splitter: Optional[str] = None
2024-02-18 06:29:52 +00:00
chunk_size: int
chunk_overlap: int
2024-05-08 17:47:05 +00:00
class YoutubeLoaderConfig(BaseModel):
2024-08-14 12:46:31 +00:00
language: list[str]
2024-05-08 17:47:05 +00:00
translation: Optional[str] = None
2024-06-02 02:03:56 +00:00
class WebSearchConfig(BaseModel):
2024-06-02 03:08:08 +00:00
enabled: bool
2024-06-02 02:40:48 +00:00
engine: Optional[str] = None
2024-06-02 02:03:56 +00:00
searxng_query_url: Optional[str] = None
google_pse_api_key: Optional[str] = None
google_pse_engine_id: Optional[str] = None
2024-06-02 02:40:48 +00:00
brave_search_api_key: Optional[str] = None
2024-06-02 02:03:56 +00:00
serpstack_api_key: Optional[str] = None
serpstack_https: Optional[bool] = None
serper_api_key: Optional[str] = None
serply_api_key: Optional[str] = None
tavily_api_key: Optional[str] = None
searchapi_api_key: Optional[str] = None
searchapi_engine: Optional[str] = None
2024-11-04 01:07:24 +00:00
jina_api_key: Optional[str] = None
bing_search_v7_endpoint: Optional[str] = None
bing_search_v7_subscription_key: Optional[str] = None
2024-06-02 02:03:56 +00:00
result_count: Optional[int] = None
concurrent_requests: Optional[int] = None
2024-06-02 02:40:48 +00:00
class WebConfig(BaseModel):
search: WebSearchConfig
web_loader_ssl_verification: Optional[bool] = None
2024-03-10 20:32:34 +00:00
class ConfigUpdateForm(BaseModel):
pdf_extract_images: Optional[bool] = None
2024-08-27 15:05:24 +00:00
file: Optional[FileConfig] = None
2024-07-02 00:11:09 +00:00
content_extraction: Optional[ContentExtractionConfig] = None
chunk: Optional[ChunkParamUpdateForm] = None
2024-05-08 17:47:05 +00:00
youtube: Optional[YoutubeLoaderConfig] = None
2024-06-02 02:40:48 +00:00
web: Optional[WebConfig] = None
2024-03-10 20:32:34 +00:00
@app.post("/config/update")
async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
app.state.config.PDF_EXTRACT_IMAGES = (
form_data.pdf_extract_images
if form_data.pdf_extract_images is not None
else app.state.config.PDF_EXTRACT_IMAGES
)
2024-08-27 15:05:24 +00:00
if form_data.file is not None:
app.state.config.FILE_MAX_SIZE = form_data.file.max_size
app.state.config.FILE_MAX_COUNT = form_data.file.max_count
2024-07-02 00:11:09 +00:00
if form_data.content_extraction is not None:
log.info(f"Updating text settings: {form_data.content_extraction}")
app.state.config.CONTENT_EXTRACTION_ENGINE = form_data.content_extraction.engine
app.state.config.TIKA_SERVER_URL = form_data.content_extraction.tika_server_url
2024-06-02 02:40:48 +00:00
if form_data.chunk is not None:
2024-10-13 11:24:13 +00:00
app.state.config.TEXT_SPLITTER = form_data.chunk.text_splitter
2024-06-02 02:40:48 +00:00
app.state.config.CHUNK_SIZE = form_data.chunk.chunk_size
app.state.config.CHUNK_OVERLAP = form_data.chunk.chunk_overlap
2024-06-02 02:40:48 +00:00
if form_data.youtube is not None:
app.state.config.YOUTUBE_LOADER_LANGUAGE = form_data.youtube.language
app.state.YOUTUBE_LOADER_TRANSLATION = form_data.youtube.translation
2024-06-02 02:40:48 +00:00
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
2024-06-02 02:40:48 +00:00
)
2024-05-08 17:47:05 +00:00
2024-06-02 03:08:08 +00:00
app.state.config.ENABLE_RAG_WEB_SEARCH = form_data.web.search.enabled
2024-06-02 02:40:48 +00:00
app.state.config.RAG_WEB_SEARCH_ENGINE = form_data.web.search.engine
app.state.config.SEARXNG_QUERY_URL = form_data.web.search.searxng_query_url
app.state.config.GOOGLE_PSE_API_KEY = form_data.web.search.google_pse_api_key
app.state.config.GOOGLE_PSE_ENGINE_ID = (
form_data.web.search.google_pse_engine_id
)
app.state.config.BRAVE_SEARCH_API_KEY = (
form_data.web.search.brave_search_api_key
)
app.state.config.SERPSTACK_API_KEY = form_data.web.search.serpstack_api_key
app.state.config.SERPSTACK_HTTPS = form_data.web.search.serpstack_https
app.state.config.SERPER_API_KEY = form_data.web.search.serper_api_key
app.state.config.SERPLY_API_KEY = form_data.web.search.serply_api_key
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
2024-08-27 22:10:27 +00:00
app.state.config.SEARCHAPI_ENGINE = form_data.web.search.searchapi_engine
2024-11-04 01:07:24 +00:00
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
)
2024-06-02 02:40:48 +00:00
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
)
2024-05-08 17:47:05 +00:00
2024-02-18 06:29:52 +00:00
return {
"status": True,
"pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
2024-08-27 13:51:40 +00:00
"file": {
"max_size": app.state.config.FILE_MAX_SIZE,
"max_count": app.state.config.FILE_MAX_COUNT,
},
2024-07-02 00:11:09 +00:00
"content_extraction": {
"engine": app.state.config.CONTENT_EXTRACTION_ENGINE,
"tika_server_url": app.state.config.TIKA_SERVER_URL,
},
2024-03-10 20:32:34 +00:00
"chunk": {
2024-10-13 11:24:13 +00:00
"text_splitter": app.state.config.TEXT_SPLITTER,
"chunk_size": app.state.config.CHUNK_SIZE,
"chunk_overlap": app.state.config.CHUNK_OVERLAP,
2024-03-10 20:32:34 +00:00
},
2024-05-08 17:47:05 +00:00
"youtube": {
"language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
2024-05-08 17:47:05 +00:00
"translation": app.state.YOUTUBE_LOADER_TRANSLATION,
},
2024-06-02 02:40:48 +00:00
"web": {
"web_loader_ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
2024-06-02 02:40:48 +00:00
"search": {
2024-06-02 03:08:08 +00:00
"enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
2024-06-02 02:40:48 +00:00
"engine": app.state.config.RAG_WEB_SEARCH_ENGINE,
"searxng_query_url": app.state.config.SEARXNG_QUERY_URL,
"google_pse_api_key": app.state.config.GOOGLE_PSE_API_KEY,
"google_pse_engine_id": app.state.config.GOOGLE_PSE_ENGINE_ID,
"brave_search_api_key": app.state.config.BRAVE_SEARCH_API_KEY,
"serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
"serpstack_https": app.state.config.SERPSTACK_HTTPS,
"serper_api_key": app.state.config.SERPER_API_KEY,
"serply_api_key": app.state.config.SERPLY_API_KEY,
"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,
2024-11-04 01:07:24 +00:00
"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,
2024-06-02 02:40:48 +00:00
"result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
"concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
},
},
2024-02-18 06:29:52 +00:00
}
2024-01-07 06:59:22 +00:00
2024-02-18 06:41:03 +00:00
@app.get("/template")
2024-06-27 18:29:59 +00:00
async def get_rag_template(user=Depends(get_verified_user)):
2024-02-18 06:41:03 +00:00
return {
"status": True,
"template": app.state.config.RAG_TEMPLATE,
2024-02-18 06:41:03 +00:00
}
2024-03-03 02:56:57 +00:00
@app.get("/query/settings")
async def get_query_settings(user=Depends(get_admin_user)):
return {
"status": True,
"template": app.state.config.RAG_TEMPLATE,
"k": app.state.config.TOP_K,
"r": app.state.config.RELEVANCE_THRESHOLD,
"hybrid": app.state.config.ENABLE_RAG_HYBRID_SEARCH,
2024-03-03 02:56:57 +00:00
}
2024-02-18 06:41:03 +00:00
2024-03-03 02:56:57 +00:00
class QuerySettingsForm(BaseModel):
k: Optional[int] = None
r: Optional[float] = None
2024-03-03 02:56:57 +00:00
template: Optional[str] = None
2024-04-25 22:31:21 +00:00
hybrid: Optional[bool] = None
2024-03-03 02:56:57 +00:00
@app.post("/query/settings/update")
async def update_query_settings(
form_data: QuerySettingsForm, user=Depends(get_admin_user)
):
2024-10-13 11:24:13 +00:00
app.state.config.RAG_TEMPLATE = form_data.template
app.state.config.TOP_K = form_data.k if form_data.k else 4
app.state.config.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
2024-10-13 11:24:13 +00:00
app.state.config.ENABLE_RAG_HYBRID_SEARCH = (
2024-05-18 02:53:38 +00:00
form_data.hybrid if form_data.hybrid else False
)
2024-04-25 22:31:21 +00:00
return {
"status": True,
"template": app.state.config.RAG_TEMPLATE,
"k": app.state.config.TOP_K,
"r": app.state.config.RELEVANCE_THRESHOLD,
"hybrid": app.state.config.ENABLE_RAG_HYBRID_SEARCH,
2024-04-25 22:31:21 +00:00
}
2024-01-07 06:59:22 +00:00
2024-09-28 00:23:09 +00:00
####################################
#
# Document process and retrieval
#
####################################
2024-02-01 21:35:41 +00:00
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)
2024-09-28 00:38:59 +00:00
def save_docs_to_vector_db(
docs,
collection_name,
metadata: Optional[dict] = None,
overwrite: bool = False,
split: bool = True,
2024-10-03 13:44:17 +00:00
add: bool = False,
2024-09-28 00:23:09 +00:00
) -> bool:
log.info(
f"save_docs_to_vector_db: document {_get_docs_info(docs)} {collection_name}"
)
2024-09-28 00:23:09 +00:00
2024-10-03 13:53:21 +00:00
# Check if entries with the same hash (metadata.hash) already exist
if metadata and "hash" in metadata:
2024-10-04 06:06:47 +00:00
result = VECTOR_DB_CLIENT.query(
2024-10-03 13:53:21 +00:00
collection_name=collection_name,
filter={"hash": metadata["hash"]},
)
2024-10-04 06:06:47 +00:00
2024-10-07 21:03:42 +00:00
if result is not None:
2024-10-04 06:06:47 +00:00
existing_doc_ids = result.ids[0]
if existing_doc_ids:
log.info(f"Document with hash {metadata['hash']} already exists")
raise ValueError(ERROR_MESSAGES.DUPLICATE_CONTENT)
2024-10-03 13:53:21 +00:00
2024-09-28 00:38:59 +00:00
if split:
2024-10-13 09:07:50 +00:00
if app.state.config.TEXT_SPLITTER in ["", "character"]:
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=app.state.config.CHUNK_SIZE,
chunk_overlap=app.state.config.CHUNK_OVERLAP,
add_start_index=True,
)
elif app.state.config.TEXT_SPLITTER == "token":
2024-10-26 05:23:21 +00:00
log.info(
f"Using token text splitter: {app.state.config.TIKTOKEN_ENCODING_NAME}"
)
tiktoken.get_encoding(str(app.state.config.TIKTOKEN_ENCODING_NAME))
2024-10-13 09:07:50 +00:00
text_splitter = TokenTextSplitter(
2024-10-26 04:47:47 +00:00
encoding_name=str(app.state.config.TIKTOKEN_ENCODING_NAME),
2024-10-13 09:07:50 +00:00
chunk_size=app.state.config.CHUNK_SIZE,
chunk_overlap=app.state.config.CHUNK_OVERLAP,
add_start_index=True,
)
else:
raise ValueError(ERROR_MESSAGES.DEFAULT("Invalid text splitter"))
2024-09-28 00:38:59 +00:00
docs = text_splitter.split_documents(docs)
2024-09-28 00:23:09 +00:00
2024-09-28 00:38:59 +00:00
if len(docs) == 0:
raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
2024-09-28 00:23:09 +00:00
texts = [doc.page_content for doc in docs]
2024-10-13 10:25:11 +00:00
metadatas = [
{
**doc.metadata,
**(metadata if metadata else {}),
"embedding_config": json.dumps(
{
"engine": app.state.config.RAG_EMBEDDING_ENGINE,
"model": app.state.config.RAG_EMBEDDING_MODEL,
}
),
}
for doc in docs
]
2024-09-28 00:23:09 +00:00
# ChromaDB does not like datetime formats
# for meta-data so convert them to string.
for metadata in metadatas:
for key, value in metadata.items():
if isinstance(value, datetime):
metadata[key] = str(value)
2024-04-14 21:55:00 +00:00
try:
2024-09-28 00:23:09 +00:00
if VECTOR_DB_CLIENT.has_collection(collection_name=collection_name):
log.info(f"collection {collection_name} already exists")
2024-10-03 13:44:17 +00:00
if overwrite:
VECTOR_DB_CLIENT.delete_collection(collection_name=collection_name)
log.info(f"deleting existing collection {collection_name}")
2024-10-17 20:08:10 +00:00
elif add is False:
2024-10-21 00:45:37 +00:00
log.info(
f"collection {collection_name} already exists, overwrite is False and add is False"
)
2024-10-03 13:44:17 +00:00
return True
2024-04-14 23:15:39 +00:00
2024-10-03 13:44:17 +00:00
log.info(f"adding to collection {collection_name}")
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.RAG_EMBEDDING_BATCH_SIZE,
2024-10-03 13:44:17 +00:00
)
embeddings = embedding_function(
list(map(lambda x: x.replace("\n", " "), texts))
)
2024-04-14 23:15:39 +00:00
2024-10-04 07:23:14 +00:00
items = [
{
"id": str(uuid.uuid4()),
"text": text,
"vector": embeddings[idx],
2024-10-13 10:25:11 +00:00
"metadata": metadatas[idx],
2024-10-04 07:23:14 +00:00
}
for idx, text in enumerate(texts)
]
2024-10-04 07:46:32 +00:00
2024-10-03 13:44:17 +00:00
VECTOR_DB_CLIENT.insert(
collection_name=collection_name,
2024-10-04 07:23:14 +00:00
items=items,
2024-10-03 13:44:17 +00:00
)
return True
2024-04-14 21:55:00 +00:00
except Exception as e:
log.exception(e)
2024-09-28 00:23:09 +00:00
return False
2024-02-01 21:35:41 +00:00
2024-09-28 00:23:09 +00:00
class ProcessFileForm(BaseModel):
file_id: str
2024-10-04 07:23:14 +00:00
content: Optional[str] = None
2024-09-28 00:23:09 +00:00
collection_name: Optional[str] = None
2024-09-28 00:23:09 +00:00
@app.post("/process/file")
def process_file(
form_data: ProcessFileForm,
user=Depends(get_verified_user),
):
try:
file = Files.get_file_by_id(form_data.file_id)
2024-06-12 08:37:53 +00:00
2024-09-28 00:38:59 +00:00
collection_name = form_data.collection_name
2024-10-05 16:58:46 +00:00
2024-09-28 00:38:59 +00:00
if collection_name is None:
2024-10-04 05:22:22 +00:00
collection_name = f"file-{file.id}"
2024-09-28 00:38:59 +00:00
2024-10-04 07:23:14 +00:00
if form_data.content:
2024-10-05 17:08:48 +00:00
# Update the content in the file
# Usage: /files/{file_id}/data/content/update
2024-10-05 17:05:12 +00:00
VECTOR_DB_CLIENT.delete(
collection_name=f"file-{file.id}",
filter={"file_id": file.id},
)
2024-10-04 07:23:14 +00:00
docs = [
Document(
page_content=form_data.content,
metadata={
"name": file.meta.get("name", file.filename),
"created_by": file.user_id,
2024-10-05 17:05:12 +00:00
"file_id": file.id,
2024-10-04 07:23:14 +00:00
**file.meta,
},
)
]
text_content = form_data.content
2024-10-05 16:58:46 +00:00
elif form_data.collection_name:
2024-10-05 17:08:48 +00:00
# Check if the file has already been processed and save the content
# Usage: /knowledge/{id}/file/add, /knowledge/{id}/file/update
2024-10-05 16:58:46 +00:00
result = VECTOR_DB_CLIENT.query(
collection_name=f"file-{file.id}", filter={"file_id": file.id}
)
2024-10-07 21:03:42 +00:00
if result is not None and len(result.ids[0]) > 0:
2024-10-05 16:58:46 +00:00
docs = [
Document(
page_content=result.documents[0][idx],
metadata=result.metadatas[0][idx],
)
for idx, id in enumerate(result.ids[0])
]
else:
docs = [
Document(
page_content=file.data.get("content", ""),
metadata={
"name": file.meta.get("name", file.filename),
"created_by": file.user_id,
2024-10-05 17:05:12 +00:00
"file_id": file.id,
2024-10-05 16:58:46 +00:00
**file.meta,
},
)
]
2024-10-04 07:23:14 +00:00
text_content = file.data.get("content", "")
else:
2024-10-05 17:08:48 +00:00
# Process the file and save the content
# Usage: /files/
2024-10-21 00:45:37 +00:00
file_path = file.path
2024-10-04 07:23:14 +00:00
if file_path:
2024-10-21 06:45:15 +00:00
file_path = Storage.get_file(file_path)
2024-10-05 16:58:46 +00:00
loader = Loader(
engine=app.state.config.CONTENT_EXTRACTION_ENGINE,
TIKA_SERVER_URL=app.state.config.TIKA_SERVER_URL,
PDF_EXTRACT_IMAGES=app.state.config.PDF_EXTRACT_IMAGES,
)
2024-10-04 07:23:14 +00:00
docs = loader.load(
file.filename, file.meta.get("content_type"), file_path
)
else:
docs = [
Document(
page_content=file.data.get("content", ""),
metadata={
"name": file.filename,
"created_by": file.user_id,
2024-10-05 17:05:12 +00:00
"file_id": file.id,
2024-10-04 07:23:14 +00:00
**file.meta,
},
)
]
text_content = " ".join([doc.page_content for doc in docs])
2024-10-03 13:44:17 +00:00
2024-09-28 08:53:25 +00:00
log.debug(f"text_content: {text_content}")
2024-10-04 05:22:22 +00:00
Files.update_file_data_by_id(
2024-10-03 13:44:17 +00:00
file.id,
2024-10-01 20:13:39 +00:00
{"content": text_content},
2024-09-28 00:56:56 +00:00
)
2024-09-28 00:23:09 +00:00
2024-10-04 07:23:14 +00:00
hash = calculate_sha256_string(text_content)
2024-10-04 05:22:22 +00:00
Files.update_file_hash_by_id(file.id, hash)
2024-09-28 00:23:09 +00:00
try:
2024-09-28 00:38:59 +00:00
result = save_docs_to_vector_db(
2024-10-03 13:44:17 +00:00
docs=docs,
collection_name=collection_name,
metadata={
2024-10-04 07:23:14 +00:00
"file_id": file.id,
2024-09-28 00:23:09 +00:00
"name": file.meta.get("name", file.filename),
2024-10-03 13:44:17 +00:00
"hash": hash,
2024-09-28 00:23:09 +00:00
},
2024-10-03 13:44:17 +00:00
add=(True if form_data.collection_name else False),
2024-09-28 00:23:09 +00:00
)
if result:
2024-10-04 05:22:22 +00:00
Files.update_file_metadata_by_id(
file.id,
{
"collection_name": collection_name,
},
)
2024-09-28 00:23:09 +00:00
return {
"status": True,
"collection_name": collection_name,
"filename": file.meta.get("name", file.filename),
2024-09-28 08:53:25 +00:00
"content": text_content,
2024-09-28 00:23:09 +00:00
}
except Exception as e:
2024-10-04 04:10:33 +00:00
raise e
2024-09-28 00:23:09 +00:00
except Exception as e:
log.exception(e)
if "No pandoc was found" in str(e):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.PANDOC_NOT_INSTALLED,
)
else:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
2024-10-04 04:10:33 +00:00
detail=str(e),
2024-09-28 00:23:09 +00:00
)
2024-09-28 00:29:08 +00:00
class ProcessTextForm(BaseModel):
2024-09-28 00:23:09 +00:00
name: str
content: str
collection_name: Optional[str] = None
2024-09-28 00:29:08 +00:00
@app.post("/process/text")
def process_text(
form_data: ProcessTextForm,
2024-09-28 00:23:09 +00:00
user=Depends(get_verified_user),
):
collection_name = form_data.collection_name
if collection_name is None:
collection_name = calculate_sha256_string(form_data.content)
2024-09-28 00:38:59 +00:00
docs = [
Document(
page_content=form_data.content,
metadata={"name": form_data.name, "created_by": user.id},
)
]
2024-09-29 16:55:26 +00:00
text_content = form_data.content
log.debug(f"text_content: {text_content}")
2024-09-28 00:38:59 +00:00
result = save_docs_to_vector_db(docs, collection_name)
2024-09-28 00:23:09 +00:00
if result:
2024-09-29 16:55:26 +00:00
return {
"status": True,
"collection_name": collection_name,
"content": text_content,
}
2024-09-28 00:23:09 +00:00
else:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=ERROR_MESSAGES.DEFAULT(),
)
@app.post("/process/youtube")
2024-09-28 00:29:08 +00:00
def process_youtube_video(form_data: ProcessUrlForm, user=Depends(get_verified_user)):
2024-09-28 00:23:09 +00:00
try:
2024-09-28 00:38:59 +00:00
collection_name = form_data.collection_name
if not collection_name:
collection_name = calculate_sha256_string(form_data.url)[:63]
2024-09-28 00:23:09 +00:00
loader = YoutubeLoader.from_youtube_url(
form_data.url,
2024-11-03 11:18:17 +00:00
add_video_info=False,
2024-09-28 00:23:09 +00:00
language=app.state.config.YOUTUBE_LOADER_LANGUAGE,
translation=app.state.YOUTUBE_LOADER_TRANSLATION,
)
2024-09-28 00:38:59 +00:00
docs = loader.load()
2024-10-07 02:44:02 +00:00
content = " ".join([doc.page_content for doc in docs])
log.debug(f"text_content: {content}")
2024-09-28 00:38:59 +00:00
save_docs_to_vector_db(docs, collection_name, overwrite=True)
2024-09-28 00:29:08 +00:00
2024-09-28 00:23:09 +00:00
return {
"status": True,
"collection_name": collection_name,
"filename": form_data.url,
2024-10-07 02:44:02 +00:00
"file": {
"data": {
"content": content,
},
"meta": {
"name": form_data.url,
},
},
2024-09-28 00:23:09 +00:00
}
except Exception as e:
log.exception(e)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
@app.post("/process/web")
2024-09-28 00:29:08 +00:00
def process_web(form_data: ProcessUrlForm, user=Depends(get_verified_user)):
2024-09-28 00:23:09 +00:00
try:
2024-09-28 00:38:59 +00:00
collection_name = form_data.collection_name
if not collection_name:
collection_name = calculate_sha256_string(form_data.url)[:63]
2024-09-28 00:23:09 +00:00
loader = get_web_loader(
form_data.url,
verify_ssl=app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
requests_per_second=app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
)
2024-09-28 00:38:59 +00:00
docs = loader.load()
2024-10-07 02:44:02 +00:00
content = " ".join([doc.page_content for doc in docs])
log.debug(f"text_content: {content}")
2024-09-28 00:38:59 +00:00
save_docs_to_vector_db(docs, collection_name, overwrite=True)
2024-09-28 00:29:08 +00:00
2024-09-28 00:23:09 +00:00
return {
"status": True,
"collection_name": collection_name,
"filename": form_data.url,
2024-10-07 02:44:02 +00:00
"file": {
"data": {
"content": content,
},
"meta": {
"name": form_data.url,
},
},
2024-09-28 00:23:09 +00:00
}
except Exception as e:
log.exception(e)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
2024-06-12 18:08:05 +00:00
2024-06-02 02:52:12 +00:00
def search_web(engine: str, query: str) -> list[SearchResult]:
"""Search the web using a search engine and return the results as a list of SearchResult objects.
Will look for a search engine API key in environment variables in the following order:
- SEARXNG_QUERY_URL
- GOOGLE_PSE_API_KEY + GOOGLE_PSE_ENGINE_ID
- BRAVE_SEARCH_API_KEY
- SERPSTACK_API_KEY
- SERPER_API_KEY
- SERPLY_API_KEY
- TAVILY_API_KEY
- SEARCHAPI_API_KEY + SEARCHAPI_ENGINE (by default `google`)
2024-06-02 02:52:12 +00:00
Args:
query (str): The query to search for
"""
# TODO: add playwright to search the web
if engine == "searxng":
if app.state.config.SEARXNG_QUERY_URL:
2024-06-02 02:57:00 +00:00
return search_searxng(
app.state.config.SEARXNG_QUERY_URL,
query,
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
2024-06-17 21:32:23 +00:00
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
2024-06-02 02:57:00 +00:00
)
2024-06-02 02:52:12 +00:00
else:
raise Exception("No SEARXNG_QUERY_URL found in environment variables")
elif engine == "google_pse":
if (
app.state.config.GOOGLE_PSE_API_KEY
and app.state.config.GOOGLE_PSE_ENGINE_ID
):
return search_google_pse(
app.state.config.GOOGLE_PSE_API_KEY,
app.state.config.GOOGLE_PSE_ENGINE_ID,
query,
2024-06-02 02:57:00 +00:00
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
2024-06-17 21:32:23 +00:00
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
2024-06-02 02:52:12 +00:00
)
else:
raise Exception(
"No GOOGLE_PSE_API_KEY or GOOGLE_PSE_ENGINE_ID found in environment variables"
)
elif engine == "brave":
if app.state.config.BRAVE_SEARCH_API_KEY:
2024-06-02 02:57:00 +00:00
return search_brave(
app.state.config.BRAVE_SEARCH_API_KEY,
query,
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
2024-06-17 21:32:23 +00:00
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
2024-06-02 02:57:00 +00:00
)
2024-06-02 02:52:12 +00:00
else:
raise Exception("No BRAVE_SEARCH_API_KEY found in environment variables")
elif engine == "serpstack":
if app.state.config.SERPSTACK_API_KEY:
return search_serpstack(
app.state.config.SERPSTACK_API_KEY,
query,
2024-06-02 02:57:00 +00:00
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
2024-06-02 02:52:12 +00:00
https_enabled=app.state.config.SERPSTACK_HTTPS,
)
else:
raise Exception("No SERPSTACK_API_KEY found in environment variables")
elif engine == "serper":
if app.state.config.SERPER_API_KEY:
2024-06-02 02:57:00 +00:00
return search_serper(
app.state.config.SERPER_API_KEY,
query,
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
2024-06-17 21:32:23 +00:00
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
2024-06-02 02:57:00 +00:00
)
2024-06-02 02:52:12 +00:00
else:
raise Exception("No SERPER_API_KEY found in environment variables")
elif engine == "serply":
if app.state.config.SERPLY_API_KEY:
return search_serply(
app.state.config.SERPLY_API_KEY,
query,
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
2024-06-17 21:32:23 +00:00
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
)
else:
raise Exception("No SERPLY_API_KEY found in environment variables")
elif engine == "duckduckgo":
2024-06-17 21:32:23 +00:00
return search_duckduckgo(
query,
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
)
elif engine == "tavily":
if app.state.config.TAVILY_API_KEY:
return search_tavily(
app.state.config.TAVILY_API_KEY,
query,
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
)
else:
raise Exception("No TAVILY_API_KEY found in environment variables")
elif engine == "searchapi":
if app.state.config.SEARCHAPI_API_KEY:
return search_searchapi(
app.state.config.SEARCHAPI_API_KEY,
app.state.config.SEARCHAPI_ENGINE,
query,
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
)
else:
raise Exception("No SEARCHAPI_API_KEY found in environment variables")
elif engine == "jina":
2024-11-04 01:07:24 +00:00
return search_jina(
app.state.config.JINA_API_KEY,
query,
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
)
2024-10-28 09:33:52 +00:00
elif engine == "bing":
return search_bing(
2024-11-04 01:07:24 +00:00
app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
app.state.config.BING_SEARCH_V7_ENDPOINT,
2024-10-28 09:33:52 +00:00
str(DEFAULT_LOCALE),
query,
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
)
2024-06-02 02:52:12 +00:00
else:
raise Exception("No search engine API key found in environment variables")
2024-09-28 00:23:09 +00:00
@app.post("/process/web/search")
def process_web_search(form_data: SearchForm, user=Depends(get_verified_user)):
try:
2024-06-12 07:18:22 +00:00
logging.info(
f"trying to web search with {app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query}"
)
2024-06-02 02:52:12 +00:00
web_results = search_web(
app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query
)
except Exception as e:
log.exception(e)
print(e)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.WEB_SEARCH_ERROR(e),
)
try:
collection_name = form_data.collection_name
if collection_name == "":
collection_name = calculate_sha256_string(form_data.query)[:63]
2024-09-28 00:38:59 +00:00
urls = [result.link for result in web_results]
loader = get_web_loader(urls, verify_ssl=app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION)
2024-09-28 00:38:59 +00:00
docs = loader.load()
2024-09-29 16:55:26 +00:00
2024-09-28 00:38:59 +00:00
save_docs_to_vector_db(docs, collection_name, overwrite=True)
return {
"status": True,
"collection_name": collection_name,
"filenames": urls,
}
except Exception as e:
log.exception(e)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
2024-09-28 00:23:09 +00:00
class QueryDocForm(BaseModel):
collection_name: str
query: str
k: Optional[int] = None
r: Optional[float] = None
hybrid: Optional[bool] = None
2024-03-24 07:40:27 +00:00
2024-09-28 00:23:09 +00:00
@app.post("/query/doc")
def query_doc_handler(
form_data: QueryDocForm,
user=Depends(get_verified_user),
):
try:
if app.state.config.ENABLE_RAG_HYBRID_SEARCH:
return query_doc_with_hybrid_search(
collection_name=form_data.collection_name,
query=form_data.query,
embedding_function=app.state.EMBEDDING_FUNCTION,
k=form_data.k if form_data.k else app.state.config.TOP_K,
reranking_function=app.state.sentence_transformer_rf,
r=(
form_data.r if form_data.r else app.state.config.RELEVANCE_THRESHOLD
),
)
else:
return query_doc(
collection_name=form_data.collection_name,
query=form_data.query,
embedding_function=app.state.EMBEDDING_FUNCTION,
k=form_data.k if form_data.k else app.state.config.TOP_K,
)
except Exception as e:
log.exception(e)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
2024-03-24 07:40:27 +00:00
2024-09-28 00:23:09 +00:00
class QueryCollectionsForm(BaseModel):
collection_names: list[str]
query: str
k: Optional[int] = None
r: Optional[float] = None
hybrid: Optional[bool] = None
2024-03-26 06:47:08 +00:00
2024-03-24 07:40:27 +00:00
2024-09-28 00:23:09 +00:00
@app.post("/query/collection")
def query_collection_handler(
form_data: QueryCollectionsForm,
2024-06-27 18:29:59 +00:00
user=Depends(get_verified_user),
2024-01-07 10:46:12 +00:00
):
2024-01-07 06:59:22 +00:00
try:
2024-09-28 00:23:09 +00:00
if app.state.config.ENABLE_RAG_HYBRID_SEARCH:
return query_collection_with_hybrid_search(
collection_names=form_data.collection_names,
query=form_data.query,
embedding_function=app.state.EMBEDDING_FUNCTION,
k=form_data.k if form_data.k else app.state.config.TOP_K,
reranking_function=app.state.sentence_transformer_rf,
r=(
form_data.r if form_data.r else app.state.config.RELEVANCE_THRESHOLD
),
2024-01-13 13:46:56 +00:00
)
else:
2024-09-28 00:23:09 +00:00
return query_collection(
collection_names=form_data.collection_names,
query=form_data.query,
embedding_function=app.state.EMBEDDING_FUNCTION,
k=form_data.k if form_data.k else app.state.config.TOP_K,
2024-07-15 11:05:38 +00:00
)
2024-06-18 20:50:18 +00:00
except Exception as e:
log.exception(e)
2024-03-24 07:40:27 +00:00
raise HTTPException(
2024-09-28 00:23:09 +00:00
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
2024-03-24 07:40:27 +00:00
)
2024-09-28 00:23:09 +00:00
####################################
#
# Vector DB operations
#
####################################
2024-02-18 05:06:08 +00:00
2024-10-03 13:44:17 +00:00
class DeleteForm(BaseModel):
collection_name: str
file_id: str
@app.post("/delete")
def delete_entries_from_collection(form_data: DeleteForm, user=Depends(get_admin_user)):
try:
if VECTOR_DB_CLIENT.has_collection(collection_name=form_data.collection_name):
file = Files.get_file_by_id(form_data.file_id)
hash = file.hash
VECTOR_DB_CLIENT.delete(
collection_name=form_data.collection_name,
metadata={"hash": hash},
)
return {"status": True}
else:
return {"status": False}
except Exception as e:
log.exception(e)
return {"status": False}
@app.post("/reset/db")
def reset_vector_db(user=Depends(get_admin_user)):
2024-09-10 01:27:50 +00:00
VECTOR_DB_CLIENT.reset()
2024-10-13 10:02:02 +00:00
Knowledges.delete_all_knowledge()
2024-01-07 09:40:36 +00:00
@app.post("/reset/uploads")
2024-06-04 04:45:36 +00:00
def reset_upload_dir(user=Depends(get_admin_user)) -> bool:
folder = f"{UPLOAD_DIR}"
try:
# Check if the directory exists
if os.path.exists(folder):
# Iterate over all the files and directories in the specified directory
for filename in os.listdir(folder):
file_path = os.path.join(folder, filename)
try:
if os.path.isfile(file_path) or os.path.islink(file_path):
os.unlink(file_path) # Remove the file or link
elif os.path.isdir(file_path):
shutil.rmtree(file_path) # Remove the directory
except Exception as e:
print(f"Failed to delete {file_path}. Reason: {e}")
else:
print(f"The directory {folder} does not exist")
except Exception as e:
print(f"Failed to process the directory {folder}. Reason: {e}")
return True
2024-05-19 13:51:32 +00:00
2024-06-12 07:18:22 +00:00
2024-05-19 13:51:32 +00:00
if ENV == "dev":
@app.get("/ef")
async def get_embeddings():
return {"result": app.state.EMBEDDING_FUNCTION("hello world")}
@app.get("/ef/{text}")
async def get_embeddings_text(text: str):
return {"result": app.state.EMBEDDING_FUNCTION(text)}