[Store] Add a way to configure the setup options for mariadb store

Author: lyrixx

src/ai-bundle/config/options.php

@@ -598,6 +598,11 @@
                                 ->stringNode('table_name')->cannotBeEmpty()->end()
                                 ->stringNode('index_name')->cannotBeEmpty()->end()
                                 ->stringNode('vector_field_name')->cannotBeEmpty()->end()
+                                ->arrayNode('setup_options')
+                                    ->children()
+                                        ->integerNode('dimensions')->end()
+                                    ->end()
+                                ->end()
                             ->end()
                         ->end()
                     ->end()

src/ai-bundle/config/services.php

@@ -205,6 +205,7 @@
         ->set('ai.command.setup_store', SetupStoreCommand::class)
             ->args([
                 tagged_locator('ai.store', 'name'),
+                abstract_arg('setup store options'),
             ])
             ->tag('console.command')
         ->set('ai.command.drop_store', DropStoreCommand::class)

src/ai-bundle/src/AiBundle.php

@@ -169,9 +169,11 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
             $this->processMultiAgentConfig($multiAgentName, $multiAgent, $builder);
         }
 
+        $setupStoresOptions = [];
         foreach ($config['store'] ?? [] as $type => $store) {
-            $this->processStoreConfig($type, $store, $builder);
+            $this->processStoreConfig($type, $store, $builder, $setupStoresOptions);
         }
+        $builder->getDefinition('ai.command.setup_store')->setArgument(1, $setupStoresOptions);
 
         $stores = array_keys($builder->findTaggedServiceIds('ai.store'));
 
@@ -898,8 +900,9 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
 
     /**
      * @param array<string, mixed> $stores
+     * @param array<string, mixed> $setupStoresOptions
      */
-    private function processStoreConfig(string $type, array $stores, ContainerBuilder $container): void
+    private function processStoreConfig(string $type, array $stores, ContainerBuilder $container, &$setupStoresOptions): void
     {
         if ('azure_search' === $type) {
             foreach ($stores as $name => $store) {
@@ -1093,6 +1096,8 @@ private function processStoreConfig(string $type, array $stores, ContainerBuilde
                 $container->setDefinition($serviceId, $definition);
                 $container->registerAliasForArgument($serviceId, StoreInterface::class, $name);
                 $container->registerAliasForArgument($serviceId, StoreInterface::class, $type.'_'.$name);
+
+                $setupStoresOptions[$serviceId] = $store['setup_options'] ?? [];
             }
         }
 

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

@@ -123,7 +123,8 @@ public function testStoreCommandsAreDefined()
         $this->assertTrue($container->hasDefinition('ai.command.setup_store'));
 
         $setupStoreCommandDefinition = $container->getDefinition('ai.command.setup_store');
-        $this->assertCount(1, $setupStoreCommandDefinition->getArguments());
+        $this->assertCount(2, $setupStoreCommandDefinition->getArguments());
+        $this->assertSame(['ai.store.mariadb.my_mariadb_store' => ['dimensions' => 1024]], $setupStoreCommandDefinition->getArgument(1));
         $this->assertArrayHasKey('console.command', $setupStoreCommandDefinition->getTags());
 
         $this->assertTrue($container->hasDefinition('ai.command.drop_store'));
@@ -3502,6 +3503,9 @@ private function getFullConfig(): array
                             'table_name' => 'vector_table',
                             'index_name' => 'vector_idx',
                             'vector_field_name' => 'vector',
+                            'setup_options' => [
+                                'dimensions' => 1024,
+                            ],
                         ],
                     ],
                     'meilisearch' => [

src/store/src/Command/SetupStoreCommand.php

@@ -31,9 +31,11 @@ final class SetupStoreCommand extends Command
 {
     /**
      * @param ServiceLocator<ManagedStoreInterface> $stores
+     * @param array<string, mixed>                  $setupStoresOptions
      */
     public function __construct(
         private readonly ServiceLocator $stores,
+        private readonly array $setupStoresOptions = [],
     ) {
         parent::__construct();
     }
@@ -84,7 +86,7 @@ protected function execute(InputInterface $input, OutputInterface $output): int
         $store = $this->stores->get($storeName);
 
         try {
-            $store->setup();
+            $store->setup($this->setupStoresOptions[$storeName] ?? []);
             $io->success(\sprintf('The "%s" store was set up successfully.', $storeName));
         } catch (\Exception $e) {
             throw new RuntimeException(\sprintf('An error occurred while setting up the "%s" store: ', $storeName).$e->getMessage(), previous: $e);

src/store/tests/Command/SetupStoreCommandTest.php

@@ -90,11 +90,16 @@ public function testCommandCannotSetupStoreWithException()
     public function testCommandCanSetupDefinedStore()
     {
         $store = $this->createMock(ManagedStoreInterface::class);
-        $store->expects($this->once())->method('setup');
-
-        $command = new SetupStoreCommand(new ServiceLocator([
-            'foo' => static fn (): object => $store,
-        ]));
+        $store->expects($this->once())->method('setup')->with(['some_option' => 'some_value']);
+
+        $command = new SetupStoreCommand(
+            new ServiceLocator([
+                'foo' => static fn (): object => $store,
+            ]),
+            [
+                'foo' => ['some_option' => 'some_value'],
+            ]
+        );
 
         $tester = new CommandTester($command);