| from abc import abstractmethod | from abc import abstractmethod | ||||
| from typing import Any, Optional | |||||
| from typing import Optional | |||||
| from msal_extensions.persistence import ABC # type: ignore | from msal_extensions.persistence import ABC # type: ignore | ||||
| from pydantic import BaseModel, ConfigDict | from pydantic import BaseModel, ConfigDict | ||||
| model_config = ConfigDict(arbitrary_types_allowed=True) | model_config = ConfigDict(arbitrary_types_allowed=True) | ||||
| @abstractmethod | @abstractmethod | ||||
| def _run( | |||||
| self, | |||||
| *args: Any, | |||||
| **kwargs: Any, | |||||
| ) -> Any: | |||||
| def _run(self, query: str) -> str: | |||||
| """Use the tool. | """Use the tool. | ||||
| Add run_manager: Optional[CallbackManagerForToolRun] = None | Add run_manager: Optional[CallbackManagerForToolRun] = None |
| expected_value = self.graph_runtime_state.variable_pool.convert_template( | expected_value = self.graph_runtime_state.variable_pool.convert_template( | ||||
| expected_value | expected_value | ||||
| ).value[0] | ).value[0] | ||||
| if expected_value.value_type == "number": # type: ignore | |||||
| if expected_value.value_type in {"number", "integer", "float"}: # type: ignore | |||||
| expected_value = expected_value.value # type: ignore | expected_value = expected_value.value # type: ignore | ||||
| elif expected_value.value_type == "string": # type: ignore | elif expected_value.value_type == "string": # type: ignore | ||||
| expected_value = re.sub(r"[\r\n\t]+", " ", expected_value.text).strip() # type: ignore | expected_value = re.sub(r"[\r\n\t]+", " ", expected_value.text).strip() # type: ignore |
| retriever_resources=original_retriever_resource, context=context_str.strip() | retriever_resources=original_retriever_resource, context=context_str.strip() | ||||
| ) | ) | ||||
| def _convert_to_original_retriever_resource(self, context_dict: dict): | |||||
| def _convert_to_original_retriever_resource(self, context_dict: dict) -> RetrievalSourceMetadata | None: | |||||
| if ( | if ( | ||||
| "metadata" in context_dict | "metadata" in context_dict | ||||
| and "_source" in context_dict["metadata"] | and "_source" in context_dict["metadata"] |