KeyError in populate_stage when LLM hallucinates dataset ID
- Dominant language
- Python
- Stars
- 26
- Forks
- 1
- Avg merge
- 1d 17h
- Merged PRs (30d)
- 25
Description
### StatGPT Backend version
0.4.0
### Short description
During indicator selection, the LLM occasionally hallucinates a dataset UUID that doesn't exist in `datasets_dict`. The `populate_stage()` method (debug stage formatting) runs on the raw LLM response **before** `_remove_hallucinations()` gets a chance to filter them out. When `format_multidataset_queries()` tries to look up the hallucinated dataset ID, it crashes with a `KeyError`.
### What steps will reproduce the bug?
Run Data Query. Occasionally LLM will hallucinate dataset UUIDs (in my case it just misspelled one letter in UUID) - then this bug will happen
### What is the expected behavior?
_No response_
### What do you see instead?
Error calling tool Query_Data:
...
File "pis_v1_v2.py", line 267, in async_lambda
return await inputs[self.PARSED_RESPONSE_KEY].populate_stage(inputs)
File "pis_v1_v2.py", line 210, in populate_stage
line = await formatter.format_multidataset_queries(queries=self.get_queries())
File "format_utils.py", line 112, in format_multidataset_queries
dataset = self.datasets[dataset_id]
KeyError: 'UUID'
### Additional information
## Summary
During indicator selection, the LLM occasionally hallucinates a dataset UUID that doesn't exist in `datasets_dict`. The `populate_stage()` method (debug stage formatting) runs on the raw LLM response **before** `_remove_hallucinations()` gets a chance to filter them out. When `format_multidataset_queries()` tries to look up the hallucinated dataset ID, it crashes with a `KeyError`.
## Occurrence
Observed once during a full eval run. Rare — the LLM usually returns valid dataset IDs from the provided candidates, but occasionally fabricates or garbles a UUID.
## Root Cause
The chain execution order in `_create_chain_inner()` (`pis_v1_v2.py:269-276`):
```
format_candidates → LLM call → populate_stage() → _remove_hallucinations()
^ ERROR ^ would have caught it
```
`populate_stage()` formats the raw LLM response for debug display and accesses `self.datasets[dataset_id]` (via `format_multidataset_queries`), which raises `KeyError` for a hallucinated dataset ID. The `_remove_hallucinations()` step, which is specifically designed to catch this, runs one step too late.
## Affected Files
- `statgpt/app/chains/data_query/query_builder/indicator_selection/format_utils.py:112` — `self.datasets[dataset_id]` lookup
- `statgpt/app/chains/data_query/query_builder/indicator_selection/semantic/packed_indicators_selection/pis_v1_v2.py:210,522,529,542` — `populate_stage` callers and dataset references section
## Logs
### LLM Input (candidates presented to LLM)
https://epa.ms/jmQ9j3
### LLM Output (raw response with hallucinated ID)
https://epa.ms/Nd1VN5
Contributor guide
Assessment
This issue has not been assessed yet.