Skip to content

Commit e1aebb8

Browse files
fix: LLMMetadataExtractor serialisation (#136)
* fixing serialization * fixing tests * adding page_range (already expanded) to the serialisation * fixing serialisation tests
1 parent 61cd118 commit e1aebb8

File tree

2 files changed

+7
-3
lines changed

2 files changed

+7
-3
lines changed

haystack_experimental/components/extractors/llm_metadata_extractor.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -228,11 +228,12 @@ def to_dict(self) -> Dict[str, Any]:
228228
return default_to_dict(
229229
self,
230230
prompt=self.prompt,
231-
input_text=self.prompt_variable,
231+
prompt_variable=self.prompt_variable,
232232
expected_keys=self.expected_keys,
233233
raise_on_failure=self.raise_on_failure,
234234
generator_api=self.generator_api.value,
235235
generator_api_params=llm_provider["init_parameters"],
236+
page_range=self.expanded_range
236237
)
237238

238239
@classmethod

test/components/extractors/test_llm_metadata_extractor.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ def test_init_missing_prompt_variable(self, monkeypatch):
5959
prompt_variable="test2"
6060
)
6161

62-
def test_to_dict(self, monkeypatch):
62+
def test_to_dict_default_params(self, monkeypatch):
6363
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
6464
extractor = LLMMetadataExtractor(
6565
prompt="some prompt that was used with the LLM {{test}}",
@@ -68,15 +68,18 @@ def test_to_dict(self, monkeypatch):
6868
prompt_variable="test",
6969
generator_api_params={'model': 'gpt-4o-mini', 'generation_kwargs': {"temperature": 0.5}},
7070
raise_on_failure=True)
71+
7172
extractor_dict = extractor.to_dict()
73+
7274
assert extractor_dict == {
7375
'type': 'haystack_experimental.components.extractors.llm_metadata_extractor.LLMMetadataExtractor',
7476
'init_parameters': {
7577
'prompt': 'some prompt that was used with the LLM {{test}}',
7678
'expected_keys': ['key1', 'key2'],
7779
'raise_on_failure': True,
78-
'input_text': 'test',
80+
'prompt_variable': 'test',
7981
'generator_api': 'openai',
82+
'page_range': None,
8083
'generator_api_params': {
8184
'api_base_url': None,
8285
'api_key': {'env_vars': ['OPENAI_API_KEY'],'strict': True,'type': 'env_var'},

0 commit comments

Comments
 (0)