Timothy Jaeryang Baek 3 months ago
parent
commit
d542881ee4
1 changed files with 10 additions and 5 deletions
  1. 10 5
      backend/open_webui/routers/retrieval.py

+ 10 - 5
backend/open_webui/routers/retrieval.py

@@ -82,8 +82,8 @@ from open_webui.config import (
     RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
     RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
     UPLOAD_DIR,
     UPLOAD_DIR,
     DEFAULT_LOCALE,
     DEFAULT_LOCALE,
-    RAG_EMBEDDING_PASSAGE_PREFIX,
-    RAG_EMBEDDING_QUERY_PREFIX
+    RAG_EMBEDDING_CONTENT_PREFIX,
+    RAG_EMBEDDING_QUERY_PREFIX,
 )
 )
 from open_webui.env import (
 from open_webui.env import (
     SRC_LOG_LEVELS,
     SRC_LOG_LEVELS,
@@ -892,7 +892,9 @@ def save_docs_to_vector_db(
         )
         )
 
 
         embeddings = embedding_function(
         embeddings = embedding_function(
-            list(map(lambda x: x.replace("\n", " "), texts)), prefix=RAG_EMBEDDING_PASSAGE_PREFIX, user=user
+            list(map(lambda x: x.replace("\n", " "), texts)),
+            prefix=RAG_EMBEDDING_CONTENT_PREFIX,
+            user=user,
         )
         )
 
 
         items = [
         items = [
@@ -1536,7 +1538,6 @@ def query_doc_handler(
                 query_embedding=request.app.state.EMBEDDING_FUNCTION(
                 query_embedding=request.app.state.EMBEDDING_FUNCTION(
                     form_data.query, prefix=RAG_EMBEDDING_QUERY_PREFIX, user=user
                     form_data.query, prefix=RAG_EMBEDDING_QUERY_PREFIX, user=user
                 ),
                 ),
-
                 k=form_data.k if form_data.k else request.app.state.config.TOP_K,
                 k=form_data.k if form_data.k else request.app.state.config.TOP_K,
                 user=user,
                 user=user,
             )
             )
@@ -1663,7 +1664,11 @@ if ENV == "dev":
 
 
     @router.get("/ef/{text}")
     @router.get("/ef/{text}")
     async def get_embeddings(request: Request, text: Optional[str] = "Hello World!"):
     async def get_embeddings(request: Request, text: Optional[str] = "Hello World!"):
-        return {"result": request.app.state.EMBEDDING_FUNCTION(text, RAG_EMBEDDING_QUERY_PREFIX)}
+        return {
+            "result": request.app.state.EMBEDDING_FUNCTION(
+                text, RAG_EMBEDDING_QUERY_PREFIX
+            )
+        }
 
 
 
 
 class BatchProcessFilesForm(BaseModel):
 class BatchProcessFilesForm(BaseModel):