360 lines
11 KiB
Python
360 lines
11 KiB
Python
import logging
|
|
import math
|
|
import re
|
|
from datetime import datetime
|
|
from typing import Optional
|
|
import uuid
|
|
|
|
|
|
from open_webui.utils.misc import get_last_user_message, get_messages_content
|
|
|
|
from open_webui.env import SRC_LOG_LEVELS
|
|
from open_webui.config import DEFAULT_RAG_TEMPLATE
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
|
|
|
|
|
def get_task_model_id(
|
|
default_model_id: str, task_model: str, task_model_external: str, models
|
|
) -> str:
|
|
# Set the task model
|
|
task_model_id = default_model_id
|
|
# Check if the user has a custom task model and use that model
|
|
if models[task_model_id].get("connection_type") == "local":
|
|
if task_model and task_model in models:
|
|
task_model_id = task_model
|
|
else:
|
|
if task_model_external and task_model_external in models:
|
|
task_model_id = task_model_external
|
|
|
|
return task_model_id
|
|
|
|
|
|
def prompt_variables_template(template: str, variables: dict[str, str]) -> str:
|
|
for variable, value in variables.items():
|
|
template = template.replace(variable, value)
|
|
return template
|
|
|
|
|
|
def prompt_template(
|
|
template: str, user_name: Optional[str] = None, user_location: Optional[str] = None
|
|
) -> str:
|
|
# Get the current date
|
|
current_date = datetime.now()
|
|
|
|
# Format the date to YYYY-MM-DD
|
|
formatted_date = current_date.strftime("%Y-%m-%d")
|
|
formatted_time = current_date.strftime("%I:%M:%S %p")
|
|
formatted_weekday = current_date.strftime("%A")
|
|
|
|
template = template.replace("{{CURRENT_DATE}}", formatted_date)
|
|
template = template.replace("{{CURRENT_TIME}}", formatted_time)
|
|
template = template.replace(
|
|
"{{CURRENT_DATETIME}}", f"{formatted_date} {formatted_time}"
|
|
)
|
|
template = template.replace("{{CURRENT_WEEKDAY}}", formatted_weekday)
|
|
|
|
if user_name:
|
|
# Replace {{USER_NAME}} in the template with the user's name
|
|
template = template.replace("{{USER_NAME}}", user_name)
|
|
else:
|
|
# Replace {{USER_NAME}} in the template with "Unknown"
|
|
template = template.replace("{{USER_NAME}}", "Unknown")
|
|
|
|
if user_location:
|
|
# Replace {{USER_LOCATION}} in the template with the current location
|
|
template = template.replace("{{USER_LOCATION}}", user_location)
|
|
else:
|
|
# Replace {{USER_LOCATION}} in the template with "Unknown"
|
|
template = template.replace("{{USER_LOCATION}}", "Unknown")
|
|
|
|
return template
|
|
|
|
|
|
def replace_prompt_variable(template: str, prompt: str) -> str:
|
|
def replacement_function(match):
|
|
full_match = match.group(
|
|
0
|
|
).lower() # Normalize to lowercase for consistent handling
|
|
start_length = match.group(1)
|
|
end_length = match.group(2)
|
|
middle_length = match.group(3)
|
|
|
|
if full_match == "{{prompt}}":
|
|
return prompt
|
|
elif start_length is not None:
|
|
return prompt[: int(start_length)]
|
|
elif end_length is not None:
|
|
return prompt[-int(end_length) :]
|
|
elif middle_length is not None:
|
|
middle_length = int(middle_length)
|
|
if len(prompt) <= middle_length:
|
|
return prompt
|
|
start = prompt[: math.ceil(middle_length / 2)]
|
|
end = prompt[-math.floor(middle_length / 2) :]
|
|
return f"{start}...{end}"
|
|
return ""
|
|
|
|
# Updated regex pattern to make it case-insensitive with the `(?i)` flag
|
|
pattern = r"(?i){{prompt}}|{{prompt:start:(\d+)}}|{{prompt:end:(\d+)}}|{{prompt:middletruncate:(\d+)}}"
|
|
template = re.sub(pattern, replacement_function, template)
|
|
return template
|
|
|
|
|
|
def replace_messages_variable(
|
|
template: str, messages: Optional[list[dict]] = None
|
|
) -> str:
|
|
def replacement_function(match):
|
|
full_match = match.group(0)
|
|
start_length = match.group(1)
|
|
end_length = match.group(2)
|
|
middle_length = match.group(3)
|
|
# If messages is None, handle it as an empty list
|
|
if messages is None:
|
|
return ""
|
|
|
|
# Process messages based on the number of messages required
|
|
if full_match == "{{MESSAGES}}":
|
|
return get_messages_content(messages)
|
|
elif start_length is not None:
|
|
return get_messages_content(messages[: int(start_length)])
|
|
elif end_length is not None:
|
|
return get_messages_content(messages[-int(end_length) :])
|
|
elif middle_length is not None:
|
|
mid = int(middle_length)
|
|
|
|
if len(messages) <= mid:
|
|
return get_messages_content(messages)
|
|
# Handle middle truncation: split to get start and end portions of the messages list
|
|
half = mid // 2
|
|
start_msgs = messages[:half]
|
|
end_msgs = messages[-half:] if mid % 2 == 0 else messages[-(half + 1) :]
|
|
formatted_start = get_messages_content(start_msgs)
|
|
formatted_end = get_messages_content(end_msgs)
|
|
return f"{formatted_start}\n{formatted_end}"
|
|
return ""
|
|
|
|
template = re.sub(
|
|
r"{{MESSAGES}}|{{MESSAGES:START:(\d+)}}|{{MESSAGES:END:(\d+)}}|{{MESSAGES:MIDDLETRUNCATE:(\d+)}}",
|
|
replacement_function,
|
|
template,
|
|
)
|
|
|
|
return template
|
|
|
|
|
|
# {{prompt:middletruncate:8000}}
|
|
|
|
|
|
def rag_template(template: str, context: str, query: str):
|
|
if template.strip() == "":
|
|
template = DEFAULT_RAG_TEMPLATE
|
|
|
|
template = prompt_template(template)
|
|
|
|
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(
|
|
"WARNING: Potential prompt injection attack: the RAG "
|
|
"context contains '<context>' and '</context>'. This might be "
|
|
"nothing, or the user might be trying to hack something."
|
|
)
|
|
|
|
query_placeholders = []
|
|
if "[query]" in context:
|
|
query_placeholder = "{{QUERY" + str(uuid.uuid4()) + "}}"
|
|
template = template.replace("[query]", query_placeholder)
|
|
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)
|
|
|
|
return template
|
|
|
|
|
|
def title_generation_template(
|
|
template: str, messages: list[dict], user: Optional[dict] = None
|
|
) -> str:
|
|
prompt = get_last_user_message(messages)
|
|
template = replace_prompt_variable(template, prompt)
|
|
template = replace_messages_variable(template, messages)
|
|
|
|
template = prompt_template(
|
|
template,
|
|
**(
|
|
{"user_name": user.get("name"), "user_location": user.get("location")}
|
|
if user
|
|
else {}
|
|
),
|
|
)
|
|
|
|
return template
|
|
|
|
|
|
def follow_up_generation_template(
|
|
template: str, messages: list[dict], user: Optional[dict] = None
|
|
) -> str:
|
|
prompt = get_last_user_message(messages)
|
|
template = replace_prompt_variable(template, prompt)
|
|
template = replace_messages_variable(template, messages)
|
|
|
|
template = prompt_template(
|
|
template,
|
|
**(
|
|
{"user_name": user.get("name"), "user_location": user.get("location")}
|
|
if user
|
|
else {}
|
|
),
|
|
)
|
|
return template
|
|
|
|
|
|
def tags_generation_template(
|
|
template: str, messages: list[dict], user: Optional[dict] = None
|
|
) -> str:
|
|
prompt = get_last_user_message(messages)
|
|
template = replace_prompt_variable(template, prompt)
|
|
template = replace_messages_variable(template, messages)
|
|
|
|
template = prompt_template(
|
|
template,
|
|
**(
|
|
{"user_name": user.get("name"), "user_location": user.get("location")}
|
|
if user
|
|
else {}
|
|
),
|
|
)
|
|
return template
|
|
|
|
|
|
def image_prompt_generation_template(
|
|
template: str, messages: list[dict], user: Optional[dict] = None
|
|
) -> str:
|
|
prompt = get_last_user_message(messages)
|
|
template = replace_prompt_variable(template, prompt)
|
|
template = replace_messages_variable(template, messages)
|
|
|
|
template = prompt_template(
|
|
template,
|
|
**(
|
|
{"user_name": user.get("name"), "user_location": user.get("location")}
|
|
if user
|
|
else {}
|
|
),
|
|
)
|
|
return template
|
|
|
|
|
|
def emoji_generation_template(
|
|
template: str, prompt: str, user: Optional[dict] = None
|
|
) -> str:
|
|
template = replace_prompt_variable(template, prompt)
|
|
template = prompt_template(
|
|
template,
|
|
**(
|
|
{"user_name": user.get("name"), "user_location": user.get("location")}
|
|
if user
|
|
else {}
|
|
),
|
|
)
|
|
|
|
return template
|
|
|
|
|
|
def autocomplete_generation_template(
|
|
template: str,
|
|
prompt: str,
|
|
messages: Optional[list[dict]] = None,
|
|
type: Optional[str] = None,
|
|
user: Optional[dict] = None,
|
|
) -> str:
|
|
template = template.replace("{{TYPE}}", type if type else "")
|
|
template = replace_prompt_variable(template, prompt)
|
|
template = replace_messages_variable(template, messages)
|
|
|
|
template = prompt_template(
|
|
template,
|
|
**(
|
|
{"user_name": user.get("name"), "user_location": user.get("location")}
|
|
if user
|
|
else {}
|
|
),
|
|
)
|
|
return template
|
|
|
|
|
|
def query_generation_template(
|
|
template: str, messages: list[dict], user: Optional[dict] = None
|
|
) -> str:
|
|
prompt = get_last_user_message(messages)
|
|
template = replace_prompt_variable(template, prompt)
|
|
template = replace_messages_variable(template, messages)
|
|
|
|
template = prompt_template(
|
|
template,
|
|
**(
|
|
{"user_name": user.get("name"), "user_location": user.get("location")}
|
|
if user
|
|
else {}
|
|
),
|
|
)
|
|
return template
|
|
|
|
|
|
def moa_response_generation_template(
|
|
template: str, prompt: str, responses: list[str]
|
|
) -> str:
|
|
def replacement_function(match):
|
|
full_match = match.group(0)
|
|
start_length = match.group(1)
|
|
end_length = match.group(2)
|
|
middle_length = match.group(3)
|
|
|
|
if full_match == "{{prompt}}":
|
|
return prompt
|
|
elif start_length is not None:
|
|
return prompt[: int(start_length)]
|
|
elif end_length is not None:
|
|
return prompt[-int(end_length) :]
|
|
elif middle_length is not None:
|
|
middle_length = int(middle_length)
|
|
if len(prompt) <= middle_length:
|
|
return prompt
|
|
start = prompt[: math.ceil(middle_length / 2)]
|
|
end = prompt[-math.floor(middle_length / 2) :]
|
|
return f"{start}...{end}"
|
|
return ""
|
|
|
|
template = re.sub(
|
|
r"{{prompt}}|{{prompt:start:(\d+)}}|{{prompt:end:(\d+)}}|{{prompt:middletruncate:(\d+)}}",
|
|
replacement_function,
|
|
template,
|
|
)
|
|
|
|
responses = [f'"""{response}"""' for response in responses]
|
|
responses = "\n\n".join(responses)
|
|
|
|
template = template.replace("{{responses}}", responses)
|
|
return template
|
|
|
|
|
|
def tools_function_calling_generation_template(template: str, tools_specs: str) -> str:
|
|
template = template.replace("{{TOOLS}}", tools_specs)
|
|
return template
|