chore(graph): Remove obsolete linear nodes and legacy tests

backend/tests/test_coder.py

@@ -1,62 +0,0 @@
-import pytest
-from unittest.mock import MagicMock, patch
-from ea_chatbot.graph.nodes.coder import coder_node
-from ea_chatbot.graph.nodes.error_corrector import error_corrector_node
-
-@pytest.fixture
-def mock_state():
-    return {
-        "messages": [],
-        "question": "Show me results for New Jersey",
-        "plan": "Step 1: Load data\nStep 2: Filter by NJ",
-        "code": None,
-        "error": None,
-        "plots": [],
-        "dfs": {},
-        "next_action": "plan"
-    }
-
-@patch("ea_chatbot.graph.nodes.coder.get_llm_model")
-@patch("ea_chatbot.utils.database_inspection.get_data_summary")
-def test_coder_node(mock_get_summary, mock_get_llm, mock_state):
-    """Test coder node generates code from plan."""
-    mock_get_summary.return_value = "Column: Name, Type: text"
-    
-    mock_llm = MagicMock()
-    mock_get_llm.return_value = mock_llm
-    
-    from ea_chatbot.schemas import CodeGenerationResponse
-    mock_response = CodeGenerationResponse(
-        code="import pandas as pd\nprint('Hello')",
-        explanation="Generated code"
-    )
-    mock_llm.with_structured_output.return_value.invoke.return_value = mock_response
-    
-    result = coder_node(mock_state)
-    
-    assert "code" in result
-    assert "import pandas as pd" in result["code"]
-    assert "error" in result
-    assert result["error"] is None
-
-@patch("ea_chatbot.graph.nodes.error_corrector.get_llm_model")
-def test_error_corrector_node(mock_get_llm, mock_state):
-    """Test error corrector node fixes code."""
-    mock_state["code"] = "import pandas as pd\nprint(undefined_var)"
-    mock_state["error"] = "NameError: name 'undefined_var' is not defined"
-    
-    mock_llm = MagicMock()
-    mock_get_llm.return_value = mock_llm
-    
-    from ea_chatbot.schemas import CodeGenerationResponse
-    mock_response = CodeGenerationResponse(
-        code="import pandas as pd\nprint('Defined')",
-        explanation="Fixed variable"
-    )
-    mock_llm.with_structured_output.return_value.invoke.return_value = mock_response
-    
-    result = error_corrector_node(mock_state)
-    
-    assert "code" in result
-    assert "print('Defined')" in result["code"]
-    assert result["error"] is None

backend/tests/test_executor.py

@@ -1,122 +0,0 @@
-import pytest
-import pandas as pd
-from unittest.mock import MagicMock, patch
-from matplotlib.figure import Figure
-from ea_chatbot.graph.nodes.executor import executor_node
-
-@pytest.fixture
-def mock_settings():
-    with patch("ea_chatbot.graph.nodes.executor.Settings") as MockSettings:
-        mock_settings_instance = MagicMock()
-        mock_settings_instance.db_host = "localhost"
-        mock_settings_instance.db_port = 5432
-        mock_settings_instance.db_user = "user"
-        mock_settings_instance.db_pswd = "pass"
-        mock_settings_instance.db_name = "test_db"
-        mock_settings_instance.db_table = "test_table"
-        MockSettings.return_value = mock_settings_instance
-        yield mock_settings_instance
-
-@pytest.fixture
-def mock_db_client():
-    with patch("ea_chatbot.graph.nodes.executor.DBClient") as MockDBClient:
-        mock_client_instance = MagicMock()
-        MockDBClient.return_value = mock_client_instance
-        yield mock_client_instance
-
-def test_executor_node_success_simple_print(mock_settings, mock_db_client):
-    """Test executing simple code that prints to stdout."""
-    state = {
-        "code": "print('Hello, World!')",
-        "question": "test",
-        "messages": []
-    }
-    
-    result = executor_node(state)
-    
-    assert "code_output" in result
-    assert "Hello, World!" in result["code_output"]
-    assert result["error"] is None
-    assert result["plots"] == []
-    assert result["dfs"] == {}
-
-def test_executor_node_success_dataframe(mock_settings, mock_db_client):
-    """Test executing code that creates a DataFrame."""
-    code = """
-import pandas as pd
-df = pd.DataFrame({'a': [1, 2], 'b': [3, 4]})
-print(df)
-"""
-    state = {
-        "code": code,
-        "question": "test",
-        "messages": []
-    }
-    
-    result = executor_node(state)
-    
-    assert "code_output" in result
-    assert "a  b" in result["code_output"] # Check part of DF string representation
-    assert "dfs" in result
-    assert "df" in result["dfs"]
-    assert isinstance(result["dfs"]["df"], pd.DataFrame)
-
-def test_executor_node_success_plot(mock_settings, mock_db_client):
-    """Test executing code that generates a plot."""
-    code = """
-import matplotlib.pyplot as plt
-fig = plt.figure()
-plots.append(fig)
-print('Plot generated')
-"""
-    state = {
-        "code": code,
-        "question": "test",
-        "messages": []
-    }
-    
-    result = executor_node(state)
-    
-    assert "Plot generated" in result["code_output"]
-    assert "plots" in result
-    assert len(result["plots"]) == 1
-    assert isinstance(result["plots"][0], Figure)
-
-def test_executor_node_error_syntax(mock_settings, mock_db_client):
-    """Test executing code with a syntax error."""
-    state = {
-        "code": "print('Hello World", # Missing closing quote
-        "question": "test",
-        "messages": []
-    }
-    
-    result = executor_node(state)
-    
-    assert result["error"] is not None
-    assert "SyntaxError" in result["error"]
-
-def test_executor_node_error_runtime(mock_settings, mock_db_client):
-    """Test executing code with a runtime error."""
-    state = {
-        "code": "print(1 / 0)",
-        "question": "test",
-        "messages": []
-    }
-    
-    result = executor_node(state)
-    
-    assert result["error"] is not None
-    assert "ZeroDivisionError" in result["error"]
-
-def test_executor_node_no_code(mock_settings, mock_db_client):
-    """Test handling when no code is provided."""
-    state = {
-        "code": None,
-        "question": "test",
-        "messages": []
-    }
-    
-    result = executor_node(state)
-    
-    assert "error" in result
-    assert "No code provided" in result["error"]

backend/tests/test_executor_vfs.py

@@ -1,36 +0,0 @@
-from ea_chatbot.graph.nodes.executor import executor_node
-from ea_chatbot.graph.state import AgentState
-
-def test_executor_with_vfs():
-    """Verify that the executor node provides VFS access to the code."""
-    state = AgentState(
-        messages=[],
-        question="test",
-        analysis={},
-        next_action="test",
-        iterations=0,
-        checklist=[],
-        current_step=0,
-        vfs={},
-        plots=[],
-        dfs={}
-    )
-    
-    # Code that uses the 'vfs' helper
-    code = """
-vfs.write("output.txt", "Execution Result", metadata={"type": "text"})
-print("VFS Write Complete")
-"""
-    state["code"] = code
-    
-    result = executor_node(state)
-    
-    # Check if the execution was successful
-    assert result["error"] is None
-    assert "VFS Write Complete" in result["code_output"]
-    
-    # Verify that the VFS state was updated
-    # Note: executor_node returns a dict of updates, which should include the updated 'vfs'
-    assert "vfs" in result
-    assert "output.txt" in result["vfs"]
-    assert result["vfs"]["output.txt"]["content"] == "Execution Result"

backend/tests/test_researcher.py

@@ -1,34 +0,0 @@
-import pytest
-from unittest.mock import MagicMock, patch
-from langchain_core.messages import AIMessage
-from langchain_openai import ChatOpenAI
-from ea_chatbot.graph.nodes.researcher import researcher_node
-
-@pytest.fixture
-def mock_llm():
-    with patch("ea_chatbot.graph.nodes.researcher.get_llm_model") as mock_get_llm:
-        mock_llm_instance = MagicMock(spec=ChatOpenAI)
-        mock_get_llm.return_value = mock_llm_instance
-        yield mock_llm_instance
-
-def test_researcher_node_success(mock_llm):
-    """Test that researcher_node invokes LLM with web_search tool and returns messages."""
-    state = {
-        "question": "What is the capital of France?",
-        "messages": []
-    }
-    
-    mock_llm_with_tools = MagicMock()
-    mock_llm.bind_tools.return_value = mock_llm_with_tools
-    mock_llm_with_tools.invoke.return_value = AIMessage(content="The capital of France is Paris.")
-    
-    result = researcher_node(state)
-    
-    assert mock_llm.bind_tools.called
-    # Check that it was called with web_search
-    args, kwargs = mock_llm.bind_tools.call_args
-    assert {"type": "web_search"} in args[0]
-    
-    assert mock_llm_with_tools.invoke.called
-    assert "messages" in result
-    assert result["messages"][0].content == "The capital of France is Paris."

backend/tests/test_researcher_search_tools.py

@@ -1,62 +0,0 @@
-import pytest
-from unittest.mock import MagicMock, patch
-from langchain_core.messages import AIMessage
-from langchain_openai import ChatOpenAI
-from langchain_google_genai import ChatGoogleGenerativeAI
-from ea_chatbot.graph.nodes.researcher import researcher_node
-
-@pytest.fixture
-def base_state():
-    return {
-        "question": "Who won the 2024 election?",
-        "messages": [],
-        "summary": ""
-    }
-
-@patch("ea_chatbot.graph.nodes.researcher.get_llm_model")
-def test_researcher_binds_openai_search(mock_get_llm, base_state):
-    """Test that OpenAI LLM binds 'web_search' tool."""
-    mock_llm = MagicMock(spec=ChatOpenAI)
-    mock_get_llm.return_value = mock_llm
-    
-    mock_llm_with_tools = MagicMock()
-    mock_llm.bind_tools.return_value = mock_llm_with_tools
-    mock_llm_with_tools.invoke.return_value = AIMessage(content="OpenAI Search Result")
-    
-    result = researcher_node(base_state)
-    
-    # Verify bind_tools called with correct OpenAI tool
-    mock_llm.bind_tools.assert_called_once_with([{"type": "web_search"}])
-    assert result["messages"][0].content == "OpenAI Search Result"
-
-@patch("ea_chatbot.graph.nodes.researcher.get_llm_model")
-def test_researcher_binds_google_search(mock_get_llm, base_state):
-    """Test that Google LLM binds 'google_search' tool."""
-    mock_llm = MagicMock(spec=ChatGoogleGenerativeAI)
-    mock_get_llm.return_value = mock_llm
-    
-    mock_llm_with_tools = MagicMock()
-    mock_llm.bind_tools.return_value = mock_llm_with_tools
-    mock_llm_with_tools.invoke.return_value = AIMessage(content="Google Search Result")
-    
-    result = researcher_node(base_state)
-    
-    # Verify bind_tools called with correct Google tool
-    mock_llm.bind_tools.assert_called_once_with([{"google_search": {}}])
-    assert result["messages"][0].content == "Google Search Result"
-
-@patch("ea_chatbot.graph.nodes.researcher.get_llm_model")
-def test_researcher_fallback_on_bind_error(mock_get_llm, base_state):
-    """Test that researcher falls back to basic LLM if bind_tools fails."""
-    mock_llm = MagicMock(spec=ChatOpenAI)
-    mock_get_llm.return_value = mock_llm
-    
-    # Simulate bind_tools failing (e.g. model doesn't support it)
-    mock_llm.bind_tools.side_effect = Exception("Not supported")
-    mock_llm.invoke.return_value = AIMessage(content="Basic Result")
-    
-    result = researcher_node(base_state)
-    
-    # Should still succeed using the base LLM
-    assert result["messages"][0].content == "Basic Result"
-    mock_llm.invoke.assert_called_once()

backend/tests/test_summarizer.py

@@ -1,47 +0,0 @@
-import pytest
-from unittest.mock import MagicMock, patch
-from langchain_core.messages import AIMessage
-from ea_chatbot.graph.nodes.summarizer import summarizer_node
-
-@pytest.fixture
-def mock_llm():
-    with patch("ea_chatbot.graph.nodes.summarizer.get_llm_model") as mock_get_llm:
-        mock_llm_instance = MagicMock()
-        mock_get_llm.return_value = mock_llm_instance
-        yield mock_llm_instance
-
-def test_summarizer_node_success(mock_llm):
-    """Test that summarizer_node invokes LLM with correct inputs and returns messages."""
-    state = {
-        "question": "What is the total count?",
-        "plan": "1. Run query\n2. Sum results",
-        "code_output": "The total is 100",
-        "messages": []
-    }
-    
-    mock_llm.invoke.return_value = AIMessage(content="The final answer is 100.")
-    
-    result = summarizer_node(state)
-    
-    # Verify LLM was called
-    assert mock_llm.invoke.called
-    
-    # Verify result structure
-    assert "messages" in result
-    assert len(result["messages"]) == 1
-    assert isinstance(result["messages"][0], AIMessage)
-    assert result["messages"][0].content == "The final answer is 100."
-
-def test_summarizer_node_empty_state(mock_llm):
-    """Test handling of empty or minimal state."""
-    state = {
-        "question": "Empty?",
-        "messages": []
-    }
-    
-    mock_llm.invoke.return_value = AIMessage(content="No data provided.")
-    
-    result = summarizer_node(state)
-    
-    assert "messages" in result
-    assert result["messages"][0].content == "No data provided."

backend/tests/test_workflow_e2e.py

@@ -12,7 +12,8 @@ def mock_llms():
          patch("ea_chatbot.graph.workers.data_analyst.nodes.coder.get_llm_model") as mock_coder, \
          patch("ea_chatbot.graph.workers.data_analyst.nodes.summarizer.get_llm_model") as mock_worker_summarizer, \
          patch("ea_chatbot.graph.nodes.synthesizer.get_llm_model") as mock_synthesizer, \
-         patch("ea_chatbot.graph.nodes.researcher.get_llm_model") as mock_researcher, \
+         patch("ea_chatbot.graph.workers.researcher.nodes.searcher.get_llm_model") as mock_researcher, \
+         patch("ea_chatbot.graph.workers.researcher.nodes.summarizer.get_llm_model") as mock_res_summarizer, \
          patch("ea_chatbot.graph.nodes.reflector.get_llm_model") as mock_reflector:
         yield {
             "qa": mock_qa,
@@ -21,6 +22,7 @@ def mock_llms():
             "worker_summarizer": mock_worker_summarizer,
             "synthesizer": mock_synthesizer,
             "researcher": mock_researcher,
+            "res_summarizer": mock_res_summarizer,
             "reflector": mock_reflector
         }
 
@@ -113,17 +115,22 @@ def test_workflow_research_flow(mock_llms):
         checklist=[ChecklistTask(task="Search Web", worker="researcher")]
     )
     
-    # 3. Mock Researcher
+    # 3. Mock Researcher Searcher
     mock_res_instance = MagicMock()
     mock_llms["researcher"].return_value = mock_res_instance
     mock_res_instance.invoke.return_value = AIMessage(content="Research Result")
     
-    # 4. Mock Reflector
+    # 4. Mock Researcher Summarizer
+    mock_rs_instance = MagicMock()
+    mock_llms["res_summarizer"].return_value = mock_rs_instance
+    mock_rs_instance.invoke.return_value = AIMessage(content="Researcher Summary")
+
+    # 5. Mock Reflector
     mock_reflector_instance = MagicMock()
     mock_llms["reflector"].return_value = mock_reflector_instance
     mock_reflector_instance.with_structured_output.return_value.invoke.return_value = MagicMock(satisfied=True, reasoning="Good.")
 
-    # 5. Mock Synthesizer
+    # 6. Mock Synthesizer
     mock_syn_instance = MagicMock()
     mock_llms["synthesizer"].return_value = mock_syn_instance
     mock_syn_instance.invoke.return_value = AIMessage(content="Final Research Summary")