Skip to content

Commit 6c7526e

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

25 files changed

+231
-16
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: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@
137137
use Symfony\Component\Security\Core\Authorization\AuthorizationCheckerInterface;
138138
use Symfony\Component\Translation\TranslatableMessage;
139139
use Symfony\Contracts\HttpClient\HttpClientInterface;
140+
use Symfony\Contracts\Service\ResetInterface;
140141

141142
use function Symfony\Component\String\u;
142143

@@ -182,7 +183,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
182183
$traceablePlatformDefinition = (new Definition(TraceablePlatform::class))
183184
->setDecoratedService($platform, priority: -1024)
184185
->setArguments([new Reference('.inner')])
185-
->addTag('ai.traceable_platform');
186+
->addTag('ai.traceable_platform')
187+
->addTag('kernel.reset', ['method' => 'reset']);
186188
$suffix = u($platform)->after('ai.platform.')->toString();
187189
$builder->setDefinition('ai.traceable_platform.'.$suffix, $traceablePlatformDefinition);
188190
}
@@ -259,7 +261,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
259261
new Reference('.inner'),
260262
new Reference(ClockInterface::class),
261263
])
262-
->addTag('ai.traceable_message_store');
264+
->addTag('ai.traceable_message_store')
265+
->addTag('kernel.reset', ['method' => 'reset']);
263266
$suffix = u($messageStore)->afterLast('.')->toString();
264267
$builder->setDefinition('ai.traceable_message_store.'.$suffix, $traceableMessageStoreDefinition);
265268
}
@@ -294,7 +297,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
294297
new Reference('.inner'),
295298
new Reference(ClockInterface::class),
296299
])
297-
->addTag('ai.traceable_chat');
300+
->addTag('ai.traceable_chat')
301+
->addTag('kernel.reset', ['method' => 'reset']);
298302
$suffix = u($chat)->afterLast('.')->toString();
299303
$builder->setDefinition('ai.traceable_chat.'.$suffix, $traceableChatDefinition);
300304
}
@@ -1121,7 +1125,8 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
11211125
->setClass(TraceableToolbox::class)
11221126
->setArguments([new Reference('.inner')])
11231127
->setDecoratedService('ai.toolbox.'.$name, priority: -1024)
1124-
->addTag('ai.traceable_toolbox');
1128+
->addTag('ai.traceable_toolbox')
1129+
->addTag('kernel.reset', ['method' => 'reset']);
11251130
$container->setDefinition('ai.traceable_toolbox.'.$name, $traceableToolboxDefinition);
11261131
}
11271132

@@ -1512,7 +1517,9 @@ private function processStoreConfig(string $type, array $stores, ContainerBuilde
15121517
->setArguments($arguments)
15131518
->addTag('proxy', ['interface' => StoreInterface::class])
15141519
->addTag('proxy', ['interface' => ManagedStoreInterface::class])
1515-
->addTag('ai.store');
1520+
->addTag('proxy', ['interface' => ResetInterface::class])
1521+
->addTag('ai.store')
1522+
->addTag('kernel.reset', ['method' => 'reset']);
15161523

15171524
$container->setDefinition('ai.store.'.$type.'.'.$name, $definition);
15181525
$container->registerAliasForArgument('ai.store.'.$type.'.'.$name, StoreInterface::class, $name);
@@ -2064,7 +2071,9 @@ private function processMessageStoreConfig(string $type, array $messageStores, C
20642071
->setArgument(0, $messageStore['identifier'])
20652072
->addTag('proxy', ['interface' => MessageStoreInterface::class])
20662073
->addTag('proxy', ['interface' => ManagedMessageStoreInterface::class])
2067-
->addTag('ai.message_store');
2074+
->addTag('proxy', ['interface' => ResetInterface::class])
2075+
->addTag('ai.message_store')
2076+
->addTag('kernel.reset', ['method' => 'reset']);
20682077

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

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

Lines changed: 7 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,9 @@ public function submit(UserMessage $message): AssistantMessage
6667

6768
return $this->chat->submit($message);
6869
}
70+
71+
public function reset(): void
72+
{
73+
$this->calls = [];
74+
}
6975
}

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

Lines changed: 7 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,9 @@ public function drop(): void
6970

7071
$this->messageStore->drop();
7172
}
73+
74+
public function reset(): void
75+
{
76+
$this->calls = [];
77+
}
7278
}

‎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: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@
9595
use Symfony\Component\RateLimiter\Storage\InMemoryStorage;
9696
use Symfony\Component\Serializer\Serializer;
9797
use Symfony\Contracts\HttpClient\HttpClientInterface;
98+
use Symfony\Contracts\Service\ResetInterface;
9899

99100
class AiBundleTest extends TestCase
100101
{
@@ -1539,6 +1540,7 @@ public function testInMemoryStoreWithoutCustomStrategyCanBeConfigured()
15391540
$this->assertSame([
15401541
['interface' => StoreInterface::class],
15411542
['interface' => ManagedStoreInterface::class],
1543+
['interface' => ResetInterface::class],
15421544
], $definition->getTag('proxy'));
15431545
$this->assertTrue($definition->hasTag('ai.store'));
15441546

@@ -1578,6 +1580,7 @@ public function testInMemoryStoreWithCustomStrategyCanBeConfigured()
15781580
$this->assertSame([
15791581
['interface' => StoreInterface::class],
15801582
['interface' => ManagedStoreInterface::class],
1583+
['interface' => ResetInterface::class],
15811584
], $definition->getTag('proxy'));
15821585
$this->assertTrue($definition->hasTag('ai.store'));
15831586

@@ -6715,6 +6718,7 @@ public function testMemoryMessageStoreCanBeConfiguredWithCustomKey()
67156718
$this->assertSame([
67166719
['interface' => MessageStoreInterface::class],
67176720
['interface' => ManagedMessageStoreInterface::class],
6721+
['interface' => ResetInterface::class],
67186722
], $definition->getTag('proxy'));
67196723
$this->assertTrue($definition->hasTag('ai.message_store'));
67206724
}

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