From 9095ddb9bb2d310911ac8bec2a0fda537d3641a1 Mon Sep 17 00:00:00 2001
From: Iyadh Faleh <iyadh.faleh@gmail.com>
Date: Mon, 16 Dec 2024 19:34:02 +0100
Subject: [PATCH] Set chain of the processor

---
 src/Chain.php                                 |  3 +-
 src/Chain/ChainAwareProcessor.php             |  4 +--
 src/Chain/ChainAwareTrait.php                 | 31 +++++++++++++------
 src/Chain/ToolBox/ChainProcessor.php          |  4 ++-
 ...ToolWrong.php => ToolWithoutAttribute.php} |  2 +-
 5 files changed, 30 insertions(+), 14 deletions(-)
 rename tests/Fixture/Tool/{ToolWrong.php => ToolWithoutAttribute.php} (91%)

diff --git a/src/Chain.php b/src/Chain.php
index c8a43831..5d7e9f34 100644
--- a/src/Chain.php
+++ b/src/Chain.php
@@ -28,8 +28,9 @@ public function __construct(
      */
     public function process(MessageBag $messages, array $options = [], ?ChainAwareProcessor $chainProcessor = null): ResponseInterface
     {
-        $input = new Input($this->llm, $messages, $options);
+        $chainProcessor?->setChain($this);
 
+        $input = new Input($this->llm, $messages, $options);
         if ($chainProcessor) {
             array_map(fn (InputProcessor $processor) => $processor->processInput($input), $chainProcessor->getInputProcessors());
         }
diff --git a/src/Chain/ChainAwareProcessor.php b/src/Chain/ChainAwareProcessor.php
index 422a888a..8a7e5612 100644
--- a/src/Chain/ChainAwareProcessor.php
+++ b/src/Chain/ChainAwareProcessor.php
@@ -4,11 +4,11 @@
 
 namespace PhpLlm\LlmChain\Chain;
 
-use PhpLlm\LlmChain\Chain;
+use PhpLlm\LlmChain\ChainInterface;
 
 interface ChainAwareProcessor
 {
-    public function setChain(Chain $chain): void;
+    public function setChain(ChainInterface $chain): void;
 
     public function addOutputProcessor(OutputProcessor $outputProcessor): self;
 
diff --git a/src/Chain/ChainAwareTrait.php b/src/Chain/ChainAwareTrait.php
index df7dd24e..ad3e13ea 100644
--- a/src/Chain/ChainAwareTrait.php
+++ b/src/Chain/ChainAwareTrait.php
@@ -4,26 +4,27 @@
 
 namespace PhpLlm\LlmChain\Chain;
 
-use PhpLlm\LlmChain\Chain;
 use PhpLlm\LlmChain\Chain\ToolBox\ToolBoxInterface;
+use PhpLlm\LlmChain\ChainInterface;
+use PhpLlm\LlmChain\Exception\InvalidArgumentException;
 
 trait ChainAwareTrait
 {
-    private Chain $chain;
+    private ChainInterface $chain;
 
     private ToolBoxInterface $toolBox;
 
     /**
      * @var InputProcessor[]
      */
-    private array $inputProcessors;
+    private array $inputProcessors = [];
 
     /**
      * @var OutputProcessor[]
      */
-    private array $outputProcessors;
+    private array $outputProcessors = [];
 
-    public function setChain(Chain $chain): void
+    public function setChain(ChainInterface $chain): void
     {
         $this->chain = $chain;
     }
@@ -56,16 +57,28 @@ public function addInputProcessor(InputProcessor $outputProcessor): self
         return $this;
     }
 
-    public function setOutputProcessors(array $outputProcessors): self
+    public function setOutputProcessors(iterable $outputProcessors): self
     {
-        $this->outputProcessors = $outputProcessors;
+        foreach ($outputProcessors as $processor) {
+            if (!$processor instanceof OutputProcessor) {
+                throw new InvalidArgumentException(sprintf('Processor %s must implement %s interface.', $processor::class, OutputProcessor::class));
+            }
+
+            $this->addOutputProcessor($processor);
+        }
 
         return $this;
     }
 
-    public function setInputProcessors(array $inputProcessors): self
+    public function setInputProcessors(iterable $inputProcessors): self
     {
-        $this->inputProcessors = $inputProcessors;
+        foreach ($inputProcessors as $processor) {
+            if (!$processor instanceof InputProcessor) {
+                throw new InvalidArgumentException(sprintf('Processor %s must implement %s interface.', $processor::class, InputProcessor::class));
+            }
+
+            $this->addInputProcessor($processor);
+        }
 
         return $this;
     }
diff --git a/src/Chain/ToolBox/ChainProcessor.php b/src/Chain/ToolBox/ChainProcessor.php
index cea1b6ef..b35a5a53 100644
--- a/src/Chain/ToolBox/ChainProcessor.php
+++ b/src/Chain/ToolBox/ChainProcessor.php
@@ -7,12 +7,14 @@
 use PhpLlm\LlmChain\Chain\ChainAwareProcessor;
 use PhpLlm\LlmChain\Chain\ChainAwareTrait;
 use PhpLlm\LlmChain\Chain\Input;
+use PhpLlm\LlmChain\Chain\InputProcessor;
 use PhpLlm\LlmChain\Chain\Output;
+use PhpLlm\LlmChain\Chain\OutputProcessor;
 use PhpLlm\LlmChain\Exception\MissingModelSupport;
 use PhpLlm\LlmChain\Model\Message\Message;
 use PhpLlm\LlmChain\Model\Response\ToolCallResponse;
 
-final class ChainProcessor implements ChainAwareProcessor
+final class ChainProcessor implements InputProcessor, OutputProcessor, ChainAwareProcessor
 {
     use ChainAwareTrait;
 
diff --git a/tests/Fixture/Tool/ToolWrong.php b/tests/Fixture/Tool/ToolWithoutAttribute.php
similarity index 91%
rename from tests/Fixture/Tool/ToolWrong.php
rename to tests/Fixture/Tool/ToolWithoutAttribute.php
index f1964a70..7ab6d315 100644
--- a/tests/Fixture/Tool/ToolWrong.php
+++ b/tests/Fixture/Tool/ToolWithoutAttribute.php
@@ -4,7 +4,7 @@
 
 namespace PhpLlm\LlmChain\Tests\Fixture\Tool;
 
-final class ToolWrong
+final class ToolWithoutAttribute
 {
     /**
      * @param string $text   The text given to the tool