Add metadata propagation from source streams

Author: oesteve

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

@@ -56,8 +56,7 @@ public function invoke(string $model, array|string|object $input, array $options
         }
 
         if ($options['stream'] ?? false) {
-            $originalStream = $deferredResult->asStream();
-            $deferredResult = new DeferredResult(new PlainConverter($this->createTraceableStreamResult($originalStream)), $deferredResult->getRawResult(), $options);
+            $deferredResult = new DeferredResult(new PlainConverter($this->createTraceableStreamResult($deferredResult)), $deferredResult->getRawResult(), $options);
         }
 
         $this->calls[] = [
@@ -75,16 +74,20 @@ public function getModelCatalog(): ModelCatalogInterface
         return $this->platform->getModelCatalog();
     }
 
-    private function createTraceableStreamResult(\Generator $originalStream): StreamResult
+    private function createTraceableStreamResult(DeferredResult $sourceResult): StreamResult
     {
-        return $result = new StreamResult((function () use (&$result, $originalStream) {
+        return $result = new StreamResult((function () use (&$result, $sourceResult) {
             $this->resultCache[$result] = '';
-            foreach ($originalStream as $chunk) {
+            foreach ($sourceResult->asStream() as $chunk) {
                 yield $chunk;
                 if (\is_string($chunk)) {
                     $this->resultCache[$result] .= $chunk;
                 }
             }
+
+            foreach ($sourceResult->getResult()->getMetadata() as $key => $value) {
+                $result->getMetadata()->add($key, $value);
+            }
         })());
     }
 }

src/platform/src/Result/DeferredResult.php

@@ -119,7 +119,13 @@ public function asVectors(): array
      */
     public function asStream(): \Generator
     {
-        yield from $this->as(StreamResult::class)->getContent();
+        $streamResult = $this->as(StreamResult::class);
+
+        yield from $streamResult->getContent();
+
+        foreach ($streamResult->getMetadata() as $key => $value) {
+            $this->getResult()->getMetadata()->add($key, $value);
+        }
     }
 
     /**