Adding Post-Processing Pipeline to Haystack Agents

CustomAgent.py

class CustomAgent(Agent):

    # Post-Processing Pipeline
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def add_post_processing_pipeline(self, pipeline):
        self.post_processing_pipeline = pipeline
        self.callback_manager.on_agent_finish += self.on_result_post_process

    def on_result_post_process(self, last_step):
        if last_step.is_last():
            pipeline_result = self.post_processing_pipeline.run(
                query=last_step.prompt_node_response
            )

            last_step.prompt_node_response = last_step.prompt_node_response.replace(
                last_step.prompt_node_response, pipeline_result["results"]
            )

            return last_step

running-agent.py

# Adding the pre-processing pipeline to Agent
post_processing_pipeline = Pipeline()
post_processing_pipeline.add_node(
    component=add_navigation_link_node, name="Navigation", inputs=["Query"])
agent.add_post_processing_pipeline(post_processing_pipeline)