Skip to content

Instantly share code, notes, and snippets.

@koganei
Last active April 16, 2023 15:52
Show Gist options
  • Save koganei/263b9890a05843d5483b23c44fbfd1f0 to your computer and use it in GitHub Desktop.
Save koganei/263b9890a05843d5483b23c44fbfd1f0 to your computer and use it in GitHub Desktop.
Adding Post-Processing Pipeline to Haystack Agents
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
# 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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment