refactor: remove legacy LangChain chains and migrate to streaming agent

lguibr · lguibr · commit f01ec7c6c80e · 2026-02-04T14:24:39.000-03:00
- Delete `commitai/chains.py` as part of the transition to the new agent-based architecture.
- Update the CLI to utilize the streaming interface for commit message generation.
- Migrate TODO scanning and context summarization logic to the agent module.
- Add comprehensive unit tests for the migrated logic and streaming pipeline in `tests/test_agent_v3.py`.
- Update existing CLI tests to reflect the switch from `invoke()` to `stream()`.
diff --git a/commitai/chains.py b/commitai/chains.py
diff --git a/commitai/cli.py b/commitai/cli.py
@@ -265,9 +265,7 @@ def generate_message(  # noqa: C901
             inputs["template"] = final_template_content
 
         # Invoke the Agent Pipeline (which now returns a generator)
-        stream_gen = agent_pipeline.invoke(inputs)
-
-        # Use UI to handle streaming visualization
+        stream_gen = agent_pipeline.stream(inputs)
         commit_message = ui.stream_response(stream_gen)
 
     except Exception as e:
diff --git a/tests/test_agent_v3.py b/tests/test_agent_v3.py
@@ -1,20 +1,17 @@
 from unittest.mock import MagicMock, patch
 
-from commitai.agent import create_commit_agent
+from langchain_core.messages import AIMessage
+
+from commitai.agent import create_commit_agent, scan_todos, summarize_context
 
 
 def test_agent_v3_initialization():
     """Verify that the V3 agent can be initialized with the new middlewares."""
     mock_llm = MagicMock()
 
-    # We need to mock create_agent because it will try to set up real tools/middlewares
-    # and we don't want to run them or depend on OS/Git state during this unit test.
-    # However, we want to verify that create_agent CALL was correct (args passed).
-
     with patch("commitai.agent.create_agent") as mock_create_agent:
         mock_create_agent.return_value = MagicMock()  # Return a mock graph
 
-        # Also mock the middlewares to check if they were instantiated
         with (
             patch("commitai.agent.SummarizationMiddleware") as mock_summ,
             patch("commitai.agent.FilesystemFileSearchMiddleware") as mock_files,
@@ -38,3 +35,103 @@ def test_agent_v3_initialization():
             assert len(kwargs["middleware"]) == 5
 
             assert agent_runnable is not None
+
+
+def test_scan_todos():
+    """Test TODO scanning logic."""
+    diff_with_todo = "+ # TODO: Fix this later\n+ function foo() {\n+ # FIXME: Old bug"
+    result = scan_todos(diff_with_todo)
+
+    assert "todos" in result
+    assert "todo_str" in result
+    assert len(result["todos"]) == 2
+    assert "Fix this later" in result["todos"][0]
+    assert "Old bug" in result["todos"][1]
+    # Check for the raw line content as captured
+    assert "- # TODO: Fix this later" in result["todo_str"]
+
+
+def test_scan_todos_none():
+    """Test TODO scanning with no todos."""
+    diff_clean = "+ function foo() {}"
+    result = scan_todos(diff_clean)
+    assert not result["todos"]
+    assert result["todo_str"] == "None"
+
+
+def test_summarize_context_short():
+    """Test summarization (even short diffs trigger LLM now)."""
+    mock_llm = MagicMock()
+    mock_llm.invoke.return_value = AIMessage(content="Summary")
+    diff = "short diff"
+    summary = summarize_context(mock_llm, diff)
+
+    # It returns content directly now
+    assert summary == "Summary"
+    mock_llm.invoke.assert_called_once()
+
+
+def test_summarize_context_long():
+    """Test summarization triggered for long diffs."""
+    mock_llm = MagicMock()
+    mock_llm.invoke.return_value = AIMessage(content="Summary of changes")
+
+    # Create diff > 3000 chars
+    diff = "a" * 3005
+    summary = summarize_context(mock_llm, diff)
+
+    assert summary == "Summary of changes"
+    mock_llm.invoke.assert_called_once()
+
+
+def test_run_pipeline_streaming():
+    """Test the streaming pipeline execution."""
+    mock_llm = MagicMock()
+    # Mock LLM response for summarize_context call inside pipeline
+    mock_llm.invoke.return_value = AIMessage(content="Summary")
+
+    # Mock the agent graph created inside factory
+    mock_graph = MagicMock()
+
+    # Mock stream output from graph
+    final_message = AIMessage(content="feat: new feature")
+    events = [
+        {
+            "messages": [
+                AIMessage(
+                    content="", tool_calls=[{"name": "git_log", "args": {}, "id": "1"}]
+                )
+            ]
+        },
+        {"messages": [final_message]},
+    ]
+    mock_graph.stream.return_value = iter(events)
+
+    with patch("commitai.agent.create_agent", return_value=mock_graph):
+        with (
+            patch("commitai.agent.SummarizationMiddleware"),
+            patch("commitai.agent.FilesystemFileSearchMiddleware"),
+            patch("commitai.agent.ShellToolMiddleware"),
+            patch("commitai.agent.HumanInTheLoopMiddleware"),
+            patch("commitai.agent.LLMToolSelectorMiddleware"),
+        ):
+            pipeline_func = create_commit_agent(mock_llm)
+
+            inputs = {"diff": "some diff", "explanation": "expl"}
+            # Use stream() to ensure we get the generator/iterator properly
+            stream_gen = pipeline_func.stream(inputs)
+
+            results = list(stream_gen)
+
+            # Verify we received dicts
+            for r in results:
+                assert isinstance(r, dict), f"Received non-dict result: {r}"
+
+            # Check for thought event (from tool call)
+            assert any(r["type"] == "thought" for r in results)
+
+            # Check for token events (from final message)
+            token_events = [r for r in results if r["type"] == "token"]
+            assert len(token_events) > 0
+            full_text = "".join(t["content"] for t in token_events)
+            assert "feat: new feature" in full_text
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -60,7 +60,7 @@ def mock_generate_deps(tmp_path):
 
         # Agent Mock (RunnableLambda now)
         mock_agent_runnable = MagicMock()
-        mock_agent_runnable.invoke.return_value = iter(
+        mock_agent_runnable.stream.return_value = iter(
             [{"type": "token", "content": "Generated commit message"}]
         )
         mock_create_agent.return_value = mock_agent_runnable
@@ -130,7 +130,7 @@ def test_generate_default_gemini(mock_generate_deps):
         model="gemini-3-flash-preview",
         google_api_key="fake_google_key",
     )
-    mock_generate_deps["agent_instance"].invoke.assert_called_once()
+    mock_generate_deps["agent_instance"].stream.assert_called_once()
     mock_generate_deps["commit"].assert_called_once_with("Generated commit message")
     mock_generate_deps["ui"].render_header.assert_called_once()
 
@@ -213,7 +213,7 @@ def test_generate_no_staged_changes(mock_generate_deps):
         "⚠️ Warning: No staged changes found. Exiting."
     )
 
-    mock_generate_deps["agent_instance"].invoke.assert_not_called()
+    mock_generate_deps["agent_instance"].stream.assert_not_called()
     mock_generate_deps["commit"].assert_not_called()
 
 
@@ -270,7 +270,7 @@ def test_generate_no_explanation(mock_generate_deps):
     result = runner.invoke(cli, ["generate", "--no-review"])
 
     assert result.exit_code == 0, result.output
-    mock_generate_deps["agent_instance"].invoke.assert_called_once()
+    mock_generate_deps["agent_instance"].stream.assert_called_once()
     mock_generate_deps["commit"].assert_called_once()
 
 
@@ -298,7 +298,7 @@ def getenv_side_effect_with_template(key, default=None):
     assert result.exit_code == 0, result.output
 
     # Verify agent invocation has correct args
-    call_args = mock_generate_deps["agent_instance"].invoke.call_args
+    call_args = mock_generate_deps["agent_instance"].stream.call_args
     assert call_args is not None, "agent invoke was not called"
     invoked_args = call_args[0][0]
     assert invoked_args["explanation"] == "Test explanation"
@@ -324,7 +324,7 @@ def test_generate_with_local_template(mock_get_template, mock_generate_deps):
     assert result.exit_code == 0, result.output
     mock_get_template.assert_called_once()  # Verify get_commit_template was called
     # Verify agent invocation has correct args
-    call_args = mock_generate_deps["agent_instance"].invoke.call_args
+    call_args = mock_generate_deps["agent_instance"].stream.call_args
     assert call_args is not None, "agent invoke was not called"
     invoked_args = call_args[0][0]
     assert invoked_args["explanation"] == "Test explanation"
@@ -353,7 +353,7 @@ def test_generate_with_deprecated_template_option(mock_generate_deps):
     mock_generate_deps["ui"].console.print.assert_any_call(
         "[warning]⚠️ --template/-t is deprecated.[/warning]"
     )
-    mock_generate_deps["agent_instance"].invoke.assert_called_once()
+    mock_generate_deps["agent_instance"].stream.assert_called_once()
     mock_generate_deps["commit"].assert_called_once()
 
 
@@ -443,7 +443,7 @@ def test_generate_google_module_not_installed(mock_generate_deps):
 def test_generate_llm_invoke_error(mock_generate_deps):
     """Test generate command handling error during llm.invoke."""
     runner = CliRunner()
-    mock_generate_deps["agent_instance"].invoke.side_effect = Exception("AI API Error")
+    mock_generate_deps["agent_instance"].stream.side_effect = Exception("AI API Error")
     result = runner.invoke(cli, ["generate", "--no-review", "Test explanation"])
 
     assert result.exit_code == 1
