Skip to content

Commit 25a63c9

Browse files
santysisisantysisi
committed
[AI Bundle][MCP Bundle][Chat][Store] add ResetInterface to prevent memory leaks
1 parent 417e6a1 commit 25a63c9

25 files changed

+239
-39
lines changed

‎src/ai-bundle/CHANGELOG.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
CHANGELOG
22
=========
33

4+
0.4
5+
---
6+
7+
* Add `ResetInterface` support to `TraceableChat`, `TraceableMessageStore`, `TraceablePlatform` and `TraceableToolbox` to clear collected data between requests
8+
49
0.2
510
---
611

‎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
@@ -182,7 +182,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
182182
$traceablePlatformDefinition = (new Definition(TraceablePlatform::class))
183183
->setDecoratedService($platform, priority: -1024)
184184
->setArguments([new Reference('.inner')])
185-
->addTag('ai.traceable_platform');
185+
->addTag('ai.traceable_platform')
186+
->addTag('kernel.reset', ['method' => 'reset']);
186187
$suffix = u($platform)->after('ai.platform.')->toString();
187188
$builder->setDefinition('ai.traceable_platform.'.$suffix, $traceablePlatformDefinition);
188189
}
@@ -259,7 +260,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
259260
new Reference('.inner'),
260261
new Reference(ClockInterface::class),
261262
])
262-
->addTag('ai.traceable_message_store');
263+
->addTag('ai.traceable_message_store')
264+
->addTag('kernel.reset', ['method' => 'reset']);
263265
$suffix = u($messageStore)->afterLast('.')->toString();
264266
$builder->setDefinition('ai.traceable_message_store.'.$suffix, $traceableMessageStoreDefinition);
265267
}
@@ -294,7 +296,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
294296
new Reference('.inner'),
295297
new Reference(ClockInterface::class),
296298
])
297-
->addTag('ai.traceable_chat');
299+
->addTag('ai.traceable_chat')
300+
->addTag('kernel.reset', ['method' => 'reset']);
298301
$suffix = u($chat)->afterLast('.')->toString();
299302
$builder->setDefinition('ai.traceable_chat.'.$suffix, $traceableChatDefinition);
300303
}
@@ -1121,7 +1124,8 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
11211124
->setClass(TraceableToolbox::class)
11221125
->setArguments([new Reference('.inner')])
11231126
->setDecoratedService('ai.toolbox.'.$name, priority: -1024)
1124-
->addTag('ai.traceable_toolbox');
1127+
->addTag('ai.traceable_toolbox')
1128+
->addTag('kernel.reset', ['method' => 'reset']);
11251129
$container->setDefinition('ai.traceable_toolbox.'.$name, $traceableToolboxDefinition);
11261130
}
11271131

@@ -1508,11 +1512,9 @@ private function processStoreConfig(string $type, array $stores, ContainerBuilde
15081512

15091513
$definition = new Definition(InMemoryStore::class);
15101514
$definition
1511-
->setLazy(true)
15121515
->setArguments($arguments)
1513-
->addTag('proxy', ['interface' => StoreInterface::class])
1514-
->addTag('proxy', ['interface' => ManagedStoreInterface::class])
1515-
->addTag('ai.store');
1516+
->addTag('ai.store')
1517+
->addTag('kernel.reset', ['method' => 'reset']);
15161518

15171519
$container->setDefinition('ai.store.'.$type.'.'.$name, $definition);
15181520
$container->registerAliasForArgument('ai.store.'.$type.'.'.$name, StoreInterface::class, $name);
@@ -2060,11 +2062,9 @@ private function processMessageStoreConfig(string $type, array $messageStores, C
20602062
foreach ($messageStores as $name => $messageStore) {
20612063
$definition = new Definition(InMemoryMessageStore::class);
20622064
$definition
2063-
->setLazy(true)
20642065
->setArgument(0, $messageStore['identifier'])
2065-
->addTag('proxy', ['interface' => MessageStoreInterface::class])
2066-
->addTag('proxy', ['interface' => ManagedMessageStoreInterface::class])
2067-
->addTag('ai.message_store');
2066+
->addTag('ai.message_store')
2067+
->addTag('kernel.reset', ['method' => 'reset']);
20682068

20692069
$container->setDefinition('ai.message_store.'.$type.'.'.$name, $definition);
20702070
$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
@@ -1530,17 +1530,14 @@ public function testInMemoryStoreWithoutCustomStrategyCanBeConfigured()
15301530
$definition = $container->getDefinition('ai.store.memory.my_memory_store_with_custom_strategy');
15311531
$this->assertSame(InMemoryStore::class, $definition->getClass());
15321532

1533-
$this->assertTrue($definition->isLazy());
1533+
$this->assertFalse($definition->isLazy());
15341534
$this->assertCount(1, $definition->getArguments());
15351535
$this->assertInstanceOf(Definition::class, $definition->getArgument(0));
15361536
$this->assertSame(DistanceCalculator::class, $definition->getArgument(0)->getClass());
15371537

1538-
$this->assertTrue($definition->hasTag('proxy'));
1539-
$this->assertSame([
1540-
['interface' => StoreInterface::class],
1541-
['interface' => ManagedStoreInterface::class],
1542-
], $definition->getTag('proxy'));
1538+
$this->assertFalse($definition->hasTag('proxy'));
15431539
$this->assertTrue($definition->hasTag('ai.store'));
1540+
$this->assertTrue($definition->hasTag('kernel.reset'));
15441541

15451542
$this->assertTrue($container->hasAlias('.Symfony\AI\Store\StoreInterface $my_memory_store_with_custom_strategy'));
15461543
$this->assertTrue($container->hasAlias('Symfony\AI\Store\StoreInterface $myMemoryStoreWithCustomStrategy'));
@@ -1569,17 +1566,14 @@ public function testInMemoryStoreWithCustomStrategyCanBeConfigured()
15691566
$definition = $container->getDefinition('ai.store.memory.my_memory_store_with_custom_strategy');
15701567
$this->assertSame(InMemoryStore::class, $definition->getClass());
15711568

1572-
$this->assertTrue($definition->isLazy());
1569+
$this->assertFalse($definition->isLazy());
15731570
$this->assertCount(1, $definition->getArguments());
15741571
$this->assertInstanceOf(Reference::class, $definition->getArgument(0));
15751572
$this->assertSame('ai.store.distance_calculator.my_memory_store_with_custom_strategy', (string) $definition->getArgument(0));
15761573

1577-
$this->assertTrue($definition->hasTag('proxy'));
1578-
$this->assertSame([
1579-
['interface' => StoreInterface::class],
1580-
['interface' => ManagedStoreInterface::class],
1581-
], $definition->getTag('proxy'));
1574+
$this->assertFalse($definition->hasTag('proxy'));
15821575
$this->assertTrue($definition->hasTag('ai.store'));
1576+
$this->assertTrue($definition->hasTag('kernel.reset'));
15831577

15841578
$this->assertTrue($container->hasAlias('.Symfony\AI\Store\StoreInterface $my_memory_store_with_custom_strategy'));
15851579
$this->assertTrue($container->hasAlias('Symfony\AI\Store\StoreInterface $myMemoryStoreWithCustomStrategy'));
@@ -6709,14 +6703,12 @@ public function testMemoryMessageStoreCanBeConfiguredWithCustomKey()
67096703

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

6712-
$this->assertTrue($definition->isLazy());
6706+
$this->assertFalse($definition->isLazy());
67136707
$this->assertSame('foo', $definition->getArgument(0));
67146708

6715-
$this->assertSame([
6716-
['interface' => MessageStoreInterface::class],
6717-
['interface' => ManagedMessageStoreInterface::class],
6718-
], $definition->getTag('proxy'));
6709+
$this->assertFalse($definition->hasTag('proxy'));
67196710
$this->assertTrue($definition->hasTag('ai.message_store'));
6711+
$this->assertTrue($definition->hasTag('kernel.reset'));
67206712
}
67216713

67226714
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)