fix: Fix pipeline with join node

[sjrl]

Related Issues

• fixes #issue-number

Proposed Changes:

We have started to develop multi prompt node pipelines for clients. Some of the most popular ones include rewriting the query during the pipeline execution. Here is a simple example using a PromptNode as a spell checker for the incoming user query.

from haystack import Pipeline
from haystack.nodes import BM25Retriever, EmbeddingRetriever, PromptNode, Shaper, JoinDocuments, PromptTemplate
from haystack.document_stores import InMemoryDocumentStore

document_store = InMemoryDocumentStore(use_bm25=True)
dicts = [{"content": "The capital of Germany is Berlin."}, {"content": "The capital of France is Paris."}]
document_store.write_documents(dicts)

query_prompt_node = PromptNode(
    model_name_or_path="gpt-3.5-turbo",
    api_key="",
    default_prompt_template=PromptTemplate("You are a spell checker. Given a user query return the same query with all spelling errors fixed.\nUser Query: {query}\nSpell Checked Query:")
)
shaper = Shaper(
    func="join_strings",
    inputs={"strings": "results"},
    outputs=["query"],
)
qa_prompt_node = PromptNode(
    model_name_or_path="gpt-3.5-turbo",
    api_key="",
    default_prompt_template=PromptTemplate("Answer the user query. Query: {query}")
)
sparse_retriever = BM25Retriever(
    document_store=document_store,
    top_k=2
)
dense_retriever = EmbeddingRetriever(
    document_store=document_store,
    embedding_model="intfloat/e5-base-v2",
    model_format="sentence_transformers",
    top_k=2
)
document_store.update_embeddings(dense_retriever)

pipeline = Pipeline()
pipeline.add_node(component=query_prompt_node, name="QueryPromptNode", inputs=["Query"])
pipeline.add_node(component=shaper, name="ListToString", inputs=["QueryPromptNode"])
pipeline.add_node(component=sparse_retriever, name="BM25", inputs=["ListToString"])
pipeline.add_node(component=dense_retriever, name="Embedding", inputs=["ListToString"])
pipeline.add_node(
    component=JoinDocuments(join_mode="concatenate"), name="Join", inputs=["BM25", "Embedding"]
)
pipeline.add_node(component=qa_prompt_node, name="QAPromptNode", inputs=["Join"])

out = pipeline.run(query="What is the captial of Grmny?", debug=True)
print(out["invocation_context"])
# Current v1.x branch
# {'query': 'What is the captial of Grmny?',  <-- Original Query!!
#   'results': ['The capital of Germany is Berlin.'],
#   'prompts': ['Answer the user query. Query: What is the captial of Grmny?'],  <-- Original Query!!
# After this PR
# {'query': 'What is the capital of Germany?',  <-- Rewritten Query!!
#   'results': ['The capital of Germany is Berlin.'],
#   'prompts': ['Answer the user query. Query: What is the capital of Germany?'],  <-- Rewritten Query!!

The issue is that the query used in the second prompt node (qa_prompt_node) is the original query and not the one rewritten by the first prompt node (query_prompt_node). This has to do with how information the Pipeline.run method passes through JoinNodes specifically.

This PR fixes this issue by updating the Pipeline.run method to properly pass node_output information through a JoinNode.

How did you test it?

Notes for the reviewer

Checklist

• I have read the contributors guidelines and the code of conduct
• I have updated the related issue with new insights and changes
• I added unit tests and updated the docstrings
• I've used one of the conventional commit types for my PR title: fix:, feat:, build:, chore:, ci:, docs:, style:, refactor:, perf:, test:.
• I documented my code
• I ran pre-commit hooks and fixed any issue

[sjrl]

It seems like the mypy error is unrelated to changes made in this PR

haystack/nodes/file_converter/image.py:122: error: Argument 1 to "_image_to_text" of "ImageToTextConverter" has incompatible type "Image"; expected "PpmImageFile"  [arg-type]

[anakin87]

@sjrl I'm trying to fix the mypy error.

I would prefer that this PR be reviewed by someone who knows the 1.x Pipeline run logic better than I do.

[sjrl]

@tstadel I added a invocation_context.copy() to both the PromptNode and the Shaper to avoid overwriting values in the _debug dict. Checked that with the changes that running the example in the main description and printing

out["_debug"]["QueryPromptNode"]
{'input': {'query': 'What is the captial of Grmny?', 'debug': True},
 'output': {'results': ['What is the capital of Germany?'],
  'invocation_context': {'query': 'What is the captial of Grmny?',
   'results': ['What is the capital of Germany?'],
   'prompts': ['You are a spell checker. Given a user query return the same query with all spelling errors fixed.\nUser Query: What is the captial of Grmny?\nSpell Checked Query:']}},
 'runtime': {'prompts_used': ['You are a spell checker. Given a user query return the same query with all spelling errors fixed.\nUser Query: What is the captial of Grmny?\nSpell Checked Query:']},
 'exec_time_ms': 540.78}

gives the expected debug output. Without the shaper change in particular you get

out["_debug"]["QueryPromptNode"]
{'input': {'query': 'What is the captial of Grmny?', 'debug': True},
 'output': {'results': ['What is the capital of Germany?'],
  'invocation_context': {'query': 'What is the capital of Germany?',  <--- HERE: This is wrong
   'results': ['What is the capital of Germany?'],
   'prompts': ['You are a spell checker. Given a user query return the same query with all spelling errors fixed.\nUser Query: What is the captial of Grmny?\nSpell Checked Query:']}},
 'runtime': {'prompts_used': ['You are a spell checker. Given a user query return the same query with all spelling errors fixed.\nUser Query: What is the captial of Grmny?\nSpell Checked Query:']},
 'exec_time_ms': 535.65}

[tstadel]

Thanks for the additional invocation_context fix and test! Looking very good!

[anakin87]

@tstadel thanks for your help!

I left a comment about the release note.
When all is good, feel free to merge.

[tstadel]

Nice! I just checked. This PR will fix the issue regarding batch debug output of nodes previous to JoinNode reported in #2962 (comment) 🎉