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

Remove "main" chain, Improve debug logging, Remove max_steps from send_to, Add retries to recursive web loader

Browse Source
This commit is contained in:
Nuno Campos
2023-10-23 10:26:10 +01:00
parent 9dfffe5fc0
commit 274600436b
11 changed files with 176 additions and 232 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -16,7 +16,7 @@ grow_value = (
)
app = Pregel(
grow_value,
chains={"grow_value": grow_value},
channels={"value": channels.LastValue(str)},
input="value",
output="value",
examples/draft-revise-loop.py
+39 -24
View File
@@ -7,7 +7,6 @@ from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough
from permchain import Pregel, channels
from permchain.pregel import PregelIO
# prompts
@@ -81,32 +80,46 @@ reviser_chain = reviser_prompt | gpt3 | StrOutputParser()
# application
drafter = (
# subscribe to question channel as a dict with a single key, "question"
Pregel.subscribe_to(["question"])
| drafter_chain
| Pregel.send_to("draft")
)
editor = (
# subscribe to draft channel as a dict with a single key, "draft"
Pregel.subscribe_to(["draft"])
| editor_chain
| Pregel.send_to(
# send to "notes" channel if the editor does not accept the draft
notes=lambda x: x["arguments"]["notes"]
if x["name"] == "revise"
else None
)
)
reviser = (
# subscribe to new values of "notes" channel,
# and join them with the input value (question) and "draft"
Pregel.subscribe_to(["notes"]).join(["question", "draft"])
| reviser_chain
| Pregel.send_to("draft")
)
draft_revise_loop = Pregel(
drafter_chain | Pregel.send_to("draft"),
chains={
"editor": (
# subscribe to draft channel as a dict with a single key, "draft"
Pregel.subscribe_to(["draft"])
| editor_chain
| Pregel.send_to(
# send to "notes" channel if the editor does not accept the draft
notes=lambda x: x["arguments"]["notes"]
if x["name"] == "revise"
else None
)
),
"reviser": (
# subscribe to new values of "notes" channel,
# and join them with the input value (question) and "draft"
Pregel.subscribe_to(["notes"]).join([PregelIO.IN, "draft"])
| reviser_chain
| Pregel.send_to("draft")
),
},
channels={
"question": channels.LastValue(str),
"draft": channels.LastValue(str),
"notes": channels.LastValue(str),
},
chains={
"drafter": drafter,
"editor": editor,
"reviser": reviser,
},
# input will be a dict with a single key, "question"
input=["question"],
# output will be a dict with keys "draft" and "notes"
output=["draft", "notes"],
# debug logging
@@ -115,13 +128,15 @@ draft_revise_loop = Pregel(
# run
for draft in draft_revise_loop.stream("What food do turtles eat?"):
for draft in draft_revise_loop.stream({"question": "What food do turtles eat?"}):
print(draft)
print("---")
async def main():
async for draft in draft_revise_loop.astream("What food do turtles eat?"):
async for draft in draft_revise_loop.astream(
{"question": "What food do turtles eat?"}
):
print(draft)
print("---")
examples/recursive-web-loader.py
+23 -5
View File
@@ -1,4 +1,5 @@
from typing import Callable, FrozenSet, Optional, TypedDict
from contextlib import asynccontextmanager, contextmanager
from typing import AsyncGenerator, Callable, FrozenSet, Generator, Optional, TypedDict
import httpx
from langchain.schema import Document
@@ -10,6 +11,14 @@ from permchain import Pregel, channels
# Load url with sync httpx client
@contextmanager
def httpx_client() -> Generator[httpx.Client, None, None]:
with httpx.HTTPTransport(retries=3) as transport, httpx.Client(
transport=transport
) as client:
yield client
class LoadUrlInput(TypedDict):
url: str
visited: FrozenSet[str]
@@ -24,6 +33,14 @@ def load_url(input: LoadUrlInput) -> str:
# Same as above but with async httpx client
@asynccontextmanager
async def httpx_aclient() -> AsyncGenerator[httpx.AsyncClient, None]:
async with httpx.AsyncHTTPTransport(retries=3) as transport, httpx.AsyncClient(
transport=transport
) as client:
yield client
class LoadUrlInputAsync(TypedDict):
url: str
visited: FrozenSet[str]
@@ -92,7 +109,6 @@ def recursive_web_loader(
)
if url not in x["visited"] and url != x["url"]
],
_max_steps=max_depth,
)
)
return Pregel(
@@ -108,17 +124,19 @@ def recursive_web_loader(
"next_urls": channels.UniqueInbox(str),
"documents": channels.Stream(Document),
"visited": channels.Set(str),
"client": channels.ContextManager(httpx.Client, httpx.AsyncClient),
"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
output=["documents", "visited"],
)
# debug logging
debug=True,
).with_config({"recursion_limit": max_depth + 1})
loader = recursive_web_loader(max_depth=3)
documents = loader.invoke("https://docs.python.org/3.9/")
print(documents)
print(len(documents["documents"]))
permchain/pregel/__init__.py
+57 -141
View File
@@ -2,8 +2,6 @@ from __future__ import annotations
import asyncio
import concurrent.futures
import enum
import logging
from collections import defaultdict, deque
from typing import Any, AsyncIterator, Iterator, Mapping, Optional, Sequence, Type, cast
@@ -12,7 +10,7 @@ from langchain.callbacks.manager import (
CallbackManagerForChainRun,
)
from langchain.globals import get_debug
from langchain.pydantic_v1 import BaseModel, create_model
from langchain.pydantic_v1 import BaseModel, Field, create_model, root_validator
from langchain.schema.runnable import (
Runnable,
RunnablePassthrough,
@@ -31,34 +29,14 @@ from permchain.channels.base import (
ChannelsManager,
EmptyChannelError,
)
from permchain.channels.last_value import LastValue
from permchain.pregel.constants import (
CHAINS_MAIN,
CONFIG_KEY_READ,
CONFIG_KEY_SEND,
CONFIG_KEY_STEP,
)
from permchain.pregel.debug import print_step_start
from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND
from permchain.pregel.debug import print_checkpoint, print_step_start
from permchain.pregel.io import map_input, map_output
from permchain.pregel.log import logger
from permchain.pregel.read import PregelBatch, PregelInvoke
from permchain.pregel.validate import validate_chains_channels
from permchain.pregel.write import PregelSink
logger = logging.getLogger(__name__)
# Before Python 3.11 native StrEnum is not available
class StrEnum(str, enum.Enum):
"""A string enum."""
pass
class PregelIO(StrEnum):
"""Pregel IO channels."""
IN = "__pregel_input"
OUT = "__pregel_output"
class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
channels: Mapping[str, Channel]
@@ -71,53 +49,17 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
step_timeout: Optional[float] = None
debug: bool
debug: bool = Field(default_factory=get_debug)
class Config:
arbitrary_types_allowed = True
def __init__(
self,
main: Optional[Runnable] = None,
*,
chains: Mapping[str, PregelInvoke | PregelBatch],
channels: Mapping[str, Channel],
output: str | Sequence[str] = PregelIO.OUT,
input: str | Sequence[str] = PregelIO.IN,
step_timeout: Optional[float] = None,
debug: Optional[bool] = None,
) -> None:
chains = {**chains}
channels = {**channels}
if main is not None:
chains[CHAINS_MAIN] = (
Pregel.subscribe_to(PregelIO.IN) | main | Pregel.send_to(PregelIO.OUT)
)
elif output is PregelIO.OUT:
raise ValueError(
f"When no main runnable is provided, output must be one or more of the channels in {channels.keys()}"
)
if input is PregelIO.IN:
channels[PregelIO.IN] = LastValue(
main.input_schema if main is not None else Any # type: ignore[arg-type]
)
if output is PregelIO.OUT:
channels[PregelIO.OUT] = LastValue(
main.output_schema if main is not None else Any # type: ignore[arg-type]
)
validate_chains_channels(chains, channels, input, output)
super().__init__(
chains=chains,
channels=channels,
output=output,
input=input,
step_timeout=step_timeout,
debug=debug if debug is not None else get_debug(),
@root_validator(skip_on_failure=True)
def validate_pregel(cls, values: dict[str, Any]) -> dict[str, Any]:
validate_chains_channels(
values["chains"], values["channels"], values["input"], values["output"]
)
return values
@property
def InputType(self) -> Any:
@@ -176,7 +118,6 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
def send_to(
cls,
*channels: str,
_max_steps: Optional[int] = None,
**kwargs: RunnableLike,
) -> PregelSink:
"""Writes to channels the result of the lambda, or None to skip writing."""
@@ -184,8 +125,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
channels=(
[(c, RunnablePassthrough()) for c in channels]
+ [(k, coerce_to_runnable(v)) for k, v in kwargs.items()]
),
max_steps=_max_steps,
)
)
def _transform(
@@ -202,14 +142,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
next_tasks = _apply_writes_and_prepare_next_tasks(
processes,
channels,
deque((self.input, chunk) for chunk in input)
if isinstance(self.input, str)
else deque(
(k, v)
for chunk in input
for k, v in chunk.items()
if k in self.input
),
deque(w for c in input for w in map_input(self.input, c)),
)
def read(chan: str) -> Any:
@@ -244,7 +177,6 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_READ: read,
CONFIG_KEY_STEP: step,
},
),
)
@@ -254,36 +186,20 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
timeout=self.step_timeout,
)
while done:
# if any task failed
if exc := done.pop().exception():
# cancel all pending tasks
while inflight:
inflight.pop().cancel()
# raise the exception
raise exc
# TODO this is where retry of an entire step would happen
if inflight:
# if we got here means we timed out
while inflight:
# cancel all pending tasks
inflight.pop().cancel()
# raise timeout error
raise TimeoutError(f"Timed out at step {step}")
# interrupt on failure or timeout
_interrupt_or_proceed(done, inflight, step)
# apply writes to channels, decide on next step
next_tasks = _apply_writes_and_prepare_next_tasks(
processes, channels, pending_writes
)
# if any write to output channel in this step, yield current value
if isinstance(self.output, str):
if any(chan == self.output for chan, _ in pending_writes):
yield channels[self.output].get()
else:
if updated := {c for c, _ in pending_writes if c in self.output}:
yield {chan: channels[chan].get() for chan in updated}
if self.debug:
print_checkpoint(step, channels)
# if any write to output channels in this step, yield current value
for output in map_output(self.output, pending_writes, channels):
yield output
# TODO this is where we'd save checkpoint
@@ -303,16 +219,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
next_tasks = _apply_writes_and_prepare_next_tasks(
processes,
channels,
deque([(self.input, chunk) async for chunk in input])
if isinstance(self.input, str)
else deque(
[
(k, v)
async for chunk in input
for k, v in chunk.items()
if k in self.input
]
),
deque([w async for c in input for w in map_input(self.input, c)]),
)
def read(chan: str) -> Any:
@@ -349,7 +256,6 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_READ: read,
CONFIG_KEY_STEP: step,
},
),
)
@@ -360,36 +266,22 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
timeout=self.step_timeout,
)
while done:
# if any task failed
if exc := done.pop().exception():
# cancel all pending tasks
while inflight:
inflight.pop().cancel()
# raise the exception
raise exc
# TODO this is where retry of an entire step would happen
if inflight:
# if we got here means we timed out
while inflight:
# cancel all pending tasks
inflight.pop().cancel()
# raise timeout error
raise TimeoutError(f"Timed out at step {step}")
# interrupt on failure or timeout
_interrupt_or_proceed(done, inflight, step)
# apply writes to channels, decide on next step
next_tasks = _apply_writes_and_prepare_next_tasks(
processes, channels, pending_writes
)
# if any write to output channel in this step, yield current value
if isinstance(self.output, str):
if any(chan == self.output for chan, _ in pending_writes):
yield channels[self.output].get()
else:
if updated := {c for c, _ in pending_writes if c in self.output}:
yield {chan: channels[chan].get() for chan in updated}
if self.debug:
print_checkpoint(step, channels)
# if any write to output channels in this step, yield current value
for output in map_output(self.output, pending_writes, channels):
yield output
# TODO this is where we'd save checkpoint
# if no more tasks, we're done
if not next_tasks:
@@ -459,6 +351,30 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
yield chunk
def _interrupt_or_proceed(
done: set[concurrent.futures.Future[Any]] | set[asyncio.Task[Any]],
inflight: set[concurrent.futures.Future[Any]] | set[asyncio.Task[Any]],
step: int,
) -> None:
while done:
# if any task failed
if exc := done.pop().exception():
# cancel all pending tasks
while inflight:
inflight.pop().cancel()
# raise the exception
raise exc
# TODO this is where retry of an entire step would happen
if inflight:
# if we got here means we timed out
while inflight:
# cancel all pending tasks
inflight.pop().cancel()
# raise timeout error
raise TimeoutError(f"Timed out at step {step}")
def _apply_writes_and_prepare_next_tasks(
processes: Mapping[str, PregelInvoke | PregelBatch],
channels: Mapping[str, Channel],
@@ -475,7 +391,7 @@ def _apply_writes_and_prepare_next_tasks(
if chan in channels:
channels[chan].update(vals)
updated_channels.add(chan)
elif chan != PregelIO.OUT:
else:
logger.warning(f"Skipping write for channel {chan} which has no readers")
tasks: list[tuple[Runnable, Any, str]] = []
permchain/pregel/constants.py
-3
View File
@@ -1,5 +1,2 @@
CHAINS_MAIN = "__pregel_main"
CONFIG_KEY_STEP = "__pregel_step"
CONFIG_KEY_SEND = "__pregel_send"
CONFIG_KEY_READ = "__pregel_read"
permchain/pregel/debug.py
+15 -9
View File
@@ -1,10 +1,10 @@
from pprint import pformat
from typing import Any
from typing import Any, Iterator, Mapping
from langchain.schema.runnable import Runnable
from langchain.utils.input import get_bolded_text, get_colored_text
from permchain.pregel.constants import CHAINS_MAIN
from permchain.channels.base import Channel, EmptyChannelError
def print_step_start(step: int, next_tasks: list[tuple[Runnable, Any, str]]) -> None:
@@ -14,15 +14,21 @@ def print_step_start(step: int, next_tasks: list[tuple[Runnable, Any, str]]) ->
+ get_bolded_text(
f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Next tasks:\n"
)
+ pformat(
{name if name != CHAINS_MAIN else _main: val for _, val, name in next_tasks}
)
+ "\n".join(f"- {name}({pformat(val)})" for _, val, name in next_tasks)
)
class MainChainStr:
def __repr__(cls) -> str:
return "main"
def print_checkpoint(step: int, channels: Mapping[str, Channel]) -> None:
print(
f"{get_colored_text('[pregel/checkpoint]', color='blue')} "
+ get_bolded_text(f"Finishing step {step}. Channel values:\n")
+ pformat({name: val for name, val in _read_channels(channels)}, depth=1)
)
_main = MainChainStr()
def _read_channels(channels: Mapping[str, Channel]) -> Iterator[tuple[str, Any]]:
for name, channel in channels.items():
try:
yield (name, channel.get())
except EmptyChannelError:
pass
permchain/pregel/io.py
+34
View File
@@ -0,0 +1,34 @@
from typing import Any, Iterator, Mapping, Sequence
from permchain.channels.base import Channel
from permchain.pregel.log import logger
def map_input(
input_channels: str | Sequence[str], chunk: dict[str, Any] | Any
) -> Iterator[tuple[str, Any]]:
"""Map input chunk to a sequence of pending writes in the form (channel, value)."""
if isinstance(input_channels, str):
yield (input_channels, chunk)
else:
if not isinstance(chunk, dict):
raise TypeError(f"Expected chunk to be a dict, got {type(chunk).__name__}")
for k in chunk:
if k in input_channels:
yield (k, chunk[k])
else:
logger.warning(f"Input channel {k} not found in {input_channels}")
def map_output(
output_channels: str | Sequence[str],
pending_writes: Sequence[tuple[str, Any]],
channels: Mapping[str, Channel],
) -> Iterator[dict[str, Any] | Any]:
"""Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
if isinstance(output_channels, str):
if any(chan == output_channels for chan, _ in pending_writes):
yield channels[output_channels].get()
else:
if updated := {c for c, _ in pending_writes if c in output_channels}:
yield {chan: channels[chan].get() for chan in updated}
permchain/pregel/log.py
+3
View File
@@ -0,0 +1,3 @@
import logging
logger = logging.getLogger(__name__)
permchain/pregel/write.py
+2 -23
View File
@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import Any, Callable, Optional, Sequence
from typing import Any, Callable, Sequence
from langchain.schema.runnable import (
Runnable,
@@ -9,7 +9,7 @@ from langchain.schema.runnable import (
)
from langchain.schema.runnable.utils import ConfigurableFieldSpec
from permchain.pregel.constants import CONFIG_KEY_SEND, CONFIG_KEY_STEP
from permchain.pregel.constants import CONFIG_KEY_SEND
TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
@@ -22,28 +22,17 @@ class PregelSink(RunnableLambda):
or None to skip writing.
"""
max_steps: Optional[int]
def __init__(
self,
*,
channels: Sequence[tuple[str, Runnable]],
max_steps: Optional[int] = None,
):
super().__init__(func=self._write, afunc=self._awrite) # type: ignore[arg-type]
self.channels = channels
self.max_steps = max_steps
@property
def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
return [
ConfigurableFieldSpec(
id=CONFIG_KEY_STEP,
name=CONFIG_KEY_STEP,
description=None,
default=None,
annotation=int,
),
ConfigurableFieldSpec(
id=CONFIG_KEY_SEND,
name=CONFIG_KEY_SEND,
@@ -54,11 +43,6 @@ class PregelSink(RunnableLambda):
]
def _write(self, input: Any, config: RunnableConfig) -> None:
step: int = config["configurable"][CONFIG_KEY_STEP]
if self.max_steps is not None and step >= self.max_steps:
return
write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
values = [(chan, r.invoke(input, config)) for chan, r in self.channels]
@@ -68,11 +52,6 @@ class PregelSink(RunnableLambda):
return input
async def _awrite(self, input: Any, config: RunnableConfig) -> None:
step: int = config["configurable"][CONFIG_KEY_STEP]
if self.max_steps is not None and step >= self.max_steps:
return
write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
values = [(chan, await r.ainvoke(input, config)) for chan, r in self.channels]
tests/test_pregel.py
+1 -13
View File
@@ -4,7 +4,7 @@ from contextlib import contextmanager
from typing import Generator
import pytest
from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
from langchain.schema.runnable import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Pregel, channels
@@ -31,18 +31,6 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
assert app.invoke(2) == 3
def test_invoke_single_process_in_out_implicit(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = RunnableLambda(add_one)
app = Pregel(chain, channels={}, chains={})
# TODO add back when langchain is updated with with_types()
# assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
# assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
assert app.invoke(2) == 3
def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
tests/test_pregel_async.py
+1 -13
View File
@@ -3,7 +3,7 @@ from contextlib import asynccontextmanager, contextmanager
from typing import Any, AsyncGenerator, AsyncIterator, Generator
import pytest
from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
from langchain.schema.runnable import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Pregel, channels
@@ -28,18 +28,6 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
assert await app.ainvoke(2) == 3
async def test_invoke_single_process_in_out_implicit(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = RunnableLambda(add_one)
app = Pregel(chain, channels={}, chains={})
# TODO add back when langchain is updated with with_types()
# assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
# assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
assert await app.ainvoke(2) == 3
async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")