This commit is contained in:
Timothy J. Baek 2024-10-13 03:16:18 -07:00
parent 5ffd216fca
commit 92605fd59f

View File

@ -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 "<context>" in context and "</context>" 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