codekingpro/portable-devtools

1from __future__ import annotations2 3from collections.abc import Sequence4from dataclasses import dataclass5from typing import Any, Generic, Protocol, TypeAlias6 7from langchain_core.runnables import Runnable, RunnableConfig8from langgraph.store.base import BaseStore9 10from langgraph._internal._typing import EMPTY_SEQ11from langgraph.runtime import Runtime12from langgraph.types import CachePolicy, RetryPolicy, StreamWriter, TimeoutPolicy13from langgraph.typing import ContextT, NodeInputT, NodeInputT_contra14 15 16class _Node(Protocol[NodeInputT_contra]):17    def __call__(self, state: NodeInputT_contra) -> Any: ...18 19 20class _NodeWithConfig(Protocol[NodeInputT_contra]):21    def __call__(self, state: NodeInputT_contra, config: RunnableConfig) -> Any: ...22 23 24class _NodeWithWriter(Protocol[NodeInputT_contra]):25    def __call__(self, state: NodeInputT_contra, *, writer: StreamWriter) -> Any: ...26 27 28class _NodeWithStore(Protocol[NodeInputT_contra]):29    def __call__(self, state: NodeInputT_contra, *, store: BaseStore) -> Any: ...30 31 32class _NodeWithWriterStore(Protocol[NodeInputT_contra]):33    def __call__(34        self, state: NodeInputT_contra, *, writer: StreamWriter, store: BaseStore35    ) -> Any: ...36 37 38class _NodeWithConfigWriter(Protocol[NodeInputT_contra]):39    def __call__(40        self, state: NodeInputT_contra, *, config: RunnableConfig, writer: StreamWriter41    ) -> Any: ...42 43 44class _NodeWithConfigStore(Protocol[NodeInputT_contra]):45    def __call__(46        self, state: NodeInputT_contra, *, config: RunnableConfig, store: BaseStore47    ) -> Any: ...48 49 50class _NodeWithConfigWriterStore(Protocol[NodeInputT_contra]):51    def __call__(52        self,53        state: NodeInputT_contra,54        *,55        config: RunnableConfig,56        writer: StreamWriter,57        store: BaseStore,58    ) -> Any: ...59 60 61class _NodeWithRuntime(Protocol[NodeInputT_contra, ContextT]):62    def __call__(63        self, state: NodeInputT_contra, *, runtime: Runtime[ContextT]64    ) -> Any: ...65 66 67# TODO: we probably don't want to explicitly support the config / store signatures once68# we move to adding a context arg. Maybe what we do is we add support for kwargs with param spec69# this is purely for typing purposes though, so can easily change in the coming weeks.70StateNode: TypeAlias = (71    _Node[NodeInputT]72    | _NodeWithConfig[NodeInputT]73    | _NodeWithWriter[NodeInputT]74    | _NodeWithStore[NodeInputT]75    | _NodeWithWriterStore[NodeInputT]76    | _NodeWithConfigWriter[NodeInputT]77    | _NodeWithConfigStore[NodeInputT]78    | _NodeWithConfigWriterStore[NodeInputT]79    | _NodeWithRuntime[NodeInputT, ContextT]80    | Runnable[NodeInputT, Any]81)82 83 84@dataclass(slots=True)85class StateNodeSpec(Generic[NodeInputT, ContextT]):86    runnable: StateNode[NodeInputT, ContextT]87    metadata: dict[str, Any] | None88    input_schema: type[NodeInputT]89    retry_policy: RetryPolicy | Sequence[RetryPolicy] | None90    cache_policy: CachePolicy | None91    is_error_handler: bool = False92    error_handler_node: str | None = None93    ends: tuple[str, ...] | dict[str, str] | None = EMPTY_SEQ94    defer: bool = False95    timeout: TimeoutPolicy | None = None96