mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-22 23:52:23 +02:00
125 lines
3.2 KiB
Python
125 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
from abc import ABC
|
|
|
|
from typing import (
|
|
Any,
|
|
Callable,
|
|
Generic,
|
|
Mapping,
|
|
Optional,
|
|
Sequence,
|
|
TypeVar,
|
|
)
|
|
|
|
from langchain.load.serializable import Serializable
|
|
from langchain.schema.runnable import (
|
|
Runnable,
|
|
RunnableBinding,
|
|
RunnableConfig,
|
|
RunnablePassthrough,
|
|
RunnableSequence,
|
|
)
|
|
from langchain.schema.runnable.base import Other, coerce_to_runnable
|
|
from pydantic import Field
|
|
|
|
T = TypeVar("T")
|
|
T_in = TypeVar("T_in")
|
|
T_out = TypeVar("T_out")
|
|
|
|
|
|
INPUT_TOPIC = "__in__"
|
|
OUTPUT_TOPIC = "__out__"
|
|
|
|
|
|
class Topic(Serializable, Generic[T], ABC):
|
|
name: str
|
|
|
|
def __init__(self, name: str):
|
|
super().__init__(name=name)
|
|
|
|
def subscribe(self) -> RunnableSubscriber[T]:
|
|
return RunnableSubscriber(topic=self)
|
|
|
|
def current(self) -> RunnableCurrentValue[T]:
|
|
return RunnableCurrentValue(topic=self)
|
|
|
|
def publish(self) -> RunnablePublisher[T]:
|
|
return RunnablePublisher(topic=self)
|
|
|
|
def publish_each(self) -> Runnable[T, T]:
|
|
return RunnablePublisherEach(topic=self)
|
|
|
|
@classmethod
|
|
@property
|
|
def IN(cls) -> Topic[T_in]:
|
|
return cls[T_in](INPUT_TOPIC)
|
|
|
|
@classmethod
|
|
@property
|
|
def OUT(cls) -> Topic[T_out]:
|
|
return cls[T_out](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]],
|
|
) -> RunnableSequence[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]],
|
|
) -> RunnableSequence[Other, Any]:
|
|
raise NotImplementedError()
|
|
|
|
|
|
class RunnablePublisher(RunnablePassthrough[T]):
|
|
topic: Topic[T]
|
|
|
|
def invoke(self, input: T, config: Optional[RunnableConfigForPubSub] = None) -> T:
|
|
send = config.get("send", None)
|
|
if send is not None:
|
|
send(self.topic.name, input)
|
|
return super().invoke(input, config)
|
|
|
|
|
|
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)
|
|
|
|
|
|
class RunnableCurrentValue(Serializable, Runnable[Any, T]):
|
|
topic: Topic[T]
|
|
|
|
def invoke(self, input: T, config: Optional[RunnableConfigForPubSub] = None) -> T:
|
|
get = config.get("get", None)
|
|
if get is not None:
|
|
return get(self.topic.name)
|
|
else:
|
|
raise ValueError("Cannot get value in this context")
|