From 64a6ec39b0fa0f431efb4a9451e74867dff63325 Mon Sep 17 00:00:00 2001
From: Oskar Stark <oskarstark@googlemail.com>
Date: Thu, 26 Sep 2024 15:33:43 +0200
Subject: [PATCH] Better config

---
 composer.json                                 |  2 +-
 src/DependencyInjection/Configuration.php     | 37 ++++++++++++++-----
 src/DependencyInjection/LlmChainExtension.php |  3 +-
 3 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/composer.json b/composer.json
index 5336004..1129bf4 100644
--- a/composer.json
+++ b/composer.json
@@ -11,7 +11,7 @@
     ],
     "require": {
         "php": ">=8.2",
-        "php-llm/llm-chain": "^0.2",
+        "php-llm/llm-chain": "^0.3",
         "symfony/config": "^6.4 || ^7.0",
         "symfony/dependency-injection": "^6.4 || ^7.0",
         "symfony/framework-bundle": "^6.4 || ^7.0"
diff --git a/src/DependencyInjection/Configuration.php b/src/DependencyInjection/Configuration.php
index 7c366f7..f1a0360 100644
--- a/src/DependencyInjection/Configuration.php
+++ b/src/DependencyInjection/Configuration.php
@@ -47,19 +47,36 @@ public function getConfigTreeBuilder(): TreeBuilder
                         ->end()
                     ->end()
                 ->end()
-                ->arrayNode('stores')
-                    ->normalizeKeys(false)
-                    ->useAttributeAsKey('name')
-                    ->arrayPrototype()
+            ->arrayNode('stores')
+                ->children()
+                    ->arrayNode('chroma_db')
+                        ->children()
+                            ->scalarNode('engine')->defaultValue('chroma-db')->end()
+                            ->scalarNode('collection_name')->isRequired()->end()
+                        ->end()
+                    ->end()
+                    ->arrayNode('azure_search')
                         ->children()
-                            ->enumNode('engine')->values(['chroma-db', 'azure-search'])->isRequired()->end()
-                            ->scalarNode('collection_name')->end()
-                            ->scalarNode('api_key')->end()
-                            ->scalarNode('endpoint')->end()
-                            ->scalarNode('index_name')->end()
-                            ->scalarNode('api_version')->end()
+                            ->scalarNode('engine')->defaultValue('azure-search')->end()
+                            ->scalarNode('api_key')->isRequired()->end()
+                            ->scalarNode('endpoint')->isRequired()->end()
+                            ->scalarNode('index_name')->isRequired()->end()
+                            ->scalarNode('api_version')->defaultValue('2024-07-01')->end()
+                            ->scalarNode('vector_field_name')->defaultValue('vector')->end()
                         ->end()
                     ->end()
+//                    ->arrayNode('mongodb')
+//                        ->children()
+//                            ->scalarNode('engine')->defaultValue('mongodb')->end()
+//                            ->scalarNode('uri')->isRequired()->end()
+//                            ->scalarNode('database_name')->isRequired()->end()
+//                            ->scalarNode('collection_name')->isRequired()->end()
+//                            ->scalarNode('index_name')->end()
+//                            ->scalarNode('vector_field_name')->defaultValue('vector')->end()
+//                            ->booleanNode('bulk_write')->defaultFalse()->end()
+//                        ->end()
+//                    ->end()
+                ->end()
             ->end()
         ;
 
diff --git a/src/DependencyInjection/LlmChainExtension.php b/src/DependencyInjection/LlmChainExtension.php
index c1bced7..2ef1aec 100644
--- a/src/DependencyInjection/LlmChainExtension.php
+++ b/src/DependencyInjection/LlmChainExtension.php
@@ -160,7 +160,8 @@ private function processStoreConfig(string $name, array $stores, ContainerBuilde
                 ->replaceArgument('$endpointUrl', $stores['endpoint'])
                 ->replaceArgument('$apiKey', $stores['api_key'])
                 ->replaceArgument('$indexName', $stores['index_name'])
-                ->replaceArgument('$apiVersion', $stores['api_version']);
+                ->replaceArgument('$apiVersion', $stores['api_version'])
+                ->replaceArgument('$vectorFieldName', $stores['vector_field_name']);
 
             $container->setDefinition('llm_chain.store.'.$name, $definition);
         }