diff --git a/examples/runnable-pregel.py b/examples/runnable-pregel.py index e4ec8c8be..0bb94a660 100644 --- a/examples/runnable-pregel.py +++ b/examples/runnable-pregel.py @@ -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("---") diff --git a/permchain/__init__.py b/permchain/__init__.py index 770967f30..f6e3f8dfb 100644 --- a/permchain/__init__.py +++ b/permchain/__init__.py @@ -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"] diff --git a/permchain/channels.py b/permchain/channels.py index daafb7068..c2708851d 100644 --- a/permchain/channels.py +++ b/permchain/channels.py @@ -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)) diff --git a/permchain/connection.py b/permchain/connection.py deleted file mode 100644 index 9c8528816..000000000 --- a/permchain/connection.py +++ /dev/null @@ -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 - ) diff --git a/permchain/connection_inmemory.py b/permchain/connection_inmemory.py deleted file mode 100644 index 53b288ddf..000000000 --- a/permchain/connection_inmemory.py +++ /dev/null @@ -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] diff --git a/permchain/constants.py b/permchain/constants.py deleted file mode 100644 index 4c9c12ee7..000000000 --- a/permchain/constants.py +++ /dev/null @@ -1,2 +0,0 @@ -CONFIG_GET_KEY = "pubsub_get" -CONFIG_SEND_KEY = "pubsub_send" diff --git a/permchain/pregel.py b/permchain/pregel.py index 04224df13..6e7a8ea55 100644 --- a/permchain/pregel.py +++ b/permchain/pregel.py @@ -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() diff --git a/permchain/pubsub.py b/permchain/pubsub.py deleted file mode 100644 index 8688467b5..000000000 --- a/permchain/pubsub.py +++ /dev/null @@ -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"] diff --git a/permchain/topic.py b/permchain/topic.py deleted file mode 100644 index ae7f5ac67..000000000 --- a/permchain/topic.py +++ /dev/null @@ -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")