bug #1560 [Agent][Chat] Add metadata to AssistantMessage (janssensglenn)

This PR was squashed before being merged into the main branch.

Discussion
----------

[Agent][Chat] Add `metadata` to `AssistantMessage`

| Q             | A
| ------------- | ---
| Bug fix?      | yes
| New feature?  | no
| Docs?         | no
| Issues        | Fix #1558
| License       | MIT

This PR adds the sources of a TextResult to the metadata of the AssistantMessage within the [Chat](https://github.com/symfony/ai-chat) component.

Commits
-------

337f0f3 [Agent][Chat] Add `metadata` to `AssistantMessage`

Author: chr-hertel

src/agent/CHANGELOG.md

@@ -5,6 +5,7 @@ CHANGELOG
 ---
 
  * [BC BREAK] Rename `Symfony\AI\Agent\Toolbox\Tool\Agent` to `Symfony\AI\Agent\Toolbox\Tool\Subagent`
+ * Add `MetaDataAwareTrait` to `MockResponse`, the metadata will also be set on the returned `TextResult` when calling the `toResult` function
 
 0.3
 ---

src/agent/src/MockResponse.php

@@ -11,6 +11,7 @@
 
 namespace Symfony\AI\Agent;
 
+use Symfony\AI\Platform\Metadata\MetadataAwareTrait;
 use Symfony\AI\Platform\Result\ResultInterface;
 use Symfony\AI\Platform\Result\TextResult;
 
@@ -24,14 +25,19 @@
  */
 final class MockResponse
 {
+    use MetadataAwareTrait;
+
     public function __construct(
         private readonly string $content = '',
     ) {
     }
 
     public function toResult(): ResultInterface
     {
-        return new TextResult($this->content);
+        $result = new TextResult($this->content);
+        $result->getMetadata()->merge($this->getMetadata());
+
+        return $result;
     }
 
     public function getContent(): string

src/agent/tests/MockResponseTest.php

@@ -34,10 +34,13 @@ public function testConstructorWithEmptyContent()
     public function testToResult()
     {
         $response = new MockResponse('Response content');
+        $response->getMetadata()->add('sources', ['https://example.com']);
+
         $result = $response->toResult();
 
         $this->assertInstanceOf(TextResult::class, $result);
         $this->assertSame('Response content', $result->getContent());
+        $this->assertEquals($response->getMetadata(), $result->getMetadata());
     }
 
     public function testCreate()

src/chat/CHANGELOG.md

@@ -5,3 +5,4 @@ CHANGELOG
 ---
 
  * Add the component
+ * Add `metadata` from `TextResult` to `AssistantMessage`

src/chat/src/Chat.php

@@ -45,6 +45,7 @@ public function submit(UserMessage $message): AssistantMessage
         \assert($result instanceof TextResult);
 
         $assistantMessage = Message::ofAssistant($result->getContent());
+        $assistantMessage->getMetadata()->merge($result->getMetadata());
         $messages->add($assistantMessage);
 
         $this->store->save($messages);

src/chat/tests/ChatTest.php

@@ -11,25 +11,24 @@
 
 namespace Symfony\AI\Chat\Tests;
 
-use PHPUnit\Framework\MockObject\MockObject;
 use PHPUnit\Framework\TestCase;
-use Symfony\AI\Agent\AgentInterface;
+use Symfony\AI\Agent\MockAgent;
+use Symfony\AI\Agent\MockResponse;
 use Symfony\AI\Chat\Chat;
 use Symfony\AI\Chat\InMemory\Store as InMemoryStore;
 use Symfony\AI\Platform\Message\AssistantMessage;
 use Symfony\AI\Platform\Message\Message;
 use Symfony\AI\Platform\Message\MessageBag;
-use Symfony\AI\Platform\Result\TextResult;
 
 final class ChatTest extends TestCase
 {
-    private AgentInterface&MockObject $agent;
+    private MockAgent $agent;
     private InMemoryStore $store;
     private Chat $chat;
 
     protected function setUp(): void
     {
-        $this->agent = $this->createMock(AgentInterface::class);
+        $this->agent = new MockAgent();
         $this->store = new InMemoryStore();
         $this->chat = new Chat($this->agent, $this->store);
     }
@@ -41,29 +40,30 @@ public function testItInitiatesChatByClearingAndSavingMessages()
         $this->chat->initiate($messages);
 
         $this->assertCount(0, $this->store->load());
+
+        $this->agent->assertNotCalled();
     }
 
     public function testItSubmitsUserMessageAndReturnsAssistantMessage()
     {
-        $userMessage = Message::ofUser('Hello, how are you?');
+        $userMessage = Message::ofUser($userPrompt = 'Hello, how are you?');
         $assistantContent = 'I am doing well, thank you!';
+        $assistantSources = ['https://example.com'];
 
-        $textResult = new TextResult($assistantContent);
-
-        $this->agent->expects($this->once())
-            ->method('call')
-            ->with($this->callback(static function (MessageBag $messages) use ($userMessage) {
-                $messagesArray = $messages->getMessages();
+        $response = new MockResponse($assistantContent);
+        $response->getMetadata()->add('sources', $assistantSources);
 
-                return end($messagesArray) === $userMessage;
-            }))
-            ->willReturn($textResult);
+        $this->agent->addResponse($userPrompt, $response);
 
         $result = $this->chat->submit($userMessage);
 
         $this->assertInstanceOf(AssistantMessage::class, $result);
         $this->assertSame($assistantContent, $result->getContent());
+        $this->assertSame($assistantSources, $result->getMetadata()->get('sources', []));
         $this->assertCount(2, $this->store->load());
+
+        $this->agent->assertCallCount(1);
+        $this->agent->assertCalledWith($userPrompt);
     }
 
     public function testItAppendsMessagesToExistingConversation()
@@ -75,42 +75,37 @@ public function testItAppendsMessagesToExistingConversation()
         $existingMessages->add($existingUserMessage);
         $existingMessages->add($existingAssistantMessage);
 
-        $newUserMessage = Message::ofUser('Can you help with programming?');
-        $newAssistantContent = 'Yes, I can help with programming!';
+        $this->store->save($existingMessages);
 
-        $textResult = new TextResult($newAssistantContent);
+        $newUserMessage = Message::ofUser($newUserPrompt = 'Can you help with programming?');
+        $newAssistantContent = 'Yes, I can help with programming!';
 
-        $this->agent->expects($this->once())
-            ->method('call')
-            ->willReturn($textResult);
+        $this->agent->addResponse($newUserPrompt, $newAssistantContent);
 
         $result = $this->chat->submit($newUserMessage);
 
         $this->assertInstanceOf(AssistantMessage::class, $result);
         $this->assertSame($newAssistantContent, $result->getContent());
-        $this->assertCount(2, $this->store->load());
+        $this->assertCount(4, $this->store->load());
+
+        $this->agent->assertCallCount(1);
+        $this->agent->assertCalledWith($newUserPrompt);
     }
 
     public function testItHandlesEmptyMessageStore()
     {
-        $userMessage = Message::ofUser('First message');
+        $userMessage = Message::ofUser($userPrompt = 'First message');
         $assistantContent = 'First response';
 
-        $textResult = new TextResult($assistantContent);
-
-        $this->agent->expects($this->once())
-            ->method('call')
-            ->with($this->callback(static function (MessageBag $messages) {
-                $messagesArray = $messages->getMessages();
-
-                return 1 === \count($messagesArray);
-            }))
-            ->willReturn($textResult);
+        $this->agent->addResponse($userPrompt, $assistantContent);
 
         $result = $this->chat->submit($userMessage);
 
         $this->assertInstanceOf(AssistantMessage::class, $result);
         $this->assertSame($assistantContent, $result->getContent());
         $this->assertCount(2, $this->store->load());
+
+        $this->agent->assertCallCount(1);
+        $this->agent->assertCalledWith($userPrompt);
     }
 }