diff --git a/.gitignore b/.gitignore
index 90cac6b3..fc623910 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
 vendor
 .php-cs-fixer.cache
 .phpunit.cache
+coverage
diff --git a/Makefile b/Makefile
index e50455d2..823a36e4 100644
--- a/Makefile
+++ b/Makefile
@@ -9,3 +9,6 @@ qa-lowest:
 	vendor/bin/phpstan
 	vendor/bin/phpunit
 	git restore composer.lock
+
+coverage:
+	XDEBUG_MODE=coverage vendor/bin/phpunit --coverage-html=coverage
diff --git a/phpstan.dist.neon b/phpstan.dist.neon
index 138b9b5e..ae5c8505 100644
--- a/phpstan.dist.neon
+++ b/phpstan.dist.neon
@@ -4,3 +4,8 @@ parameters:
         - examples/
         - src/
         - tests/
+    ignoreErrors:
+        -
+            message: '#no value type specified in iterable type array#'
+            path: tests/*
+
diff --git a/src/Message/Message.php b/src/Message/Message.php
index 6f7d11d7..57fd59de 100644
--- a/src/Message/Message.php
+++ b/src/Message/Message.php
@@ -6,7 +6,7 @@
 
 use PhpLlm\LlmChain\Response\ToolCall;
 
-final class Message
+final readonly class Message implements \JsonSerializable
 {
     /**
      * @param ?ToolCall[] $toolCalls
@@ -46,8 +46,52 @@ public function isSystem(): bool
         return Role::System === $this->role;
     }
 
+    public function isAssistant(): bool
+    {
+        return Role::Assistant === $this->role;
+    }
+
+    public function isUser(): bool
+    {
+        return Role::User === $this->role;
+    }
+
+    public function isToolCall(): bool
+    {
+        return Role::ToolCall === $this->role;
+    }
+
     public function hasToolCalls(): bool
     {
         return null !== $this->toolCalls && 0 !== count($this->toolCalls);
     }
+
+    /**
+     * @return array{
+     *     role: 'system'|'assistant'|'user'|'tool',
+     *     content: ?string,
+     *     tool_calls?: ToolCall[],
+     *     tool_call_id?: string
+     * }
+     */
+    public function jsonSerialize(): array
+    {
+        $array = [
+            'role' => $this->role->value,
+        ];
+
+        if (null !== $this->content) {
+            $array['content'] = $this->content;
+        }
+
+        if ($this->hasToolCalls() && $this->isToolCall()) {
+            $array['tool_call_id'] = $this->toolCalls[0]->id;
+        }
+
+        if ($this->hasToolCalls() && $this->isAssistant()) {
+            $array['tool_calls'] = $this->toolCalls;
+        }
+
+        return $array;
+    }
 }
diff --git a/src/Message/MessageBag.php b/src/Message/MessageBag.php
index 5b30ed07..a08b6aa3 100644
--- a/src/Message/MessageBag.php
+++ b/src/Message/MessageBag.php
@@ -6,13 +6,6 @@
 
 /**
  * @template-extends \ArrayObject<int, Message>
- *
- * @phpstan-type MessageBagData array<int, array{
- *      role: 'system'|'assistant'|'user'|'function',
- *      content: ?string,
- *      function_call?: array{name: string, arguments: string},
- *      name?: string
- *  }>
  */
 final class MessageBag extends \ArrayObject implements \JsonSerializable
 {
@@ -59,39 +52,10 @@ public function prepend(Message $message): self
     }
 
     /**
-     * @return MessageBagData
-     */
-    public function toArray(): array
-    {
-        return array_map(
-            function (Message $message) {
-                $array = [
-                    'role' => $message->role->value,
-                ];
-
-                if (null !== $message->content) {
-                    $array['content'] = $message->content;
-                }
-
-                if (null !== $message->hasToolCalls() && Role::ToolCall === $message->role) {
-                    $array['tool_call_id'] = $message->toolCalls[0]->id;
-                }
-
-                if (null !== $message->hasToolCalls() && Role::Assistant === $message->role) {
-                    $array['tool_calls'] = $message->toolCalls;
-                }
-
-                return $array;
-            },
-            $this->getArrayCopy(),
-        );
-    }
-
-    /**
-     * @return MessageBagData
+     * @return Message[]
      */
     public function jsonSerialize(): array
     {
-        return $this->toArray();
+        return $this->getArrayCopy();
     }
 }
diff --git a/src/Response/ToolCall.php b/src/Response/ToolCall.php
index 37e10eee..431c86dd 100644
--- a/src/Response/ToolCall.php
+++ b/src/Response/ToolCall.php
@@ -12,7 +12,7 @@
     public function __construct(
         public string $id,
         public string $name,
-        public array $arguments,
+        public array $arguments = [],
     ) {
     }
 
diff --git a/tests/Message/MessageBagTest.php b/tests/Message/MessageBagTest.php
new file mode 100644
index 00000000..9e4e175b
--- /dev/null
+++ b/tests/Message/MessageBagTest.php
@@ -0,0 +1,108 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Message;
+
+use PhpLlm\LlmChain\Message\Message;
+use PhpLlm\LlmChain\Message\MessageBag;
+use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\Small;
+use PHPUnit\Framework\Attributes\UsesClass;
+use PHPUnit\Framework\TestCase;
+
+#[CoversClass(MessageBag::class)]
+#[UsesClass(Message::class)]
+#[Small]
+final class MessageBagTest extends TestCase
+{
+    public function testGetSystemMessage(): void
+    {
+        $messageBag = new MessageBag(
+            Message::forSystem('My amazing system prompt.'),
+            Message::ofAssistant('It is time to sleep.'),
+            Message::ofUser('Hello, world!'),
+        );
+
+        $systemMessage = $messageBag->getSystemMessage();
+
+        self::assertSame('My amazing system prompt.', $systemMessage->content);
+    }
+
+    public function testGetSystemMessageWithoutSystemMessage(): void
+    {
+        $messageBag = new MessageBag(
+            Message::ofAssistant('It is time to sleep.'),
+            Message::ofUser('Hello, world!'),
+        );
+
+        self::assertNull($messageBag->getSystemMessage());
+    }
+
+    public function testWith(): void
+    {
+        $messageBag = new MessageBag(
+            Message::forSystem('My amazing system prompt.'),
+            Message::ofAssistant('It is time to sleep.'),
+            Message::ofUser('Hello, world!'),
+        );
+
+        $newMessage = Message::ofAssistant('It is time to wake up.');
+        $newMessageBag = $messageBag->with($newMessage);
+
+        self::assertCount(3, $messageBag);
+        self::assertCount(4, $newMessageBag);
+        self::assertSame('It is time to wake up.', $newMessageBag[3]->content);
+    }
+
+    public function testWithoutSystemMessage(): void
+    {
+        $messageBag = new MessageBag(
+            Message::forSystem('My amazing system prompt.'),
+            Message::ofAssistant('It is time to sleep.'),
+            Message::ofUser('Hello, world!'),
+        );
+
+        $newMessageBag = $messageBag->withoutSystemMessage();
+
+        self::assertCount(3, $messageBag);
+        self::assertCount(2, $newMessageBag);
+        self::assertSame('It is time to sleep.', $newMessageBag[0]->content);
+    }
+
+    public function testPrepend(): void
+    {
+        $messageBag = new MessageBag(
+            Message::ofAssistant('It is time to sleep.'),
+            Message::ofUser('Hello, world!'),
+        );
+
+        $newMessage = Message::forSystem('My amazing system prompt.');
+        $newMessageBag = $messageBag->prepend($newMessage);
+
+        self::assertCount(2, $messageBag);
+        self::assertCount(3, $newMessageBag);
+        self::assertSame('My amazing system prompt.', $newMessageBag[0]->content);
+    }
+
+    public function testJsonSerialize(): void
+    {
+        $messageBag = new MessageBag(
+            Message::forSystem('My amazing system prompt.'),
+            Message::ofAssistant('It is time to sleep.'),
+            Message::ofUser('Hello, world!'),
+        );
+
+        $json = json_encode($messageBag);
+
+        self::assertJson($json);
+        self::assertJsonStringEqualsJsonString(
+            json_encode([
+                ['role' => 'system', 'content' => 'My amazing system prompt.'],
+                ['role' => 'assistant', 'content' => 'It is time to sleep.'],
+                ['role' => 'user', 'content' => 'Hello, world!'],
+            ]),
+            $json
+        );
+    }
+}
diff --git a/tests/Message/MessageTest.php b/tests/Message/MessageTest.php
new file mode 100644
index 00000000..667a2a5a
--- /dev/null
+++ b/tests/Message/MessageTest.php
@@ -0,0 +1,136 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Message;
+
+use PhpLlm\LlmChain\Message\Message;
+use PhpLlm\LlmChain\Response\ToolCall;
+use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\DataProvider;
+use PHPUnit\Framework\Attributes\Small;
+use PHPUnit\Framework\Attributes\UsesClass;
+use PHPUnit\Framework\TestCase;
+
+#[CoversClass(Message::class)]
+#[UsesClass(ToolCall::class)]
+#[Small]
+final class MessageTest extends TestCase
+{
+    public function testCreateSystemMessage(): void
+    {
+        $message = Message::forSystem('My amazing system prompt.');
+
+        self::assertSame('My amazing system prompt.', $message->content);
+        self::assertTrue($message->isSystem());
+        self::assertFalse($message->isAssistant());
+        self::assertFalse($message->isUser());
+        self::assertFalse($message->isToolCall());
+        self::assertFalse($message->hasToolCalls());
+    }
+
+    public function testCreateAssistantMessage(): void
+    {
+        $message = Message::ofAssistant('It is time to sleep.');
+
+        self::assertSame('It is time to sleep.', $message->content);
+        self::assertFalse($message->isSystem());
+        self::assertTrue($message->isAssistant());
+        self::assertFalse($message->isUser());
+        self::assertFalse($message->isToolCall());
+        self::assertFalse($message->hasToolCalls());
+    }
+
+    public function testCreateAssistantMessageWithToolCalls(): void
+    {
+        $toolCalls = [
+            new ToolCall('call_123456', 'my_tool', ['foo' => 'bar']),
+            new ToolCall('call_456789', 'my_faster_tool'),
+        ];
+        $message = Message::ofAssistant(toolCalls: $toolCalls);
+
+        self::assertCount(2, $message->toolCalls);
+        self::assertFalse($message->isSystem());
+        self::assertTrue($message->isAssistant());
+        self::assertFalse($message->isUser());
+        self::assertFalse($message->isToolCall());
+        self::assertTrue($message->hasToolCalls());
+    }
+
+    public function testCreateUserMessage(): void
+    {
+        $message = Message::ofUser('Hi, my name is John.');
+
+        self::assertSame('Hi, my name is John.', $message->content);
+        self::assertFalse($message->isSystem());
+        self::assertFalse($message->isAssistant());
+        self::assertTrue($message->isUser());
+        self::assertFalse($message->isToolCall());
+        self::assertFalse($message->hasToolCalls());
+    }
+
+    public function testCreateToolCallMessage(): void
+    {
+        $toolCall = new ToolCall('call_123456', 'my_tool', ['foo' => 'bar']);
+        $message = Message::ofToolCall($toolCall, 'Foo bar.');
+
+        self::assertSame('Foo bar.', $message->content);
+        self::assertCount(1, $message->toolCalls);
+        self::assertFalse($message->isSystem());
+        self::assertFalse($message->isAssistant());
+        self::assertFalse($message->isUser());
+        self::assertTrue($message->isToolCall());
+        self::assertTrue($message->hasToolCalls());
+    }
+
+    #[DataProvider('provideJsonScenarios')]
+    public function testJsonSerialize(Message $message, array $expected): void
+    {
+        self::assertSame($expected, $message->jsonSerialize());
+    }
+
+    public static function provideJsonScenarios(): array
+    {
+        $toolCall1 = new ToolCall('call_123456', 'my_tool', ['foo' => 'bar']);
+        $toolCall2 = new ToolCall('call_456789', 'my_faster_tool');
+
+        return [
+            'system' => [
+                Message::forSystem('My amazing system prompt.'),
+                [
+                    'role' => 'system',
+                    'content' => 'My amazing system prompt.',
+                ],
+            ],
+            'assistant' => [
+                Message::ofAssistant('It is time to sleep.'),
+                [
+                    'role' => 'assistant',
+                    'content' => 'It is time to sleep.',
+                ],
+            ],
+            'assistant_with_tool_calls' => [
+                Message::ofAssistant(toolCalls: [$toolCall1, $toolCall2]),
+                [
+                    'role' => 'assistant',
+                    'tool_calls' => [$toolCall1, $toolCall2],
+                ],
+            ],
+            'user' => [
+                Message::ofUser('Hi, my name is John.'),
+                [
+                    'role' => 'user',
+                    'content' => 'Hi, my name is John.',
+                ],
+            ],
+            'tool_call' => [
+                Message::ofToolCall($toolCall1, 'Foo bar.'),
+                [
+                    'role' => 'tool',
+                    'content' => 'Foo bar.',
+                    'tool_call_id' => 'call_123456',
+                ],
+            ],
+        ];
+    }
+}
diff --git a/tests/Response/ChoiceTest.php b/tests/Response/ChoiceTest.php
new file mode 100644
index 00000000..2e9a249c
--- /dev/null
+++ b/tests/Response/ChoiceTest.php
@@ -0,0 +1,54 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Tests\Response;
+
+use PhpLlm\LlmChain\Response\Choice;
+use PhpLlm\LlmChain\Response\ToolCall;
+use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\Small;
+use PHPUnit\Framework\Attributes\UsesClass;
+use PHPUnit\Framework\TestCase;
+
+#[CoversClass(Choice::class)]
+#[UsesClass(ToolCall::class)]
+#[Small]
+final class ChoiceTest extends TestCase
+{
+    public function testChoiceEmpty(): void
+    {
+        $choice = new Choice();
+        self::assertFalse($choice->hasContent());
+        self::assertNull($choice->getContent());
+        self::assertFalse($choice->hasToolCall());
+        self::assertCount(0, $choice->getToolCalls());
+    }
+
+    public function testChoiceWithContent(): void
+    {
+        $choice = new Choice('content');
+        self::assertTrue($choice->hasContent());
+        self::assertSame('content', $choice->getContent());
+        self::assertFalse($choice->hasToolCall());
+        self::assertCount(0, $choice->getToolCalls());
+    }
+
+    public function testChoiceWithToolCall(): void
+    {
+        $choice = new Choice(null, [new ToolCall('name', 'arguments')]);
+        self::assertFalse($choice->hasContent());
+        self::assertNull($choice->getContent());
+        self::assertTrue($choice->hasToolCall());
+        self::assertCount(1, $choice->getToolCalls());
+    }
+
+    public function testChoiceWithContentAndToolCall(): void
+    {
+        $choice = new Choice('content', [new ToolCall('name', 'arguments')]);
+        self::assertTrue($choice->hasContent());
+        self::assertSame('content', $choice->getContent());
+        self::assertTrue($choice->hasToolCall());
+        self::assertCount(1, $choice->getToolCalls());
+    }
+}
diff --git a/tests/Response/ResponseTest.php b/tests/Response/ResponseTest.php
new file mode 100644
index 00000000..23d27759
--- /dev/null
+++ b/tests/Response/ResponseTest.php
@@ -0,0 +1,98 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Tests\Response;
+
+use PhpLlm\LlmChain\Response\Choice;
+use PhpLlm\LlmChain\Response\Response;
+use PhpLlm\LlmChain\Response\ToolCall;
+use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\Small;
+use PHPUnit\Framework\Attributes\UsesClass;
+use PHPUnit\Framework\TestCase;
+
+#[CoversClass(Response::class)]
+#[UsesClass(Choice::class)]
+#[UsesClass(ToolCall::class)]
+#[Small]
+final class ResponseTest extends TestCase
+{
+    public function testGetChoices(): void
+    {
+        $response = new Response(
+            new Choice('content', [new ToolCall('call_123456', 'name', ['foo' => 'bar'])]),
+            new Choice('content', [new ToolCall('call_234567', 'name', ['foo' => 'bar'])]),
+        );
+
+        self::assertCount(2, $response->getChoices());
+    }
+
+    public function testConstructorThrowsException(): void
+    {
+        $this->expectException(\LogicException::class);
+        $this->expectExceptionMessage('Response must have at least one choice');
+
+        $response = new Response();
+    }
+
+    public function testGetContent(): void
+    {
+        $response = new Response(
+            new Choice('content', [new ToolCall('call_123456', 'name', ['foo' => 'bar'])]),
+        );
+
+        self::assertSame('content', $response->getContent());
+    }
+
+    public function testGetContentThrowsException(): void
+    {
+        $response = new Response(
+            new Choice('content', [new ToolCall('call_123456', 'name', ['foo' => 'bar'])]),
+            new Choice('content', [new ToolCall('call_123456', 'name', ['foo' => 'bar'])]),
+        );
+
+        $this->expectException(\LogicException::class);
+        $this->expectExceptionMessage('Response has more than one choice');
+
+        $response->getContent();
+    }
+
+    public function testGetToolCalls(): void
+    {
+        $response = new Response(
+            new Choice('content', [new ToolCall('call_123456', 'name', ['foo' => 'bar'])]),
+        );
+
+        self::assertCount(1, $response->getToolCalls());
+    }
+
+    public function testGetToolCallsThrowsException(): void
+    {
+        $response = new Response(
+            new Choice('content', [new ToolCall('call_123456', 'name', ['foo' => 'bar'])]),
+            new Choice('content', [new ToolCall('call_123456', 'name', ['foo' => 'bar'])]),
+        );
+
+        $this->expectException(\LogicException::class);
+        $this->expectExceptionMessage('Response has more than one choice');
+
+        $response->getToolCalls();
+    }
+
+    public function testHasToolCalls(): void
+    {
+        $response = new Response(
+            new Choice('content', [new ToolCall('call_123456', 'name', ['foo' => 'bar'])]),
+        );
+
+        self::assertTrue($response->hasToolCalls());
+    }
+
+    public function testHasToolCallsReturnsFalse(): void
+    {
+        $response = new Response(new Choice('content'));
+
+        self::assertFalse($response->hasToolCalls());
+    }
+}
diff --git a/tests/Response/ToolCallTest.php b/tests/Response/ToolCallTest.php
new file mode 100644
index 00000000..a64fb7d5
--- /dev/null
+++ b/tests/Response/ToolCallTest.php
@@ -0,0 +1,39 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Tests\Response;
+
+use PhpLlm\LlmChain\Response\ToolCall;
+use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\Small;
+use PHPUnit\Framework\TestCase;
+
+#[CoversClass(ToolCall::class)]
+#[Small]
+final class ToolCallTest extends TestCase
+{
+    public function testToolCall(): void
+    {
+        $toolCall = new ToolCall('id', 'name', ['foo' => 'bar']);
+        self::assertSame('id', $toolCall->id);
+        self::assertSame('name', $toolCall->name);
+        self::assertSame(['foo' => 'bar'], $toolCall->arguments);
+    }
+
+    public function testToolCallJsonSerialize(): void
+    {
+        $toolCall = new ToolCall('id', 'name', ['foo' => 'bar']);
+        self::assertSame(
+            [
+                'id' => 'id',
+                'type' => 'function',
+                'function' => [
+                    'name' => 'name',
+                    'arguments' => '{"foo":"bar"}',
+                ],
+            ],
+            $toolCall->jsonSerialize()
+        );
+    }
+}
diff --git a/tests/ToolBox/ParameterAnalyzerTest.php b/tests/ToolBox/ParameterAnalyzerTest.php
index ea91761f..f49215f2 100644
--- a/tests/ToolBox/ParameterAnalyzerTest.php
+++ b/tests/ToolBox/ParameterAnalyzerTest.php
@@ -7,11 +7,17 @@
 use PhpLlm\LlmChain\Tests\ToolBox\Tool\ToolNoParams;
 use PhpLlm\LlmChain\Tests\ToolBox\Tool\ToolOptionalParam;
 use PhpLlm\LlmChain\Tests\ToolBox\Tool\ToolRequiredParams;
+use PhpLlm\LlmChain\ToolBox\AsTool;
+use PhpLlm\LlmChain\ToolBox\Metadata;
 use PhpLlm\LlmChain\ToolBox\ParameterAnalyzer;
 use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\UsesClass;
 use PHPUnit\Framework\TestCase;
 
 #[CoversClass(ParameterAnalyzer::class)]
+#[UsesClass(AsTool::class)]
+#[UsesClass(Metadata::class)]
+#[UsesClass(ParameterAnalyzer::class)]
 final class ParameterAnalyzerTest extends TestCase
 {
     private ParameterAnalyzer $analyzer;
diff --git a/tests/ToolBox/RegistryTest.php b/tests/ToolBox/RegistryTest.php
index 42436573..a3231649 100644
--- a/tests/ToolBox/RegistryTest.php
+++ b/tests/ToolBox/RegistryTest.php
@@ -8,13 +8,21 @@
 use PhpLlm\LlmChain\Tests\ToolBox\Tool\ToolNoParams;
 use PhpLlm\LlmChain\Tests\ToolBox\Tool\ToolOptionalParam;
 use PhpLlm\LlmChain\Tests\ToolBox\Tool\ToolRequiredParams;
+use PhpLlm\LlmChain\ToolBox\AsTool;
+use PhpLlm\LlmChain\ToolBox\Metadata;
 use PhpLlm\LlmChain\ToolBox\ParameterAnalyzer;
 use PhpLlm\LlmChain\ToolBox\Registry;
 use PhpLlm\LlmChain\ToolBox\ToolAnalyzer;
 use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\UsesClass;
 use PHPUnit\Framework\TestCase;
 
 #[CoversClass(Registry::class)]
+#[UsesClass(ToolCall::class)]
+#[UsesClass(AsTool::class)]
+#[UsesClass(Metadata::class)]
+#[UsesClass(ParameterAnalyzer::class)]
+#[UsesClass(ToolAnalyzer::class)]
 final class RegistryTest extends TestCase
 {
     private Registry $registry;
diff --git a/tests/ToolBox/ToolAnalyzerTest.php b/tests/ToolBox/ToolAnalyzerTest.php
index 6deace6b..eba9fb02 100644
--- a/tests/ToolBox/ToolAnalyzerTest.php
+++ b/tests/ToolBox/ToolAnalyzerTest.php
@@ -2,19 +2,25 @@
 
 declare(strict_types=1);
 
-namespace PhpLlm\LlmChain\Tests;
+namespace PhpLlm\LlmChain\Tests\ToolBox;
 
 use PhpLlm\LlmChain\Exception\InvalidToolImplementation;
 use PhpLlm\LlmChain\Tests\ToolBox\Tool\ToolMultiple;
 use PhpLlm\LlmChain\Tests\ToolBox\Tool\ToolRequiredParams;
 use PhpLlm\LlmChain\Tests\ToolBox\Tool\ToolWrong;
+use PhpLlm\LlmChain\ToolBox\AsTool;
 use PhpLlm\LlmChain\ToolBox\Metadata;
 use PhpLlm\LlmChain\ToolBox\ParameterAnalyzer;
 use PhpLlm\LlmChain\ToolBox\ToolAnalyzer;
 use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\UsesClass;
 use PHPUnit\Framework\TestCase;
 
 #[CoversClass(ToolAnalyzer::class)]
+#[UsesClass(AsTool::class)]
+#[UsesClass(Metadata::class)]
+#[UsesClass(ParameterAnalyzer::class)]
+#[UsesClass(InvalidToolImplementation::class)]
 final class ToolAnalyzerTest extends TestCase
 {
     private ToolAnalyzer $toolAnalyzer;