epam / epam/statgpt-backend

NER step crashes Query_Data when the LLM returns prose instead of JSON

Open
#595 1 comment 0 reactions 1 assignee Claimed by @kryachkow View on GitHub
bug
Dominant language
Python
Stars
26
Forks
1
Avg merge
1d 17h
Merged PRs (30d)
25

Description

### Short description

Several steps in the Data Query pipeline parse the LLM output with a bare `PydanticOutputParser` and rely on a `{format_instructions}` prompt placeholder to coax JSON out of the model. When the model returns prose instead of JSON, the parse throws and the whole `Query_Data` tool call fails with an opaque error. This was first observed in the named-entity-recognition (NER) step, but the same fragile pattern is used by the indicator-selection and non-indicator dimension-selection chains.

### What steps will reproduce the bug?

1. Call the `Query_Data` tool (e.g. via MCP `__query_datasets`) with a natural-language query such as "nominal GDP growth for the USA, 2022–2026".
2. The NER step's LLM (GPT-4.1 family) occasionally answers the query conversationally (a hallucinated markdown table) instead of emitting the entity JSON.
3. The `PydanticOutputParser` fails to parse the prose as JSON.

### What is the expected behavior?

Each of these steps always yields a schema-valid response, and the model cannot derail the pipeline by returning free-form text. A transient bad response should not abort the entire tool call.

### What do you see instead?

`OutputParserException: Invalid json output` propagates up through the pipeline into the tool; the MCP provider's catch-all converts it into a generic `" tool failed to execute"`, so the caller (e.g. the Deep Research agent) gets no usable result.

Error log (trimmed)

```text
[chain/error] [tool:Query_Data > ... > parser:PydanticOutputParser] [32ms]
Parser run errored with error:
"OutputParserException('Invalid json output: Here is the nominal Gross Domestic Product (GDP)
growth level information for the United States of America (USA) for the years 2022 to 2026 ...
| Year | Nominal GDP Growth (%) | Label |
| 2022 | 9.1 | Historical |
...')"
json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)

[tool/error] [tool:Query_Data] [10.56s] Tool run errored with error: OutputParserException(...)
ERROR | statgpt.app.mcp.provider | Error executing MCP tool Query_Data
```

The LLM call itself succeeded (~5.5s); only the JSON parse failed (~32ms) — the model returned well-formed prose rather than schema JSON.

### Additional information

**Root cause:** three of the Data Query search-prep chains still use the legacy `{format_instructions}` + `PydanticOutputParser` pattern, which *trusts* the model to voluntarily return JSON:

- `NamedEntitiesChain` (`statgpt/app/chains/data_query/query_builder/misc/named_entities.py`)
- `CandidatesSelectionMappingChainFactory` (`statgpt/app/chains/candidates_selection_mapping.py`, indicator selection)
- `CandidatesSelectionSimpleChainFactory` (`statgpt/app/chains/candidates_selection_simple.py`, non-indicator dimension selection)

Their siblings already constrain output at the API level:

- `DateTimeDimensionChain` -> `.with_structured_output(schema=..., method='json_schema')`
- `DataSetsSelectionChain` -> `.with_structured_output(..., method='json_schema')`

**Suggested fix:** switch all three chains to `.with_structured_output(..., method='json_schema')`, drop the parsers, and remove the now-redundant `{format_instructions}` placeholders from the corresponding prompts in `data_query.yaml` (`namedEntitiesPrompt`, `indicatorsSelectionSystemPrompt`, `validationSystemPrompt`). This makes prose output impossible and eliminates the failure mode.

**Follow-up (separate, optional):** the MCP provider catch-all (`statgpt/app/mcp/provider.py`) still turns *any* internal failure into the same opaque `"tool failed to execute"` message; consider surfacing more actionable errors.

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.