[AIBundle] Fix stream result profiling

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

@@ -18,14 +18,15 @@
 use Symfony\Component\DependencyInjection\Attribute\TaggedIterator;
 use Symfony\Component\HttpFoundation\Request;
 use Symfony\Component\HttpFoundation\Response;
+use Symfony\Component\HttpKernel\DataCollector\LateDataCollectorInterface;
 
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
  *
  * @phpstan-import-type PlatformCallData from TraceablePlatform
  * @phpstan-import-type ToolCallData from TraceableToolbox
  */
-final class DataCollector extends AbstractDataCollector
+final class DataCollector extends AbstractDataCollector implements LateDataCollectorInterface
 {
     /**
      * @var TraceablePlatform[]
@@ -53,6 +54,11 @@ public function __construct(
     }
 
     public function collect(Request $request, Response $response, ?\Throwable $exception = null): void
+    {
+        $this->lateCollect();
+    }
+
+    public function lateCollect(): void
     {
         $this->data = [
             'tools' => $this->defaultToolBox->getTools(),
@@ -102,7 +108,14 @@ private function awaitCallResults(TraceablePlatform $platform): array
     {
         $calls = $platform->calls;
         foreach ($calls as $key => $call) {
-            $call['result'] = $call['result']->await()->getContent();
+            $result = $call['result']->await();
+
+            if (isset($platform->resultCache[$result])) {
+                $call['result'] = $platform->resultCache[$result];
+            } else {
+                $call['result'] = $result->getContent();
+            }
+
             $calls[$key] = $call;
         }
 

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

@@ -14,7 +14,9 @@
 use Symfony\AI\Platform\Message\Content\File;
 use Symfony\AI\Platform\Model;
 use Symfony\AI\Platform\PlatformInterface;
+use Symfony\AI\Platform\Result\ResultInterface;
 use Symfony\AI\Platform\Result\ResultPromise;
+use Symfony\AI\Platform\Result\StreamResult;
 
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
@@ -32,27 +34,50 @@ final class TraceablePlatform implements PlatformInterface
      * @var PlatformCallData[]
      */
     public array $calls = [];
+    /**
+     * @var \WeakMap<ResultInterface, string>
+     */
+    public \WeakMap $resultCache;
 
     public function __construct(
         private readonly PlatformInterface $platform,
     ) {
+        $this->resultCache = new \WeakMap();
     }
 
     public function invoke(Model $model, array|string|object $input, array $options = []): ResultPromise
     {
-        $result = $this->platform->invoke($model, $input, $options);
+        $resultPromise = $this->platform->invoke($model, $input, $options);
 
         if ($input instanceof File) {
             $input = $input::class.': '.$input->getFormat();
         }
 
+        if ($options['stream'] ?? false) {
+            $originalStream = $resultPromise->asStream();
+            $resultPromise = new ResultPromise(fn () => $this->createTraceableStreamResult($originalStream), $resultPromise->getRawResult(), $options);
+        }
+
         $this->calls[] = [
             'model' => $model,
             'input' => \is_object($input) ? clone $input : $input,
             'options' => $options,
-            'result' => $result,
+            'result' => $resultPromise,
         ];
 
-        return $result;
+        return $resultPromise;
+    }
+
+    private function createTraceableStreamResult(\Generator $originalStream): StreamResult
+    {
+        return $result = new StreamResult((function () use (&$result, $originalStream) {
+            $this->resultCache[$result] = '';
+            foreach ($originalStream as $chunk) {
+                yield $chunk;
+                if (\is_string($chunk)) {
+                    $this->resultCache[$result] .= $chunk;
+                }
+            }
+        })());
     }
 }

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

@@ -0,0 +1,77 @@
+<?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\Profiler;
+
+use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\UsesClass;
+use PHPUnit\Framework\TestCase;
+use Symfony\AI\Agent\Toolbox\ToolboxInterface;
+use Symfony\AI\AIBundle\Profiler\DataCollector;
+use Symfony\AI\AIBundle\Profiler\TraceablePlatform;
+use Symfony\AI\Platform\Message\Content\Text;
+use Symfony\AI\Platform\Message\Message;
+use Symfony\AI\Platform\Message\MessageBag;
+use Symfony\AI\Platform\Model;
+use Symfony\AI\Platform\PlatformInterface;
+use Symfony\AI\Platform\Result\RawResultInterface;
+use Symfony\AI\Platform\Result\ResultPromise;
+use Symfony\AI\Platform\Result\StreamResult;
+use Symfony\AI\Platform\Result\TextResult;
+
+#[CoversClass(DataCollector::class)]
+#[UsesClass(TraceablePlatform::class)]
+#[UsesClass(ResultPromise::class)]
+class DataCollectorTest extends TestCase
+{
+    public function testCollectsDataForNonStreamingResponse()
+    {
+        $platform = $this->createMock(PlatformInterface::class);
+        $traceablePlatform = new TraceablePlatform($platform);
+        $messageBag = new MessageBag(Message::ofUser(new Text('Hello')));
+        $result = new TextResult('Assistant response');
+
+        $platform->method('invoke')->willReturn(new ResultPromise(static fn () => $result, $this->createStub(RawResultInterface::class)));
+
+        $result = $traceablePlatform->invoke($this->createStub(Model::class), $messageBag, ['stream' => false]);
+        $this->assertSame('Assistant response', $result->asText());
+
+        $dataCollector = new DataCollector([$traceablePlatform], $this->createStub(ToolboxInterface::class), []);
+        $dataCollector->lateCollect();
+
+        $this->assertCount(1, $dataCollector->getPlatformCalls());
+        $this->assertSame('Assistant response', $dataCollector->getPlatformCalls()[0]['result']);
+    }
+
+    public function testCollectsDataForStreamingResponse()
+    {
+        $platform = $this->createMock(PlatformInterface::class);
+        $traceablePlatform = new TraceablePlatform($platform);
+        $messageBag = new MessageBag(Message::ofUser(new Text('Hello')));
+        $result = new StreamResult(
+            (function () {
+                yield 'Assistant ';
+                yield 'response';
+            })(),
+        );
+
+        $platform->method('invoke')->willReturn(new ResultPromise(static fn () => $result, $this->createStub(RawResultInterface::class)));
+
+        $result = $traceablePlatform->invoke($this->createStub(Model::class), $messageBag, ['stream' => true]);
+        $this->assertSame('Assistant response', implode('', iterator_to_array($result->asStream())));
+
+        $dataCollector = new DataCollector([$traceablePlatform], $this->createStub(ToolboxInterface::class), []);
+        $dataCollector->lateCollect();
+
+        $this->assertCount(1, $dataCollector->getPlatformCalls());
+        $this->assertSame('Assistant response', $dataCollector->getPlatformCalls()[0]['result']);
+    }
+}