codekingpro/portable-devtools

1"""Chain that hits a URL and then uses an LLM to parse results."""2 3from __future__ import annotations4 5from typing import Any, Dict, List, Optional6 7from langchain_classic.chains import LLMChain8from langchain_classic.chains.base import Chain9from langchain_core.callbacks import CallbackManagerForChainRun10from pydantic import ConfigDict, Field, model_validator11 12from langchain_community.utilities.requests import TextRequestsWrapper13 14DEFAULT_HEADERS = {15    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.88 Safari/537.36"  # noqa: E50116}17 18 19class LLMRequestsChain(Chain):20    """Chain that requests a URL and then uses an LLM to parse results.21 22    **Security Note**: This chain can make GET requests to arbitrary URLs,23        including internal URLs.24 25        Control access to who can run this chain and what network access26        this chain has.27 28        See https://python.langchain.com/docs/security for more information.29    """30 31    llm_chain: LLMChain32    requests_wrapper: TextRequestsWrapper = Field(33        default_factory=lambda: TextRequestsWrapper(headers=DEFAULT_HEADERS),34        exclude=True,35    )36    text_length: int = 800037    requests_key: str = "requests_result"  #: :meta private:38    input_key: str = "url"  #: :meta private:39    output_key: str = "output"  #: :meta private:40 41    model_config = ConfigDict(42        arbitrary_types_allowed=True,43        extra="forbid",44    )45 46    @property47    def input_keys(self) -> List[str]:48        """Will be whatever keys the prompt expects.49 50        :meta private:51        """52        return [self.input_key]53 54    @property55    def output_keys(self) -> List[str]:56        """Will always return text key.57 58        :meta private:59        """60        return [self.output_key]61 62    @model_validator(mode="before")63    @classmethod64    def validate_environment(cls, values: Dict) -> Any:65        """Validate that api key and python package exists in environment."""66        try:67            from bs4 import BeautifulSoup  # noqa: F40168 69        except ImportError:70            raise ImportError(71                "Could not import bs4 python package. "72                "Please install it with `pip install bs4`."73            )74        return values75 76    def _call(77        self,78        inputs: Dict[str, Any],79        run_manager: Optional[CallbackManagerForChainRun] = None,80    ) -> Dict[str, Any]:81        from bs4 import BeautifulSoup82 83        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()84        # Other keys are assumed to be needed for LLM prediction85        other_keys = {k: v for k, v in inputs.items() if k != self.input_key}86        url = inputs[self.input_key]87        res = self.requests_wrapper.get(url)88        # extract the text from the html89        soup = BeautifulSoup(res, "html.parser")  # type: ignore[arg-type]90        other_keys[self.requests_key] = soup.get_text()[: self.text_length]91        result = self.llm_chain.predict(92            callbacks=_run_manager.get_child(), **other_keys93        )94        return {self.output_key: result}95 96    @property97    def _chain_type(self) -> str:98        return "llm_requests_chain"99