fix: gracefully handle hallucinated dataset IDs in populate_stage to prevent KeyError #152

[veryniceuser]

Skip unknown dataset IDs with a warning log in format_multidataset_queries and V2 populate_stage dataset references, instead of crashing. This prevents a KeyError when the LLM hallucinates a dataset UUID before _remove_hallucinations gets a chance to filter it out.

Applicable issues

• fixes KeyError in populate_stage when LLM hallucinates dataset ID #152

Description of changes

Checklist

• Title of the pull request follows Conventional Commits specification
• Deployed and tested in a Review environment.

By submitting this pull request, I confirm that my contribution is made under the terms of the MIT license.

[navalnica]

the issue is not only with formatting. if llm selected incorrect uuid, the whole query is unsable. we need to delete it inside indicator selection chain

[navalnica]

also, it's better not to pass uuid to llm. instead, we can pass:

• dataset urn (imf:weo, for example)
• source_id/entity_id
• 1-based index

we also need to make sure there are no duplicates. for example if there are duplicated urns (for some reason. though they should not) we can add index prefix: "1-imf:weo"

[veryniceuser]

the issue is not only with formatting. if llm selected incorrect uuid, the whole query is unsable. we need to delete it inside indicator selection chain

I think it gets deleted in the latter stage at _remove_hallucinations()

[veryniceuser]

also, it's better not to pass uuid to llm. instead, we can pass:

• dataset urn (imf:weo, for example)
• source_id/entity_id
• 1-based index

we also need to make sure there are no duplicates. for example if there are duplicated urns (for some reason. though they should not) we can add index prefix: "1-imf:weo"

Sound like a bigger change. Should we merge this PR as a hotfix, and create a follow-up feature request to make these changes?

[navalnica]

I don't like the current proposed change. it states that, when populating stage or formatting query, the query can contain dataset queries to unavailable datasets. this seems incorrect - we should remove any hallucinations from LLM response and call formatting only afterwards

[veryniceuser]

also, it's better not to pass uuid to llm. instead, we can pass:

• dataset urn (imf:weo, for example)
• source_id/entity_id
• 1-based index

we also need to make sure there are no duplicates. for example if there are duplicated urns (for some reason. though they should not) we can add index prefix: "1-imf:weo"

Created a separate feature request: #172, and pushed changes in PR #173. It's done to un-block this hotfix merge, and to have review and discussion for the in-depth fix in the separate place.

[veryniceuser]

I don't like the current proposed change. it states that, when populating stage or formatting query, the query can contain dataset queries to unavailable datasets. this seems incorrect - we should remove any hallucinations from LLM response and call formatting only afterwards

I disagree that proposed change is wrong, populate stage should not make implicit assumptions, and must not fail if the code that calls it will change in the future.

On the other hand, swapping remove_hallucinations and populate_stage makes sense. But because of the code structure, such change seem to require refactoring. And so should be prioritized separately