feat(aibundle): add support for traceable*

Author: Guikingone

src/ai-bundle/config/services.php

@@ -24,6 +24,8 @@
 use Symfony\AI\AiBundle\Command\AgentCallCommand;
 use Symfony\AI\AiBundle\Command\PlatformInvokeCommand;
 use Symfony\AI\AiBundle\Profiler\DataCollector;
+use Symfony\AI\AiBundle\Profiler\TraceableChat;
+use Symfony\AI\AiBundle\Profiler\TraceableMessageStore;
 use Symfony\AI\AiBundle\Profiler\TraceableToolbox;
 use Symfony\AI\AiBundle\Security\EventListener\IsGrantedToolAttributeListener;
 use Symfony\AI\Chat\Command\DropStoreCommand as DropMessageStoreCommand;
@@ -180,6 +182,8 @@
                 tagged_iterator('ai.traceable_platform'),
                 service('ai.toolbox'),
                 tagged_iterator('ai.traceable_toolbox'),
+                tagged_iterator('ai.traceable_message_store'),
+                tagged_iterator('ai.traceable_chat'),
             ])
             ->tag('data_collector')
         ->set('ai.traceable_toolbox', TraceableToolbox::class)
@@ -188,6 +192,18 @@
                 service('.inner'),
             ])
             ->tag('ai.traceable_toolbox')
+        ->set('ai.traceable_message_store', TraceableMessageStore::class)
+            ->decorate('ai.message_store')
+            ->args([
+                service('.inner'),
+            ])
+            ->tag('ai.traceable_message_store')
+        ->set('ai.traceable_chat', TraceableChat::class)
+            ->decorate('ai.chat')
+            ->args([
+                service('.inner'),
+            ])
+            ->tag('ai.traceable_chat')
 
         // token usage processors
         ->set('ai.platform.token_usage_processor.anthropic', AnthropicTokenOutputProcessor::class)

src/ai-bundle/src/AiBundle.php

@@ -31,6 +31,7 @@
 use Symfony\AI\Agent\Toolbox\ToolFactory\MemoryToolFactory;
 use Symfony\AI\AiBundle\DependencyInjection\ProcessorCompilerPass;
 use Symfony\AI\AiBundle\Exception\InvalidArgumentException;
+use Symfony\AI\AiBundle\Profiler\TraceableMessageStore;
 use Symfony\AI\AiBundle\Profiler\TraceablePlatform;
 use Symfony\AI\AiBundle\Profiler\TraceableToolbox;
 use Symfony\AI\AiBundle\Security\Attribute\IsGrantedTool;
@@ -175,6 +176,17 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
             $builder->setAlias(MessageStoreInterface::class, reset($messageStores));
         }
 
+        if ($builder->getParameter('kernel.debug')) {
+            foreach ($messageStores as $messageStore) {
+                $traceablePlatformDefinition = (new Definition(TraceableMessageStore::class))
+                    ->setDecoratedService($messageStore)
+                    ->setArguments([new Reference('.inner')])
+                    ->addTag('ai.traceable_message_store');
+                $suffix = u($messageStore)->afterLast('.')->toString();
+                $builder->setDefinition('ai.traceable_message_store.'.$suffix, $traceablePlatformDefinition);
+            }
+        }
+
         if ([] === $messageStores) {
             $builder->removeDefinition('ai.command.setup_message_store');
             $builder->removeDefinition('ai.command.drop_message_store');

src/ai-bundle/src/Profiler/DataCollector.php

@@ -12,7 +12,6 @@
 namespace Symfony\AI\AiBundle\Profiler;
 
 use Symfony\AI\Agent\Toolbox\ToolboxInterface;
-use Symfony\AI\Platform\Model;
 use Symfony\AI\Platform\Tool\Tool;
 use Symfony\Bundle\FrameworkBundle\DataCollector\AbstractDataCollector;
 use Symfony\Component\HttpFoundation\Request;
@@ -24,6 +23,7 @@
  *
  * @phpstan-import-type PlatformCallData from TraceablePlatform
  * @phpstan-import-type ToolCallData from TraceableToolbox
+ * @phpstan-import-type MessageStoreData from TraceableMessageStore
  */
 final class DataCollector extends AbstractDataCollector implements LateDataCollectorInterface
 {
@@ -37,17 +37,33 @@ final class DataCollector extends AbstractDataCollector implements LateDataColle
      */
     private readonly array $toolboxes;
 
+    /**
+     * @var TraceableMessageStore[]
+     */
+    private readonly array $messageStores;
+
+    /**
+     * @var TraceableChat[]
+     */
+    private readonly array $chats;
+
     /**
      * @param TraceablePlatform[] $platforms
      * @param TraceableToolbox[]  $toolboxes
+     * @param TraceableMessageStore[] $messageStores
+     * @param TraceableChat[] $chats
      */
     public function __construct(
         iterable $platforms,
         private readonly ToolboxInterface $defaultToolBox,
         iterable $toolboxes,
+        iterable $messageStores,
+        iterable $chats,
     ) {
         $this->platforms = $platforms instanceof \Traversable ? iterator_to_array($platforms) : $platforms;
         $this->toolboxes = $toolboxes instanceof \Traversable ? iterator_to_array($toolboxes) : $toolboxes;
+        $this->messageStores = $messageStores instanceof \Traversable ? iterator_to_array($messageStores) : $messageStores;
+        $this->chats = $chats instanceof \Traversable ? iterator_to_array($chats) : $chats;
     }
 
     public function collect(Request $request, Response $response, ?\Throwable $exception = null): void
@@ -61,6 +77,7 @@ public function lateCollect(): void
             'tools' => $this->defaultToolBox->getTools(),
             'platform_calls' => array_merge(...array_map($this->awaitCallResults(...), $this->platforms)),
             'tool_calls' => array_merge(...array_map(fn (TraceableToolbox $toolbox) => $toolbox->calls, $this->toolboxes)),
+            'messages' => array_merge(...array_map(static fn (TraceableMessageStore $messageStore): array => $messageStore->calls, $this->messageStores)),
         ];
     }
 
@@ -93,6 +110,14 @@ public function getToolCalls(): array
         return $this->data['tool_calls'] ?? [];
     }
 
+    /**
+     * @return MessageStoreData[]
+     */
+    public function getMessages(): array
+    {
+        return $this->data['messages'] ?? [];
+    }
+
     /**
      * @return array{
      *     model: string,

src/ai-bundle/src/Profiler/TraceableChat.php

@@ -0,0 +1,47 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AiBundle\Profiler;
+
+use Symfony\AI\Chat\ChatInterface;
+use Symfony\AI\Platform\Message\AssistantMessage;
+use Symfony\AI\Platform\Message\MessageBag;
+use Symfony\AI\Platform\Message\UserMessage;
+
+/**
+ * @author Guillaume Loulier <[email protected]>
+ *
+ * @phpstan-type MessageStoreData array{
+ *      bag: MessageBag,
+ *  }
+ */
+final class TraceableChat implements ChatInterface
+{
+    /**
+     * @var array
+     */
+    public array $calls = [];
+
+    public function __construct(
+        private readonly ChatInterface $chat,
+    ) {
+    }
+
+    public function initiate(MessageBag $messages): void
+    {
+        // TODO: Implement initiate() method.
+    }
+
+    public function submit(UserMessage $message): AssistantMessage
+    {
+        // TODO: Implement submit() method.
+    }
+}

src/ai-bundle/src/Profiler/TraceableMessageStore.php

@@ -0,0 +1,68 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AiBundle\Profiler;
+
+use Symfony\AI\Chat\ManagedStoreInterface;
+use Symfony\AI\Chat\MessageStoreInterface;
+use Symfony\AI\Platform\Message\MessageBag;
+
+/**
+ * @author Guillaume Loulier <[email protected]>
+ *
+ * @phpstan-type MessageStoreData array{
+ *      bag: MessageBag,
+ *  }
+ */
+final class TraceableMessageStore implements ManagedStoreInterface, MessageStoreInterface
+{
+    /**
+     * @var MessageStoreData[]
+     */
+    public array $calls = [];
+
+    public function __construct(
+        private readonly MessageStoreInterface|ManagedStoreInterface $messageStore,
+    ) {
+    }
+
+    public function save(MessageBag $messages): void
+    {
+        $this->calls[] = [
+            'bag' => $messages,
+        ];
+
+        $this->messageStore->save($messages);
+    }
+
+    public function load(): MessageBag
+    {
+        return $this->messageStore->load();
+    }
+
+    public function setup(array $options = []): void
+    {
+        if (!$this->messageStore instanceof ManagedStoreInterface) {
+            return;
+        }
+
+        $this->messageStore->setup($options);
+    }
+
+    public function drop(): void
+    {
+        if (!$this->messageStore instanceof ManagedStoreInterface) {
+            return;
+        }
+
+        $this->messageStore->drop();
+    }
+}

src/ai-bundle/tests/Profiler/TraceableChatTest.php

@@ -0,0 +1,19 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AiBundle\Tests\Profiler;
+
+use PHPUnit\Framework\TestCase;
+
+final class TraceableChatTest extends TestCase
+{
+
+}

src/ai-bundle/tests/Profiler/TraceableMessageStoreTest.php

@@ -0,0 +1,19 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AiBundle\Tests\Profiler;
+
+use PHPUnit\Framework\TestCase;
+
+final class TraceableMessageStoreTest extends TestCase
+{
+
+}