Skip to content

Router

Bases: QueryComponent

Selector component.

Source code in llama-index-core/llama_index/core/query_pipeline/components/router.py
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
class SelectorComponent(QueryComponent):
    """Selector component."""

    model_config = ConfigDict(arbitrary_types_allowed=True)
    selector: BaseSelector = Field(..., description="Selector")

    def set_callback_manager(self, callback_manager: CallbackManager) -> None:
        """Set callback manager."""

    def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
        """Validate component inputs during run_component."""
        if "choices" not in input:
            raise ValueError("Input must have key 'choices'")
        if not isinstance(input["choices"], list):
            raise ValueError("Input choices must be a list")

        for idx, choice in enumerate(input["choices"]):
            # make stringable
            input["choices"][idx] = validate_and_convert_stringable(choice)

        # make sure `query` is stringable
        if "query" not in input:
            raise ValueError("Input must have key 'query'")
        input["query"] = validate_and_convert_stringable(input["query"])

        return input

    def _run_component(self, **kwargs: Any) -> Any:
        """Run component."""
        output = self.selector.select(kwargs["choices"], kwargs["query"])
        return {"output": output.selections}

    async def _arun_component(self, **kwargs: Any) -> Any:
        """Run component (async)."""
        # NOTE: no native async for postprocessor
        return self._run_component(**kwargs)

    @property
    def input_keys(self) -> InputKeys:
        """Input keys."""
        return InputKeys.from_keys({"choices", "query"})

    @property
    def output_keys(self) -> OutputKeys:
        """Output keys."""
        return OutputKeys.from_keys({"output"})

input_keys property #

input_keys: InputKeys

Input keys.

output_keys property #

output_keys: OutputKeys

Output keys.

set_callback_manager #

set_callback_manager(callback_manager: CallbackManager) -> None

Set callback manager.

Source code in llama-index-core/llama_index/core/query_pipeline/components/router.py
30
31
def set_callback_manager(self, callback_manager: CallbackManager) -> None:
    """Set callback manager."""

Bases: QueryComponent

Router Component.

Routes queries to different query components based on a selector.

Assumes a single query component is selected.

Source code in llama-index-core/llama_index/core/query_pipeline/components/router.py
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
class RouterComponent(QueryComponent):
    """Router Component.

    Routes queries to different query components based on a selector.

    Assumes a single query component is selected.

    """

    model_config = ConfigDict(arbitrary_types_allowed=True)
    selector: SerializeAsAny[BaseSelector] = Field(..., description="Selector")
    choices: List[str] = Field(
        ..., description="Choices (must correspond to components)"
    )
    components: List[SerializeAsAny[QueryComponent]] = Field(
        ..., description="Components (must correspond to choices)"
    )
    verbose: bool = Field(default=False, description="Verbose")

    _query_keys: List[str] = PrivateAttr()

    def __init__(
        self,
        selector: BaseSelector,
        choices: List[str],
        components: List[QUERY_COMPONENT_TYPE],
        verbose: bool = False,
    ) -> None:
        """Init."""
        new_components = []
        query_keys = []
        for component in components:
            if isinstance(component, ChainableMixin):
                new_component = component.as_query_component()
            else:
                new_component = component

            # validate component has one input key
            if len(new_component.free_req_input_keys) != 1:
                raise ValueError("Expected one required input key")
            query_keys.append(next(iter(new_component.free_req_input_keys)))
            new_components.append(new_component)
        super().__init__(
            selector=selector,
            choices=choices,
            components=new_components,
            verbose=verbose,
        )
        self._query_keys = query_keys

    def set_callback_manager(self, callback_manager: CallbackManager) -> None:
        """Set callback manager."""
        for component in self.components:
            component.set_callback_manager(callback_manager)

    def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
        """Validate component inputs during run_component."""
        # make sure `query` is stringable
        if "query" not in input:
            raise ValueError("Input must have key 'query'")
        input["query"] = validate_and_convert_stringable(input["query"])

        return input

    def validate_component_outputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
        """Validate component inputs during run_component."""
        return input

    def _validate_component_outputs(self, output: Dict[str, Any]) -> Dict[str, Any]:
        raise NotImplementedError

    def _run_component(self, **kwargs: Any) -> Any:
        """Run component."""
        # for the output selection, run the corresponding component, aggregate into list
        sel_output = self.selector.select(self.choices, kwargs["query"])
        # assume one selection
        if len(sel_output.selections) != 1:
            raise ValueError("Expected one selection")
        component = self.components[sel_output.ind]
        log_str = f"Selecting component {sel_output.ind}: " f"{sel_output.reason}."
        if self.verbose:
            print_text(log_str + "\n", color="pink")
        # run component
        # run with input_keys of component
        return component.run_component(
            **{self._query_keys[sel_output.ind]: kwargs["query"]}
        )

    async def _arun_component(self, **kwargs: Any) -> Any:
        """Run component (async)."""
        # for the output selection, run the corresponding component, aggregate into list
        sel_output = await self.selector.aselect(self.choices, kwargs["query"])
        # assume one selection
        if len(sel_output.selections) != 1:
            raise ValueError("Expected one selection")
        component = self.components[sel_output.ind]
        log_str = f"Selecting component {sel_output.ind}: " f"{sel_output.reason}."
        if self.verbose:
            print_text(log_str + "\n", color="pink")
        # run component
        return await component.arun_component(
            **{self._query_keys[sel_output.ind]: kwargs["query"]}
        )

    @property
    def input_keys(self) -> InputKeys:
        """Input keys."""
        return InputKeys.from_keys({"query"})

    @property
    def output_keys(self) -> OutputKeys:
        """Output keys."""
        # not used
        return OutputKeys.from_keys(set())

    @property
    def sub_query_components(self) -> List["QueryComponent"]:
        """Get sub query components.

        Certain query components may have sub query components, e.g. a
        query pipeline will have sub query components, and so will
        an IfElseComponent.

        """
        return self.components

input_keys property #

input_keys: InputKeys

Input keys.

output_keys property #

output_keys: OutputKeys

Output keys.

sub_query_components property #

sub_query_components: List[QueryComponent]

Get sub query components.

Certain query components may have sub query components, e.g. a query pipeline will have sub query components, and so will an IfElseComponent.

set_callback_manager #

set_callback_manager(callback_manager: CallbackManager) -> None

Set callback manager.

Source code in llama-index-core/llama_index/core/query_pipeline/components/router.py
122
123
124
125
def set_callback_manager(self, callback_manager: CallbackManager) -> None:
    """Set callback manager."""
    for component in self.components:
        component.set_callback_manager(callback_manager)

validate_component_outputs #

validate_component_outputs(input: Dict[str, Any]) -> Dict[str, Any]

Validate component inputs during run_component.

Source code in llama-index-core/llama_index/core/query_pipeline/components/router.py
136
137
138
def validate_component_outputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
    """Validate component inputs during run_component."""
    return input