fixing LLMProvider

deepset-ai · Sep 27, 2024 · 359e1aa · 359e1aa
1 parent 1973bd3
commit 359e1aa
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/haystack_experimental/components/extractors/llm_metadata_extractor.py b/haystack_experimental/components/extractors/llm_metadata_extractor.py
@@ -140,15 +140,16 @@ def __init__( # pylint: disable=R0917
         self.builder = PromptBuilder(prompt, required_variables=[input_text])
         self.raise_on_failure = raise_on_failure
         self.expected_keys = expected_keys
-        self.generator_api = generator_api
+        self.generator_api = generator_api if isinstance(generator_api, LLMProvider) \
+            else LLMProvider.from_str(generator_api)
         self.generator_api_params = generator_api_params or {}
         self.llm_provider = self._init_generator(generator_api, self.generator_api_params)
         if self.input_text not in self.prompt:
             raise ValueError(f"Input text '{self.input_text}' must be in the prompt.")
 
     @staticmethod
     def _init_generator(
-            generator_api: LLMProvider,
+            generator_api: Union[str,LLMProvider],
             generator_api_params: Optional[Dict[str, Any]]
     ) -> Union[OpenAIGenerator, AzureOpenAIGenerator, AmazonBedrockGenerator, VertexAIGeminiGenerator]:
         """