mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-19 22:25:44 +02:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
837212b969 | ||
|
|
70918557ca | ||
|
|
1e44bda48f |
@@ -77,6 +77,10 @@ __pypackages__/
|
||||
# Environments
|
||||
.env
|
||||
.envrc
|
||||
*.crt
|
||||
*.key
|
||||
*.pem
|
||||
credentials.json
|
||||
.venv
|
||||
.venvs
|
||||
env/
|
||||
@@ -98,6 +102,7 @@ dmypy.json
|
||||
|
||||
.vercel
|
||||
.turbo
|
||||
node_modules/
|
||||
.editorconfig
|
||||
.scratch
|
||||
.worktrees/
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import inspect
|
||||
import dis
|
||||
import re
|
||||
import textwrap
|
||||
from collections.abc import Callable, Sequence
|
||||
from functools import partial
|
||||
from types import CodeType, FunctionType
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import (
|
||||
@@ -17,7 +16,6 @@ from langchain_core.runnables import (
|
||||
from langchain_core.runnables.base import RunnableBindingBase
|
||||
from langchain_core.runnables.config import run_in_executor
|
||||
from langgraph.checkpoint.base import ChannelVersions
|
||||
from typing_extensions import override
|
||||
|
||||
from langgraph._internal._runnable import RunnableCallable, RunnableSeq
|
||||
from langgraph._internal._timeout import sync_timeout_unsupported
|
||||
@@ -137,155 +135,87 @@ def validate_timeout_supported(runnable: Runnable, *, name: str) -> None:
|
||||
raise sync_timeout_unsupported(name)
|
||||
|
||||
|
||||
# Values treated as dead ends when deciding whether to walk a function's
|
||||
# bytecode. A container can hold a graph, but `find_subgraph_pregel` does not
|
||||
# look inside one, so skipping it costs nothing while that holds. Matched by
|
||||
# exact type, since a subclass of a builtin can carry attributes.
|
||||
_LEAF_TYPES = frozenset(
|
||||
{
|
||||
int,
|
||||
float,
|
||||
complex,
|
||||
bool,
|
||||
str,
|
||||
bytes,
|
||||
bytearray,
|
||||
list,
|
||||
tuple,
|
||||
dict,
|
||||
set,
|
||||
frozenset,
|
||||
type(None),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def get_function_nonlocals(func: Callable) -> list[Any]:
|
||||
"""Get the nonlocal variables accessed by a function.
|
||||
"""Get the values a function reaches from outside its own scope.
|
||||
|
||||
Args:
|
||||
func: The function to check.
|
||||
|
||||
Returns:
|
||||
List[Any]: The nonlocal variables accessed by the function.
|
||||
Every captured cell value, the globals the function names, and each
|
||||
value along an attribute path it loads. Over-approximates: a value can
|
||||
come back without the function reaching it at runtime.
|
||||
"""
|
||||
try:
|
||||
code = inspect.getsource(func)
|
||||
tree = ast.parse(textwrap.dedent(code))
|
||||
visitor = FunctionNonLocals()
|
||||
visitor.visit(tree)
|
||||
values: list[Any] = []
|
||||
closure = (
|
||||
inspect.getclosurevars(func.__wrapped__)
|
||||
if hasattr(func, "__wrapped__") and callable(func.__wrapped__)
|
||||
else inspect.getclosurevars(func)
|
||||
)
|
||||
candidates = {**closure.globals, **closure.nonlocals}
|
||||
for k, v in candidates.items():
|
||||
if k in visitor.nonlocals:
|
||||
values.append(v)
|
||||
for kk in visitor.nonlocals:
|
||||
if "." in kk and kk.startswith(k):
|
||||
vv = v
|
||||
for part in kk.split(".")[1:]:
|
||||
if vv is None:
|
||||
break
|
||||
else:
|
||||
try:
|
||||
vv = getattr(vv, part)
|
||||
except AttributeError:
|
||||
break
|
||||
else:
|
||||
values.append(vv)
|
||||
except (SyntaxError, TypeError, OSError, SystemError):
|
||||
func = getattr(func, "__func__", func) # bound method -> function
|
||||
wrapped = getattr(func, "__wrapped__", None)
|
||||
if callable(wrapped):
|
||||
func = getattr(wrapped, "__func__", wrapped)
|
||||
if not isinstance(func, FunctionType):
|
||||
return []
|
||||
code = func.__code__
|
||||
|
||||
cells: dict[str, Any] = {}
|
||||
for name, cell in zip(code.co_freevars, func.__closure__ or ()):
|
||||
try:
|
||||
cells[name] = cell.cell_contents
|
||||
except ValueError:
|
||||
continue # empty cell: a recursive def not yet bound
|
||||
|
||||
# Every captured value counts, referenced or not: over-declaring costs an
|
||||
# introspection entry, under-declaring drops the subgraph's checkpoints and
|
||||
# stream events. Checking each cell against the bytecode would cost more and
|
||||
# only trade the cheap error for the expensive one.
|
||||
values: list[Any] = list(cells.values())
|
||||
global_ns = func.__globals__
|
||||
globals_ = {name: global_ns[name] for name in code.co_names if name in global_ns}
|
||||
if all(type(v) in _LEAF_TYPES for v in (*cells.values(), *globals_.values())):
|
||||
return values
|
||||
|
||||
# Nested code objects hold the references made by inner defs, lambdas and
|
||||
# comprehensions, which resolve against the namespaces gathered above.
|
||||
codes = [code]
|
||||
for c in codes:
|
||||
codes.extend(k for k in c.co_consts if isinstance(k, CodeType))
|
||||
value: Any = None
|
||||
for instruction in dis.get_instructions(c):
|
||||
opname = instruction.opname
|
||||
if opname == "LOAD_GLOBAL":
|
||||
value = globals_.get(instruction.argval)
|
||||
elif opname == "LOAD_DEREF":
|
||||
value = cells.get(instruction.argval)
|
||||
elif opname in ("LOAD_ATTR", "LOAD_METHOD"):
|
||||
value = getattr(value, instruction.argval, None)
|
||||
else:
|
||||
value = None # anything else ends the chain: `a, b.c` is not `a.c`
|
||||
continue
|
||||
if value is not None:
|
||||
values.append(value)
|
||||
return values
|
||||
|
||||
|
||||
class FunctionNonLocals(ast.NodeVisitor):
|
||||
"""Get the nonlocal variables accessed of a function."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.nonlocals: set[str] = set()
|
||||
|
||||
@override
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
|
||||
"""Visit a function definition.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
visitor = NonLocals()
|
||||
visitor.visit(node)
|
||||
self.nonlocals.update(visitor.loads - visitor.stores)
|
||||
|
||||
@override
|
||||
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:
|
||||
"""Visit an async function definition.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
visitor = NonLocals()
|
||||
visitor.visit(node)
|
||||
self.nonlocals.update(visitor.loads - visitor.stores)
|
||||
|
||||
@override
|
||||
def visit_Lambda(self, node: ast.Lambda) -> Any:
|
||||
"""Visit a lambda function.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
visitor = NonLocals()
|
||||
visitor.visit(node)
|
||||
self.nonlocals.update(visitor.loads - visitor.stores)
|
||||
|
||||
|
||||
class NonLocals(ast.NodeVisitor):
|
||||
"""Get nonlocal variables accessed."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.loads: set[str] = set()
|
||||
self.stores: set[str] = set()
|
||||
|
||||
@override
|
||||
def visit_Name(self, node: ast.Name) -> Any:
|
||||
"""Visit a name node.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
if isinstance(node.ctx, ast.Load):
|
||||
self.loads.add(node.id)
|
||||
elif isinstance(node.ctx, ast.Store):
|
||||
self.stores.add(node.id)
|
||||
|
||||
@override
|
||||
def visit_Attribute(self, node: ast.Attribute) -> Any:
|
||||
"""Visit an attribute node.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
if isinstance(node.ctx, ast.Load):
|
||||
parent = node.value
|
||||
attr_expr = node.attr
|
||||
while isinstance(parent, ast.Attribute):
|
||||
attr_expr = parent.attr + "." + attr_expr
|
||||
parent = parent.value
|
||||
if isinstance(parent, ast.Name):
|
||||
self.loads.add(parent.id + "." + attr_expr)
|
||||
self.loads.discard(parent.id)
|
||||
elif isinstance(parent, ast.Call):
|
||||
if isinstance(parent.func, ast.Name):
|
||||
self.loads.add(parent.func.id)
|
||||
else:
|
||||
parent = parent.func
|
||||
attr_expr = ""
|
||||
while isinstance(parent, ast.Attribute):
|
||||
if attr_expr:
|
||||
attr_expr = parent.attr + "." + attr_expr
|
||||
else:
|
||||
attr_expr = parent.attr
|
||||
parent = parent.value
|
||||
if isinstance(parent, ast.Name):
|
||||
self.loads.add(parent.id + "." + attr_expr)
|
||||
|
||||
|
||||
def is_xxh3_128_hexdigest(value: str) -> bool:
|
||||
"""Check if the given string matches the format of xxh3_128_hexdigest."""
|
||||
return bool(re.fullmatch(r"[0-9a-f]{32}", value))
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
"""Tests for subgraph auto-detection (`pregel/_utils.py`).
|
||||
|
||||
Detection failing is silent — the graph still runs, only introspection goes
|
||||
quiet — so every shape a node can hold a graph in is pinned here. The expected
|
||||
values are what the source-parsing implementation this replaced produced for
|
||||
the same shapes, except for `sourceless`, whose source it could not read,
|
||||
`empty_closure_cell`, on which it raised, and `unreachable_attribute_chain`,
|
||||
where it reported a graph that dropped code could never invoke.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import operator
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.pregel._utils import get_function_nonlocals
|
||||
|
||||
|
||||
class State(TypedDict):
|
||||
log: Annotated[list, operator.add]
|
||||
|
||||
|
||||
def _leaf(tag: str) -> Any:
|
||||
"""Return a compiled graph that reports itself as `tag`."""
|
||||
builder = StateGraph(State)
|
||||
builder.add_node(tag, lambda s: {"log": [tag]})
|
||||
builder.add_edge(START, tag)
|
||||
builder.add_edge(tag, END)
|
||||
compiled = builder.compile()
|
||||
compiled.name = tag
|
||||
return compiled
|
||||
|
||||
|
||||
def _detect(node: Any) -> str | None:
|
||||
"""Return the name of the subgraph detected for `node`, or None."""
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("n", node)
|
||||
builder.add_edge(START, "n")
|
||||
builder.add_edge("n", END)
|
||||
subgraphs = builder.compile().nodes["n"].subgraphs
|
||||
return getattr(subgraphs[0], "name", "?") if subgraphs else None
|
||||
|
||||
|
||||
class _Box:
|
||||
def __init__(self, payload: Any) -> None:
|
||||
self.payload = payload
|
||||
|
||||
|
||||
class _ListSubclass(list):
|
||||
pass
|
||||
|
||||
|
||||
class _MethodHolder:
|
||||
def __init__(self) -> None:
|
||||
self.graph = _leaf("via_self")
|
||||
|
||||
def as_node(self, state: State) -> Any:
|
||||
return self.graph.invoke(state)
|
||||
|
||||
|
||||
MODULE_GRAPH = _leaf("module_global")
|
||||
CHAIN = _Box(_Box(_leaf("attr_chain")))
|
||||
GRAPH_IN_PLAIN_LIST = [_leaf("in_list")]
|
||||
METHOD_HOLDER = _MethodHolder()
|
||||
|
||||
|
||||
def closure_capture() -> Any:
|
||||
sub = _leaf("closure")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return sub.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def module_global() -> Any:
|
||||
def node(state: State) -> Any:
|
||||
return MODULE_GRAPH.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def attribute_chain() -> Any:
|
||||
def node(state: State) -> Any:
|
||||
return CHAIN.payload.payload.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def nested_def_captured_attribute() -> Any:
|
||||
"""A chain on a captured holder, named only inside a nested code object.
|
||||
|
||||
The captured value is the holder, not the graph, so the chain itself has to
|
||||
be recovered from the nested scope.
|
||||
"""
|
||||
holder = _Box(_leaf("nested_captured"))
|
||||
|
||||
def node(state: State) -> Any:
|
||||
def inner() -> Any:
|
||||
return holder.payload.invoke(state)
|
||||
|
||||
return inner()
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def unreachable_branch() -> Any:
|
||||
"""A captured graph referenced only from code the compiler removes."""
|
||||
sub = _leaf("unreachable")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
if False:
|
||||
sub.invoke(state)
|
||||
return {"log": []}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def unreachable_attribute_chain() -> Any:
|
||||
"""A graph named only along an attribute path the compiler dropped.
|
||||
|
||||
The closure keeps `holder`, but the `.payload` load is gone. The source
|
||||
parser reported this one; dropped code cannot invoke anything, so that was
|
||||
a phantom rather than a detection.
|
||||
"""
|
||||
holder = _Box(_leaf("unreachable_attr"))
|
||||
|
||||
def node(state: State) -> Any:
|
||||
if False:
|
||||
holder.payload.invoke(state)
|
||||
return {"log": []}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def wrapper_referencing_nothing() -> Any:
|
||||
"""A wrapper whose own scope holds nothing, so only `__wrapped__` leads on.
|
||||
|
||||
`functools.wraps` would leave the wrapper closing over the inner function;
|
||||
setting the attribute by hand does not.
|
||||
"""
|
||||
sub = _leaf("via_wrapped")
|
||||
|
||||
def inner(state: State) -> Any:
|
||||
return sub.invoke(state)
|
||||
|
||||
def wrapper(state: State) -> Any:
|
||||
return {"log": []}
|
||||
|
||||
wrapper.__wrapped__ = inner
|
||||
return wrapper
|
||||
|
||||
|
||||
def captured_list_subclass() -> Any:
|
||||
"""A `list` subclass is not a leaf: it can carry a graph as an attribute."""
|
||||
holder = _ListSubclass()
|
||||
holder.payload = _leaf("list_subclass")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return holder.payload.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def empty_closure_cell() -> Any:
|
||||
"""An unassigned closure variable leaves a cell that cannot be read."""
|
||||
sub = _leaf("beside_empty_cell")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return unassigned, sub.invoke(state)
|
||||
|
||||
return node
|
||||
unassigned = 1 # never runs, so the cell it creates is never filled
|
||||
|
||||
|
||||
def sourceless() -> Any:
|
||||
"""A node compiled without a source file, which `getsource` could not read."""
|
||||
namespace: dict[str, Any] = {"SOURCELESS": _leaf("sourceless")}
|
||||
exec(
|
||||
compile(
|
||||
"def node(state):\n return SOURCELESS.invoke(state)", "<test>", "exec"
|
||||
),
|
||||
namespace,
|
||||
)
|
||||
return namespace["node"]
|
||||
|
||||
|
||||
async def _async_node(state: State) -> Any:
|
||||
return await MODULE_GRAPH.ainvoke(state)
|
||||
|
||||
|
||||
def async_node() -> Any:
|
||||
return _async_node
|
||||
|
||||
|
||||
def no_subgraph() -> Any:
|
||||
"""Nothing but leaf values in reach, so the bytecode walk is skipped."""
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return {"log": [len("abc") + 1]}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def recombined_names() -> Any:
|
||||
"""Loads `CHAIN.payload` and `local.payload`, never `CHAIN.payload.payload`."""
|
||||
|
||||
def node(state: State) -> Any:
|
||||
local = _Box("not a graph")
|
||||
return {"log": [CHAIN.payload, local.payload]}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def broken_attribute_chain() -> Any:
|
||||
holder = _Box("a string, so `.payload.missing` cannot resolve")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return holder.payload.missing.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def nested_def_global() -> Any:
|
||||
"""A global named only in a nested code object: out of reach, as before."""
|
||||
|
||||
def node(state: State) -> Any:
|
||||
def inner() -> Any:
|
||||
return MODULE_GRAPH.invoke(state)
|
||||
|
||||
return inner()
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def graph_in_plain_list() -> Any:
|
||||
def node(state: State) -> Any:
|
||||
return GRAPH_IN_PLAIN_LIST[0].invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def bound_method_self() -> Any:
|
||||
return METHOD_HOLDER.as_node
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("factory", "expected"),
|
||||
[
|
||||
(closure_capture, "closure"),
|
||||
(module_global, "module_global"),
|
||||
(attribute_chain, "attr_chain"),
|
||||
(nested_def_captured_attribute, "nested_captured"),
|
||||
(unreachable_branch, "unreachable"),
|
||||
(wrapper_referencing_nothing, "via_wrapped"),
|
||||
(captured_list_subclass, "list_subclass"),
|
||||
(empty_closure_cell, "beside_empty_cell"),
|
||||
(sourceless, "sourceless"),
|
||||
(async_node, "module_global"),
|
||||
# Shapes no reference chain reaches: a subscript, an instance attribute
|
||||
# of `self`, a global named only in a nested scope, and an attribute
|
||||
# path the compiler dropped.
|
||||
(no_subgraph, None),
|
||||
(recombined_names, None),
|
||||
(broken_attribute_chain, None),
|
||||
(nested_def_global, None),
|
||||
(graph_in_plain_list, None),
|
||||
(bound_method_self, None),
|
||||
(unreachable_attribute_chain, None),
|
||||
],
|
||||
ids=lambda value: value.__name__ if callable(value) else str(value),
|
||||
)
|
||||
def test_subgraph_detection(factory: Any, expected: str | None) -> None:
|
||||
assert _detect(factory()) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"candidate",
|
||||
[functools.partial(lambda state, extra: {"log": [extra]}, extra="x"), len],
|
||||
ids=["partial", "builtin"],
|
||||
)
|
||||
def test_callables_without_a_code_object_are_handled(candidate: Any) -> None:
|
||||
assert get_function_nonlocals(candidate) == []
|
||||
@@ -1,8 +1,15 @@
|
||||
from langgraph_sdk.auth import Auth
|
||||
from langgraph_sdk.client import get_client, get_sync_client
|
||||
from langgraph_sdk.encryption import Encryption
|
||||
from langgraph_sdk.encryption.types import EncryptionContext
|
||||
from langgraph_sdk.encryption.types import DecryptResult, EncryptionContext
|
||||
|
||||
__version__ = "0.4.2"
|
||||
__version__ = "0.4.3"
|
||||
|
||||
__all__ = ["Auth", "Encryption", "EncryptionContext", "get_client", "get_sync_client"]
|
||||
__all__ = [
|
||||
"Auth",
|
||||
"DecryptResult",
|
||||
"Encryption",
|
||||
"EncryptionContext",
|
||||
"get_client",
|
||||
"get_sync_client",
|
||||
]
|
||||
|
||||
@@ -18,6 +18,9 @@ import warnings
|
||||
|
||||
from langgraph_sdk.encryption import types
|
||||
|
||||
_BlobDecryptorT = typing.TypeVar("_BlobDecryptorT", bound=types.BlobDecryptor)
|
||||
_JsonDecryptorT = typing.TypeVar("_JsonDecryptorT", bound=types.JsonDecryptor)
|
||||
|
||||
|
||||
class LangGraphBetaWarning(UserWarning):
|
||||
"""Warning for beta features in LangGraph SDK."""
|
||||
@@ -141,7 +144,7 @@ class _DecryptDecorators:
|
||||
def __init__(self, parent: Encryption):
|
||||
self._parent = parent
|
||||
|
||||
def blob(self, fn: types.BlobDecryptor) -> types.BlobDecryptor:
|
||||
def blob(self, fn: _BlobDecryptorT) -> _BlobDecryptorT:
|
||||
"""Register a blob decryption handler.
|
||||
|
||||
The handler will be called to decrypt opaque data like checkpoint blobs.
|
||||
@@ -149,7 +152,9 @@ class _DecryptDecorators:
|
||||
Example:
|
||||
```python
|
||||
@encryption.decrypt.blob
|
||||
async def decrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes:
|
||||
async def decrypt_blob(
|
||||
ctx: EncryptionContext, blob: bytes
|
||||
) -> bytes | DecryptResult[bytes]:
|
||||
# Decrypt the blob using your encryption service
|
||||
return decrypted_blob
|
||||
```
|
||||
@@ -170,13 +175,15 @@ class _DecryptDecorators:
|
||||
self._parent._blob_decryptor = fn
|
||||
return fn
|
||||
|
||||
def json(self, fn: types.JsonDecryptor) -> types.JsonDecryptor:
|
||||
def json(self, fn: _JsonDecryptorT) -> _JsonDecryptorT:
|
||||
"""Register the JSON decryption handler.
|
||||
|
||||
Example:
|
||||
```python
|
||||
@encryption.decrypt.json
|
||||
async def decrypt_json(ctx: EncryptionContext, data: dict) -> dict:
|
||||
async def decrypt_json(
|
||||
ctx: EncryptionContext, data: dict
|
||||
) -> dict | DecryptResult[dict]:
|
||||
# Decrypt the data
|
||||
return decrypt_data(data)
|
||||
```
|
||||
@@ -369,7 +376,7 @@ class Encryption:
|
||||
"""Reference to encryption type definitions.
|
||||
|
||||
Provides access to all type definitions used in the encryption system,
|
||||
including EncryptionContext, BlobEncryptor, BlobDecryptor,
|
||||
including EncryptionContext, DecryptResult, BlobEncryptor, BlobDecryptor,
|
||||
JsonEncryptor, and JsonDecryptor.
|
||||
"""
|
||||
|
||||
|
||||
@@ -9,10 +9,30 @@ from __future__ import annotations
|
||||
|
||||
import typing
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
|
||||
Json = dict[str, typing.Any]
|
||||
"""JSON-serializable dictionary type for structured data encryption."""
|
||||
|
||||
T = typing.TypeVar("T")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DecryptResult(typing.Generic[T]):
|
||||
"""Decrypted data and optional replacement ciphertext.
|
||||
|
||||
Return this from a decrypt handler when encrypted data should be replaced,
|
||||
such as after rotating its encryption key. Returning plaintext directly
|
||||
remains supported when no replacement is needed.
|
||||
|
||||
Attributes:
|
||||
plaintext: Decrypted data returned to the caller
|
||||
replacement: New encrypted data to persist in place of the input
|
||||
"""
|
||||
|
||||
plaintext: T
|
||||
replacement: T | None = None
|
||||
|
||||
|
||||
class EncryptionContext:
|
||||
"""Context passed to encryption/decryption handlers.
|
||||
@@ -57,7 +77,9 @@ Returns:
|
||||
Awaitable that resolves to encrypted bytes
|
||||
"""
|
||||
|
||||
BlobDecryptor = Callable[[EncryptionContext, bytes], Awaitable[bytes]]
|
||||
BlobDecryptor = Callable[
|
||||
[EncryptionContext, bytes], Awaitable[bytes | DecryptResult[bytes]]
|
||||
]
|
||||
"""Handler for decrypting opaque blob data like checkpoints.
|
||||
|
||||
Note: Must be an async function. Decryption typically involves I/O operations
|
||||
@@ -68,7 +90,8 @@ Args:
|
||||
blob: The encrypted bytes to decrypt
|
||||
|
||||
Returns:
|
||||
Awaitable that resolves to decrypted bytes
|
||||
Awaitable that resolves to decrypted bytes, or a DecryptResult containing
|
||||
decrypted bytes and replacement ciphertext
|
||||
"""
|
||||
|
||||
JsonEncryptor = Callable[[EncryptionContext, Json], Awaitable[Json]]
|
||||
@@ -101,7 +124,9 @@ Returns:
|
||||
Awaitable that resolves to encrypted JSON dictionary
|
||||
"""
|
||||
|
||||
JsonDecryptor = Callable[[EncryptionContext, Json], Awaitable[Json]]
|
||||
JsonDecryptor = Callable[
|
||||
[EncryptionContext, Json], Awaitable[Json | DecryptResult[Json]]
|
||||
]
|
||||
"""Handler for decrypting structured JSON data.
|
||||
|
||||
Note: Must be an async function. Decryption typically involves I/O operations
|
||||
@@ -115,7 +140,8 @@ Args:
|
||||
data: The encrypted JSON dictionary
|
||||
|
||||
Returns:
|
||||
Awaitable that resolves to decrypted JSON dictionary
|
||||
Awaitable that resolves to a decrypted JSON dictionary, or a DecryptResult
|
||||
containing decrypted JSON and replacement ciphertext
|
||||
"""
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
|
||||
@@ -1,8 +1,40 @@
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
import pytest
|
||||
|
||||
from langgraph_sdk import DecryptResult, EncryptionContext
|
||||
from langgraph_sdk.encryption import DuplicateHandlerError, Encryption
|
||||
|
||||
|
||||
def test_decrypt_result():
|
||||
result = DecryptResult(plaintext=b"plain", replacement=b"rotated")
|
||||
|
||||
assert result.plaintext == b"plain"
|
||||
assert result.replacement == b"rotated"
|
||||
assert DecryptResult(plaintext={"plain": True}).replacement is None
|
||||
|
||||
|
||||
def test_decrypt_decorators_preserve_return_types():
|
||||
encryption = Encryption()
|
||||
|
||||
@encryption.decrypt.blob
|
||||
async def blob_dec(_ctx: EncryptionContext, data: bytes) -> bytes:
|
||||
return data
|
||||
|
||||
@encryption.decrypt.json
|
||||
async def json_dec(
|
||||
_ctx: EncryptionContext, data: dict[str, object]
|
||||
) -> dict[str, object]:
|
||||
return data
|
||||
|
||||
blob_handler: Callable[[EncryptionContext, bytes], Awaitable[bytes]] = blob_dec
|
||||
json_handler: Callable[
|
||||
[EncryptionContext, dict[str, object]], Awaitable[dict[str, object]]
|
||||
] = json_dec
|
||||
assert blob_handler is blob_dec
|
||||
assert json_handler is json_dec
|
||||
|
||||
|
||||
class TestHandlerValidation:
|
||||
"""Test duplicate handler and signature validation."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user