wehub-resource-sync
c56bef871b
Sync docs with Docusaurus / sync (push) Waiting to run
Tests / Check if changed (push) Waiting to run
Tests / format (push) Blocked by required conditions
Tests / check-imports (push) Blocked by required conditions
Tests / Unit / macos-latest (push) Blocked by required conditions
Tests / Unit / ubuntu-latest (push) Blocked by required conditions
Tests / Unit / windows-latest (push) Blocked by required conditions
Tests / mypy (push) Blocked by required conditions
Tests / Integration / ubuntu-latest (push) Blocked by required conditions
Tests / Integration / macos-latest (push) Blocked by required conditions
Tests / Integration / windows-latest (push) Blocked by required conditions
Tests / notify-slack-on-failure (push) Blocked by required conditions
Tests / Mark tests as completed (push) Blocked by required conditions
Docker image release / Build base image (push) Waiting to run
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
5124 lines
244 KiB
Python
5124 lines
244 KiB
Python
# ruff: noqa: D103
|
||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||
#
|
||
# SPDX-License-Identifier: Apache-2.0
|
||
|
||
import json
|
||
import re
|
||
from copy import deepcopy
|
||
from typing import Any
|
||
from unittest.mock import ANY
|
||
|
||
import pytest
|
||
from pandas import DataFrame
|
||
from pytest_bdd import given, scenarios
|
||
|
||
from haystack import Document, Pipeline, component
|
||
from haystack.components.builders import AnswerBuilder, ChatPromptBuilder, PromptBuilder
|
||
from haystack.components.converters import (
|
||
CSVToDocument,
|
||
HTMLToDocument,
|
||
JSONConverter,
|
||
OutputAdapter,
|
||
TextFileToDocument,
|
||
)
|
||
from haystack.components.joiners import AnswerJoiner, BranchJoiner, DocumentJoiner, StringJoiner
|
||
from haystack.components.preprocessors import DocumentCleaner, DocumentSplitter
|
||
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
|
||
from haystack.components.routers import ConditionalRouter, FileTypeRouter
|
||
from haystack.components.routers.conditional_router import Route
|
||
from haystack.core.component.types import Variadic
|
||
from haystack.dataclasses import ByteStream, ChatMessage, ChatRole, GeneratedAnswer, TextContent
|
||
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
||
from haystack.document_stores.types import DuplicatePolicy
|
||
from haystack.testing.factory import component_class
|
||
from haystack.testing.sample_components import (
|
||
Accumulate,
|
||
AddFixedValue,
|
||
Double,
|
||
FString,
|
||
Greet,
|
||
Hello,
|
||
Parity,
|
||
Remainder,
|
||
Repeat,
|
||
StringListJoiner,
|
||
Subtract,
|
||
Sum,
|
||
TextSplitter,
|
||
Threshold,
|
||
)
|
||
from test.core.pipeline.features.conftest import FixedGenerator, PipelineRunData
|
||
|
||
pytestmark = [pytest.mark.usefixtures("pipeline_run_mode"), pytest.mark.integration]
|
||
|
||
scenarios("pipeline_run.feature")
|
||
|
||
|
||
@given("a pipeline that has no components", target_fixture="pipeline_data")
|
||
def pipeline_that_has_no_components():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
return pipeline, [PipelineRunData(inputs={}, expected_outputs={})]
|
||
|
||
|
||
@given("a pipeline that is linear", target_fixture="pipeline_data")
|
||
def pipeline_that_is_linear():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("first_addition", AddFixedValue(add=2))
|
||
pipeline.add_component("second_addition", AddFixedValue())
|
||
pipeline.add_component("double", Double())
|
||
pipeline.connect("first_addition", "double")
|
||
pipeline.connect("double", "second_addition")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"first_addition": {"value": 1}},
|
||
expected_outputs={"second_addition": {"result": 7}},
|
||
expected_component_calls={
|
||
("first_addition", 1): {"value": 1, "add": None},
|
||
("double", 1): {"value": 3},
|
||
("second_addition", 1): {"value": 6, "add": None},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has an infinite loop", target_fixture="pipeline_data")
|
||
def pipeline_that_has_an_infinite_loop():
|
||
routes: list[Route] = [
|
||
{"condition": "{{number > 2}}", "output": "{{number}}", "output_name": "big_number", "output_type": int},
|
||
{"condition": "{{number <= 2}}", "output": "{{number + 2}}", "output_name": "small_number", "output_type": int},
|
||
]
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
pipe.add_component("main_input", BranchJoiner(int))
|
||
pipe.add_component("first_router", ConditionalRouter(routes=routes))
|
||
pipe.add_component("second_router", ConditionalRouter(routes=routes))
|
||
|
||
pipe.connect("main_input", "first_router.number")
|
||
pipe.connect("first_router.big_number", "second_router.number")
|
||
pipe.connect("second_router.big_number", "main_input")
|
||
|
||
return pipe, [PipelineRunData({"main_input": {"value": 3}})]
|
||
|
||
|
||
@given("a pipeline that is really complex with lots of components, forks, and loops", target_fixture="pipeline_data")
|
||
def pipeline_complex():
|
||
pipeline = Pipeline(max_runs_per_component=3)
|
||
pipeline.add_component("greet_first", Greet(message="Hello, the value is {value}."))
|
||
pipeline.add_component("accumulate_1", Accumulate())
|
||
pipeline.add_component("add_two", AddFixedValue(add=2))
|
||
pipeline.add_component("parity", Parity())
|
||
pipeline.add_component("add_one", AddFixedValue(add=1))
|
||
pipeline.add_component("accumulate_2", Accumulate())
|
||
|
||
pipeline.add_component("branch_joiner", BranchJoiner(type_=int))
|
||
pipeline.add_component("below_10", Threshold(threshold=10))
|
||
pipeline.add_component("double", Double())
|
||
|
||
pipeline.add_component("greet_again", Greet(message="Hello again, now the value is {value}."))
|
||
pipeline.add_component("sum", Sum())
|
||
|
||
pipeline.add_component("greet_enumerator", Greet(message="Hello from enumerator, here the value became {value}."))
|
||
pipeline.add_component("enumerate", Repeat(outputs=["first", "second"]))
|
||
pipeline.add_component("add_three", AddFixedValue(add=3))
|
||
|
||
pipeline.add_component("diff", Subtract())
|
||
pipeline.add_component("greet_one_last_time", Greet(message="Bye bye! The value here is {value}!"))
|
||
pipeline.add_component("replicate", Repeat(outputs=["first", "second"]))
|
||
pipeline.add_component("add_five", AddFixedValue(add=5))
|
||
pipeline.add_component("add_four", AddFixedValue(add=4))
|
||
pipeline.add_component("accumulate_3", Accumulate())
|
||
|
||
pipeline.connect("greet_first", "accumulate_1")
|
||
pipeline.connect("accumulate_1", "add_two")
|
||
pipeline.connect("add_two", "parity")
|
||
|
||
pipeline.connect("parity.even", "greet_again")
|
||
pipeline.connect("greet_again", "sum.values")
|
||
pipeline.connect("sum", "diff.first_value")
|
||
pipeline.connect("diff", "greet_one_last_time")
|
||
pipeline.connect("greet_one_last_time", "replicate")
|
||
pipeline.connect("replicate.first", "add_five.value")
|
||
pipeline.connect("replicate.second", "add_four.value")
|
||
pipeline.connect("add_four", "accumulate_3")
|
||
|
||
pipeline.connect("parity.odd", "add_one.value")
|
||
pipeline.connect("add_one", "branch_joiner.value")
|
||
pipeline.connect("branch_joiner", "below_10")
|
||
|
||
pipeline.connect("below_10.below", "double")
|
||
pipeline.connect("double", "branch_joiner.value")
|
||
|
||
pipeline.connect("below_10.above", "accumulate_2")
|
||
pipeline.connect("accumulate_2", "diff.second_value")
|
||
|
||
pipeline.connect("greet_enumerator", "enumerate")
|
||
pipeline.connect("enumerate.second", "sum.values")
|
||
|
||
pipeline.connect("enumerate.first", "add_three.value")
|
||
pipeline.connect("add_three", "sum.values")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"greet_first": {"value": 1}, "greet_enumerator": {"value": 1}},
|
||
expected_outputs={"accumulate_3": {"value": -7}, "add_five": {"result": -6}},
|
||
expected_component_calls={
|
||
("greet_first", 1): {"value": 1, "log_level": None, "message": None},
|
||
("greet_enumerator", 1): {"value": 1, "log_level": None, "message": None},
|
||
("accumulate_1", 1): {"value": 1},
|
||
("add_two", 1): {"value": 1, "add": None},
|
||
("parity", 1): {"value": 3},
|
||
("add_one", 1): {"value": 3, "add": None},
|
||
("branch_joiner", 1): {"value": [4]},
|
||
("below_10", 1): {"value": 4, "threshold": None},
|
||
("double", 1): {"value": 4},
|
||
("branch_joiner", 2): {"value": [8]},
|
||
("below_10", 2): {"value": 8, "threshold": None},
|
||
("double", 2): {"value": 8},
|
||
("branch_joiner", 3): {"value": [16]},
|
||
("below_10", 3): {"value": 16, "threshold": None},
|
||
("accumulate_2", 1): {"value": 16},
|
||
("enumerate", 1): {"value": 1},
|
||
("add_three", 1): {"value": 1, "add": None},
|
||
("sum", 1): {"values": [1, 4]},
|
||
("diff", 1): {"first_value": 5, "second_value": 16},
|
||
("greet_one_last_time", 1): {"value": -11, "log_level": None, "message": None},
|
||
("replicate", 1): {"value": -11},
|
||
("add_five", 1): {"value": -11, "add": None},
|
||
("add_four", 1): {"value": -11, "add": None},
|
||
("accumulate_3", 1): {"value": -7},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has a single component with a default input", target_fixture="pipeline_data")
|
||
def pipeline_that_has_a_single_component_with_a_default_input():
|
||
@component
|
||
class WithDefault:
|
||
@component.output_types(b=int)
|
||
def run(self, a: int, b: int = 2) -> dict[str, int]:
|
||
return {"c": a + b}
|
||
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("with_defaults", WithDefault())
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"with_defaults": {"a": 40, "b": 30}},
|
||
expected_outputs={"with_defaults": {"c": 70}},
|
||
expected_component_calls={("with_defaults", 1): {"a": 40, "b": 30}},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"with_defaults": {"a": 40}},
|
||
expected_outputs={"with_defaults": {"c": 42}},
|
||
expected_component_calls={("with_defaults", 1): {"a": 40, "b": 2}},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has two loops of identical lengths", target_fixture="pipeline_data")
|
||
def pipeline_that_has_two_loops_of_identical_lengths():
|
||
pipeline = Pipeline(max_runs_per_component=2)
|
||
pipeline.add_component("branch_joiner", BranchJoiner(type_=int))
|
||
pipeline.add_component("remainder", Remainder(divisor=3))
|
||
pipeline.add_component("add_one", AddFixedValue(add=1))
|
||
pipeline.add_component("add_two", AddFixedValue(add=2))
|
||
|
||
pipeline.connect("branch_joiner.value", "remainder.value")
|
||
pipeline.connect("remainder.remainder_is_1", "add_two.value")
|
||
pipeline.connect("remainder.remainder_is_2", "add_one.value")
|
||
pipeline.connect("add_two", "branch_joiner.value")
|
||
pipeline.connect("add_one", "branch_joiner.value")
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 0}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 0}},
|
||
expected_component_calls={("branch_joiner", 1): {"value": [0]}, ("remainder", 1): {"value": 0}},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 3}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 3}},
|
||
expected_component_calls={("branch_joiner", 1): {"value": [3]}, ("remainder", 1): {"value": 3}},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 4}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 6}},
|
||
expected_component_calls={
|
||
("branch_joiner", 1): {"value": [4]},
|
||
("remainder", 1): {"value": 4},
|
||
("add_two", 1): {"value": 4, "add": None},
|
||
("branch_joiner", 2): {"value": [6]},
|
||
("remainder", 2): {"value": 6},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 5}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 6}},
|
||
expected_component_calls={
|
||
("branch_joiner", 1): {"value": [5]},
|
||
("remainder", 1): {"value": 5},
|
||
("add_one", 1): {"value": 5, "add": None},
|
||
("branch_joiner", 2): {"value": [6]},
|
||
("remainder", 2): {"value": 6},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 6}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 6}},
|
||
expected_component_calls={("branch_joiner", 1): {"value": [6]}, ("remainder", 1): {"value": 6}},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has two loops of different lengths", target_fixture="pipeline_data")
|
||
def pipeline_that_has_two_loops_of_different_lengths():
|
||
pipeline = Pipeline(max_runs_per_component=2)
|
||
pipeline.add_component("branch_joiner", BranchJoiner(type_=int))
|
||
pipeline.add_component("remainder", Remainder(divisor=3))
|
||
pipeline.add_component("add_one", AddFixedValue(add=1))
|
||
pipeline.add_component("add_two_1", AddFixedValue(add=1))
|
||
pipeline.add_component("add_two_2", AddFixedValue(add=1))
|
||
|
||
pipeline.connect("branch_joiner.value", "remainder.value")
|
||
pipeline.connect("remainder.remainder_is_1", "add_two_1.value")
|
||
pipeline.connect("add_two_1", "add_two_2.value")
|
||
pipeline.connect("add_two_2", "branch_joiner")
|
||
pipeline.connect("remainder.remainder_is_2", "add_one.value")
|
||
pipeline.connect("add_one", "branch_joiner")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 0}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 0}},
|
||
expected_component_calls={("branch_joiner", 1): {"value": [0]}, ("remainder", 1): {"value": 0}},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 3}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 3}},
|
||
expected_component_calls={("branch_joiner", 1): {"value": [3]}, ("remainder", 1): {"value": 3}},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 4}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 6}},
|
||
expected_component_calls={
|
||
("branch_joiner", 1): {"value": [4]},
|
||
("remainder", 1): {"value": 4},
|
||
("add_two_1", 1): {"value": 4, "add": None},
|
||
("add_two_2", 1): {"value": 5, "add": None},
|
||
("branch_joiner", 2): {"value": [6]},
|
||
("remainder", 2): {"value": 6},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 5}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 6}},
|
||
expected_component_calls={
|
||
("branch_joiner", 1): {"value": [5]},
|
||
("remainder", 1): {"value": 5},
|
||
("add_one", 1): {"value": 5, "add": None},
|
||
("branch_joiner", 2): {"value": [6]},
|
||
("remainder", 2): {"value": 6},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"branch_joiner": {"value": 6}},
|
||
expected_outputs={"remainder": {"remainder_is_0": 6}},
|
||
expected_component_calls={("branch_joiner", 1): {"value": [6]}, ("remainder", 1): {"value": 6}},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has a single loop with two conditional branches", target_fixture="pipeline_data")
|
||
def pipeline_that_has_a_single_loop_with_two_conditional_branches():
|
||
pipeline = Pipeline(max_runs_per_component=3)
|
||
|
||
pipeline.add_component("add_one", AddFixedValue(add=1))
|
||
pipeline.add_component("branch_joiner", BranchJoiner(type_=int))
|
||
pipeline.add_component("below_10", Threshold(threshold=10))
|
||
pipeline.add_component("below_5", Threshold(threshold=5))
|
||
pipeline.add_component("add_three", AddFixedValue(add=3))
|
||
pipeline.add_component("accumulator", Accumulate())
|
||
pipeline.add_component("add_two", AddFixedValue(add=2))
|
||
|
||
pipeline.connect("add_one.result", "branch_joiner")
|
||
pipeline.connect("branch_joiner.value", "below_10.value")
|
||
pipeline.connect("below_10.below", "accumulator.value")
|
||
pipeline.connect("accumulator.value", "below_5.value")
|
||
pipeline.connect("below_5.above", "add_three.value")
|
||
pipeline.connect("below_5.below", "branch_joiner")
|
||
pipeline.connect("add_three.result", "branch_joiner")
|
||
pipeline.connect("below_10.above", "add_two.value")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"add_one": {"value": 3}},
|
||
expected_outputs={"add_two": {"result": 13}},
|
||
expected_component_calls={
|
||
("accumulator", 1): {"value": 4},
|
||
("accumulator", 2): {"value": 4},
|
||
("add_one", 1): {"add": None, "value": 3},
|
||
("add_three", 1): {"add": None, "value": 8},
|
||
("add_two", 1): {"add": None, "value": 11},
|
||
("below_10", 1): {"threshold": None, "value": 4},
|
||
("below_10", 2): {"threshold": None, "value": 4},
|
||
("below_10", 3): {"threshold": None, "value": 11},
|
||
("below_5", 1): {"threshold": None, "value": 4},
|
||
("below_5", 2): {"threshold": None, "value": 8},
|
||
("branch_joiner", 1): {"value": [4]},
|
||
("branch_joiner", 2): {"value": [4]},
|
||
("branch_joiner", 3): {"value": [11]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has a component with dynamic inputs defined in init", target_fixture="pipeline_data")
|
||
def pipeline_that_has_a_component_with_dynamic_inputs_defined_in_init():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("hello", Hello())
|
||
pipeline.add_component("fstring", FString(template="This is the greeting: {greeting}!", variables=["greeting"]))
|
||
pipeline.add_component("splitter", TextSplitter())
|
||
pipeline.connect("hello.output", "fstring.greeting")
|
||
pipeline.connect("fstring.string", "splitter.sentence")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"hello": {"word": "Alice"}},
|
||
expected_outputs={"splitter": {"output": ["This", "is", "the", "greeting:", "Hello,", "Alice!!"]}},
|
||
expected_component_calls={
|
||
("fstring", 1): {"greeting": "Hello, Alice!", "template": None},
|
||
("hello", 1): {"word": "Alice"},
|
||
("splitter", 1): {"sentence": "This is the greeting: Hello, Alice!!"},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"hello": {"word": "Alice"}, "fstring": {"template": "Received: {greeting}"}},
|
||
expected_outputs={"splitter": {"output": ["Received:", "Hello,", "Alice!"]}},
|
||
expected_component_calls={
|
||
("fstring", 1): {"greeting": "Hello, Alice!", "template": "Received: {greeting}"},
|
||
("hello", 1): {"word": "Alice"},
|
||
("splitter", 1): {"sentence": "Received: Hello, Alice!"},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has two branches that don't merge", target_fixture="pipeline_data")
|
||
def pipeline_that_has_two_branches_that_dont_merge():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("add_one", AddFixedValue(add=1))
|
||
pipeline.add_component("parity", Parity())
|
||
pipeline.add_component("add_ten", AddFixedValue(add=10))
|
||
pipeline.add_component("double", Double())
|
||
pipeline.add_component("add_three", AddFixedValue(add=3))
|
||
|
||
pipeline.connect("add_one.result", "parity.value")
|
||
pipeline.connect("parity.even", "add_ten.value")
|
||
pipeline.connect("parity.odd", "double.value")
|
||
pipeline.connect("add_ten.result", "add_three.value")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"add_one": {"value": 1}},
|
||
expected_outputs={"add_three": {"result": 15}},
|
||
expected_component_calls={
|
||
("add_one", 1): {"add": None, "value": 1},
|
||
("add_ten", 1): {"add": None, "value": 2},
|
||
("add_three", 1): {"add": None, "value": 12},
|
||
("parity", 1): {"value": 2},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"add_one": {"value": 2}},
|
||
expected_outputs={"double": {"value": 6}},
|
||
expected_component_calls={
|
||
("add_one", 1): {"add": None, "value": 2},
|
||
("double", 1): {"value": 3},
|
||
("parity", 1): {"value": 3},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has three branches that don't merge", target_fixture="pipeline_data")
|
||
def pipeline_that_has_three_branches_that_dont_merge():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("add_one", AddFixedValue(add=1))
|
||
pipeline.add_component("repeat", Repeat(outputs=["first", "second"]))
|
||
pipeline.add_component("add_ten", AddFixedValue(add=10))
|
||
pipeline.add_component("double", Double())
|
||
pipeline.add_component("add_three", AddFixedValue(add=3))
|
||
pipeline.add_component("add_one_again", AddFixedValue(add=1))
|
||
|
||
pipeline.connect("add_one.result", "repeat.value")
|
||
pipeline.connect("repeat.first", "add_ten.value")
|
||
pipeline.connect("repeat.second", "double.value")
|
||
pipeline.connect("repeat.second", "add_three.value")
|
||
pipeline.connect("add_three.result", "add_one_again.value")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"add_one": {"value": 1}},
|
||
expected_outputs={"add_one_again": {"result": 6}, "add_ten": {"result": 12}, "double": {"value": 4}},
|
||
expected_component_calls={
|
||
("add_one", 1): {"add": None, "value": 1},
|
||
("add_one_again", 1): {"add": None, "value": 5},
|
||
("add_ten", 1): {"add": None, "value": 2},
|
||
("add_three", 1): {"add": None, "value": 2},
|
||
("double", 1): {"value": 2},
|
||
("repeat", 1): {"value": 2},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has two branches that merge", target_fixture="pipeline_data")
|
||
def pipeline_that_has_two_branches_that_merge():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("first_addition", AddFixedValue(add=2))
|
||
pipeline.add_component("second_addition", AddFixedValue(add=2))
|
||
pipeline.add_component("third_addition", AddFixedValue(add=2))
|
||
pipeline.add_component("diff", Subtract())
|
||
pipeline.add_component("fourth_addition", AddFixedValue(add=1))
|
||
|
||
pipeline.connect("first_addition.result", "second_addition.value")
|
||
pipeline.connect("second_addition.result", "diff.first_value")
|
||
pipeline.connect("third_addition.result", "diff.second_value")
|
||
pipeline.connect("diff", "fourth_addition.value")
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"first_addition": {"value": 1}, "third_addition": {"value": 1}},
|
||
expected_outputs={"fourth_addition": {"result": 3}},
|
||
expected_component_calls={
|
||
("diff", 1): {"first_value": 5, "second_value": 3},
|
||
("first_addition", 1): {"add": None, "value": 1},
|
||
("fourth_addition", 1): {"add": None, "value": 2},
|
||
("second_addition", 1): {"add": None, "value": 3},
|
||
("third_addition", 1): {"add": None, "value": 1},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has different combinations of branches that merge and do not merge", target_fixture="pipeline_data"
|
||
)
|
||
def pipeline_that_has_different_combinations_of_branches_that_merge_and_do_not_merge():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("add_one", AddFixedValue())
|
||
pipeline.add_component("parity", Parity())
|
||
pipeline.add_component("add_ten", AddFixedValue(add=10))
|
||
pipeline.add_component("double", Double())
|
||
pipeline.add_component("add_four", AddFixedValue(add=4))
|
||
pipeline.add_component("add_two", AddFixedValue())
|
||
pipeline.add_component("add_two_as_well", AddFixedValue())
|
||
pipeline.add_component("diff", Subtract())
|
||
|
||
pipeline.connect("add_one.result", "parity.value")
|
||
pipeline.connect("parity.even", "add_four.value")
|
||
pipeline.connect("parity.odd", "double.value")
|
||
pipeline.connect("add_ten.result", "diff.first_value")
|
||
pipeline.connect("double.value", "diff.second_value")
|
||
pipeline.connect("parity.odd", "add_ten.value")
|
||
pipeline.connect("add_four.result", "add_two.value")
|
||
pipeline.connect("add_four.result", "add_two_as_well.value")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"add_one": {"value": 1}, "add_two": {"add": 2}, "add_two_as_well": {"add": 2}},
|
||
expected_outputs={"add_two": {"result": 8}, "add_two_as_well": {"result": 8}},
|
||
expected_component_calls={
|
||
("add_four", 1): {"add": None, "value": 2},
|
||
("add_one", 1): {"add": None, "value": 1},
|
||
("add_two", 1): {"add": 2, "value": 6},
|
||
("add_two_as_well", 1): {"add": 2, "value": 6},
|
||
("parity", 1): {"value": 2},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"add_one": {"value": 2}, "add_two": {"add": 2}, "add_two_as_well": {"add": 2}},
|
||
expected_outputs={"diff": {"difference": 7}},
|
||
expected_component_calls={
|
||
("add_one", 1): {"add": None, "value": 2},
|
||
("add_ten", 1): {"add": None, "value": 3},
|
||
("diff", 1): {"first_value": 13, "second_value": 6},
|
||
("double", 1): {"value": 3},
|
||
("parity", 1): {"value": 3},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has two branches, one of which loops back", target_fixture="pipeline_data")
|
||
def pipeline_that_has_two_branches_one_of_which_loops_back():
|
||
pipeline = Pipeline(max_runs_per_component=10)
|
||
pipeline.add_component("add_zero", AddFixedValue(add=0))
|
||
pipeline.add_component("branch_joiner", BranchJoiner(type_=int))
|
||
pipeline.add_component("sum", Sum())
|
||
pipeline.add_component("below_10", Threshold(threshold=10))
|
||
pipeline.add_component("add_one", AddFixedValue(add=1))
|
||
pipeline.add_component("counter", Accumulate())
|
||
pipeline.add_component("add_two", AddFixedValue(add=2))
|
||
|
||
pipeline.connect("add_zero", "branch_joiner.value")
|
||
pipeline.connect("branch_joiner", "below_10.value")
|
||
pipeline.connect("below_10.below", "add_one.value")
|
||
pipeline.connect("add_one.result", "counter.value")
|
||
pipeline.connect("counter.value", "branch_joiner.value")
|
||
pipeline.connect("below_10.above", "add_two.value")
|
||
pipeline.connect("add_two.result", "sum.values")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"add_zero": {"value": 8}, "sum": {"values": 2}},
|
||
expected_outputs={"sum": {"total": 23}},
|
||
expected_component_calls={
|
||
("add_one", 1): {"add": None, "value": 8},
|
||
("add_one", 2): {"add": None, "value": 9},
|
||
("add_two", 1): {"add": None, "value": 19},
|
||
("add_zero", 1): {"add": None, "value": 8},
|
||
("below_10", 1): {"threshold": None, "value": 8},
|
||
("below_10", 2): {"threshold": None, "value": 9},
|
||
("below_10", 3): {"threshold": None, "value": 19},
|
||
("branch_joiner", 1): {"value": [8]},
|
||
("branch_joiner", 2): {"value": [9]},
|
||
("branch_joiner", 3): {"value": [19]},
|
||
("counter", 1): {"value": 9},
|
||
("counter", 2): {"value": 10},
|
||
("sum", 1): {"values": [2, 21]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has a component with mutable input", target_fixture="pipeline_data")
|
||
def pipeline_that_has_a_component_with_mutable_input():
|
||
@component
|
||
class InputMangler:
|
||
@component.output_types(mangled_list=list[str])
|
||
def run(self, input_list: list[str]) -> dict[str, list[str]]:
|
||
input_list.append("extra_item")
|
||
return {"mangled_list": input_list}
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
pipe.add_component("mangler1", InputMangler())
|
||
pipe.add_component("mangler2", InputMangler())
|
||
pipe.add_component("concat1", StringListJoiner())
|
||
pipe.add_component("concat2", StringListJoiner())
|
||
pipe.connect("mangler1", "concat1")
|
||
pipe.connect("mangler2", "concat2")
|
||
|
||
input_list = ["foo", "bar"]
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"mangler1": {"input_list": input_list}, "mangler2": {"input_list": input_list}},
|
||
expected_outputs={
|
||
"concat1": {"output": ["foo", "bar", "extra_item"]},
|
||
"concat2": {"output": ["foo", "bar", "extra_item"]},
|
||
},
|
||
expected_component_calls={
|
||
("concat1", 1): {"inputs": [["foo", "bar", "extra_item"]]},
|
||
("concat2", 1): {"inputs": [["foo", "bar", "extra_item"]]},
|
||
("mangler1", 1): {"input_list": ["foo", "bar"]},
|
||
("mangler2", 1): {"input_list": ["foo", "bar"]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has a component with mutable output sent to multiple inputs", target_fixture="pipeline_data")
|
||
def pipeline_that_has_a_component_with_mutable_output_sent_to_multiple_inputs():
|
||
@component
|
||
class PassThroughPromptBuilder:
|
||
# This is a pass-through component that returns the same input
|
||
@component.output_types(prompt=list[ChatMessage])
|
||
def run(self, prompt_source: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||
return {"prompt": prompt_source}
|
||
|
||
@component
|
||
class MessageMerger:
|
||
@component.output_types(merged_message=str)
|
||
def run(self, messages: list[ChatMessage], metadata: dict | None = None) -> dict[str, str]:
|
||
return {"merged_message": "\n".join(t.text or "" for t in messages)}
|
||
|
||
@component
|
||
class FakeGenerator:
|
||
# This component is a fake generator that always returns the same message
|
||
@component.output_types(replies=list[ChatMessage])
|
||
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||
return {"replies": [ChatMessage.from_assistant("Fake message")]}
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
pipe.add_component("prompt_builder", PassThroughPromptBuilder())
|
||
pipe.add_component("llm", FakeGenerator())
|
||
pipe.add_component("mm1", MessageMerger())
|
||
pipe.add_component("mm2", MessageMerger())
|
||
|
||
pipe.connect("prompt_builder.prompt", "llm.messages")
|
||
pipe.connect("prompt_builder.prompt", "mm1")
|
||
pipe.connect("llm.replies", "mm2")
|
||
|
||
messages = [
|
||
ChatMessage.from_system("Always respond in English even if some input data is in other languages."),
|
||
ChatMessage.from_user("Tell me about Berlin"),
|
||
]
|
||
params = {"metadata": {"metadata_key": "metadata_value", "meta2": "value2"}}
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"mm1": params, "mm2": params, "prompt_builder": {"prompt_source": messages}},
|
||
expected_outputs={
|
||
"mm1": {
|
||
"merged_message": "Always respond in English even if some input data is in other languages.\n"
|
||
"Tell me about Berlin"
|
||
},
|
||
"mm2": {"merged_message": "Fake message"},
|
||
},
|
||
expected_component_calls={
|
||
("llm", 1): {"messages": messages},
|
||
("mm1", 1): {
|
||
"messages": messages,
|
||
"metadata": {"meta2": "value2", "metadata_key": "metadata_value"},
|
||
},
|
||
("mm2", 1): {
|
||
"messages": [ChatMessage.from_assistant(text="Fake message")],
|
||
"metadata": {"meta2": "value2", "metadata_key": "metadata_value"},
|
||
},
|
||
("prompt_builder", 1): {"prompt_source": messages},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has a greedy and variadic component after a component with default input",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_that_has_a_greedy_and_variadic_component_after_a_component_with_default_input():
|
||
# ruff: noqa: D205
|
||
"""
|
||
This test verifies that `Pipeline.run()` executes the components in the correct order when
|
||
there's a greedy Component with variadic input right before a Component with at least one default input.
|
||
|
||
We use the `spying_tracer` fixture to simplify the code to verify the order of execution.
|
||
This creates some coupling between this test and how we trace the Pipeline execution.
|
||
A worthy tradeoff in my opinion, we will notice right away if we change either the run logic or
|
||
the tracing logic.
|
||
"""
|
||
document_store = InMemoryDocumentStore()
|
||
document_store.write_documents([Document(content="This is a simple document")])
|
||
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
template = "Given this documents: {{ documents|join(', ', attribute='content') }} Answer this question: {{ query }}"
|
||
pipeline.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))
|
||
pipeline.add_component("prompt_builder", PromptBuilder(template=template))
|
||
pipeline.add_component("branch_joiner", BranchJoiner(list[Document]))
|
||
|
||
pipeline.connect("retriever", "branch_joiner")
|
||
pipeline.connect("branch_joiner", "prompt_builder.documents")
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"query": "This is my question"},
|
||
expected_outputs={
|
||
"prompt_builder": {
|
||
"prompt": "Given this documents: This is a simple document Answer this question: This is my "
|
||
"question"
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("branch_joiner", 1): {
|
||
"value": [
|
||
[
|
||
Document(
|
||
id="328f0cbb6722c5cfa290aa2b78bcda8dc5afa09f0e2c23092afc502ba89c85e7",
|
||
content="This is a simple document",
|
||
score=0.7192051811294521,
|
||
)
|
||
]
|
||
]
|
||
},
|
||
("prompt_builder", 1): {
|
||
"documents": [
|
||
Document(
|
||
id="328f0cbb6722c5cfa290aa2b78bcda8dc5afa09f0e2c23092afc502ba89c85e7",
|
||
content="This is a simple document",
|
||
score=0.7192051811294521,
|
||
)
|
||
],
|
||
"query": "This is my question",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("retriever", 1): {
|
||
"filters": None,
|
||
"query": "This is my question",
|
||
"scale_score": None,
|
||
"top_k": None,
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has a component that doesn't return a dictionary", target_fixture="pipeline_data")
|
||
def pipeline_that_has_a_component_that_doesnt_return_a_dictionary():
|
||
BrokenComponent = component_class(
|
||
"BrokenComponent",
|
||
input_types={"a": int},
|
||
output_types={"b": int},
|
||
output=1, # type:ignore
|
||
)
|
||
pipe = Pipeline(max_runs_per_component=10)
|
||
pipe.add_component("comp", BrokenComponent())
|
||
return pipe, [PipelineRunData({"comp": {"a": 1}})]
|
||
|
||
|
||
@given("a pipeline that has a component with only default inputs", target_fixture="pipeline_data")
|
||
def pipeline_that_has_a_component_with_only_default_inputs():
|
||
FakeGenerator = component_class(
|
||
"FakeGenerator", input_types={"prompt": str}, output_types={"replies": list[str]}, output={"replies": ["Paris"]}
|
||
)
|
||
docs = [Document(content="Rome is the capital of Italy"), Document(content="Paris is the capital of France")]
|
||
doc_store = InMemoryDocumentStore()
|
||
doc_store.write_documents(docs)
|
||
template = (
|
||
"Given the following information, answer the question.\n"
|
||
"Context:\n"
|
||
"{% for document in documents %}"
|
||
" {{ document.content }}\n"
|
||
"{% endfor %}"
|
||
"Question: {{ query }}"
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
|
||
pipe.add_component("retriever", InMemoryBM25Retriever(document_store=doc_store))
|
||
pipe.add_component("prompt_builder", PromptBuilder(template=template))
|
||
pipe.add_component("generator", FakeGenerator())
|
||
pipe.add_component("answer_builder", AnswerBuilder())
|
||
|
||
pipe.connect("retriever", "prompt_builder.documents")
|
||
pipe.connect("prompt_builder.prompt", "generator.prompt")
|
||
pipe.connect("generator.replies", "answer_builder.replies")
|
||
pipe.connect("retriever.documents", "answer_builder.documents")
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"query": "What is the capital of France?"},
|
||
expected_outputs={
|
||
"answer_builder": {
|
||
"answers": [
|
||
GeneratedAnswer(
|
||
data="Paris",
|
||
query="What is the capital of France?",
|
||
documents=[
|
||
Document(
|
||
id="413dccdf51a54cca75b7ed2eddac04e6e58560bd2f0caf4106a3efc023fe3651",
|
||
content="Paris is the capital of France",
|
||
score=1.7780417596697045,
|
||
meta={"source_index": 1},
|
||
),
|
||
Document(
|
||
id="a4a874fc2ef75015da7924d709fbdd2430e46a8e94add6e0f26cd32c1c03435d",
|
||
content="Rome is the capital of Italy",
|
||
score=1.3448247718197388,
|
||
meta={"source_index": 2},
|
||
),
|
||
],
|
||
meta={"all_messages": ["Paris"]},
|
||
)
|
||
]
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("answer_builder", 1): {
|
||
"documents": [
|
||
Document(
|
||
id="413dccdf51a54cca75b7ed2eddac04e6e58560bd2f0caf4106a3efc023fe3651",
|
||
content="Paris is the capital of France",
|
||
score=1.7780417596697045,
|
||
),
|
||
Document(
|
||
id="a4a874fc2ef75015da7924d709fbdd2430e46a8e94add6e0f26cd32c1c03435d",
|
||
content="Rome is the capital of Italy",
|
||
score=1.3448247718197388,
|
||
),
|
||
],
|
||
"meta": None,
|
||
"pattern": None,
|
||
"query": "What is the capital of France?",
|
||
"expand_reference_ranges": None,
|
||
"reference_pattern": None,
|
||
"replies": ["Paris"],
|
||
},
|
||
("generator", 1): {
|
||
"prompt": "Given the following information, answer the question.\n"
|
||
"Context:\n"
|
||
" Paris is the capital of France\n"
|
||
" Rome is the capital of Italy\n"
|
||
"Question: What is the capital of France?"
|
||
},
|
||
("prompt_builder", 1): {
|
||
"documents": [
|
||
Document(
|
||
id="413dccdf51a54cca75b7ed2eddac04e6e58560bd2f0caf4106a3efc023fe3651",
|
||
content="Paris is the capital of France",
|
||
score=1.7780417596697045,
|
||
),
|
||
Document(
|
||
id="a4a874fc2ef75015da7924d709fbdd2430e46a8e94add6e0f26cd32c1c03435d",
|
||
content="Rome is the capital of Italy",
|
||
score=1.3448247718197388,
|
||
),
|
||
],
|
||
"query": "What is the capital of France?",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("retriever", 1): {
|
||
"filters": None,
|
||
"query": "What is the capital of France?",
|
||
"scale_score": None,
|
||
"top_k": None,
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has a component with only default inputs as first to run and receives inputs from a loop",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_that_has_a_component_with_only_default_inputs_as_first_to_run_and_receives_inputs_from_a_loop():
|
||
"""
|
||
This tests verifies that a Pipeline doesn't get stuck running in a loop if it has all the following characteristics:
|
||
|
||
- The first Component has all defaults for its inputs
|
||
- The first Component receives one input from the user
|
||
- The first Component receives one input from a loop in the Pipeline
|
||
- The second Component has at least one default input
|
||
"""
|
||
|
||
def fake_generator_run(
|
||
self: Any, generation_kwargs: dict[str, Any] | None = None, **kwargs: dict[str, Any]
|
||
) -> dict[str, list[str]]:
|
||
# Simple hack to simulate a model returning a different reply after the
|
||
# the first time it's called
|
||
if getattr(fake_generator_run, "called", False):
|
||
return {"replies": ["Rome"]}
|
||
fake_generator_run.called = True # type: ignore[attr-defined]
|
||
return {"replies": ["Paris"]}
|
||
|
||
FakeGenerator = component_class(
|
||
"FakeGenerator",
|
||
input_types={"prompt": str},
|
||
output_types={"replies": list[str]},
|
||
extra_fields={"run": fake_generator_run},
|
||
)
|
||
template = (
|
||
"Answer the following question.\n"
|
||
"{% if previous_replies %}\n"
|
||
"Previously you replied incorrectly this:\n"
|
||
"{% for reply in previous_replies %}\n"
|
||
" - {{ reply }}\n"
|
||
"{% endfor %}\n"
|
||
"{% endif %}\n"
|
||
"Question: {{ query }}"
|
||
)
|
||
router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{ replies == ['Rome'] }}",
|
||
"output": "{{ replies }}",
|
||
"output_name": "correct_replies",
|
||
"output_type": list[int],
|
||
},
|
||
{
|
||
"condition": "{{ replies == ['Paris'] }}",
|
||
"output": "{{ replies }}",
|
||
"output_name": "incorrect_replies",
|
||
"output_type": list[int],
|
||
},
|
||
]
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=2)
|
||
|
||
pipe.add_component("prompt_builder", PromptBuilder(template=template, required_variables=["query"]))
|
||
pipe.add_component("generator", FakeGenerator())
|
||
pipe.add_component("router", router)
|
||
|
||
pipe.connect("prompt_builder.prompt", "generator.prompt")
|
||
pipe.connect("generator.replies", "router.replies")
|
||
pipe.connect("router.incorrect_replies", "prompt_builder.previous_replies")
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"prompt_builder": {"query": "What is the capital of Italy?"}},
|
||
expected_outputs={"router": {"correct_replies": ["Rome"]}},
|
||
expected_component_calls={
|
||
("generator", 1): {
|
||
"generation_kwargs": None,
|
||
"prompt": "Answer the following question.\n\nQuestion: What is the capital of Italy?",
|
||
},
|
||
("generator", 2): {
|
||
"generation_kwargs": None,
|
||
"prompt": "Answer the following question.\n\n"
|
||
"Previously you replied incorrectly this:\n\n"
|
||
" - Paris\n\n\n"
|
||
"Question: What is the capital of Italy?",
|
||
},
|
||
("prompt_builder", 1): {
|
||
"previous_replies": "",
|
||
"query": "What is the capital of Italy?",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("prompt_builder", 2): {
|
||
"previous_replies": ["Paris"],
|
||
"query": "What is the capital of Italy?",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("router", 1): {"replies": ["Paris"]},
|
||
("router", 2): {"replies": ["Rome"]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has multiple branches that merge into a component with a single variadic input",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_that_has_multiple_branches_that_merge_into_a_component_with_a_single_variadic_input():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("add_one", AddFixedValue())
|
||
pipeline.add_component("parity", Remainder(divisor=2))
|
||
pipeline.add_component("add_ten", AddFixedValue(add=10))
|
||
pipeline.add_component("double", Double())
|
||
pipeline.add_component("add_four", AddFixedValue(add=4))
|
||
pipeline.add_component("add_one_again", AddFixedValue())
|
||
pipeline.add_component("sum", Sum())
|
||
|
||
pipeline.connect("add_one.result", "parity.value")
|
||
pipeline.connect("parity.remainder_is_0", "add_ten.value")
|
||
pipeline.connect("parity.remainder_is_1", "double.value")
|
||
pipeline.connect("add_one.result", "sum.values")
|
||
pipeline.connect("add_ten.result", "sum.values")
|
||
pipeline.connect("double.value", "sum.values")
|
||
pipeline.connect("parity.remainder_is_1", "add_four.value")
|
||
pipeline.connect("add_four.result", "add_one_again.value")
|
||
pipeline.connect("add_one_again.result", "sum.values")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"add_one": {"value": 1}},
|
||
expected_outputs={"sum": {"total": 14}},
|
||
expected_component_calls={
|
||
("add_one", 1): {"add": None, "value": 1},
|
||
("add_ten", 1): {"add": None, "value": 2},
|
||
("parity", 1): {"value": 2},
|
||
("sum", 1): {"values": [2, 12]},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"add_one": {"value": 2}},
|
||
expected_outputs={"sum": {"total": 17}},
|
||
expected_component_calls={
|
||
("add_four", 1): {"add": None, "value": 3},
|
||
("add_one", 1): {"add": None, "value": 2},
|
||
("add_one_again", 1): {"add": None, "value": 7},
|
||
("double", 1): {"value": 3},
|
||
("parity", 1): {"value": 3},
|
||
("sum", 1): {"values": [3, 6, 8]},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has multiple branches of different lengths that merge into a component with a single variadic input",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_that_has_multiple_branches_of_different_lengths_that_merge_into_a_component_with_a_single_variadic_input():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("first_addition", AddFixedValue(add=2))
|
||
pipeline.add_component("second_addition", AddFixedValue(add=2))
|
||
pipeline.add_component("third_addition", AddFixedValue(add=2))
|
||
pipeline.add_component("sum", Sum())
|
||
pipeline.add_component("fourth_addition", AddFixedValue(add=1))
|
||
|
||
pipeline.connect("first_addition.result", "second_addition.value")
|
||
pipeline.connect("first_addition.result", "sum.values")
|
||
pipeline.connect("second_addition.result", "sum.values")
|
||
pipeline.connect("third_addition.result", "sum.values")
|
||
pipeline.connect("sum.total", "fourth_addition.value")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"first_addition": {"value": 1}, "third_addition": {"value": 1}},
|
||
expected_outputs={"fourth_addition": {"result": 12}},
|
||
expected_component_calls={
|
||
("first_addition", 1): {"add": None, "value": 1},
|
||
("fourth_addition", 1): {"add": None, "value": 11},
|
||
("second_addition", 1): {"add": None, "value": 3},
|
||
("sum", 1): {"values": AnyOrder([3, 3, 5])},
|
||
("third_addition", 1): {"add": None, "value": 1},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that is linear and returns intermediate outputs", target_fixture="pipeline_data")
|
||
def pipeline_that_is_linear_and_returns_intermediate_outputs():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("first_addition", AddFixedValue(add=2))
|
||
pipeline.add_component("second_addition", AddFixedValue())
|
||
pipeline.add_component("double", Double())
|
||
pipeline.connect("first_addition", "double")
|
||
pipeline.connect("double", "second_addition")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"first_addition": {"value": 1}},
|
||
include_outputs_from={"second_addition", "double", "first_addition"},
|
||
expected_outputs={
|
||
"double": {"value": 6},
|
||
"first_addition": {"result": 3},
|
||
"second_addition": {"result": 7},
|
||
},
|
||
expected_component_calls={
|
||
("double", 1): {"value": 3},
|
||
("first_addition", 1): {"add": None, "value": 1},
|
||
("second_addition", 1): {"add": None, "value": 6},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"first_addition": {"value": 1}},
|
||
include_outputs_from={"double"},
|
||
expected_outputs={"double": {"value": 6}, "second_addition": {"result": 7}},
|
||
expected_component_calls={
|
||
("double", 1): {"value": 3},
|
||
("first_addition", 1): {"add": None, "value": 1},
|
||
("second_addition", 1): {"add": None, "value": 6},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has a loop and returns intermediate outputs from it", target_fixture="pipeline_data")
|
||
def pipeline_that_has_a_loop_and_returns_intermediate_outputs_from_it():
|
||
pipeline = Pipeline(max_runs_per_component=10)
|
||
pipeline.add_component("add_one", AddFixedValue(add=1))
|
||
pipeline.add_component("branch_joiner", BranchJoiner(type_=int))
|
||
pipeline.add_component("below_10", Threshold(threshold=10))
|
||
pipeline.add_component("below_5", Threshold(threshold=5))
|
||
pipeline.add_component("add_three", AddFixedValue(add=3))
|
||
pipeline.add_component("accumulator", Accumulate())
|
||
pipeline.add_component("add_two", AddFixedValue(add=2))
|
||
|
||
pipeline.connect("add_one.result", "branch_joiner")
|
||
pipeline.connect("branch_joiner.value", "below_10.value")
|
||
pipeline.connect("below_10.below", "accumulator.value")
|
||
pipeline.connect("accumulator.value", "below_5.value")
|
||
pipeline.connect("below_5.above", "add_three.value")
|
||
pipeline.connect("below_5.below", "branch_joiner")
|
||
pipeline.connect("add_three.result", "branch_joiner")
|
||
pipeline.connect("below_10.above", "add_two.value")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"add_one": {"value": 3}},
|
||
include_outputs_from={
|
||
"add_two",
|
||
"add_one",
|
||
"branch_joiner",
|
||
"below_10",
|
||
"accumulator",
|
||
"below_5",
|
||
"add_three",
|
||
},
|
||
expected_outputs={
|
||
"add_two": {"result": 13},
|
||
"add_one": {"result": 4},
|
||
"branch_joiner": {"value": 11},
|
||
"below_10": {"above": 11},
|
||
"accumulator": {"value": 8},
|
||
"below_5": {"above": 8},
|
||
"add_three": {"result": 11},
|
||
},
|
||
expected_component_calls={
|
||
("accumulator", 1): {"value": 4},
|
||
("accumulator", 2): {"value": 4},
|
||
("add_one", 1): {"add": None, "value": 3},
|
||
("add_three", 1): {"add": None, "value": 8},
|
||
("add_two", 1): {"add": None, "value": 11},
|
||
("below_10", 1): {"threshold": None, "value": 4},
|
||
("below_10", 2): {"threshold": None, "value": 4},
|
||
("below_10", 3): {"threshold": None, "value": 11},
|
||
("below_5", 1): {"threshold": None, "value": 4},
|
||
("below_5", 2): {"threshold": None, "value": 8},
|
||
("branch_joiner", 1): {"value": [4]},
|
||
("branch_joiner", 2): {"value": [4]},
|
||
("branch_joiner", 3): {"value": [11]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that is linear and returns intermediate outputs from multiple sockets", target_fixture="pipeline_data"
|
||
)
|
||
def pipeline_that_is_linear_and_returns_intermediate_outputs_from_multiple_sockets():
|
||
@component
|
||
class DoubleWithOriginal:
|
||
"""
|
||
Doubles the input value and returns the original value as well.
|
||
"""
|
||
|
||
@component.output_types(value=int, original=int)
|
||
def run(self, value: int) -> dict[str, int]:
|
||
return {"value": value * 2, "original": value}
|
||
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("first_addition", AddFixedValue(add=2))
|
||
pipeline.add_component("second_addition", AddFixedValue())
|
||
pipeline.add_component("double", DoubleWithOriginal())
|
||
pipeline.connect("first_addition", "double")
|
||
pipeline.connect("double.value", "second_addition")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"first_addition": {"value": 1}},
|
||
include_outputs_from={"second_addition", "double", "first_addition"},
|
||
expected_outputs={
|
||
"double": {"original": 3, "value": 6},
|
||
"first_addition": {"result": 3},
|
||
"second_addition": {"result": 7},
|
||
},
|
||
expected_component_calls={
|
||
("double", 1): {"value": 3},
|
||
("first_addition", 1): {"add": None, "value": 1},
|
||
("second_addition", 1): {"add": None, "value": 6},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"first_addition": {"value": 1}},
|
||
include_outputs_from={"double"},
|
||
expected_outputs={"double": {"original": 3, "value": 6}, "second_addition": {"result": 7}},
|
||
expected_component_calls={
|
||
("double", 1): {"value": 3},
|
||
("first_addition", 1): {"add": None, "value": 1},
|
||
("second_addition", 1): {"add": None, "value": 6},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has a component with default inputs that doesn't receive anything from its sender",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_that_has_a_component_with_default_inputs_that_doesnt_receive_anything_from_its_sender():
|
||
router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{'reisen' in sentence}}",
|
||
"output": "German",
|
||
"output_name": "language_1",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{'viajar' in sentence}}",
|
||
"output": "Spanish",
|
||
"output_name": "language_2",
|
||
"output_type": str,
|
||
},
|
||
]
|
||
)
|
||
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("router", router)
|
||
pipeline.add_component("pb", PromptBuilder(template="Ok, I know, that's {{language}}"))
|
||
pipeline.connect("router.language_2", "pb.language")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"router": {"sentence": "Wir mussen reisen"}},
|
||
expected_outputs={"router": {"language_1": "German"}},
|
||
expected_component_calls={("router", 1): {"sentence": "Wir mussen reisen"}},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"router": {"sentence": "Yo tengo que viajar"}},
|
||
expected_outputs={"pb": {"prompt": "Ok, I know, that's Spanish"}},
|
||
expected_component_calls={
|
||
("pb", 1): {"language": "Spanish", "template": None, "template_variables": None},
|
||
("router", 1): {"sentence": "Yo tengo que viajar"},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has a component with default inputs that doesn't receive anything from its sender but receives input from user",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_that_has_a_component_with_default_inputs_that_doesnt_receive_anything_from_its_sender_but_receives_input_from_user():
|
||
prompt = PromptBuilder(
|
||
template="""Please generate an SQL query. The query should answer the following Question: {{ question }};
|
||
If the question cannot be answered given the provided table and columns, return 'no_answer'
|
||
The query is to be answered for the table is called 'absenteeism' with the following
|
||
Columns: {{ columns }};
|
||
Answer:""",
|
||
required_variables=["question", "columns"],
|
||
)
|
||
|
||
@component
|
||
class FakeGenerator:
|
||
@component.output_types(replies=list[str])
|
||
def run(self, prompt: str) -> dict[str, list[str]]:
|
||
if "no_answer" in prompt:
|
||
return {"replies": ["There's simply no_answer to this question"]}
|
||
return {"replies": ["Some SQL query"]}
|
||
|
||
@component
|
||
class FakeSQLQuerier:
|
||
@component.output_types(results=str)
|
||
def run(self, query: str) -> dict[str, str]:
|
||
return {"results": "This is the query result", "query": query}
|
||
|
||
router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{'no_answer' not in replies[0]}}",
|
||
"output": "{{replies[0]}}",
|
||
"output_name": "sql",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{'no_answer' in replies[0]}}",
|
||
"output": "{{question}}",
|
||
"output_name": "go_to_fallback",
|
||
"output_type": str,
|
||
},
|
||
]
|
||
)
|
||
|
||
fallback_prompt = PromptBuilder(
|
||
template="""User entered a query that cannot be answered with the given table.
|
||
The query was: {{ question }} and the table had columns: {{ columns }}.
|
||
Let the user know why the question cannot be answered""",
|
||
required_variables=["question"],
|
||
)
|
||
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("prompt", prompt)
|
||
pipeline.add_component("llm", FakeGenerator())
|
||
pipeline.add_component("router", router)
|
||
pipeline.add_component("fallback_prompt", fallback_prompt)
|
||
pipeline.add_component("fallback_llm", FakeGenerator())
|
||
pipeline.add_component("sql_querier", FakeSQLQuerier())
|
||
|
||
pipeline.connect("prompt", "llm")
|
||
pipeline.connect("llm.replies", "router.replies")
|
||
pipeline.connect("router.sql", "sql_querier.query")
|
||
pipeline.connect("router.go_to_fallback", "fallback_prompt.question")
|
||
pipeline.connect("fallback_prompt", "fallback_llm")
|
||
|
||
columns = "Age, Absenteeism_time_in_hours, Days, Disciplinary_failure"
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={
|
||
"prompt": {"question": "This is a question with no_answer", "columns": columns},
|
||
"router": {"question": "This is a question with no_answer"},
|
||
},
|
||
expected_outputs={"fallback_llm": {"replies": ["There's simply no_answer to this question"]}},
|
||
expected_component_calls={
|
||
("fallback_llm", 1): {
|
||
"prompt": "User entered a query that cannot be answered with the given table.\n"
|
||
"The query was: This is a question with no_answer and the table had columns: .\n"
|
||
"Let the user know why the question cannot be answered"
|
||
},
|
||
("fallback_prompt", 1): {
|
||
"columns": "",
|
||
"question": "This is a question with no_answer",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("llm", 1): {
|
||
"prompt": "Please generate an SQL query. The query should answer "
|
||
"the following Question: This is a question with no_answer;\n"
|
||
"If the question cannot be answered given the provided table and columns, return 'no_answer'\n"
|
||
"The query is to be answered for the table is called 'absenteeism' with the following\n"
|
||
"Columns: Age, Absenteeism_time_in_hours, Days, Disciplinary_failure;\n"
|
||
"Answer:"
|
||
},
|
||
("prompt", 1): {
|
||
"columns": "Age, Absenteeism_time_in_hours, Days, Disciplinary_failure",
|
||
"question": "This is a question with no_answer",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("router", 1): {
|
||
"question": "This is a question with no_answer",
|
||
"replies": ["There's simply no_answer to this question"],
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has a loop and a component with default inputs that doesn't receive anything from its sender but receives input from user",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_that_has_a_loop_and_a_component_with_default_inputs_that_doesnt_receive_anything_from_its_sender_but_receives_input_from_user():
|
||
static_content = """
|
||
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.
|
||
"""
|
||
template = (
|
||
static_content
|
||
+ """Here is the comment:
|
||
{{ comment }}
|
||
. 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 -%}"""
|
||
)
|
||
prompt_builder = PromptBuilder(template=template, required_variables=None)
|
||
|
||
@component
|
||
class FakeOutputValidator:
|
||
@component.output_types(valid_replies=list[str], invalid_replies=list[str] | None, error_message=str | None)
|
||
def run(self, replies: list[str]) -> dict[str, Any]:
|
||
if not getattr(self, "called", False):
|
||
self.called = True
|
||
return {"invalid_replies": ["This is an invalid reply"], "error_message": "this is an error message"}
|
||
return {"valid_replies": replies}
|
||
|
||
@component
|
||
class FakeGenerator:
|
||
@component.output_types(replies=list[str])
|
||
def run(self, prompt: str) -> dict[str, list[str]]:
|
||
return {"replies": ["This is a valid reply"]}
|
||
|
||
llm = FakeGenerator()
|
||
validator = FakeOutputValidator()
|
||
|
||
pipeline = Pipeline(max_runs_per_component=2)
|
||
pipeline.add_component("prompt_builder", prompt_builder)
|
||
|
||
pipeline.add_component("llm", llm)
|
||
pipeline.add_component("output_validator", validator)
|
||
|
||
pipeline.connect("prompt_builder.prompt", "llm.prompt")
|
||
pipeline.connect("llm.replies", "output_validator.replies")
|
||
pipeline.connect("output_validator.invalid_replies", "prompt_builder.invalid_replies")
|
||
|
||
pipeline.connect("output_validator.error_message", "prompt_builder.error_message")
|
||
|
||
comment = "I loved the quality of the meal but the courier was rude"
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"prompt_builder": {"template_variables": {"comment": comment}}},
|
||
expected_outputs={"output_validator": {"valid_replies": ["This is a valid reply"]}},
|
||
expected_component_calls={
|
||
("llm", 1): {
|
||
"prompt": static_content
|
||
+ """Here is the comment:
|
||
I loved the quality of the meal but the courier was rude
|
||
. Just return the category names in the list. If there aren't any, return an empty list."""
|
||
},
|
||
("llm", 2): {
|
||
"prompt": static_content
|
||
+ """Here is the comment:
|
||
I loved the quality of the meal but the courier was rude
|
||
. Just return the category names in the list. If there aren't any, return an empty list.You already created the following output in a previous attempt: ['This is an invalid reply']
|
||
However, this doesn't comply with the format requirements from above and triggered this Python exception: this is an error message
|
||
Correct the output and try again. Just return the corrected output without any extra explanations."""
|
||
},
|
||
("output_validator", 1): {"replies": ["This is a valid reply"]},
|
||
("output_validator", 2): {"replies": ["This is a valid reply"]},
|
||
("prompt_builder", 1): {
|
||
"comment": "",
|
||
"error_message": "",
|
||
"invalid_replies": "",
|
||
"template": None,
|
||
"template_variables": {"comment": "I loved the quality of the meal but the courier was rude"},
|
||
},
|
||
("prompt_builder", 2): {
|
||
"comment": "",
|
||
"error_message": "this is an error message",
|
||
"invalid_replies": ["This is an invalid reply"],
|
||
"template": None,
|
||
"template_variables": {"comment": "I loved the quality of the meal but the courier was rude"},
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has multiple components with only default inputs and are added in a different order from the order of execution",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_that_has_multiple_components_with_only_default_inputs_and_are_added_in_a_different_order_from_the_order_of_execution():
|
||
prompt_builder1 = PromptBuilder(
|
||
template="""
|
||
You are a spellchecking system. Check the given query and fill in the corrected query.
|
||
|
||
Question: {{question}}
|
||
Corrected question:
|
||
"""
|
||
)
|
||
prompt_builder2 = PromptBuilder(
|
||
template="""According to these documents:
|
||
|
||
{% for doc in documents -%}
|
||
{{ doc.content }}
|
||
{%- endfor %}
|
||
|
||
Answer the given question: {{question}}
|
||
Answer:
|
||
"""
|
||
)
|
||
prompt_builder3 = PromptBuilder(
|
||
template="""
|
||
{% for ans in replies %}
|
||
{{ ans }}
|
||
{% endfor %}
|
||
"""
|
||
)
|
||
|
||
doc = Document(content="This is a document")
|
||
|
||
@component
|
||
class FakeRetriever:
|
||
@component.output_types(documents=list[Document])
|
||
def run(
|
||
self, query: str, filters: dict[str, Any] | None = None, top_k: int | None = None
|
||
) -> dict[str, list[Document]]:
|
||
return {"documents": [doc]}
|
||
|
||
@component
|
||
class FakeRanker:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, query: str, documents: list[Document], top_k: int | None = None) -> dict[str, list[Document]]:
|
||
return {"documents": documents}
|
||
|
||
@component
|
||
class FakeGenerator:
|
||
@component.output_types(replies=list[str], meta=dict[str, Any])
|
||
def run(self, prompt: str, generation_kwargs: dict[str, Any] | None = None) -> dict[str, Any]:
|
||
return {"replies": ["This is a reply"], "meta": {"meta_key": "meta_value"}}
|
||
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component(name="retriever", instance=FakeRetriever())
|
||
pipeline.add_component(name="ranker", instance=FakeRanker())
|
||
pipeline.add_component(name="prompt_builder2", instance=prompt_builder2)
|
||
pipeline.add_component(name="prompt_builder1", instance=prompt_builder1)
|
||
pipeline.add_component(name="prompt_builder3", instance=prompt_builder3)
|
||
pipeline.add_component(name="llm", instance=FakeGenerator())
|
||
pipeline.add_component(name="spellchecker", instance=FakeGenerator())
|
||
|
||
pipeline.connect("prompt_builder1", "spellchecker")
|
||
pipeline.connect("spellchecker.replies", "prompt_builder3")
|
||
pipeline.connect("prompt_builder3", "retriever.query")
|
||
pipeline.connect("prompt_builder3", "ranker.query")
|
||
pipeline.connect("retriever.documents", "ranker.documents")
|
||
pipeline.connect("ranker.documents", "prompt_builder2.documents")
|
||
pipeline.connect("prompt_builder3", "prompt_builder2.question")
|
||
pipeline.connect("prompt_builder2", "llm")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"prompt_builder1": {"question": "Wha i Acromegaly?"}},
|
||
expected_outputs={
|
||
"llm": {"replies": ["This is a reply"], "meta": {"meta_key": "meta_value"}},
|
||
"spellchecker": {"meta": {"meta_key": "meta_value"}},
|
||
},
|
||
expected_component_calls={
|
||
("llm", 1): {
|
||
"generation_kwargs": None,
|
||
"prompt": "According to these documents:\n\nThis is a document\n\nAnswer the given question: \n \n This is a reply\n \n \nAnswer:",
|
||
},
|
||
("prompt_builder1", 1): {
|
||
"question": "Wha i Acromegaly?",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("prompt_builder2", 1): {
|
||
"documents": [doc],
|
||
"question": "\n \n This is a reply\n \n ",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("prompt_builder3", 1): {
|
||
"replies": ["This is a reply"],
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("ranker", 1): {
|
||
"documents": [doc],
|
||
"query": "\n \n This is a reply\n \n ",
|
||
"top_k": None,
|
||
},
|
||
("retriever", 1): {
|
||
"filters": None,
|
||
"query": "\n \n This is a reply\n \n ",
|
||
"top_k": None,
|
||
},
|
||
("spellchecker", 1): {
|
||
"generation_kwargs": None,
|
||
"prompt": "\n"
|
||
" You are a spellchecking system. Check "
|
||
"the given query and fill in the corrected "
|
||
"query.\n"
|
||
"\n"
|
||
" Question: Wha i Acromegaly?\n"
|
||
" Corrected question:\n"
|
||
" ",
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that is linear with conditional branching and multiple joins", target_fixture="pipeline_data")
|
||
def that_is_linear_with_conditional_branching_and_multiple_joins():
|
||
pipeline = Pipeline()
|
||
|
||
doc1 = Document(content="This is a document")
|
||
doc2 = Document(content="This is another document")
|
||
|
||
@component
|
||
class FakeRouter:
|
||
@component.output_types(LEGIT=str, INJECTION=str)
|
||
def run(self, query: str) -> dict[str, str]:
|
||
if "injection" in query:
|
||
return {"INJECTION": query}
|
||
return {"LEGIT": query}
|
||
|
||
@component
|
||
class FakeEmbedder:
|
||
@component.output_types(embeddings=list[float])
|
||
def run(self, text: str) -> dict[str, list[float]]:
|
||
return {"embeddings": [1.0, 2.0, 3.0]}
|
||
|
||
@component
|
||
class FakeRanker:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, query: str, documents: list[Document]) -> dict[str, list[Document]]:
|
||
return {"documents": documents}
|
||
|
||
@component
|
||
class FakeRetriever:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, query: str) -> dict[str, list[Document]]:
|
||
if "injection" in query:
|
||
return {"documents": []}
|
||
return {"documents": [doc1]}
|
||
|
||
@component
|
||
class FakeEmbeddingRetriever:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, query_embedding: list[float]) -> dict[str, list[Document]]:
|
||
return {"documents": [doc2]}
|
||
|
||
pipeline.add_component(name="router", instance=FakeRouter())
|
||
pipeline.add_component(name="text_embedder", instance=FakeEmbedder())
|
||
pipeline.add_component(name="retriever", instance=FakeEmbeddingRetriever())
|
||
pipeline.add_component(name="emptyretriever", instance=FakeRetriever())
|
||
pipeline.add_component(name="joinerfinal", instance=DocumentJoiner())
|
||
pipeline.add_component(name="joinerhybrid", instance=DocumentJoiner())
|
||
pipeline.add_component(name="ranker", instance=FakeRanker())
|
||
pipeline.add_component(name="bm25retriever", instance=FakeRetriever())
|
||
|
||
pipeline.connect("router.INJECTION", "emptyretriever.query")
|
||
pipeline.connect("router.LEGIT", "text_embedder.text")
|
||
pipeline.connect("text_embedder", "retriever.query_embedding")
|
||
pipeline.connect("router.LEGIT", "ranker.query")
|
||
pipeline.connect("router.LEGIT", "bm25retriever.query")
|
||
pipeline.connect("bm25retriever", "joinerhybrid.documents")
|
||
pipeline.connect("retriever", "joinerhybrid.documents")
|
||
pipeline.connect("joinerhybrid.documents", "ranker.documents")
|
||
pipeline.connect("ranker", "joinerfinal.documents")
|
||
pipeline.connect("emptyretriever", "joinerfinal.documents")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"router": {"query": "I'm a legit question"}},
|
||
expected_outputs={"joinerfinal": {"documents": [doc1, doc2]}},
|
||
expected_component_calls={
|
||
("router", 1): {"query": "I'm a legit question"},
|
||
("text_embedder", 1): {"text": "I'm a legit question"},
|
||
("bm25retriever", 1): {"query": "I'm a legit question"},
|
||
("retriever", 1): {"query_embedding": [1.0, 2.0, 3.0]},
|
||
("joinerhybrid", 1): {"documents": [[doc1], [doc2]], "top_k": None},
|
||
("ranker", 1): {"query": "I'm a legit question", "documents": [doc1, doc2]},
|
||
("joinerfinal", 1): {"documents": [[doc1, doc2]], "top_k": None},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"router": {"query": "I'm a nasty prompt injection"}},
|
||
expected_outputs={"joinerfinal": {"documents": []}},
|
||
expected_component_calls={
|
||
("router", 1): {"query": "I'm a nasty prompt injection"},
|
||
("emptyretriever", 1): {"query": "I'm a nasty prompt injection"},
|
||
("joinerfinal", 1): {"documents": [[]], "top_k": None},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that is a simple agent", target_fixture="pipeline_data")
|
||
def that_is_a_simple_agent():
|
||
search_message_template = """
|
||
Given these web search results:
|
||
|
||
{% for doc in documents %}
|
||
{{ doc.content }}
|
||
{% endfor %}
|
||
|
||
Be as brief as possible, max one sentence.
|
||
Answer the question: {{search_query}}
|
||
"""
|
||
|
||
react_message_template = """
|
||
Solve a question answering task with interleaving Thought, Action, Observation steps.
|
||
|
||
Thought reasons about the current situation
|
||
|
||
Action can be:
|
||
google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use
|
||
finish - Returns the final answer (given in square brackets) and finishes the task
|
||
|
||
Observation summarizes the Action outcome and helps in formulating the next
|
||
Thought in Thought, Action, Observation interleaving triplet of steps.
|
||
|
||
After each Observation, provide the next Thought and next Action.
|
||
Don't execute multiple steps even though you know the answer.
|
||
Only generate Thought and Action, never Observation, you'll get Observation from Action.
|
||
Follow the pattern in the example below.
|
||
|
||
Example:
|
||
###########################
|
||
Question: Which magazine was started first Arthur’s Magazine or First for Women?
|
||
Thought: I need to search Arthur’s Magazine and First for Women, and find which was started
|
||
first.
|
||
Action: google_search[When was 'Arthur’s Magazine' started?]
|
||
Observation: Arthur’s Magazine was an American literary periodical ˘
|
||
published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by
|
||
Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846
|
||
it was merged into Godey’s Lady’s Book.
|
||
Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next
|
||
Action: google_search[When was 'First for Women' magazine started?]
|
||
Observation: First for Women is a woman’s magazine published by Bauer Media Group in the
|
||
USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011
|
||
the circulation of the magazine was 1,310,696 copies.
|
||
Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for
|
||
Women), so Arthur’s Magazine was started first.
|
||
Action: finish[Arthur’s Magazine]
|
||
############################
|
||
|
||
Let's start, the question is: {{query}}
|
||
|
||
Thought:
|
||
"""
|
||
|
||
@component
|
||
class FakeThoughtActionOpenAIChatGenerator:
|
||
run_counter = 0
|
||
|
||
@component.output_types(replies=list[ChatMessage])
|
||
def run(
|
||
self, messages: list[ChatMessage], generation_kwargs: dict[str, Any] | None = None
|
||
) -> dict[str, list[ChatMessage]]:
|
||
if self.run_counter == 0:
|
||
self.run_counter += 1
|
||
return {
|
||
"replies": [
|
||
ChatMessage.from_assistant(
|
||
"thinking\n Action: google_search[What is taller, Eiffel Tower or Leaning Tower of Pisa]\n"
|
||
)
|
||
]
|
||
}
|
||
|
||
return {"replies": [ChatMessage.from_assistant("thinking\n Action: finish[Eiffel Tower]\n")]}
|
||
|
||
@component
|
||
class FakeConclusionOpenAIChatGenerator:
|
||
@component.output_types(replies=list[ChatMessage])
|
||
def run(
|
||
self, messages: list[ChatMessage], generation_kwargs: dict[str, Any] | None = None
|
||
) -> dict[str, list[ChatMessage]]:
|
||
return {"replies": [ChatMessage.from_assistant("Tower of Pisa is 55 meters tall\n")]}
|
||
|
||
@component
|
||
class FakeSerperDevWebSearch:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, query: str) -> dict[str, list[Document]]:
|
||
return {
|
||
"documents": [
|
||
Document(content="Eiffel Tower is 300 meters tall"),
|
||
Document(content="Tower of Pisa is 55 meters tall"),
|
||
]
|
||
}
|
||
|
||
# main part
|
||
pipeline = Pipeline()
|
||
pipeline.add_component("main_input", BranchJoiner(list[ChatMessage]))
|
||
pipeline.add_component("prompt_builder", ChatPromptBuilder(variables=["query"]))
|
||
pipeline.add_component("llm", FakeThoughtActionOpenAIChatGenerator())
|
||
|
||
@component
|
||
class ToolExtractor:
|
||
@component.output_types(output=list[str | None])
|
||
def run(self, messages: list[ChatMessage]) -> dict[str, list[str | None]]:
|
||
prompt = messages[-1].text
|
||
assert isinstance(prompt, str)
|
||
lines = prompt.strip().split("\n")
|
||
for line in reversed(lines):
|
||
pattern = r"Action:\s*(\w+)\[(.*?)\]"
|
||
|
||
match = re.search(pattern, line)
|
||
if match:
|
||
action_name = match.group(1)
|
||
parameter = match.group(2)
|
||
return {"output": [action_name, parameter]}
|
||
return {"output": [None, None]}
|
||
|
||
pipeline.add_component("tool_extractor", ToolExtractor())
|
||
|
||
@component
|
||
class PromptConcatenator:
|
||
def __init__(self, suffix: str = ""):
|
||
self._suffix = suffix
|
||
|
||
@component.output_types(output=list[ChatMessage])
|
||
def run(self, replies: list[ChatMessage], current_prompt: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||
curr_prompt = current_prompt[-1].text if current_prompt[-1].text else ""
|
||
reply = replies[-1].text if replies[-1].text else ""
|
||
content = curr_prompt + reply + self._suffix
|
||
return {"output": [ChatMessage.from_user(content)]}
|
||
|
||
@component
|
||
class SearchOutputAdapter:
|
||
@component.output_types(output=list[ChatMessage])
|
||
def run(self, replies: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||
content = f"Observation: {replies[-1].text}\n"
|
||
return {"output": [ChatMessage.from_assistant(content)]}
|
||
|
||
pipeline.add_component("prompt_concatenator_after_action", PromptConcatenator())
|
||
|
||
pipeline.add_component(
|
||
"router",
|
||
ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{'search' in tool_id_and_param[0]}}",
|
||
"output": "{{tool_id_and_param[1]}}",
|
||
"output_name": "search",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{'finish' in tool_id_and_param[0]}}",
|
||
"output": "{{tool_id_and_param[1]}}",
|
||
"output_name": "finish",
|
||
"output_type": str,
|
||
},
|
||
]
|
||
),
|
||
)
|
||
pipeline.add_component("router_search", FakeSerperDevWebSearch())
|
||
pipeline.add_component("search_prompt_builder", ChatPromptBuilder(variables=["documents", "search_query"]))
|
||
pipeline.add_component("search_llm", FakeConclusionOpenAIChatGenerator())
|
||
|
||
pipeline.add_component("search_output_adapter", SearchOutputAdapter())
|
||
pipeline.add_component("prompt_concatenator_after_observation", PromptConcatenator(suffix="\nThought: "))
|
||
|
||
# main
|
||
pipeline.connect("main_input", "prompt_builder.template")
|
||
pipeline.connect("prompt_builder.prompt", "llm.messages")
|
||
pipeline.connect("llm.replies", "prompt_concatenator_after_action.replies")
|
||
|
||
# tools
|
||
pipeline.connect("prompt_builder.prompt", "prompt_concatenator_after_action.current_prompt")
|
||
pipeline.connect("prompt_concatenator_after_action", "tool_extractor.messages")
|
||
|
||
pipeline.connect("tool_extractor", "router")
|
||
pipeline.connect("router.search", "router_search.query")
|
||
pipeline.connect("router_search.documents", "search_prompt_builder.documents")
|
||
pipeline.connect("router.search", "search_prompt_builder.search_query")
|
||
pipeline.connect("search_prompt_builder.prompt", "search_llm.messages")
|
||
|
||
pipeline.connect("search_llm.replies", "search_output_adapter.replies")
|
||
pipeline.connect("search_output_adapter", "prompt_concatenator_after_observation.replies")
|
||
pipeline.connect("prompt_concatenator_after_action", "prompt_concatenator_after_observation.current_prompt")
|
||
pipeline.connect("prompt_concatenator_after_observation", "main_input")
|
||
|
||
search_message = [ChatMessage.from_user(search_message_template)]
|
||
messages = [ChatMessage.from_user(react_message_template)]
|
||
question = "which tower is taller: eiffel tower or tower of pisa?"
|
||
|
||
return pipeline, [
|
||
PipelineRunData(
|
||
inputs={
|
||
"main_input": {"value": messages},
|
||
"prompt_builder": {"query": question},
|
||
"search_prompt_builder": {"template": search_message},
|
||
},
|
||
expected_outputs={"router": {"finish": "Eiffel Tower"}},
|
||
expected_component_calls={
|
||
("llm", 1): {
|
||
"generation_kwargs": None,
|
||
"messages": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: which tower is taller: eiffel tower or tower of pisa?\n\n Thought:\n "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
},
|
||
("llm", 2): {
|
||
"generation_kwargs": None,
|
||
"messages": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: which tower is taller: eiffel tower or tower of pisa?\n\n Thought:\n thinking\n Action: google_search[What is taller, Eiffel Tower or Leaning Tower of Pisa]\nObservation: Tower of Pisa is 55 meters tall\n\n\nThought: "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
},
|
||
("main_input", 1): {
|
||
"value": [
|
||
[
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: {{query}}\n\n Thought:\n "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
]
|
||
]
|
||
},
|
||
("main_input", 2): {
|
||
"value": [
|
||
[
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: which tower is taller: eiffel tower or tower of pisa?\n\n Thought:\n thinking\n Action: google_search[What is taller, Eiffel Tower or Leaning Tower of Pisa]\nObservation: Tower of Pisa is 55 meters tall\n\n\nThought: "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
]
|
||
]
|
||
},
|
||
("prompt_builder", 1): {
|
||
"query": "which tower is taller: eiffel tower or tower of pisa?",
|
||
"template": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: {{query}}\n\n Thought:\n "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
"template_variables": None,
|
||
},
|
||
("prompt_builder", 2): {
|
||
"query": "which tower is taller: eiffel tower or tower of pisa?",
|
||
"template": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: which tower is taller: eiffel tower or tower of pisa?\n\n Thought:\n thinking\n Action: google_search[What is taller, Eiffel Tower or Leaning Tower of Pisa]\nObservation: Tower of Pisa is 55 meters tall\n\n\nThought: "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
"template_variables": None,
|
||
},
|
||
("prompt_concatenator_after_action", 1): {
|
||
"current_prompt": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: which tower is taller: eiffel tower or tower of pisa?\n\n Thought:\n "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
"replies": [
|
||
ChatMessage(
|
||
_role=ChatRole.ASSISTANT,
|
||
_content=[
|
||
TextContent(
|
||
text="thinking\n Action: google_search[What is taller, Eiffel Tower or Leaning Tower of Pisa]\n"
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
},
|
||
("prompt_concatenator_after_action", 2): {
|
||
"current_prompt": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: which tower is taller: eiffel tower or tower of pisa?\n\n Thought:\n thinking\n Action: google_search[What is taller, Eiffel Tower or Leaning Tower of Pisa]\nObservation: Tower of Pisa is 55 meters tall\n\n\nThought: "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
"replies": [
|
||
ChatMessage(
|
||
_role=ChatRole.ASSISTANT,
|
||
_content=[TextContent(text="thinking\n Action: finish[Eiffel Tower]\n")],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
},
|
||
("prompt_concatenator_after_observation", 1): {
|
||
"current_prompt": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: which tower is taller: eiffel tower or tower of pisa?\n\n Thought:\n thinking\n Action: google_search[What is taller, Eiffel Tower or Leaning Tower of Pisa]\n"
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
"replies": [
|
||
ChatMessage(
|
||
_role=ChatRole.ASSISTANT,
|
||
_content=[TextContent(text="Observation: Tower of Pisa is 55 meters tall\n\n")],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
},
|
||
("router", 1): {
|
||
"tool_id_and_param": ["google_search", "What is taller, Eiffel Tower or Leaning Tower of Pisa"]
|
||
},
|
||
("router", 2): {"tool_id_and_param": ["finish", "Eiffel Tower"]},
|
||
("router_search", 1): {"query": "What is taller, Eiffel Tower or Leaning Tower of Pisa"},
|
||
("search_llm", 1): {
|
||
"generation_kwargs": None,
|
||
"messages": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Given these web search results:\n\n \n Eiffel Tower is 300 meters tall\n \n Tower of Pisa is 55 meters tall\n \n\n Be as brief as possible, max one sentence.\n Answer the question: What is taller, Eiffel Tower or Leaning Tower of Pisa\n "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
},
|
||
("search_output_adapter", 1): {
|
||
"replies": [
|
||
ChatMessage(
|
||
_role=ChatRole.ASSISTANT,
|
||
_content=[TextContent(text="Tower of Pisa is 55 meters tall\n")],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
]
|
||
},
|
||
("search_prompt_builder", 1): {
|
||
"documents": [
|
||
Document(
|
||
id="c37eb19352b261b17314cac9e1539921b5996f88c99ad0b134f12effb38ed467",
|
||
content="Eiffel Tower is 300 meters tall",
|
||
),
|
||
Document(
|
||
id="c5281056a220c32e6fa1c4ae7d3f263c0f25fd620592c5e45049a9dcb778f129",
|
||
content="Tower of Pisa is 55 meters tall",
|
||
),
|
||
],
|
||
"search_query": "What is taller, Eiffel Tower or Leaning Tower of Pisa",
|
||
"template": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Given these web search results:\n\n {% for doc in documents %}\n {{ doc.content }}\n {% endfor %}\n\n Be as brief as possible, max one sentence.\n Answer the question: {{search_query}}\n "
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
],
|
||
"template_variables": None,
|
||
},
|
||
("tool_extractor", 1): {
|
||
"messages": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: which tower is taller: eiffel tower or tower of pisa?\n\n Thought:\n thinking\n Action: google_search[What is taller, Eiffel Tower or Leaning Tower of Pisa]\n"
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
]
|
||
},
|
||
("tool_extractor", 2): {
|
||
"messages": [
|
||
ChatMessage(
|
||
_role=ChatRole.USER,
|
||
_content=[
|
||
TextContent(
|
||
text="\n Solve a question answering task with interleaving Thought, Action, Observation steps.\n\n Thought reasons about the current situation\n\n Action can be:\n google_search - Searches Google for the exact concept/entity (given in square brackets) and returns the results for you to use\n finish - Returns the final answer (given in square brackets) and finishes the task\n\n Observation summarizes the Action outcome and helps in formulating the next\n Thought in Thought, Action, Observation interleaving triplet of steps.\n\n After each Observation, provide the next Thought and next Action.\n Don't execute multiple steps even though you know the answer.\n Only generate Thought and Action, never Observation, you'll get Observation from Action.\n Follow the pattern in the example below.\n\n Example:\n ###########################\n Question: Which magazine was started first Arthur’s Magazine or First for Women?\n Thought: I need to search Arthur’s Magazine and First for Women, and find which was started\n first.\n Action: google_search[When was 'Arthur’s Magazine' started?]\n Observation: Arthur’s Magazine was an American literary periodical ˘\n published in Philadelphia and founded in 1844. Edited by Timothy Shay Arthur, it featured work by\n Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others. In May 1846\n it was merged into Godey’s Lady’s Book.\n Thought: Arthur’s Magazine was started in 1844. I need to search First for Women founding date next\n Action: google_search[When was 'First for Women' magazine started?]\n Observation: First for Women is a woman’s magazine published by Bauer Media Group in the\n USA. The magazine was started in 1989. It is based in Englewood Cliffs, New Jersey. In 2011\n the circulation of the magazine was 1,310,696 copies.\n Thought: First for Women was started in 1989. 1844 (Arthur’s Magazine) ¡ 1989 (First for\n Women), so Arthur’s Magazine was started first.\n Action: finish[Arthur’s Magazine]\n ############################\n\n Let's start, the question is: which tower is taller: eiffel tower or tower of pisa?\n\n Thought:\n thinking\n Action: google_search[What is taller, Eiffel Tower or Leaning Tower of Pisa]\nObservation: Tower of Pisa is 55 meters tall\n\n\nThought: thinking\n Action: finish[Eiffel Tower]\n"
|
||
)
|
||
],
|
||
_name=None,
|
||
_meta={},
|
||
)
|
||
]
|
||
},
|
||
},
|
||
)
|
||
]
|
||
|
||
|
||
@given("a pipeline that has a variadic component that receives partial inputs", target_fixture="pipeline_data")
|
||
def that_has_a_variadic_component_that_receives_partial_inputs():
|
||
@component
|
||
class ConditionalDocumentCreator:
|
||
def __init__(self, content: str):
|
||
self._content = content
|
||
|
||
@component.output_types(documents=list[Document], noop=None)
|
||
def run(self, create_document: bool = False) -> dict[str, list[Document] | None]:
|
||
if create_document:
|
||
return {"documents": [Document(id=self._content, content=self._content)]}
|
||
return {"noop": None}
|
||
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("first_creator", ConditionalDocumentCreator(content="First document"))
|
||
pipeline.add_component("second_creator", ConditionalDocumentCreator(content="Second document"))
|
||
pipeline.add_component("third_creator", ConditionalDocumentCreator(content="Third document"))
|
||
pipeline.add_component("documents_joiner", DocumentJoiner())
|
||
|
||
pipeline.connect("first_creator.documents", "documents_joiner.documents")
|
||
pipeline.connect("second_creator.documents", "documents_joiner.documents")
|
||
pipeline.connect("third_creator.documents", "documents_joiner.documents")
|
||
|
||
doc1 = Document(id="First document", content="First document")
|
||
doc2 = Document(id="Second document", content="Second document")
|
||
doc3 = Document(id="Third document", content="Third document")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"first_creator": {"create_document": True}, "third_creator": {"create_document": True}},
|
||
expected_outputs={"second_creator": {"noop": None}, "documents_joiner": {"documents": [doc1, doc3]}},
|
||
expected_component_calls={
|
||
("documents_joiner", 1): {"documents": [[doc1], [doc3]], "top_k": None},
|
||
("first_creator", 1): {"create_document": True},
|
||
("second_creator", 1): {"create_document": False},
|
||
("third_creator", 1): {"create_document": True},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"first_creator": {"create_document": True}, "second_creator": {"create_document": True}},
|
||
expected_outputs={"third_creator": {"noop": None}, "documents_joiner": {"documents": [doc1, doc2]}},
|
||
expected_component_calls={
|
||
("documents_joiner", 1): {"documents": [[doc1], [doc2]], "top_k": None},
|
||
("first_creator", 1): {"create_document": True},
|
||
("second_creator", 1): {"create_document": True},
|
||
("third_creator", 1): {"create_document": False},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that has a variadic component that receives partial inputs in a different order",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def that_has_a_variadic_component_that_receives_partial_inputs_different_order():
|
||
@component
|
||
class ConditionalDocumentCreator:
|
||
def __init__(self, content: str):
|
||
self._content = content
|
||
|
||
@component.output_types(documents=list[Document], noop=None)
|
||
def run(self, create_document: bool = False) -> dict[str, list[Document] | None]:
|
||
if create_document:
|
||
return {"documents": [Document(id=self._content, content=self._content)]}
|
||
return {"noop": None}
|
||
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("third_creator", ConditionalDocumentCreator(content="Third document"))
|
||
pipeline.add_component("first_creator", ConditionalDocumentCreator(content="First document"))
|
||
pipeline.add_component("second_creator", ConditionalDocumentCreator(content="Second document"))
|
||
pipeline.add_component("documents_joiner", DocumentJoiner())
|
||
|
||
pipeline.connect("first_creator.documents", "documents_joiner.documents")
|
||
pipeline.connect("second_creator.documents", "documents_joiner.documents")
|
||
pipeline.connect("third_creator.documents", "documents_joiner.documents")
|
||
|
||
doc1 = Document(id="First document", content="First document")
|
||
doc2 = Document(id="Second document", content="Second document")
|
||
doc3 = Document(id="Third document", content="Third document")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"first_creator": {"create_document": True}, "third_creator": {"create_document": True}},
|
||
expected_outputs={"second_creator": {"noop": None}, "documents_joiner": {"documents": [doc1, doc3]}},
|
||
expected_component_calls={
|
||
("documents_joiner", 1): {"documents": [[doc1], [doc3]], "top_k": None},
|
||
("first_creator", 1): {"create_document": True},
|
||
("second_creator", 1): {"create_document": False},
|
||
("third_creator", 1): {"create_document": True},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={"first_creator": {"create_document": True}, "second_creator": {"create_document": True}},
|
||
expected_outputs={"third_creator": {"noop": None}, "documents_joiner": {"documents": [doc1, doc2]}},
|
||
expected_component_calls={
|
||
("documents_joiner", 1): {"documents": [[doc1], [doc2]], "top_k": None},
|
||
("first_creator", 1): {"create_document": True},
|
||
("second_creator", 1): {"create_document": True},
|
||
("third_creator", 1): {"create_document": False},
|
||
},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has an answer joiner variadic component", target_fixture="pipeline_data")
|
||
def that_has_an_answer_joiner_variadic_component():
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("answer_builder_1", AnswerBuilder())
|
||
pipeline.add_component("answer_builder_2", AnswerBuilder())
|
||
pipeline.add_component("answer_joiner", AnswerJoiner())
|
||
|
||
pipeline.connect("answer_builder_1.answers", "answer_joiner")
|
||
pipeline.connect("answer_builder_2.answers", "answer_joiner")
|
||
|
||
query = "What's Natural Language Processing?"
|
||
reply1 = "This is a test answer"
|
||
reply2 = "This is a second test answer"
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={
|
||
"answer_builder_1": {"query": query, "replies": [reply1]},
|
||
"answer_builder_2": {"query": query, "replies": [reply2]},
|
||
},
|
||
expected_outputs={
|
||
"answer_joiner": {
|
||
"answers": [
|
||
GeneratedAnswer(data=reply1, query=query, documents=[], meta={"all_messages": [reply1]}),
|
||
GeneratedAnswer(data=reply2, query=query, documents=[], meta={"all_messages": [reply2]}),
|
||
]
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("answer_builder_1", 1): {
|
||
"documents": None,
|
||
"meta": None,
|
||
"pattern": None,
|
||
"query": query,
|
||
"expand_reference_ranges": None,
|
||
"reference_pattern": None,
|
||
"replies": [reply1],
|
||
},
|
||
("answer_builder_2", 1): {
|
||
"documents": None,
|
||
"meta": None,
|
||
"pattern": None,
|
||
"query": query,
|
||
"expand_reference_ranges": None,
|
||
"reference_pattern": None,
|
||
"replies": [reply2],
|
||
},
|
||
("answer_joiner", 1): {
|
||
"answers": [
|
||
[GeneratedAnswer(data=reply1, query=query, documents=[], meta={"all_messages": [reply1]})],
|
||
[GeneratedAnswer(data=reply2, query=query, documents=[], meta={"all_messages": [reply2]})],
|
||
],
|
||
"top_k": None,
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that is linear and a component in the middle receives optional input from other components and input from the user",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def that_is_linear_and_a_component_in_the_middle_receives_optional_input_from_other_components_and_input_from_the_user():
|
||
@component
|
||
class QueryMetadataExtractor:
|
||
@component.output_types(filters=dict[str, Any])
|
||
def run(self, prompt: str) -> dict[str, dict[str, Any]]:
|
||
metadata = json.loads(prompt)
|
||
filters = []
|
||
for key, value in metadata.items():
|
||
filters.append({"field": f"meta.{key}", "operator": "==", "value": value})
|
||
|
||
return {"filters": {"operator": "AND", "conditions": filters}}
|
||
|
||
documents = [
|
||
Document(
|
||
content="some publication about Alzheimer prevention research done over 2023 patients study",
|
||
meta={"year": 2022, "disease": "Alzheimer", "author": "Michael Butter"},
|
||
id="doc1",
|
||
),
|
||
Document(
|
||
content="some text about investigation and treatment of Alzheimer disease",
|
||
meta={"year": 2023, "disease": "Alzheimer", "author": "John Bread"},
|
||
id="doc2",
|
||
),
|
||
Document(
|
||
content="A study on the effectiveness of new therapies for Parkinson's disease",
|
||
meta={"year": 2022, "disease": "Parkinson", "author": "Alice Smith"},
|
||
id="doc3",
|
||
),
|
||
Document(
|
||
content="An overview of the latest research on the genetics of Parkinson's disease and its implications for treatment",
|
||
meta={"year": 2023, "disease": "Parkinson", "author": "David Jones"},
|
||
id="doc4",
|
||
),
|
||
]
|
||
document_store = InMemoryDocumentStore(bm25_algorithm="BM25Plus")
|
||
document_store.write_documents(documents=documents, policy=DuplicatePolicy.OVERWRITE)
|
||
|
||
pipeline = Pipeline()
|
||
pipeline.add_component(instance=PromptBuilder('{"disease": "Alzheimer", "year": 2023}'), name="builder")
|
||
pipeline.add_component(instance=QueryMetadataExtractor(), name="metadata_extractor")
|
||
pipeline.add_component(instance=InMemoryBM25Retriever(document_store=document_store), name="retriever")
|
||
pipeline.add_component(instance=DocumentJoiner(), name="document_joiner")
|
||
|
||
pipeline.connect("builder.prompt", "metadata_extractor.prompt")
|
||
pipeline.connect("metadata_extractor.filters", "retriever.filters")
|
||
pipeline.connect("retriever.documents", "document_joiner.documents")
|
||
|
||
query = "publications 2023 Alzheimer's disease"
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"retriever": {"query": query}},
|
||
expected_outputs={
|
||
"document_joiner": {
|
||
"documents": [
|
||
Document(
|
||
content="some text about investigation and treatment of Alzheimer disease",
|
||
meta={"year": 2023, "disease": "Alzheimer", "author": "John Bread"},
|
||
id="doc2",
|
||
score=4.148111588215998,
|
||
)
|
||
]
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("builder", 1): {"template": None, "template_variables": None},
|
||
("document_joiner", 1): {
|
||
"documents": [
|
||
[
|
||
Document(
|
||
id="doc2",
|
||
content="some text about investigation and treatment of Alzheimer disease",
|
||
meta={"year": 2023, "disease": "Alzheimer", "author": "John Bread"},
|
||
score=4.148111588215998,
|
||
)
|
||
]
|
||
],
|
||
"top_k": None,
|
||
},
|
||
("metadata_extractor", 1): {"prompt": '{"disease": "Alzheimer", "year": 2023}'},
|
||
("retriever", 1): {
|
||
"filters": {
|
||
"conditions": [
|
||
{"field": "meta.disease", "operator": "==", "value": "Alzheimer"},
|
||
{"field": "meta.year", "operator": "==", "value": 2023},
|
||
],
|
||
"operator": "AND",
|
||
},
|
||
"query": query,
|
||
"scale_score": None,
|
||
"top_k": None,
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has a cycle that would get it stuck", target_fixture="pipeline_data")
|
||
def that_has_a_cycle_that_would_get_it_stuck():
|
||
# NOTE: The name of this test is a bit misleading since the pipeline doesn't actually get stuck in the loop.
|
||
# It doesn't even run at all b/c it gets stopped at the prompt_builder component.
|
||
# The prompt_builder component doesn't run since it requires two inputs that are never provided in the first
|
||
# loop: invalid_replies and error_message.
|
||
# ruff: noqa: E501
|
||
template = """Here is the comment: {{ comment }}
|
||
|
||
{% 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.
|
||
{% endif %}"""
|
||
prompt_builder = PromptBuilder(
|
||
template=template, required_variables=["comment", "invalid_replies", "error_message"]
|
||
)
|
||
|
||
@component
|
||
class FakeOutputValidator:
|
||
@component.output_types(valid_replies=list[str], invalid_replies=list[str], error_message=str)
|
||
def run(self, replies: list[str]) -> dict[str, list[str] | str]:
|
||
if not getattr(self, "called", False):
|
||
self.called = True
|
||
return {"invalid_replies": ["This is an invalid reply"], "error_message": "this is an error message"}
|
||
return {"valid_replies": replies}
|
||
|
||
@component
|
||
class FakeGenerator:
|
||
@component.output_types(replies=list[str])
|
||
def run(self, prompt: str) -> dict[str, list[str]]:
|
||
return {"replies": ["This is a valid reply"]}
|
||
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("prompt_builder", prompt_builder)
|
||
|
||
pipeline.add_component("llm", FakeGenerator())
|
||
pipeline.add_component("output_validator", FakeOutputValidator())
|
||
|
||
pipeline.connect("prompt_builder.prompt", "llm.prompt")
|
||
pipeline.connect("llm.replies", "output_validator.replies")
|
||
pipeline.connect("output_validator.invalid_replies", "prompt_builder.invalid_replies")
|
||
pipeline.connect("output_validator.error_message", "prompt_builder.error_message")
|
||
|
||
comment = "I loved the quality of the meal but the courier was rude"
|
||
return (pipeline, [PipelineRunData(inputs={"prompt_builder": {"comment": comment}})])
|
||
|
||
|
||
@given("a pipeline that has a loop in the middle", target_fixture="pipeline_data")
|
||
def that_has_a_loop_in_the_middle():
|
||
@component
|
||
class FakeGenerator:
|
||
@component.output_types(replies=list[str])
|
||
def run(self, prompt: str) -> dict[str, list[str]]:
|
||
replies = []
|
||
if getattr(self, "first_run", True):
|
||
self.first_run = False
|
||
replies.append("No answer")
|
||
else:
|
||
replies.append("42")
|
||
return {"replies": replies}
|
||
|
||
@component
|
||
class PromptCleaner:
|
||
@component.output_types(clean_prompt=str)
|
||
def run(self, prompt: str) -> dict[str, str]:
|
||
return {"clean_prompt": prompt.strip()}
|
||
|
||
pipeline = Pipeline(max_runs_per_component=2)
|
||
pipeline.add_component("prompt_cleaner", PromptCleaner())
|
||
pipeline.add_component(
|
||
"prompt_builder",
|
||
PromptBuilder(template="", variables=["question", "invalid_replies"], required_variables=["question"]),
|
||
)
|
||
pipeline.add_component("llm", FakeGenerator())
|
||
pipeline.add_component(
|
||
"answer_validator",
|
||
ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{ 'No answer' in replies }}",
|
||
"output": "{{ replies }}",
|
||
"output_name": "invalid_replies",
|
||
"output_type": list[str],
|
||
},
|
||
{
|
||
"condition": "{{ 'No answer' not in replies }}",
|
||
"output": "{{ replies }}",
|
||
"output_name": "valid_replies",
|
||
"output_type": list[str],
|
||
},
|
||
]
|
||
),
|
||
)
|
||
pipeline.add_component("answer_builder", AnswerBuilder())
|
||
|
||
pipeline.connect("prompt_cleaner.clean_prompt", "prompt_builder.template")
|
||
pipeline.connect("prompt_builder.prompt", "llm.prompt")
|
||
pipeline.connect("llm.replies", "answer_validator.replies")
|
||
pipeline.connect("answer_validator.invalid_replies", "prompt_builder.invalid_replies")
|
||
pipeline.connect("answer_validator.valid_replies", "answer_builder.replies")
|
||
|
||
question = "What is the answer?"
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={
|
||
"prompt_cleaner": {"prompt": "Random template"},
|
||
"prompt_builder": {"question": question},
|
||
"answer_builder": {"query": question},
|
||
},
|
||
expected_outputs={
|
||
"answer_builder": {
|
||
"answers": [
|
||
GeneratedAnswer(data="42", query=question, documents=[], meta={"all_messages": ["42"]})
|
||
]
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("answer_builder", 1): {
|
||
"documents": None,
|
||
"meta": None,
|
||
"pattern": None,
|
||
"query": "What is the answer?",
|
||
"expand_reference_ranges": None,
|
||
"reference_pattern": None,
|
||
"replies": ["42"],
|
||
},
|
||
("answer_validator", 1): {"replies": ["No answer"]},
|
||
("answer_validator", 2): {"replies": ["42"]},
|
||
("llm", 1): {"prompt": "Random template"},
|
||
("llm", 2): {"prompt": ""},
|
||
("prompt_builder", 1): {
|
||
"invalid_replies": "",
|
||
"question": "What is the answer?",
|
||
"template": "Random template",
|
||
"template_variables": None,
|
||
},
|
||
("prompt_builder", 2): {
|
||
"invalid_replies": ["No answer"],
|
||
"question": "What is the answer?",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("prompt_cleaner", 1): {"prompt": "Random template"},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has variadic component that receives a conditional input", target_fixture="pipeline_data")
|
||
def that_has_variadic_component_that_receives_a_conditional_input():
|
||
@component
|
||
class NoOp:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
|
||
return {"documents": documents}
|
||
|
||
@component
|
||
class CommaSplitter:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
|
||
res = []
|
||
current_id = 0
|
||
for doc in documents:
|
||
if doc.content is None:
|
||
continue
|
||
for split in doc.content.split(","):
|
||
res.append(Document(content=split, id=str(current_id)))
|
||
current_id += 1
|
||
return {"documents": res}
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
|
||
pipe.add_component(
|
||
"conditional_router",
|
||
ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{ documents|length > 1 }}",
|
||
"output": "{{ documents }}",
|
||
"output_name": "long",
|
||
"output_type": list[Document],
|
||
},
|
||
{
|
||
"condition": "{{ documents|length <= 1 }}",
|
||
"output": "{{ documents }}",
|
||
"output_name": "short",
|
||
"output_type": list[Document],
|
||
},
|
||
],
|
||
unsafe=True,
|
||
),
|
||
)
|
||
pipe.add_component(
|
||
"empty_lines_cleaner", DocumentCleaner(remove_empty_lines=True, remove_extra_whitespaces=False, keep_id=True)
|
||
)
|
||
pipe.add_component("comma_splitter", CommaSplitter())
|
||
pipe.add_component("document_cleaner", DocumentCleaner(keep_id=True))
|
||
pipe.add_component("document_joiner", DocumentJoiner())
|
||
|
||
pipe.add_component("noop2", NoOp())
|
||
pipe.add_component("noop3", NoOp())
|
||
|
||
pipe.connect("noop2", "noop3")
|
||
pipe.connect("noop3", "conditional_router")
|
||
|
||
pipe.connect("conditional_router.long", "empty_lines_cleaner")
|
||
pipe.connect("empty_lines_cleaner", "document_joiner")
|
||
|
||
pipe.connect("comma_splitter", "document_cleaner")
|
||
pipe.connect("document_cleaner", "document_joiner")
|
||
pipe.connect("comma_splitter", "document_joiner")
|
||
|
||
document = Document(
|
||
id="1000", content="This document has so many, sentences. Like this one, or this one. Or even this other one."
|
||
)
|
||
expected_docs = [
|
||
Document(id="0", content="This document has so many"),
|
||
Document(id="1", content=" sentences. Like this one"),
|
||
Document(id="2", content=" or this one. Or even this other one."),
|
||
]
|
||
|
||
return pipe, [
|
||
PipelineRunData(
|
||
inputs={"noop2": {"documents": [document]}, "comma_splitter": {"documents": [document]}},
|
||
expected_outputs={
|
||
"conditional_router": {
|
||
"short": [
|
||
Document(
|
||
id="1000",
|
||
content="This document has so many, sentences. Like this one, or this one. Or even "
|
||
"this other one.",
|
||
)
|
||
]
|
||
},
|
||
"document_joiner": {"documents": expected_docs},
|
||
},
|
||
expected_component_calls={
|
||
("comma_splitter", 1): {
|
||
"documents": [
|
||
Document(
|
||
id="1000",
|
||
content="This document has so many, sentences. Like this one, or this one. Or even "
|
||
"this other one.",
|
||
)
|
||
]
|
||
},
|
||
("conditional_router", 1): {
|
||
"documents": [
|
||
Document(
|
||
id="1000",
|
||
content="This document has so many, sentences. Like this one, or this one. Or even "
|
||
"this other one.",
|
||
)
|
||
]
|
||
},
|
||
("document_cleaner", 1): {"documents": expected_docs},
|
||
("document_joiner", 1): {
|
||
"documents": [
|
||
expected_docs,
|
||
[
|
||
Document(id="0", content="This document has so many"),
|
||
Document(id="1", content="sentences. Like this one"),
|
||
Document(id="2", content="or this one. Or even this other one."),
|
||
],
|
||
],
|
||
"top_k": None,
|
||
},
|
||
("noop2", 1): {"documents": [document]},
|
||
("noop3", 1): {"documents": [document]},
|
||
},
|
||
),
|
||
PipelineRunData(
|
||
inputs={
|
||
"noop2": {"documents": [document, document]},
|
||
"comma_splitter": {"documents": [document, document]},
|
||
},
|
||
expected_outputs={
|
||
"document_joiner": {
|
||
"documents": [
|
||
*expected_docs,
|
||
Document(id="3", content="This document has so many"),
|
||
Document(id="4", content=" sentences. Like this one"),
|
||
Document(id="5", content=" or this one. Or even this other one."),
|
||
document,
|
||
]
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("comma_splitter", 1): {"documents": [document, document]},
|
||
("conditional_router", 1): {"documents": [document, document]},
|
||
("document_cleaner", 1): {
|
||
"documents": [
|
||
*expected_docs,
|
||
Document(id="3", content="This document has so many"),
|
||
Document(id="4", content=" sentences. Like this one"),
|
||
Document(id="5", content=" or this one. Or even this other one."),
|
||
]
|
||
},
|
||
("document_joiner", 1): {
|
||
"documents": [
|
||
[
|
||
*expected_docs,
|
||
Document(id="3", content="This document has so many"),
|
||
Document(id="4", content=" sentences. Like this one"),
|
||
Document(id="5", content=" or this one. Or even this other one."),
|
||
],
|
||
[
|
||
Document(id="0", content="This document has so many"),
|
||
Document(id="1", content="sentences. Like this one"),
|
||
Document(id="2", content="or this one. Or even this other one."),
|
||
Document(id="3", content="This document has so many"),
|
||
Document(id="4", content="sentences. Like this one"),
|
||
Document(id="5", content="or this one. Or even this other one."),
|
||
],
|
||
[document, document],
|
||
],
|
||
"top_k": None,
|
||
},
|
||
("empty_lines_cleaner", 1): {"documents": [document, document]},
|
||
("noop2", 1): {"documents": [document, document]},
|
||
("noop3", 1): {"documents": [document, document]},
|
||
},
|
||
),
|
||
]
|
||
|
||
|
||
class AnyOrder: # noqa: PLW1641 # Object does not implement `__hash__` method but it's ok
|
||
"""List wrapper with order-insensitive equality"""
|
||
|
||
def __init__(self, items, key=None):
|
||
self.items = items
|
||
self.key = key
|
||
|
||
def __eq__(self, other):
|
||
return isinstance(other, list) and sorted(self.items, key=self.key) == sorted(other, key=self.key)
|
||
|
||
|
||
@given("a pipeline that has a string variadic component", target_fixture="pipeline_data")
|
||
def that_has_a_string_variadic_component():
|
||
string_1 = "What's Natural Language Processing?"
|
||
string_2 = "What's is life?"
|
||
|
||
pipeline = Pipeline()
|
||
pipeline.add_component("prompt_builder_1", PromptBuilder("Builder 1: {{query}}"))
|
||
pipeline.add_component("prompt_builder_2", PromptBuilder("Builder 2: {{query}}"))
|
||
pipeline.add_component("string_joiner", StringJoiner())
|
||
|
||
pipeline.connect("prompt_builder_1.prompt", "string_joiner.strings")
|
||
pipeline.connect("prompt_builder_2.prompt", "string_joiner.strings")
|
||
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"prompt_builder_1": {"query": string_1}, "prompt_builder_2": {"query": string_2}},
|
||
expected_outputs={
|
||
"string_joiner": {
|
||
"strings": AnyOrder(
|
||
["Builder 1: What's Natural Language Processing?", "Builder 2: What's is life?"]
|
||
)
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("prompt_builder_1", 1): {
|
||
"query": "What's Natural Language Processing?",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("prompt_builder_2", 1): {"query": "What's is life?", "template": None, "template_variables": None},
|
||
("string_joiner", 1): {
|
||
"strings": AnyOrder(
|
||
["Builder 1: What's Natural Language Processing?", "Builder 2: What's is life?"]
|
||
)
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that is an agent that can use RAG", target_fixture="pipeline_data")
|
||
def an_agent_that_can_use_RAG():
|
||
@component
|
||
class FakeRetriever:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, query: str) -> dict[str, list[Document]]:
|
||
return {
|
||
"documents": [
|
||
Document(content="This is a document potentially answering the question.", meta={"access_group": 1})
|
||
]
|
||
}
|
||
|
||
agent_prompt_template = """
|
||
Your task is to answer the user's question.
|
||
You can use a RAG system to find information.
|
||
Use the RAG system until you have sufficient information to answer the question.
|
||
To use the RAG system, output "search:" followed by your question.
|
||
Once you have an answer, output "answer:" followed by your answer.
|
||
|
||
Here is the question: {{query}}
|
||
"""
|
||
|
||
rag_prompt_template = """
|
||
Answer the question based on the provided documents.
|
||
Question: {{ query }}
|
||
Documents:
|
||
{% for document in documents %}
|
||
{{ document.content }}
|
||
{% endfor %}
|
||
"""
|
||
|
||
router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{ 'search:' in replies[0] }}",
|
||
"output": "{{ replies[0] }}",
|
||
"output_name": "search",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{ 'answer:' in replies[0] }}",
|
||
"output": "{{ replies }}",
|
||
"output_name": "answer",
|
||
"output_type": list[str],
|
||
},
|
||
]
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=2)
|
||
|
||
pipe.add_component("joiner", BranchJoiner(type_=str))
|
||
pipe.add_component("rag_llm", FixedGenerator(replies=["This is all the information I found!"]))
|
||
pipe.add_component("rag_prompt", PromptBuilder(template=rag_prompt_template))
|
||
pipe.add_component("agent_prompt", PromptBuilder(template=agent_prompt_template))
|
||
pipe.add_component("agent_llm", FixedGenerator(replies=["search: Can you help me?", "answer: here is my answer"]))
|
||
pipe.add_component("router", router)
|
||
pipe.add_component(
|
||
"concatenator", OutputAdapter(template="{{current_prompt + '\n' + rag_answer[0]}}", output_type=str)
|
||
)
|
||
pipe.add_component("retriever", FakeRetriever())
|
||
pipe.add_component("answer_builder", AnswerBuilder())
|
||
|
||
pipe.connect("agent_prompt.prompt", "joiner.value")
|
||
pipe.connect("joiner.value", "agent_llm.prompt")
|
||
pipe.connect("agent_llm.replies", "router.replies")
|
||
pipe.connect("router.search", "retriever.query")
|
||
pipe.connect("router.answer", "answer_builder.replies")
|
||
pipe.connect("retriever.documents", "rag_prompt.documents")
|
||
pipe.connect("rag_prompt.prompt", "rag_llm.prompt")
|
||
pipe.connect("rag_llm.replies", "concatenator.rag_answer")
|
||
pipe.connect("joiner.value", "concatenator.current_prompt")
|
||
pipe.connect("concatenator.output", "joiner.value")
|
||
|
||
query = "Does this run reliably?"
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={
|
||
"agent_prompt": {"query": query},
|
||
"rag_prompt": {"query": query},
|
||
"answer_builder": {"query": query},
|
||
},
|
||
expected_outputs={
|
||
"answer_builder": {
|
||
"answers": [
|
||
GeneratedAnswer(
|
||
data="answer: here is my answer",
|
||
query=query,
|
||
documents=[],
|
||
meta={"all_messages": ["answer: here is my answer"]},
|
||
)
|
||
]
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("agent_llm", 1): {
|
||
"prompt": "\n"
|
||
"Your task is to answer the user's question.\n"
|
||
"You can use a RAG system to find information.\n"
|
||
"Use the RAG system until you have sufficient "
|
||
"information to answer the question.\n"
|
||
'To use the RAG system, output "search:" '
|
||
"followed by your question.\n"
|
||
'Once you have an answer, output "answer:" '
|
||
"followed by your answer.\n"
|
||
"\n"
|
||
"Here is the question: Does this run reliably?\n"
|
||
" "
|
||
},
|
||
("agent_llm", 2): {
|
||
"prompt": "\n"
|
||
"Your task is to answer the user's question.\n"
|
||
"You can use a RAG system to find information.\n"
|
||
"Use the RAG system until you have sufficient "
|
||
"information to answer the question.\n"
|
||
'To use the RAG system, output "search:" '
|
||
"followed by your question.\n"
|
||
'Once you have an answer, output "answer:" '
|
||
"followed by your answer.\n"
|
||
"\n"
|
||
"Here is the question: Does this run reliably?\n"
|
||
" \n"
|
||
"This is all the information I found!"
|
||
},
|
||
("agent_prompt", 1): {
|
||
"query": "Does this run reliably?",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("answer_builder", 1): {
|
||
"documents": None,
|
||
"meta": None,
|
||
"pattern": None,
|
||
"query": "Does this run reliably?",
|
||
"expand_reference_ranges": None,
|
||
"reference_pattern": None,
|
||
"replies": ["answer: here is my answer"],
|
||
},
|
||
("concatenator", 1): {
|
||
"current_prompt": "\n"
|
||
"Your task is to answer the user's "
|
||
"question.\n"
|
||
"You can use a RAG system to find "
|
||
"information.\n"
|
||
"Use the RAG system until you have "
|
||
"sufficient information to answer the "
|
||
"question.\n"
|
||
"To use the RAG system, output "
|
||
'"search:" followed by your '
|
||
"question.\n"
|
||
"Once you have an answer, output "
|
||
'"answer:" followed by your answer.\n'
|
||
"\n"
|
||
"Here is the question: Does this run "
|
||
"reliably?\n"
|
||
" ",
|
||
"rag_answer": ["This is all the information I found!"],
|
||
},
|
||
("joiner", 1): {
|
||
"value": [
|
||
"\n"
|
||
"Your task is to answer the user's question.\n"
|
||
"You can use a RAG system to find information.\n"
|
||
"Use the RAG system until you have sufficient "
|
||
"information to answer the question.\n"
|
||
'To use the RAG system, output "search:" followed '
|
||
"by your question.\n"
|
||
'Once you have an answer, output "answer:" followed '
|
||
"by your answer.\n"
|
||
"\n"
|
||
"Here is the question: Does this run reliably?\n"
|
||
" "
|
||
]
|
||
},
|
||
("joiner", 2): {
|
||
"value": [
|
||
"\n"
|
||
"Your task is to answer the user's question.\n"
|
||
"You can use a RAG system to find information.\n"
|
||
"Use the RAG system until you have sufficient "
|
||
"information to answer the question.\n"
|
||
'To use the RAG system, output "search:" followed '
|
||
"by your question.\n"
|
||
'Once you have an answer, output "answer:" followed '
|
||
"by your answer.\n"
|
||
"\n"
|
||
"Here is the question: Does this run reliably?\n"
|
||
" \n"
|
||
"This is all the information I found!"
|
||
]
|
||
},
|
||
("rag_llm", 1): {
|
||
"prompt": "\n"
|
||
"Answer the question based on the provided "
|
||
"documents.\n"
|
||
"Question: Does this run reliably?\n"
|
||
"Documents:\n"
|
||
"\n"
|
||
"This is a document potentially answering the "
|
||
"question.\n"
|
||
"\n"
|
||
" "
|
||
},
|
||
("rag_prompt", 1): {
|
||
"documents": [
|
||
Document(
|
||
id="366a10745500c26f1177f434c74513daacaa7f9d2e09ba892cfcd48652eb80c1",
|
||
content="This is a document potentially answering the question.",
|
||
meta={"access_group": 1},
|
||
)
|
||
],
|
||
"query": "Does this run reliably?",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("retriever", 1): {"query": "search: Can you help me?"},
|
||
("router", 1): {"replies": ["search: Can you help me?"]},
|
||
("router", 2): {"replies": ["answer: here is my answer"]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has a feedback loop", target_fixture="pipeline_data")
|
||
def has_feedback_loop():
|
||
code_prompt_template = """
|
||
Generate code to solve the task: {{ task }}
|
||
|
||
{% if feedback %}
|
||
Here is your initial attempt and some feedback:
|
||
{{ feedback }}
|
||
{% endif %}
|
||
"""
|
||
feedback_prompt_template = """
|
||
Check if this code is valid and can run: {{ code[0] }}
|
||
Return "PASS" if it passes and "FAIL" if it fails.
|
||
Provide additional feedback on why it fails.
|
||
"""
|
||
router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{ 'FAIL' in replies[0] }}",
|
||
"output": "{{ replies[0] }}",
|
||
"output_name": "fail",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{ 'PASS' in replies[0] }}",
|
||
"output": "{{ code }}",
|
||
"output_name": "pass",
|
||
"output_type": list[str],
|
||
},
|
||
]
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=100)
|
||
|
||
pipe.add_component("code_llm", FixedGenerator(replies=["invalid code", "valid code"]))
|
||
pipe.add_component("code_prompt", PromptBuilder(template=code_prompt_template, required_variables=["task"]))
|
||
pipe.add_component("feedback_prompt", PromptBuilder(template=feedback_prompt_template, required_variables=["code"]))
|
||
pipe.add_component("feedback_llm", FixedGenerator(replies=["FAIL", "PASS"]))
|
||
pipe.add_component("router", router)
|
||
pipe.add_component(
|
||
"concatenator", OutputAdapter(template="{{current_prompt[0] + '\n' + feedback[0]}}", output_type=str)
|
||
)
|
||
pipe.add_component("answer_builder", AnswerBuilder())
|
||
|
||
pipe.connect("code_prompt.prompt", "code_llm.prompt")
|
||
pipe.connect("code_llm.replies", "feedback_prompt.code")
|
||
pipe.connect("feedback_llm.replies", "router.replies")
|
||
pipe.connect("router.fail", "concatenator.feedback")
|
||
pipe.connect("router.pass", "answer_builder.replies")
|
||
pipe.connect("code_llm.replies", "router.code")
|
||
pipe.connect("feedback_prompt.prompt", "feedback_llm.prompt")
|
||
pipe.connect("code_llm.replies", "concatenator.current_prompt")
|
||
pipe.connect("concatenator.output", "code_prompt.feedback")
|
||
|
||
task = "Generate code to generate christmas ascii-art"
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"code_prompt": {"task": task}, "answer_builder": {"query": task}},
|
||
expected_outputs={
|
||
"answer_builder": {
|
||
"answers": [
|
||
GeneratedAnswer(
|
||
data="valid code", query=task, documents=[], meta={"all_messages": ["valid code"]}
|
||
)
|
||
]
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("answer_builder", 1): {
|
||
"documents": None,
|
||
"meta": None,
|
||
"pattern": None,
|
||
"query": "Generate code to generate christmas ascii-art",
|
||
"expand_reference_ranges": None,
|
||
"reference_pattern": None,
|
||
"replies": ["valid code"],
|
||
},
|
||
("code_llm", 1): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"\n"
|
||
" "
|
||
},
|
||
("code_llm", 2): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"\n"
|
||
"Here is your initial attempt and some feedback:\n"
|
||
"invalid code\n"
|
||
"F\n"
|
||
"\n"
|
||
" "
|
||
},
|
||
("code_prompt", 1): {
|
||
"feedback": "",
|
||
"task": "Generate code to generate christmas ascii-art",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("code_prompt", 2): {
|
||
"feedback": "invalid code\nF",
|
||
"task": "Generate code to generate christmas ascii-art",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("concatenator", 1): {"current_prompt": ["invalid code"], "feedback": "FAIL"},
|
||
("feedback_llm", 1): {
|
||
"prompt": "\n"
|
||
"Check if this code is valid and can run: "
|
||
"invalid code\n"
|
||
'Return "PASS" if it passes and "FAIL" if it '
|
||
"fails.\n"
|
||
"Provide additional feedback on why it "
|
||
"fails.\n"
|
||
" "
|
||
},
|
||
("feedback_llm", 2): {
|
||
"prompt": "\n"
|
||
"Check if this code is valid and can run: "
|
||
"valid code\n"
|
||
'Return "PASS" if it passes and "FAIL" if it '
|
||
"fails.\n"
|
||
"Provide additional feedback on why it "
|
||
"fails.\n"
|
||
" "
|
||
},
|
||
("feedback_prompt", 1): {"code": ["invalid code"], "template": None, "template_variables": None},
|
||
("feedback_prompt", 2): {"code": ["valid code"], "template": None, "template_variables": None},
|
||
("router", 1): {"code": ["invalid code"], "replies": ["FAIL"]},
|
||
("router", 2): {"code": ["valid code"], "replies": ["PASS"]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline created in a non-standard order that has a loop", target_fixture="pipeline_data")
|
||
def has_non_standard_order_loop():
|
||
code_prompt_template = """
|
||
Generate code to solve the task: {{ task }}
|
||
|
||
{% if feedback %}
|
||
Here is your initial attempt and some feedback:
|
||
{{ feedback }}
|
||
{% endif %}
|
||
"""
|
||
feedback_prompt_template = """
|
||
Check if this code is valid and can run: {{ code[0] }}
|
||
Return "PASS" if it passes and "FAIL" if it fails.
|
||
Provide additional feedback on why it fails.
|
||
"""
|
||
router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{ 'FAIL' in replies[0] }}",
|
||
"output": "{{ replies[0] }}",
|
||
"output_name": "fail",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{ 'PASS' in replies[0] }}",
|
||
"output": "{{ code }}",
|
||
"output_name": "pass",
|
||
"output_type": list[str],
|
||
},
|
||
]
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=2)
|
||
|
||
pipe.add_component(
|
||
"concatenator", OutputAdapter(template="{{current_prompt[0] + '\n' + feedback[0]}}", output_type=str)
|
||
)
|
||
pipe.add_component("code_llm", FixedGenerator(replies=["invalid code", "valid code"]))
|
||
pipe.add_component("code_prompt", PromptBuilder(template=code_prompt_template, required_variables=["task"]))
|
||
pipe.add_component("feedback_prompt", PromptBuilder(template=feedback_prompt_template, required_variables=["code"]))
|
||
pipe.add_component("feedback_llm", FixedGenerator(replies=["FAIL", "PASS"]))
|
||
pipe.add_component("router", router)
|
||
pipe.add_component("answer_builder", AnswerBuilder())
|
||
|
||
pipe.connect("concatenator.output", "code_prompt.feedback")
|
||
pipe.connect("code_prompt.prompt", "code_llm.prompt")
|
||
pipe.connect("code_llm.replies", "feedback_prompt.code")
|
||
pipe.connect("feedback_llm.replies", "router.replies")
|
||
pipe.connect("router.fail", "concatenator.feedback")
|
||
pipe.connect("feedback_prompt.prompt", "feedback_llm.prompt")
|
||
pipe.connect("router.pass", "answer_builder.replies")
|
||
pipe.connect("code_llm.replies", "router.code")
|
||
pipe.connect("code_llm.replies", "concatenator.current_prompt")
|
||
|
||
task = "Generate code to generate christmas ascii-art"
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"code_prompt": {"task": task}, "answer_builder": {"query": task}},
|
||
expected_outputs={
|
||
"answer_builder": {
|
||
"answers": [
|
||
GeneratedAnswer(
|
||
data="valid code", query=task, documents=[], meta={"all_messages": ["valid code"]}
|
||
)
|
||
]
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("answer_builder", 1): {
|
||
"documents": None,
|
||
"meta": None,
|
||
"pattern": None,
|
||
"query": "Generate code to generate christmas ascii-art",
|
||
"expand_reference_ranges": None,
|
||
"reference_pattern": None,
|
||
"replies": ["valid code"],
|
||
},
|
||
("code_llm", 1): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"\n"
|
||
" "
|
||
},
|
||
("code_llm", 2): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"\n"
|
||
"Here is your initial attempt and some feedback:\n"
|
||
"invalid code\n"
|
||
"F\n"
|
||
"\n"
|
||
" "
|
||
},
|
||
("code_prompt", 1): {
|
||
"feedback": "",
|
||
"task": "Generate code to generate christmas ascii-art",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("code_prompt", 2): {
|
||
"feedback": "invalid code\nF",
|
||
"task": "Generate code to generate christmas ascii-art",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("concatenator", 1): {"current_prompt": ["invalid code"], "feedback": "FAIL"},
|
||
("feedback_llm", 1): {
|
||
"prompt": "\n"
|
||
"Check if this code is valid and can run: "
|
||
"invalid code\n"
|
||
'Return "PASS" if it passes and "FAIL" if it '
|
||
"fails.\n"
|
||
"Provide additional feedback on why it "
|
||
"fails.\n"
|
||
" "
|
||
},
|
||
("feedback_llm", 2): {
|
||
"prompt": "\n"
|
||
"Check if this code is valid and can run: "
|
||
"valid code\n"
|
||
'Return "PASS" if it passes and "FAIL" if it '
|
||
"fails.\n"
|
||
"Provide additional feedback on why it "
|
||
"fails.\n"
|
||
" "
|
||
},
|
||
("feedback_prompt", 1): {"code": ["invalid code"], "template": None, "template_variables": None},
|
||
("feedback_prompt", 2): {"code": ["valid code"], "template": None, "template_variables": None},
|
||
("router", 1): {"code": ["invalid code"], "replies": ["FAIL"]},
|
||
("router", 2): {"code": ["valid code"], "replies": ["PASS"]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has an agent with a feedback cycle", target_fixture="pipeline_data")
|
||
def agent_with_feedback_cycle():
|
||
@component
|
||
class FakeFileEditor:
|
||
@component.output_types(files=str)
|
||
def run(self, replies: list[str]) -> dict[str, str]:
|
||
return {"files": "This is the edited file content."}
|
||
|
||
code_prompt_template = """
|
||
Generate code to solve the task: {{ task }}
|
||
|
||
You can edit files by returning:
|
||
Edit: file_name
|
||
|
||
Once you solved the task, respond with:
|
||
Task finished!
|
||
|
||
{% if feedback %}
|
||
Here is your initial attempt and some feedback:
|
||
{{ feedback }}
|
||
{% endif %}
|
||
"""
|
||
feedback_prompt_template = """
|
||
{% if task_finished %}
|
||
Check if this code is valid and can run: {{ code }}
|
||
Return "PASS" if it passes and "FAIL" if it fails.
|
||
Provide additional feedback on why it fails.
|
||
{% endif %}
|
||
"""
|
||
feedback_router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{ 'FAIL' in replies[0] }}",
|
||
"output": "{{ current_prompt + '\n' + replies[0] }}",
|
||
"output_name": "fail",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{ 'PASS' in replies[0] }}",
|
||
"output": "{{ replies }}",
|
||
"output_name": "pass",
|
||
"output_type": list[str],
|
||
},
|
||
]
|
||
)
|
||
tool_use_router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{ 'Edit:' in replies[0] }}",
|
||
"output": "{{ replies }}",
|
||
"output_name": "edit",
|
||
"output_type": list[str],
|
||
},
|
||
{
|
||
"condition": "{{ 'Task finished!' in replies[0] }}",
|
||
"output": "{{ replies }}",
|
||
"output_name": "done",
|
||
"output_type": list[str],
|
||
},
|
||
]
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=8)
|
||
|
||
pipe.add_component("code_prompt", PromptBuilder(template=code_prompt_template, required_variables=["task"]))
|
||
pipe.add_component("joiner", BranchJoiner(type_=str))
|
||
pipe.add_component(
|
||
"code_llm", FixedGenerator(replies=["Edit: file_1.py", "Edit: file_2.py", "Edit: file_3.py", "Task finished!"])
|
||
)
|
||
pipe.add_component("tool_use_router", tool_use_router)
|
||
pipe.add_component("file_editor", FakeFileEditor())
|
||
pipe.add_component(
|
||
"agent_concatenator", OutputAdapter(template="{{current_prompt + '\n' + files}}", output_type=str)
|
||
)
|
||
pipe.add_component(
|
||
"feedback_prompt", PromptBuilder(template=feedback_prompt_template, required_variables=["task_finished"])
|
||
)
|
||
pipe.add_component("feedback_llm", FixedGenerator(replies=["FAIL", "PASS"]))
|
||
pipe.add_component("feedback_router", feedback_router)
|
||
|
||
# Main Agent
|
||
pipe.connect("code_prompt.prompt", "joiner.value")
|
||
pipe.connect("joiner.value", "code_llm.prompt")
|
||
pipe.connect("code_llm.replies", "tool_use_router.replies")
|
||
pipe.connect("tool_use_router.edit", "file_editor.replies")
|
||
pipe.connect("file_editor.files", "agent_concatenator.files")
|
||
pipe.connect("joiner.value", "agent_concatenator.current_prompt")
|
||
pipe.connect("agent_concatenator.output", "joiner.value")
|
||
|
||
# Feedback Cycle
|
||
pipe.connect("tool_use_router.done", "feedback_prompt.task_finished")
|
||
pipe.connect("agent_concatenator.output", "feedback_prompt.code")
|
||
pipe.connect("feedback_prompt.prompt", "feedback_llm.prompt")
|
||
pipe.connect("feedback_llm.replies", "feedback_router.replies")
|
||
pipe.connect("agent_concatenator.output", "feedback_router.current_prompt")
|
||
pipe.connect("feedback_router.fail", "joiner.value")
|
||
|
||
task = "Generate code to generate christmas ascii-art"
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"code_prompt": {"task": task}},
|
||
expected_outputs={"feedback_router": {"pass": ["PASS"]}},
|
||
expected_component_calls={
|
||
("agent_concatenator", 1): {
|
||
"current_prompt": "\n"
|
||
"Generate code to solve the "
|
||
"task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by "
|
||
"returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, "
|
||
"respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" ",
|
||
"files": "This is the edited file content.",
|
||
},
|
||
("agent_concatenator", 2): {
|
||
"current_prompt": "\n"
|
||
"Generate code to solve the "
|
||
"task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by "
|
||
"returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, "
|
||
"respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file "
|
||
"content.",
|
||
"files": "This is the edited file content.",
|
||
},
|
||
("agent_concatenator", 3): {
|
||
"current_prompt": "\n"
|
||
"Generate code to solve the "
|
||
"task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by "
|
||
"returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, "
|
||
"respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"This is the edited file "
|
||
"content.",
|
||
"files": "This is the edited file content.",
|
||
},
|
||
("agent_concatenator", 4): {
|
||
"current_prompt": "\n"
|
||
"Generate code to solve the "
|
||
"task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by "
|
||
"returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, "
|
||
"respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"FAIL",
|
||
"files": "This is the edited file content.",
|
||
},
|
||
("agent_concatenator", 5): {
|
||
"current_prompt": "\n"
|
||
"Generate code to solve the "
|
||
"task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by "
|
||
"returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, "
|
||
"respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file "
|
||
"content.",
|
||
"files": "This is the edited file content.",
|
||
},
|
||
("agent_concatenator", 6): {
|
||
"current_prompt": "\n"
|
||
"Generate code to solve the "
|
||
"task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by "
|
||
"returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, "
|
||
"respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file "
|
||
"content.\n"
|
||
"This is the edited file "
|
||
"content.",
|
||
"files": "This is the edited file content.",
|
||
},
|
||
("code_llm", 1): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" "
|
||
},
|
||
("code_llm", 2): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content."
|
||
},
|
||
("code_llm", 3): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content."
|
||
},
|
||
("code_llm", 4): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content."
|
||
},
|
||
("code_llm", 5): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL"
|
||
},
|
||
("code_llm", 6): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file content."
|
||
},
|
||
("code_llm", 7): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content."
|
||
},
|
||
("code_llm", 8): {
|
||
"prompt": "\n"
|
||
"Generate code to solve the task: Generate code "
|
||
"to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content."
|
||
},
|
||
("code_prompt", 1): {"feedback": "", "task": task, "template": None, "template_variables": None},
|
||
("feedback_llm", 1): {
|
||
"prompt": "\n"
|
||
"\n"
|
||
"Check if this code is valid and can run: \n"
|
||
"Generate code to solve the task: Generate "
|
||
"code to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
'Return "PASS" if it passes and "FAIL" if it '
|
||
"fails.\n"
|
||
"Provide additional feedback on why it "
|
||
"fails.\n"
|
||
"\n"
|
||
" "
|
||
},
|
||
("feedback_llm", 2): {
|
||
"prompt": "\n"
|
||
"\n"
|
||
"Check if this code is valid and can run: \n"
|
||
"Generate code to solve the task: Generate "
|
||
"code to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
'Return "PASS" if it passes and "FAIL" if it '
|
||
"fails.\n"
|
||
"Provide additional feedback on why it "
|
||
"fails.\n"
|
||
"\n"
|
||
" "
|
||
},
|
||
("feedback_prompt", 1): {
|
||
"code": "\n"
|
||
"Generate code to solve the task: Generate "
|
||
"code to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.",
|
||
"task_finished": ["Task finished!"],
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("feedback_prompt", 2): {
|
||
"code": "\n"
|
||
"Generate code to solve the task: Generate "
|
||
"code to generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.",
|
||
"task_finished": ["Task finished!"],
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("feedback_router", 1): {
|
||
"current_prompt": "\n"
|
||
"Generate code to solve the task: "
|
||
"Generate code to generate "
|
||
"christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond "
|
||
"with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.",
|
||
"replies": ["FAIL"],
|
||
},
|
||
("feedback_router", 2): {
|
||
"current_prompt": "\n"
|
||
"Generate code to solve the task: "
|
||
"Generate code to generate "
|
||
"christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond "
|
||
"with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.",
|
||
"replies": ["PASS"],
|
||
},
|
||
("file_editor", 1): {"replies": ["Edit: file_1.py"]},
|
||
("file_editor", 2): {"replies": ["Edit: file_2.py"]},
|
||
("file_editor", 3): {"replies": ["Edit: file_3.py"]},
|
||
("file_editor", 4): {"replies": ["Edit: file_1.py"]},
|
||
("file_editor", 5): {"replies": ["Edit: file_2.py"]},
|
||
("file_editor", 6): {"replies": ["Edit: file_3.py"]},
|
||
("joiner", 1): {
|
||
"value": [
|
||
"\n"
|
||
"Generate code to solve the task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" "
|
||
]
|
||
},
|
||
("joiner", 2): {
|
||
"value": [
|
||
"\n"
|
||
"Generate code to solve the task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content."
|
||
]
|
||
},
|
||
("joiner", 3): {
|
||
"value": [
|
||
"\n"
|
||
"Generate code to solve the task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content."
|
||
]
|
||
},
|
||
("joiner", 4): {
|
||
"value": [
|
||
"\n"
|
||
"Generate code to solve the task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content."
|
||
]
|
||
},
|
||
("joiner", 5): {
|
||
"value": [
|
||
"\n"
|
||
"Generate code to solve the task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL"
|
||
]
|
||
},
|
||
("joiner", 6): {
|
||
"value": [
|
||
"\n"
|
||
"Generate code to solve the task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file content."
|
||
]
|
||
},
|
||
("joiner", 7): {
|
||
"value": [
|
||
"\n"
|
||
"Generate code to solve the task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content."
|
||
]
|
||
},
|
||
("joiner", 8): {
|
||
"value": [
|
||
"\n"
|
||
"Generate code to solve the task: Generate code to "
|
||
"generate christmas ascii-art\n"
|
||
"\n"
|
||
"You can edit files by returning:\n"
|
||
"Edit: file_name\n"
|
||
"\n"
|
||
"Once you solved the task, respond with:\n"
|
||
"Task finished!\n"
|
||
"\n"
|
||
"\n"
|
||
" \n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"FAIL\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content.\n"
|
||
"This is the edited file content."
|
||
]
|
||
},
|
||
("tool_use_router", 1): {"replies": ["Edit: file_1.py"]},
|
||
("tool_use_router", 2): {"replies": ["Edit: file_2.py"]},
|
||
("tool_use_router", 3): {"replies": ["Edit: file_3.py"]},
|
||
("tool_use_router", 4): {"replies": ["Task finished!"]},
|
||
("tool_use_router", 5): {"replies": ["Edit: file_1.py"]},
|
||
("tool_use_router", 6): {"replies": ["Edit: file_2.py"]},
|
||
("tool_use_router", 7): {"replies": ["Edit: file_3.py"]},
|
||
("tool_use_router", 8): {"replies": ["Task finished!"]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that passes outputs that are consumed in cycle to outside the cycle", target_fixture="pipeline_data")
|
||
def passes_outputs_outside_cycle():
|
||
@component
|
||
class AnswerBuilderWithPrompt:
|
||
@component.output_types(answers=list[GeneratedAnswer])
|
||
def run(self, replies: list[str], query: str, prompt: str | None = None) -> dict[str, Any]:
|
||
answer = GeneratedAnswer(data=replies[0], query=query, documents=[], meta={"all_messages": replies})
|
||
if prompt is not None:
|
||
answer.meta["prompt"] = prompt
|
||
return {"answers": [answer]}
|
||
|
||
feedback_prompt_template = """
|
||
Check if this code is valid and can run: {{ code[0] }}
|
||
Return "PASS" if it passes and "FAIL" if it fails.
|
||
Provide additional feedback on why it fails.
|
||
"""
|
||
valid_response = """
|
||
def generate_santa_sleigh():
|
||
'''
|
||
Returns ASCII art of Santa Claus on his sleigh with Rudolph leading the way.
|
||
'''
|
||
# implementation goes here.
|
||
return art
|
||
"""
|
||
router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{ 'FAIL' in replies[0] }}",
|
||
"output": "{{ replies[0] }}",
|
||
"output_name": "fail",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{ 'PASS' in replies[0] }}",
|
||
"output": "{{ code }}",
|
||
"output_name": "pass",
|
||
"output_type": list[str],
|
||
},
|
||
]
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=3)
|
||
|
||
pipe.add_component(
|
||
"concatenator",
|
||
OutputAdapter(template="{{code_prompt + '\n' + generated_code[0] + '\n' + feedback}}", output_type=str),
|
||
)
|
||
pipe.add_component("code_llm", FixedGenerator(replies=["invalid code", "invalid code", valid_response]))
|
||
pipe.add_component("code_prompt", PromptBuilder(template="{{task}}"))
|
||
pipe.add_component("feedback_prompt", PromptBuilder(template=feedback_prompt_template))
|
||
pipe.add_component("feedback_llm", FixedGenerator(replies=["FAIL", "FAIL, come on, try again.", "PASS"]))
|
||
pipe.add_component("router", router)
|
||
pipe.add_component("joiner", BranchJoiner(type_=str))
|
||
pipe.add_component("answer_builder", AnswerBuilderWithPrompt())
|
||
|
||
pipe.connect("concatenator.output", "joiner.value")
|
||
pipe.connect("joiner.value", "code_prompt.task")
|
||
pipe.connect("code_prompt.prompt", "code_llm.prompt")
|
||
pipe.connect("code_prompt.prompt", "concatenator.code_prompt")
|
||
pipe.connect("code_llm.replies", "feedback_prompt.code")
|
||
pipe.connect("feedback_llm.replies", "router.replies")
|
||
pipe.connect("router.fail", "concatenator.feedback")
|
||
pipe.connect("feedback_prompt.prompt", "feedback_llm.prompt")
|
||
pipe.connect("router.pass", "answer_builder.replies")
|
||
pipe.connect("code_llm.replies", "router.code")
|
||
pipe.connect("code_llm.replies", "concatenator.generated_code")
|
||
pipe.connect("concatenator.output", "answer_builder.prompt")
|
||
|
||
task = "Generate code to generate christmas ascii-art"
|
||
|
||
expected_prompt = """Generate code to generate christmas ascii-art
|
||
invalid code
|
||
FAIL
|
||
invalid code
|
||
FAIL, come on, try again."""
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"joiner": {"value": task}, "answer_builder": {"query": task}},
|
||
expected_outputs={
|
||
"answer_builder": {
|
||
"answers": [
|
||
GeneratedAnswer(
|
||
data=valid_response,
|
||
query=task,
|
||
documents=[],
|
||
meta={"prompt": expected_prompt, "all_messages": [valid_response]},
|
||
)
|
||
]
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("answer_builder", 1): {"prompt": expected_prompt, "query": task, "replies": [valid_response]},
|
||
("code_llm", 1): {"prompt": task},
|
||
("code_llm", 2): {"prompt": "Generate code to generate christmas ascii-art\ninvalid code\nFAIL"},
|
||
("code_llm", 3): {"prompt": expected_prompt},
|
||
("code_prompt", 1): {"task": task, "template": None, "template_variables": None},
|
||
("code_prompt", 2): {
|
||
"task": "Generate code to generate christmas ascii-art\ninvalid code\nFAIL",
|
||
"template": None,
|
||
"template_variables": None,
|
||
},
|
||
("code_prompt", 3): {"task": expected_prompt, "template": None, "template_variables": None},
|
||
("concatenator", 1): {"code_prompt": task, "feedback": "FAIL", "generated_code": ["invalid code"]},
|
||
("concatenator", 2): {
|
||
"code_prompt": "Generate code to generate christmas ascii-art\ninvalid code\nFAIL",
|
||
"feedback": "FAIL, come on, try again.",
|
||
"generated_code": ["invalid code"],
|
||
},
|
||
("feedback_llm", 1): {
|
||
"prompt": "\n"
|
||
"Check if this code is valid and can run: "
|
||
"invalid code\n"
|
||
'Return "PASS" if it passes and "FAIL" if it '
|
||
"fails.\n"
|
||
"Provide additional feedback on why it "
|
||
"fails.\n"
|
||
" "
|
||
},
|
||
("feedback_llm", 2): {
|
||
"prompt": "\n"
|
||
"Check if this code is valid and can run: "
|
||
"invalid code\n"
|
||
'Return "PASS" if it passes and "FAIL" if it '
|
||
"fails.\n"
|
||
"Provide additional feedback on why it "
|
||
"fails.\n"
|
||
" "
|
||
},
|
||
("feedback_llm", 3): {
|
||
"prompt": "\n"
|
||
"Check if this code is valid and can run: \n"
|
||
"def generate_santa_sleigh():\n"
|
||
" '''\n"
|
||
" Returns ASCII art of Santa Claus on his "
|
||
"sleigh with Rudolph leading the way.\n"
|
||
" '''\n"
|
||
" # implementation goes here.\n"
|
||
" return art\n"
|
||
" \n"
|
||
'Return "PASS" if it passes and "FAIL" if it '
|
||
"fails.\n"
|
||
"Provide additional feedback on why it "
|
||
"fails.\n"
|
||
" "
|
||
},
|
||
("feedback_prompt", 1): {"code": ["invalid code"], "template": None, "template_variables": None},
|
||
("feedback_prompt", 2): {"code": ["invalid code"], "template": None, "template_variables": None},
|
||
("feedback_prompt", 3): {"code": [valid_response], "template": None, "template_variables": None},
|
||
("joiner", 1): {"value": [task]},
|
||
("joiner", 2): {"value": ["Generate code to generate christmas ascii-art\ninvalid code\nFAIL"]},
|
||
("joiner", 3): {"value": [expected_prompt]},
|
||
("router", 1): {"code": ["invalid code"], "replies": ["FAIL"]},
|
||
("router", 2): {"code": ["invalid code"], "replies": ["FAIL, come on, try again."]},
|
||
("router", 3): {"code": [valid_response], "replies": ["PASS"]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline with a component that has dynamic default inputs", target_fixture="pipeline_data")
|
||
def pipeline_with_dynamic_defaults():
|
||
@component
|
||
class ParrotWithDynamicDefaultInputs:
|
||
def __init__(self, input_variable: str):
|
||
self.input_variable = input_variable
|
||
component.set_input_type(self, input_variable, str, default="Parrot doesn't only parrot!")
|
||
|
||
@component.output_types(response=str)
|
||
def run(self, **kwargs):
|
||
return {"response": kwargs[self.input_variable]}
|
||
|
||
parrot = ParrotWithDynamicDefaultInputs("parrot")
|
||
pipeline = Pipeline()
|
||
pipeline.add_component("parrot", parrot)
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"parrot": {"parrot": "Are you a parrot?"}},
|
||
expected_outputs={"parrot": {"response": "Are you a parrot?"}},
|
||
expected_component_calls={("parrot", 1): {"parrot": "Are you a parrot?"}},
|
||
),
|
||
PipelineRunData(
|
||
inputs={},
|
||
expected_outputs={"parrot": {"response": "Parrot doesn't only parrot!"}},
|
||
expected_component_calls={("parrot", 1): {"parrot": "Parrot doesn't only parrot!"}},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline with a component that has variadic dynamic default inputs", target_fixture="pipeline_data")
|
||
def pipeline_with_variadic_dynamic_defaults():
|
||
@component
|
||
class ParrotWithVariadicDynamicDefaultInputs:
|
||
def __init__(self, input_variable: str):
|
||
self.input_variable = input_variable
|
||
component.set_input_type(self, input_variable, Variadic[str], default="Parrot doesn't only parrot!")
|
||
|
||
@component.output_types(response=list[str])
|
||
def run(self, **kwargs):
|
||
return {"response": kwargs[self.input_variable]}
|
||
|
||
parrot = ParrotWithVariadicDynamicDefaultInputs("parrot")
|
||
pipeline = Pipeline(max_runs_per_component=1)
|
||
pipeline.add_component("parrot", parrot)
|
||
return (
|
||
pipeline,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"parrot": {"parrot": "Are you a parrot?"}},
|
||
expected_outputs={"parrot": {"response": ["Are you a parrot?"]}},
|
||
expected_component_calls={("parrot", 1): {"parrot": ["Are you a parrot?"]}},
|
||
),
|
||
PipelineRunData(
|
||
inputs={},
|
||
expected_outputs={"parrot": {"response": ["Parrot doesn't only parrot!"]}},
|
||
expected_component_calls={("parrot", 1): {"parrot": ["Parrot doesn't only parrot!"]}},
|
||
),
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that is a file conversion pipeline with two joiners", target_fixture="pipeline_data")
|
||
def pipeline_that_converts_files():
|
||
csv_data = """
|
||
some,header,row
|
||
0,1,0
|
||
"""
|
||
txt_data = "Text file content for testing this."
|
||
sources = [
|
||
ByteStream.from_string(text=csv_data, mime_type="text/csv", meta={"file_type": "csv"}),
|
||
ByteStream.from_string(text=txt_data, mime_type="text/plain", meta={"file_type": "txt"}),
|
||
ByteStream.from_string(
|
||
text='{"content": "Some test content"}', mime_type="application/json", meta={"file_type": "json"}
|
||
),
|
||
]
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
|
||
pipe.add_component("router", FileTypeRouter(mime_types=["text/csv", "text/plain", "application/json"]))
|
||
pipe.add_component("splitter", DocumentSplitter(split_by="word", split_length=3, split_overlap=0))
|
||
pipe.add_component("txt_converter", TextFileToDocument())
|
||
pipe.add_component("csv_converter", CSVToDocument())
|
||
pipe.add_component("json_converter", JSONConverter(content_key="content"))
|
||
pipe.add_component("b_joiner", DocumentJoiner())
|
||
pipe.add_component("a_joiner", DocumentJoiner())
|
||
|
||
pipe.connect("router.text/plain", "txt_converter.sources")
|
||
pipe.connect("router.application/json", "json_converter.sources")
|
||
pipe.connect("router.text/csv", "csv_converter.sources")
|
||
pipe.connect("txt_converter.documents", "b_joiner.documents")
|
||
pipe.connect("json_converter.documents", "b_joiner.documents")
|
||
pipe.connect("csv_converter.documents", "a_joiner.documents")
|
||
pipe.connect("b_joiner.documents", "splitter.documents")
|
||
pipe.connect("splitter.documents", "a_joiner.documents")
|
||
|
||
expected_pre_split_docs = [
|
||
Document(content="Some test content", meta={"file_type": "json"}),
|
||
Document(content=txt_data, meta={"file_type": "txt"}),
|
||
]
|
||
expected_splits_docs = [
|
||
Document(
|
||
content="Some test content",
|
||
meta={
|
||
"file_type": "json",
|
||
"source_id": "7eead7200d4ecead81a174a1da6512d8955f3a23acdc3f8431885d4793a63a74",
|
||
"page_number": 1,
|
||
"split_id": 0,
|
||
"split_idx_start": 0,
|
||
},
|
||
),
|
||
Document(
|
||
content="Text file content ",
|
||
meta={
|
||
"file_type": "txt",
|
||
"source_id": "696d5c046b58b24bf806ff94f6b529fb3a08f068b6bf39e572683537736a0c27",
|
||
"page_number": 1,
|
||
"split_id": 0,
|
||
"split_idx_start": 0,
|
||
},
|
||
),
|
||
Document(
|
||
content="for testing this.",
|
||
meta={
|
||
"file_type": "txt",
|
||
"source_id": "696d5c046b58b24bf806ff94f6b529fb3a08f068b6bf39e572683537736a0c27",
|
||
"page_number": 1,
|
||
"split_id": 1,
|
||
"split_idx_start": 18,
|
||
},
|
||
),
|
||
]
|
||
expected_csv_docs = [Document(content=csv_data, meta={"file_type": "csv"})]
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"router": {"sources": sources}},
|
||
expected_outputs={"a_joiner": {"documents": expected_csv_docs + expected_splits_docs}},
|
||
expected_component_calls={
|
||
("router", 1): {"sources": sources, "meta": None},
|
||
("csv_converter", 1): {"sources": [sources[0]], "meta": None, "content_column": None},
|
||
("txt_converter", 1): {"sources": [sources[1]], "meta": None},
|
||
("json_converter", 1): {"sources": [sources[2]], "meta": None},
|
||
("b_joiner", 1): {
|
||
"documents": AnyOrder(
|
||
[[expected_pre_split_docs[0]], [expected_pre_split_docs[1]]], key=lambda d: d[0].id
|
||
),
|
||
"top_k": None,
|
||
},
|
||
("splitter", 1): {"documents": expected_pre_split_docs},
|
||
("a_joiner", 1): {
|
||
"documents": AnyOrder([expected_csv_docs, expected_splits_docs], key=lambda d: d[0].id),
|
||
"top_k": None,
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that is a file conversion pipeline with three joiners", target_fixture="pipeline_data")
|
||
def pipeline_that_converts_files_with_three_joiners():
|
||
# What does this test?
|
||
# When a component does not produce outputs, and the successors only receive inputs from this component,
|
||
# then the successors will not run.
|
||
# The successor of the successor would never know that its predecessor did not run.
|
||
# This tests an edge case where multiple lazy variadic components wait for input and the execution order is
|
||
# determined by a topological sort.
|
||
html_data = """
|
||
<html><body>Some content</body></html>
|
||
"""
|
||
txt_data = "Text file content"
|
||
sources = [
|
||
ByteStream.from_string(text=txt_data, mime_type="text/plain", meta={"file_type": "txt"}),
|
||
ByteStream.from_string(text=html_data, mime_type="text/html", meta={"file_type": "html"}),
|
||
]
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
|
||
pipe.add_component("router", FileTypeRouter(mime_types=["text/csv", "text/plain", "application/json", "text/html"]))
|
||
pipe.add_component("splitter", DocumentSplitter(split_by="word", split_length=3, split_overlap=0))
|
||
pipe.add_component("txt_converter", TextFileToDocument())
|
||
pipe.add_component("csv_converter", CSVToDocument())
|
||
pipe.add_component("json_converter", JSONConverter(content_key="content"))
|
||
pipe.add_component("html_converter", HTMLToDocument())
|
||
pipe.add_component("joiner", DocumentJoiner())
|
||
pipe.add_component("DocumentJoiner_1", DocumentJoiner())
|
||
pipe.add_component("DocumentJoiner_2", DocumentJoiner())
|
||
pipe.add_component("page_splitter", DocumentSplitter(split_by="page", split_length=1, split_overlap=0))
|
||
|
||
pipe.connect("router.text/plain", "txt_converter.sources")
|
||
pipe.connect("router.application/json", "json_converter.sources")
|
||
pipe.connect("router.text/csv", "csv_converter.sources")
|
||
pipe.connect("router.text/html", "html_converter.sources")
|
||
pipe.connect("txt_converter.documents", "joiner.documents")
|
||
pipe.connect("json_converter.documents", "joiner.documents")
|
||
pipe.connect("csv_converter.documents", "DocumentJoiner_1.documents")
|
||
pipe.connect("html_converter.documents", "DocumentJoiner_1.documents")
|
||
pipe.connect("joiner.documents", "splitter.documents")
|
||
pipe.connect("DocumentJoiner_1.documents", "page_splitter.documents")
|
||
pipe.connect("splitter.documents", "DocumentJoiner_2.documents")
|
||
pipe.connect("page_splitter.documents", "DocumentJoiner_2.documents")
|
||
|
||
expected_html_doc = Document(content="Some content", meta={"file_type": "html"})
|
||
# Need to use ANY for id since not sure how it's calculated in splitter
|
||
expected_html_doc_aft_splitter = Document(
|
||
id=ANY,
|
||
content="Some content",
|
||
meta={
|
||
"file_type": "html",
|
||
"page_number": 1,
|
||
"split_id": 0,
|
||
"split_idx_start": 0,
|
||
"source_id": expected_html_doc.id,
|
||
},
|
||
)
|
||
expected_txt_doc = Document(content=txt_data, meta={"file_type": "txt"})
|
||
# Need to use ANY for id since not sure how it's calculated in splitter
|
||
expected_txt_doc_aft_splitter = Document(
|
||
id=ANY,
|
||
content=txt_data,
|
||
meta={
|
||
"file_type": "txt",
|
||
"page_number": 1,
|
||
"split_id": 0,
|
||
"split_idx_start": 0,
|
||
"source_id": expected_txt_doc.id,
|
||
},
|
||
)
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"router": {"sources": sources}},
|
||
expected_outputs={
|
||
"DocumentJoiner_2": {
|
||
"documents": AnyOrder(
|
||
[expected_html_doc_aft_splitter, expected_txt_doc_aft_splitter], key=lambda d: d.content
|
||
)
|
||
}
|
||
},
|
||
expected_component_calls={
|
||
("router", 1): {"sources": sources, "meta": None},
|
||
("html_converter", 1): {"sources": [sources[1]], "meta": None, "extraction_kwargs": None},
|
||
("txt_converter", 1): {"sources": [sources[0]], "meta": None},
|
||
("joiner", 1): {"documents": [[expected_txt_doc]], "top_k": None},
|
||
("DocumentJoiner_1", 1): {"documents": [[expected_html_doc]], "top_k": None},
|
||
("DocumentJoiner_2", 1): {
|
||
"documents": AnyOrder(
|
||
[[expected_html_doc_aft_splitter], [expected_txt_doc_aft_splitter]],
|
||
key=lambda d: d[0].content,
|
||
),
|
||
"top_k": None,
|
||
},
|
||
("splitter", 1): {"documents": [expected_txt_doc]},
|
||
("page_splitter", 1): {"documents": [expected_html_doc]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that is a file conversion pipeline with three joiners and a loop", target_fixture="pipeline_data")
|
||
def pipeline_that_converts_files_with_three_joiners_and_a_loop():
|
||
# What does this test?
|
||
# When a component does not produce outputs, and the successors only receive inputs from this component,
|
||
# then the successors will not run.
|
||
# The successor of the successor would never know that its predecessor did not run.
|
||
# This tests an edge case where multiple lazy variadic components wait for input and the execution order is
|
||
# determined by a topological sort.
|
||
@component
|
||
class FakeDataExtractor:
|
||
def __init__(self, metas):
|
||
self.metas = metas
|
||
self.current_idx = 0
|
||
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
|
||
sorted_docs = sorted(documents, key=lambda doc: doc.meta["file_type"])
|
||
if self.current_idx >= len(sorted_docs):
|
||
self.current_idx = 0
|
||
|
||
for doc in sorted_docs:
|
||
doc.meta = {**doc.meta, **self.metas[self.current_idx]}
|
||
self.current_idx += 1
|
||
|
||
return {"documents": sorted_docs}
|
||
|
||
html_data = """
|
||
<html><body>Some content</body></html>
|
||
"""
|
||
txt_data = "Text file content"
|
||
sources = [
|
||
ByteStream.from_string(text=txt_data, mime_type="text/plain", meta={"file_type": "txt"}),
|
||
ByteStream.from_string(text=html_data, mime_type="text/html", meta={"file_type": "html"}),
|
||
]
|
||
extraction_router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{documents[0].meta['iteration'] == 1}}",
|
||
"output_name": "continue",
|
||
"output": "{{documents}}",
|
||
"output_type": list[Document],
|
||
},
|
||
{
|
||
"condition": "{{documents[0].meta['iteration'] == 2}}",
|
||
"output_name": "stop",
|
||
"output": "{{documents}}",
|
||
"output_type": list[Document],
|
||
},
|
||
],
|
||
unsafe=True,
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=2)
|
||
|
||
pipe.add_component("router", FileTypeRouter(mime_types=["text/csv", "text/plain", "application/json", "text/html"]))
|
||
pipe.add_component("splitter", DocumentSplitter(split_by="word", split_length=3, split_overlap=0))
|
||
pipe.add_component("txt_converter", TextFileToDocument())
|
||
pipe.add_component("csv_converter", CSVToDocument())
|
||
pipe.add_component("json_converter", JSONConverter(content_key="content"))
|
||
pipe.add_component("html_converter", HTMLToDocument())
|
||
pipe.add_component("joiner", DocumentJoiner())
|
||
pipe.add_component("DocumentJoiner_1", DocumentJoiner())
|
||
pipe.add_component("DocumentJoiner_2", DocumentJoiner())
|
||
pipe.add_component("page_splitter", DocumentSplitter(split_by="page", split_length=1, split_overlap=0))
|
||
pipe.add_component("metadata_generator", FakeDataExtractor(metas=[{"iteration": 1}, {"iteration": 2}]))
|
||
pipe.add_component("extraction_router", extraction_router)
|
||
pipe.add_component("branch_joiner", BranchJoiner(type_=list[Document]))
|
||
|
||
pipe.connect("router.text/plain", "txt_converter.sources")
|
||
pipe.connect("router.application/json", "json_converter.sources")
|
||
pipe.connect("router.text/csv", "csv_converter.sources")
|
||
pipe.connect("router.text/html", "html_converter.sources")
|
||
pipe.connect("txt_converter.documents", "joiner.documents")
|
||
pipe.connect("json_converter.documents", "joiner.documents")
|
||
pipe.connect("csv_converter.documents", "DocumentJoiner_1.documents")
|
||
pipe.connect("html_converter.documents", "DocumentJoiner_1.documents")
|
||
pipe.connect("joiner.documents", "splitter.documents")
|
||
pipe.connect("DocumentJoiner_1.documents", "page_splitter.documents")
|
||
pipe.connect("splitter.documents", "DocumentJoiner_2.documents")
|
||
pipe.connect("page_splitter.documents", "DocumentJoiner_2.documents")
|
||
pipe.connect("DocumentJoiner_2.documents", "branch_joiner.value")
|
||
pipe.connect("branch_joiner.value", "metadata_generator.documents")
|
||
pipe.connect("metadata_generator.documents", "extraction_router.documents")
|
||
pipe.connect("extraction_router.continue", "branch_joiner.value")
|
||
|
||
expected_html_doc = Document(content="Some content", meta={"file_type": "html"})
|
||
expected_txt_doc = Document(content=txt_data, meta={"file_type": "txt"})
|
||
expected_docs = [
|
||
Document(
|
||
content=expected_html_doc.content,
|
||
meta={
|
||
"file_type": "html",
|
||
"source_id": expected_html_doc.id,
|
||
"page_number": 1,
|
||
"split_id": 0,
|
||
"split_idx_start": 0,
|
||
},
|
||
),
|
||
Document(
|
||
content=txt_data,
|
||
meta={
|
||
"file_type": "txt",
|
||
"source_id": expected_txt_doc.id,
|
||
"page_number": 1,
|
||
"split_id": 0,
|
||
"split_idx_start": 0,
|
||
},
|
||
),
|
||
]
|
||
|
||
expected_docs_iteration_1 = []
|
||
expected_docs_iteration_2 = []
|
||
for doc in expected_docs:
|
||
doc_1 = deepcopy(doc)
|
||
doc_1.meta["iteration"] = 1
|
||
doc_2 = deepcopy(doc)
|
||
doc_2.meta["iteration"] = 2
|
||
expected_docs_iteration_1.append(doc_1)
|
||
expected_docs_iteration_2.append(doc_2)
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"router": {"sources": sources}},
|
||
expected_outputs={"extraction_router": {"stop": expected_docs_iteration_2}},
|
||
expected_component_calls={
|
||
("router", 1): {"sources": sources, "meta": None},
|
||
("html_converter", 1): {"sources": [sources[1]], "meta": None, "extraction_kwargs": None},
|
||
("txt_converter", 1): {"sources": [sources[0]], "meta": None},
|
||
("joiner", 1): {"documents": [[expected_txt_doc]], "top_k": None},
|
||
("DocumentJoiner_1", 1): {"documents": [[expected_html_doc]], "top_k": None},
|
||
# ANY because we can't know the order of documents for the async run path
|
||
("DocumentJoiner_2", 1): {"documents": ANY, "top_k": None},
|
||
("splitter", 1): {"documents": [expected_txt_doc]},
|
||
("page_splitter", 1): {"documents": [expected_html_doc]},
|
||
# Same as above
|
||
("branch_joiner", 1): {"value": ANY},
|
||
# Same as above
|
||
("metadata_generator", 1): {"documents": ANY},
|
||
("extraction_router", 1): {"documents": expected_docs_iteration_1},
|
||
("branch_joiner", 2): {"value": [expected_docs_iteration_1]},
|
||
("metadata_generator", 2): {"documents": expected_docs_iteration_1},
|
||
("extraction_router", 2): {"documents": expected_docs_iteration_2},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has components returning dataframes", target_fixture="pipeline_data")
|
||
def pipeline_has_components_returning_dataframes():
|
||
def get_df():
|
||
return DataFrame({"a": [1, 2], "b": [1, 2]})
|
||
|
||
@component
|
||
class DataFramer:
|
||
@component.output_types(dataframe=DataFrame)
|
||
def run(self, dataframe: DataFrame) -> dict[str, Any]:
|
||
return {"dataframe": get_df()}
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
pipe.add_component("df_1", DataFramer())
|
||
pipe.add_component("df_2", DataFramer())
|
||
pipe.connect("df_1", "df_2")
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"df_1": {"dataframe": get_df()}},
|
||
expected_outputs={"df_2": {"dataframe": get_df()}},
|
||
expected_component_calls={("df_1", 1): {"dataframe": get_df()}, ("df_2", 1): {"dataframe": get_df()}},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline where a single component connects multiple sockets to the same receiver socket",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_single_component_many_sockets_same_target():
|
||
router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{query == 'route_1'}}",
|
||
"output": "{{query}}",
|
||
"output_name": "route_1",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{query == 'route_2'}}",
|
||
"output": "{{query}}",
|
||
"output_name": "route_2",
|
||
"output_type": str,
|
||
},
|
||
]
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
pipe.add_component("joiner", BranchJoiner(type_=str))
|
||
pipe.add_component("router", router)
|
||
pipe.connect("router.route_1", "joiner.value")
|
||
pipe.connect("router.route_2", "joiner.value")
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"router": {"query": "route_1"}},
|
||
expected_outputs={"joiner": {"value": "route_1"}},
|
||
expected_component_calls={("router", 1): {"query": "route_1"}, ("joiner", 1): {"value": ["route_1"]}},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline where a component in a cycle provides inputs for a component outside the cycle in one iteration "
|
||
"and no input in another iteration",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_component_cycle_input_no_input():
|
||
router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{query == 'iterationiterationiterationiteration'}}",
|
||
"output": "{{query}}",
|
||
"output_name": "exit",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{query != 'iterationiterationiterationiteration'}}",
|
||
"output": "{{query}}",
|
||
"output_name": "continue",
|
||
"output_type": str,
|
||
},
|
||
]
|
||
)
|
||
outside_router = ConditionalRouter(
|
||
routes=[
|
||
{
|
||
"condition": "{{query == 'iterationiteration'}}",
|
||
"output": "{{query}}",
|
||
"output_name": "cycle_output",
|
||
"output_type": str,
|
||
},
|
||
{
|
||
"condition": "{{query != 'iterationiteration'}}",
|
||
"output": "{{query}}",
|
||
"output_name": "no_output",
|
||
"output_type": str,
|
||
},
|
||
]
|
||
)
|
||
|
||
pipe = Pipeline(max_runs_per_component=2)
|
||
|
||
pipe.add_component("joiner", BranchJoiner(type_=str))
|
||
pipe.add_component("router", router)
|
||
pipe.add_component("builder", PromptBuilder(template="{{query ~ query}}"))
|
||
pipe.add_component(
|
||
"outside_builder",
|
||
PromptBuilder(
|
||
template="{{cycle_output ~ delayed_input}}", required_variables=["cycle_output", "delayed_input"]
|
||
),
|
||
)
|
||
pipe.add_component("outside_router", outside_router)
|
||
|
||
pipe.connect("joiner.value", "builder.query")
|
||
pipe.connect("builder.prompt", "router.query")
|
||
pipe.connect("router.continue", "joiner.value")
|
||
pipe.connect("builder.prompt", "outside_router.query")
|
||
pipe.connect("outside_router.cycle_output", "outside_builder.cycle_output")
|
||
pipe.connect("router.exit", "outside_builder.delayed_input")
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"joiner": {"value": "iteration"}},
|
||
expected_outputs={
|
||
"outside_builder": {"prompt": "iterationiterationiterationiterationiterationiteration"},
|
||
"outside_router": {"no_output": "iterationiterationiterationiteration"},
|
||
},
|
||
expected_component_calls={
|
||
("joiner", 1): {"value": ["iteration"]},
|
||
("builder", 1): {"query": "iteration", "template": None, "template_variables": None},
|
||
("router", 1): {"query": "iterationiteration"},
|
||
("outside_router", 1): {"query": "iterationiteration"},
|
||
("joiner", 2): {"value": ["iterationiteration"]},
|
||
("builder", 2): {"query": "iterationiteration", "template": None, "template_variables": None},
|
||
("router", 2): {"query": "iterationiterationiterationiteration"},
|
||
("outside_router", 2): {"query": "iterationiterationiterationiteration"},
|
||
("outside_builder", 1): {
|
||
"cycle_output": "iterationiteration",
|
||
"template": None,
|
||
"template_variables": None,
|
||
"delayed_input": "iterationiterationiterationiteration",
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that is blocked because not enough component inputs", target_fixture="pipeline_data")
|
||
def that_is_blocked_not_enough_component_inputs():
|
||
# This component requires both `streams` and `query` inputs to run
|
||
# If one is missing then the pipeline is blocked and cannot run.
|
||
@component
|
||
class PayloadBuilder:
|
||
@component.output_types(payload=dict[str, Any])
|
||
def run(self, streams: list[int], query: str) -> dict[str, Any]:
|
||
return {"payload": {"streams": streams, "query": query}}
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
pipe.add_component(
|
||
"router",
|
||
ConditionalRouter(
|
||
[
|
||
{
|
||
"condition": "{{streams|length < 2}}",
|
||
"output": "{{query}}",
|
||
"output_type": str,
|
||
"output_name": "query",
|
||
},
|
||
{
|
||
"condition": "{{streams|length >= 2}}",
|
||
"output": "{{streams}}",
|
||
"output_type": list[int],
|
||
"output_name": "streams",
|
||
},
|
||
]
|
||
),
|
||
)
|
||
pipe.add_component("payload_builder", PayloadBuilder())
|
||
|
||
# Since the router only outputs either `streams` or `query` it's impossible to run PayloadBuilder.
|
||
pipe.connect("router.streams", "payload_builder.streams")
|
||
pipe.connect("router.query", "payload_builder.query")
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"router": {"streams": [1, 2, 3], "query": "Haystack"}},
|
||
expected_outputs={},
|
||
expected_component_calls={
|
||
("router", 1): {"streams": [1, 2, 3], "query": "Haystack"}
|
||
# PayloadBuilder will not run because it requires both inputs
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that is a file conversion pipeline with three auto joiners", target_fixture="pipeline_data")
|
||
def pipeline_that_converts_files_with_three_auto_joiners():
|
||
html_data = """
|
||
<html><body>Some content</body></html>
|
||
"""
|
||
txt_data = "Text file content"
|
||
|
||
sources = [
|
||
ByteStream.from_string(text=txt_data, mime_type="text/plain", meta={"file_type": "txt"}),
|
||
ByteStream.from_string(text=html_data, mime_type="text/html", meta={"file_type": "html"}),
|
||
]
|
||
|
||
# We use a mock writer to avoid needing to do any clean up if using a document store
|
||
@component
|
||
class MockDocumentWriter:
|
||
@component.output_types(documents_written=int)
|
||
def run(self, documents: list[Document]) -> dict[str, int]:
|
||
return {"documents_written": len(documents)}
|
||
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
|
||
pipe.add_component("router", FileTypeRouter(mime_types=["text/csv", "text/plain", "application/json", "text/html"]))
|
||
pipe.add_component("splitter", DocumentSplitter(split_by="word", split_length=3, split_overlap=0))
|
||
pipe.add_component("txt_converter", TextFileToDocument())
|
||
pipe.add_component("csv_converter", CSVToDocument())
|
||
pipe.add_component("json_converter", JSONConverter(content_key="content"))
|
||
pipe.add_component("html_converter", HTMLToDocument())
|
||
pipe.add_component("writer", MockDocumentWriter())
|
||
pipe.add_component("page_splitter", DocumentSplitter(split_by="page", split_length=1, split_overlap=0))
|
||
|
||
pipe.connect("router.text/plain", "txt_converter.sources")
|
||
pipe.connect("router.application/json", "json_converter.sources")
|
||
pipe.connect("router.text/csv", "csv_converter.sources")
|
||
pipe.connect("router.text/html", "html_converter.sources")
|
||
pipe.connect("txt_converter.documents", "splitter.documents")
|
||
pipe.connect("json_converter.documents", "splitter.documents")
|
||
pipe.connect("csv_converter.documents", "page_splitter.documents")
|
||
pipe.connect("html_converter.documents", "page_splitter.documents")
|
||
pipe.connect("splitter.documents", "writer.documents")
|
||
pipe.connect("page_splitter.documents", "writer.documents")
|
||
|
||
expected_txt_document = Document(
|
||
id=ANY,
|
||
content=txt_data,
|
||
meta={"file_type": "txt", "source_id": ANY, "page_number": 1, "split_id": 0, "split_idx_start": 0},
|
||
)
|
||
expected_html_document = Document(
|
||
id=ANY,
|
||
content="Some content",
|
||
meta={"file_type": "html", "source_id": ANY, "page_number": 1, "split_id": 0, "split_idx_start": 0},
|
||
)
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"router": {"sources": sources}},
|
||
expected_outputs={"writer": {"documents_written": 2}},
|
||
expected_component_calls={
|
||
("router", 1): {"sources": sources, "meta": None},
|
||
("html_converter", 1): {"sources": [sources[1]], "meta": None, "extraction_kwargs": None},
|
||
("txt_converter", 1): {"sources": [sources[0]], "meta": None},
|
||
("writer", 1): {
|
||
"documents": AnyOrder([expected_txt_document, expected_html_document], key=lambda d: d.content)
|
||
},
|
||
("splitter", 1): {"documents": [Document(content=txt_data, meta={"file_type": "txt"})]},
|
||
("page_splitter", 1): {"documents": [Document(content="Some content", meta={"file_type": "html"})]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that has an auto joiner that takes in user inputs", target_fixture="pipeline_data")
|
||
def pipeline_that_has_an_auto_joiner_that_takes_in_user_inputs():
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
|
||
@component
|
||
class FakeRetriever:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, query: str) -> dict[str, list[Document]]:
|
||
doc = Document(content=f"Document for query: {query}")
|
||
return {"documents": [doc]}
|
||
|
||
@component
|
||
class FakeRanker:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
|
||
return {"documents": documents}
|
||
|
||
pipe.add_component("retriever", FakeRetriever())
|
||
pipe.add_component("ranker", FakeRanker())
|
||
pipe.connect("retriever.documents", "ranker.documents")
|
||
|
||
user_doc = Document(content="User document")
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"retriever": {"query": "test query"}, "ranker": {"documents": [user_doc]}},
|
||
expected_outputs={
|
||
"ranker": {"documents": [user_doc, Document(content="Document for query: test query")]}
|
||
},
|
||
expected_component_calls={
|
||
("retriever", 1): {"query": "test query"},
|
||
("ranker", 1): {"documents": [user_doc, Document(content="Document for query: test query")]},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that performs automatic conversion between list of ChatMessage and str", target_fixture="pipeline_data"
|
||
)
|
||
def pipeline_that_performs_automatic_conversion_between_list_of_ChatMessage_and_str():
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
|
||
@component
|
||
class FakeChatGenerator:
|
||
@component.output_types(replies=list[ChatMessage])
|
||
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||
return {"replies": [ChatMessage.from_assistant("Pisa tower height")]}
|
||
|
||
@component
|
||
class FakeRetriever:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, query: str) -> dict[str, list[Document]]:
|
||
return {"documents": [Document(content="Pisa tower height is 55 meters")]}
|
||
|
||
pipe.add_component("chat_generator", FakeChatGenerator())
|
||
pipe.add_component("retriever", FakeRetriever())
|
||
pipe.connect("chat_generator", "retriever")
|
||
|
||
user_query = "Rephrase the following question for a search engine: how tall is the tower of Pisa?"
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"chat_generator": {"messages": [ChatMessage.from_user(user_query)]}},
|
||
expected_outputs={"retriever": {"documents": [Document(content="Pisa tower height is 55 meters")]}},
|
||
expected_component_calls={
|
||
("chat_generator", 1): {"messages": [ChatMessage.from_user(user_query)]},
|
||
("retriever", 1): {"query": "Pisa tower height"},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given(
|
||
"a pipeline that performs automatic conversion wrapping ChatMessage for a Union receiver",
|
||
target_fixture="pipeline_data",
|
||
)
|
||
def pipeline_that_performs_automatic_conversion_wrapping_ChatMessage_for_Union_receiver():
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
|
||
@component
|
||
class FakeAgent:
|
||
@component.output_types(last_message=ChatMessage)
|
||
def run(self, messages: list[ChatMessage]) -> dict[str, ChatMessage]:
|
||
return {"last_message": ChatMessage.from_assistant("Pisa tower is 55 meters tall")}
|
||
|
||
@component
|
||
class FakeAnswerBuilder:
|
||
@component.output_types(answers=list[str])
|
||
def run(self, query: str, replies: list[str] | list[ChatMessage]) -> dict[str, list[str]]:
|
||
return {"answers": [r.text or "" if isinstance(r, ChatMessage) else r for r in replies]}
|
||
|
||
pipe.add_component("agent", FakeAgent())
|
||
pipe.add_component("answer_builder", FakeAnswerBuilder())
|
||
pipe.connect("agent.last_message", "answer_builder.replies")
|
||
|
||
user_query = "How tall is the tower of Pisa?"
|
||
messages = [ChatMessage.from_user(user_query)]
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"agent": {"messages": messages}, "answer_builder": {"query": user_query}},
|
||
expected_outputs={"answer_builder": {"answers": ["Pisa tower is 55 meters tall"]}},
|
||
expected_component_calls={
|
||
("agent", 1): {"messages": messages},
|
||
("answer_builder", 1): {
|
||
"query": user_query,
|
||
"replies": [ChatMessage.from_assistant("Pisa tower is 55 meters tall")],
|
||
},
|
||
},
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
@given("a pipeline that fails automatic conversion between list of ChatMessage and str", target_fixture="pipeline_data")
|
||
def pipeline_that_fails_automatic_conversion_between_list_of_ChatMessage_and_str():
|
||
pipe = Pipeline(max_runs_per_component=1)
|
||
|
||
@component
|
||
class FakeChatGenerator:
|
||
@component.output_types(replies=list[ChatMessage])
|
||
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||
return {"replies": [ChatMessage.from_assistant()]}
|
||
|
||
@component
|
||
class FakeRetriever:
|
||
@component.output_types(documents=list[Document])
|
||
def run(self, query: str) -> dict[str, list[Document]]:
|
||
return {"documents": [Document(content="A relevant document")]}
|
||
|
||
pipe.add_component("chat_generator", FakeChatGenerator())
|
||
pipe.add_component("retriever", FakeRetriever())
|
||
pipe.connect("chat_generator", "retriever")
|
||
|
||
user_query = "Find info on Milano Cortina Olympic Winter Games"
|
||
|
||
return (
|
||
pipe,
|
||
[
|
||
PipelineRunData(
|
||
inputs={"chat_generator": {"messages": [ChatMessage.from_user(user_query)]}},
|
||
expected_component_calls={("chat_generator", 1): {"messages": [ChatMessage.from_user(user_query)]}},
|
||
)
|
||
],
|
||
)
|