class RetrieverComponent(QueryComponent):
"""Retriever component."""
model_config = ConfigDict(arbitrary_types_allowed=True)
retriever: BaseRetriever = Field(..., description="Retriever")
def set_callback_manager(self, callback_manager: CallbackManager) -> None:
"""Set callback manager."""
self.retriever.callback_manager = callback_manager
def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
"""Validate component inputs during run_component."""
# make sure input is a string
input["input"] = validate_and_convert_stringable(input["input"])
return input
def _run_component(self, **kwargs: Any) -> Any:
"""Run component."""
output = self.retriever.retrieve(kwargs["input"])
return {"output": output}
async def _arun_component(self, **kwargs: Any) -> Any:
"""Run component."""
output = await self.retriever.aretrieve(kwargs["input"])
return {"output": output}
@property
def input_keys(self) -> InputKeys:
"""Input keys."""
return InputKeys.from_keys({"input"})
@property
def output_keys(self) -> OutputKeys:
"""Output keys."""
return OutputKeys.from_keys({"output"})
