WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-22 17:45:09 +02:00
Code Issues Packages Projects Releases Wiki Activity

Implement reduce chain

Browse Source
This commit is contained in:
Nuno Campos
2023-10-23 12:01:35 +01:00
parent 1e08395810
commit 49ff5e2eab
11 changed files with 290 additions and 242 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/combine_docs.ipynb
+106 -86
View File
@@ -12,26 +12,20 @@
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": 1,
"id": "624c452c-ddd5-4390-9065-7ec55dc64b96",
"metadata": {},
"outputs": [],
"source": [
"from operator import itemgetter\n",
"\n",
"from langchain.chat_models.openai import ChatOpenAI\n",
"from langchain.prompts import (\n",
" SystemMessagePromptTemplate,\n",
" ChatPromptTemplate,\n",
" PromptTemplate,\n",
")\n",
"from langchain.prompts import ChatPromptTemplate, PromptTemplate\n",
"from langchain.schema.output_parser import StrOutputParser\n",
"from langchain.schema.runnable import Runnable, RunnablePassthrough\n",
"from langchain.schema.output_parser import StrOutputParser\n",
"from langchain.runnables.openai_functions import OpenAIFunctionsRouter\n",
"from langchain.schema.runnable import RunnableMap, RunnablePassthrough\n",
"from langchain.schema.document import Document\n",
"from langchain.schema import format_document\n",
"\n",
"from permchain import Pregel, channels"
"from permchain import Pregel, PregelRead, channels\n"
]
},
{
@@ -51,7 +45,7 @@
"metadata": {},
"outputs": [],
"source": [
"from langchain.schema.runnable import RunnableLambda"
"from langchain.schema.runnable import RunnableLambda\n"
]
},
{
@@ -65,7 +59,7 @@
"\n",
"_combine_documents = RunnableLambda(\n",
" lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)\n",
").map() | (lambda x: \"\\n\\n\".join(x))"
").map() | (lambda x: \"\\n\\n\".join(x))\n"
]
},
{
@@ -78,7 +72,7 @@
"docs = [\n",
" Document(page_content=\"Harrison used to work at Kensho\"),\n",
" Document(page_content=\"Ankush worked at Facebook\"),\n",
"]"
"]\n"
]
},
{
@@ -103,7 +97,8 @@
" ]\n",
" )\n",
" | ChatOpenAI()\n",
")"
" | StrOutputParser()\n",
")\n"
]
},
{
@@ -115,7 +110,7 @@
{
"data": {
"text/plain": [
"AIMessage(content='Harrison used to work at Kensho.')"
"'Harrison used to work at Kensho.'"
]
},
"execution_count": 6,
@@ -124,7 +119,7 @@
}
],
"source": [
"stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})"
"stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})\n"
]
},
{
@@ -144,7 +139,7 @@
"metadata": {},
"outputs": [],
"source": [
"many_docs = docs * 5"
"many_docs = docs * 5\n"
]
},
{
@@ -169,7 +164,7 @@
" new_result_doc_list.append(_sub_result_docs[:-1])\n",
" _sub_result_docs = _sub_result_docs[-1:]\n",
" new_result_doc_list.append(_sub_result_docs)\n",
" return new_result_doc_list"
" return new_result_doc_list\n"
]
},
{
@@ -201,7 +196,7 @@
"source": [
"# Just to show what its like split\n",
"split_docs = _split_list_of_docs(many_docs)\n",
"split_docs"
"split_docs\n"
]
},
{
@@ -209,110 +204,135 @@
"execution_count": 10,
"id": "8d524ba6-0939-4a5d-8db0-4fa1ef06eaeb",
"metadata": {},
"outputs": [
{
"ename": "TypeError",
"evalue": "LastValue() takes no arguments",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[10], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m input_inbox \u001b[38;5;241m=\u001b[39m \u001b[43mchannels\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mLastValue\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;28;43mstr\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43minput_inbox\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m reduce_inbox \u001b[38;5;241m=\u001b[39m channels\u001b[38;5;241m.\u001b[39mLastValue[\u001b[38;5;28mstr\u001b[39m](\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mreduce_inbox\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 3\u001b[0m collapse_inbox \u001b[38;5;241m=\u001b[39m channels\u001b[38;5;241m.\u001b[39mLastValue[\u001b[38;5;28mstr\u001b[39m](\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcollapse_inbox\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
"File \u001b[0;32m/opt/homebrew/Cellar/python@3.11/3.11.5/Frameworks/Python.framework/Versions/3.11/lib/python3.11/typing.py:1268\u001b[0m, in \u001b[0;36m_BaseGenericAlias.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1265\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_inst:\n\u001b[1;32m 1266\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mType \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_name\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m cannot be instantiated; \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 1267\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124muse \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__origin__\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m() instead\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m-> 1268\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__origin__\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1269\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 1270\u001b[0m result\u001b[38;5;241m.\u001b[39m__orig_class__ \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\n",
"\u001b[0;31mTypeError\u001b[0m: LastValue() takes no arguments"
]
}
],
"outputs": [],
"source": [
"input_inbox = channels.LastValue[str](\"input_inbox\")\n",
"reduce_inbox = channels.LastValue[str](\"reduce_inbox\")\n",
"collapse_inbox = channels.LastValue[str](\"collapse_inbox\")\n",
"output_inbox = channels.LastValue[str](\"output_inbox\")"
"chans = {\n",
" # input\n",
" \"question\": channels.LastValue(str),\n",
" \"docs\": channels.Inbox(Document),\n",
" # intermediate\n",
" \"docs_to_finalize\": channels.Inbox(Document),\n",
" # output\n",
" \"answer\": channels.LastValue(str),\n",
"}\n"
]
},
{
"cell_type": "code",
"execution_count": 22,
"execution_count": 11,
"id": "67370694-86f4-4b64-9d4f-38b2e306abeb",
"metadata": {},
"outputs": [],
"source": [
"# Decide if should finish or should reduce one more step\n",
"def decide_end(plan):\n",
" if len(plan[\"docs\"]) > 1:\n",
" return Pregel.send_to(\"reduce_inbox\")\n",
"def decide(docs: list[Document]) -> Runnable:\n",
" if len(_split_list_of_docs(docs)) > 1:\n",
" # send back to the beginning if we still need to collapse more\n",
" return Pregel.write_to(\"docs\")\n",
" else:\n",
" return stuff_chain | Pregel.send_to(\"output_inbox\")\n",
" # send to the finalizer if we're ready to produce final answer\n",
" return Pregel.write_to(\"docs_to_finalize\")\n",
"\n",
"\n",
"# Chain that collapses documents then chooses end\n",
"collapse_chain = (\n",
" Pregel.subscribe_to(docs=collapse_inbox, question=\"question\")\n",
" | RunnablePassthrough.assign(docs=lambda x: _split_list_of_docs(x[\"docs\"]))\n",
" | decide_end\n",
"collapse = (\n",
" Pregel.subscribe_to(\"docs\")\n",
" | _split_list_of_docs\n",
" | {\"docs_list\": RunnablePassthrough(), \"question\": PregelRead(\"question\")}\n",
" # {docs: list[list[Doc]], question: str} -> list[{docs: list[Doc], question: str}]\n",
" | (lambda x: [{\"docs\": docs, \"question\": x[\"question\"]} for docs in x[\"docs_list\"]])\n",
" | stuff_chain.map() # Collapse each list of docs to a single string\n",
" | (lambda x: [Document(page_content=s) for s in x]) # A new (smaller) list of docs\n",
" | decide\n",
")\n",
"\n",
"\n",
"reduce_chain = (\n",
" Pregel.subscribe_to(input=input_inbox)\n",
" | (lambda x: [{\"docs\": d, \"question\": x[\"question\"]} for d in x[\"docs\"]])\n",
" | stuff_chain.map()\n",
" | Pregel.send_to(\n",
" {\n",
" \"collapse_inbox\": {\n",
" \"docs\": lambda x: [Document(page_content=m.content) for m in x],\n",
" }\n",
" }\n",
" )\n",
")"
"# Convert final set of docs to an answer\n",
"finalize = (\n",
" Pregel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n",
" | stuff_chain\n",
" | Pregel.write_to(\"answer\")\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 23,
"execution_count": 12,
"id": "3019e7d2-ab7f-4868-b43c-ad898d824a26",
"metadata": {},
"outputs": [
{
"ename": "ValidationError",
"evalue": "6 validation errors for Pregel\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mValidationError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[23], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m pubsub \u001b[38;5;241m=\u001b[39m \u001b[43mPregel\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreduce_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcollapse_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_inbox\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moutput\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_inbox\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/workplace/permchain/permchain/pregel.py:244\u001b[0m, in \u001b[0;36mPregel.__init__\u001b[0;34m(self, input, output, step_timeout, *processes, **kwargs)\u001b[0m\n\u001b[1;32m 236\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\n\u001b[1;32m 237\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 238\u001b[0m \u001b[38;5;241m*\u001b[39mprocesses: PregelInvoke \u001b[38;5;241m|\u001b[39m PregelBatch,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 242\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[1;32m 243\u001b[0m ):\n\u001b[0;32m--> 244\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 245\u001b[0m \u001b[43m \u001b[49m\u001b[43mprocesses\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprocesses\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 246\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 247\u001b[0m \u001b[43m \u001b[49m\u001b[43moutput\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 248\u001b[0m \u001b[43m \u001b[49m\u001b[43mstep_timeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstep_timeout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 249\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 250\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/.pyenv/versions/3.10.1/envs/permchain/lib/python3.10/site-packages/langchain/load/serializable.py:90\u001b[0m, in \u001b[0;36mSerializable.__init__\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m 89\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m---> 90\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 91\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_lc_kwargs \u001b[38;5;241m=\u001b[39m kwargs\n",
"File \u001b[0;32m~/.pyenv/versions/3.10.1/envs/permchain/lib/python3.10/site-packages/pydantic/main.py:341\u001b[0m, in \u001b[0;36mpydantic.main.BaseModel.__init__\u001b[0;34m()\u001b[0m\n",
"\u001b[0;31mValidationError\u001b[0m: 6 validation errors for Pregel\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 0\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 1\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)\nprocesses -> 2\n value is not a valid dict (type=type_error.dict)"
]
}
],
"outputs": [],
"source": [
"pubsub = Pregel(\n",
" input_inbox, reduce_inbox, collapse_inbox, input=input_inbox, output=output_inbox\n",
")"
"reduce_chain = Pregel(\n",
" chains={\n",
" \"collapse\": collapse,\n",
" \"finalize\": finalize,\n",
" },\n",
" channels=chans,\n",
" input=[\"question\", \"docs\"],\n",
" output=\"answer\",\n",
" debug=True,\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 101,
"execution_count": 13,
"id": "69fcb829-3dae-432a-8db3-11bbb179a7d2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 0 with 1 task. Next tasks:\n",
"\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook'),\n",
" Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook'),\n",
" Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook'),\n",
" Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook'),\n",
" Document(page_content='Harrison used to work at Kensho'),\n",
" Document(page_content='Ankush worked at Facebook')))\n",
"\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 0. Channel values:\n",
"\u001b[0m{'docs': (...), 'question': 'where did harrison work'}\n",
"\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 1 with 1 task. Next tasks:\n",
"\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.')))\n",
"\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 1. Channel values:\n",
"\u001b[0m{'docs': (...), 'question': 'where did harrison work'}\n",
"\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 2 with 1 task. Next tasks:\n",
"\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.')))\n",
"\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 2. Channel values:\n",
"\u001b[0m{'docs': (...),\n",
" 'docs_to_finalize': (...),\n",
" 'question': 'where did harrison work'}\n",
"\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 3 with 1 task. Next tasks:\n",
"\u001b[0m- finalize({'docs': (Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'))})\n",
"\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 3. Channel values:\n",
"\u001b[0m{'answer': 'Harrison used to work at Kensho.',\n",
" 'docs': (...),\n",
" 'docs_to_finalize': (...),\n",
" 'question': 'where did harrison work'}\n"
]
},
{
"data": {
"text/plain": [
"[AIMessage(content='Harrison used to work at Kensho.', additional_kwargs={}, example=False)]"
"'Harrison used to work at Kensho.'"
]
},
"execution_count": 101,
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"reduce_agent.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})"
"reduce_chain.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})\n"
]
},
{
examples/draft-revise-loop.py
+8 -6
View File
@@ -5,7 +5,7 @@ from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
from langchain.prompts import SystemMessagePromptTemplate
from langchain.schema.output_parser import StrOutputParser
from permchain import Pregel, channels
from permchain import Channels, Pregel
# prompts
@@ -74,6 +74,12 @@ reviser_chain = reviser_prompt | gpt3 | StrOutputParser()
# application
channels = {
"question": Channels.LastValue(str),
"draft": Channels.LastValue(str),
"notes": Channels.LastValue(str),
}
drafter = (
# subscribe to question channel as a dict with a single key, "question"
Pregel.subscribe_to(["question"])
@@ -102,11 +108,7 @@ reviser = (
)
draft_revise_loop = Pregel(
channels={
"question": channels.LastValue(str),
"draft": channels.LastValue(str),
"notes": channels.LastValue(str),
},
channels=channels,
chains={
"drafter": drafter,
"editor": editor,
examples/readme.py
+2 -2
View File
@@ -1,4 +1,4 @@
from permchain import Pregel, channels
from permchain import Channels, Pregel
grow_value = (
Pregel.subscribe_to("value")
@@ -8,7 +8,7 @@ grow_value = (
app = Pregel(
chains={"grow_value": grow_value},
channels={"value": channels.LastValue(str)},
channels={"value": Channels.LastValue(str)},
input="value",
output="value",
)
examples/recursive-web-loader.py
+10 -9
View File
@@ -6,7 +6,7 @@ from langchain.schema import Document
from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
from langchain.utils.html import extract_sub_links
from permchain import Pregel, channels
from permchain import Channels, Pregel
# Load url with sync httpx client
@@ -82,6 +82,14 @@ def recursive_web_loader(
# assign default extractors
extractor = extractor or (lambda x: x)
metadata_extractor = metadata_extractor or _metadata_extractor
# define the channels
channels = {
"base_url": Channels.LastValue(str),
"next_urls": Channels.UniqueInbox(str),
"documents": Channels.Stream(Document),
"visited": Channels.Set(str),
"client": Channels.ContextManager(httpx_client, httpx_aclient),
}
# the main chain that gets executed recursively
visitor = (
# while there are urls in next_urls
@@ -112,20 +120,13 @@ def recursive_web_loader(
)
)
return Pregel(
channels=channels,
chains={
# use the base_url as the first url to visit
"input": Pregel.subscribe_to("base_url") | Pregel.write_to("next_urls"),
# add the main chain
"visitor": visitor,
},
# define the channels
channels={
"base_url": channels.LastValue(str),
"next_urls": channels.UniqueInbox(str),
"documents": channels.Stream(Document),
"visited": channels.Set(str),
"client": channels.ContextManager(httpx_client, httpx_aclient),
},
# this will accept a string as input
input="base_url",
# and return a dict with documents and visited set