feat(profiler): trace and display AgentInterface calls

Author: Griffon-Weglot

src/ai-bundle/config/services.php

@@ -11,6 +11,7 @@
 
 namespace Symfony\Component\DependencyInjection\Loader\Configurator;
 
+use Symfony\AI\Agent\AgentInterface;
 use Symfony\AI\Agent\StructuredOutput\AgentProcessor as StructureOutputProcessor;
 use Symfony\AI\Agent\StructuredOutput\ResponseFormatFactory;
 use Symfony\AI\Agent\StructuredOutput\ResponseFormatFactoryInterface;
@@ -23,6 +24,7 @@
 use Symfony\AI\Agent\Toolbox\ToolResultConverter;
 use Symfony\AI\AiBundle\Command\ChatCommand;
 use Symfony\AI\AiBundle\Profiler\DataCollector;
+use Symfony\AI\AiBundle\Profiler\TraceableAgent;
 use Symfony\AI\AiBundle\Profiler\TraceableToolbox;
 use Symfony\AI\AiBundle\Security\EventListener\IsGrantedToolAttributeListener;
 use Symfony\AI\Platform\Bridge\Anthropic\Contract\AnthropicContract;
@@ -127,6 +129,13 @@
             ])
             ->tag('kernel.event_listener')
 
+        ->set('ai.traceable_agent', TraceableAgent::class)
+        ->decorate(AgentInterface::class, priority: 5)
+        ->args([
+            service('.inner'),
+            service('ai.data_collector'),
+        ])
+
         // profiler
         ->set('ai.data_collector', DataCollector::class)
             ->args([

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

@@ -37,6 +37,11 @@ final class DataCollector extends AbstractDataCollector implements LateDataColle
      */
     private readonly array $toolboxes;
 
+    /**
+     * @var list<array{method: string, duration: float, input: mixed, result: mixed, error: ?\Throwable}>
+     */
+    private array $collectedChatCalls = [];
+
     /**
      * @param TraceablePlatform[] $platforms
      * @param TraceableToolbox[]  $toolboxes
@@ -52,7 +57,6 @@ public function __construct(
 
     public function collect(Request $request, Response $response, ?\Throwable $exception = null): void
     {
-        $this->lateCollect();
     }
 
     public function lateCollect(): void
@@ -61,6 +65,18 @@ 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)),
+            'chat_calls' => $this->cloneVar($this->collectedChatCalls),
+        ];
+    }
+
+    public function collectChatCall(string $method, float $duration, mixed $input, mixed $result, ?\Throwable $error): void
+    {
+        $this->collectedChatCalls[] = [
+            'method' => $method,
+            'duration' => $duration,
+            'input' => $input,
+            'result' => $result,
+            'error' => $error,
         ];
     }
 
@@ -93,6 +109,27 @@ public function getToolCalls(): array
         return $this->data['tool_calls'] ?? [];
     }
 
+    /**
+     * @return list<array{method: string, duration: float, input: mixed, result: mixed, error: ?\Throwable}>
+     */
+    public function getChatCalls(): array
+    {
+        if (!isset($this->data['chat_calls'])) {
+            return [];
+        }
+
+        /** @var list<array{method: string, duration: float, input: mixed, result: mixed, error: ?\Throwable}> $chatCalls */
+        $chatCalls = $this->data['chat_calls']->getValue(true);
+
+        return $chatCalls;
+    }
+
+    public function reset(): void
+    {
+        $this->data = [];
+        $this->collectedChatCalls = [];
+    }
+
     /**
      * @return array{
      *     model: Model,

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

@@ -0,0 +1,60 @@
+<?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\Agent\AgentInterface;
+use Symfony\AI\Platform\Message\MessageBag;
+use Symfony\AI\Platform\Result\ResultInterface;
+use Symfony\Contracts\Service\ResetInterface;
+
+final class TraceableAgent implements AgentInterface, ResetInterface
+{
+    public function __construct(
+        private readonly AgentInterface $decorated,
+        private readonly DataCollector $collector,
+    ) {
+    }
+
+    public function call(MessageBag $messages, array $options = []): ResultInterface
+    {
+        $startTime = microtime(true);
+        $error = null;
+        $response = null;
+
+        try {
+            return $response = $this->decorated->call($messages, $options);
+        } catch (\Throwable $e) {
+            $error = $e;
+            throw $e;
+        } finally {
+            $this->collector->collectChatCall(
+                'call',
+                microtime(true) - $startTime,
+                $messages,
+                $response,
+                $error
+            );
+        }
+    }
+
+    public function reset(): void
+    {
+        if ($this->decorated instanceof ResetInterface) {
+            $this->decorated->reset();
+        }
+    }
+
+    public function getName(): string
+    {
+        return 'TraceableAgent';
+    }
+}