Skip to content

Commit 1618bc9

Browse files
OskarStarkOskarStark
committed
feature #1533 [AI Bundle] add ResetInterface to prevent memory leaks (santysisi)
This PR was merged into the main branch. Discussion ---------- [AI Bundle] add `ResetInterface` to prevent memory leaks | Q | A | ------------- | --- | Bug fix? | yes | New feature? | yes | Docs? | no | Issues | Fix #1531 | License | MIT This PR introduces the `ResetInterface` and `reset` tags to the tracers and in memory stores. These changes help optimize memory usage by preventing unnecessary memory increases with each request. Additionally, they ensure that erroneous information is no longer displayed in the profiler, improving accuracy and performance monitoring. Commits ------- 274862c [AI Bundle][MCP Bundle][Chat][Store] add `ResetInterface` to prevent memory leaks
2 parents bc5f2da + 274862c commit 1618bc9

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)