LLM Call response of JudgeAgent not always returning <Yes> or <No>

[sschrijver-pon]

System Info

macos = 14.5
python = 3.10.13
pandasai = 2.2.12

🐛 Describe the bug

Using AzureOpenAI agent in combination with JudgeAgent.

llm = AzureOpenAI(...)

judge = JudgeAgent(config={"llm": llm, "verbose": True})

agent = Agent("filepath", config={"llm": llm, "verbose": True}, judge=judge)

The logs show that the following is added to the prompt by the JudgeAgent:

Reason step by step and at the end answer:
1. Explain what the code does
2. Explain what the user query asks for
3. Strictly compare the query with the code that is generated
Always return <Yes> or <No> if exactly meets the requirements

But, the actual answers of the LLM responses do not contain <Yes> or <No> and only answers the questions 1, 2 and 3, so

pandas-ai/pandasai/ee/agents/judge_agent/pipeline/llm_call.py

Lines 44 to 50 in e011e8f

	result = False
	if "<Yes>" in response:
	result = True
	elif "<No>" in response:
	result = False
	else:
	raise InvalidOutputValueMismatch("Invalid response of LLM Call")

throws a pandasai.exceptions.InvalidOutputValueMismatch: Invalid response of LLM Call

[dosubot]

Hello, @sschrijver-pon! I'm here to help you with any bugs, questions, or contributions you have. Let's work together to resolve your issue.

To ensure that the LLM responses include the required <Yes> or <No> to avoid the pandasai.exceptions.InvalidOutputValueMismatch: Invalid response of LLM Call error when using the JudgeAgent with the AzureOpenAI agent, you need to make sure that the LLM is explicitly instructed to include these tags in its responses. This can be done by modifying the prompt given to the LLM to include clear instructions.

Here is an example of how you might modify the prompt:

prompt = """
Please analyze the following data and respond with either <Yes> or <No> based on your analysis.

<Data>
{data}
</Data>

Your response should be either <Yes> or <No>.
"""
response = pipeline_context.config.llm.call(prompt, pipeline_context)

By ensuring that the prompt explicitly asks for <Yes> or <No>, you can help guide the LLM to produce the required output format, thus avoiding the InvalidOutputValueMismatch error ^[1].

Additionally, you can implement unit tests to verify that the LLM responses are correctly formatted. For example:

def test_llm_call(self, sample_df, context, logger, config):
    input_validator = LLMCall()

    config["llm"].call = MagicMock(return_value="<Yes>")

    context = PipelineContext([sample_df], config)

    result = input_validator.execute(input="test", context=context, logger=logger)

    assert isinstance(result, LogicUnitOutput)
    assert result.output is True

def test_llm_call_no(self, sample_df, context, logger, config):
    input_validator = LLMCall()

    config["llm"].call = MagicMock(return_value="<No>")

    context = PipelineContext([sample_df], config)

    result = input_validator.execute(input="test", context=context, logger=logger)

    assert isinstance(result, LogicUnitOutput)
    assert result.output is False

def test_llm_call_with_no_tags(self, sample_df, context, logger, config):
    input_validator = LLMCall()

    config["llm"].call = MagicMock(return_value="yes")

    context = PipelineContext([sample_df], config)

    with pytest.raises(InvalidOutputValueMismatch):
        input_validator.execute(input="test", context=context, logger=logger)

These tests ensure that the LLM responses are correctly formatted and help catch any issues early in the development process ^[1].

To continue talking to Dosu, mention @dosu.

[sschrijver-pon]

Possible solution, change

pandas-ai/pandasai/ee/agents/judge_agent/prompts/templates/judge_agent_prompt.tmpl

Lines 7 to 11 in e011e8f

	Reason step by step and at the end answer:
	1. Explain what the code does
	2. Explain what the user query asks for
	3. Strictly compare the query with the code that is generated
	Always return <Yes> or <No> if exactly meets the requirements

to:

Reason step by step and at the end answer:
1. Explain what the code does
2. Explain what the user query asks for
3. Strictly compare the query with the code that is generated
4. Always return Yes or No if exactly meets the requirements

and

pandas-ai/pandasai/ee/agents/judge_agent/pipeline/llm_call.py

Lines 45 to 48 in e011e8f

	if "<Yes>" in response:
	result = True
	elif "<No>" in response:
	result = False

to

if "4. Yes" in response:
    result = True
elif "4. No" in response:
    result = False

Edit: I also tried the solution below, but it seems to skip the <> in the answer:

4. Always return <Yes> or <No> if exactly meets the requirements