[AI Bundle] add ResetInterface to prevent memory leaks

src/ai-bundle/CHANGELOG.md

@@ -6,6 +6,7 @@ CHANGELOG
 
  * Add `chats` data from `DataCollector` to the `data_collector.html.twig` template
  * [BC BREAK] Rename service ID prefix `ai.toolbox.{agent}.agent_wrapper.` to `ai.toolbox.{agent}.subagent.`
+ * Add `ResetInterface` support to `TraceableChat`, `TraceableMessageStore`, `TraceablePlatform` and `TraceableToolbox` to clear collected data 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);
             }
 
@@ -1508,11 +1512,9 @@ private function processStoreConfig(string $type, array $stores, ContainerBuilde
 
                 $definition = new Definition(InMemoryStore::class);
                 $definition
-                    ->setLazy(true)
                     ->setArguments($arguments)
-                    ->addTag('proxy', ['interface' => StoreInterface::class])
-                    ->addTag('proxy', ['interface' => ManagedStoreInterface::class])
-                    ->addTag('ai.store');
+                    ->addTag('ai.store')
+                    ->addTag('kernel.reset', ['method' => 'reset']);
 
                 $container->setDefinition('ai.store.'.$type.'.'.$name, $definition);
                 $container->registerAliasForArgument('ai.store.'.$type.'.'.$name, StoreInterface::class, $name);
@@ -2060,11 +2062,9 @@ private function processMessageStoreConfig(string $type, array $messageStores, C
             foreach ($messageStores as $name => $messageStore) {
                 $definition = new Definition(InMemoryMessageStore::class);
                 $definition
-                    ->setLazy(true)
                     ->setArgument(0, $messageStore['identifier'])
-                    ->addTag('proxy', ['interface' => MessageStoreInterface::class])
-                    ->addTag('proxy', ['interface' => ManagedMessageStoreInterface::class])
-                    ->addTag('ai.message_store');
+                    ->addTag('ai.message_store')
+                    ->addTag('kernel.reset', ['method' => 'reset']);
 
                 $container->setDefinition('ai.message_store.'.$type.'.'.$name, $definition);
                 $container->registerAliasForArgument('ai.message_store.'.$type.'.'.$name, MessageStoreInterface::class, $name);

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,12 @@ public function submit(UserMessage $message): AssistantMessage
 
         return $this->chat->submit($message);
     }
+
+    public function reset(): void
+    {
+        if ($this->chat instanceof ResetInterface) {
+            $this->chat->reset();
+        }
+        $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,12 @@ public function drop(): void
 
         $this->messageStore->drop();
     }
+
+    public function reset(): void
+    {
+        if ($this->messageStore instanceof ResetInterface) {
+            $this->messageStore->reset();
+        }
+        $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/DependencyInjection/AiBundleTest.php

@@ -1530,17 +1530,14 @@ public function testInMemoryStoreWithoutCustomStrategyCanBeConfigured()
         $definition = $container->getDefinition('ai.store.memory.my_memory_store_with_custom_strategy');
         $this->assertSame(InMemoryStore::class, $definition->getClass());
 
-        $this->assertTrue($definition->isLazy());
+        $this->assertFalse($definition->isLazy());
         $this->assertCount(1, $definition->getArguments());
         $this->assertInstanceOf(Definition::class, $definition->getArgument(0));
         $this->assertSame(DistanceCalculator::class, $definition->getArgument(0)->getClass());
 
-        $this->assertTrue($definition->hasTag('proxy'));
-        $this->assertSame([
-            ['interface' => StoreInterface::class],
-            ['interface' => ManagedStoreInterface::class],
-        ], $definition->getTag('proxy'));
+        $this->assertFalse($definition->hasTag('proxy'));
         $this->assertTrue($definition->hasTag('ai.store'));
+        $this->assertTrue($definition->hasTag('kernel.reset'));
 
         $this->assertTrue($container->hasAlias('.Symfony\AI\Store\StoreInterface $my_memory_store_with_custom_strategy'));
         $this->assertTrue($container->hasAlias('Symfony\AI\Store\StoreInterface $myMemoryStoreWithCustomStrategy'));
@@ -1569,17 +1566,14 @@ public function testInMemoryStoreWithCustomStrategyCanBeConfigured()
         $definition = $container->getDefinition('ai.store.memory.my_memory_store_with_custom_strategy');
         $this->assertSame(InMemoryStore::class, $definition->getClass());
 
-        $this->assertTrue($definition->isLazy());
+        $this->assertFalse($definition->isLazy());
         $this->assertCount(1, $definition->getArguments());
         $this->assertInstanceOf(Reference::class, $definition->getArgument(0));
         $this->assertSame('ai.store.distance_calculator.my_memory_store_with_custom_strategy', (string) $definition->getArgument(0));
 
-        $this->assertTrue($definition->hasTag('proxy'));
-        $this->assertSame([
-            ['interface' => StoreInterface::class],
-            ['interface' => ManagedStoreInterface::class],
-        ], $definition->getTag('proxy'));
+        $this->assertFalse($definition->hasTag('proxy'));
         $this->assertTrue($definition->hasTag('ai.store'));
+        $this->assertTrue($definition->hasTag('kernel.reset'));
 
         $this->assertTrue($container->hasAlias('.Symfony\AI\Store\StoreInterface $my_memory_store_with_custom_strategy'));
         $this->assertTrue($container->hasAlias('Symfony\AI\Store\StoreInterface $myMemoryStoreWithCustomStrategy'));
@@ -6709,14 +6703,12 @@ public function testMemoryMessageStoreCanBeConfiguredWithCustomKey()
 
         $definition = $container->getDefinition('ai.message_store.memory.custom');
 
-        $this->assertTrue($definition->isLazy());
+        $this->assertFalse($definition->isLazy());
         $this->assertSame('foo', $definition->getArgument(0));
 
-        $this->assertSame([
-            ['interface' => MessageStoreInterface::class],
-            ['interface' => ManagedMessageStoreInterface::class],
-        ], $definition->getTag('proxy'));
+        $this->assertFalse($definition->hasTag('proxy'));
         $this->assertTrue($definition->hasTag('ai.message_store'));
+        $this->assertTrue($definition->hasTag('kernel.reset'));
     }
 
     public function testMongoDbMessageStoreIsConfigured()

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);
+    }
+}

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

@@ -47,6 +47,18 @@ public function testExecute()
         $this->assertSame('tool_result', $traceableToolbox->calls[0]->getResult());
     }
 
+    public function testResetClearsCalls()
+    {
+        $toolbox = $this->createToolbox([]);
+        $traceableToolbox = new TraceableToolbox($toolbox);
+
+        $traceableToolbox->execute(new ToolCall('foo', '__invoke'));
+        $this->assertCount(1, $traceableToolbox->calls);
+
+        $traceableToolbox->reset();
+        $this->assertCount(0, $traceableToolbox->calls);
+    }
+
     /**
      * @param Tool[] $tools
      */

src/chat/CHANGELOG.md

@@ -1,6 +1,11 @@
 CHANGELOG
 =========
 
+0.4
+---
+
+ * Add `ResetInterface` support to in-memory store
+
 0.1
 ---
 

src/chat/composer.json

@@ -26,7 +26,8 @@
     "require": {
         "php": ">=8.2",
         "symfony/ai-agent": "^0.3",
-        "symfony/ai-platform": "^0.3"
+        "symfony/ai-platform": "^0.3",
+        "symfony/service-contracts": "^2.5|^3"
     },
     "require-dev": {
         "phpstan/phpstan": "^2.1",