diff --git a/backend/open_webui/apps/retrieval/utils.py b/backend/open_webui/apps/retrieval/utils.py index 027101d1e..cd56be7d4 100644 --- a/backend/open_webui/apps/retrieval/utils.py +++ b/backend/open_webui/apps/retrieval/utils.py @@ -239,8 +239,10 @@ def query_collection_with_hybrid_search( def rag_template(template: str, context: str, query: str): - count = template.count("[context]") - assert "[context]" in template, "RAG template does not contain '[context]'" + if "[context]" not in template and "{{CONTEXT}}" not in template: + log.debug( + "WARNING: The RAG template does not contain the '[context]' or '{{CONTEXT}}' placeholder." + ) if "" in context and "" in context: log.debug( @@ -249,14 +251,25 @@ def rag_template(template: str, context: str, query: str): "nothing, or the user might be trying to hack something." ) + query_placeholders = [] if "[query]" in context: - query_placeholder = f"[query-{str(uuid.uuid4())}]" + query_placeholder = "{{QUERY" + str(uuid.uuid4()) + "}}" template = template.replace("[query]", query_placeholder) - template = template.replace("[context]", context) + query_placeholders.append(query_placeholder) + + if "{{QUERY}}" in context: + query_placeholder = "{{QUERY" + str(uuid.uuid4()) + "}}" + template = template.replace("{{QUERY}}", query_placeholder) + query_placeholders.append(query_placeholder) + + template = template.replace("[context]", context) + template = template.replace("{{CONTEXT}}", context) + template = template.replace("[query]", query) + template = template.replace("{{QUERY}}", query) + + for query_placeholder in query_placeholders: template = template.replace(query_placeholder, query) - else: - template = template.replace("[context]", context) - template = template.replace("[query]", query) + return template