[AI Bundle] Trace and display AgentInterface calls in Profiler

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\AgentCallCommand;
 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\AiMlApi\ModelCatalog as AiMlApiModelCatalog;
@@ -170,6 +172,12 @@
             ->tag('kernel.event_listener')
 
         // profiler
+        ->set('ai.traceable_agent', TraceableAgent::class)
+            ->decorate(AgentInterface::class, priority: 5)
+            ->args([
+                service('.inner'),
+                service('ai.data_collector'),
+            ])
         ->set('ai.data_collector', DataCollector::class)
             ->args([
                 tagged_iterator('ai.traceable_platform'),

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

@@ -18,11 +18,11 @@
 use Symfony\Component\HttpFoundation\Request;
 use Symfony\Component\HttpFoundation\Response;
 use Symfony\Component\HttpKernel\DataCollector\LateDataCollectorInterface;
+use Symfony\Component\VarDumper\Cloner\Data;
 
 /**
  * @author Christopher Hertel <[email protected]>
  *
- * @phpstan-import-type PlatformCallData from TraceablePlatform
  * @phpstan-import-type ToolCallData from TraceableToolbox
  */
 final class DataCollector extends AbstractDataCollector implements LateDataCollectorInterface
@@ -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 $collectedCalls = [];
+
     /**
      * @param TraceablePlatform[] $platforms
      * @param TraceableToolbox[]  $toolboxes
@@ -52,15 +57,55 @@ public function __construct(
 
     public function collect(Request $request, Response $response, ?\Throwable $exception = null): void
     {
-        $this->lateCollect();
     }
 
     public function lateCollect(): void
     {
+        $platformCalls = [];
+        foreach ($this->platforms as $platform) {
+            $calls = $platform->calls;
+            foreach ($calls as $call) {
+                $result = $call['result']->await();
+                if (isset($platform->resultCache[$result])) {
+                    $call['result'] = $platform->resultCache[$result];
+                } else {
+                    $call['result'] = $result->getContent();
+                }
+
+                $call['model'] = $this->cloneVar($call['model']);
+                $call['input'] = $this->cloneVar($call['input']);
+                $call['options'] = $this->cloneVar($call['options']);
+                $call['result'] = $this->cloneVar($call['result']);
+
+                $platformCalls[] = $call;
+            }
+        }
+
+        $toolCalls = [];
+        foreach ($this->toolboxes as $toolbox) {
+            foreach ($toolbox->calls as $call) {
+                $call['call'] = $this->cloneVar($call['call']);
+                $call['result'] = $this->cloneVar($call['result']);
+                $toolCalls[] = $call;
+            }
+        }
+
         $this->data = [
             '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)),
+            'platform_calls' => $platformCalls,
+            'tool_calls' => $toolCalls,
+            'agent_calls' => $this->collectedCalls,
+        ];
+    }
+
+    public function collectAgentCall(string $method, float $duration, mixed $input, mixed $result, ?\Throwable $error): void
+    {
+        $this->collectedCalls[] = [
+            'method' => $method,
+            'duration' => $duration,
+            'input' => $this->cloneVar($input),
+            'result' => $this->cloneVar($result),
+            'error' => $this->cloneVar($error),
         ];
     }
 
@@ -70,7 +115,12 @@ public static function getTemplate(): string
     }
 
     /**
-     * @return PlatformCallData[]
+     * @return array{
+     *     model: Data,
+     *     input: Data,
+     *     options: Data,
+     *     result: Data
+     * }[]
      */
     public function getPlatformCalls(): array
     {
@@ -94,28 +144,16 @@ public function getToolCalls(): array
     }
 
     /**
-     * @return array{
-     *     model: string,
-     *     input: array<mixed>|string|object,
-     *     options: array<string, mixed>,
-     *     result: string|iterable<mixed>|object|null
-     * }[]
+     * @return list<array{method: string, duration: float, input: mixed, result: mixed, error: ?\Throwable}>
      */
-    private function awaitCallResults(TraceablePlatform $platform): array
+    public function getAgentCalls(): array
     {
-        $calls = $platform->calls;
-        foreach ($calls as $key => $call) {
-            $result = $call['result']->await();
-
-            if (isset($platform->resultCache[$result])) {
-                $call['result'] = $platform->resultCache[$result];
-            } else {
-                $call['result'] = $result->getContent();
-            }
-
-            $calls[$key] = $call;
-        }
+        return $this->data['agent_calls'] ?? [];
+    }
 
-        return $calls;
+    public function reset(): void
+    {
+        $this->data = [];
+        $this->collectedCalls = [];
     }
 }

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->collectAgentCall(
+                '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';
+    }
+}