[AI Bundle] clean tracers memory after each request

Author: santysisi

src/ai-bundle/CHANGELOG.md

@@ -1,6 +1,11 @@
 CHANGELOG
 =========
 
+0.3
+---
+
+* Add `kernel.reset` support to tracers to clear call logs between requests
+
 0.2
 ---
 

src/ai-bundle/composer.json

@@ -25,6 +25,7 @@
         "symfony/console": "^7.3|^8.0",
         "symfony/dependency-injection": "^7.3|^8.0",
         "symfony/framework-bundle": "^7.3|^8.0",
+        "symfony/service-contracts": "^2.5|^3",
         "symfony/string": "^7.3|^8.0"
     },
     "require-dev": {

src/ai-bundle/src/AiBundle.php

@@ -182,7 +182,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
                 $traceablePlatformDefinition = (new Definition(TraceablePlatform::class))
                     ->setDecoratedService($platform, priority: -1024)
                     ->setArguments([new Reference('.inner')])
-                    ->addTag('ai.traceable_platform');
+                    ->addTag('ai.traceable_platform')
+                    ->addTag('kernel.reset', ['method' => 'reset']);
                 $suffix = u($platform)->after('ai.platform.')->toString();
                 $builder->setDefinition('ai.traceable_platform.'.$suffix, $traceablePlatformDefinition);
             }
@@ -259,7 +260,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
                         new Reference('.inner'),
                         new Reference(ClockInterface::class),
                     ])
-                    ->addTag('ai.traceable_message_store');
+                    ->addTag('ai.traceable_message_store')
+                    ->addTag('kernel.reset', ['method' => 'reset']);
                 $suffix = u($messageStore)->afterLast('.')->toString();
                 $builder->setDefinition('ai.traceable_message_store.'.$suffix, $traceableMessageStoreDefinition);
             }
@@ -294,7 +296,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
                         new Reference('.inner'),
                         new Reference(ClockInterface::class),
                     ])
-                    ->addTag('ai.traceable_chat');
+                    ->addTag('ai.traceable_chat')
+                    ->addTag('kernel.reset', ['method' => 'reset']);
                 $suffix = u($chat)->afterLast('.')->toString();
                 $builder->setDefinition('ai.traceable_chat.'.$suffix, $traceableChatDefinition);
             }
@@ -1121,7 +1124,8 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
                     ->setClass(TraceableToolbox::class)
                     ->setArguments([new Reference('.inner')])
                     ->setDecoratedService('ai.toolbox.'.$name, priority: -1024)
-                    ->addTag('ai.traceable_toolbox');
+                    ->addTag('ai.traceable_toolbox')
+                    ->addTag('kernel.reset', ['method' => 'reset']);
                 $container->setDefinition('ai.traceable_toolbox.'.$name, $traceableToolboxDefinition);
             }
 

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

@@ -16,6 +16,7 @@
 use Symfony\AI\Platform\Message\MessageBag;
 use Symfony\AI\Platform\Message\UserMessage;
 use Symfony\Component\Clock\ClockInterface;
+use Symfony\Contracts\Service\ResetInterface;
 
 /**
  * @author Guillaume Loulier <personal@guillaumeloulier.fr>
@@ -27,7 +28,7 @@
  *      saved_at: \DateTimeImmutable,
  *  }
  */
-final class TraceableChat implements ChatInterface
+final class TraceableChat implements ChatInterface, ResetInterface
 {
     /**
      * @var array<int, array{
@@ -66,4 +67,9 @@ public function submit(UserMessage $message): AssistantMessage
 
         return $this->chat->submit($message);
     }
+
+    public function reset(): void
+    {
+        $this->calls = [];
+    }
 }

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

@@ -15,6 +15,7 @@
 use Symfony\AI\Chat\MessageStoreInterface;
 use Symfony\AI\Platform\Message\MessageBag;
 use Symfony\Component\Clock\ClockInterface;
+use Symfony\Contracts\Service\ResetInterface;
 
 /**
  * @author Guillaume Loulier <personal@guillaumeloulier.fr>
@@ -24,7 +25,7 @@
  *      saved_at: \DateTimeImmutable,
  *  }
  */
-final class TraceableMessageStore implements ManagedStoreInterface, MessageStoreInterface
+final class TraceableMessageStore implements ManagedStoreInterface, MessageStoreInterface, ResetInterface
 {
     /**
      * @var MessageStoreData[]
@@ -69,4 +70,9 @@ public function drop(): void
 
         $this->messageStore->drop();
     }
+
+    public function reset(): void
+    {
+        $this->calls = [];
+    }
 }

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

@@ -19,6 +19,7 @@
 use Symfony\AI\Platform\Result\DeferredResult;
 use Symfony\AI\Platform\Result\ResultInterface;
 use Symfony\AI\Platform\Result\StreamResult;
+use Symfony\Contracts\Service\ResetInterface;
 
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
@@ -30,7 +31,7 @@
  *     result: DeferredResult,
  * }
  */
-final class TraceablePlatform implements PlatformInterface
+final class TraceablePlatform implements PlatformInterface, ResetInterface
 {
     /**
      * @var PlatformCallData[]
@@ -74,6 +75,12 @@ public function getModelCatalog(): ModelCatalogInterface
         return $this->platform->getModelCatalog();
     }
 
+    public function reset(): void
+    {
+        $this->calls = [];
+        $this->resultCache = new \WeakMap();
+    }
+
     private function createTraceableStreamResult(DeferredResult $originalStream): StreamResult
     {
         return $result = new StreamResult((function () use (&$result, $originalStream) {

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

@@ -14,11 +14,12 @@
 use Symfony\AI\Agent\Toolbox\ToolboxInterface;
 use Symfony\AI\Agent\Toolbox\ToolResult;
 use Symfony\AI\Platform\Result\ToolCall;
+use Symfony\Contracts\Service\ResetInterface;
 
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
  */
-final class TraceableToolbox implements ToolboxInterface
+final class TraceableToolbox implements ToolboxInterface, ResetInterface
 {
     /**
      * @var ToolResult[]
@@ -39,4 +40,9 @@ public function execute(ToolCall $toolCall): ToolResult
     {
         return $this->calls[] = $this->toolbox->execute($toolCall);
     }
+
+    public function reset(): void
+    {
+        $this->calls = [];
+    }
 }

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

@@ -60,4 +60,17 @@ public function testInitializationMessageBagCanBeRetrieved()
         $this->assertInstanceOf(UserMessage::class, $traceableChat->calls[1]['message']);
         $this->assertInstanceOf(\DateTimeImmutable::class, $traceableChat->calls[1]['saved_at']);
     }
+
+    public function testResetClearsCalls()
+    {
+        $agent = $this->createStub(AgentInterface::class);
+        $chat = new Chat($agent, new InMemoryStore());
+        $traceableChat = new TraceableChat($chat, new MonotonicClock());
+
+        $traceableChat->initiate(new MessageBag());
+        $this->assertCount(1, $traceableChat->calls);
+
+        $traceableChat->reset();
+        $this->assertCount(0, $traceableChat->calls);
+    }
 }

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

@@ -42,4 +42,16 @@ public function testSubmittedMessageBagCanBeRetrieved()
         $this->assertCount(1, $calls[0]['bag']);
         $this->assertInstanceOf(\DateTimeImmutable::class, $calls[0]['saved_at']);
     }
+
+    public function testResetClearsCalls()
+    {
+        $messageStore = new InMemoryStore();
+        $traceableMessageStore = new TraceableMessageStore($messageStore, new MonotonicClock());
+
+        $traceableMessageStore->save(new MessageBag());
+        $this->assertCount(1, $traceableMessageStore->calls);
+
+        $traceableMessageStore->reset();
+        $this->assertCount(0, $traceableMessageStore->calls);
+    }
 }

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

@@ -0,0 +1,45 @@
+<?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\TestCase;
+use Symfony\AI\AiBundle\Profiler\TraceablePlatform;
+use Symfony\AI\Platform\PlainConverter;
+use Symfony\AI\Platform\PlatformInterface;
+use Symfony\AI\Platform\Result\DeferredResult;
+use Symfony\AI\Platform\Result\RawResultInterface;
+use Symfony\AI\Platform\Result\TextResult;
+
+final class TraceablePlatformTest extends TestCase
+{
+    public function testResetClearsCallsAndResultCache()
+    {
+        $platform = $this->createStub(PlatformInterface::class);
+        $traceablePlatform = new TraceablePlatform($platform);
+        $result = new TextResult('Assistant response');
+
+        $platform->method('invoke')->willReturn(new DeferredResult(new PlainConverter($result), $this->createStub(RawResultInterface::class)));
+
+        $traceablePlatform->invoke('gpt-4o', 'Hello');
+        $this->assertCount(1, $traceablePlatform->calls);
+        $this->assertSame('gpt-4o', $traceablePlatform->calls[0]['model']);
+        $this->assertSame('Hello', $traceablePlatform->calls[0]['input']);
+
+        $oldCache = $traceablePlatform->resultCache;
+
+        $traceablePlatform->reset();
+
+        $this->assertCount(0, $traceablePlatform->calls);
+        $this->assertNotSame($oldCache, $traceablePlatform->resultCache);
+        $this->assertInstanceOf(\WeakMap::class, $traceablePlatform->resultCache);
+    }
+}