WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-25 17:12:26 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1582 from langchain-ai/nc/2sep/rm-pydantic-base-class

Browse Source 
Remove pydantic base class from Pregel/PregelNode
This commit is contained in:
Nuno Campos
2024-09-03 09:39:03 -07:00
committed by GitHub
parent 43d5dbf174 4ddb4ae5a0
commit b4944fa0a0
8 changed files with 171 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
+11 -10
View File
@@ -165,14 +165,15 @@ class SqliteSaver(BaseCheckpointSaver):
Yields:
sqlite3.Cursor: A cursor for the SQLite database.
"""
self.setup()
cur = self.conn.cursor()
try:
yield cur
finally:
if transaction:
self.conn.commit()
cur.close()
with self.lock:
self.setup()
cur = self.conn.cursor()
try:
yield cur
finally:
if transaction:
self.conn.commit()
cur.close()
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database.
@@ -398,7 +399,7 @@ class SqliteSaver(BaseCheckpointSaver):
checkpoint_ns = config["configurable"]["checkpoint_ns"]
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
serialized_metadata = self.jsonplus_serde.dumps(metadata)
with self.lock, self.cursor() as cur:
with self.cursor() as cur:
cur.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
(
@@ -434,7 +435,7 @@ class SqliteSaver(BaseCheckpointSaver):
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
"""
with self.lock, self.cursor() as cur:
with self.cursor() as cur:
cur.executemany(
"INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
[
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
+6 -4
View File
@@ -276,7 +276,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
"""
await self.setup()
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
async with self.conn.cursor() as cur:
async with self.lock, self.conn.cursor() as cur:
# find the latest checkpoint for the thread_id
if checkpoint_id := get_checkpoint_id(config):
await cur.execute(
@@ -371,7 +371,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur:
async with self.lock, self.conn.execute(
query, params
) as cur, self.conn.cursor() as wcur:
async for (
thread_id,
checkpoint_ns,
@@ -438,7 +440,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
checkpoint_ns = config["configurable"]["checkpoint_ns"]
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
serialized_metadata = self.jsonplus_serde.dumps(metadata)
async with self.conn.execute(
async with self.lock, self.conn.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
(
str(config["configurable"]["thread_id"]),
@@ -475,7 +477,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
task_id (str): Identifier for the task creating the writes.
"""
await self.setup()
async with self.conn.cursor() as cur:
async with self.lock, self.conn.cursor() as cur:
await cur.executemany(
"INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
[
libs/langgraph/langgraph/graph/graph.py
+4
View File
@@ -424,6 +424,10 @@ class Graph:
class CompiledGraph(Pregel):
builder: Graph
def __init__(self, *, builder: Graph, **kwargs):
super().__init__(**kwargs)
self.builder = builder
def attach_node(self, key: str, node: NodeSpec) -> None:
self.channels[key] = EphemeralValue(Any)
self.nodes[key] = (
libs/langgraph/langgraph/graph/state.py
+1 -1
View File
@@ -17,10 +17,10 @@ from typing import (
overload,
)
from langchain_core.pydantic_v1 import BaseModel
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.base import RunnableLike
from langchain_core.runnables.utils import create_model
from pydantic import BaseModel
from langgraph.channels.base import BaseChannel
from langgraph.channels.binop import BinaryOperatorAggregate
libs/langgraph/langgraph/pregel/__init__.py
+48 -41
View File
@@ -25,12 +25,10 @@ from uuid import UUID, uuid5
from langchain_core.globals import get_debug
from langchain_core.load.dump import dumpd
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
from langchain_core.runnables import (
Runnable,
RunnableLambda,
RunnableSequence,
RunnableSerializable,
)
from langchain_core.runnables.base import Input, Output, coerce_to_runnable
from langchain_core.runnables.config import (
@@ -48,6 +46,7 @@ from langchain_core.runnables.utils import (
get_unique_config_specs,
)
from langchain_core.tracers._streaming import _StreamingCallbackHandler
from pydantic import BaseModel
from typing_extensions import Self
from langgraph.channels.base import (
@@ -186,16 +185,10 @@ class Channel:
)
class Pregel(
RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]
):
class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
nodes: Mapping[str, PregelNode]
channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = Field(
default_factory=dict
)
auto_validate: bool = True
channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
stream_mode: StreamMode = "values"
"""Mode to stream output, defaults to 'values'."""
@@ -205,16 +198,16 @@ class Pregel(
stream_channels: Optional[Union[str, Sequence[str]]] = None
"""Channels to stream, defaults to all channels not in reserved channels"""
interrupt_after_nodes: Union[All, Sequence[str]] = Field(default_factory=list)
interrupt_after_nodes: Union[All, Sequence[str]]
interrupt_before_nodes: Union[All, Sequence[str]] = Field(default_factory=list)
interrupt_before_nodes: Union[All, Sequence[str]]
input_channels: Union[str, Sequence[str]]
step_timeout: Optional[float] = None
"""Maximum time to wait for a step to complete, in seconds. Defaults to None."""
debug: bool = Field(default_factory=get_debug)
debug: bool
"""Whether to print debug information during execution. Defaults to False."""
checkpointer: Optional[BaseCheckpointSaver] = None
@@ -232,36 +225,50 @@ class Pregel(
name: str = "LangGraph"
class Config:
arbitrary_types_allowed = True
def __init__(
self,
*,
nodes: Mapping[str, PregelNode],
channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = None,
auto_validate: bool = True,
stream_mode: StreamMode = "values",
output_channels: Union[str, Sequence[str]],
stream_channels: Optional[Union[str, Sequence[str]]] = None,
interrupt_after_nodes: Union[All, Sequence[str]] = (),
interrupt_before_nodes: Union[All, Sequence[str]] = (),
input_channels: Union[str, Sequence[str]],
step_timeout: Optional[float] = None,
debug: Optional[bool] = None,
checkpointer: Optional[BaseCheckpointSaver] = None,
store: Optional[BaseStore] = None,
retry_policy: Optional[RetryPolicy] = None,
config_type: Optional[Type[Any]] = None,
config: Optional[RunnableConfig] = None,
name: str = "LangGraph",
) -> None:
self.nodes = nodes
self.channels = channels or {}
self.stream_mode = stream_mode
self.output_channels = output_channels
self.stream_channels = stream_channels
self.interrupt_after_nodes = interrupt_after_nodes
self.interrupt_before_nodes = interrupt_before_nodes
self.input_channels = input_channels
self.step_timeout = step_timeout
self.debug = debug if debug is not None else get_debug()
self.checkpointer = checkpointer
self.store = store
self.retry_policy = retry_policy
self.config_type = config_type
self.config = config
self.name = name
if auto_validate:
self.validate()
def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self:
return self.copy(
update={"config": cast(RunnableConfig, {**(config or {}), **kwargs})}
)
@classmethod
def is_lc_serializable(cls) -> bool:
"""Return whether the graph can be serialized by Langchain."""
return True
@root_validator(skip_on_failure=True)
def validate_on_init(cls, values: dict[str, Any]) -> dict[str, Any]:
if not values["auto_validate"]:
return values
validate_graph(
values["nodes"],
values["channels"],
values["input_channels"],
values["output_channels"],
values["stream_channels"],
values["interrupt_after_nodes"],
values["interrupt_before_nodes"],
)
if values["interrupt_after_nodes"] or values["interrupt_before_nodes"]:
if not values["checkpointer"]:
raise ValueError("Interrupts require a checkpointer")
return values
attrs = {**self.__dict__}
attrs["config"] = merge_configs(self.config, config, kwargs)
return self.__class__(**attrs)
def validate(self) -> Self:
validate_graph(
libs/langgraph/langgraph/pregel/read.py
+84 -42
View File
@@ -1,8 +1,16 @@
from __future__ import annotations
from typing import Any, Callable, Mapping, Optional, Sequence, Union
from typing import (
Any,
AsyncIterator,
Callable,
Iterator,
Mapping,
Optional,
Sequence,
Union,
)
from langchain_core.pydantic_v1 import Field
from langchain_core.runnables import (
Runnable,
RunnableConfig,
@@ -10,7 +18,7 @@ from langchain_core.runnables import (
RunnableSequence,
RunnableSerializable,
)
from langchain_core.runnables.base import Other, RunnableBindingBase, coerce_to_runnable
from langchain_core.runnables.base import Input, Other, Output, coerce_to_runnable
from langchain_core.runnables.config import merge_configs
from langchain_core.runnables.utils import ConfigurableFieldSpec
@@ -99,20 +107,47 @@ class ChannelRead(RunnableCallable):
DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough()
class PregelNode(RunnableBindingBase):
class PregelNode(Runnable):
channels: Union[list[str], Mapping[str, str]]
triggers: list[str] = Field(default_factory=list)
triggers: list[str]
mapper: Optional[Callable[[Any], Any]] = None
mapper: Optional[Callable[[Any], Any]]
writers: list[Runnable] = Field(default_factory=list)
writers: list[Runnable]
bound: Runnable[Any, Any] = Field(default=DEFAULT_BOUND)
bound: Runnable[Any, Any]
kwargs: Mapping[str, Any] = Field(default_factory=dict)
retry_policy: Optional[RetryPolicy]
retry_policy: Optional[RetryPolicy] = None
config: RunnableConfig
def __init__(
self,
*,
channels: Union[list[str], Mapping[str, str]],
triggers: Sequence[str],
mapper: Optional[Callable[[Any], Any]] = None,
writers: Optional[list[Runnable]] = None,
tags: Optional[list[str]] = None,
metadata: Optional[Mapping[str, Any]] = None,
bound: Optional[Runnable[Any, Any]] = None,
retry_policy: Optional[RetryPolicy] = None,
config: Optional[RunnableConfig] = None,
) -> None:
self.channels = channels
self.triggers = list(triggers)
self.mapper = mapper
self.writers = writers or []
self.bound = bound if bound is not None else DEFAULT_BOUND
self.retry_policy = retry_policy
self.config = merge_configs(
config, {"tags": tags or [], "metadata": metadata or {}}
)
def copy(self, update: dict[str, Any]) -> PregelNode:
attrs = {**self.__dict__, **update}
return PregelNode(**attrs)
def get_writers(self) -> list[Runnable]:
"""Get writers with optimizations applied."""
@@ -145,38 +180,6 @@ class PregelNode(RunnableBindingBase):
else:
return self.bound
def __init__(
self,
*,
channels: Union[list[str], Mapping[str, str]],
triggers: Sequence[str],
mapper: Optional[Callable[[Any], Any]] = None,
writers: Optional[list[Runnable]] = None,
tags: Optional[list[str]] = None,
metadata: Optional[Mapping[str, Any]] = None,
bound: Optional[Runnable[Any, Any]] = None,
kwargs: Optional[Mapping[str, Any]] = None,
config: Optional[RunnableConfig] = None,
retry_policy: Optional[RetryPolicy] = None,
**other_kwargs: Any,
) -> None:
super().__init__(
channels=channels,
triggers=triggers,
mapper=mapper,
writers=writers or [],
bound=bound or DEFAULT_BOUND,
kwargs=kwargs or {},
retry_policy=retry_policy,
config=merge_configs(
config, {"tags": tags or [], "metadata": metadata or {}}
),
**other_kwargs,
)
def __repr_args__(self) -> Any:
return [(k, v) for k, v in super().__repr_args__() if k != "bound"]
def join(self, channels: Sequence[str]) -> PregelNode:
assert isinstance(channels, list) or isinstance(
channels, tuple
@@ -226,3 +229,42 @@ class PregelNode(RunnableBindingBase):
],
) -> RunnableSerializable:
raise NotImplementedError()
def invoke(
self,
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> Output:
return self.bound.invoke(input, merge_configs(self.config, config), **kwargs)
async def ainvoke(
self,
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> Output:
return await self.bound.ainvoke(
input, merge_configs(self.config, config), **kwargs
)
def stream(
self,
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> Iterator[Output]:
yield from self.bound.stream(
input, merge_configs(self.config, config), **kwargs
)
async def astream(
self,
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> AsyncIterator[Output]:
async for item in self.bound.astream(
input, merge_configs(self.config, config), **kwargs
):
yield item
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+10 -30
View File
@@ -3304,7 +3304,7 @@
'query',
'inner',
]),
'title': 'Input',
'title': 'LangGraphInput',
'type': 'object',
})
# ---
@@ -3323,11 +3323,7 @@
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'title': 'LangGraphOutput',
'type': 'object',
})
# ---
@@ -3374,7 +3370,7 @@
'query',
'inner',
]),
'title': 'Input',
'title': 'LangGraphInput',
'type': 'object',
})
# ---
@@ -3393,11 +3389,7 @@
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'title': 'LangGraphOutput',
'type': 'object',
})
# ---
@@ -3444,7 +3436,7 @@
'query',
'inner',
]),
'title': 'Input',
'title': 'LangGraphInput',
'type': 'object',
})
# ---
@@ -3463,11 +3455,7 @@
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'title': 'LangGraphOutput',
'type': 'object',
})
# ---
@@ -3514,7 +3502,7 @@
'query',
'inner',
]),
'title': 'Input',
'title': 'LangGraphInput',
'type': 'object',
})
# ---
@@ -3533,11 +3521,7 @@
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'title': 'LangGraphOutput',
'type': 'object',
})
# ---
@@ -3584,7 +3568,7 @@
'query',
'inner',
]),
'title': 'Input',
'title': 'LangGraphInput',
'type': 'object',
})
# ---
@@ -3603,11 +3587,7 @@
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'title': 'LangGraphOutput',
'type': 'object',
})
# ---
libs/langgraph/tests/test_pregel.py
+7 -5
View File
@@ -1903,7 +1903,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
one = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("between") | add_one
with pytest.raises(ValueError):
with pytest.raises(TypeError):
Pregel(nodes={"one": one, "two": two})
@@ -9943,10 +9943,12 @@ def test_send_to_nested_graphs(
graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"})
# continue past interrupt
assert graph.invoke(None, config=config) == {
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
}
assert sorted(
graph.stream(None, config=config), key=lambda d: d["generate_joke"]["jokes"][0]
) == [
{"generate_joke": {"jokes": ["Joke about cats - hohoho"]}},
{"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}},
]
actual_snapshot = graph.get_state(config)
expected_snapshot = StateSnapshot(