epam / epam/statgpt-backend

KeyError in populate_stage when LLM hallucinates dataset ID

Open
#152 0 comments 0 reactions 1 assignee Claimed by @veryniceuser View on GitHub
bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.