coco-framework 0.1.0__py3-none-any.whl

coco_graph/checkpoint/sqlite/utils.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import json
+import re
+from collections.abc import Sequence
+from typing import Any
+
+from coco_core.runnables import RunnableConfig
+from coco_graph.checkpoint.base import get_checkpoint_id
+
+_FILTER_PATTERN = re.compile(r"^[a-zA-Z0-9_.-]+$")
+
+
+def _validate_filter_key(key: str) -> None:
+    """Validate that a filter key is safe for use in SQL queries.
+
+    Args:
+        key: The filter key to validate
+
+    Raises:
+        ValueError: If the key contains invalid characters that could enable SQL injection
+    """
+    # Allow alphanumeric characters, underscores, dots, and hyphens
+    # This covers typical JSON property names while preventing SQL injection
+    if not _FILTER_PATTERN.match(key):
+        raise ValueError(
+            f"Invalid filter key: '{key}'. Filter keys must contain only alphanumeric characters, underscores, dots, and hyphens."
+        )
+
+
+def _metadata_predicate(
+    metadata_filter: dict[str, Any],
+) -> tuple[Sequence[str], Sequence[Any]]:
+    """Return WHERE clause predicates for (a)search() given metadata filter.
+
+    This method returns a tuple of a string and a tuple of values. The string
+    is the parametered WHERE clause predicate (excluding the WHERE keyword):
+    "column1 = ? AND column2 IS ?". The tuple of values contains the values
+    for each of the corresponding parameters.
+    """
+
+    def _where_value(query_value: Any) -> tuple[str, Any]:
+        """Return tuple of operator and value for WHERE clause predicate."""
+        if query_value is None:
+            return ("IS ?", None)
+        elif (
+            isinstance(query_value, str)
+            or isinstance(query_value, int)
+            or isinstance(query_value, float)
+        ):
+            return ("= ?", query_value)
+        elif isinstance(query_value, bool):
+            return ("= ?", 1 if query_value else 0)
+        elif isinstance(query_value, dict) or isinstance(query_value, list):
+            # query value for JSON object cannot have trailing space after separators (, :)
+            # SQLite json_extract() returns JSON string without whitespace
+            return ("= ?", json.dumps(query_value, separators=(",", ":")))
+        else:
+            return ("= ?", str(query_value))
+
+    predicates = []
+    param_values = []
+
+    # process metadata query
+    for query_key, query_value in metadata_filter.items():
+        _validate_filter_key(query_key)
+        operator, param_value = _where_value(query_value)
+        predicates.append(
+            f"json_extract(CAST(metadata AS TEXT), '$.{query_key}') {operator}"
+        )
+        param_values.append(param_value)
+
+    return (predicates, param_values)
+
+
+def search_where(
+    config: RunnableConfig | None,
+    filter: dict[str, Any] | None,
+    before: RunnableConfig | None = None,
+) -> tuple[str, Sequence[Any]]:
+    """Return WHERE clause predicates for (a)search() given metadata filter
+    and `before` config.
+
+    This method returns a tuple of a string and a tuple of values. The string
+    is the parametered WHERE clause predicate (including the WHERE keyword):
+    "WHERE column1 = ? AND column2 IS ?". The tuple of values contains the
+    values for each of the corresponding parameters.
+    """
+    wheres = []
+    param_values = []
+
+    # construct predicate for config filter
+    if config is not None:
+        wheres.append("thread_id = ?")
+        param_values.append(config["configurable"]["thread_id"])
+        checkpoint_ns = config["configurable"].get("checkpoint_ns")
+        if checkpoint_ns is not None:
+            wheres.append("checkpoint_ns = ?")
+            param_values.append(checkpoint_ns)
+
+        if checkpoint_id := get_checkpoint_id(config):
+            wheres.append("checkpoint_id = ?")
+            param_values.append(checkpoint_id)
+
+    # construct predicate for metadata filter
+    if filter:
+        metadata_predicates, metadata_values = _metadata_predicate(filter)
+        wheres.extend(metadata_predicates)
+        param_values.extend(metadata_values)
+
+    # construct predicate for `before`
+    if before is not None:
+        wheres.append("checkpoint_id < ?")
+        param_values.append(get_checkpoint_id(before))
+
+    return ("WHERE " + " AND ".join(wheres) if wheres else "", param_values)

coco_graph/config.py

@@ -0,0 +1,196 @@
+import asyncio
+import sys
+from typing import Any
+
+from coco_core.runnables import RunnableConfig
+from coco_core.runnables.config import var_child_runnable_config
+from coco_graph.store.base import BaseStore
+
+from coco_graph._internal._constants import CONF, CONFIG_KEY_RUNTIME
+from coco_graph.types import StreamWriter
+
+
+def _no_op_stream_writer(c: Any) -> None:
+    pass
+
+
+def get_config() -> RunnableConfig:
+    if sys.version_info < (3, 11):
+        try:
+            if asyncio.current_task():
+                raise RuntimeError(
+                    "Python 3.11 or later required to use this in an async context"
+                )
+        except RuntimeError:
+            pass
+    if var_config := var_child_runnable_config.get():
+        return var_config
+    else:
+        raise RuntimeError("Called get_config outside of a runnable context")
+
+
+def get_store() -> BaseStore:
+    """Access CocoGraph store from inside a graph node or entrypoint task at runtime.
+
+    Can be called from inside any [`StateGraph`][coco_graph.graph.StateGraph] node or
+    functional API [`task`][coco_graph.func.task], as long as the `StateGraph` or the [`entrypoint`][coco_graph.func.entrypoint]
+    was initialized with a store, e.g.:
+
+    ```python
+    # with StateGraph
+    graph = (
+        StateGraph(...)
+        ...
+        .compile(store=store)
+    )
+
+    # or with entrypoint
+    @entrypoint(store=store)
+    def workflow(inputs):
+        ...
+    ```
+
+    !!! warning "Async with Python < 3.11"
+
+        If you are using Python < 3.11 and are running CocoGraph asynchronously,
+        `get_store()` won't work since it uses [`contextvar`](https://docs.python.org/3/library/contextvars.html) propagation (only available in [Python >= 3.11](https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task)).
+
+
+    Example: Using with `StateGraph`
+        ```python
+        from typing_extensions import TypedDict
+        from coco_graph.graph import StateGraph, START
+        from coco_graph.store.memory import InMemoryStore
+        from coco_graph.config import get_store
+
+        store = InMemoryStore()
+        store.put(("values",), "foo", {"bar": 2})
+
+
+        class State(TypedDict):
+            foo: int
+
+
+        def my_node(state: State):
+            my_store = get_store()
+            stored_value = my_store.get(("values",), "foo").value["bar"]
+            return {"foo": stored_value + 1}
+
+
+        graph = (
+            StateGraph(State)
+            .add_node(my_node)
+            .add_edge(START, "my_node")
+            .compile(store=store)
+        )
+
+        graph.invoke({"foo": 1})
+        ```
+
+        ```pycon
+        {"foo": 3}
+        ```
+
+    Example: Using with functional API
+        ```python
+        from coco_graph.func import entrypoint, task
+        from coco_graph.store.memory import InMemoryStore
+        from coco_graph.config import get_store
+
+        store = InMemoryStore()
+        store.put(("values",), "foo", {"bar": 2})
+
+
+        @task
+        def my_task(value: int):
+            my_store = get_store()
+            stored_value = my_store.get(("values",), "foo").value["bar"]
+            return stored_value + 1
+
+
+        @entrypoint(store=store)
+        def workflow(value: int):
+            return my_task(value).result()
+
+
+        workflow.invoke(1)
+        ```
+
+        ```pycon
+        3
+        ```
+    """
+    return get_config()[CONF][CONFIG_KEY_RUNTIME].store
+
+
+def get_stream_writer() -> StreamWriter:
+    """Access CocoGraph [`StreamWriter`][coco_graph.types.StreamWriter] from inside a graph node or entrypoint task at runtime.
+
+    Can be called from inside any [`StateGraph`][coco_graph.graph.StateGraph] node or
+    functional API [`task`][coco_graph.func.task].
+
+    !!! warning "Async with Python < 3.11"
+
+        If you are using Python < 3.11 and are running CocoGraph asynchronously,
+        `get_stream_writer()` won't work since it uses [`contextvar`](https://docs.python.org/3/library/contextvars.html) propagation (only available in [Python >= 3.11](https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task)).
+
+    Example: Using with `StateGraph`
+        ```python
+        from typing_extensions import TypedDict
+        from coco_graph.graph import StateGraph, START
+        from coco_graph.config import get_stream_writer
+
+
+        class State(TypedDict):
+            foo: int
+
+
+        def my_node(state: State):
+            my_stream_writer = get_stream_writer()
+            my_stream_writer({"custom_data": "Hello!"})
+            return {"foo": state["foo"] + 1}
+
+
+        graph = (
+            StateGraph(State)
+            .add_node(my_node)
+            .add_edge(START, "my_node")
+            .compile(store=store)
+        )
+
+        for chunk in graph.stream({"foo": 1}, stream_mode="custom"):
+            print(chunk)
+        ```
+
+        ```pycon
+        {"custom_data": "Hello!"}
+        ```
+
+    Example: Using with functional API
+        ```python
+        from coco_graph.func import entrypoint, task
+        from coco_graph.config import get_stream_writer
+
+
+        @task
+        def my_task(value: int):
+            my_stream_writer = get_stream_writer()
+            my_stream_writer({"custom_data": "Hello!"})
+            return value + 1
+
+
+        @entrypoint(store=store)
+        def workflow(value: int):
+            return my_task(value).result()
+
+
+        for chunk in workflow.stream(1, stream_mode="custom"):
+            print(chunk)
+        ```
+
+        ```pycon
+        {"custom_data": "Hello!"}
+        ```
+    """
+    runtime = get_config()[CONF][CONFIG_KEY_RUNTIME]
+    return runtime.stream_writer

coco_graph/constants.py

@@ -0,0 +1,64 @@
+import sys
+from typing import Any
+from warnings import warn
+
+from coco_graph._internal._constants import (
+    CONF,
+    CONFIG_KEY_CHECKPOINTER,
+    TASKS,
+)
+from coco_graph.warnings import CocoGraphDeprecatedSinceV10
+
+__all__ = (
+    "TAG_NOSTREAM",
+    "TAG_HIDDEN",
+    "START",
+    "END",
+    # retained for backwards compatibility (mostly trustcoco-graph-api), should be removed in v2 (or earlier)
+    "CONF",
+    "TASKS",
+    "CONFIG_KEY_CHECKPOINTER",
+)
+
+# --- Public constants ---
+TAG_NOSTREAM = sys.intern("nostream")
+"""Tag to disable streaming for a chat model."""
+TAG_HIDDEN = sys.intern("langsmith:hidden")
+"""Tag to hide a node/edge from certain tracing/streaming environments."""
+END = sys.intern("__end__")
+"""The last (maybe virtual) node in graph-style Pregel."""
+START = sys.intern("__start__")
+"""The first (maybe virtual) node in graph-style Pregel."""
+
+
+def __getattr__(name: str) -> Any:
+    if name in ["Send", "Interrupt"]:
+        warn(
+            f"Importing {name} from coco_graph.constants is deprecated. "
+            f"Please use 'from coco_graph.types import {name}' instead.",
+            CocoGraphDeprecatedSinceV10,
+            stacklevel=2,
+        )
+
+        from importlib import import_module
+
+        module = import_module("coco_graph.types")
+        return getattr(module, name)
+
+    try:
+        from importlib import import_module
+
+        private_constants = import_module("coco_graph._internal._constants")
+        attr = getattr(private_constants, name)
+        warn(
+            f"Importing {name} from coco_graph.constants is deprecated. "
+            f"This constant is now private and should not be used directly. "
+            "Please let the CocoGraph team know if you need this value.",
+            CocoGraphDeprecatedSinceV10,
+            stacklevel=2,
+        )
+        return attr
+    except AttributeError:
+        pass
+
+    raise AttributeError(f"module has no attribute '{name}'")

coco_graph/errors.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Literal
+from warnings import warn
+
+# EmptyChannelError is re-exported from coco_graph.channels.base
+from coco_graph.checkpoint.base import EmptyChannelError
+from typing_extensions import deprecated
+
+from coco_graph.types import Command, Interrupt
+from coco_graph.warnings import CocoGraphDeprecatedSinceV10
+
+__all__ = (
+    "EmptyChannelError",
+    "ErrorCode",
+    "GraphDrained",
+    "GraphRecursionError",
+    "InvalidUpdateError",
+    "GraphBubbleUp",
+    "GraphInterrupt",
+    "NodeCancelledError",
+    "NodeError",
+    "NodeInterrupt",
+    "NodeTimeoutError",
+    "ParentCommand",
+    "EmptyInputError",
+    "TaskNotFound",
+)
+
+
+class ErrorCode(Enum):
+    GRAPH_RECURSION_LIMIT = "GRAPH_RECURSION_LIMIT"
+    INVALID_CONCURRENT_GRAPH_UPDATE = "INVALID_CONCURRENT_GRAPH_UPDATE"
+    INVALID_GRAPH_NODE_RETURN_VALUE = "INVALID_GRAPH_NODE_RETURN_VALUE"
+    MULTIPLE_SUBGRAPHS = "MULTIPLE_SUBGRAPHS"
+    INVALID_CHAT_HISTORY = "INVALID_CHAT_HISTORY"
+
+
+def create_error_message(*, message: str, error_code: ErrorCode) -> str:
+    return (
+        f"{message}\n"
+        "For troubleshooting, visit: https://docs.langchain.com/oss/python/langgraph/"
+        f"errors/{error_code.value}"
+    )
+
+
+class GraphBubbleUp(Exception):
+    pass
+
+
+class GraphDrained(GraphBubbleUp):
+    """Raised when a graph run exits early due to a drain request.
+
+    This indicates the graph stopped cooperatively at a superstep boundary
+    because `RunControl.request_drain()` was called (e.g., in response to
+    SIGTERM). The checkpoint is saved and the run can be resumed later.
+    """
+
+    def __init__(self, reason: str = "shutdown") -> None:
+        self.reason = reason
+        super().__init__(f"Graph drained: {reason}")
+
+
+class GraphRecursionError(RecursionError):
+    """Raised when the graph has exhausted the maximum number of steps.
+
+    This prevents infinite loops. To increase the maximum number of steps,
+    run your graph with a config specifying a higher `recursion_limit`.
+
+    Troubleshooting guides:
+
+    - [`GRAPH_RECURSION_LIMIT`](https://docs.langchain.com/oss/python/langgraph/GRAPH_RECURSION_LIMIT)
+
+    Examples:
+
+        graph = builder.compile()
+        graph.invoke(
+            {"messages": [("user", "Hello, world!")]},
+            # The config is the second positional argument
+            {"recursion_limit": 1000},
+        )
+    """
+
+    pass
+
+
+class InvalidUpdateError(Exception):
+    """Raised when attempting to update a channel with an invalid set of updates.
+
+    Troubleshooting guides:
+
+    - [`INVALID_CONCURRENT_GRAPH_UPDATE`](https://docs.langchain.com/oss/python/langgraph/INVALID_CONCURRENT_GRAPH_UPDATE)
+    - [`INVALID_GRAPH_NODE_RETURN_VALUE`](https://docs.langchain.com/oss/python/langgraph/INVALID_GRAPH_NODE_RETURN_VALUE)
+    """
+
+    pass
+
+
+class GraphInterrupt(GraphBubbleUp):
+    """Raised when a subgraph is interrupted, suppressed by the root graph.
+    Never raised directly, or surfaced to the user."""
+
+    def __init__(self, interrupts: Sequence[Interrupt] = ()) -> None:
+        super().__init__(interrupts)
+
+
+@deprecated(
+    "NodeInterrupt is deprecated. Please use [`interrupt`][coco_graph.types.interrupt] instead.",
+    category=None,
+)
+class NodeInterrupt(GraphInterrupt):
+    """Raised by a node to interrupt execution."""
+
+    def __init__(self, value: Any, id: str | None = None) -> None:
+        warn(
+            "NodeInterrupt is deprecated. Please use `coco_graph.types.interrupt` instead.",
+            CocoGraphDeprecatedSinceV10,
+            stacklevel=2,
+        )
+        if id is None:
+            super().__init__([Interrupt(value=value)])
+        else:
+            super().__init__([Interrupt(value=value, id=id)])
+
+
+class ParentCommand(GraphBubbleUp):
+    args: tuple[Command]
+
+    def __init__(self, command: Command) -> None:
+        super().__init__(command)
+
+
+class EmptyInputError(Exception):
+    """Raised when graph receives an empty input."""
+
+    pass
+
+
+class TaskNotFound(Exception):
+    """Raised when the executor is unable to find a task (for distributed mode)."""
+
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class NodeError:
+    """Failure context passed to a node-level error handler.
+
+    Inject by adding a parameter typed `NodeError` to a handler registered via
+    `StateGraph.add_node(..., error_handler=...)`:
+
+    ```python
+    def handler(state: State, error: NodeError) -> Command:
+        return Command(update={"status": f"recovered from {error.node}: {error.error}"})
+    ```
+    """
+
+    node: str
+    """Name of the node whose execution failed."""
+
+    error: BaseException
+    """Exception raised by the failed node."""
+
+
+class NodeCancelledError(Exception):
+    """Raised when a node body raises ``asyncio.CancelledError`` itself.
+
+    ``asyncio.CancelledError`` is a ``BaseException`` and the pregel runner
+    treats cancelled task futures as silent tear-down (e.g. when it stops
+    sibling tasks after a peer fails). That is the correct behaviour for
+    *framework-initiated* cancellation, but a user node that raises
+    ``asyncio.CancelledError`` from its own body should surface as a node
+    failure, the same way any other exception would.
+
+    The retry layer converts user-raised ``asyncio.CancelledError`` into this
+    type so it flows through the normal error path and the run reports as
+    ``error`` instead of silently succeeding.
+    """
+
+    node: str
+
+    def __init__(self, node: str, message: str | None = None) -> None:
+        super().__init__(message or f"Node {node!r} raised asyncio.CancelledError")
+        self.node = node
+
+
+class NodeTimeoutError(Exception):
+    """Raised when a node invocation exceeds one of its configured timeouts.
+
+    Does **not** inherit from the built-in `TimeoutError` (a subclass of
+    `OSError`) so that the default `RetryPolicy` treats it as retryable.
+
+    Both `idle_timeout` and `run_timeout` reflect the configured policy at the
+    time of the failure (each is `None` if not configured). `kind` and
+    `timeout` identify which one fired.
+    """
+
+    node: str
+    timeout: float
+    run_timeout: float | None
+    idle_timeout: float | None
+    elapsed: float
+    kind: Literal["idle", "run"]
+
+    def __init__(
+        self,
+        node: str,
+        elapsed: float,
+        *,
+        kind: Literal["idle", "run"],
+        idle_timeout: float | None = None,
+        run_timeout: float | None = None,
+    ) -> None:
+        if kind == "idle":
+            if idle_timeout is None:
+                raise ValueError("idle_timeout is required when kind='idle'")
+            message = (
+                f"Node '{node}' exceeded its idle timeout of "
+                f"{idle_timeout:.3f}s without making progress "
+                f"(elapsed: {elapsed:.3f}s)."
+            )
+            self.timeout = idle_timeout
+        elif kind == "run":
+            if run_timeout is None:
+                raise ValueError("run_timeout is required when kind='run'")
+            message = (
+                f"Node '{node}' exceeded its run timeout of "
+                f"{run_timeout:.3f}s (elapsed: {elapsed:.3f}s)."
+            )
+            self.timeout = run_timeout
+        else:
+            raise ValueError("kind must be 'idle' or 'run'")
+        super().__init__(message)
+        self.node = node
+        self.elapsed = elapsed
+        self.kind = kind
+        self.idle_timeout = idle_timeout
+        self.run_timeout = run_timeout