feat: add option to keep tool messages (#321)

Fixes #321

Author: philipheimboeck

README.md

@@ -341,6 +341,34 @@ $eventDispatcher->addListener(ToolCallsExecuted::class, function (ToolCallsExecu
 });
 ```
 
+#### Keeping Tool Messages
+
+Sometimes you might wish to keep the tool messages (`AssistantMessage` containing the `toolCalls` and `ToolCallMessage` containing the response) in the context.
+Enable the `keepToolMessages` flag to ensure those messages will be added to your `MessageBag`.
+
+```php
+use PhpLlm\LlmChain\Chain\Toolbox\ChainProcessor;
+use PhpLlm\LlmChain\Chain\Toolbox\Toolbox;
+
+// Platform & LLM instantiation
+$messages = new MessageBag(
+    Message::forSystem(<<<PROMPT
+        Please answer all user questions only using the similary_search tool. Do not add information and if you cannot
+        find an answer, say so.
+        PROMPT),
+    Message::ofUser('...') // The user's question.
+);
+
+$yourTool = new YourTool();
+
+$toolbox = Toolbox::create($yourTool);
+$toolProcessor = new ChainProcessor($toolbox, keepToolMessages: true);
+
+$chain = new Chain($platform, $llm, inputProcessor: [$toolProcessor], outputProcessor: [$toolProcessor]);
+$response = $chain->call($messages);
+// $messages will now include the tool messages
+```
+
 #### Code Examples (with built-in tools)
 
 1. [Brave Tool](examples/toolbox/brave.php)

src/Chain/Toolbox/ChainProcessor.php

@@ -28,6 +28,7 @@ public function __construct(
         private readonly ToolboxInterface $toolbox,
         private readonly ToolResultConverter $resultConverter = new ToolResultConverter(),
         private readonly ?EventDispatcherInterface $eventDispatcher = null,
+        private readonly bool $keepToolMessages = false,
     ) {
     }
 
@@ -81,7 +82,7 @@ private function isFlatStringArray(array $tools): bool
     private function handleToolCallsCallback(Output $output): \Closure
     {
         return function (ToolCallResponse $response, ?AssistantMessage $streamedAssistantResponse = null) use ($output): ResponseInterface {
-            $messages = clone $output->messages;
+            $messages = $this->keepToolMessages ? $output->messages : clone $output->messages;
 
             if (null !== $streamedAssistantResponse && '' !== $streamedAssistantResponse->content) {
                 $messages->add($streamedAssistantResponse);

tests/Chain/Toolbox/ChainProcessorTest.php

@@ -4,18 +4,26 @@
 
 namespace PhpLlm\LlmChain\Tests\Chain\Toolbox;
 
+use PhpLlm\LlmChain\Chain;
 use PhpLlm\LlmChain\Chain\Input;
+use PhpLlm\LlmChain\Chain\Output;
 use PhpLlm\LlmChain\Chain\Toolbox\ChainProcessor;
 use PhpLlm\LlmChain\Chain\Toolbox\ExecutionReference;
 use PhpLlm\LlmChain\Chain\Toolbox\Metadata;
 use PhpLlm\LlmChain\Chain\Toolbox\ToolboxInterface;
 use PhpLlm\LlmChain\Exception\MissingModelSupport;
 use PhpLlm\LlmChain\Model\LanguageModel;
+use PhpLlm\LlmChain\Model\Message\AssistantMessage;
 use PhpLlm\LlmChain\Model\Message\MessageBag;
+use PhpLlm\LlmChain\Model\Message\ToolCallMessage;
+use PhpLlm\LlmChain\Model\Response\ToolCall;
+use PhpLlm\LlmChain\Model\Response\ToolCallResponse;
+use PhpLlm\LlmChain\PlatformInterface;
 use PHPUnit\Framework\Attributes\CoversClass;
 use PHPUnit\Framework\Attributes\Test;
 use PHPUnit\Framework\Attributes\UsesClass;
 use PHPUnit\Framework\TestCase;
+use Symfony\Contracts\EventDispatcher\EventDispatcherInterface;
 
 #[CoversClass(ChainProcessor::class)]
 #[UsesClass(Input::class)]
@@ -93,4 +101,60 @@ public function processInputWithUnsupportedToolCallingWillThrowException(): void
 
         $chainProcessor->processInput($input);
     }
+
+    #[Test]
+    public function processOutputWithToolCallResponseKeepingMessages(): void
+    {
+        $toolbox = $this->createMock(ToolboxInterface::class);
+        $toolbox->expects($this->once())->method('execute')->willReturn('Test response');
+
+        $llm = $this->createStub(LanguageModel::class);
+
+        $eventDispatcher = $this->createMock(EventDispatcherInterface::class);
+        $eventDispatcher->expects($this->once())->method('dispatch');
+
+        $messageBag = new MessageBag();
+
+        $response = new ToolCallResponse(new ToolCall('id1', 'tool1', ['arg1' => 'value1']));
+
+        $chain = new Chain($this->createStub(PlatformInterface::class), $llm);
+
+        $chainProcessor = new ChainProcessor($toolbox, eventDispatcher: $eventDispatcher, keepToolMessages: true);
+        $chainProcessor->setChain($chain);
+
+        $output = new Output($llm, $response, $messageBag, []);
+
+        $chainProcessor->processOutput($output);
+
+        self::assertCount(2, $messageBag);
+        self::assertInstanceOf(AssistantMessage::class, $messageBag->getMessages()[0]);
+        self::assertInstanceOf(ToolCallMessage::class, $messageBag->getMessages()[1]);
+    }
+
+    #[Test]
+    public function processOutputWithToolCallResponseForgettingMessages(): void
+    {
+        $toolbox = $this->createMock(ToolboxInterface::class);
+        $toolbox->expects($this->once())->method('execute')->willReturn('Test response');
+
+        $llm = $this->createStub(LanguageModel::class);
+
+        $eventDispatcher = $this->createMock(EventDispatcherInterface::class);
+        $eventDispatcher->expects($this->once())->method('dispatch');
+
+        $messageBag = new MessageBag();
+
+        $response = new ToolCallResponse(new ToolCall('id1', 'tool1', ['arg1' => 'value1']));
+
+        $chain = new Chain($this->createStub(PlatformInterface::class), $llm);
+
+        $chainProcessor = new ChainProcessor($toolbox, eventDispatcher: $eventDispatcher, keepToolMessages: false);
+        $chainProcessor->setChain($chain);
+
+        $output = new Output($llm, $response, $messageBag, []);
+
+        $chainProcessor->processOutput($output);
+
+        self::assertCount(0, $messageBag);
+    }
 }