various fixes from Yuhong's list

backend/onyx/agents/agent_search/deep_search_a/answer_initial_sub_question/nodes/format_answer.py

@@ -21,6 +21,7 @@ def format_answer(state: AnswerQuestionState) -> AnswerQuestionOutput:
                 answer=state.answer,
                 expanded_retrieval_results=state.expanded_retrieval_results,
                 documents=state.documents,
+                context_documents=state.context_documents,
                 sub_question_retrieval_stats=state.sub_question_retrieval_stats,
             )
         ],

backend/onyx/agents/agent_search/deep_search_a/expanded_retrieval/nodes/format_results.py

@@ -75,7 +75,7 @@ def format_results(
     return ExpandedRetrievalUpdate(
         expanded_retrieval_result=ExpandedRetrievalResult(
             expanded_queries_results=state.expanded_retrieval_results,
-            all_documents=state.reranked_documents,
+            all_documents=stream_documents,
             context_documents=state.reranked_documents,
             sub_question_retrieval_stats=sub_question_retrieval_stats,
         ),

backend/onyx/agents/agent_search/deep_search_a/main/graph_builder.py

@@ -32,6 +32,9 @@ from onyx.agents.agent_search.deep_search_a.main.nodes.agent_path_routing import
 from onyx.agents.agent_search.deep_search_a.main.nodes.agent_search_start import (
     agent_search_start,
 )
+from onyx.agents.agent_search.deep_search_a.main.nodes.answer_comparison import (
+    answer_comparison,
+)
 from onyx.agents.agent_search.deep_search_a.main.nodes.direct_llm_handling import (
     direct_llm_handling,
 )
@@ -184,7 +187,10 @@ def main_graph_builder(test_mode: bool = False) -> StateGraph:
         node="refined_answer_decision",
         action=refined_answer_decision,
     )
-
+    graph.add_node(
+        node="answer_comparison",
+        action=answer_comparison,
+    )
     graph.add_node(
         node="logging_node",
         action=agent_logging,
@@ -313,6 +319,10 @@ def main_graph_builder(test_mode: bool = False) -> StateGraph:
 
     graph.add_edge(
         start_key="generate_refined_answer",
+        end_key="answer_comparison",
+    )
+    graph.add_edge(
+        start_key="answer_comparison",
         end_key="logging_node",
     )
 

backend/onyx/agents/agent_search/deep_search_a/main/nodes/answer_comparison.py

@@ -0,0 +1,61 @@
+from datetime import datetime
+from typing import cast
+
+from langchain_core.callbacks.manager import dispatch_custom_event
+from langchain_core.messages import HumanMessage
+from langchain_core.runnables import RunnableConfig
+
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import AnswerComparison
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
+from onyx.agents.agent_search.models import AgentSearchConfig
+from onyx.agents.agent_search.shared_graph_utils.prompts import ANSWER_COMPARISON_PROMPT
+from onyx.chat.models import RefinedAnswerImprovement
+
+
+def answer_comparison(state: MainState, config: RunnableConfig) -> AnswerComparison:
+    now_start = datetime.now()
+
+    agent_a_config = cast(AgentSearchConfig, config["metadata"]["config"])
+    question = agent_a_config.search_request.query
+    initial_answer = state.initial_answer
+    refined_answer = state.refined_answer
+
+    logger.debug(f"--------{now_start}--------ANSWER COMPARISON STARTED--")
+
+    answer_comparison_prompt = ANSWER_COMPARISON_PROMPT.format(
+        question=question, initial_answer=initial_answer, refined_answer=refined_answer
+    )
+
+    msg = [HumanMessage(content=answer_comparison_prompt)]
+
+    # Get the rewritten queries in a defined format
+    model = agent_a_config.fast_llm
+
+    # no need to stream this
+    resp = model.invoke(msg)
+
+    if isinstance(resp.content, str) and "yes" in resp.content.lower():
+        refined_answer_improvement = True
+    else:
+        refined_answer_improvement = False
+
+    dispatch_custom_event(
+        "refined_answer_improvement",
+        RefinedAnswerImprovement(
+            refined_answer_improvement=refined_answer_improvement,
+        ),
+    )
+
+    now_end = datetime.now()
+
+    logger.debug(
+        f"--------{now_end}--{now_end - now_start}--------ANSWER COMPARISON COMPLETED---"
+    )
+
+    return AnswerComparison(
+        refined_answer_improvement=refined_answer_improvement,
+        log_messages=[
+            f"{now_start} -- Answer comparison: {refined_answer_improvement},  Time taken: {now_end - now_start}"
+        ],
+    )

backend/onyx/agents/agent_search/deep_search_a/main/nodes/generate_initial_answer.py

@@ -45,6 +45,7 @@ from onyx.agents.agent_search.shared_graph_utils.prompts import (
 from onyx.agents.agent_search.shared_graph_utils.prompts import UNKNOWN_ANSWER
 from onyx.agents.agent_search.shared_graph_utils.utils import format_docs
 from onyx.agents.agent_search.shared_graph_utils.utils import get_persona_prompt
+from onyx.agents.agent_search.shared_graph_utils.utils import get_today_prompt
 from onyx.agents.agent_search.shared_graph_utils.utils import parse_question_id
 from onyx.chat.models import AgentAnswerPiece
 from onyx.chat.models import ExtendedToolResponse
@@ -63,8 +64,10 @@ def generate_initial_answer(
     persona_prompt = get_persona_prompt(agent_a_config.search_request.persona)
 
     history = build_history_prompt(agent_a_config.prompt_builder)
+    date_str = get_today_prompt()
 
-    sub_question_docs = state.documents
+    # sub_question_docs = state.documents
+    sub_question_docs = state.context_documents
     all_original_question_documents = state.all_original_question_documents
 
     relevant_docs = dedup_inference_sections(
@@ -166,7 +169,11 @@ def generate_initial_answer(
         doc_context = trim_prompt_piece(
             model.config,
             doc_context,
-            base_prompt + sub_question_answer_str + persona_specification + history,
+            base_prompt
+            + sub_question_answer_str
+            + persona_specification
+            + history
+            + date_str,
         )
 
         msg = [
@@ -179,6 +186,7 @@ def generate_initial_answer(
                     relevant_docs=format_docs(relevant_docs),
                     persona_specification=persona_specification,
                     history=history,
+                    date_prompt=date_str,
                 )
             )
         ]

backend/onyx/agents/agent_search/deep_search_a/main/nodes/generate_refined_answer.py

@@ -42,6 +42,7 @@ from onyx.agents.agent_search.shared_graph_utils.prompts import (
 from onyx.agents.agent_search.shared_graph_utils.prompts import UNKNOWN_ANSWER
 from onyx.agents.agent_search.shared_graph_utils.utils import format_docs
 from onyx.agents.agent_search.shared_graph_utils.utils import get_persona_prompt
+from onyx.agents.agent_search.shared_graph_utils.utils import get_today_prompt
 from onyx.agents.agent_search.shared_graph_utils.utils import parse_question_id
 from onyx.chat.models import AgentAnswerPiece
 from onyx.chat.models import ExtendedToolResponse
@@ -60,7 +61,7 @@ def generate_refined_answer(
     persona_prompt = get_persona_prompt(agent_a_config.search_request.persona)
 
     history = build_history_prompt(agent_a_config.prompt_builder)
-
+    date_str = get_today_prompt()
     initial_documents = state.documents
     revised_documents = state.refined_documents
 
@@ -189,6 +190,7 @@ def generate_refined_answer(
                 relevant_docs=relevant_docs,
                 initial_answer=remove_document_citations(initial_answer),
                 persona_specification=persona_specification,
+                date_prompt=date_str,
             )
         )
     ]

backend/onyx/agents/agent_search/deep_search_a/main/nodes/ingest_initial_base_retrieval.py

@@ -31,7 +31,7 @@ def ingest_initial_base_retrieval(
 
     return ExpandedRetrievalUpdate(
         original_question_retrieval_results=state.base_expanded_retrieval_result.expanded_queries_results,
-        all_original_question_documents=state.base_expanded_retrieval_result.all_documents,
+        all_original_question_documents=state.base_expanded_retrieval_result.context_documents,
         original_question_retrieval_stats=sub_question_retrieval_stats,
         log_messages=[
             f"{now_end} -- Main - Ingestion base retrieval,  Time taken: {now_end - now_start}"

backend/onyx/agents/agent_search/deep_search_a/main/nodes/ingest_initial_sub_question_answers.py

@@ -17,9 +17,11 @@ def ingest_initial_sub_question_answers(
 
     logger.debug(f"--------{now_start}--------INGEST ANSWERS---")
     documents = []
+    context_documents = []
     answer_results = state.answer_results if hasattr(state, "answer_results") else []
     for answer_result in answer_results:
         documents.extend(answer_result.documents)
+        context_documents.extend(answer_result.context_documents)
 
     now_end = datetime.now()
 
@@ -31,6 +33,7 @@ def ingest_initial_sub_question_answers(
         # Deduping is done by the documents operator for the main graph
         # so we might not need to dedup here
         documents=dedup_inference_sections(documents, []),
+        context_documents=dedup_inference_sections(context_documents, []),
         decomp_answer_results=answer_results,
         log_messages=[
             f"{now_end} -- Main - Ingest initial processed sub questions,  Time taken: {now_end - now_start}"

backend/onyx/agents/agent_search/deep_search_a/main/states.py

@@ -53,6 +53,14 @@ class BaseDecompUpdateBase(BaseModel):
     initial_decomp_questions: list[str] = []
 
 
+class AnswerComparisonBase(BaseModel):
+    refined_answer_improvement: bool = False
+
+
+class AnswerComparison(AnswerComparisonBase, LoggerUpdate):
+    pass
+
+
 class RoutingDecisionBase(BaseModel):
     routing: str = ""
     sample_doc_str: str = ""
@@ -108,6 +116,7 @@ class RequireRefinedAnswerUpdate(LoggingUpdate):
 
 class DecompAnswersUpdate(LoggingUpdate):
     documents: Annotated[list[InferenceSection], dedup_inference_sections] = []
+    context_documents: Annotated[list[InferenceSection], dedup_inference_sections] = []
     decomp_answer_results: Annotated[
         list[QuestionAnswerResults], dedup_question_answer_results
     ] = []
@@ -175,6 +184,7 @@ class MainState(
     RefinedAgentStartStats,
     RefinedAgentEndStats,
     RoutingDecisionBase,
+    AnswerComparisonBase,
 ):
     # expanded_retrieval_result: Annotated[list[ExpandedRetrievalResult], add]
     base_raw_search_result: Annotated[list[ExpandedRetrievalResult], add]

backend/onyx/agents/agent_search/run_graph.py

@@ -19,6 +19,7 @@ from onyx.chat.models import AgentAnswerPiece
 from onyx.chat.models import AnswerPacket
 from onyx.chat.models import AnswerStream
 from onyx.chat.models import ExtendedToolResponse
+from onyx.chat.models import RefinedAnswerImprovement
 from onyx.chat.models import StreamStopInfo
 from onyx.chat.models import SubQueryPiece
 from onyx.chat.models import SubQuestionPiece
@@ -75,6 +76,8 @@ def _parse_agent_event(
             return cast(ToolResponse, event["data"])
         elif event["name"] == "basic_response":
             return cast(AnswerPacket, event["data"])
+        elif event["name"] == "refined_answer_improvement":
+            return cast(RefinedAnswerImprovement, event["data"])
     return None
 
 
@@ -262,6 +265,8 @@ if __name__ == "__main__":
                 logger.info(
                     f"   ---------- FA {output.level} - {output.level_question_nr}  {output.answer_piece} | "
                 )
+            elif isinstance(output, RefinedAnswerImprovement):
+                logger.info(f"   ---------- RE {output.refined_answer_improvement} | ")
 
         # for tool_response in tool_responses:
         #    logger.debug(tool_response)

backend/onyx/agents/agent_search/shared_graph_utils/agent_prompt_ops.py

@@ -5,6 +5,7 @@ from langchain_core.messages.tool import ToolMessage
 
 from onyx.agents.agent_search.shared_graph_utils.prompts import BASE_RAG_PROMPT_v2
 from onyx.agents.agent_search.shared_graph_utils.prompts import HISTORY_PROMPT
+from onyx.agents.agent_search.shared_graph_utils.utils import get_today_prompt
 from onyx.chat.prompt_builder.answer_prompt_builder import AnswerPromptBuilder
 from onyx.context.search.models import InferenceSection
 from onyx.llm.interfaces import LLMConfig
@@ -24,6 +25,8 @@ def build_sub_question_answer_prompt(
         content=persona_specification,
     )
 
+    date_str = get_today_prompt()
+
     docs_format_list = [
         f"""Document Number: [D{doc_nr + 1}]\n
                              Content: {doc.combined_content}\n\n"""
@@ -33,11 +36,14 @@ def build_sub_question_answer_prompt(
     docs_str = "\n\n".join(docs_format_list)
 
     docs_str = trim_prompt_piece(
-        config, docs_str, BASE_RAG_PROMPT_v2 + question + original_question
+        config, docs_str, BASE_RAG_PROMPT_v2 + question + original_question + date_str
     )
     human_message = HumanMessage(
         content=BASE_RAG_PROMPT_v2.format(
-            question=question, original_question=original_question, context=docs_str
+            question=question,
+            original_question=original_question,
+            context=docs_str,
+            date_prompt=date_str,
         )
     )
 
@@ -74,17 +80,21 @@ def build_history_prompt(prompt_builder: AnswerPromptBuilder | None) -> str:
     if prompt_builder.single_message_history is not None:
         history = prompt_builder.single_message_history
     else:
-        history = ""
+        history_components = []
         previous_message_type = None
         for message in prompt_builder.raw_message_history:
             if "user" in message.message_type:
-                history += f"User: {message.message}\n"
+                history_components.append(f"User: {message.message}\n")
                 previous_message_type = "user"
             elif "assistant" in message.message_type:
-                # only use the initial agent answer for the history
+                # only use the last agent answer for the history
                 if previous_message_type != "assistant":
-                    history += f"You/Agent: {message.message}\n"
+                    history_components.append(f"You/Agent: {message.message}\n")
+                else:
+                    history_components = history_components[:-1]
+                    history_components.append(f"You/Agent: {message.message}\n")
                 previous_message_type = "assistant"
             else:
                 continue
+        history = "\n".join(history_components)
     return HISTORY_PROMPT.format(history=history) if history else ""

backend/onyx/agents/agent_search/shared_graph_utils/models.py

@@ -102,6 +102,7 @@ class QuestionAnswerResults(BaseModel):
     quality: str
     expanded_retrieval_results: list[QueryResult]
     documents: list[InferenceSection]
+    context_documents: list[InferenceSection]
     sub_question_retrieval_stats: AgentChunkStats
 
 

backend/onyx/agents/agent_search/shared_graph_utils/prompts.py

@@ -2,6 +2,8 @@ UNKNOWN_ANSWER = "I do not have enough information to answer this question."
 
 NO_RECOVERED_DOCS = "No relevant documents recovered"
 
+DATE_PROMPT = """Today is {date}.\n\n"""
+
 HISTORY_PROMPT = """\n
 For more context, here is the history of the conversation so far that preceeded this question:
 \n ------- \n
@@ -37,6 +39,7 @@ REWRITE_PROMPT_MULTI = """ \n
 BASE_RAG_PROMPT = (
     """ \n
     {persona_specification}
+    {date_prompt}
     Use the context provided below - and only the
     provided context - to answer the given question. (Note that the answer is in service of anserwing a broader
     question, given below as 'motivation'.)
@@ -44,8 +47,9 @@ BASE_RAG_PROMPT = (
     Again, only use the provided context and do not use your internal knowledge! If you cannot answer the
     question based on the context, say """
     + f'"{UNKNOWN_ANSWER}"'
-    + """. It is a matter of life and death that you do NOT
-    use your internal knowledge, just the provided information!
+    + """.
+    It is a matter of life and death that you do NOT use your internal knowledge, just the provided
+    information!
 
     Make sure that you keep all relevant information, specifically as it concerns to the ultimate goal.
     (But keep other details as well.)
@@ -61,6 +65,7 @@ BASE_RAG_PROMPT = (
 
 BASE_RAG_PROMPT_v2 = (
     """ \n
+    {date_prompt}
     Use the context provided below - and only the
     provided context - to answer the given question. (Note that the answer is in service of answering a broader
     question, given below as 'motivation'.)
@@ -74,7 +79,8 @@ BASE_RAG_PROMPT_v2 = (
     Make sure that you keep all relevant information, specifically as it concerns to the ultimate goal.
     (But keep other details as well.)
 
-    Please remember to provide inline citations in the format [[D1]](), [[D2]](), [[D3]](), etc.
+    Please remember to provide inline citations in the format [[D1]](), [[D2]](), [[D3]](), etc!
+     It is important that the citation is close to the information it supports.
     Proper citations are very important to the user!\n\n\n
 
     For your general information, here is the ultimate motivation:
@@ -534,7 +540,7 @@ INITIAL_DECOMPOSITION_PROMPT = """ \n
 
 INITIAL_RAG_BASE_PROMPT = (
     """ \n
-You are an assistant for question-answering tasks. Use the information provided below - and only the
+ You are an assistant for question-answering tasks. Use the information provided below - and only the
 provided information - to answer the provided question.
 
 The information provided below consists ofa number of documents that were deemed relevant for the question.
@@ -542,7 +548,7 @@ The information provided below consists ofa number of documents that were deemed
 IMPORTANT RULES:
 - If you cannot reliably answer the question solely using the provided information, say that you cannot reliably answer.
 You may give some additional facts you learned, but do not try to invent an answer.
-- If the information is empty or irrelevant, just say """
+- If the information is irrelevant, just say """
     + f'"{UNKNOWN_ANSWER}"'
     + """.
 - If the information is relevant but not fully conclusive, specify that the information is not conclusive and say why.
@@ -643,7 +649,7 @@ SUB_QUESTION_SEARCH_RESULTS_TEMPLATE = """
 INITIAL_RAG_PROMPT_SUB_QUESTION_SEARCH = (
     """ \n
 {persona_specification}
-
+{date_prompt}
 Use the information provided below - and only the provided information - to answer the main question that will be provided.
 
 The information provided below consists of:
@@ -660,7 +666,8 @@ IMPORTANT RULES:
  - If the information is relevant but not fully conclusive, specify that the information is not conclusive and say why.
  - The answers to the subquestions should help you to structure your thoughts in order to answer the question.
 
-Please provide inline citations of documentsin the format [[D1]](), [[D2]](), [[D3]](), etc.,  If you have multiple citations,
+Please provide inline citations of documentsin the format [[D1]](), [[D2]](), [[D3]](), etc.!
+It is important that the citation is close to the information it supports.  If you have multiple citations,
 please cite for example as [[D1]]()[[D3]](), or [[D2]]()[[D4]](), etc. Feel free to cite documents in addition
 to the sub-questions! Proper citations are important for the final answer to be verifiable! \n\n\n
 
@@ -700,7 +707,7 @@ Answer:"""
 INITIAL_RAG_PROMPT = (
     """ \n
 {persona_specification}
-
+ {date_prompt}
 Use the information provided below - and only the provided information - to answer the provided main question.
 
 The information provided below consists of:
@@ -710,12 +717,13 @@ answer
 
 {history}
 
-Please provide inline citations to documents in the format [[D1]](), [[D2]](), [[D3]](), etc. If you have multiple
-citations, please cite for example as [[D1]]()[[D3]](), or [[D2]]()[[D4]](), etc.
+Please provide inline citations to documents in the format [[D1]](), [[D2]](), [[D3]](), etc.! It is important that the citation
+is close to the information it supports. If you have multiple citations that support a fact, please cite for example
+as [[D1]]()[[D3]](), or [[D2]]()[[D4]](), etc.
 Feel free to also cite sub-questions in addition to documents, but make sure that you have documents cited with the sub-question
 citation. If you want to cite both a document and a sub-question, please use [[D1]]()[[Q3]](), or [[D2]]()[[D7]]()[[Q4]](), etc.
-Again, please do not cite sub-questions without a document citation!
-Citations are very important for the user!
+Again, please do NEVER cite sub-questions without a document citation!
+Proper citations are very important for the user!
 
 IMPORTANT RULES:
  - If you cannot reliably answer the question solely using the provided information, say that you cannot reliably answer.
@@ -752,6 +760,7 @@ Answer:"""
 INITIAL_RAG_PROMPT_NO_SUB_QUESTIONS = (
     """{answered_sub_questions}
 {persona_specification}
+{date_prompt}
 Use the information provided below
 - and only the provided information - to answer the provided question.
 The information provided below consists of a number of documents that were deemed relevant for the question.
@@ -767,7 +776,8 @@ IMPORTANT RULES:
 
 Again, you should be sure that the answer is supported by the information provided!
 
-Please provide inline citations to documents in the format [[D1]](), [[D2]](), [[D3]](), etc. If you have multiple
+Please provide inline citations to documents in the format [[D1]](), [[D2]](), [[D3]](), etc!  It is important that the citation
+is close to the information it supports. If you have multiple
 citations, please cite for example as [[D1]]()[[D3]](), or [[D2]]()[[D4]](), etc. Citations are very important for the
 user!
 
@@ -789,6 +799,7 @@ Answer:"""
 REVISED_RAG_PROMPT = (
     """\n
 {persona_specification}
+{date_prompt}
 Use the information provided below - and only the provided information - to answer the provided main question.
 
 The information provided below consists of:
@@ -800,12 +811,13 @@ the main question. Note that the sub-questions have a type, 'initial' and 'revis
     3) a number of documents that were deemed relevant for the question. This the is the context that you use largey for
 citations (see below).
 
-Please provide inline citations to documents in the format [[D1]](), [[D2]](), [[D3]](), etc. If you have multiple
+Please provide inline citations to documents in the format [[D1]](), [[D2]](), [[D3]](), etc!
+ It is important that the citation is close to the information it supports. If you have multiple
 citations, please cite for example as [[D1]]()[[D3]](), or [[D2]]()[[D4]](), etc.
 Feel free to also cite sub-questions in addition to documents, but make sure that you have documents cited with the sub-question
 citation. If you want to cite both a document and a sub-question, please use [[D1]]()[[Q3]](), or [[D2]]()[[D7]]()[[Q4]](), etc.
-Again, please do not cite sub-questions without a document citation!
-Citations are very important for the user!\n\n
+Again, please do NEVER cite sub-questions without a document citation!
+Proper citations are very important for the user!\n\n
 
 {history}
 
@@ -853,6 +865,7 @@ Answer:"""
 REVISED_RAG_PROMPT_NO_SUB_QUESTIONS = (
     """{answered_sub_questions}\n
 {persona_specification}
+{date_prompt}
 Use the information provided below - and only the
 provided information - to answer the provided question.
 
@@ -860,7 +873,8 @@ The information provided below consists of:
     1) an initial answer that was given but found to be lacking in some way.
     2) a number of documents that were also deemed relevant for the question.
 
-Please provide inline citations to documents in the format [[D1]](), [[D2]](), [[D3]](), etc. If you have multiple
+Please provide inline citations to documents in the format [[D1]](), [[D2]](), [[D3]](), etc!
+ It is important that the citation is close to the information it supports. If you have multiple
 citations, please cite for example as [[D1]]()[[D3]](), or [[D2]]()[[D4]](), etc. Citations are very important for the user!\n\n
 
 {history}
@@ -933,4 +947,27 @@ ENTITY_TERM_PROMPT = """ \n
         }}]
     }}
     }}
+
+   """
+ANSWER_COMPARISON_PROMPT = """
+    For the given question, please compare the initial answer and the refined answer and determine if
+the refined answer is substantially better than the initial answer. Better could mean:
+ - additional information
+ - more comprehensive information
+ - more concise information
+ - more structured information
+ - substantially more document citations ([[D1]](), [[D2]](), [[D3]](), etc.)
+
+ Here is the question:
+ {question}
+
+ Here is the initial answer:
+ {initial_answer}
+
+ Here is the refined answer:
+ {refined_answer}
+
+ With these criteriain mind, is the refined answer substantially better than the initial answer?
+
+ Please answer with a simple 'yes' or 'no'.
     """

backend/onyx/agents/agent_search/shared_graph_utils/utils.py

@@ -18,6 +18,7 @@ from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.models import (
     EntityRelationshipTermExtraction,
 )
+from onyx.agents.agent_search.shared_graph_utils.prompts import DATE_PROMPT
 from onyx.chat.models import AnswerStyleConfig
 from onyx.chat.models import CitationConfig
 from onyx.chat.models import DocumentPruningConfig
@@ -276,3 +277,7 @@ def dispatch_separated(
         streamed_tokens.append(content)
 
     return streamed_tokens
+
+
+def get_today_prompt() -> str:
+    return DATE_PROMPT.format(date=datetime.now().strftime("%A, %B %d, %Y"))

backend/onyx/chat/models.py

@@ -335,12 +335,20 @@ class ExtendedToolResponse(ToolResponse):
     level_question_nr: int
 
 
-ProSearchPacket = (
-    SubQuestionPiece | AgentAnswerPiece | SubQueryPiece | ExtendedToolResponse
+class RefinedAnswerImprovement(BaseModel):
+    refined_answer_improvement: bool
+
+
+AgentSearchPacket = (
+    SubQuestionPiece
+    | AgentAnswerPiece
+    | SubQueryPiece
+    | ExtendedToolResponse
+    | RefinedAnswerImprovement
 )
 
 AnswerPacket = (
-    AnswerQuestionPossibleReturn | ProSearchPacket | ToolCallKickoff | ToolResponse
+    AnswerQuestionPossibleReturn | AgentSearchPacket | ToolCallKickoff | ToolResponse
 )
 
 
@@ -351,7 +359,7 @@ ResponsePart = (
     | ToolResponse
     | ToolCallFinalResult
     | StreamStopInfo
-    | ProSearchPacket
+    | AgentSearchPacket
 )
 
 AnswerStream = Iterator[AnswerPacket]

backend/onyx/chat/process_message.py

@@ -12,6 +12,7 @@ from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.chat.answer import Answer
 from onyx.chat.chat_utils import create_chat_chain
 from onyx.chat.chat_utils import create_temporary_persona
+from onyx.chat.models import AgentSearchPacket
 from onyx.chat.models import AllCitations
 from onyx.chat.models import AnswerStyleConfig
 from onyx.chat.models import ChatOnyxBotResponse
@@ -28,7 +29,6 @@ from onyx.chat.models import MessageSpecificCitations
 from onyx.chat.models import OnyxAnswerPiece
 from onyx.chat.models import OnyxContexts
 from onyx.chat.models import PromptConfig
-from onyx.chat.models import ProSearchPacket
 from onyx.chat.models import QADocsResponse
 from onyx.chat.models import StreamingError
 from onyx.chat.models import StreamStopInfo
@@ -304,7 +304,7 @@ ChatPacket = (
     | MessageSpecificCitations
     | MessageResponseIDInfo
     | StreamStopInfo
-    | ProSearchPacket
+    | AgentSearchPacket
 )
 ChatPacketStream = Iterator[ChatPacket]