fix: address reviewer import and streaming feedback

Author: heliamoh

src/agent/profiles/react_to_me.py

@@ -25,9 +25,7 @@ def __init__(
         super().__init__(llm, embedding)
 
         # Create runnables (tasks & tools)
-        streaming_llm = llm.model_copy(update={"streaming": True})
-
-        self.unsafe_answer_generator = create_unsafe_answer_generator(streaming_llm)
+        self.unsafe_answer_generator = create_unsafe_answer_generator(llm)
         self.reactome_rag: Runnable = create_reactome_rag(
             llm, embedding, streaming=True
         )
@@ -75,21 +73,25 @@ async def generate_unsafe_response(
             else str(final_answer_message)
         )
 
-        updated_state = dict(state)
-        updated_state.update(
-            chat_history=[
+        history = list(state.get("chat_history", []))
+        history.extend(
+            [
                 HumanMessage(state["user_input"]),
                 (
                     final_answer_message
                     if hasattr(final_answer_message, "content")
                     else AIMessage(final_answer)
                 ),
-            ],
+            ]
+        )
+
+        return ReactToMeState(
+            **state,
+            chat_history=history,
             answer=final_answer,
             safety=SAFETY_UNSAFE,
             additional_content={"search_results": []},
         )
-        return ReactToMeState(**updated_state)
 
     async def call_model(
         self, state: ReactToMeState, config: RunnableConfig
@@ -106,15 +108,18 @@ async def call_model(
             },
             config,
         )
-        updated_state = dict(state)
-        updated_state.update(
-            chat_history=[
+        history = list(state.get("chat_history", []))
+        history.extend(
+            [
                 HumanMessage(state["user_input"]),
                 AIMessage(result["answer"]),
-            ],
+            ]
+        )
+        return ReactToMeState(
+            **state,
+            chat_history=history,
             answer=result["answer"],
         )
-        return ReactToMeState(**updated_state)
 
 
 def create_reactome_graph(

src/agent/tasks/unsafe_answer.py

@@ -33,6 +33,8 @@ def create_unsafe_answer_generator(llm: BaseChatModel) -> Runnable:
 
 You must not provide any workaround, implicit answer, or redirection toward unsafe content.
 """
+    streaming_llm = llm.model_copy(update={"streaming": True})
+
     prompt = ChatPromptTemplate.from_messages(
         [
             ("system", system_prompt),
@@ -43,4 +45,4 @@ def create_unsafe_answer_generator(llm: BaseChatModel) -> Runnable:
         ]
     )
 
-    return prompt | llm
+    return prompt | streaming_llm

src/tools/preprocessing/__init__.py

@@ -2,5 +2,7 @@
 Preprocessing utilities with reusable workflows and state definitions.
 """
 
-from .state import PreprocessingState  # noqa: F401
-from .workflow import create_preprocessing_workflow  # noqa: F401
+from tools.preprocessing.state import PreprocessingState
+from tools.preprocessing.workflow import create_preprocessing_workflow
+
+__all__ = ["PreprocessingState", "create_preprocessing_workflow"]