diff --git a/examples/recursive-web-loader.py b/examples/recursive-web-loader.py index 2a06fba35..74adcbfc8 100644 --- a/examples/recursive-web-loader.py +++ b/examples/recursive-web-loader.py @@ -8,6 +8,8 @@ from langchain.utils.html import extract_sub_links from permchain import Pregel, channels from permchain.pregel import PregelRead +# Load url with sync httpx client + class LoadUrlInput(TypedDict): url: str @@ -15,22 +17,28 @@ class LoadUrlInput(TypedDict): client: httpx.Client +def load_url(input: LoadUrlInput) -> str: + response = input["client"].get(input["url"]) + return response.text + + +# Same as above but with async httpx client + + class LoadUrlInputAsync(TypedDict): url: str visited: FrozenSet[str] client: httpx.AsyncClient -def load_url(input: LoadUrlInput) -> str: - response = input["client"].get(input["url"]) - return response.text - - async def load_url_async(input: LoadUrlInputAsync) -> str: response = await input["client"].get(input["url"]) return response.text +# default metadata extractor copied from langchain.document_loaders + + def _metadata_extractor(raw_html: str, url: str) -> dict: """Extract metadata from raw html using BeautifulSoup.""" metadata = {"source": url} @@ -55,19 +63,29 @@ def recursive_web_loader( extractor: Optional[Callable[[str], str]] = None, metadata_extractor: Optional[Callable[[str, str], dict]] = None, ) -> Pregel: + # assign default extractors extractor = extractor or (lambda x: x) metadata_extractor = metadata_extractor or _metadata_extractor + # the main chain that gets executed recursively chain = ( + # while there are urls in next_urls + # run the chain below for each url in next_urls + # adding the current values of visited set, base_url and httpx client Pregel.subscribe_to_each("next_urls", key="url").join( ["visited", "client", "base_url"] ) + # load the url (with sync and async implementations) | RunnablePassthrough.assign(body=RunnableLambda(load_url, load_url_async)) | Pregel.send_to( + # send this url to the visited set visited=lambda x: x["url"], + # send a new document to the documents stream documents=lambda x: Document( page_content=extractor(x["body"]), metadata=metadata_extractor(x["body"], x["url"]), ), + # send the next urls to the next_urls set + # only if not visited already next_urls=lambda x: [ url for url in extract_sub_links( @@ -79,8 +97,11 @@ def recursive_web_loader( ) ) return Pregel( + # use the base_url as the first url to visit Pregel.subscribe_to("base_url") | Pregel.send_to("next_urls"), + # add the main chain chain, + # define the channels channels={ "base_url": channels.LastValue(str), "next_urls": channels.UniqueInbox(str), @@ -90,8 +111,10 @@ def recursive_web_loader( httpx.Client | httpx.AsyncClient, httpx.Client, httpx.AsyncClient ), }, + # this will accept a string as input input="base_url", - output="visited", + # and return a dict with documents and visited set + output=["documents", "visited"], )