fix: LLMMetadataExtractor serialisation (#136)

* fixing serialization * fixing tests * adding page_range (already expanded) to the serialisation * fixing serialisation tests
deepset-ai · Dec 1, 2024 · e1aebb8 · e1aebb8
1 parent 61cd118
commit e1aebb8
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/haystack_experimental/components/extractors/llm_metadata_extractor.py b/haystack_experimental/components/extractors/llm_metadata_extractor.py
@@ -228,11 +228,12 @@ def to_dict(self) -> Dict[str, Any]:
         return default_to_dict(
             self,
             prompt=self.prompt,
-            input_text=self.prompt_variable,
+            prompt_variable=self.prompt_variable,
             expected_keys=self.expected_keys,
             raise_on_failure=self.raise_on_failure,
             generator_api=self.generator_api.value,
             generator_api_params=llm_provider["init_parameters"],
+            page_range=self.expanded_range
         )
 
     @classmethod

diff --git a/test/components/extractors/test_llm_metadata_extractor.py b/test/components/extractors/test_llm_metadata_extractor.py
@@ -59,7 +59,7 @@ def test_init_missing_prompt_variable(self, monkeypatch):
                 prompt_variable="test2"
             )
 
-    def test_to_dict(self, monkeypatch):
+    def test_to_dict_default_params(self, monkeypatch):
         monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
         extractor = LLMMetadataExtractor(
             prompt="some prompt that was used with the LLM {{test}}",
@@ -68,15 +68,18 @@ def test_to_dict(self, monkeypatch):
             prompt_variable="test",
             generator_api_params={'model': 'gpt-4o-mini', 'generation_kwargs': {"temperature": 0.5}},
             raise_on_failure=True)
+
         extractor_dict = extractor.to_dict()
+
         assert extractor_dict == {
             'type': 'haystack_experimental.components.extractors.llm_metadata_extractor.LLMMetadataExtractor',
             'init_parameters': {
                 'prompt': 'some prompt that was used with the LLM {{test}}',
                 'expected_keys': ['key1', 'key2'],
                 'raise_on_failure': True,
-                'input_text': 'test',
+                'prompt_variable': 'test',
                 'generator_api': 'openai',
+                'page_range': None,
                 'generator_api_params': {
                       'api_base_url': None,
                       'api_key': {'env_vars': ['OPENAI_API_KEY'],'strict': True,'type': 'env_var'},