Browse Source

refac: "rag" endpoints renamed to "retrieval"

Timothy J. Baek 7 months ago
parent
commit
e1103305f5
27 changed files with 41 additions and 34 deletions
  1. 2 1
      backend/open_webui/apps/retrieval/main.py
  2. 0 0
      backend/open_webui/apps/retrieval/search/brave.py
  3. 0 0
      backend/open_webui/apps/retrieval/search/duckduckgo.py
  4. 0 0
      backend/open_webui/apps/retrieval/search/google_pse.py
  5. 0 0
      backend/open_webui/apps/retrieval/search/jina_search.py
  6. 0 0
      backend/open_webui/apps/retrieval/search/main.py
  7. 0 0
      backend/open_webui/apps/retrieval/search/searchapi.py
  8. 0 0
      backend/open_webui/apps/retrieval/search/searxng.py
  9. 0 0
      backend/open_webui/apps/retrieval/search/serper.py
  10. 0 0
      backend/open_webui/apps/retrieval/search/serply.py
  11. 0 0
      backend/open_webui/apps/retrieval/search/serpstack.py
  12. 0 0
      backend/open_webui/apps/retrieval/search/tavily.py
  13. 0 0
      backend/open_webui/apps/retrieval/search/testdata/brave.json
  14. 0 0
      backend/open_webui/apps/retrieval/search/testdata/google_pse.json
  15. 0 0
      backend/open_webui/apps/retrieval/search/testdata/searchapi.json
  16. 0 0
      backend/open_webui/apps/retrieval/search/testdata/searxng.json
  17. 0 0
      backend/open_webui/apps/retrieval/search/testdata/serper.json
  18. 0 0
      backend/open_webui/apps/retrieval/search/testdata/serply.json
  19. 0 0
      backend/open_webui/apps/retrieval/search/testdata/serpstack.json
  20. 0 0
      backend/open_webui/apps/retrieval/utils.py
  21. 0 0
      backend/open_webui/apps/retrieval/vector/connector.py
  22. 0 0
      backend/open_webui/apps/retrieval/vector/dbs/chroma.py
  23. 0 0
      backend/open_webui/apps/retrieval/vector/dbs/milvus.py
  24. 0 0
      backend/open_webui/apps/retrieval/vector/main.py
  25. 37 28
      backend/open_webui/main.py
  26. 1 4
      src/lib/components/chat/MessageInput.svelte
  27. 1 1
      src/lib/constants.ts

+ 2 - 1
backend/open_webui/apps/rag/main.py → backend/open_webui/apps/retrieval/main.py

@@ -1061,7 +1061,7 @@ def store_data_in_vector_db(
 
     if len(docs) > 0:
         log.info(f"store_data_in_vector_db {docs}")
-        return store_docs_in_vector_db(docs, collection_name, metadata, overwrite), None
+        return store_docs_in_vector_db(docs, collection_name, metadata, overwrite)
     else:
         raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
 
@@ -1377,6 +1377,7 @@ def process_doc(
             )
 
             if result:
+
                 return {
                     "status": True,
                     "collection_name": collection_name,

+ 0 - 0
backend/open_webui/apps/rag/search/brave.py → backend/open_webui/apps/retrieval/search/brave.py


+ 0 - 0
backend/open_webui/apps/rag/search/duckduckgo.py → backend/open_webui/apps/retrieval/search/duckduckgo.py


+ 0 - 0
backend/open_webui/apps/rag/search/google_pse.py → backend/open_webui/apps/retrieval/search/google_pse.py


+ 0 - 0
backend/open_webui/apps/rag/search/jina_search.py → backend/open_webui/apps/retrieval/search/jina_search.py


+ 0 - 0
backend/open_webui/apps/rag/search/main.py → backend/open_webui/apps/retrieval/search/main.py


+ 0 - 0
backend/open_webui/apps/rag/search/searchapi.py → backend/open_webui/apps/retrieval/search/searchapi.py


+ 0 - 0
backend/open_webui/apps/rag/search/searxng.py → backend/open_webui/apps/retrieval/search/searxng.py


+ 0 - 0
backend/open_webui/apps/rag/search/serper.py → backend/open_webui/apps/retrieval/search/serper.py


+ 0 - 0
backend/open_webui/apps/rag/search/serply.py → backend/open_webui/apps/retrieval/search/serply.py


+ 0 - 0
backend/open_webui/apps/rag/search/serpstack.py → backend/open_webui/apps/retrieval/search/serpstack.py


+ 0 - 0
backend/open_webui/apps/rag/search/tavily.py → backend/open_webui/apps/retrieval/search/tavily.py


+ 0 - 0
backend/open_webui/apps/rag/search/testdata/brave.json → backend/open_webui/apps/retrieval/search/testdata/brave.json


+ 0 - 0
backend/open_webui/apps/rag/search/testdata/google_pse.json → backend/open_webui/apps/retrieval/search/testdata/google_pse.json


+ 0 - 0
backend/open_webui/apps/rag/search/testdata/searchapi.json → backend/open_webui/apps/retrieval/search/testdata/searchapi.json


+ 0 - 0
backend/open_webui/apps/rag/search/testdata/searxng.json → backend/open_webui/apps/retrieval/search/testdata/searxng.json


+ 0 - 0
backend/open_webui/apps/rag/search/testdata/serper.json → backend/open_webui/apps/retrieval/search/testdata/serper.json


+ 0 - 0
backend/open_webui/apps/rag/search/testdata/serply.json → backend/open_webui/apps/retrieval/search/testdata/serply.json


+ 0 - 0
backend/open_webui/apps/rag/search/testdata/serpstack.json → backend/open_webui/apps/retrieval/search/testdata/serpstack.json


+ 0 - 0
backend/open_webui/apps/rag/utils.py → backend/open_webui/apps/retrieval/utils.py


+ 0 - 0
backend/open_webui/apps/rag/vector/connector.py → backend/open_webui/apps/retrieval/vector/connector.py


+ 0 - 0
backend/open_webui/apps/rag/vector/dbs/chroma.py → backend/open_webui/apps/retrieval/vector/dbs/chroma.py


+ 0 - 0
backend/open_webui/apps/rag/vector/dbs/milvus.py → backend/open_webui/apps/retrieval/vector/dbs/milvus.py


+ 0 - 0
backend/open_webui/apps/rag/vector/main.py → backend/open_webui/apps/retrieval/vector/main.py


+ 37 - 28
backend/open_webui/main.py

@@ -16,37 +16,45 @@ from typing import Optional
 import aiohttp
 import requests
 
-
-from open_webui.apps.audio.main import app as audio_app
-from open_webui.apps.images.main import app as images_app
-from open_webui.apps.ollama.main import app as ollama_app
 from open_webui.apps.ollama.main import (
-    GenerateChatCompletionForm,
+    app as ollama_app,
+    get_all_models as get_ollama_models,
     generate_chat_completion as generate_ollama_chat_completion,
     generate_openai_chat_completion as generate_ollama_openai_chat_completion,
+    GenerateChatCompletionForm,
 )
-from open_webui.apps.ollama.main import get_all_models as get_ollama_models
-from open_webui.apps.openai.main import app as openai_app
 from open_webui.apps.openai.main import (
+    app as openai_app,
     generate_chat_completion as generate_openai_chat_completion,
+    get_all_models as get_openai_models,
 )
-from open_webui.apps.openai.main import get_all_models as get_openai_models
-from open_webui.apps.rag.main import app as rag_app
-from open_webui.apps.rag.utils import get_rag_context, rag_template
-from open_webui.apps.socket.main import app as socket_app, periodic_usage_pool_cleanup
-from open_webui.apps.socket.main import get_event_call, get_event_emitter
-from open_webui.apps.webui.internal.db import Session
-from open_webui.apps.webui.main import app as webui_app
+
+from open_webui.apps.retrieval.main import app as retrieval_app
+from open_webui.apps.retrieval.utils import get_rag_context, rag_template
+
+from open_webui.apps.socket.main import (
+    app as socket_app,
+    periodic_usage_pool_cleanup,
+    get_event_call,
+    get_event_emitter,
+)
+
 from open_webui.apps.webui.main import (
+    app as webui_app,
     generate_function_chat_completion,
     get_pipe_models,
 )
+from open_webui.apps.webui.internal.db import Session
+
 from open_webui.apps.webui.models.auths import Auths
 from open_webui.apps.webui.models.functions import Functions
 from open_webui.apps.webui.models.models import Models
 from open_webui.apps.webui.models.users import UserModel, Users
+
 from open_webui.apps.webui.utils import load_function_module_by_id
 
+from open_webui.apps.audio.main import app as audio_app
+from open_webui.apps.images.main import app as images_app
 
 from authlib.integrations.starlette_client import OAuth
 from authlib.oidc.core import UserInfo
@@ -491,11 +499,11 @@ async def chat_completion_files_handler(body) -> tuple[dict, dict[str, list]]:
         contexts, citations = get_rag_context(
             files=files,
             messages=body["messages"],
-            embedding_function=rag_app.state.EMBEDDING_FUNCTION,
-            k=rag_app.state.config.TOP_K,
-            reranking_function=rag_app.state.sentence_transformer_rf,
-            r=rag_app.state.config.RELEVANCE_THRESHOLD,
-            hybrid_search=rag_app.state.config.ENABLE_RAG_HYBRID_SEARCH,
+            embedding_function=retrieval_app.state.EMBEDDING_FUNCTION,
+            k=retrieval_app.state.config.TOP_K,
+            reranking_function=retrieval_app.state.sentence_transformer_rf,
+            r=retrieval_app.state.config.RELEVANCE_THRESHOLD,
+            hybrid_search=retrieval_app.state.config.ENABLE_RAG_HYBRID_SEARCH,
         )
 
         log.debug(f"rag_contexts: {contexts}, citations: {citations}")
@@ -608,7 +616,7 @@ class ChatCompletionMiddleware(BaseHTTPMiddleware):
             if prompt is None:
                 raise Exception("No user message found")
             if (
-                rag_app.state.config.RELEVANCE_THRESHOLD == 0
+                retrieval_app.state.config.RELEVANCE_THRESHOLD == 0
                 and context_string.strip() == ""
             ):
                 log.debug(
@@ -620,14 +628,14 @@ class ChatCompletionMiddleware(BaseHTTPMiddleware):
             if model["owned_by"] == "ollama":
                 body["messages"] = prepend_to_first_user_message_content(
                     rag_template(
-                        rag_app.state.config.RAG_TEMPLATE, context_string, prompt
+                        retrieval_app.state.config.RAG_TEMPLATE, context_string, prompt
                     ),
                     body["messages"],
                 )
             else:
                 body["messages"] = add_or_update_system_message(
                     rag_template(
-                        rag_app.state.config.RAG_TEMPLATE, context_string, prompt
+                        retrieval_app.state.config.RAG_TEMPLATE, context_string, prompt
                     ),
                     body["messages"],
                 )
@@ -849,7 +857,7 @@ async def check_url(request: Request, call_next):
 async def update_embedding_function(request: Request, call_next):
     response = await call_next(request)
     if "/embedding/update" in request.url.path:
-        webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION
+        webui_app.state.EMBEDDING_FUNCTION = retrieval_app.state.EMBEDDING_FUNCTION
     return response
 
 
@@ -877,11 +885,12 @@ app.mount("/openai", openai_app)
 
 app.mount("/images/api/v1", images_app)
 app.mount("/audio/api/v1", audio_app)
-app.mount("/rag/api/v1", rag_app)
+app.mount("/retrieval/api/v1", retrieval_app)
 
 app.mount("/api/v1", webui_app)
 
-webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION
+
+webui_app.state.EMBEDDING_FUNCTION = retrieval_app.state.EMBEDDING_FUNCTION
 
 
 async def get_all_models():
@@ -2066,7 +2075,7 @@ async def get_app_config(request: Request):
             "enable_login_form": webui_app.state.config.ENABLE_LOGIN_FORM,
             **(
                 {
-                    "enable_web_search": rag_app.state.config.ENABLE_RAG_WEB_SEARCH,
+                    "enable_web_search": retrieval_app.state.config.ENABLE_RAG_WEB_SEARCH,
                     "enable_image_generation": images_app.state.config.ENABLED,
                     "enable_community_sharing": webui_app.state.config.ENABLE_COMMUNITY_SHARING,
                     "enable_message_rating": webui_app.state.config.ENABLE_MESSAGE_RATING,
@@ -2092,8 +2101,8 @@ async def get_app_config(request: Request):
                     },
                 },
                 "file": {
-                    "max_size": rag_app.state.config.FILE_MAX_SIZE,
-                    "max_count": rag_app.state.config.FILE_MAX_COUNT,
+                    "max_size": retrieval_app.state.config.FILE_MAX_SIZE,
+                    "max_count": retrieval_app.state.config.FILE_MAX_COUNT,
                 },
                 "permissions": {**webui_app.state.config.USER_PERMISSIONS},
             }

+ 1 - 4
src/lib/components/chat/MessageInput.svelte

@@ -159,16 +159,13 @@
 	const processFileItem = async (fileItem) => {
 		try {
 			const res = await processDocToVectorDB(localStorage.token, fileItem.id);
-
 			if (res) {
 				fileItem.status = 'processed';
 				fileItem.collection_name = res.collection_name;
 				files = files;
 			}
 		} catch (e) {
-			// Remove the failed doc from the files array
-			// files = files.filter((f) => f.id !== fileItem.id);
-			toast.error(e);
+			// We keep the file in the files list even if it fails to process
 			fileItem.status = 'processed';
 			files = files;
 		}

+ 1 - 1
src/lib/constants.ts

@@ -11,7 +11,7 @@ export const OLLAMA_API_BASE_URL = `${WEBUI_BASE_URL}/ollama`;
 export const OPENAI_API_BASE_URL = `${WEBUI_BASE_URL}/openai`;
 export const AUDIO_API_BASE_URL = `${WEBUI_BASE_URL}/audio/api/v1`;
 export const IMAGES_API_BASE_URL = `${WEBUI_BASE_URL}/images/api/v1`;
-export const RAG_API_BASE_URL = `${WEBUI_BASE_URL}/rag/api/v1`;
+export const RAG_API_BASE_URL = `${WEBUI_BASE_URL}/retrieval/api/v1`;
 
 export const WEBUI_VERSION = APP_VERSION;
 export const WEBUI_BUILD_HASH = APP_BUILD_HASH;