refac: rag context handling

This commit is contained in:
Timothy Jaeryang Baek 2025-09-24 09:38:14 -05:00
parent 72d19e44dc
commit f096e99059
2 changed files with 36 additions and 39 deletions

View file

@ -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 = [

View file

@ -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