feat: add support for azure ai search store

Author: chr-hertel

README.md

@@ -13,32 +13,42 @@ composer require php-llm/llm-chain-bundle:@dev php-llm/llm-chain:@dev
 ```yaml
 # config/packages/llm_chain.yaml
 llm_chain:
-  runtimes:
-    azure_gpt:
-      type: 'azure'
-      base_url: '%env(AZURE_OPENAI_BASEURL)%'
-      deployment: '%env(AZURE_OPENAI_GPT)%'
-      api_key: '%env(AZURE_OPENAI_KEY)%'
-      version: '%env(AZURE_OPENAI_VERSION)%'
-    azure_embeddings:
-      type: 'azure'
-      base_url: '%env(AZURE_OPENAI_BASEURL)%'
-      deployment: '%env(AZURE_OPENAI_EMBEDDINGS)%'
-      api_key: '%env(AZURE_OPENAI_KEY)%'
-      version: '%env(AZURE_OPENAI_VERSION)%'
-    openai:
-      type: 'openai'
-      api_key: '%env(OPENAI_API_KEY)%'
-  llms:
-    azure_gpt:
-      runtime: 'azure_gpt'
-    original_gpt:
-      runtime: 'openai'
-  embeddings:
-    azure_embeddings:
-      runtime: 'azure_embeddings'
-    original_embeddings:
-      runtime: 'openai'
+    runtimes:
+        azure_gpt:
+            type: 'azure'
+            base_url: '%env(AZURE_OPENAI_BASEURL)%'
+            deployment: '%env(AZURE_OPENAI_GPT)%'
+            api_key: '%env(AZURE_OPENAI_KEY)%'
+            version: '%env(AZURE_OPENAI_VERSION)%'
+        azure_embeddings:
+            type: 'azure'
+            base_url: '%env(AZURE_OPENAI_BASEURL)%'
+            deployment: '%env(AZURE_OPENAI_EMBEDDINGS)%'
+            api_key: '%env(AZURE_OPENAI_KEY)%'
+            version: '%env(AZURE_OPENAI_VERSION)%'
+        openai:
+            type: 'openai'
+            api_key: '%env(OPENAI_API_KEY)%'
+    llms:
+        azure_gpt:
+            runtime: 'azure_gpt'
+        original_gpt:
+            runtime: 'openai'
+    embeddings:
+        azure_embeddings:
+            runtime: 'azure_embeddings'
+        original_embeddings:
+            runtime: 'openai'
+    stores:
+        chroma_db:
+            type: 'chroma-db'
+            collection_name: '%env(CHROMA_COLLECTION)%'
+        azure_search:
+            type: 'azure-search'
+            api_key: '%env(AZURE_SEARCH_KEY)%' 
+            endpoint: '%env(AZURE_SEARCH_ENDPOINT)%'
+            index_name: '%env(AZURE_SEARCH_INDEX)%' 
+            api_version: '2024-07-01'
 ```
 
 ## Usage

src/DependencyInjection/Configuration.php

@@ -52,8 +52,12 @@ public function getConfigTreeBuilder(): TreeBuilder
                     ->useAttributeAsKey('name')
                     ->arrayPrototype()
                         ->children()
-                            ->scalarNode('engine')->isRequired()->end()
+                            ->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()
                         ->end()
                     ->end()
             ->end()

src/DependencyInjection/LlmChainExtension.php

@@ -11,6 +11,7 @@
 use PhpLlm\LlmChain\OpenAI\Runtime;
 use PhpLlm\LlmChain\OpenAI\Runtime\Azure as AzureRuntime;
 use PhpLlm\LlmChain\OpenAI\Runtime\OpenAI as OpenAIRuntime;
+use PhpLlm\LlmChain\Store\Azure\SearchStore as AzureSearchStore;
 use PhpLlm\LlmChain\Store\ChromaDb\Store as ChromaDbStore;
 use PhpLlm\LlmChain\Store\StoreInterface;
 use PhpLlm\LlmChain\Store\VectorStoreInterface;
@@ -122,5 +123,16 @@ private function processStoreConfig(string $name, mixed $stores, ContainerBuilde
 
             $container->setDefinition('llm_chain.store.'.$name, $definition);
         }
+
+        if ('azure-search' === $stores['engine']) {
+            $definition = new ChildDefinition(AzureSearchStore::class);
+            $definition
+                ->replaceArgument('$endpointUrl', $stores['endpoint'])
+                ->replaceArgument('$apiKey', $stores['api_key'])
+                ->replaceArgument('$indexName', $stores['index_name'])
+                ->replaceArgument('$apiVersion', $stores['api_version']);
+
+            $container->setDefinition('llm_chain.store.'.$name, $definition);
+        }
     }
 }

src/Resources/config/services.php

@@ -11,11 +11,12 @@
 use PhpLlm\LlmChain\OpenAI\Runtime;
 use PhpLlm\LlmChain\OpenAI\Runtime\Azure as AzureRuntime;
 use PhpLlm\LlmChain\OpenAI\Runtime\OpenAI as OpenAIRuntime;
-use PhpLlm\LlmChain\RetrievalChain;
+use PhpLlm\LlmChain\Store\Azure\SearchStore;
 use PhpLlm\LlmChain\ToolBox\ParameterAnalyzer;
 use PhpLlm\LlmChain\ToolBox\Registry;
 use PhpLlm\LlmChain\ToolBox\ToolAnalyzer;
 use PhpLlm\LlmChain\ToolChain;
+use PhpLlm\LlmChain\Store\Azure\SearchStore as AzureSearchStore;
 use PhpLlm\LlmChain\Store\ChromaDb\Store as ChromaDbStore;
 
 return static function (ContainerConfigurator $container) {
@@ -26,7 +27,6 @@
 
         // chains
         ->set(Chat::class)
-        ->set(RetrievalChain::class)
         ->set(ToolChain::class)
 
         // runtimes
@@ -57,6 +57,14 @@
             ])
 
         // stores
+        ->set(AzureSearchStore::class)
+            ->abstract()
+            ->args([
+                '$endpointUrl' => abstract_arg('Endpoint URL for Azure AI Search API'),
+                '$apiKey' => abstract_arg('API key for Azure AI Search API'),
+                '$indexName' => abstract_arg('Name of Azure AI Search index'),
+                '$apiVersion' => abstract_arg('API version for Azure AI Search API'),
+            ])
         ->set(ChromaDbStore::class)
             ->abstract()
             ->args([