From 9661fee55428b81c3f65bc2863efbcd70c7e403a Mon Sep 17 00:00:00 2001
From: thiswillbeyourgithub
 <26625900+thiswillbeyourgithub@users.noreply.github.com>
Date: Thu, 12 Sep 2024 15:19:40 +0200
Subject: [PATCH] fix: handle case where [query] happens in the RAG context

---
 backend/open_webui/apps/rag/utils.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/backend/open_webui/apps/rag/utils.py b/backend/open_webui/apps/rag/utils.py
index 5bddf0a96..41fc4e2d4 100644
--- a/backend/open_webui/apps/rag/utils.py
+++ b/backend/open_webui/apps/rag/utils.py
@@ -1,5 +1,6 @@
 import logging
 import os
+import uuid
 from typing import Optional, Union
 
 import requests
@@ -197,8 +198,15 @@ def rag_template(template: str, context: str, query: str):
         f"RAG template contains an unexpected number of '[context]' : {count}"
     )
     assert "[context]" in template, "RAG template does not contain '[context]'"
-    template = template.replace("[context]", context)
-    template = template.replace("[query]", query)
+
+    if "[query]" in context:
+        query_placeholder = str(uuid.uuid4())
+        template = template.replace("[QUERY]", query_placeholder)
+        template = template.replace("[context]", context)
+        template = template.replace(query_placeholder, query)
+    else:
+        template = template.replace("[context]", context)
+        template = template.replace("[query]", query)
     return template