feature #1653 [AI Bundle] Move debug service decorating to compiler pass (HypeMC)

This PR was merged into the main branch.

Discussion
----------

[AI Bundle] Move debug service decorating to compiler pass

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

Currently, user-defined services are not decorated with traceable classes in debug mode. By moving the decoration to a compiler pass, these cases are also covered.

My use case requires multiple instances of the OpenAI platform with different API keys in order to track usage per feature. Since OpenAI's configuration does not support multiple instances, I registered and tagged the services manually.

Commits
-------

304a1276 [AI Bundle] Move debug service decorating to compiler pass

src/ai-bundle/CHANGELOG.md

@@ -1,6 +1,11 @@
 CHANGELOG
 =========
 
+0.6
+---
+
+ * Move debug service decorating to compiler pass to cover user-defined services
+
 0.5
 ---
 

src/ai-bundle/src/AiBundle.php

@@ -30,12 +30,9 @@ use Symfony\AI\Agent\Toolbox\FaultTolerantToolbox;
 use Symfony\AI\Agent\Toolbox\Tool\Subagent;
 use Symfony\AI\Agent\Toolbox\ToolFactory\ChainFactory;
 use Symfony\AI\Agent\Toolbox\ToolFactory\MemoryToolFactory;
+use Symfony\AI\AiBundle\DependencyInjection\DebugCompilerPass;
 use Symfony\AI\AiBundle\DependencyInjection\ProcessorCompilerPass;
 use Symfony\AI\AiBundle\Exception\InvalidArgumentException;
-use Symfony\AI\AiBundle\Profiler\TraceableChat;
-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;
 use Symfony\AI\Chat\Bridge\Cache\MessageStore as CacheMessageStore;
 use Symfony\AI\Chat\Bridge\Cloudflare\MessageStore as CloudflareMessageStore;
@@ -142,8 +139,6 @@ use Symfony\Component\Security\Core\Authorization\AuthorizationCheckerInterface;
 use Symfony\Component\Translation\TranslatableMessage;
 use Symfony\Contracts\HttpClient\HttpClientInterface;
 
-use function Symfony\Component\String\u;
-
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
  */
@@ -153,6 +148,7 @@ final class AiBundle extends AbstractBundle
     {
         parent::build($container);
 
+        $container->addCompilerPass(new DebugCompilerPass());
         $container->addCompilerPass(new ProcessorCompilerPass());
     }
 
@@ -181,17 +177,6 @@ final class AiBundle extends AbstractBundle
         if (1 === \count($platforms)) {
             $builder->setAlias(PlatformInterface::class, reset($platforms));
         }
-        if ($builder->getParameter('kernel.debug')) {
-            foreach ($platforms as $platform) {
-                $traceablePlatformDefinition = (new Definition(TraceablePlatform::class))
-                    ->setDecoratedService($platform, priority: -1024)
-                    ->setArguments([new Reference('.inner')])
-                    ->addTag('ai.traceable_platform')
-                    ->addTag('kernel.reset', ['method' => 'reset']);
-                $suffix = u($platform)->after('ai.platform.')->toString();
-                $builder->setDefinition('ai.traceable_platform.'.$suffix, $traceablePlatformDefinition);
-            }
-        }
 
         if ([] !== ($config['agent'] ?? [])) {
             if (!ContainerBuilder::willBeAvailable('symfony/ai-agent', Agent::class, ['symfony/ai-bundle'])) {
@@ -256,21 +241,6 @@ final class AiBundle extends AbstractBundle
             $builder->setAlias(MessageStoreInterface::class, reset($messageStores));
         }
 
-        if ($builder->getParameter('kernel.debug')) {
-            foreach ($messageStores as $messageStore) {
-                $traceableMessageStoreDefinition = (new Definition(TraceableMessageStore::class))
-                    ->setDecoratedService($messageStore, priority: -1024)
-                    ->setArguments([
-                        new Reference('.inner'),
-                        new Reference(ClockInterface::class),
-                    ])
-                    ->addTag('ai.traceable_message_store')
-                    ->addTag('kernel.reset', ['method' => 'reset']);
-                $suffix = u($messageStore)->afterLast('.')->toString();
-                $builder->setDefinition('ai.traceable_message_store.'.$suffix, $traceableMessageStoreDefinition);
-            }
-        }
-
         if ([] === $messageStores) {
             $builder->removeDefinition('ai.command.setup_message_store');
             $builder->removeDefinition('ai.command.drop_message_store');
@@ -292,21 +262,6 @@ final class AiBundle extends AbstractBundle
             $builder->setAlias(ChatInterface::class, reset($chats));
         }
 
-        if ($builder->getParameter('kernel.debug')) {
-            foreach ($chats as $chat) {
-                $traceableChatDefinition = (new Definition(TraceableChat::class))
-                    ->setDecoratedService($chat, priority: -1024)
-                    ->setArguments([
-                        new Reference('.inner'),
-                        new Reference(ClockInterface::class),
-                    ])
-                    ->addTag('ai.traceable_chat')
-                    ->addTag('kernel.reset', ['method' => 'reset']);
-                $suffix = u($chat)->afterLast('.')->toString();
-                $builder->setDefinition('ai.traceable_chat.'.$suffix, $traceableChatDefinition);
-            }
-        }
-
         if ([] !== ($config['vectorizer'] ?? [])) {
             if (!ContainerBuilder::willBeAvailable('symfony/ai-store', StoreInterface::class, ['symfony/ai-bundle'])) {
                 throw new RuntimeException('Vectorizer configuration requires "symfony/ai-store" package. Try running "composer require symfony/ai-store".');
@@ -1189,16 +1144,6 @@ final class AiBundle extends AbstractBundle
                     ->setDecoratedService('ai.toolbox.'.$name, priority: -1024);
             }
 
-            if ($container->getParameter('kernel.debug')) {
-                $traceableToolboxDefinition = (new Definition('ai.traceable_toolbox.'.$name))
-                    ->setClass(TraceableToolbox::class)
-                    ->setArguments([new Reference('.inner')])
-                    ->setDecoratedService('ai.toolbox.'.$name, priority: -1024)
-                    ->addTag('ai.traceable_toolbox')
-                    ->addTag('kernel.reset', ['method' => 'reset']);
-                $container->setDefinition('ai.traceable_toolbox.'.$name, $traceableToolboxDefinition);
-            }
-
             $toolProcessorDefinition = (new ChildDefinition('ai.tool.agent_processor.abstract'))
                 ->replaceArgument(0, new Reference('ai.toolbox.'.$name))
                 ->replaceArgument(3, $config['keep_tool_messages'])

src/ai-bundle/src/DependencyInjection/DebugCompilerPass.php

@@ -0,0 +1,80 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AiBundle\DependencyInjection;
+
+use Symfony\AI\AiBundle\Profiler\TraceableChat;
+use Symfony\AI\AiBundle\Profiler\TraceableMessageStore;
+use Symfony\AI\AiBundle\Profiler\TraceablePlatform;
+use Symfony\AI\AiBundle\Profiler\TraceableToolbox;
+use Symfony\Component\Clock\ClockInterface;
+use Symfony\Component\DependencyInjection\Compiler\CompilerPassInterface;
+use Symfony\Component\DependencyInjection\ContainerBuilder;
+use Symfony\Component\DependencyInjection\Definition;
+use Symfony\Component\DependencyInjection\Reference;
+
+use function Symfony\Component\String\u;
+
+class DebugCompilerPass implements CompilerPassInterface
+{
+    public function process(ContainerBuilder $container): void
+    {
+        if (!$container->getParameter('kernel.debug')) {
+            return;
+        }
+
+        foreach (array_keys($container->findTaggedServiceIds('ai.platform')) as $platform) {
+            $traceablePlatformDefinition = (new Definition(TraceablePlatform::class))
+                ->setDecoratedService($platform, priority: -1024)
+                ->setArguments([new Reference('.inner')])
+                ->addTag('ai.traceable_platform')
+                ->addTag('kernel.reset', ['method' => 'reset']);
+            $suffix = u($platform)->after('ai.platform.')->toString();
+            $container->setDefinition('ai.traceable_platform.'.$suffix, $traceablePlatformDefinition);
+        }
+
+        foreach (array_keys($container->findTaggedServiceIds('ai.message_store')) as $messageStore) {
+            $traceableMessageStoreDefinition = (new Definition(TraceableMessageStore::class))
+                ->setDecoratedService($messageStore, priority: -1024)
+                ->setArguments([
+                    new Reference('.inner'),
+                    new Reference(ClockInterface::class),
+                ])
+                ->addTag('ai.traceable_message_store')
+                ->addTag('kernel.reset', ['method' => 'reset']);
+            $suffix = u($messageStore)->afterLast('.')->toString();
+            $container->setDefinition('ai.traceable_message_store.'.$suffix, $traceableMessageStoreDefinition);
+        }
+
+        foreach (array_keys($container->findTaggedServiceIds('ai.chat')) as $chat) {
+            $traceableChatDefinition = (new Definition(TraceableChat::class))
+                ->setDecoratedService($chat, priority: -1024)
+                ->setArguments([
+                    new Reference('.inner'),
+                    new Reference(ClockInterface::class),
+                ])
+                ->addTag('ai.traceable_chat')
+                ->addTag('kernel.reset', ['method' => 'reset']);
+            $suffix = u($chat)->afterLast('.')->toString();
+            $container->setDefinition('ai.traceable_chat.'.$suffix, $traceableChatDefinition);
+        }
+
+        foreach (array_keys($container->findTaggedServiceIds('ai.toolbox')) as $toolbox) {
+            $traceableToolboxDefinition = (new Definition(TraceableToolbox::class))
+                ->setDecoratedService($toolbox, priority: -1024)
+                ->setArguments([new Reference('.inner')])
+                ->addTag('ai.traceable_toolbox')
+                ->addTag('kernel.reset', ['method' => 'reset']);
+            $suffix = u($toolbox)->afterLast('.')->toString();
+            $container->setDefinition('ai.traceable_toolbox.'.$suffix, $traceableToolboxDefinition);
+        }
+    }
+}

src/ai-bundle/tests/DependencyInjection/AiBundleTest.php

@@ -27,6 +27,7 @@ use Symfony\AI\Agent\Memory\StaticMemoryProvider;
 use Symfony\AI\Agent\MultiAgent\Handoff;
 use Symfony\AI\Agent\MultiAgent\MultiAgent;
 use Symfony\AI\AiBundle\AiBundle;
+use Symfony\AI\AiBundle\DependencyInjection\DebugCompilerPass;
 use Symfony\AI\AiBundle\Exception\InvalidArgumentException;
 use Symfony\AI\Chat\ChatInterface;
 use Symfony\AI\Chat\ManagedStoreInterface as ManagedMessageStoreInterface;
@@ -7558,6 +7559,8 @@ class AiBundleTest extends TestCase
         $extension = (new AiBundle())->getContainerExtension();
         $extension->load($configuration, $container);
 
+        (new DebugCompilerPass())->process($container);
+
         return $container;
     }
 

src/ai-bundle/tests/DependencyInjection/DebugCompilerPassTest.php

@@ -0,0 +1,83 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AiBundle\Tests\DependencyInjection;
+
+use PHPUnit\Framework\TestCase;
+use Symfony\AI\AiBundle\DependencyInjection\DebugCompilerPass;
+use Symfony\AI\AiBundle\Profiler\TraceableChat;
+use Symfony\AI\AiBundle\Profiler\TraceableMessageStore;
+use Symfony\AI\AiBundle\Profiler\TraceablePlatform;
+use Symfony\AI\AiBundle\Profiler\TraceableToolbox;
+use Symfony\Component\Clock\ClockInterface;
+use Symfony\Component\DependencyInjection\ContainerBuilder;
+use Symfony\Component\DependencyInjection\Reference;
+
+class DebugCompilerPassTest extends TestCase
+{
+    public function testProcessAddsTraceableDefinitionsInDebug()
+    {
+        $container = new ContainerBuilder();
+        $container->setParameter('kernel.debug', true);
+
+        $container->register('ai.platform.azure.eu', \stdClass::class)->addTag('ai.platform');
+        $container->register('ai.message_store.memory.main', \stdClass::class)->addTag('ai.message_store');
+        $container->register('ai.chat.main', \stdClass::class)->addTag('ai.chat');
+        $container->register('ai.toolbox.my_agent', \stdClass::class)->addTag('ai.toolbox');
+
+        (new DebugCompilerPass())->process($container);
+
+        $traceablePlatform = $container->getDefinition('ai.traceable_platform.azure.eu');
+        $this->assertSame(TraceablePlatform::class, $traceablePlatform->getClass());
+        $this->assertSame(['ai.platform.azure.eu', null, -1024], $traceablePlatform->getDecoratedService());
+        $this->assertEquals([new Reference('.inner')], $traceablePlatform->getArguments());
+        $this->assertTrue($traceablePlatform->hasTag('ai.traceable_platform'));
+        $this->assertSame([['method' => 'reset']], $traceablePlatform->getTag('kernel.reset'));
+
+        $traceableMessageStore = $container->getDefinition('ai.traceable_message_store.main');
+        $this->assertSame(TraceableMessageStore::class, $traceableMessageStore->getClass());
+        $this->assertSame(['ai.message_store.memory.main', null, -1024], $traceableMessageStore->getDecoratedService());
+        $this->assertEquals([new Reference('.inner'), new Reference(ClockInterface::class)], $traceableMessageStore->getArguments());
+        $this->assertTrue($traceableMessageStore->hasTag('ai.traceable_message_store'));
+        $this->assertSame([['method' => 'reset']], $traceableMessageStore->getTag('kernel.reset'));
+
+        $traceableChat = $container->getDefinition('ai.traceable_chat.main');
+        $this->assertSame(TraceableChat::class, $traceableChat->getClass());
+        $this->assertSame(['ai.chat.main', null, -1024], $traceableChat->getDecoratedService());
+        $this->assertEquals([new Reference('.inner'), new Reference(ClockInterface::class)], $traceableChat->getArguments());
+        $this->assertTrue($traceableChat->hasTag('ai.traceable_chat'));
+        $this->assertSame([['method' => 'reset']], $traceableChat->getTag('kernel.reset'));
+
+        $traceableToolbox = $container->getDefinition('ai.traceable_toolbox.my_agent');
+        $this->assertSame(TraceableToolbox::class, $traceableToolbox->getClass());
+        $this->assertSame(['ai.toolbox.my_agent', null, -1024], $traceableToolbox->getDecoratedService());
+        $this->assertEquals([new Reference('.inner')], $traceableToolbox->getArguments());
+        $this->assertTrue($traceableToolbox->hasTag('ai.traceable_toolbox'));
+        $this->assertSame([['method' => 'reset']], $traceableToolbox->getTag('kernel.reset'));
+    }
+
+    public function testProcessSkipsWhenDebugDisabled()
+    {
+        $container = new ContainerBuilder();
+        $container->setParameter('kernel.debug', false);
+        $container->register('ai.platform.anthropic', \stdClass::class)->addTag('ai.platform');
+        $container->register('ai.message_store.memory.main', \stdClass::class)->addTag('ai.message_store');
+        $container->register('ai.chat.main', \stdClass::class)->addTag('ai.chat');
+        $container->register('ai.toolbox.my_agent', \stdClass::class)->addTag('ai.toolbox');
+
+        (new DebugCompilerPass())->process($container);
+
+        $this->assertFalse($container->hasDefinition('ai.traceable_platform.anthropic'));
+        $this->assertFalse($container->hasDefinition('ai.traceable_message_store.main'));
+        $this->assertFalse($container->hasDefinition('ai.traceable_chat.main'));
+        $this->assertFalse($container->hasDefinition('ai.traceable_toolbox.my_agent'));
+    }
+}