fix: address reviewer feedback and resolve state management issues

Author: heliamoh

src/agent/profiles/base.py

@@ -65,9 +65,8 @@ async def preprocess(self, state: BaseState, config: RunnableConfig) -> BaseStat
 
     async def postprocess(self, state: BaseState, config: RunnableConfig) -> BaseState:
         search_results: list[WebSearchResult] = []
-        if (
-            state.get("safety") == SAFETY_SAFE
-            and config["configurable"]["enable_postprocess"]
+        if state.get("safety") == SAFETY_SAFE and config.get("configurable", {}).get(
+            "enable_postprocess", False
         ):
             result: SearchState = await self.search_workflow.ainvoke(
                 SearchState(

src/agent/profiles/cross_database.py

@@ -15,16 +15,12 @@
     create_uniprot_rewriter_w_reactome
 from agent.tasks.cross_database.summarize_reactome_uniprot import \
     create_reactome_uniprot_summarizer
-from agent.tasks.detect_language import create_language_detector
 from agent.tasks.safety_checker import SafetyCheck, create_safety_checker
 from retrievers.reactome.rag import create_reactome_rag
 from retrievers.uniprot.rag import create_uniprot_rag
 
 
 class CrossDatabaseState(BaseState):
-    safety: str  # LLM-assessed safety level of the user input
-    query_language: str  # language of the user input
-
     reactome_query: str  # LLM-generated query for Reactome
     reactome_answer: str  # LLM-generated answer from Reactome
     reactome_completeness: str  # LLM-assessed completeness of the Reactome answer
@@ -48,7 +44,6 @@ def __init__(
 
         self.safety_checker = create_safety_checker(llm)
         self.completeness_checker = create_completeness_grader(llm)
-        self.detect_language = create_language_detector(llm)
         self.write_reactome_query = create_reactome_rewriter_w_uniprot(llm)
         self.write_uniprot_query = create_uniprot_rewriter_w_reactome(llm)
         self.summarize_final_answer = create_reactome_uniprot_summarizer(
@@ -60,7 +55,6 @@ def __init__(
         # Set up nodes
         state_graph.add_node("check_question_safety", self.check_question_safety)
         state_graph.add_node("preprocess_question", self.preprocess)
-        state_graph.add_node("identify_query_language", self.identify_query_language)
         state_graph.add_node("conduct_research", self.conduct_research)
         state_graph.add_node("generate_reactome_answer", self.generate_reactome_answer)
         state_graph.add_node("rewrite_reactome_query", self.rewrite_reactome_query)
@@ -74,7 +68,6 @@ def __init__(
         state_graph.add_node("postprocess", self.postprocess)
         # Set up edges
         state_graph.set_entry_point("preprocess_question")
-        state_graph.add_edge("preprocess_question", "identify_query_language")
         state_graph.add_edge("preprocess_question", "check_question_safety")
         state_graph.add_conditional_edges(
             "check_question_safety",
@@ -112,7 +105,7 @@ async def check_question_safety(
             config,
         )
         if result.binary_score == "No":
-            inappropriate_input = f"This is the user's question and it is NOT appropriate for you to answer: {state["user_input"]}. \n\n explain that you are unable to answer the question but you can answer questions about topics related to the Reactome Pathway Knowledgebase or UniProt Knowledgebas."
+            inappropriate_input = f"This is the user's question and it is NOT appropriate for you to answer: {state["user_input"]}. \n\n explain that you are unable to answer the question but you can answer questions about topics related to the Reactome Pathway Knowledgebase or UniProt Knowledgebase."
             return CrossDatabaseState(
                 safety=result.binary_score,
                 user_input=inappropriate_input,
@@ -130,14 +123,6 @@ async def proceed_with_research(
         else:
             return "Finish"
 
-    async def identify_query_language(
-        self, state: CrossDatabaseState, config: RunnableConfig
-    ) -> CrossDatabaseState:
-        query_language: str = await self.detect_language.ainvoke(
-            {"user_input": state["user_input"]}, config
-        )
-        return CrossDatabaseState(query_language=query_language)
-
     async def conduct_research(
         self, state: CrossDatabaseState, config: RunnableConfig
     ) -> CrossDatabaseState:
@@ -256,7 +241,6 @@ async def generate_final_response(
         final_response: str = await self.summarize_final_answer.ainvoke(
             {
                 "input": state["rephrased_input"],
-                "query_language": state["query_language"],
                 "reactome_answer": state["reactome_answer"],
                 "uniprot_answer": state["uniprot_answer"],
             },

src/agent/profiles/react_to_me.py

@@ -6,8 +6,8 @@
 from langchain_core.runnables import Runnable, RunnableConfig
 from langgraph.graph.state import StateGraph
 
-from agent.profiles.base import (DEFAULT_LANGUAGE, SAFETY_SAFE, SAFETY_UNSAFE,
-                                 BaseGraphBuilder, BaseState)
+from agent.profiles.base import (SAFETY_SAFE, SAFETY_UNSAFE, BaseGraphBuilder,
+                                 BaseState)
 from agent.tasks.unsafe_answer import create_unsafe_answer_generator
 from retrievers.reactome.rag import create_reactome_rag
 
@@ -60,7 +60,6 @@ async def generate_unsafe_response(
     ) -> ReactToMeState:
         final_answer_message = await self.unsafe_answer_generator.ainvoke(
             {
-                "language": state.get("detected_language", DEFAULT_LANGUAGE),
                 "user_input": state.get("rephrased_input", state["user_input"]),
                 "reason_unsafe": state.get("reason_unsafe", ""),
             },
@@ -86,7 +85,6 @@ async def generate_unsafe_response(
         )
 
         return ReactToMeState(
-            **state,
             chat_history=history,
             answer=final_answer,
             safety=SAFETY_UNSAFE,
@@ -99,7 +97,6 @@ async def call_model(
         result: dict[str, Any] = await self.reactome_rag.ainvoke(
             {
                 "input": state["rephrased_input"],
-                "expanded_queries": state.get("expanded_queries", []),
                 "chat_history": (
                     state.get("chat_history")
                     if state.get("chat_history")
@@ -116,7 +113,6 @@ async def call_model(
             ]
         )
         return ReactToMeState(
-            **state,
             chat_history=history,
             answer=result["answer"],
         )

src/agent/tasks/unsafe_answer.py

@@ -18,17 +18,15 @@ def create_unsafe_answer_generator(llm: BaseChatModel) -> Runnable:
 
 You have advanced training in scientific ethics, dual-use research concerns, and responsible AI use.
 
-You will receive three inputs: 
+You will receive two inputs: 
 1. The user's question.
 2. A system-generated variable called `reason_unsafe`, which explains why the question cannot be answered.
-3. The user's preferred language (as a language code or name).
 
 Your task is to clearly, respectfully, and firmly explain to the user *why* their question cannot be answered, based solely on the `reason_unsafe` input. Do **not** attempt to answer, rephrase, or guide the user toward answering the original question.
 
 You must:
-- Respond in the user’s preferred language.
 - Politely explain the refusal, grounded in the `reason_unsafe`.
-- Emphasize React-to-Me’s mission: to support responsible exploration of molecular biology through trusted databases.
+- Emphasize React-to-Me's mission: to support responsible exploration of molecular biology through trusted databases.
 - Suggest examples of appropriate topics (e.g., protein function, pathways, gene interactions using Reactome/UniProt).
 
 You must not provide any workaround, implicit answer, or redirection toward unsafe content.
@@ -40,7 +38,7 @@ def create_unsafe_answer_generator(llm: BaseChatModel) -> Runnable:
             ("system", system_prompt),
             (
                 "user",
-                "Language:{language}\n\nQuestion:{user_input}\n\n Reason for unsafe or out of scope: {reason_unsafe}",
+                "Question:{user_input}\n\n Reason for unsafe or out of scope: {reason_unsafe}",
             ),
         ]
     )

src/tools/preprocessing/state.py

@@ -14,5 +14,3 @@ class PreprocessingState(TypedDict, total=False):
     rephrased_input: str
     safety: str
     reason_unsafe: str
-    expanded_queries: list[str]
-    detected_language: str