LLM Outputs Disappear After Connecting Output Validator as a Loop to Prompt Builder

[dafajon]

I am using AmazonBedrockChatGenerator to be able to use Claude models. As result my prompt builder component is a DynamicChatPromptBuilder.

As a runtime variable I am defining invalid_replies and error_message to guide the prompt in order to fix errors if there are any. The if condition is stated in a Jinja template in the prompt.

When I connect the OutputValidator outputs to runtime_variables, the pipeline outputs: {}. However when there is no loop and I initialize the runtime_variables explicitly in the data argument of run and break the loop, my pipeline works.

Haystack version: 2.1.2

Below is the code:

system_message = """You are an experienced and accurate Turkish CX speacialist that classifies customer comments into pre-defined categories below:\n
Negative experience labels:
- Late delivery
- Rotten/spoilt item
- Bad Courier behavior

Positive experience labels:
- Good courier behavior
- Thanks & appreciation
- Love message to courier
- Fast delivery
- Quality of products

Create a JSON object as a response. The fields are: 'positive_experience', 'negative_experience'.
Assign at least one of the pre-defined labels to the given customer comment under positive and negative experience fields. If the comment has a positive experience, list the label under positive_experience field. If the comments has a negative_experience, list it under the 'negative_experience' field.
"""
user_message = """Here is the comment:\n{{comment}}\n. Just return the category names in the list. If there aren't any, return an empty list.

{% if invalid_replies and error_message %}
  You already created the following output in a previous attempt: {{invalid_replies}}
  However, this doesn't comply with the format requirements from above and triggered this Python exception: {{error_message}}
  Correct the output and try again. Just return the corrected output without any extra explanations.
{% endif %}
"""

messages = [ChatMessage.from_system(system_message),
            ChatMessage.from_user(user_message)]

prompt_builder = DynamicChatPromptBuilder(runtime_variables=["invalid_replies", "error_message"])

llm = AmazonBedrockChatGenerator(model='anthropic.claude-3-sonnet-20240229-v1:0', 
                                aws_region_name=Secret.from_token(env.get("aws_region")),
                                aws_profile_name=Secret.from_token(env.get("aws_profile")))

class ClassifierOutput(BaseModel):
    positive_experience: List[Optional[str]]
    negative_experience: List[Optional[str]]

@component
class OutputValidator:
    def __init__(self, pydantic_model: BaseModel):
        self.pydantic_model = pydantic_model
        self.iteration_counter = 0


    @component.output_types(valid_replies=List[ChatMessage], invalid_replies=Optional[List[ChatMessage]], error_message=Optional[str])
    def run(self, replies: List[ChatMessage]):

        self.iteration_counter += 1

        print(Fore.BLUE + "Running output through validator")

        try:
            output_dict = json.loads(replies[0].content)
            self.pydantic_model.parse_obj(output_dict)

            print(
                Fore.GREEN
                + f"OutputValidator at Iteration {self.iteration_counter}: Valid JSON from LLM - No need for looping: {replies[0].content}"
            )
            return {"valid_replies": replies}

        except(ValueError, ValidationError) as e:
            print(
                Fore.RED
                + f"OutputValidator at Iteration {self.iteration_counter}: Invalid JSON from LLM - Let's try again.\n"
                f"Output from LLM:\n {replies[0].content} \n"
                f"Error from OutputValidator: {e}"
            )
            return {"invalid_replies": replies, "error_message": str(e)}

pipeline = Pipeline()
pipeline.add_component("prompt_builder", prompt_builder)

pipeline.add_component("llm", llm)
pipeline.add_component("output_validator", validator)

pipeline.connect("prompt_builder.prompt", "llm.messages")
pipeline.connect("llm.replies", "output_validator.replies")

# Output becomes empty dict when I make this loop connection.
pipeline.connect("output_validator.invalid_replies", "prompt_builder.invalid_replies")
pipeline.connect("output_validator.error_message", "prompt_builder.error_message")

res = pipeline.run(data={"prompt_builder": {"template_variables": {"comment": comment}, 
                                            "prompt_source": messages,
                                            # Explicitly inject runtime_variables when there is no loop connection of the validator.
                                            # "invalid_replies": {},
                                            # "error_message":""
                                           }
                        },
                   include_outputs_from=["prompt_builder", "llm", "output_validator"])

Output:

{}

Output without Loop:

{'output_validator': {'valid_replies': [ChatMessage(content='{\n    "positive_experience": ["Quality of products"],\n    "negative_experience": ["Bad Courier behavior"]\n}', role=<ChatRole.ASSISTANT: 'assistant'>, name=None, meta={'id': 'msg_01Sa75CDq79Ynv52RUx59hRb', 'model': 'claude-3-sonnet-28k-20240229', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 226, 'output_tokens': 31}})]},
 'prompt_builder': {'prompt': [ChatMessage(content="You are an experienced and accurate Turkish CX speacialist that classifies customer comments into pre-defined categories below:\n\nNegative experience labels:\n- Late delivery\n- Rotten/spoilt item\n- Bad Courier behavior\n\nPositive experience labels:\n- Good courier behavior\n- Thanks & appreciation\n- Love message to courier\n- Fast delivery\n- Quality of products\n\nCreate a JSON object as a response. The fields are: 'positive_experience', 'negative_experience'.\nAssign at least one of the pre-defined labels to the given customer comment under positive and negative experience fields. If the comment has a positive experience, list the label under positive_experience field. If the comments has a negative_experience, list it under the 'negative_experience' field.", role=<ChatRole.SYSTEM: 'system'>, name=None, meta={}),
   ChatMessage(content="Here is the comment:\nürünler iyiydi ama kurye kafama atsaydı keşke.\n. Just return the category names in the list. If there aren't any, return an empty list.\n\n", role=<ChatRole.USER: 'user'>, name=None, meta={})]},
 'llm': {'replies': [ChatMessage(content='{\n    "positive_experience": ["Quality of products"],\n    "negative_experience": ["Bad Courier behavior"]\n}', role=<ChatRole.ASSISTANT: 'assistant'>, name=None, meta={'id': 'msg_01Sa75CDq79Ynv52RUx59hRb', 'model': 'claude-3-sonnet-28k-20240229', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 226, 'output_tokens': 31}})]}}

[bilgeyucel]

I could reproduce this issue. It seems like the combination of optional outputs and looping pipelines confuses the pipeline.run
Here's the related issue: #7533