tdpeuter/open-webui
tdpeuter/open-webui
1
Fork 0
forked from open-webui/open-webui
Code Issues Pull requests 1 Releases 2 Packages Activity Actions

feat: toggle hybrid search

Browse source
This commit is contained in:
Steven Kreitzer 2024-04-25 17:31:21 -05:00
parent 984dbf13ab
commit 9755cd5baa
4 changed files with 133 additions and 88 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
backend/apps/rag/main.py
View file

@ -70,6 +70,7 @@ from config import (
RAG_EMBEDDING_MODEL, RAG_EMBEDDING_MODEL,
RAG_EMBEDDING_MODEL_AUTO_UPDATE, RAG_EMBEDDING_MODEL_AUTO_UPDATE,
RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE, RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
RAG_HYBRID,
RAG_RERANKING_MODEL, RAG_RERANKING_MODEL,
RAG_RERANKING_MODEL_AUTO_UPDATE, RAG_RERANKING_MODEL_AUTO_UPDATE,
RAG_RERANKING_MODEL_TRUST_REMOTE_CODE, RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
@ -91,6 +92,8 @@ app = FastAPI()
app.state.TOP_K = RAG_TOP_K app.state.TOP_K = RAG_TOP_K
app.state.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD app.state.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD
app.state.HYBRID = RAG_HYBRID
app.state.CHUNK_SIZE = CHUNK_SIZE app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
@ -321,6 +324,7 @@ async def get_query_settings(user=Depends(get_admin_user)):
"template": app.state.RAG_TEMPLATE, "template": app.state.RAG_TEMPLATE,
"k": app.state.TOP_K, "k": app.state.TOP_K,
"r": app.state.RELEVANCE_THRESHOLD, "r": app.state.RELEVANCE_THRESHOLD,
"hybrid": app.state.HYBRID,
} }
@ -328,6 +332,7 @@ class QuerySettingsForm(BaseModel):
k: Optional[int] = None k: Optional[int] = None
r: Optional[float] = None r: Optional[float] = None
template: Optional[str] = None template: Optional[str] = None
hybrid: Optional[bool] = None
@app.post("/query/settings/update") @app.post("/query/settings/update")
@ -337,7 +342,14 @@ async def update_query_settings(
app.state.RAG_TEMPLATE = form_data.template if form_data.template else RAG_TEMPLATE app.state.RAG_TEMPLATE = form_data.template if form_data.template else RAG_TEMPLATE
app.state.TOP_K = form_data.k if form_data.k else 4 app.state.TOP_K = form_data.k if form_data.k else 4
app.state.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0 app.state.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
return {"status": True, "template": app.state.RAG_TEMPLATE} app.state.HYBRID = form_data.hybrid if form_data.hybrid else False
return {
"status": True,
"template": app.state.RAG_TEMPLATE,
"k": app.state.TOP_K,
"r": app.state.RELEVANCE_THRESHOLD,
"hybrid": app.state.HYBRID,
}
class QueryDocForm(BaseModel): class QueryDocForm(BaseModel):
@ -345,6 +357,7 @@ class QueryDocForm(BaseModel):
query: str query: str
k: Optional[int] = None k: Optional[int] = None
r: Optional[float] = None r: Optional[float] = None
hybrid: Optional[bool] = None
@app.post("/query/doc") @app.post("/query/doc")
@ -368,6 +381,7 @@ def query_doc_handler(
r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD, r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
embeddings_function=embeddings_function, embeddings_function=embeddings_function,
reranking_function=app.state.sentence_transformer_rf, reranking_function=app.state.sentence_transformer_rf,
hybrid=form_data.hybrid if form_data.hybrid else app.state.HYBRID,
) )
except Exception as e: except Exception as e:
log.exception(e) log.exception(e)
@ -382,6 +396,7 @@ class QueryCollectionsForm(BaseModel):
query: str query: str
k: Optional[int] = None k: Optional[int] = None
r: Optional[float] = None r: Optional[float] = None
hybrid: Optional[bool] = None
@app.post("/query/collection") @app.post("/query/collection")
@ -405,6 +420,7 @@ def query_collection_handler(
r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD, r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
embeddings_function=embeddings_function, embeddings_function=embeddings_function,
reranking_function=app.state.sentence_transformer_rf, reranking_function=app.state.sentence_transformer_rf,
hybrid=form_data.hybrid if form_data.hybrid else app.state.HYBRID,
) )
except Exception as e: except Exception as e:
log.exception(e) log.exception(e)

11
backend/apps/rag/utils.py
View file

@ -32,13 +32,13 @@ def query_embeddings_doc(
collection_name: str, collection_name: str,
query: str, query: str,
embeddings_function, embeddings_function,
reranking_function,
k: int, k: int,
reranking_function: Optional[CrossEncoder] = None,
r: Optional[float] = None, r: Optional[float] = None,
hybrid: Optional[bool] = False,
): ):
try: try:
if hybrid:
if reranking_function:
# if you use docker use the model from the environment variable # if you use docker use the model from the environment variable
collection = CHROMA_CLIENT.get_collection(name=collection_name) collection = CHROMA_CLIENT.get_collection(name=collection_name)
@ -142,6 +142,7 @@ def query_embeddings_collection(
r: float, r: float,
embeddings_function, embeddings_function,
reranking_function, reranking_function,
hybrid: bool,
): ):
results = [] results = []
@ -155,6 +156,7 @@ def query_embeddings_collection(
r=r, r=r,
embeddings_function=embeddings_function, embeddings_function=embeddings_function,
reranking_function=reranking_function, reranking_function=reranking_function,
hybrid=hybrid,
) )
results.append(result) results.append(result)
except: except:
@ -211,6 +213,7 @@ def rag_messages(
template, template,
k, k,
r, r,
hybrid,
embedding_engine, embedding_engine,
embedding_model, embedding_model,
embedding_function, embedding_function,
@ -283,6 +286,7 @@ def rag_messages(
r=r, r=r,
embeddings_function=embeddings_function, embeddings_function=embeddings_function,
reranking_function=reranking_function, reranking_function=reranking_function,
hybrid=hybrid,
) )
else: else:
context = query_embeddings_doc( context = query_embeddings_doc(
@ -292,6 +296,7 @@ def rag_messages(
r=r, r=r,
embeddings_function=embeddings_function, embeddings_function=embeddings_function,
reranking_function=reranking_function, reranking_function=reranking_function,
hybrid=hybrid,
) )
except Exception as e: except Exception as e:
log.exception(e) log.exception(e)

1
backend/config.py
View file

@ -422,6 +422,7 @@ CHROMA_DATA_PATH = f"{DATA_DIR}/vector_db"
RAG_TOP_K = int(os.environ.get("RAG_TOP_K", "5")) RAG_TOP_K = int(os.environ.get("RAG_TOP_K", "5"))
RAG_RELEVANCE_THRESHOLD = float(os.environ.get("RAG_RELEVANCE_THRESHOLD", "0.0")) RAG_RELEVANCE_THRESHOLD = float(os.environ.get("RAG_RELEVANCE_THRESHOLD", "0.0"))
RAG_HYBRID = os.environ.get("RAG_HYBRID", "").lower() == "true"
RAG_EMBEDDING_ENGINE = os.environ.get("RAG_EMBEDDING_ENGINE", "") RAG_EMBEDDING_ENGINE = os.environ.get("RAG_EMBEDDING_ENGINE", "")

191
src/lib/components/documents/Settings/General.svelte
View file

@ -43,7 +43,8 @@
let querySettings = { let querySettings = {
template: '', template: '',
r: 0.0, r: 0.0,
k: 4 k: 4,
hybrid: false
}; };
const scanHandler = async () => { const scanHandler = async () => {
@ -174,6 +175,12 @@
} }
}; };
const toggleHybridSearch = async () => {
querySettings.hybrid = !querySettings.hybrid;
querySettings = await updateQuerySettings(localStorage.token, querySettings);
};
onMount(async () => { onMount(async () => {
const res = await getRAGConfig(localStorage.token); const res = await getRAGConfig(localStorage.token);
@ -202,6 +209,24 @@
<div> <div>
<div class=" mb-2 text-sm font-medium">{$i18n.t('General Settings')}</div> <div class=" mb-2 text-sm font-medium">{$i18n.t('General Settings')}</div>
<div class=" flex w-full justify-between">
<div class=" self-center text-xs font-medium">{$i18n.t('Hybrid Search')}</div>
<button
class="p-1 px-3 text-xs flex rounded transition"
on:click={() => {
toggleHybridSearch();
}}
type="button"
>
{#if querySettings.hybrid === true}
<span class="ml-2 self-center">{$i18n.t('On')}</span>
{:else}
<span class="ml-2 self-center">{$i18n.t('Off')}</span>
{/if}
</button>
</div>
<div class=" flex w-full justify-between"> <div class=" flex w-full justify-between">
<div class=" self-center text-xs font-medium">{$i18n.t('Embedding Model Engine')}</div> <div class=" self-center text-xs font-medium">{$i18n.t('Embedding Model Engine')}</div>
<div class="flex items-center relative"> <div class="flex items-center relative">
@ -386,78 +411,74 @@
<hr class=" dark:border-gray-700 my-3" /> <hr class=" dark:border-gray-700 my-3" />
<div class=" "> {#if querySettings.hybrid === true}
<div class=" mb-2 text-sm font-medium">{$i18n.t('Update Reranking Model')}</div> <div class=" ">
<div class=" mb-2 text-sm font-medium">{$i18n.t('Update Reranking Model')}</div>
<div class="flex w-full"> <div class="flex w-full">
<div class="flex-1 mr-2"> <div class="flex-1 mr-2">
<input <input
class="w-full rounded-lg py-2 px-4 text-sm dark:text-gray-300 dark:bg-gray-850 outline-none" class="w-full rounded-lg py-2 px-4 text-sm dark:text-gray-300 dark:bg-gray-850 outline-none"
placeholder={$i18n.t('Update reranking model (e.g. {{model}})', { placeholder={$i18n.t('Update reranking model (e.g. {{model}})', {
model: rerankingModel.slice(-40) model: rerankingModel.slice(-40)
})} })}
bind:value={rerankingModel} bind:value={rerankingModel}
/> />
</div> </div>
<button <button
class="px-2.5 bg-gray-100 hover:bg-gray-200 text-gray-800 dark:bg-gray-850 dark:hover:bg-gray-800 dark:text-gray-100 rounded-lg transition" class="px-2.5 bg-gray-100 hover:bg-gray-200 text-gray-800 dark:bg-gray-850 dark:hover:bg-gray-800 dark:text-gray-100 rounded-lg transition"
on:click={() => { on:click={() => {
rerankingModelUpdateHandler(); rerankingModelUpdateHandler();
}} }}
disabled={updateRerankingModelLoading} disabled={updateRerankingModelLoading}
> >
{#if updateRerankingModelLoading} {#if updateRerankingModelLoading}
<div class="self-center"> <div class="self-center">
<svg <svg
class=" w-4 h-4" class=" w-4 h-4"
viewBox="0 0 24 24" viewBox="0 0 24 24"
fill="currentColor" fill="currentColor"
xmlns="http://www.w3.org/2000/svg" xmlns="http://www.w3.org/2000/svg"
><style> ><style>
.spinner_ajPY { .spinner_ajPY {
transform-origin: center; transform-origin: center;
animation: spinner_AtaB 0.75s infinite linear; animation: spinner_AtaB 0.75s infinite linear;
}
@keyframes spinner_AtaB {
100% {
transform: rotate(360deg);
} }
} @keyframes spinner_AtaB {
</style><path 100% {
d="M12,1A11,11,0,1,0,23,12,11,11,0,0,0,12,1Zm0,19a8,8,0,1,1,8-8A8,8,0,0,1,12,20Z" transform: rotate(360deg);
opacity=".25" }
/><path }
d="M10.14,1.16a11,11,0,0,0-9,8.92A1.59,1.59,0,0,0,2.46,12,1.52,1.52,0,0,0,4.11,10.7a8,8,0,0,1,6.66-6.61A1.42,1.42,0,0,0,12,2.69h0A1.57,1.57,0,0,0,10.14,1.16Z" </style><path
class="spinner_ajPY" d="M12,1A11,11,0,1,0,23,12,11,11,0,0,0,12,1Zm0,19a8,8,0,1,1,8-8A8,8,0,0,1,12,20Z"
/></svg opacity=".25"
/><path
d="M10.14,1.16a11,11,0,0,0-9,8.92A1.59,1.59,0,0,0,2.46,12,1.52,1.52,0,0,0,4.11,10.7a8,8,0,0,1,6.66-6.61A1.42,1.42,0,0,0,12,2.69h0A1.57,1.57,0,0,0,10.14,1.16Z"
class="spinner_ajPY"
/></svg
>
</div>
{:else}
<svg
xmlns="http://www.w3.org/2000/svg"
viewBox="0 0 16 16"
fill="currentColor"
class="w-4 h-4"
> >
</div> <path
{:else} d="M8.75 2.75a.75.75 0 0 0-1.5 0v5.69L5.03 6.22a.75.75 0 0 0-1.06 1.06l3.5 3.5a.75.75 0 0 0 1.06 0l3.5-3.5a.75.75 0 0 0-1.06-1.06L8.75 8.44V2.75Z"
<svg />
xmlns="http://www.w3.org/2000/svg" <path
viewBox="0 0 16 16" d="M3.5 9.75a.75.75 0 0 0-1.5 0v1.5A2.75 2.75 0 0 0 4.75 14h6.5A2.75 2.75 0 0 0 14 11.25v-1.5a.75.75 0 0 0-1.5 0v1.5c0 .69-.56 1.25-1.25 1.25h-6.5c-.69 0-1.25-.56-1.25-1.25v-1.5Z"
fill="currentColor" />
class="w-4 h-4" </svg>
> {/if}
<path </button>
d="M8.75 2.75a.75.75 0 0 0-1.5 0v5.69L5.03 6.22a.75.75 0 0 0-1.06 1.06l3.5 3.5a.75.75 0 0 0 1.06 0l3.5-3.5a.75.75 0 0 0-1.06-1.06L8.75 8.44V2.75Z" </div>
/>
<path
d="M3.5 9.75a.75.75 0 0 0-1.5 0v1.5A2.75 2.75 0 0 0 4.75 14h6.5A2.75 2.75 0 0 0 14 11.25v-1.5a.75.75 0 0 0-1.5 0v1.5c0 .69-.56 1.25-1.25 1.25h-6.5c-.69 0-1.25-.56-1.25-1.25v-1.5Z"
/>
</svg>
{/if}
</button>
</div> </div>
</div>
<div class="mt-2 mb-1 text-xs text-gray-400 dark:text-gray-500"> <hr class=" dark:border-gray-700 my-3" />
{$i18n.t( {/if}
'Note: If you choose a reranking model, it will use that to score and rerank instead of the embedding model.'
)}
</div>
<hr class=" dark:border-gray-700 my-3" />
<div class=" flex w-full justify-between"> <div class=" flex w-full justify-between">
<div class=" self-center text-xs font-medium"> <div class=" self-center text-xs font-medium">
@ -583,25 +604,27 @@
</div> </div>
</div> </div>
<div class=" flex"> {#if querySettings.hybrid === true}
<div class=" flex w-full justify-between"> <div class=" flex">
<div class="self-center text-xs font-medium flex-1"> <div class=" flex w-full justify-between">
{$i18n.t('Relevance Threshold')} <div class="self-center text-xs font-medium flex-1">
</div> {$i18n.t('Relevance Threshold')}
</div>
<div class="self-center p-3"> <div class="self-center p-3">
<input <input
class=" w-full rounded-lg py-1.5 px-4 text-sm dark:text-gray-300 dark:bg-gray-850 outline-none" class=" w-full rounded-lg py-1.5 px-4 text-sm dark:text-gray-300 dark:bg-gray-850 outline-none"
type="number" type="number"
step="0.01" step="0.01"
placeholder={$i18n.t('Enter Relevance Threshold')} placeholder={$i18n.t('Enter Relevance Threshold')}
bind:value={querySettings.r} bind:value={querySettings.r}
autocomplete="off" autocomplete="off"
min="0.0" min="0.0"
/> />
</div>
</div> </div>
</div> </div>
</div> {/if}
<div> <div>
<div class=" mb-2.5 text-sm font-medium">{$i18n.t('RAG Template')}</div> <div class=" mb-2.5 text-sm font-medium">{$i18n.t('RAG Template')}</div>