improvements

Author: joachim-danswer

backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py

@@ -20,6 +20,8 @@
 from onyx.agents.agent_search.shared_graph_utils.utils import (
     get_langgraph_node_log_string,
 )
+from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
+from onyx.chat.models import AgentAnswerPiece
 from onyx.kg.utils.extraction_utils import get_entity_types_str
 from onyx.kg.utils.extraction_utils import get_relationship_types_str
 from onyx.prompts.dr_prompts import FAST_DR_DECISION_PROMPT
@@ -146,6 +148,17 @@ def orchestrator(
             except Exception as e:
                 logger.error(f"Error in plan generation: {e}")
                 raise e
+
+            write_custom_event(
+                "basic_response",
+                AgentAnswerPiece(
+                    answer_piece=f"\n\nHIGH_LEVEL PLAN: {plan_information.plan}\n\n\n",
+                    level=0,
+                    level_question_num=0,
+                    answer_type="agent_level_answer",
+                ),
+                writer,
+            )
         else:
             plan_information = state.plan_of_record[-1]
 

backend/onyx/agents/agent_search/dr/nodes/dr_a2_closer.py

@@ -1,20 +1,19 @@
 from datetime import datetime
 
-from langchain_core.messages import HumanMessage
 from langchain_core.runnables import RunnableConfig
 from langgraph.types import StreamWriter
 
 from onyx.agents.agent_search.dr.states import FinalUpdate
 from onyx.agents.agent_search.dr.states import MainState
 from onyx.agents.agent_search.dr.utils import aggregate_context
+from onyx.agents.agent_search.shared_graph_utils.llm import get_answer_from_llm
 from onyx.agents.agent_search.shared_graph_utils.utils import (
     get_langgraph_node_log_string,
 )
 from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
 from onyx.chat.models import AgentAnswerPiece
 from onyx.prompts.dr_prompts import FINAL_ANSWER_PROMPT
 from onyx.utils.logger import setup_logger
-from onyx.utils.threadpool_concurrency import run_with_timeout
 
 logger = setup_logger()
 
@@ -39,27 +38,13 @@ def closer(
         "---base_question---", base_question
     ).replace("---iteration_responses_string---", iteration_responses_string)
 
-    msg = [
-        HumanMessage(
-            content=final_answer_prompt,
-        )
-    ]
-
-    primary_llm = config["metadata"]["config"].tooling.primary_llm
-
-    try:
-        llm_response = run_with_timeout(
-            25,
-            primary_llm.invoke,
-            prompt=msg,
-            timeout_override=45,
-            max_tokens=1500,
-        )
-
-        final_answer = str(llm_response.content).replace("```json\n", "")
-
-    except Exception as e:
-        raise ValueError(f"Error in closer: {e}")
+    final_answer = get_answer_from_llm(
+        llm=config["metadata"]["config"].tooling.primary_llm,
+        prompt=final_answer_prompt,
+        timeout=60,
+        max_tokens=1500,
+        timeout_override=60,
+    )
 
     logger.debug(final_answer)
 

backend/onyx/agents/agent_search/dr/nodes/dr_i1_search.py

@@ -155,21 +155,33 @@ def search(
     write_custom_event(
         "basic_response",
         AgentAnswerPiece(
-            answer_piece=f"\n\nSUB-QUESTION: {search_query}\n\nANSWER: {search_answer}\n\n",
+            answer_piece=f"\n\nSUB-QUESTION (SEARCH): {search_query}\n\n-> answered!\n\n",
             level=0,
             level_question_num=0,
             answer_type="agent_level_answer",
         ),
         writer,
     )
 
+    # handle citations
+
     citation_numbers = get_cited_document_numbers(search_answer)
+    citation_number_replacement_dict = {
+        original_index: start_1_based_index + 1
+        for start_1_based_index, original_index in enumerate(citation_numbers)
+    }
 
     cited_documents: list[InferenceSection] = []
 
     for citation_number in citation_numbers:
         cited_documents.append(retrieved_docs[citation_number - 1])
 
+    # change citations from search answer
+    for original_index, replacement_index in citation_number_replacement_dict.items():
+        search_answer = search_answer.replace(
+            f"[{original_index}]", f"[{replacement_index}]"
+        )
+
     return AnswerUpdate(
         answers=[search_answer],
         iteration_responses=[

backend/onyx/prompts/dr_prompts.py

@@ -491,6 +491,9 @@
  - note that the sub-answers to the sub-questions are designed to be high-level, mostly \
 focussing on providing the citations and providing some answer facts. But the \
 main content should be in the cited documents for each sub-question.
+- make sure that a the text from a document that you use is NOT TAKEN OUT OF CONTEXT!
+- do not make anything up! Only use the information provided in the documents, or, \
+if no documents are provided for a sub-answer, in the actual sub-answer.
 - Provide a thoughtful answer that is concise and to the point, but that is detailed.
 - Please cite your sources inline in format [2], [4], etc! The numbers of the documents \
 are provided above.