-
Notifications
You must be signed in to change notification settings - Fork 5.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add stateful and loop components (#14235)
- Loading branch information
Showing
10 changed files
with
679 additions
and
189 deletions.
There are no files selected for viewing
301 changes: 130 additions & 171 deletions
301
docs/docs/examples/agent/agent_runner/query_pipeline_agent.ipynb
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
86 changes: 86 additions & 0 deletions
86
llama-index-core/llama_index/core/query_pipeline/components/loop.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,86 @@ | ||
| from llama_index.core.base.query_pipeline.query import ( | ||
| InputKeys, | ||
| OutputKeys, | ||
| QueryComponent, | ||
| ) | ||
| from llama_index.core.query_pipeline.query import QueryPipeline | ||
| from llama_index.core.bridge.pydantic import Field | ||
| from llama_index.core.callbacks.base import CallbackManager | ||
| from typing import Any, Dict, Optional, Callable | ||
|
|
||
|
|
||
| class LoopComponent(QueryComponent): | ||
| """Loop component.""" | ||
|
|
||
| pipeline: QueryPipeline = Field(..., description="Query pipeline") | ||
| should_exit_fn: Optional[Callable] = Field(..., description="Should exit function") | ||
| add_output_to_input_fn: Optional[Callable] = Field( | ||
| ..., | ||
| description="Add output to input function. If not provided, will reuse the original input for the next iteration. If provided, will call the function to combine the output into the input for the next iteration.", | ||
| ) | ||
| max_iterations: Optional[int] = Field(5, description="Max iterations") | ||
|
|
||
| class Config: | ||
| arbitrary_types_allowed = True | ||
|
|
||
| def __init__( | ||
| self, | ||
| pipeline: QueryPipeline, | ||
| should_exit_fn: Optional[Callable] = None, | ||
| add_output_to_input_fn: Optional[Callable] = None, | ||
| max_iterations: Optional[int] = 5, | ||
| ) -> None: | ||
| """Init params.""" | ||
| super().__init__( | ||
| pipeline=pipeline, | ||
| should_exit_fn=should_exit_fn, | ||
| add_output_to_input_fn=add_output_to_input_fn, | ||
| max_iterations=max_iterations, | ||
| ) | ||
|
|
||
| def set_callback_manager(self, callback_manager: CallbackManager) -> None: | ||
| """Set callback manager.""" | ||
| # TODO: implement | ||
|
|
||
| def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]: | ||
| return input | ||
|
|
||
| def _run_component(self, **kwargs: Any) -> Dict: | ||
| """Run component.""" | ||
| current_input = kwargs | ||
| for i in range(self.max_iterations): | ||
| output = self.pipeline.run_component(**current_input) | ||
| if self.should_exit_fn: | ||
| should_exit = self.should_exit_fn(output) | ||
| if should_exit: | ||
| break | ||
|
|
||
| if self.add_output_to_input_fn: | ||
| current_input = self.add_output_to_input_fn(current_input, output) | ||
|
|
||
| return output | ||
|
|
||
| async def _arun_component(self, **kwargs: Any) -> Any: | ||
| """Run component (async).""" | ||
| current_input = kwargs | ||
| for i in range(self.max_iterations): | ||
| output = await self.pipeline.arun_component(**current_input) | ||
| if self.should_exit_fn: | ||
| should_exit = self.should_exit_fn(output) | ||
| if should_exit: | ||
| break | ||
|
|
||
| if self.add_output_to_input_fn: | ||
| current_input = self.add_output_to_input_fn(current_input, output) | ||
|
|
||
| return output | ||
|
|
||
| @property | ||
| def input_keys(self) -> InputKeys: | ||
| """Input keys.""" | ||
| return self.pipeline.input_keys | ||
|
|
||
| @property | ||
| def output_keys(self) -> OutputKeys: | ||
| """Output keys.""" | ||
| return self.pipeline.output_keys |
91 changes: 91 additions & 0 deletions
91
llama-index-core/llama_index/core/query_pipeline/components/stateful.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,91 @@ | ||
| """Agent components.""" | ||
|
|
||
| from typing import Any, Callable, Dict, Optional, Set | ||
|
|
||
| from llama_index.core.base.query_pipeline.query import ( | ||
| QueryComponent, | ||
| ) | ||
| from llama_index.core.bridge.pydantic import Field | ||
| from llama_index.core.query_pipeline.components.function import ( | ||
| FnComponent, | ||
| get_parameters, | ||
| ) | ||
|
|
||
| # from llama_index.core.query_pipeline.components.input import InputComponent | ||
|
|
||
|
|
||
| class BaseStatefulComponent(QueryComponent): | ||
| """Takes in agent inputs and transforms it into desired outputs.""" | ||
|
|
||
| state: Dict[str, Any] = Field( | ||
| default_factory=dict, description="State of the pipeline." | ||
| ) | ||
|
|
||
| def reset_state(self) -> None: | ||
| """Reset state.""" | ||
| self.state = {} | ||
|
|
||
|
|
||
| class StatefulFnComponent(BaseStatefulComponent, FnComponent): | ||
| """Query component that takes in an arbitrary function. | ||
| Stateful version of `FnComponent`. Expects functions to have `state` as the first argument. | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| fn: Callable, | ||
| req_params: Optional[Set[str]] = None, | ||
| opt_params: Optional[Set[str]] = None, | ||
| state: Optional[Dict[str, Any]] = None, | ||
| **kwargs: Any | ||
| ) -> None: | ||
| """Init params.""" | ||
| # determine parameters | ||
| default_req_params, default_opt_params = get_parameters(fn) | ||
| # make sure task and step are part of the list, and remove them from the list | ||
| if "state" not in default_req_params: | ||
| raise ValueError( | ||
| "StatefulFnComponent must have 'state' as required parameters" | ||
| ) | ||
|
|
||
| default_req_params = default_req_params - {"state"} | ||
| default_opt_params = default_opt_params - {"state"} | ||
|
|
||
| if req_params is None: | ||
| req_params = default_req_params | ||
| if opt_params is None: | ||
| opt_params = default_opt_params | ||
|
|
||
| super().__init__( | ||
| fn=fn, | ||
| req_params=req_params, | ||
| opt_params=opt_params, | ||
| state=state or {}, | ||
| **kwargs | ||
| ) | ||
|
|
||
| def _run_component(self, **kwargs: Any) -> Dict: | ||
| """Run component.""" | ||
| kwargs.update({"state": self.state}) | ||
| return super()._run_component(**kwargs) | ||
|
|
||
| async def _arun_component(self, **kwargs: Any) -> Any: | ||
| """Async run component.""" | ||
| kwargs.update({"state": self.state}) | ||
| return await super()._arun_component(**kwargs) | ||
|
|
||
| # @property | ||
| # def input_keys(self) -> InputKeys: | ||
| # """Input keys.""" | ||
| # return InputKeys.from_keys( | ||
| # required_keys={"state", *self._req_params}, | ||
| # optional_keys=self._opt_params, | ||
| # ) | ||
|
|
||
| # @property | ||
| # def output_keys(self) -> OutputKeys: | ||
| # """Output keys.""" | ||
| # # output can be anything, overrode validate function | ||
| # return OutputKeys.from_keys({self.output_key}) |
Oops, something went wrong.