mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-11 20:27:54 +02:00
It works!
This commit is contained in:
Nuno Campos
+17
-16
@@ -5,9 +5,7 @@ from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
|
||||
from langchain.prompts import SystemMessagePromptTemplate
|
||||
from langchain.schema.output_parser import StrOutputParser
|
||||
|
||||
from permchain.channels import LastValue
|
||||
from permchain.pregel import Pregel
|
||||
|
||||
from permchain import Pregel, channels
|
||||
|
||||
# prompts
|
||||
|
||||
@@ -76,31 +74,32 @@ editor_chain = (
|
||||
|
||||
# state
|
||||
|
||||
question = LastValue[str]()
|
||||
question = channels.LastValue[str]("question")
|
||||
|
||||
draft = LastValue[str]()
|
||||
draft = channels.LastValue[str]("draft")
|
||||
|
||||
notes = LastValue[str]()
|
||||
notes = channels.LastValue[str]("notes")
|
||||
|
||||
# application
|
||||
|
||||
drafter_node = Pregel.read(question=question) | drafter_chain | Pregel.write(draft)
|
||||
|
||||
reviser_node = (
|
||||
Pregel.read(question=question, notes=notes, draft=draft)
|
||||
| reviser_chain
|
||||
| Pregel.write(draft)
|
||||
drafter_node = (
|
||||
Pregel.subscribe_to(question=question) | drafter_chain | Pregel.send_to(draft)
|
||||
)
|
||||
|
||||
|
||||
editor_node = (
|
||||
Pregel.read(draft=draft)
|
||||
Pregel.subscribe_to(draft=draft)
|
||||
| editor_chain
|
||||
| Pregel.write(
|
||||
| Pregel.send_to(
|
||||
{notes: lambda x: x["arguments"]["notes"] if x["name"] == "revise" else None}
|
||||
)
|
||||
)
|
||||
|
||||
reviser_node = (
|
||||
Pregel.subscribe_to(notes=notes).join(question=question, draft=draft)
|
||||
| reviser_chain
|
||||
| Pregel.send_to(draft)
|
||||
)
|
||||
|
||||
draft_revise_loop = Pregel(
|
||||
(drafter_node, reviser_node, editor_node),
|
||||
input=question,
|
||||
@@ -109,4 +108,6 @@ draft_revise_loop = Pregel(
|
||||
|
||||
# run
|
||||
|
||||
article = draft_revise_loop.invoke("What food do turtles eat?")
|
||||
for draft in draft_revise_loop.stream("What food do turtles eat?"):
|
||||
print('Draft: "' + draft + '"')
|
||||
print("---")
|
||||
|
||||
@@ -1,9 +1,4 @@
|
||||
from permchain.connection_inmemory import InMemoryPubSubConnection
|
||||
from permchain.pubsub import PubSub
|
||||
from permchain.topic import Topic
|
||||
import permchain.channels as channels
|
||||
from permchain.pregel import Pregel
|
||||
|
||||
__all__ = [
|
||||
"PubSub",
|
||||
"Topic",
|
||||
"InMemoryPubSubConnection",
|
||||
]
|
||||
__all__ = ["channels", "Pregel"]
|
||||
|
||||
+64
-13
@@ -1,5 +1,6 @@
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Callable, FrozenSet, Generic, Self, Sequence, TypeVar
|
||||
from typing import Callable, FrozenSet, Generic, Optional, Self, Sequence, TypeVar
|
||||
|
||||
Value = TypeVar("Value")
|
||||
Update = TypeVar("Update")
|
||||
@@ -14,8 +15,18 @@ class InvalidUpdateError(Exception):
|
||||
|
||||
|
||||
class Channel(Generic[Value, Update], ABC):
|
||||
def _empty(self) -> Self:
|
||||
return self.__class__()
|
||||
def __init__(self, name: str) -> None:
|
||||
self.name = name
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}({self.name})"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.name
|
||||
|
||||
@abstractmethod
|
||||
def _empty(self, checkpoint: Optional[str] = None) -> Self:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def _update(self, values: Sequence[Update]) -> None:
|
||||
@@ -25,15 +36,23 @@ class Channel(Generic[Value, Update], ABC):
|
||||
def _get(self) -> Value:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def _checkpoint(self) -> str:
|
||||
...
|
||||
|
||||
|
||||
class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]):
|
||||
def __init__(self, operator: Callable[[Value, Value], Value]):
|
||||
def __init__(self, name: str, operator: Callable[[Value, Value], Value]):
|
||||
super().__init__(name)
|
||||
self.operator = operator
|
||||
|
||||
def _empty(self) -> Self:
|
||||
return self.__class__(self.operator)
|
||||
def _empty(self, checkpoint: Optional[str] = None) -> Self:
|
||||
empty = self.__class__(self.name, self.operator)
|
||||
if checkpoint is not None:
|
||||
empty.value = json.loads(checkpoint)
|
||||
return empty
|
||||
|
||||
def _update(self, values):
|
||||
def _update(self, values: Sequence[Value]) -> None:
|
||||
if not hasattr(self, "value"):
|
||||
self.value = values[0]
|
||||
values = values[1:]
|
||||
@@ -41,40 +60,69 @@ class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]):
|
||||
for value in values:
|
||||
self.value = self.operator(self.value, value)
|
||||
|
||||
def _get(self):
|
||||
def _get(self) -> Value:
|
||||
try:
|
||||
return self.value
|
||||
except AttributeError:
|
||||
raise EmptyChannelError()
|
||||
|
||||
def _checkpoint(self) -> str:
|
||||
return json.dumps(self.value)
|
||||
|
||||
|
||||
class LastValue(Generic[Value], Channel[Value, Value]):
|
||||
def _update(self, values):
|
||||
def _empty(self, checkpoint: Optional[str] = None) -> Self:
|
||||
empty = self.__class__(self.name)
|
||||
if checkpoint is not None:
|
||||
empty.value = json.loads(checkpoint)
|
||||
return empty
|
||||
|
||||
def _update(self, values: Sequence[Value]) -> None:
|
||||
if len(values) != 1:
|
||||
raise InvalidUpdateError()
|
||||
|
||||
self.value = values[-1]
|
||||
|
||||
def _get(self):
|
||||
def _get(self) -> Value:
|
||||
try:
|
||||
return self.value
|
||||
except AttributeError:
|
||||
raise EmptyChannelError()
|
||||
|
||||
def _checkpoint(self) -> str:
|
||||
return json.dumps(self.value)
|
||||
|
||||
|
||||
class Inbox(Generic[Value], Channel[Sequence[Value], Value]):
|
||||
def _update(self, values):
|
||||
def _empty(self, checkpoint: Optional[str] = None) -> Self:
|
||||
empty = self.__class__(self.name)
|
||||
if checkpoint is not None:
|
||||
empty.queue = tuple(json.loads(checkpoint))
|
||||
return empty
|
||||
|
||||
def _update(self, values: Sequence[Value]) -> None:
|
||||
self.queue = tuple(values)
|
||||
|
||||
def _get(self):
|
||||
def _get(self) -> Sequence[Value]:
|
||||
try:
|
||||
return self.queue
|
||||
except AttributeError:
|
||||
raise EmptyChannelError()
|
||||
|
||||
def _checkpoint(self) -> str:
|
||||
return json.dumps(self.queue)
|
||||
|
||||
|
||||
class Set(Generic[Value], Channel[FrozenSet[Value], Value]):
|
||||
def _update(self, values) -> None:
|
||||
set: set[Value]
|
||||
|
||||
def _empty(self, checkpoint: Optional[str] = None) -> Self:
|
||||
empty = self.__class__(self.name)
|
||||
if checkpoint is not None:
|
||||
empty.set = set(json.loads(checkpoint))
|
||||
return empty
|
||||
|
||||
def _update(self, values: Sequence[Value]) -> None:
|
||||
if not hasattr(self, "set"):
|
||||
self.set = set()
|
||||
self.set.update(values)
|
||||
@@ -84,3 +132,6 @@ class Set(Generic[Value], Channel[FrozenSet[Value], Value]):
|
||||
return frozenset(self.set)
|
||||
except AttributeError:
|
||||
raise EmptyChannelError()
|
||||
|
||||
def _checkpoint(self) -> str:
|
||||
return json.dumps(list(self.set))
|
||||
|
||||
@@ -1,65 +0,0 @@
|
||||
import asyncio
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Callable, Iterator, TypedDict
|
||||
|
||||
|
||||
class PubSubMessage(TypedDict):
|
||||
topic: str
|
||||
value: Any
|
||||
published_at: str
|
||||
correlation_id: str
|
||||
|
||||
|
||||
PubSubListener = Callable[[PubSubMessage], None]
|
||||
|
||||
|
||||
class PubSubConnection(ABC):
|
||||
def full_name(self, prefix: str, *parts: str) -> str:
|
||||
"""Return the full topic name for a given prefix and topic name."""
|
||||
return ":".join(map(str, [prefix, *parts]))
|
||||
|
||||
@abstractmethod
|
||||
def observe(self, prefix: str) -> Iterator[PubSubMessage]:
|
||||
"""Iterate over messages for all topics under this prefix,
|
||||
without affecting listeners/iterators on each topic.
|
||||
This method waits for new messages to arrive."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def iterate(
|
||||
self, prefix: str, topic: str, *, wait: bool
|
||||
) -> Iterator[PubSubMessage]:
|
||||
"""Iterate over all currently queued messages for a topic, consuming them.
|
||||
Optionally wait for new messages to arrive."""
|
||||
...
|
||||
|
||||
# TODO add aiterate() method
|
||||
|
||||
@abstractmethod
|
||||
def listen(self, prefix: str, topic: str, listeners: list[PubSubListener]) -> None:
|
||||
...
|
||||
|
||||
async def alisten(
|
||||
self, prefix: str, topic: str, listeners: list[PubSubListener]
|
||||
) -> None:
|
||||
return await asyncio.get_event_loop().run_in_executor(
|
||||
None, self.listen, prefix, topic, listeners
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def send(self, prefix: str, topic: str, value: Any) -> None:
|
||||
...
|
||||
|
||||
async def asend(self, prefix: str, topic: str, value: Any) -> None:
|
||||
return await asyncio.get_event_loop().run_in_executor(
|
||||
None, self.send, prefix, topic, value
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def disconnect(self, name: str) -> None:
|
||||
...
|
||||
|
||||
async def adisconnect(self, name: str) -> None:
|
||||
return await asyncio.get_event_loop().run_in_executor(
|
||||
None, self.disconnect, name
|
||||
)
|
||||
@@ -1,125 +0,0 @@
|
||||
import queue
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
from datetime import datetime
|
||||
from typing import Any, Iterator, cast
|
||||
|
||||
from permchain.connection import PubSubConnection, PubSubListener, PubSubMessage
|
||||
|
||||
|
||||
class IterableQueue(queue.SimpleQueue):
|
||||
done_sentinel = object()
|
||||
|
||||
def put(
|
||||
self, item: PubSubMessage, block: bool = True, timeout: float | None = None
|
||||
) -> None:
|
||||
return super().put(item, block, timeout)
|
||||
|
||||
def get(
|
||||
self, block: bool = True, timeout: float | None = None
|
||||
) -> PubSubMessage | object:
|
||||
return super().get(block=block, timeout=timeout)
|
||||
|
||||
def __iter__(self) -> Iterator[PubSubMessage]:
|
||||
return iter(self.get, self.done_sentinel)
|
||||
|
||||
def close(self) -> None:
|
||||
self.put(self.done_sentinel)
|
||||
|
||||
|
||||
class InMemoryPubSubConnection(PubSubConnection):
|
||||
clear_on_disconnect: bool
|
||||
logs: defaultdict[str, IterableQueue]
|
||||
topics: defaultdict[str, IterableQueue]
|
||||
listeners: defaultdict[str, list[PubSubListener]]
|
||||
lock: threading.RLock
|
||||
|
||||
def __init__(self, clear_on_disconnect: bool = True) -> None:
|
||||
self.clear_on_disconnect = clear_on_disconnect
|
||||
self.logs = defaultdict(IterableQueue)
|
||||
self.topics = defaultdict(IterableQueue)
|
||||
self.listeners = defaultdict(list)
|
||||
self.lock = threading.RLock()
|
||||
|
||||
def observe(self, prefix: str) -> Iterator[PubSubMessage]:
|
||||
return iter(self.logs[str(prefix)])
|
||||
|
||||
def iterate(
|
||||
self, prefix: str, topic: str, *, wait: bool
|
||||
) -> Iterator[PubSubMessage]:
|
||||
topic = self.full_name(prefix, topic)
|
||||
|
||||
# This connection doesn't support iterating over topics with listeners connected
|
||||
with self.lock:
|
||||
if self.listeners[topic]:
|
||||
raise RuntimeError(
|
||||
f"Cannot iterate over topic {topic} while listeners are connected"
|
||||
)
|
||||
|
||||
# If wait is False, add sentinel to queue to ensure the iterator terminates
|
||||
if not wait:
|
||||
self.topics[topic].close()
|
||||
|
||||
return iter(self.topics[topic])
|
||||
|
||||
def listen(self, prefix: str, topic: str, listeners: list[PubSubListener]) -> None:
|
||||
full_name = self.full_name(prefix, topic)
|
||||
self.disconnect(full_name)
|
||||
|
||||
with self.lock:
|
||||
# Add the listeners for future messages
|
||||
self.listeners[full_name].extend(listeners)
|
||||
|
||||
# Send any pending messages to the listeners
|
||||
topic_queue = self.topics[full_name]
|
||||
while not topic_queue.empty():
|
||||
message = topic_queue.get()
|
||||
if message is not topic_queue.done_sentinel:
|
||||
for listener in self.listeners[full_name]:
|
||||
listener(cast(PubSubMessage, message))
|
||||
|
||||
def send(self, prefix: str, topic: str, value: Any) -> None:
|
||||
full_name = self.full_name(prefix, topic)
|
||||
message = PubSubMessage(
|
||||
value=value,
|
||||
topic=topic,
|
||||
correlation_id=str(prefix),
|
||||
published_at=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
# Add the message to the log
|
||||
self.logs[str(prefix)].put(message)
|
||||
with self.lock:
|
||||
listeners = self.listeners[full_name]
|
||||
if listeners:
|
||||
# Send the message to listeners if any are connected
|
||||
for listener in listeners:
|
||||
listener(message)
|
||||
else:
|
||||
# Otherwise add the message to the topic queue for later
|
||||
self.topics[full_name].put(message)
|
||||
|
||||
def disconnect(self, name: str) -> None:
|
||||
with self.lock:
|
||||
if name in self.logs:
|
||||
self.logs[name].close()
|
||||
if self.clear_on_disconnect:
|
||||
del self.logs[name]
|
||||
|
||||
to_delete = []
|
||||
for topic, q in self.topics.items():
|
||||
if topic.startswith(name):
|
||||
q.close()
|
||||
if self.clear_on_disconnect:
|
||||
to_delete.append(topic)
|
||||
# can't delete while iterating
|
||||
for topic in to_delete:
|
||||
del self.topics[topic]
|
||||
|
||||
to_delete = []
|
||||
for topic in self.listeners:
|
||||
if topic.startswith(name):
|
||||
to_delete.append(topic)
|
||||
# can't delete while iterating
|
||||
for topic in to_delete:
|
||||
del self.listeners[topic]
|
||||
@@ -1,2 +0,0 @@
|
||||
CONFIG_GET_KEY = "pubsub_get"
|
||||
CONFIG_SEND_KEY = "pubsub_send"
|
||||
+200
-52
@@ -17,21 +17,22 @@ from typing import (
|
||||
)
|
||||
|
||||
from langchain.callbacks.manager import (
|
||||
CallbackManagerForChainRun,
|
||||
AsyncCallbackManagerForChainRun,
|
||||
CallbackManagerForChainRun,
|
||||
)
|
||||
from langchain.pydantic_v1 import Field
|
||||
from langchain.schema.runnable import (
|
||||
Runnable,
|
||||
RunnableSerializable,
|
||||
RunnableBinding,
|
||||
RunnableLambda,
|
||||
RunnablePassthrough,
|
||||
RunnableSerializable,
|
||||
)
|
||||
from langchain.schema.runnable.base import (
|
||||
RunnableLike,
|
||||
Other,
|
||||
coerce_to_runnable,
|
||||
RunnableEach,
|
||||
RunnableLike,
|
||||
coerce_to_runnable,
|
||||
)
|
||||
from langchain.schema.runnable.config import (
|
||||
RunnableConfig,
|
||||
@@ -42,27 +43,78 @@ from langchain.schema.runnable.utils import ConfigurableFieldSpec, Input, Output
|
||||
|
||||
from permchain.channels import Channel, EmptyChannelError, Inbox
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
CONFIG_KEY_STEP = "__pregel_step"
|
||||
CONFIG_KEY_WRITE = "__pregel_write"
|
||||
CONFIG_KEY_SEND = "__pregel_send"
|
||||
CONFIG_KEY_READ = "__pregel_read"
|
||||
|
||||
TYPE_SEND = Callable[[Sequence[tuple[Channel, Any]]], None]
|
||||
|
||||
|
||||
class PregelRead(RunnableLambda):
|
||||
channel: Channel
|
||||
|
||||
@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: Channel) -> 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[[Channel], 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[[Channel], 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[str | None, Channel]
|
||||
channels: Mapping[None, Channel] | Mapping[str, Channel]
|
||||
|
||||
bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
|
||||
|
||||
kwargs: Mapping[str, Any] = Field(default_factory=dict)
|
||||
|
||||
def join(self, **channels: Channel) -> PregelInvoke:
|
||||
joiner = RunnablePassthrough.assign(
|
||||
**{k: PregelRead(chan) for k, chan in channels.items()}
|
||||
)
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return PregelInvoke(channels=self.channels, bound=joiner)
|
||||
else:
|
||||
return PregelInvoke(channels=self.channels, bound=self.bound | joiner)
|
||||
|
||||
def __or__(
|
||||
self,
|
||||
other: Runnable[Any, Other]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
|
||||
) -> PregelInvoke:
|
||||
) -> Runnable:
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return PregelInvoke(channels=self.channels, bound=coerce_to_runnable(other))
|
||||
else:
|
||||
@@ -73,7 +125,7 @@ class PregelInvoke(RunnableBinding):
|
||||
other: Runnable[Other, Any]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
|
||||
) -> PregelInvoke:
|
||||
) -> Runnable:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
@@ -87,7 +139,7 @@ class PregelBatch(RunnableEach):
|
||||
other: Runnable[Any, Other]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
|
||||
) -> PregelBatch:
|
||||
) -> Runnable:
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return PregelBatch(channel=self.channel, bound=coerce_to_runnable(other))
|
||||
else:
|
||||
@@ -98,12 +150,12 @@ class PregelBatch(RunnableEach):
|
||||
other: Runnable[Other, Any]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
|
||||
) -> PregelBatch:
|
||||
) -> Runnable:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class PregelSink(RunnablePassthrough):
|
||||
channels: Mapping[Channel, Runnable]
|
||||
class PregelSink(RunnableLambda):
|
||||
channels: Sequence[tuple[Channel, Runnable]]
|
||||
"""
|
||||
Mapping of write channels to Runnables that return the value to be written,
|
||||
or None to skip writing.
|
||||
@@ -111,38 +163,64 @@ class PregelSink(RunnablePassthrough):
|
||||
|
||||
max_steps: Optional[int]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
channels: Sequence[tuple[Channel, Runnable]],
|
||||
max_steps: Optional[int] = None,
|
||||
):
|
||||
super().__init__(func=self._write, afunc=self._awrite)
|
||||
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_WRITE,
|
||||
annotation=Callable[[Sequence[tuple[Channel, Any]]], None],
|
||||
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.get("configurable", {})[CONFIG_KEY_STEP]
|
||||
step: int = config["configurable"][CONFIG_KEY_STEP]
|
||||
|
||||
if step >= self.max_steps:
|
||||
if self.max_steps is not None and step >= self.max_steps:
|
||||
return
|
||||
|
||||
write: Callable[[Sequence[tuple[Channel, Any]]], None] = config.get(
|
||||
"configurable", {}
|
||||
)[CONFIG_KEY_WRITE]
|
||||
write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
|
||||
|
||||
# TODO use runnable map to run this in parallel?
|
||||
values = [(chan, r.invoke(input, config)) for chan, r in self.channels.items()]
|
||||
values = [(chan, val) for chan, val in values if val is not None]
|
||||
values = [(chan, r.invoke(input, config)) for chan, r in self.channels]
|
||||
|
||||
write(values)
|
||||
write([(chan, val) for chan, val in values if val is not None])
|
||||
|
||||
return input
|
||||
|
||||
# TODO def _awrite()
|
||||
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.items()
|
||||
]
|
||||
|
||||
write([(chan, val) for chan, val in values if val is not None])
|
||||
|
||||
return input
|
||||
|
||||
|
||||
class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
@@ -154,6 +232,9 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
|
||||
step_timeout: Optional[float] = None
|
||||
|
||||
class Config:
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
processes: Sequence[PregelInvoke | PregelBatch],
|
||||
@@ -173,51 +254,51 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
|
||||
@overload
|
||||
@classmethod
|
||||
def read(cls, __channel: Channel) -> PregelInvoke:
|
||||
def subscribe_to(cls, __channel: Channel) -> PregelInvoke:
|
||||
...
|
||||
|
||||
@overload
|
||||
def read(cls, __channel: Mapping[str, Channel], **kwargs: Channel) -> PregelInvoke:
|
||||
@classmethod
|
||||
def subscribe_to(
|
||||
cls, __channel: Mapping[str, Channel] | None = None, **kwargs: Channel
|
||||
) -> PregelInvoke:
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def read(
|
||||
cls, __channel: Channel | Mapping[str, Channel], **kwargs: Channel
|
||||
def subscribe_to(
|
||||
cls, __channel: Channel | Mapping[str, Channel] | None = None, **kwargs: Channel
|
||||
) -> PregelInvoke:
|
||||
"""Runs process.invoke() each time channels are updated."""
|
||||
return PregelInvoke(
|
||||
channels=(
|
||||
{None: __channel}
|
||||
if isinstance(__channel, Channel)
|
||||
else {**__channel, **kwargs}
|
||||
)
|
||||
__channel = __channel or {}
|
||||
return (
|
||||
PregelInvoke(channels={None: __channel})
|
||||
if isinstance(__channel, Channel)
|
||||
else PregelInvoke(channels={**__channel, **kwargs})
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def read_batch(cls, inbox: Inbox):
|
||||
def subscribe_to_each(cls, inbox: Inbox) -> PregelBatch:
|
||||
"""Runs process.batch() on the current contents of the inbox."""
|
||||
return PregelBatch(channel=inbox)
|
||||
|
||||
@classmethod
|
||||
def write(
|
||||
def send_to(
|
||||
cls,
|
||||
channels: Channel | Mapping[Channel, RunnableLike],
|
||||
*,
|
||||
max_steps: Optional[int] = None,
|
||||
):
|
||||
) -> PregelSink:
|
||||
return PregelSink(
|
||||
channels=(
|
||||
{channels: RunnablePassthrough()}
|
||||
[(channels, RunnablePassthrough())]
|
||||
if isinstance(channels, Channel)
|
||||
else {**channels}
|
||||
else [(k, coerce_to_runnable(v)) for k, v in channels.items()]
|
||||
),
|
||||
max_steps=max_steps,
|
||||
)
|
||||
|
||||
# TODO def write_each()
|
||||
|
||||
def _prepare_channels(self) -> Mapping[Channel, Channel]:
|
||||
channels = {self.output: self.output._empty()}
|
||||
channels: dict[Channel, Channel] = {self.output: self.output._empty()}
|
||||
for proc in self.processes:
|
||||
if isinstance(proc, PregelInvoke):
|
||||
for chan in proc.channels.values():
|
||||
@@ -232,10 +313,10 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
)
|
||||
|
||||
if not channels:
|
||||
return ValueError("Found 0 channels for Pregel run")
|
||||
raise ValueError("Found 0 channels for Pregel run")
|
||||
|
||||
if self.input not in channels:
|
||||
return ValueError("Input channel not being read from")
|
||||
raise ValueError("Input channel not being read from")
|
||||
|
||||
return channels
|
||||
|
||||
@@ -246,11 +327,18 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
config: RunnableConfig,
|
||||
) -> Iterator[Output]:
|
||||
processes = tuple(self.processes)
|
||||
# TODO this is where we'd restore from checkpoint
|
||||
channels = self._prepare_channels()
|
||||
next_tasks = _apply_writes_and_prepare_next_tasks(
|
||||
processes, channels, deque((self.input, chunk) for chunk in input)
|
||||
)
|
||||
|
||||
def read(chan: Channel) -> Any:
|
||||
try:
|
||||
return channels[chan]._get()
|
||||
except EmptyChannelError:
|
||||
return None
|
||||
|
||||
with get_executor_for_config(config) as executor:
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
# computation proceeds in steps, while there are channel updates
|
||||
@@ -273,7 +361,8 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
callbacks=run_manager.get_child(f"pregel:step:{step}"),
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_WRITE: pending_writes.extend,
|
||||
CONFIG_KEY_SEND: pending_writes.extend,
|
||||
CONFIG_KEY_READ: read,
|
||||
CONFIG_KEY_STEP: step,
|
||||
},
|
||||
),
|
||||
@@ -311,6 +400,8 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
if any(chan is self.output for chan, _ in pending_writes):
|
||||
yield channels[self.output]._get()
|
||||
|
||||
# TODO this is where we'd save checkpoint
|
||||
|
||||
# if no more tasks, we're done
|
||||
if not next_tasks:
|
||||
break
|
||||
@@ -328,9 +419,15 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
processes = tuple(self.processes)
|
||||
channels = self._prepare_channels()
|
||||
next_tasks = _apply_writes_and_prepare_next_tasks(
|
||||
processes, channels, [(self.input, chunk) for chunk in input]
|
||||
processes, channels, [(self.input, chunk) async for chunk in input]
|
||||
)
|
||||
|
||||
def read(chan: Channel) -> Any:
|
||||
try:
|
||||
return channels[chan]._get()
|
||||
except EmptyChannelError:
|
||||
return None
|
||||
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
# computation proceeds in steps, while there are channel updates
|
||||
# channel updates from step N are only visible in step N+1,
|
||||
@@ -338,7 +435,7 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
# channel updates being applied only at the transition between steps
|
||||
for step in range(config["recursion_limit"]):
|
||||
# collect all writes to channels, without applying them yet
|
||||
pending_writes = []
|
||||
pending_writes = deque[tuple[Channel, Any]]()
|
||||
|
||||
# execute tasks, and wait for one to fail or all to finish
|
||||
# each task is independent from all other concurrent tasks
|
||||
@@ -352,7 +449,8 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
callbacks=run_manager.get_child(f"pregel:step:{step}"),
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_WRITE: pending_writes.extend,
|
||||
CONFIG_KEY_SEND: pending_writes.extend,
|
||||
CONFIG_KEY_READ: read,
|
||||
CONFIG_KEY_STEP: step,
|
||||
},
|
||||
),
|
||||
@@ -395,10 +493,56 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
|
||||
if not next_tasks:
|
||||
break
|
||||
|
||||
# TODO invoke() consumes stream() iterator and returns last value
|
||||
# TODO ainvoke() consumes astream() iterator and returns last value
|
||||
def invoke(
|
||||
self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
|
||||
) -> Output:
|
||||
latest: Output | None = None
|
||||
for chunk in self.stream(input, config, **kwargs):
|
||||
latest = chunk
|
||||
return latest
|
||||
|
||||
# TODO do we want api to subscribe to all channels?
|
||||
def stream(
|
||||
self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
|
||||
) -> Iterator[Output]:
|
||||
return self.transform(iter([input]), config, **kwargs)
|
||||
|
||||
def transform(
|
||||
self,
|
||||
input: Iterator[Input],
|
||||
config: RunnableConfig | None = None,
|
||||
**kwargs: Any | None,
|
||||
) -> Iterator[Output]:
|
||||
return self._transform_stream_with_config(
|
||||
input, self._transform, config, **kwargs
|
||||
)
|
||||
|
||||
async def ainvoke(
|
||||
self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
|
||||
) -> Output:
|
||||
latest: Output | None = None
|
||||
async for chunk in self.astream(input, config, **kwargs):
|
||||
latest = chunk
|
||||
return latest
|
||||
|
||||
async def astream(
|
||||
self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
|
||||
) -> AsyncIterator[Output]:
|
||||
async def input_stream() -> AsyncIterator[Input]:
|
||||
yield input
|
||||
|
||||
async for chunk in self.atransform(input_stream(), config, **kwargs):
|
||||
yield chunk
|
||||
|
||||
async def atransform(
|
||||
self,
|
||||
input: AsyncIterator[Input],
|
||||
config: RunnableConfig | None = None,
|
||||
**kwargs: Any | None,
|
||||
) -> AsyncIterator[Output]:
|
||||
async for chunk in self._transform_stream_with_config(
|
||||
input, self._atransform, config, **kwargs
|
||||
):
|
||||
yield chunk
|
||||
|
||||
|
||||
def _apply_writes_and_prepare_next_tasks(
|
||||
@@ -451,3 +595,7 @@ def _apply_writes_and_prepare_next_tasks(
|
||||
tasks.append((proc, val))
|
||||
|
||||
return tasks
|
||||
|
||||
|
||||
# TODO do we want api to subscribe to all channels?
|
||||
# Do we want api to send input to multiple channels in invoke()
|
||||
|
||||
@@ -1,250 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from abc import ABC
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import CancelledError, Future
|
||||
from functools import partial
|
||||
from typing import Any, Iterator, List, Optional, Sequence, Set, TypeVar
|
||||
|
||||
from langchain.callbacks.manager import CallbackManagerForChainRun
|
||||
from langchain.schema.runnable import Runnable, RunnableConfig, patch_config
|
||||
from langchain.schema.runnable.config import get_executor_for_config
|
||||
|
||||
from permchain.connection import PubSubConnection, PubSubMessage
|
||||
from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY
|
||||
from permchain.topic import (
|
||||
INPUT_TOPIC,
|
||||
OUTPUT_TOPIC,
|
||||
RunnableReducer,
|
||||
RunnableSubscriber,
|
||||
)
|
||||
|
||||
T = TypeVar("T")
|
||||
T_in = TypeVar("T_in")
|
||||
T_out = TypeVar("T_out")
|
||||
|
||||
Process = RunnableSubscriber[T_in] | RunnableReducer[T_in]
|
||||
|
||||
|
||||
class PubSub(Runnable[Any, Any], ABC):
|
||||
processes: Sequence[Process]
|
||||
|
||||
connection: PubSubConnection
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*procs: Process | Sequence[Process],
|
||||
processes: Sequence[Process] = (),
|
||||
connection: PubSubConnection,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.lock = threading.Lock()
|
||||
self.inflight_namespaces = set()
|
||||
|
||||
self.connection = connection
|
||||
self.processes = list(processes)
|
||||
for proc in procs:
|
||||
if isinstance(proc, Sequence):
|
||||
self.processes.extend(proc)
|
||||
else:
|
||||
self.processes.append(proc)
|
||||
|
||||
def with_retry(self, **kwargs: Any) -> Runnable[Any, Any]:
|
||||
return self.__class__(
|
||||
processes=[p.with_retry(**kwargs) for p in self.processes],
|
||||
connection=self.connection,
|
||||
)
|
||||
|
||||
def _transform(
|
||||
self,
|
||||
input: Iterator[Any],
|
||||
run_manager: CallbackManagerForChainRun,
|
||||
config: RunnableConfig,
|
||||
) -> Iterator[Any]:
|
||||
# Split processes into subscribers and reducers, and group by topic
|
||||
subscribers: defaultdict[str, list[RunnableSubscriber[Any]]] = defaultdict(list)
|
||||
reducers: defaultdict[str, list[RunnableReducer[Any]]] = defaultdict(list)
|
||||
for process in self.processes:
|
||||
if isinstance(process, RunnableReducer):
|
||||
reducers[process.topic.name].append(process)
|
||||
elif isinstance(process, RunnableSubscriber):
|
||||
subscribers[process.topic.name].append(process)
|
||||
else:
|
||||
raise ValueError(f"Unknown process type: {process}")
|
||||
|
||||
# Consume input iterator into a single value
|
||||
input_value = None
|
||||
for chunk in input:
|
||||
if input_value is None:
|
||||
input_value = chunk
|
||||
else:
|
||||
input_value += chunk
|
||||
|
||||
with get_executor_for_config(config) as executor:
|
||||
# Namespace topics for each run, default to run_id, ie. isolated
|
||||
topic_prefix = str(config.get("correlation_id") or run_manager.run_id)
|
||||
|
||||
# Check if this correlation_id is currently inflight. If so, raise an error,
|
||||
# as that would make the output iterator produce incorrect results.
|
||||
with self.lock:
|
||||
if topic_prefix in self.inflight_namespaces:
|
||||
raise RuntimeError(
|
||||
f"Cannot run {self} in namespace {topic_prefix} "
|
||||
"because it is currently in use"
|
||||
)
|
||||
self.inflight_namespaces.add(topic_prefix)
|
||||
|
||||
# Track inflight futures
|
||||
inflight: Set[Future] = set()
|
||||
# Track exceptions
|
||||
exceptions: List[BaseException] = []
|
||||
|
||||
def on_idle() -> None:
|
||||
"""Called when all subscribed topics are empty.
|
||||
It first runs any topic reducers. Then, if all subscribed topics
|
||||
still empty, it closes the computation.
|
||||
"""
|
||||
if reducers:
|
||||
for topic_name, processes in reducers.items():
|
||||
# Collect all pending messages for each topic
|
||||
messages = list(
|
||||
self.connection.iterate(
|
||||
topic_prefix, topic_name, wait=False
|
||||
)
|
||||
)
|
||||
# Run each reducer once with the collected messages
|
||||
if messages:
|
||||
for process in processes:
|
||||
run_once(process, messages)
|
||||
|
||||
if not inflight:
|
||||
self.connection.disconnect(topic_prefix)
|
||||
|
||||
def check_if_idle(fut: Future) -> None:
|
||||
"""Cleanup after a process runs."""
|
||||
inflight.discard(fut)
|
||||
|
||||
try:
|
||||
exc = fut.exception()
|
||||
except CancelledError:
|
||||
exc = None
|
||||
except Exception as e:
|
||||
exc = e
|
||||
if exc is not None:
|
||||
exceptions.append(exc)
|
||||
|
||||
# Close output iterator if
|
||||
# - all processes are done, or
|
||||
# - an exception occurred
|
||||
if not inflight or exc is not None:
|
||||
on_idle()
|
||||
|
||||
def run_once(
|
||||
process: RunnableSubscriber[Any] | RunnableReducer[Any],
|
||||
messages: PubSubMessage | list[PubSubMessage],
|
||||
) -> None:
|
||||
"""Run a process once."""
|
||||
value = (
|
||||
[m["value"] for m in messages]
|
||||
if isinstance(messages, list)
|
||||
else messages["value"]
|
||||
)
|
||||
|
||||
def get(topic_name: str) -> Any:
|
||||
if topic_name == INPUT_TOPIC:
|
||||
return input_value
|
||||
elif topic_name == process.topic.name:
|
||||
return value
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Cannot get value for {topic_name} in this context"
|
||||
)
|
||||
|
||||
# Run process once in executor
|
||||
try:
|
||||
fut = executor.submit(
|
||||
process.invoke,
|
||||
value,
|
||||
config={
|
||||
**patch_config(
|
||||
config,
|
||||
callbacks=run_manager.get_child(),
|
||||
run_name=f"Topic: {process.topic.name}",
|
||||
),
|
||||
CONFIG_SEND_KEY: partial(
|
||||
self.connection.send, topic_prefix
|
||||
),
|
||||
CONFIG_GET_KEY: get,
|
||||
# TODO below doesn't work for batch calls nested inside
|
||||
# another pubsub, eg. test_invoke_join_then_call_other_pubsub
|
||||
# as all messages in each batch would share same correlation_id
|
||||
# "correlation_id": self.connection.full_name(
|
||||
# topic_prefix,
|
||||
# process.topic.name,
|
||||
# str(self.processes.index(process)),
|
||||
# ),
|
||||
},
|
||||
)
|
||||
|
||||
# Add callback to cleanup
|
||||
inflight.add(fut)
|
||||
fut.add_done_callback(check_if_idle)
|
||||
except RuntimeError:
|
||||
# If executor is now closed, just ignore this process
|
||||
# This could happen eg. if an OUT message was published durin
|
||||
# execution of run_once
|
||||
pass
|
||||
|
||||
# Listen on all subscribed topics
|
||||
for topic_name, processes in subscribers.items():
|
||||
self.connection.listen(
|
||||
topic_prefix,
|
||||
topic_name,
|
||||
[partial(run_once, process) for process in processes],
|
||||
)
|
||||
|
||||
# Send input to input processes
|
||||
self.connection.send(topic_prefix, INPUT_TOPIC, input_value)
|
||||
|
||||
try:
|
||||
if inflight:
|
||||
# Yield output until all processes are done
|
||||
# This blocks the current thread, all other work needs to go
|
||||
# through the executor
|
||||
for chunk in self.connection.observe(topic_prefix):
|
||||
yield chunk
|
||||
if chunk["topic"] == OUTPUT_TOPIC:
|
||||
# All expected output has been received, close
|
||||
self.connection.disconnect(topic_prefix)
|
||||
break
|
||||
else:
|
||||
on_idle()
|
||||
finally:
|
||||
# Cancel all inflight futures
|
||||
while inflight:
|
||||
inflight.pop().cancel()
|
||||
|
||||
# Remove namespace from inflight set
|
||||
with self.lock:
|
||||
self.inflight_namespaces.remove(topic_prefix)
|
||||
|
||||
# Raise exceptions if any
|
||||
if exceptions:
|
||||
raise exceptions[0]
|
||||
|
||||
def stream(
|
||||
self,
|
||||
input: Any,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
**kwargs: Optional[Any],
|
||||
) -> Iterator[PubSubMessage]:
|
||||
yield from self._transform_stream_with_config(
|
||||
iter([input]), self._transform, config, **kwargs
|
||||
)
|
||||
|
||||
def invoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any:
|
||||
for chunk in self.stream(input, config):
|
||||
if chunk["topic"] == OUTPUT_TOPIC:
|
||||
return chunk["value"]
|
||||
@@ -1,168 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Generic,
|
||||
Mapping,
|
||||
Optional,
|
||||
Sequence,
|
||||
TypeVar,
|
||||
)
|
||||
|
||||
from langchain.load.serializable import Serializable
|
||||
from langchain.pydantic_v1 import Field
|
||||
from langchain.schema.runnable import (
|
||||
Runnable,
|
||||
RunnableBinding,
|
||||
RunnableConfig,
|
||||
RunnablePassthrough,
|
||||
RunnableSequence,
|
||||
)
|
||||
from langchain.schema.runnable.base import Other, coerce_to_runnable
|
||||
|
||||
from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
INPUT_TOPIC = "__in__"
|
||||
OUTPUT_TOPIC = "__out__"
|
||||
|
||||
|
||||
class Topic(Serializable, Generic[T]):
|
||||
name: str
|
||||
|
||||
def __init__(self, name: str):
|
||||
super().__init__(name=name)
|
||||
|
||||
def subscribe(self) -> RunnableSubscriber[T]:
|
||||
if self.name == OUTPUT_TOPIC:
|
||||
raise ValueError("Cannot subscribe to output topic")
|
||||
|
||||
return RunnableSubscriber(topic=self)
|
||||
|
||||
def join(self) -> RunnableReducer[T]:
|
||||
if self.name == OUTPUT_TOPIC:
|
||||
raise ValueError("Cannot join on output topic")
|
||||
|
||||
return RunnableReducer(topic=self)
|
||||
|
||||
def current(self) -> RunnableCurrentValue[T]:
|
||||
if self.name == OUTPUT_TOPIC:
|
||||
raise ValueError("Cannot subscribe to output topic")
|
||||
|
||||
return RunnableCurrentValue(topic=self)
|
||||
|
||||
def publish(self) -> RunnablePublisher[T]:
|
||||
if self.name == INPUT_TOPIC:
|
||||
raise ValueError("Cannot publish to input topic")
|
||||
|
||||
return RunnablePublisher(topic=self)
|
||||
|
||||
def publish_each(self) -> RunnablePublisherEach[T]:
|
||||
if self.name == INPUT_TOPIC:
|
||||
raise ValueError("Cannot publish to input topic")
|
||||
|
||||
return RunnablePublisherEach(topic=self)
|
||||
|
||||
@classmethod
|
||||
@property
|
||||
def IN(cls) -> Topic:
|
||||
return cls(INPUT_TOPIC)
|
||||
|
||||
@classmethod
|
||||
@property
|
||||
def OUT(cls) -> Topic:
|
||||
return cls(OUTPUT_TOPIC)
|
||||
|
||||
|
||||
class RunnableConfigForPubSub(RunnableConfig):
|
||||
send: Callable[[str, Any], None]
|
||||
get: Callable[[str], Any]
|
||||
|
||||
|
||||
class RunnableSubscriber(RunnableBinding[T, Any]):
|
||||
topic: Topic[T]
|
||||
|
||||
bound: Runnable[T, Any] = Field(default_factory=RunnablePassthrough)
|
||||
|
||||
kwargs: Mapping[str, Any] = Field(default_factory=dict)
|
||||
|
||||
def __or__(
|
||||
self,
|
||||
other: Runnable[Any, Other]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
|
||||
) -> RunnableSubscriber[T, Other]:
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return RunnableSubscriber(topic=self.topic, bound=coerce_to_runnable(other))
|
||||
else:
|
||||
return RunnableSubscriber(topic=self.topic, bound=self.bound | other)
|
||||
|
||||
def __ror__(
|
||||
self,
|
||||
other: Runnable[Other, Any]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
|
||||
) -> RunnableSubscriber[Other, Any]:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class RunnableReducer(RunnableBinding[list[T], Any]):
|
||||
topic: Topic[T]
|
||||
|
||||
bound: Runnable[list[T], Any] = Field(default_factory=RunnablePassthrough)
|
||||
|
||||
kwargs: Mapping[str, Any] = Field(default_factory=dict)
|
||||
|
||||
def __or__(
|
||||
self,
|
||||
other: Runnable[Any, Other]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
|
||||
) -> RunnableSequence[list[T], Other]:
|
||||
if isinstance(self.bound, RunnablePassthrough):
|
||||
return RunnableReducer(topic=self.topic, bound=coerce_to_runnable(other))
|
||||
else:
|
||||
return RunnableReducer(topic=self.topic, bound=self.bound | other)
|
||||
|
||||
def __ror__(
|
||||
self,
|
||||
other: Runnable[Other, Any]
|
||||
| Callable[[Any], Other]
|
||||
| Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
|
||||
) -> RunnableSequence[Other, Any]:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class RunnablePublisher(Serializable, Runnable[T, T]):
|
||||
topic: Topic[T]
|
||||
|
||||
def invoke(self, input: T, config: Optional[RunnableConfigForPubSub] = None) -> T:
|
||||
send = config.get(CONFIG_SEND_KEY, None)
|
||||
if send is not None:
|
||||
send(self.topic.name, input)
|
||||
return input
|
||||
|
||||
|
||||
class RunnablePublisherEach(RunnablePublisher[Sequence[T]]):
|
||||
topic: Topic[T]
|
||||
|
||||
def invoke(
|
||||
self, input: Sequence[T], config: Optional[RunnableConfigForPubSub] = None
|
||||
) -> Sequence[T]:
|
||||
for item in input:
|
||||
super().invoke(item, config)
|
||||
return input
|
||||
|
||||
|
||||
class RunnableCurrentValue(Serializable, Runnable[Any, T]):
|
||||
topic: Topic[T]
|
||||
|
||||
def invoke(self, input: T, config: Optional[RunnableConfigForPubSub] = None) -> T:
|
||||
get: Callable[[str], None] = config.get(CONFIG_GET_KEY, None)
|
||||
if get is not None:
|
||||
return get(self.topic.name)
|
||||
else:
|
||||
raise ValueError("Cannot get value in this context")
|
||||
Reference in New Issue
Block a user