mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-12 12:47:53 +02:00
Split into smaller files
This commit is contained in:
Nuno Campos
@@ -39,6 +39,7 @@ Check `examples` for more examples.
|
||||
- [x] Test different input and output types (str, str sequence)
|
||||
- [x] Add tests for Stream, UniqueInbox
|
||||
- [ ] Add tests for subscribe_to_each().join()
|
||||
- [ ] Add optional debug logging
|
||||
- [ ] Implement checkpointing
|
||||
- [ ] Save checkpoints at end of each step
|
||||
- [ ] Load checkpoint at start of invocation
|
||||
|
||||
@@ -5,6 +5,7 @@ from typing import (
|
||||
AsyncGenerator,
|
||||
Generator,
|
||||
Generic,
|
||||
Mapping,
|
||||
Optional,
|
||||
Sequence,
|
||||
TypeVar,
|
||||
@@ -59,3 +60,27 @@ class Channel(Generic[Value, Update], ABC):
|
||||
@abstractmethod
|
||||
def checkpoint(self) -> str | None:
|
||||
...
|
||||
|
||||
|
||||
@contextmanager
|
||||
def ChannelsManager(
|
||||
channels: Mapping[str, Channel]
|
||||
) -> Generator[Mapping[str, Channel], None, None]:
|
||||
empty = {k: v.empty() for k, v in channels.items()}
|
||||
try:
|
||||
yield {k: v.__enter__() for k, v in empty.items()}
|
||||
finally:
|
||||
for v in empty.values():
|
||||
v.__exit__(None, None, None)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def AsyncChannelsManager(
|
||||
channels: Mapping[str, Channel]
|
||||
) -> AsyncGenerator[Mapping[str, Channel], None]:
|
||||
empty = {k: v.aempty() for k, v in channels.items()}
|
||||
try:
|
||||
yield {k: await v.__aenter__() for k, v in empty.items()}
|
||||
finally:
|
||||
for v in empty.values():
|
||||
await v.__aexit__(None, None, None)
|
||||
|
||||
@@ -4,283 +4,39 @@ import asyncio
|
||||
import concurrent.futures
|
||||
import logging
|
||||
from collections import defaultdict, deque
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
AsyncIterator,
|
||||
Callable,
|
||||
Generator,
|
||||
Iterator,
|
||||
Mapping,
|
||||
Optional,
|
||||
Sequence,
|
||||
Type,
|
||||
cast,
|
||||
)
|
||||
from typing import Any, AsyncIterator, Iterator, Mapping, Optional, Sequence, Type, cast
|
||||
|
||||
from langchain.callbacks.manager import (
|
||||
AsyncCallbackManagerForChainRun,
|
||||
CallbackManagerForChainRun,
|
||||
)
|
||||
from langchain.pydantic_v1 import BaseModel, Field, create_model
|
||||
from langchain.pydantic_v1 import BaseModel, create_model
|
||||
from langchain.schema.runnable import (
|
||||
Runnable,
|
||||
RunnableBinding,
|
||||
RunnableLambda,
|
||||
RunnablePassthrough,
|
||||
RunnableSerializable,
|
||||
)
|
||||
from langchain.schema.runnable.base import (
|
||||
Other,
|
||||
RunnableEach,
|
||||
RunnableLike,
|
||||
coerce_to_runnable,
|
||||
)
|
||||
from langchain.schema.runnable.base import RunnableLike, coerce_to_runnable
|
||||
from langchain.schema.runnable.config import (
|
||||
RunnableConfig,
|
||||
get_executor_for_config,
|
||||
patch_config,
|
||||
)
|
||||
from langchain.schema.runnable.utils import ConfigurableFieldSpec
|
||||
|
||||
from permchain.channels.base import Channel, EmptyChannelError
|
||||
from permchain.channels.base import (
|
||||
AsyncChannelsManager,
|
||||
Channel,
|
||||
ChannelsManager,
|
||||
EmptyChannelError,
|
||||
)
|
||||
from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STEP
|
||||
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__)
|
||||
|
||||
|
||||
CONFIG_KEY_STEP = "__pregel_step"
|
||||
CONFIG_KEY_SEND = "__pregel_send"
|
||||
CONFIG_KEY_READ = "__pregel_read"
|
||||
|
||||
TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
|
||||
|
||||
|
||||
class PregelRead(RunnableLambda):
|
||||
channel: str
|
||||
|
||||
@property
|
||||
def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
|
||||
return [
|
||||
ConfigurableFieldSpec(
|
||||
id=CONFIG_KEY_READ,
|
||||
name=CONFIG_KEY_READ,
|
||||
description=None,
|
||||
default=None,
|
||||
annotation=Callable[[Channel], Any],
|
||||
),
|
||||
]
|
||||
|
||||
def __init__(self, channel: str) -> None:
|
||||
super().__init__(func=self._read, afunc=self._aread) # type: ignore[arg-type]
|
||||
self.channel = channel
|
||||
|
||||
def _read(self, _: Any, config: RunnableConfig) -> Any:
|
||||
try:
|
||||
read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
|
||||
except KeyError:
|
||||
raise RuntimeError(
|
||||
f"Runnable {self} is not configured with a read function"
|
||||
"Make sure to call in the context of a Pregel process"
|
||||
)
|
||||
return read(self.channel)
|
||||
|
||||
async def _aread(self, _: Any, config: RunnableConfig) -> Any:
|
||||
try:
|
||||
read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
|
||||
except KeyError:
|
||||
raise RuntimeError(
|
||||
f"Runnable {self} is not configured with a read function"
|
||||
"Make sure to call in the context of a Pregel process"
|
||||
)
|
||||
return read(self.channel)
|
||||
|
||||
|
||||
class PregelInvoke(RunnableBinding):
|
||||
channels: Mapping[None, str] | Mapping[str, str]
|
||||
|
||||
bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
|
||||
|
||||
kwargs: Mapping[str, Any] = Field(default_factory=dict)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: Mapping[None, str] | Mapping[str, str],
|
||||
*,
|
||||
bound: Optional[Runnable[Any, Any]] = None,
|
||||
kwargs: Optional[Mapping[str, Any]] = None,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
**other_kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
channels=channels,
|
||||
bound=bound or RunnablePassthrough(),
|
||||
kwargs=kwargs or {},
|
||||
config=config,
|
||||
**other_kwargs,
|
||||
)
|
||||
|
||||
def __or__(
|
||||
self,
|
||||
other: Runnable[Any, Other]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
|
||||
) -> PregelInvoke:
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return PregelInvoke(channels=self.channels, bound=coerce_to_runnable(other))
|
||||
else:
|
||||
# delegate to __or__ in self.bound
|
||||
return PregelInvoke(channels=self.channels, bound=self.bound | other)
|
||||
|
||||
def __ror__(
|
||||
self,
|
||||
other: Runnable[Other, Any]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
|
||||
) -> Runnable:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class PregelBatch(RunnableEach):
|
||||
channel: str
|
||||
|
||||
key: Optional[str]
|
||||
|
||||
bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
|
||||
|
||||
def join(self, channels: Sequence[str]) -> PregelBatch:
|
||||
if self.key is None:
|
||||
raise ValueError(
|
||||
"Cannot join() additional channels without a key."
|
||||
" Pass a key arg to Pregel.subscribe_to_each()."
|
||||
)
|
||||
|
||||
joiner = RunnablePassthrough.assign(
|
||||
**{chan: PregelRead(chan) for chan in channels}
|
||||
)
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return PregelBatch(channel=self.channel, key=self.key, bound=joiner)
|
||||
else:
|
||||
return PregelBatch(
|
||||
channel=self.channel, key=self.key, bound=self.bound | joiner
|
||||
)
|
||||
|
||||
def __or__( # type: ignore[override]
|
||||
self,
|
||||
other: Runnable[Any, Other]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
|
||||
) -> PregelBatch:
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return PregelBatch(
|
||||
channel=self.channel, key=self.key, bound=coerce_to_runnable(other)
|
||||
)
|
||||
else:
|
||||
# delegate to __or__ in self.bound
|
||||
return PregelBatch(
|
||||
channel=self.channel, key=self.key, bound=self.bound | other
|
||||
)
|
||||
|
||||
def __ror__(
|
||||
self,
|
||||
other: Runnable[Other, Any]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
|
||||
) -> Runnable:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class PregelSink(RunnableLambda):
|
||||
channels: Sequence[tuple[str, Runnable]]
|
||||
"""
|
||||
Mapping of write channels to Runnables that return the value to be written,
|
||||
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,
|
||||
description=None,
|
||||
default=None,
|
||||
annotation=TYPE_SEND,
|
||||
),
|
||||
]
|
||||
|
||||
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]
|
||||
|
||||
write([(chan, val) for chan, val in values if val is not None])
|
||||
|
||||
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]
|
||||
|
||||
write([(chan, val) for chan, val in values if val is not None])
|
||||
|
||||
return input
|
||||
|
||||
|
||||
@contextmanager
|
||||
def ChannelsManager(
|
||||
channels: Mapping[str, Channel]
|
||||
) -> Generator[Mapping[str, Channel], None, None]:
|
||||
empty = {k: v.empty() for k, v in channels.items()}
|
||||
try:
|
||||
yield {k: v.__enter__() for k, v in empty.items()}
|
||||
finally:
|
||||
for v in empty.values():
|
||||
v.__exit__(None, None, None)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def AsyncChannelsManager(
|
||||
channels: Mapping[str, Channel]
|
||||
) -> AsyncGenerator[Mapping[str, Channel], None]:
|
||||
empty = {k: v.aempty() for k, v in channels.items()}
|
||||
try:
|
||||
yield {k: await v.__aenter__() for k, v in empty.items()}
|
||||
finally:
|
||||
for v in empty.values():
|
||||
await v.__aexit__(None, None, None)
|
||||
|
||||
|
||||
class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
|
||||
channels: Mapping[str, Channel]
|
||||
|
||||
@@ -428,7 +184,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
|
||||
# collect all writes to channels, without applying them yet
|
||||
pending_writes = deque[tuple[str, Any]]()
|
||||
|
||||
# execute tasks, and wait for one to fail or all to finish
|
||||
# execute tasks, and wait for one to fail or all to finish.
|
||||
# each task is independent from all other concurrent tasks
|
||||
done, inflight = concurrent.futures.wait(
|
||||
(
|
||||
@@ -528,7 +284,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
|
||||
# collect all writes to channels, without applying them yet
|
||||
pending_writes = deque[tuple[str, Any]]()
|
||||
|
||||
# execute tasks, and wait for one to fail or all to finish
|
||||
# execute tasks, and wait for one to fail or all to finish.
|
||||
# each task is independent from all other concurrent tasks
|
||||
done, inflight = await asyncio.wait(
|
||||
[
|
||||
@@ -704,43 +460,3 @@ def _apply_writes_and_prepare_next_tasks(
|
||||
tasks.append((proc, val))
|
||||
|
||||
return tasks
|
||||
|
||||
|
||||
def validate_chains_channels(
|
||||
chains: Sequence[PregelInvoke | PregelBatch],
|
||||
channels: Mapping[str, Channel],
|
||||
input: str | Sequence[str],
|
||||
output: str | Sequence[str],
|
||||
) -> None:
|
||||
subscribed_channels = set[str]()
|
||||
for chain in chains:
|
||||
if isinstance(chain, PregelInvoke):
|
||||
subscribed_channels.update(chain.channels.values())
|
||||
elif isinstance(chain, PregelBatch):
|
||||
subscribed_channels.add(chain.channel)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Invalid chain type {type(chain)}, expected Pregel.subscribe_to() or Pregel.subscribe_to_each()"
|
||||
)
|
||||
|
||||
for chan in subscribed_channels:
|
||||
if chan not in channels:
|
||||
raise ValueError(f"Channel {chan} is subscribed to, but not initialized")
|
||||
|
||||
if isinstance(input, str):
|
||||
if input not in subscribed_channels:
|
||||
raise ValueError(f"Input channel {input} is not subscribed to by any chain")
|
||||
else:
|
||||
for chan in input:
|
||||
if chan not in subscribed_channels:
|
||||
raise ValueError(
|
||||
f"Input channel {chan} is not subscribed to by any chain"
|
||||
)
|
||||
|
||||
if isinstance(output, str):
|
||||
if output not in channels:
|
||||
raise ValueError(f"Output channel {output} is not initialized")
|
||||
else:
|
||||
for chan in output:
|
||||
if chan not in channels:
|
||||
raise ValueError(f"Output channel {chan} is not initialized")
|
||||
@@ -0,0 +1,3 @@
|
||||
CONFIG_KEY_STEP = "__pregel_step"
|
||||
CONFIG_KEY_SEND = "__pregel_send"
|
||||
CONFIG_KEY_READ = "__pregel_read"
|
||||
@@ -0,0 +1,152 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable, Mapping, Optional, Sequence
|
||||
|
||||
from langchain.pydantic_v1 import Field
|
||||
from langchain.schema.runnable import (
|
||||
Runnable,
|
||||
RunnableBinding,
|
||||
RunnableConfig,
|
||||
RunnableLambda,
|
||||
RunnablePassthrough,
|
||||
)
|
||||
from langchain.schema.runnable.base import Other, RunnableEach, coerce_to_runnable
|
||||
from langchain.schema.runnable.utils import ConfigurableFieldSpec
|
||||
|
||||
from permchain.channels.base import Channel
|
||||
from permchain.pregel.constants import CONFIG_KEY_READ
|
||||
|
||||
|
||||
class PregelRead(RunnableLambda):
|
||||
channel: str
|
||||
|
||||
@property
|
||||
def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
|
||||
return [
|
||||
ConfigurableFieldSpec(
|
||||
id=CONFIG_KEY_READ,
|
||||
name=CONFIG_KEY_READ,
|
||||
description=None,
|
||||
default=None,
|
||||
annotation=Callable[[Channel], Any],
|
||||
),
|
||||
]
|
||||
|
||||
def __init__(self, channel: str) -> None:
|
||||
# TODO remove type ignore after updating langchain
|
||||
super().__init__(func=self._read, afunc=self._aread) # type: ignore[arg-type]
|
||||
self.channel = channel
|
||||
|
||||
def _read(self, _: Any, config: RunnableConfig) -> Any:
|
||||
try:
|
||||
read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
|
||||
except KeyError:
|
||||
raise RuntimeError(
|
||||
f"Runnable {self} is not configured with a read function"
|
||||
"Make sure to call in the context of a Pregel process"
|
||||
)
|
||||
return read(self.channel)
|
||||
|
||||
async def _aread(self, _: Any, config: RunnableConfig) -> Any:
|
||||
try:
|
||||
read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
|
||||
except KeyError:
|
||||
raise RuntimeError(
|
||||
f"Runnable {self} is not configured with a read function"
|
||||
"Make sure to call in the context of a Pregel process"
|
||||
)
|
||||
return read(self.channel)
|
||||
|
||||
|
||||
class PregelInvoke(RunnableBinding):
|
||||
channels: Mapping[None, str] | Mapping[str, str]
|
||||
|
||||
bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
|
||||
|
||||
kwargs: Mapping[str, Any] = Field(default_factory=dict)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: Mapping[None, str] | Mapping[str, str],
|
||||
*,
|
||||
bound: Optional[Runnable[Any, Any]] = None,
|
||||
kwargs: Optional[Mapping[str, Any]] = None,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
**other_kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
channels=channels,
|
||||
bound=bound or RunnablePassthrough(),
|
||||
kwargs=kwargs or {},
|
||||
config=config,
|
||||
**other_kwargs,
|
||||
)
|
||||
|
||||
def __or__(
|
||||
self,
|
||||
other: Runnable[Any, Other]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
|
||||
) -> PregelInvoke:
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return PregelInvoke(channels=self.channels, bound=coerce_to_runnable(other))
|
||||
else:
|
||||
# delegate to __or__ in self.bound
|
||||
return PregelInvoke(channels=self.channels, bound=self.bound | other)
|
||||
|
||||
def __ror__(
|
||||
self,
|
||||
other: Runnable[Other, Any]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
|
||||
) -> Runnable:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class PregelBatch(RunnableEach):
|
||||
channel: str
|
||||
|
||||
key: Optional[str]
|
||||
|
||||
bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
|
||||
|
||||
def join(self, channels: Sequence[str]) -> PregelBatch:
|
||||
if self.key is None:
|
||||
raise ValueError(
|
||||
"Cannot join() additional channels without a key."
|
||||
" Pass a key arg to Pregel.subscribe_to_each()."
|
||||
)
|
||||
|
||||
joiner = RunnablePassthrough.assign(
|
||||
**{chan: PregelRead(chan) for chan in channels}
|
||||
)
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return PregelBatch(channel=self.channel, key=self.key, bound=joiner)
|
||||
else:
|
||||
return PregelBatch(
|
||||
channel=self.channel, key=self.key, bound=self.bound | joiner
|
||||
)
|
||||
|
||||
def __or__( # type: ignore[override]
|
||||
self,
|
||||
other: Runnable[Any, Other]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
|
||||
) -> PregelBatch:
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return PregelBatch(
|
||||
channel=self.channel, key=self.key, bound=coerce_to_runnable(other)
|
||||
)
|
||||
else:
|
||||
# delegate to __or__ in self.bound
|
||||
return PregelBatch(
|
||||
channel=self.channel, key=self.key, bound=self.bound | other
|
||||
)
|
||||
|
||||
def __ror__(
|
||||
self,
|
||||
other: Runnable[Other, Any]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
|
||||
) -> Runnable:
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,44 @@
|
||||
from typing import Mapping, Sequence
|
||||
|
||||
from permchain.channels.base import Channel
|
||||
from permchain.pregel.read import PregelBatch, PregelInvoke
|
||||
|
||||
|
||||
def validate_chains_channels(
|
||||
chains: Sequence[PregelInvoke | PregelBatch],
|
||||
channels: Mapping[str, Channel],
|
||||
input: str | Sequence[str],
|
||||
output: str | Sequence[str],
|
||||
) -> None:
|
||||
subscribed_channels = set[str]()
|
||||
for chain in chains:
|
||||
if isinstance(chain, PregelInvoke):
|
||||
subscribed_channels.update(chain.channels.values())
|
||||
elif isinstance(chain, PregelBatch):
|
||||
subscribed_channels.add(chain.channel)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Invalid chain type {type(chain)}, expected Pregel.subscribe_to() or Pregel.subscribe_to_each()"
|
||||
)
|
||||
|
||||
for chan in subscribed_channels:
|
||||
if chan not in channels:
|
||||
raise ValueError(f"Channel {chan} is subscribed to, but not initialized")
|
||||
|
||||
if isinstance(input, str):
|
||||
if input not in subscribed_channels:
|
||||
raise ValueError(f"Input channel {input} is not subscribed to by any chain")
|
||||
else:
|
||||
for chan in input:
|
||||
if chan not in subscribed_channels:
|
||||
raise ValueError(
|
||||
f"Input channel {chan} is not subscribed to by any chain"
|
||||
)
|
||||
|
||||
if isinstance(output, str):
|
||||
if output not in channels:
|
||||
raise ValueError(f"Output channel {output} is not initialized")
|
||||
else:
|
||||
for chan in output:
|
||||
if chan not in channels:
|
||||
raise ValueError(f"Output channel {chan} is not initialized")
|
||||
@@ -0,0 +1,82 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable, Optional, Sequence
|
||||
|
||||
from langchain.schema.runnable import (
|
||||
Runnable,
|
||||
RunnableConfig,
|
||||
RunnableLambda,
|
||||
)
|
||||
from langchain.schema.runnable.utils import ConfigurableFieldSpec
|
||||
|
||||
from permchain.pregel.constants import CONFIG_KEY_SEND, CONFIG_KEY_STEP
|
||||
|
||||
TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
|
||||
|
||||
|
||||
# TODO switch to RunnablePassthrough after updating langchain
|
||||
class PregelSink(RunnableLambda):
|
||||
channels: Sequence[tuple[str, Runnable]]
|
||||
"""
|
||||
Mapping of write channels to Runnables that return the value to be written,
|
||||
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,
|
||||
description=None,
|
||||
default=None,
|
||||
annotation=TYPE_SEND,
|
||||
),
|
||||
]
|
||||
|
||||
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]
|
||||
|
||||
write([(chan, val) for chan, val in values if val is not None])
|
||||
|
||||
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]
|
||||
|
||||
write([(chan, val) for chan, val in values if val is not None])
|
||||
|
||||
return input
|
||||
Reference in New Issue
Block a user