Update langchain

permchain/channels/base.py

@@ -54,7 +54,7 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
 
     @asynccontextmanager
     async def aempty(
-        self, checkpoint: Optional[str] = None
+        self, checkpoint: Optional[Checkpoint] = None
     ) -> AsyncGenerator[Self, None]:
         """Return a new identical channel, optionally initialized from a checkpoint."""
         with self.empty(checkpoint) as value:

permchain/channels/topic.py

@@ -43,7 +43,7 @@ class Topic(
         return Sequence[self.typ]  # type: ignore[name-defined]
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> Any:
         """The type of the update received by the channel."""
         return Union[self.typ, list[self.typ]]  # type: ignore[name-defined]
 

permchain/checkpoint/base.py

@@ -1,6 +1,6 @@
 import asyncio
 from abc import ABC, abstractmethod
-from typing import Any, Mapping, Sequence
+from typing import Any, Mapping
 
 from langchain.load.serializable import Serializable
 from langchain.schema.runnable import RunnableConfig
@@ -18,7 +18,7 @@ class BaseCheckpointAdapter(Serializable, ABC):
     at: CheckpointAt = CheckpointAt.END_OF_RUN
 
     @property
-    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+    def config_specs(self) -> list[ConfigurableFieldSpec]:
         return []
 
     @abstractmethod

permchain/checkpoint/memory.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Mapping, Sequence
+from typing import Any, Dict, Mapping
 
 from langchain.pydantic_v1 import Field
 from langchain.schema.runnable import RunnableConfig
@@ -11,7 +11,7 @@ class MemoryCheckpoint(BaseCheckpointAdapter):
     storage: Dict[str, Mapping[str, Any]] = Field(default_factory=dict)
 
     @property
-    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+    def config_specs(self) -> list[ConfigurableFieldSpec]:
         return [
             ConfigurableFieldSpec(
                 id="thread_id",

permchain/pregel/__init__.py

@@ -27,7 +27,6 @@ from langchain.globals import get_debug
 from langchain.pydantic_v1 import BaseModel, Field, create_model, root_validator
 from langchain.schema.runnable import (
     Runnable,
-    RunnablePassthrough,
     RunnableSerializable,
 )
 from langchain.schema.runnable.base import Input, Output, coerce_to_runnable

permchain/pregel/validate.py

@@ -14,7 +14,7 @@ FORBIDDEN_CHANNEL_NAMES = {
 
 def validate_chains_channels(
     chains: Mapping[str, ChannelInvoke | ChannelBatch],
-    channels: Mapping[str, BaseChannel],
+    channels: dict[str, BaseChannel],
     input: str | Sequence[str],
     output: str | Sequence[str],
 ) -> None:
@@ -31,17 +31,17 @@ def validate_chains_channels(
 
     for chan in subscribed_channels:
         if chan not in channels:
-            channels[chan] = LastValue(Any)
+            channels[chan] = LastValue(Any)  # type: ignore[arg-type]
 
     if isinstance(input, str):
         if input not in channels:
-            channels[input] = LastValue(Any)
+            channels[input] = LastValue(Any)  # type: ignore[arg-type]
         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 channels:
-                channels[chan] = LastValue(Any)
+                channels[chan] = LastValue(Any)  # type: ignore[arg-type]
         if all(chan not in subscribed_channels for chan in input):
             raise ValueError(
                 f"None of the input channels {input} are subscribed to by any chain"
@@ -49,11 +49,11 @@ def validate_chains_channels(
 
     if isinstance(output, str):
         if output not in channels:
-            channels[output] = LastValue(Any)
+            channels[output] = LastValue(Any)  # type: ignore[arg-type]
     else:
         for chan in output:
             if chan not in channels:
-                channels[chan] = LastValue(Any)
+                channels[chan] = LastValue(Any)  # type: ignore[arg-type]
 
     for name in FORBIDDEN_CHANNEL_NAMES:
         if name in channels:
@@ -61,4 +61,4 @@ def validate_chains_channels(
 
     for chan in ReservedChannels:
         if chan not in channels:
-            channels[chan] = LastValue(Any)
+            channels[chan] = LastValue(Any)  # type: ignore[arg-type]

tests/test_pregel.py

@@ -281,27 +281,39 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
         | raise_if_above_10
     )
 
+    memory = MemoryCheckpoint()
+
     app = Pregel(
         chains={"chain_one": chain_one},
         channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        checkpoint=MemoryCheckpoint(),
+        checkpoint=memory,
     )
 
     # total starts out as 0, so output is 0+2=2
     assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2
-    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 2
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 2
     # total is now 2, so output is 2+3=5
     assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5
-    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
     # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
     with pytest.raises(ValueError):
         app.invoke(4, {"configurable": {"thread_id": "1"}})
     # checkpoint is not updated
-    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
     # on a new thread, total starts out as 0, so output is 0+5=5
     assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5
-    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
-    assert app.checkpoint.get({"configurable": {"thread_id": "2"}}).get("total") == 5
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
+    checkpoint = memory.get({"configurable": {"thread_id": "2"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 5
 
 
 def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:

tests/test_pregel_async.py

@@ -294,37 +294,39 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
         | raise_if_above_10
     )
 
+    memory = MemoryCheckpoint()
+
     app = Pregel(
         chains={"chain_one": chain_one},
         channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        checkpoint=MemoryCheckpoint(),
+        checkpoint=memory,
     )
 
     # total starts out as 0, so output is 0+2=2
     assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
-        "total"
-    ) == 2
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 2
     # total is now 2, so output is 2+3=5
     assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
-        "total"
-    ) == 7
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
     # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
     with pytest.raises(ValueError):
         await app.ainvoke(4, {"configurable": {"thread_id": "1"}})
     # checkpoint is not updated
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
-        "total"
-    ) == 7
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
     # on a new thread, total starts out as 0, so output is 0+5=5
     assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
-        "total"
-    ) == 7
-    assert (await app.checkpoint.aget({"configurable": {"thread_id": "2"}})).get(
-        "total"
-    ) == 5
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 7
+    checkpoint = await memory.aget({"configurable": {"thread_id": "2"}})
+    assert checkpoint is not None
+    assert checkpoint.get("total") == 5
 
 
 async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: