fix chat history management

Author: ilvalerione

src/Agent/Nodes/ChatNode.php

@@ -41,19 +41,18 @@ public function __invoke(AIInferenceEvent $event, AgentState $state): StopEvent|
         $lastMessage = $chatHistory->getLastMessage();
 
         $this->emit('inference-start', new InferenceStart($lastMessage));
-
         $response = $this->inference($event, $chatHistory->getMessages());
-
         $this->emit('inference-stop', new InferenceStop($lastMessage, $response));
 
-        // Add the response to chat history
-        $this->addToChatHistory($state, $response);
-
-        // If the response is a tool call, route to tool execution
+        // If the response is a tool call, route to the tool node.
+        // It will be responsible to add the tool call message to the chat history.
         if ($response instanceof ToolCallMessage) {
             return new ToolCallEvent($response, $event);
         }
 
+        // Add the final response to chat history (after tool loop)
+        $this->addToChatHistory($state, $response);
+
         return new StopEvent();
     }
 

src/Agent/Nodes/ToolNode.php

@@ -6,6 +6,7 @@
 
 use Generator;
 use NeuronAI\Agent\AgentState;
+use NeuronAI\Agent\ChatHistoryHelper;
 use NeuronAI\Agent\Events\ToolCallEvent;
 use NeuronAI\Chat\Messages\Stream\Chunks\ToolCallChunk;
 use NeuronAI\Chat\Messages\Stream\Chunks\ToolResultChunk;
@@ -23,6 +24,8 @@
  */
 class ToolNode extends Node
 {
+    use ChatHistoryHelper;
+
     public function __construct(
         protected int $maxTries = 10
     ) {
@@ -34,6 +37,10 @@ public function __construct(
      */
     public function __invoke(ToolCallEvent $event, AgentState $state): Generator
     {
+        // Adding the tool call message to the chat history here allows the middleware to hook
+        // the ToolNode before the tool call is added to the history.
+        $this->addToChatHistory($state, $event->toolCallMessage);
+
         $toolCallResult = yield from $this->executeTools($event->toolCallMessage, $state);
 
         // Only carry the tool result message as the next turn in the conversation