mirror of
https://github.com/open-webui/open-webui.git
synced 2025-12-12 04:15:25 +00:00
refac: rag context handling
This commit is contained in:
parent
72d19e44dc
commit
f096e99059
2 changed files with 36 additions and 39 deletions
|
|
@ -1171,26 +1171,15 @@ async def process_chat_payload(request, form_data, user, metadata, model):
|
||||||
raise Exception("No user message found")
|
raise Exception("No user message found")
|
||||||
|
|
||||||
if context_string != "":
|
if context_string != "":
|
||||||
# Workaround for Ollama 2.0+ system prompt issue
|
form_data["messages"] = add_or_update_user_message(
|
||||||
# TODO: replace with add_or_update_system_message
|
rag_template(
|
||||||
if model.get("owned_by") == "ollama":
|
request.app.state.config.RAG_TEMPLATE,
|
||||||
form_data["messages"] = prepend_to_first_user_message_content(
|
context_string,
|
||||||
rag_template(
|
prompt,
|
||||||
request.app.state.config.RAG_TEMPLATE,
|
),
|
||||||
context_string,
|
form_data["messages"],
|
||||||
prompt,
|
append=False,
|
||||||
),
|
)
|
||||||
form_data["messages"],
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
form_data["messages"] = add_or_update_system_message(
|
|
||||||
rag_template(
|
|
||||||
request.app.state.config.RAG_TEMPLATE,
|
|
||||||
context_string,
|
|
||||||
prompt,
|
|
||||||
),
|
|
||||||
form_data["messages"],
|
|
||||||
)
|
|
||||||
|
|
||||||
# If there are citations, add them to the data_items
|
# If there are citations, add them to the data_items
|
||||||
sources = [
|
sources = [
|
||||||
|
|
|
||||||
|
|
@ -120,19 +120,20 @@ def pop_system_message(messages: list[dict]) -> tuple[Optional[dict], list[dict]
|
||||||
return get_system_message(messages), remove_system_message(messages)
|
return get_system_message(messages), remove_system_message(messages)
|
||||||
|
|
||||||
|
|
||||||
def prepend_to_first_user_message_content(
|
def update_message_content(message: dict, content: str, append: bool = True) -> dict:
|
||||||
content: str, messages: list[dict]
|
if isinstance(message["content"], list):
|
||||||
) -> list[dict]:
|
for item in message["content"]:
|
||||||
for message in messages:
|
if item["type"] == "text":
|
||||||
if message["role"] == "user":
|
if append:
|
||||||
if isinstance(message["content"], list):
|
item["text"] = f"{item['text']}\n{content}"
|
||||||
for item in message["content"]:
|
else:
|
||||||
if item["type"] == "text":
|
item["text"] = f"{content}\n{item['text']}"
|
||||||
item["text"] = f"{content}\n{item['text']}"
|
else:
|
||||||
else:
|
if append:
|
||||||
message["content"] = f"{content}\n{message['content']}"
|
message["content"] = f"{message['content']}\n{content}"
|
||||||
break
|
else:
|
||||||
return messages
|
message["content"] = f"{content}\n{message['content']}"
|
||||||
|
return message
|
||||||
|
|
||||||
|
|
||||||
def add_or_update_system_message(
|
def add_or_update_system_message(
|
||||||
|
|
@ -148,10 +149,7 @@ def add_or_update_system_message(
|
||||||
"""
|
"""
|
||||||
|
|
||||||
if messages and messages[0].get("role") == "system":
|
if messages and messages[0].get("role") == "system":
|
||||||
if append:
|
messages[0] = update_message_content(messages[0], content, append)
|
||||||
messages[0]["content"] = f"{messages[0]['content']}\n{content}"
|
|
||||||
else:
|
|
||||||
messages[0]["content"] = f"{content}\n{messages[0]['content']}"
|
|
||||||
else:
|
else:
|
||||||
# Insert at the beginning
|
# Insert at the beginning
|
||||||
messages.insert(0, {"role": "system", "content": content})
|
messages.insert(0, {"role": "system", "content": content})
|
||||||
|
|
@ -159,7 +157,7 @@ def add_or_update_system_message(
|
||||||
return messages
|
return messages
|
||||||
|
|
||||||
|
|
||||||
def add_or_update_user_message(content: str, messages: list[dict]):
|
def add_or_update_user_message(content: str, messages: list[dict], append: bool = True):
|
||||||
"""
|
"""
|
||||||
Adds a new user message at the end of the messages list
|
Adds a new user message at the end of the messages list
|
||||||
or updates the existing user message at the end.
|
or updates the existing user message at the end.
|
||||||
|
|
@ -170,7 +168,7 @@ def add_or_update_user_message(content: str, messages: list[dict]):
|
||||||
"""
|
"""
|
||||||
|
|
||||||
if messages and messages[-1].get("role") == "user":
|
if messages and messages[-1].get("role") == "user":
|
||||||
messages[-1]["content"] = f"{messages[-1]['content']}\n{content}"
|
messages[-1] = update_message_content(messages[-1], content, append)
|
||||||
else:
|
else:
|
||||||
# Insert at the end
|
# Insert at the end
|
||||||
messages.append({"role": "user", "content": content})
|
messages.append({"role": "user", "content": content})
|
||||||
|
|
@ -178,6 +176,16 @@ def add_or_update_user_message(content: str, messages: list[dict]):
|
||||||
return messages
|
return messages
|
||||||
|
|
||||||
|
|
||||||
|
def prepend_to_first_user_message_content(
|
||||||
|
content: str, messages: list[dict]
|
||||||
|
) -> list[dict]:
|
||||||
|
for message in messages:
|
||||||
|
if message["role"] == "user":
|
||||||
|
message = update_message_content(message, content, append=False)
|
||||||
|
break
|
||||||
|
return messages
|
||||||
|
|
||||||
|
|
||||||
def append_or_update_assistant_message(content: str, messages: list[dict]):
|
def append_or_update_assistant_message(content: str, messages: list[dict]):
|
||||||
"""
|
"""
|
||||||
Adds a new assistant message at the end of the messages list
|
Adds a new assistant message at the end of the messages list
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue