Repository: uktrade/redbox
Pull request #197: Feature/improve evaluator
Author: login: nora-errouhly
Created: 2025-04-16T14:06:17Z | Merged: 2025-04-17T12:17:35Z

- changing graph (newroute state management)
- changing evaluator and planner prompts
- adding env variable to switch on/off citations in new route. value is set to False in env file (switching off citations in newroute)



Diff:
diff --git a/.env.example b/.env.example
index 6c449423c..982ff59df 100644
--- a/.env.example
+++ b/.env.example
@@ -102,5 +102,6 @@ ENABLE_METADATA_EXTRACTION = True
 SELF_ROUTE_ENABLED = False
 USE_CLAM_AV = False
 
+ENABLE_CITATION_NEWROUTE = False
 # Transcription
-AWS_TRANSCRIBE_ROLE_ARN = 
\ No newline at end of file
+AWS_TRANSCRIBE_ROLE_ARN = 
diff --git a/redbox-core/redbox/chains/runnables.py b/redbox-core/redbox/chains/runnables.py
index 661fdbe51..3eeeabe27 100644
--- a/redbox-core/redbox/chains/runnables.py
+++ b/redbox-core/redbox/chains/runnables.py
@@ -72,16 +72,18 @@ def _chat_prompt_from_messages(state: RedboxState) -> Runnable:
             | _additional_variables
         )
 
-        return ChatPromptTemplate(
+        chatprompt = ChatPromptTemplate(
             messages=(
                 [("system", system_prompt_message)]
                 + [(msg["role"], msg["text"]) for msg in truncated_history]
                 + [MessagesPlaceholder("messages")]
-                + [task_question_prompt + "\n\n{format_instructions}"]
+                + [task_question_prompt]
             ),
             partial_variables={"format_instructions": format_instructions},
         ).invoke(prompt_template_context)
 
+        return chatprompt
+
     return _chat_prompt_from_messages
 
 
diff --git a/redbox-core/redbox/graph/nodes/processes.py b/redbox-core/redbox/graph/nodes/processes.py
index 27a1bae62..74575c17f 100644
--- a/redbox-core/redbox/graph/nodes/processes.py
+++ b/redbox-core/redbox/graph/nodes/processes.py
@@ -27,11 +27,16 @@
 from redbox.models.chain import AgentTask, DocumentState, MultiAgentPlan, PromptSet, RedboxState, RequestMetadata
 from redbox.models.graph import ROUTE_NAME_TAG, SOURCE_DOCUMENTS_TAG, RedboxActivityEvent, RedboxEventType
 from redbox.models.prompts import PLANNER_PROMPT
-from redbox.transform import combine_documents, flatten_document_state
+from redbox.transform import combine_agents_state, combine_documents, flatten_document_state
+from langchain_core.messages import RemoveMessage
+from redbox.models.settings import get_settings
+
 
 log = logging.getLogger(__name__)
 re_keyword_pattern = re.compile(r"@(\w+)")
 
+env = get_settings()
+
 
 # Patterns: functions that build processes
 
@@ -390,22 +395,39 @@ def _build_agent(state: RedboxState):
         )
         ai_msg = worker_agent.invoke(state)
         result = run_tools_parallel(ai_msg, tools, state)
-        result = f"<{agent_name}_Result>{result}</{agent_name}_Result>"
-        return {"messages": result}
+        result_content = "".join([res.content for res in result])
+        result = f"<{agent_name}_Result>{result_content}</{agent_name}_Result>"
+        return {"agents_results": result}
 
     return _build_agent
 
 
 def create_evaluator():
-    citation_parser, format_instructions = get_structured_response_with_citations_parser()
-    evaluator_agent = build_stuff_pattern(
-        prompt_set=PromptSet.NewRoute,
-        tools=None,
-        output_parser=citation_parser,
-        format_instructions=format_instructions,
-        final_response_chain=False,
-    )
-    return evaluator_agent
+    def _create_evaluator(state: RedboxState):
+        _additional_variables = {"agents_results": combine_agents_state(state.agents_results)}
+        ENABLE_CITATION_NEWROUTE = env.enable_citation_newroute
+        if ENABLE_CITATION_NEWROUTE:
+            citation_parser, format_instructions = get_structured_response_with_citations_parser()
+            evaluator_agent = build_stuff_pattern(
+                prompt_set=PromptSet.NewRoute,
+                tools=None,
+                output_parser=citation_parser,
+                format_instructions=format_instructions,
+                final_response_chain=False,
+                additional_variables=_additional_variables,
+            )
+        else:
+            evaluator_agent = build_stuff_pattern(
+                prompt_set=PromptSet.NewRouteNoCitations,
+                tools=None,
+                output_parser=None,
+                format_instructions="",
+                final_response_chain=True,
+                additional_variables=_additional_variables,
+            )
+        return evaluator_agent
+
+    return _create_evaluator
 
 
 def invoke_custom_state(
@@ -432,6 +454,15 @@ def _invoke_custom_state(state: RedboxState):
         result = f"<{agent_name}_Result>{result}</{agent_name}_Result>"
 
         # transform response back to the parent state
-        return {"messages": result}
+        return {"agents_results": result}
 
     return _invoke_custom_state
+
+
+def delete_plan_message():
+    @RunnableLambda
+    def _delete_plan_message(state: RedboxState):
+        last_message = state.messages[-1]
+        return {"messages": [RemoveMessage(id=last_message.id)]}
+
+    return _delete_plan_message
diff --git a/redbox-core/redbox/graph/root.py b/redbox-core/redbox/graph/root.py
index 4907355be..02cd5bc7c 100644
--- a/redbox-core/redbox/graph/root.py
+++ b/redbox-core/redbox/graph/root.py
@@ -39,6 +39,7 @@
     lm_choose_route,
     report_sources_process,
     invoke_custom_state,
+    delete_plan_message,
 )
 from redbox.graph.nodes.sends import (
     build_document_chunk_send,
@@ -852,6 +853,8 @@ def build_new_graph(
     )
 
     builder.add_node("Evaluator_Agent", create_evaluator())
+    builder.add_node("clear_tasks", delete_plan_message())
+    builder.add_node("pass_user_prompt_to_LLM_message", build_passthrough_pattern())
     builder.add_node(
         "report_citations",
         report_sources_process,
@@ -860,9 +863,11 @@ def build_new_graph(
 
     builder.add_edge(START, "planner")
     builder.add_conditional_edges("planner", sending_task_to_agent)
-    builder.add_edge("Document_Agent", "Evaluator_Agent")
-    builder.add_edge("External_Data_Agent", "Evaluator_Agent")
-    builder.add_edge("Summarisation_Agent", "Evaluator_Agent")
+    builder.add_edge("Document_Agent", "clear_tasks")
+    builder.add_edge("External_Data_Agent", "clear_tasks")
+    builder.add_edge("Summarisation_Agent", "clear_tasks")
+    builder.add_edge("clear_tasks", "pass_user_prompt_to_LLM_message")
+    builder.add_edge("pass_user_prompt_to_LLM_message", "Evaluator_Agent")
     builder.add_edge("Evaluator_Agent", "report_citations")
     builder.add_edge("report_citations", END)
 
diff --git a/redbox-core/redbox/models/chain.py b/redbox-core/redbox/models/chain.py
index d831e5cfa..20346aeb8 100644
--- a/redbox-core/redbox/models/chain.py
+++ b/redbox-core/redbox/models/chain.py
@@ -66,6 +66,8 @@ class AISettings(BaseModel):
     chat_map_question_prompt: str = prompts.CHAT_MAP_QUESTION_PROMPT
     reduce_system_prompt: str = prompts.REDUCE_SYSTEM_PROMPT
     new_route_retrieval_system_prompt: str = prompts.NEW_ROUTE_RETRIEVAL_SYSTEM_PROMPT
+    new_route_retrieval_system_prompt_nocitations: str = prompts.NEW_ROUTE_RETRIEVAL_SYSTEM_PROMPT_NOCITATIONS
+    new_route_retrieval_question_prompt: str = prompts.NEW_ROUTE_RETRIEVAL_QUESTION_PROMPT
     llm_decide_route_prompt: str = prompts.LLM_DECIDE_ROUTE
 
     # Elasticsearch RAG and boost values
@@ -255,6 +257,7 @@ class RedboxState(BaseModel):
     citations: list[Citation] | None = None
     steps_left: Annotated[int | None, RemainingStepsManager] = None
     messages: Annotated[list[AnyMessage], add_messages] = Field(default_factory=list)
+    agents_results: Annotated[list[AnyMessage], add_messages] = Field(default_factory=list)
 
     @property
     def last_message(self) -> AnyMessage:
@@ -273,6 +276,7 @@ class PromptSet(StrEnum):
     SelfRoute = "self_route"
     CondenseQuestion = "condense_question"
     NewRoute = "new_route"
+    NewRouteNoCitations = "new_route_no_citations"
 
 
 def get_prompts(state: RedboxState, prompt_set: PromptSet) -> tuple[str, str]:
@@ -302,7 +306,10 @@ def get_prompts(state: RedboxState, prompt_set: PromptSet) -> tuple[str, str]:
         question_prompt = state.request.ai_settings.condense_question_prompt
     elif prompt_set == PromptSet.NewRoute:
         system_prompt = state.request.ai_settings.new_route_retrieval_system_prompt
-        question_prompt = state.request.ai_settings.agentic_retrieval_question_prompt
+        question_prompt = state.request.ai_settings.new_route_retrieval_question_prompt
+    elif prompt_set == PromptSet.NewRouteNoCitations:
+        system_prompt = state.request.ai_settings.new_route_retrieval_system_prompt_nocitations
+        question_prompt = state.request.ai_settings.new_route_retrieval_question_prompt
     return (system_prompt, question_prompt)
 
 
diff --git a/redbox-core/redbox/models/prompts.py b/redbox-core/redbox/models/prompts.py
index a2f1a6cd6..654d16344 100644
--- a/redbox-core/redbox/models/prompts.py
+++ b/redbox-core/redbox/models/prompts.py
@@ -40,43 +40,38 @@
 
    """
 
-NEW_ROUTE_RETRIEVAL_SYSTEM_PROMPT = """Expert Answer Evaluation Protocol:
-
-1. Comprehensive Analysis:
-   - Evaluate question against full message history
-   - Assess contextual completeness
-   - Determine answer feasibility
-
-2. Quality Criteria:
-   - Technical accuracy
-   - Depth of explanation
-   - Contextual relevance
-   - Clarity of communication
-
-3. Response Strategy:
-   a) If Answer Meets Criteria:
-      - Analyse the answer and check whether the answer need to be synthesised. 
-      - If needed, synthesise all responses into a comprehensive final response by incorporating the key information from the original sources
-      - If you are providing a summary, simply extract the summary from the original source by keeping all information within the original summary. Make sure the summary is comprehensive containing all details from the original source.
-      - Provide:
-        * Relevant background information
-        * Practical examples
-        * Potential use cases
-        * Step-by-step explanation
-      - Return your final response for user question in a given format {format_instructions}.
-
-   b) If Answer Insufficient:
-      - Respond "not satisfactory"
-      - Specify:
-        * Missing information
-        * Areas needing clarification
-
-4. Output Formats:
-   - Satisfactory: Comprehensive technical response using format <FORMAT>{format_instructions}</FORMAT>. Do not start your answer by saying: Here is the JSON instance.
-   - Unsatisfactory: Explicit improvement guidance
-
-Core Principle: Deliver maximum insight with precision and technical depth by synthesising responses and keeping original summaries intact.
-User question:<Question>{question}</Question>."""
+NEW_ROUTE_RETRIEVAL_SYSTEM_PROMPT = """
+Answer the question using only the context provided. The context may include results from document agent, external document agent, or summarisation agent.
+   To answer the question:
+   - check if the question is a single question or not. If the question consists of multiple sub-questions, answer each-sub question separately.
+   - always synthesise the raw results from the document agent and external document agent.
+   - If the question requests a summary of the document, you must include the entire summary generated by the summarisation agent, word-for-word, without any redactions or modifications.
+   
+   Here are few example question and expected outputs:
+   - Question: provide a summary of the document and explain the definiton of international trade.
+   - Expected answer: [Include the full summary from the summarisation agent here, verbatim] Then, explain the definition of international trade by synthesizing the results from the document agent.
+
+When including the summary from the summarisation agent, do not use any placeholders or redactions. Provide the complete summary exactly as generated by the summarisation agent.
+
+Output Formats:
+
+   Comprehensive technical response using format {format_instructions}. Do not start your answer by saying: Here is the JSON instance.
+
+"""
+NEW_ROUTE_RETRIEVAL_SYSTEM_PROMPT_NOCITATIONS = """
+Answer the question using only the context provided. The context may include results from document agent, external document agent, or summarisation agent.
+   To answer the question:
+   - check if the question is a single question or not. If the question consists of multiple sub-questions, answer each-sub question separately.
+   - always synthesise the raw results from the document agent and external document agent.
+   - If the question requests a summary of the document, you must include the entire summary generated by the summarisation agent, word-for-word, without any redactions or modifications.
+   
+   Here are few example question and expected outputs:
+   - Question: provide a summary of the document and explain the definiton of international trade.
+   - Expected answer: [Include the full summary from the summarisation agent here, verbatim] Then, explain the definition of international trade by synthesizing the results from the document agent.
+
+When including the summary from the summarisation agent, do not use any placeholders or redactions. Provide the complete summary exactly as generated by the summarisation agent.
+
+"""
 
 AGENTIC_RETRIEVAL_SYSTEM_PROMPT = (
     "You are an advanced problem-solving assistant. Your primary goal is to carefully "
@@ -145,6 +140,8 @@
     "is to provide the best possible answer with the resources available.\n\n"
     "Remember: While your priority is to answer the question, sometimes the best assistance involves "
     "guiding the user in providing the information needed for a complete solution."
+    "6. **Response Format**: When providing the answer, ALWAYS use this exact JSON schema:\n"
+    "{format_instructions}"
 )
 
 SELF_ROUTE_SYSTEM_PROMPT = """
@@ -265,6 +262,10 @@
 
 AGENTIC_RETRIEVAL_QUESTION_PROMPT = "{question}"
 
+NEW_ROUTE_RETRIEVAL_QUESTION_PROMPT = (
+    "Question: {question}. \n\n <Context>: \n\n {agents_results} \n\n <\Context> \n\n Answer: let's think step by step"
+)
+
 AGENTIC_GIVE_UP_QUESTION_PROMPT = "{question}"
 
 CHAT_MAP_QUESTION_PROMPT = "Question: {question}. \n Documents: \n {formatted_documents} \n\n Answer: "
@@ -342,7 +343,7 @@
 
 When creating your execution plan, you have access to the following specialised agents:
 
-1. **Document_Agent**: Retrieves, synthesises, and summarises information from user's uploaded documents.
+1. **Document_Agent**: Retrieves information from user's uploaded documents. It does NOT summarise the document.
 2. **External_Data_Agent**: Retrieves information from external data sources including Wikipedia, Gov.UK, and legislation.gov.uk.
 3. **Summarisation_Agent**: Summarises entire user's uploaded documents. It does not summarise outputs from other agents. 
 
diff --git a/redbox-core/redbox/models/settings.py b/redbox-core/redbox/models/settings.py
index 9e00cb28b..b3f8d6984 100644
--- a/redbox-core/redbox/models/settings.py
+++ b/redbox-core/redbox/models/settings.py
@@ -136,6 +136,7 @@ class Settings(BaseSettings):
     model_config = SettingsConfigDict(env_file=".env", env_nested_delimiter="__", extra="allow", frozen=True)
 
     enable_metadata_extraction: bool = env.bool("ENABLE_METADATA_EXTRACTION", default=True)
+    enable_citation_newroute: bool = env.bool("ENABLE_CITATION_NEWROUTE", default=True)
 
     ## Prompts
     metadata_prompt: tuple = (
diff --git a/redbox-core/redbox/transform.py b/redbox-core/redbox/transform.py
index 1d8d63a9d..7a6ffa107 100644
--- a/redbox-core/redbox/transform.py
+++ b/redbox-core/redbox/transform.py
@@ -5,7 +5,7 @@
 
 from langchain_core.callbacks.manager import dispatch_custom_event
 from langchain_core.documents import Document
-from langchain_core.messages import AIMessage
+from langchain_core.messages import AIMessage, AnyMessage
 from langchain_core.runnables import RunnableLambda
 
 from redbox.models.chain import DocumentMapping, DocumentState, LLMCallMetadata, RedboxState, RequestMetadata
@@ -274,3 +274,13 @@ def max_score(group: list[Document]) -> float:
 
     # Step 4: Flatten the list of blocks back into a single list
     return list(itertools.chain.from_iterable(all_sorted_blocks_by_max_score))
+
+
+def combine_agents_state(agents_results: list[AnyMessage]):
+    """Combine a list of agent results into a string."""
+    if not agents_results:
+        return []
+    print("agents_results")
+    print(agents_results)
+    flatten_agent_results = "\n\n".join([msg.content for msg in agents_results])
+    return flatten_agent_results


Reviews:
user: login: saisakul
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T16:26:30Z
user: login: nora-errouhly
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-17T09:25:54Z
user: login: nora-errouhly
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-17T09:27:55Z
user: login: nora-errouhly
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-17T11:35:13Z
user: login: ugo-daintta
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-17T12:17:03Z

Review comments:
diff_hunk: @@ -72,16 +72,18 @@ def _chat_prompt_from_messages(state: RedboxState) -> Runnable:
             | _additional_variables
         )
 
-        return ChatPromptTemplate(
+        chatprompt = ChatPromptTemplate(
             messages=(
                 [("system", system_prompt_message)]
                 + [(msg["role"], msg["text"]) for msg in truncated_history]
                 + [MessagesPlaceholder("messages")]
-                + [task_question_prompt + "\n\n{format_instructions}"]
path: redbox-core/redbox/chains/runnables.py
user: login: saisakul
body: this affects other routes, please modify prompts for other routes to include format_instructions.
created_at: 2025-04-16T16:26:30Z
author_association: COLLABORATOR
diff_hunk: @@ -72,16 +72,18 @@ def _chat_prompt_from_messages(state: RedboxState) -> Runnable:
             | _additional_variables
         )
 
-        return ChatPromptTemplate(
+        chatprompt = ChatPromptTemplate(
             messages=(
                 [("system", system_prompt_message)]
                 + [(msg["role"], msg["text"]) for msg in truncated_history]
                 + [MessagesPlaceholder("messages")]
-                + [task_question_prompt + "\n\n{format_instructions}"]
path: redbox-core/redbox/chains/runnables.py
user: login: nora-errouhly
body: can you specify which route/ prompts it may affects. I can then append the prompt with format instructions. any prompt modification should be in prompts.py, not in runnables.py
created_at: 2025-04-17T09:25:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -72,16 +72,18 @@ def _chat_prompt_from_messages(state: RedboxState) -> Runnable:
             | _additional_variables
         )
 
-        return ChatPromptTemplate(
+        chatprompt = ChatPromptTemplate(
             messages=(
                 [("system", system_prompt_message)]
                 + [(msg["role"], msg["text"]) for msg in truncated_history]
                 + [MessagesPlaceholder("messages")]
-                + [task_question_prompt + "\n\n{format_instructions}"]
path: redbox-core/redbox/chains/runnables.py
user: login: nora-errouhly
body: I can see format_instructions are in many system prompts. if they are defined in system prompts, they should not be required in question prompts
created_at: 2025-04-17T09:27:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -72,16 +72,18 @@ def _chat_prompt_from_messages(state: RedboxState) -> Runnable:
             | _additional_variables
         )
 
-        return ChatPromptTemplate(
+        chatprompt = ChatPromptTemplate(
             messages=(
                 [("system", system_prompt_message)]
                 + [(msg["role"], msg["text"]) for msg in truncated_history]
                 + [MessagesPlaceholder("messages")]
-                + [task_question_prompt + "\n\n{format_instructions}"]
path: redbox-core/redbox/chains/runnables.py
user: login: nora-errouhly
body: I added format instruction to one of the prompts (AGENTIC_GIVE_UP_SYSTEM_PROMPT) . all other system prompts requiring citations are fine as they already contain format_instructions
created_at: 2025-04-17T11:35:12Z
author_association: CONTRIBUTOR

PR comments:
user: login: saisakul
created_at: 2025-04-17T08:11:45Z
body: I don't think we should create a separate route for response without citations. It should use the same route but can handle different settings.
author_association: COLLABORATOR
user: login: nora-errouhly
created_at: 2025-04-17T12:03:49Z
body: > I don't think we should create a separate route for response without citations. It should use the same route but can handle different settings.

I did not add another route. it is still within same @newroute
author_association: CONTRIBUTOR