From 6d2f87e9044800320656c98a501302f2f6a3f56a Mon Sep 17 00:00:00 2001 From: jayteaftw Date: Wed, 5 Feb 2025 14:03:16 -0800 Subject: [PATCH] Added server side Prefixing --- backend/open_webui/config.py | 2 +- backend/open_webui/retrieval/utils.py | 25 +++++++++++++++++++++++-- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index f1b1c14a5..5635b70a6 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -1339,7 +1339,7 @@ RAG_EMBEDDING_PASSAGE_PREFIX = ( ) RAG_EMBEDDING_PREFIX_FIELD_NAME = ( - os.environ.get("RAG_EMBEDDING_PREFIX_FIELD_NAME", "input_type") + os.environ.get("RAG_EMBEDDING_PREFIX_FIELD_NAME", None) ) RAG_RERANKING_MODEL = PersistentConfig( diff --git a/backend/open_webui/retrieval/utils.py b/backend/open_webui/retrieval/utils.py index 544a65a89..7a9be9ea9 100644 --- a/backend/open_webui/retrieval/utils.py +++ b/backend/open_webui/retrieval/utils.py @@ -418,14 +418,22 @@ def get_model_path(model: str, update_model: bool = False): def generate_openai_batch_embeddings( model: str, texts: list[str], url: str = "https://api.openai.com/v1", key: str = "", prefix: str = None ) -> Optional[list[list[float]]]: + try: + json_data = { + "input": texts, + "model": model + } + if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME,str) and isinstance(prefix,str): + json_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix + r = requests.post( f"{url}/embeddings", headers={ "Content-Type": "application/json", "Authorization": f"Bearer {key}", }, - json={"input": texts, "model": model} if not prefix else {"input": texts, "model": model, RAG_EMBEDDING_PREFIX_FIELD_NAME: prefix}, + json=json_data, ) r.raise_for_status() data = r.json() @@ -442,13 +450,20 @@ def generate_ollama_batch_embeddings( model: str, texts: list[str], url: str, key: str = "", prefix: str = None ) -> Optional[list[list[float]]]: try: + json_data = { + "input": texts, + "model": model + } + if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME,str) and isinstance(prefix,str): + json_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix + r = requests.post( f"{url}/api/embed", headers={ "Content-Type": "application/json", "Authorization": f"Bearer {key}", }, - json={"input": texts, "model": model} if not prefix else {"input": texts, "model": model, RAG_EMBEDDING_PREFIX_FIELD_NAME: prefix}, + json=json_data, ) r.raise_for_status() data = r.json() @@ -466,6 +481,12 @@ def generate_embeddings(engine: str, model: str, text: Union[str, list[str]], pr url = kwargs.get("url", "") key = kwargs.get("key", "") + if prefix is not None and RAG_EMBEDDING_PREFIX_FIELD_NAME is None: + if isinstance(text, list): + text = [f'{prefix}{text_element}' for text_element in text] + else: + text = f'{prefix}{text}' + if engine == "ollama": if isinstance(text, list): embeddings = generate_ollama_batch_embeddings(