Skip to content

Commit a749f10

Browse files
santysisisantysisi
committed
[AI Bundle][MCP Bundle][Chat][Store] add ResetInterface to prevent memory leaks
1 parent ecca0bb commit a749f10

25 files changed

+235
-39
lines changed

‎src/ai-bundle/CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ CHANGELOG
88
* [BC BREAK] Rename service ID prefix `ai.toolbox.{agent}.agent_wrapper.` to `ai.toolbox.{agent}.subagent.`
99
* Add support for `DocumentIndexer` when no loader is configured for an indexer
1010
* [BC BREAK] The `host_url` configuration key for `Ollama` has been renamed `endpoint`
11+
* Add `ResetInterface` support to `TraceableChat`, `TraceableMessageStore`, `TraceablePlatform` and `TraceableToolbox` to clear collected data between requests
1112

1213
0.2
1314
---

‎src/ai-bundle/composer.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
"symfony/console": "^7.3|^8.0",
2626
"symfony/dependency-injection": "^7.3|^8.0",
2727
"symfony/framework-bundle": "^7.3|^8.0",
28+
"symfony/service-contracts": "^2.5|^3",
2829
"symfony/string": "^7.3|^8.0"
2930
},
3031
"require-dev": {

‎src/ai-bundle/src/AiBundle.php‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
185185
$traceablePlatformDefinition = (new Definition(TraceablePlatform::class))
186186
->setDecoratedService($platform, priority: -1024)
187187
->setArguments([new Reference('.inner')])
188-
->addTag('ai.traceable_platform');
188+
->addTag('ai.traceable_platform')
189+
->addTag('kernel.reset', ['method' => 'reset']);
189190
$suffix = u($platform)->after('ai.platform.')->toString();
190191
$builder->setDefinition('ai.traceable_platform.'.$suffix, $traceablePlatformDefinition);
191192
}
@@ -262,7 +263,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
262263
new Reference('.inner'),
263264
new Reference(ClockInterface::class),
264265
])
265-
->addTag('ai.traceable_message_store');
266+
->addTag('ai.traceable_message_store')
267+
->addTag('kernel.reset', ['method' => 'reset']);
266268
$suffix = u($messageStore)->afterLast('.')->toString();
267269
$builder->setDefinition('ai.traceable_message_store.'.$suffix, $traceableMessageStoreDefinition);
268270
}
@@ -297,7 +299,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
297299
new Reference('.inner'),
298300
new Reference(ClockInterface::class),
299301
])
300-
->addTag('ai.traceable_chat');
302+
->addTag('ai.traceable_chat')
303+
->addTag('kernel.reset', ['method' => 'reset']);
301304
$suffix = u($chat)->afterLast('.')->toString();
302305
$builder->setDefinition('ai.traceable_chat.'.$suffix, $traceableChatDefinition);
303306
}
@@ -1150,7 +1153,8 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
11501153
->setClass(TraceableToolbox::class)
11511154
->setArguments([new Reference('.inner')])
11521155
->setDecoratedService('ai.toolbox.'.$name, priority: -1024)
1153-
->addTag('ai.traceable_toolbox');
1156+
->addTag('ai.traceable_toolbox')
1157+
->addTag('kernel.reset', ['method' => 'reset']);
11541158
$container->setDefinition('ai.traceable_toolbox.'.$name, $traceableToolboxDefinition);
11551159
}
11561160

@@ -1537,11 +1541,9 @@ private function processStoreConfig(string $type, array $stores, ContainerBuilde
15371541

15381542
$definition = new Definition(InMemoryStore::class);
15391543
$definition
1540-
->setLazy(true)
15411544
->setArguments($arguments)
1542-
->addTag('proxy', ['interface' => StoreInterface::class])
1543-
->addTag('proxy', ['interface' => ManagedStoreInterface::class])
1544-
->addTag('ai.store');
1545+
->addTag('ai.store')
1546+
->addTag('kernel.reset', ['method' => 'reset']);
15451547

15461548
$container->setDefinition('ai.store.'.$type.'.'.$name, $definition);
15471549
$container->registerAliasForArgument('ai.store.'.$type.'.'.$name, StoreInterface::class, $name);
@@ -2089,11 +2091,9 @@ private function processMessageStoreConfig(string $type, array $messageStores, C
20892091
foreach ($messageStores as $name => $messageStore) {
20902092
$definition = new Definition(InMemoryMessageStore::class);
20912093
$definition
2092-
->setLazy(true)
20932094
->setArgument(0, $messageStore['identifier'])
2094-
->addTag('proxy', ['interface' => MessageStoreInterface::class])
2095-
->addTag('proxy', ['interface' => ManagedMessageStoreInterface::class])
2096-
->addTag('ai.message_store');
2095+
->addTag('ai.message_store')
2096+
->addTag('kernel.reset', ['method' => 'reset']);
20972097

20982098
$container->setDefinition('ai.message_store.'.$type.'.'.$name, $definition);
20992099
$container->registerAliasForArgument('ai.message_store.'.$type.'.'.$name, MessageStoreInterface::class, $name);

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
use Symfony\AI\Platform\Message\MessageBag;
1717
use Symfony\AI\Platform\Message\UserMessage;
1818
use Symfony\Component\Clock\ClockInterface;
19+
use Symfony\Contracts\Service\ResetInterface;
1920

2021
/**
2122
* @author Guillaume Loulier <personal@guillaumeloulier.fr>
@@ -27,7 +28,7 @@
2728
* saved_at: \DateTimeImmutable,
2829
* }
2930
*/
30-
final class TraceableChat implements ChatInterface
31+
final class TraceableChat implements ChatInterface, ResetInterface
3132
{
3233
/**
3334
* @var array<int, array{
@@ -66,4 +67,12 @@ public function submit(UserMessage $message): AssistantMessage
6667

6768
return $this->chat->submit($message);
6869
}
70+
71+
public function reset(): void
72+
{
73+
if ($this->chat instanceof ResetInterface) {
74+
$this->chat->reset();
75+
}
76+
$this->calls = [];
77+
}
6978
}

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
use Symfony\AI\Chat\MessageStoreInterface;
1616
use Symfony\AI\Platform\Message\MessageBag;
1717
use Symfony\Component\Clock\ClockInterface;
18+
use Symfony\Contracts\Service\ResetInterface;
1819

1920
/**
2021
* @author Guillaume Loulier <personal@guillaumeloulier.fr>
@@ -24,7 +25,7 @@
2425
* saved_at: \DateTimeImmutable,
2526
* }
2627
*/
27-
final class TraceableMessageStore implements ManagedStoreInterface, MessageStoreInterface
28+
final class TraceableMessageStore implements ManagedStoreInterface, MessageStoreInterface, ResetInterface
2829
{
2930
/**
3031
* @var MessageStoreData[]
@@ -69,4 +70,12 @@ public function drop(): void
6970

7071
$this->messageStore->drop();
7172
}
73+
74+
public function reset(): void
75+
{
76+
if ($this->messageStore instanceof ResetInterface) {
77+
$this->messageStore->reset();
78+
}
79+
$this->calls = [];
80+
}
7281
}

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
use Symfony\AI\Platform\Result\DeferredResult;
2020
use Symfony\AI\Platform\Result\ResultInterface;
2121
use Symfony\AI\Platform\Result\StreamResult;
22+
use Symfony\Contracts\Service\ResetInterface;
2223

2324
/**
2425
* @author Christopher Hertel <mail@christopher-hertel.de>
@@ -30,7 +31,7 @@
3031
* result: DeferredResult,
3132
* }
3233
*/
33-
final class TraceablePlatform implements PlatformInterface
34+
final class TraceablePlatform implements PlatformInterface, ResetInterface
3435
{
3536
/**
3637
* @var PlatformCallData[]
@@ -74,6 +75,12 @@ public function getModelCatalog(): ModelCatalogInterface
7475
return $this->platform->getModelCatalog();
7576
}
7677

78+
public function reset(): void
79+
{
80+
$this->calls = [];
81+
$this->resultCache = new \WeakMap();
82+
}
83+
7784
private function createTraceableStreamResult(DeferredResult $originalStream): StreamResult
7885
{
7986
return $result = new StreamResult((function () use (&$result, $originalStream) {

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,12 @@
1414
use Symfony\AI\Agent\Toolbox\ToolboxInterface;
1515
use Symfony\AI\Agent\Toolbox\ToolResult;
1616
use Symfony\AI\Platform\Result\ToolCall;
17+
use Symfony\Contracts\Service\ResetInterface;
1718

1819
/**
1920
* @author Christopher Hertel <mail@christopher-hertel.de>
2021
*/
21-
final class TraceableToolbox implements ToolboxInterface
22+
final class TraceableToolbox implements ToolboxInterface, ResetInterface
2223
{
2324
/**
2425
* @var ToolResult[]
@@ -39,4 +40,9 @@ public function execute(ToolCall $toolCall): ToolResult
3940
{
4041
return $this->calls[] = $this->toolbox->execute($toolCall);
4142
}
43+
44+
public function reset(): void
45+
{
46+
$this->calls = [];
47+
}
4248
}

‎src/ai-bundle/tests/DependencyInjection/AiBundleTest.php‎

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1535,17 +1535,14 @@ public function testInMemoryStoreWithoutCustomStrategyCanBeConfigured()
15351535
$definition = $container->getDefinition('ai.store.memory.my_memory_store_with_custom_strategy');
15361536
$this->assertSame(InMemoryStore::class, $definition->getClass());
15371537

1538-
$this->assertTrue($definition->isLazy());
1538+
$this->assertFalse($definition->isLazy());
15391539
$this->assertCount(1, $definition->getArguments());
15401540
$this->assertInstanceOf(Definition::class, $definition->getArgument(0));
15411541
$this->assertSame(DistanceCalculator::class, $definition->getArgument(0)->getClass());
15421542

1543-
$this->assertTrue($definition->hasTag('proxy'));
1544-
$this->assertSame([
1545-
['interface' => StoreInterface::class],
1546-
['interface' => ManagedStoreInterface::class],
1547-
], $definition->getTag('proxy'));
1543+
$this->assertFalse($definition->hasTag('proxy'));
15481544
$this->assertTrue($definition->hasTag('ai.store'));
1545+
$this->assertTrue($definition->hasTag('kernel.reset'));
15491546

15501547
$this->assertTrue($container->hasAlias('.'.StoreInterface::class.' $my_memory_store_with_custom_strategy'));
15511548
$this->assertTrue($container->hasAlias(StoreInterface::class.' $myMemoryStoreWithCustomStrategy'));
@@ -1574,17 +1571,14 @@ public function testInMemoryStoreWithCustomStrategyCanBeConfigured()
15741571
$definition = $container->getDefinition('ai.store.memory.my_memory_store_with_custom_strategy');
15751572
$this->assertSame(InMemoryStore::class, $definition->getClass());
15761573

1577-
$this->assertTrue($definition->isLazy());
1574+
$this->assertFalse($definition->isLazy());
15781575
$this->assertCount(1, $definition->getArguments());
15791576
$this->assertInstanceOf(Reference::class, $definition->getArgument(0));
15801577
$this->assertSame('ai.store.distance_calculator.my_memory_store_with_custom_strategy', (string) $definition->getArgument(0));
15811578

1582-
$this->assertTrue($definition->hasTag('proxy'));
1583-
$this->assertSame([
1584-
['interface' => StoreInterface::class],
1585-
['interface' => ManagedStoreInterface::class],
1586-
], $definition->getTag('proxy'));
1579+
$this->assertFalse($definition->hasTag('proxy'));
15871580
$this->assertTrue($definition->hasTag('ai.store'));
1581+
$this->assertTrue($definition->hasTag('kernel.reset'));
15881582

15891583
$this->assertTrue($container->hasAlias('.'.StoreInterface::class.' $my_memory_store_with_custom_strategy'));
15901584
$this->assertTrue($container->hasAlias(StoreInterface::class.' $myMemoryStoreWithCustomStrategy'));
@@ -6906,14 +6900,12 @@ public function testMemoryMessageStoreCanBeConfiguredWithCustomKey()
69066900

69076901
$definition = $container->getDefinition('ai.message_store.memory.custom');
69086902

6909-
$this->assertTrue($definition->isLazy());
6903+
$this->assertFalse($definition->isLazy());
69106904
$this->assertSame('foo', $definition->getArgument(0));
69116905

6912-
$this->assertSame([
6913-
['interface' => MessageStoreInterface::class],
6914-
['interface' => ManagedMessageStoreInterface::class],
6915-
], $definition->getTag('proxy'));
6906+
$this->assertFalse($definition->hasTag('proxy'));
69166907
$this->assertTrue($definition->hasTag('ai.message_store'));
6908+
$this->assertTrue($definition->hasTag('kernel.reset'));
69176909
}
69186910

69196911
public function testMongoDbMessageStoreIsConfigured()

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,4 +60,17 @@ public function testInitializationMessageBagCanBeRetrieved()
6060
$this->assertInstanceOf(UserMessage::class, $traceableChat->calls[1]['message']);
6161
$this->assertInstanceOf(\DateTimeImmutable::class, $traceableChat->calls[1]['saved_at']);
6262
}
63+
64+
public function testResetClearsCalls()
65+
{
66+
$agent = $this->createStub(AgentInterface::class);
67+
$chat = new Chat($agent, new InMemoryStore());
68+
$traceableChat = new TraceableChat($chat, new MonotonicClock());
69+
70+
$traceableChat->initiate(new MessageBag());
71+
$this->assertCount(1, $traceableChat->calls);
72+
73+
$traceableChat->reset();
74+
$this->assertCount(0, $traceableChat->calls);
75+
}
6376
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,4 +42,16 @@ public function testSubmittedMessageBagCanBeRetrieved()
4242
$this->assertCount(1, $calls[0]['bag']);
4343
$this->assertInstanceOf(\DateTimeImmutable::class, $calls[0]['saved_at']);
4444
}
45+
46+
public function testResetClearsCalls()
47+
{
48+
$messageStore = new InMemoryStore();
49+
$traceableMessageStore = new TraceableMessageStore($messageStore, new MonotonicClock());
50+
51+
$traceableMessageStore->save(new MessageBag());
52+
$this->assertCount(1, $traceableMessageStore->calls);
53+
54+
$traceableMessageStore->reset();
55+
$this->assertCount(0, $traceableMessageStore->calls);
56+
}
4557
}

0 commit comments

Comments
 (0)