coco-framework 0.1.0__py3-none-any.whl

coco_graph/func/__init__.py

@@ -0,0 +1,620 @@
+from __future__ import annotations
+
+import functools
+import inspect
+import warnings
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import (
+    Any,
+    Generic,
+    TypeVar,
+    cast,
+    get_args,
+    get_origin,
+    overload,
+)
+
+from coco_graph.cache.base import BaseCache
+from coco_graph.checkpoint.base import BaseCheckpointSaver
+from coco_graph.store.base import BaseStore
+from typing_extensions import Unpack
+
+from coco_graph._internal import _serde
+from coco_graph._internal._constants import CACHE_NS_WRITES, PREVIOUS
+from coco_graph._internal._runnable import is_async_callable
+from coco_graph._internal._timeout import (
+    coerce_timeout_policy,
+    sync_timeout_unsupported,
+)
+from coco_graph._internal._typing import MISSING, DeprecatedKwargs
+from coco_graph.channels.ephemeral_value import EphemeralValue
+from coco_graph.channels.last_value import LastValue
+from coco_graph.constants import END, START
+from coco_graph.pregel import Pregel
+from coco_graph.pregel._call import (
+    P,
+    SyncAsyncFuture,
+    T,
+    _call_with_options,
+    get_runnable_for_entrypoint,
+    identifier,
+)
+from coco_graph.pregel._read import PregelNode
+from coco_graph.pregel._write import ChannelWrite, ChannelWriteEntry
+from coco_graph.types import (
+    _DC_KWARGS,
+    CachePolicy,
+    RetryPolicy,
+    StreamMode,
+    TimeoutPolicy,
+)
+from coco_graph.typing import ContextT
+from coco_graph.warnings import CocoGraphDeprecatedSinceV05, CocoGraphDeprecatedSinceV10
+
+__all__ = ("task", "entrypoint")
+
+
+class _TaskFunction(Generic[P, T]):
+    def __init__(
+        self,
+        func: Callable[P, Awaitable[T]] | Callable[P, T],
+        *,
+        retry_policy: Sequence[RetryPolicy],
+        cache_policy: CachePolicy[Callable[P, str | bytes]] | None = None,
+        timeout: TimeoutPolicy | None = None,
+        name: str | None = None,
+    ) -> None:
+        if name is not None:
+            if hasattr(func, "__func__"):
+                # handle class methods
+                # NOTE: we're modifying the instance method to avoid modifying
+                # the original class method in case it's shared across multiple tasks
+                instance_method = functools.partial(func.__func__, func.__self__)  # type: ignore [union-attr]
+                instance_method.__name__ = name  # type: ignore [attr-defined]
+                func = instance_method
+            else:
+                # handle regular functions / partials / callable classes, etc.
+                func.__name__ = name
+        self.func = func
+        self.retry_policy = retry_policy
+        self.cache_policy = cache_policy
+        self.timeout = timeout
+        functools.update_wrapper(self, func)
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> SyncAsyncFuture[T]:
+        return _call_with_options(
+            self.func,
+            args,
+            kwargs,
+            retry_policy=self.retry_policy,
+            cache_policy=self.cache_policy,
+            timeout=self.timeout,
+        )
+
+    def clear_cache(self, cache: BaseCache) -> None:
+        """Clear the cache for this task."""
+        if self.cache_policy is not None:
+            cache.clear(((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),))
+
+    async def aclear_cache(self, cache: BaseCache) -> None:
+        """Clear the cache for this task."""
+        if self.cache_policy is not None:
+            await cache.aclear(
+                ((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),)
+            )
+
+
+@overload
+def task(
+    __func_or_none__: None = None,
+    *,
+    name: str | None = None,
+    retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
+    cache_policy: CachePolicy[Callable[P, str | bytes]] | None = None,
+    timeout: float | timedelta | TimeoutPolicy | None = None,
+    **kwargs: Unpack[DeprecatedKwargs],
+) -> Callable[
+    [Callable[P, Awaitable[T]] | Callable[P, T]],
+    _TaskFunction[P, T],
+]: ...
+
+
+@overload
+def task(__func_or_none__: Callable[P, Awaitable[T]]) -> _TaskFunction[P, T]: ...
+
+
+@overload
+def task(__func_or_none__: Callable[P, T]) -> _TaskFunction[P, T]: ...
+
+
+def task(
+    __func_or_none__: Callable[P, Awaitable[T]] | Callable[P, T] | None = None,
+    *,
+    name: str | None = None,
+    retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
+    cache_policy: CachePolicy[Callable[P, str | bytes]] | None = None,
+    timeout: float | timedelta | TimeoutPolicy | None = None,
+    **kwargs: Unpack[DeprecatedKwargs],
+) -> (
+    Callable[[Callable[P, Awaitable[T]] | Callable[P, T]], _TaskFunction[P, T]]
+    | _TaskFunction[P, T]
+):
+    """Define a CocoGraph task using the `task` decorator.
+
+    !!! important "Requires python 3.11 or higher for async functions"
+        The `task` decorator supports both sync and async functions. To use async
+        functions, ensure that you are using Python 3.11 or higher.
+
+    Tasks can only be called from within an [`entrypoint`][coco_graph.func.entrypoint] or
+    from within a `StateGraph`. A task can be called like a regular function with the
+    following differences:
+
+    - When a checkpointer is enabled, the function inputs and outputs must be serializable.
+    - The decorated function can only be called from within an entrypoint or `StateGraph`.
+    - Calling the function produces a future. This makes it easy to parallelize tasks.
+
+    Args:
+        name: An optional name for the task. If not provided, the function name will be used.
+        retry_policy: An optional retry policy (or list of policies) to use for the task in case of a failure.
+        cache_policy: An optional cache policy to use for the task. This allows caching of the task results.
+        timeout: Timeout for each task attempt. A number or `timedelta` is a hard
+            wall-clock cap and is not refreshed. Use `TimeoutPolicy` to configure
+            both a wall-clock `run_timeout` and an `idle_timeout` refreshed by
+            progress signals. For long-running work that doesn't naturally emit
+            progress, call `runtime.heartbeat()` from inside the task. When the
+            timeout fires, `NodeTimeoutError` is raised and the retry policy (if
+            any) decides whether to retry. Supported only for async tasks; sync
+            tasks cannot be safely cancelled in-process.
+
+    Returns:
+        A callable function when used as a decorator.
+
+    Example: Sync Task
+        ```python
+        from coco_graph.func import entrypoint, task
+
+
+        @task
+        def add_one_task(a: int) -> int:
+            return a + 1
+
+
+        @entrypoint()
+        def add_one(numbers: list[int]) -> list[int]:
+            futures = [add_one_task(n) for n in numbers]
+            results = [f.result() for f in futures]
+            return results
+
+
+        # Call the entrypoint
+        add_one.invoke([1, 2, 3])  # Returns [2, 3, 4]
+        ```
+
+    Example: Async Task
+        ```python
+        import asyncio
+        from coco_graph.func import entrypoint, task
+
+
+        @task
+        async def add_one_task(a: int) -> int:
+            return a + 1
+
+
+        @entrypoint()
+        async def add_one(numbers: list[int]) -> list[int]:
+            futures = [add_one_task(n) for n in numbers]
+            return asyncio.gather(*futures)
+
+
+        # Call the entrypoint
+        await add_one.ainvoke([1, 2, 3])  # Returns [2, 3, 4]
+        ```
+    """
+    if (retry := kwargs.get("retry", MISSING)) is not MISSING:
+        warnings.warn(
+            "`retry` is deprecated and will be removed. Please use `retry_policy` instead.",
+            category=CocoGraphDeprecatedSinceV05,
+            stacklevel=2,
+        )
+        if retry_policy is None:
+            retry_policy = retry  # type: ignore[assignment]
+    timeout_policy = coerce_timeout_policy(timeout)
+
+    retry_policies: Sequence[RetryPolicy] = (
+        ()
+        if retry_policy is None
+        else (retry_policy,)
+        if isinstance(retry_policy, RetryPolicy)
+        else retry_policy
+    )
+
+    def decorator(
+        func: Callable[P, Awaitable[T]] | Callable[P, T],
+    ) -> Callable[P, SyncAsyncFuture[T]]:
+        if timeout_policy is not None and not is_async_callable(func):
+            name_ = name or getattr(func, "__name__", func.__class__.__name__)
+            raise sync_timeout_unsupported(str(name_), kind="Task")
+        return _TaskFunction(
+            func,
+            retry_policy=retry_policies,
+            cache_policy=cache_policy,
+            timeout=timeout_policy,
+            name=name,
+        )
+
+    if __func_or_none__ is not None:
+        return decorator(__func_or_none__)
+
+    return decorator
+
+
+R = TypeVar("R")
+S = TypeVar("S")
+
+
+# The decorator was wrapped in a class to support the `final` attribute.
+# In this form, the `final` attribute should play nicely with IDE autocompletion,
+# and type checking tools.
+# In addition, we'll be able to surface this information in the API Reference.
+class entrypoint(Generic[ContextT]):
+    """Define a CocoGraph workflow using the `entrypoint` decorator.
+
+    ### Function signature
+
+    The decorated function must accept a **single parameter**, which serves as the input
+    to the function. This input parameter can be of any type. Use a dictionary
+    to pass **multiple parameters** to the function.
+
+    ### Injectable parameters
+
+    The decorated function can request access to additional parameters
+    that will be injected automatically at run time. These parameters include:
+
+    | Parameter        | Description                                                                                          |
+    |------------------|------------------------------------------------------------------------------------------------------|
+    | **`config`**     | A configuration object (aka `RunnableConfig`) that holds run-time configuration values.              |
+    | **`previous`**   | The previous return value for the given thread (available only when a checkpointer is provided).     |
+    | **`runtime`**    | A `Runtime` object that contains information about the current run, including context, store, writer |
+
+    The entrypoint decorator can be applied to sync functions or async functions.
+
+    ### State management
+
+    The **`previous`** parameter can be used to access the return value of the previous
+    invocation of the entrypoint on the same thread id. This value is only available
+    when a checkpointer is provided.
+
+    If you want **`previous`** to be different from the return value, you can use the
+    `entrypoint.final` object to return a value while saving a different value to the
+    checkpoint.
+
+    Args:
+        checkpointer: Specify a checkpointer to create a workflow that can persist
+            its state across runs.
+        store: A generalized key-value store. Some implementations may support
+            semantic search capabilities through an optional `index` configuration.
+        cache: A cache to use for caching the results of the workflow.
+        context_schema: Specifies the schema for the context object that will be
+            passed to the workflow.
+        cache_policy: A cache policy to use for caching the results of the workflow.
+        retry_policy: A retry policy (or list of policies) to use for the workflow in case of a failure.
+        timeout: Timeout for each workflow attempt. A number or `timedelta` is a
+            hard wall-clock cap and is not refreshed. Use `TimeoutPolicy` to
+            configure both a wall-clock `run_timeout` and an `idle_timeout`
+            refreshed by progress signals. For long-running work that doesn't
+            naturally emit progress, call `runtime.heartbeat()` from inside the
+            workflow. When the timeout fires, `NodeTimeoutError` is raised and
+            the retry policy (if any) decides whether to retry. Supported only
+            for async workflows; sync workflows cannot be safely cancelled
+            in-process.
+
+    !!! warning "`config_schema` Deprecated"
+        The `config_schema` parameter is deprecated in v0.6.0 and support will be removed in v2.0.0.
+        Please use `context_schema` instead to specify the schema for run-scoped context.
+
+
+    Example: Using entrypoint and tasks
+        ```python
+        import time
+
+        from coco_graph.func import entrypoint, task
+        from coco_graph.types import interrupt, Command
+        from coco_graph.checkpoint.memory import InMemorySaver
+
+        @task
+        def compose_essay(topic: str) -> str:
+            time.sleep(1.0)  # Simulate slow operation
+            return f"An essay about {topic}"
+
+        @entrypoint(checkpointer=InMemorySaver())
+        def review_workflow(topic: str) -> dict:
+            \"\"\"Manages the workflow for generating and reviewing an essay.
+
+            The workflow includes:
+            1. Generating an essay about the given topic.
+            2. Interrupting the workflow for human review of the generated essay.
+
+            Upon resuming the workflow, compose_essay task will not be re-executed
+            as its result is cached by the checkpointer.
+
+            Args:
+                topic: The subject of the essay.
+
+            Returns:
+                dict: A dictionary containing the generated essay and the human review.
+            \"\"\"
+            essay_future = compose_essay(topic)
+            essay = essay_future.result()
+            human_review = interrupt({
+                \"question\": \"Please provide a review\",
+                \"essay\": essay
+            })
+            return {
+                \"essay\": essay,
+                \"review\": human_review,
+            }
+
+        # Example configuration for the workflow
+        config = {
+            \"configurable\": {
+                \"thread_id\": \"some_thread\"
+            }
+        }
+
+        # Topic for the essay
+        topic = \"cats\"
+
+        # Stream the workflow to generate the essay and await human review
+        for result in review_workflow.stream(topic, config):
+            print(result)
+
+        # Example human review provided after the interrupt
+        human_review = \"This essay is great.\"
+
+        # Resume the workflow with the provided human review
+        for result in review_workflow.stream(Command(resume=human_review), config):
+            print(result)
+        ```
+
+    Example: Accessing the previous return value
+        When a checkpointer is enabled the function can access the previous return value
+        of the previous invocation on the same thread id.
+
+        ```python
+        from typing import Optional
+
+        from coco_graph.checkpoint.memory import MemorySaver
+
+        from coco_graph.func import entrypoint
+
+
+        @entrypoint(checkpointer=InMemorySaver())
+        def my_workflow(input_data: str, previous: Optional[str] = None) -> str:
+            return "world"
+
+
+        config = {"configurable": {"thread_id": "some_thread"}}
+        my_workflow.invoke("hello", config)
+        ```
+
+    Example: Using `entrypoint.final` to save a value
+        The `entrypoint.final` object allows you to return a value while saving
+        a different value to the checkpoint. This value will be accessible
+        in the next invocation of the entrypoint via the `previous` parameter, as
+        long as the same thread id is used.
+
+        ```python
+        from typing import Any
+
+        from coco_graph.checkpoint.memory import MemorySaver
+
+        from coco_graph.func import entrypoint
+
+
+        @entrypoint(checkpointer=InMemorySaver())
+        def my_workflow(
+            number: int,
+            *,
+            previous: Any = None,
+        ) -> entrypoint.final[int, int]:
+            previous = previous or 0
+            # This will return the previous value to the caller, saving
+            # 2 * number to the checkpoint, which will be used in the next invocation
+            # for the `previous` parameter.
+            return entrypoint.final(value=previous, save=2 * number)
+
+
+        config = {"configurable": {"thread_id": "some_thread"}}
+
+        my_workflow.invoke(3, config)  # 0 (previous was None)
+        my_workflow.invoke(1, config)  # 6 (previous was 3 * 2 from the previous invocation)
+        ```
+    """
+
+    def __init__(
+        self,
+        checkpointer: BaseCheckpointSaver | None = None,
+        store: BaseStore | None = None,
+        cache: BaseCache | None = None,
+        context_schema: type[ContextT] | None = None,
+        cache_policy: CachePolicy | None = None,
+        retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
+        timeout: float | timedelta | TimeoutPolicy | None = None,
+        **kwargs: Unpack[DeprecatedKwargs],
+    ) -> None:
+        """Initialize the entrypoint decorator."""
+        if (config_schema := kwargs.get("config_schema", MISSING)) is not MISSING:
+            warnings.warn(
+                "`config_schema` is deprecated and will be removed. Please use `context_schema` instead.",
+                category=CocoGraphDeprecatedSinceV10,
+                stacklevel=2,
+            )
+            if context_schema is None:
+                context_schema = cast(type[ContextT], config_schema)
+
+        if (retry := kwargs.get("retry", MISSING)) is not MISSING:
+            warnings.warn(
+                "`retry` is deprecated and will be removed. Please use `retry_policy` instead.",
+                category=CocoGraphDeprecatedSinceV05,
+                stacklevel=2,
+            )
+            if retry_policy is None:
+                retry_policy = cast("RetryPolicy | Sequence[RetryPolicy]", retry)
+
+        self.checkpointer = checkpointer
+        self.store = store
+        self.cache = cache
+        self.cache_policy = cache_policy
+        self.retry_policy = retry_policy
+        self.timeout = coerce_timeout_policy(timeout)
+        self.context_schema = context_schema
+
+    @dataclass(**_DC_KWARGS)
+    class final(Generic[R, S]):
+        """A primitive that can be returned from an entrypoint.
+
+        This primitive allows to save a value to the checkpointer distinct from the
+        return value from the entrypoint.
+
+        Example: Decoupling the return value and the save value
+            ```python
+            from coco_graph.checkpoint.memory import InMemorySaver
+            from coco_graph.func import entrypoint
+
+
+            @entrypoint(checkpointer=InMemorySaver())
+            def my_workflow(
+                number: int,
+                *,
+                previous: Any = None,
+            ) -> entrypoint.final[int, int]:
+                previous = previous or 0
+                # This will return the previous value to the caller, saving
+                # 2 * number to the checkpoint, which will be used in the next invocation
+                # for the `previous` parameter.
+                return entrypoint.final(value=previous, save=2 * number)
+
+
+            config = {"configurable": {"thread_id": "1"}}
+
+            my_workflow.invoke(3, config)  # 0 (previous was None)
+            my_workflow.invoke(1, config)  # 6 (previous was 3 * 2 from the previous invocation)
+            ```
+        """
+
+        value: R
+        """Value to return. A value will always be returned even if it is `None`."""
+        save: S
+        """The value for the state for the next checkpoint.
+
+        A value will always be saved even if it is `None`.
+        """
+
+    def __call__(self, func: Callable[..., Any]) -> Pregel:
+        """Convert a function into a Pregel graph.
+
+        Args:
+            func: The function to convert. Support both sync and async functions.
+
+        Returns:
+            A Pregel graph.
+        """
+        # wrap generators in a function that writes to StreamWriter
+        if inspect.isgeneratorfunction(func) or inspect.isasyncgenfunction(func):
+            raise NotImplementedError(
+                "Generators are not supported in the Functional API."
+            )
+
+        bound = get_runnable_for_entrypoint(func)
+        stream_mode: StreamMode = "updates"
+
+        # get input and output types
+        sig = inspect.signature(func)
+        first_parameter_name = next(iter(sig.parameters.keys()), None)
+        if not first_parameter_name:
+            raise ValueError("Entrypoint function must have at least one parameter")
+        input_type = (
+            sig.parameters[first_parameter_name].annotation
+            if sig.parameters[first_parameter_name].annotation
+            is not inspect.Signature.empty
+            else Any
+        )
+
+        def _pluck_return_value(value: Any) -> Any:
+            """Extract the return_ value the entrypoint.final object or passthrough."""
+            return value.value if isinstance(value, entrypoint.final) else value
+
+        def _pluck_save_value(value: Any) -> Any:
+            """Get save value from the entrypoint.final object or passthrough."""
+            return value.save if isinstance(value, entrypoint.final) else value
+
+        output_type, save_type = Any, Any
+        if sig.return_annotation is not inspect.Signature.empty:
+            # User does not parameterize entrypoint.final properly
+            if (
+                sig.return_annotation is entrypoint.final
+            ):  # Un-parameterized entrypoint.final
+                output_type = save_type = Any
+            else:
+                origin = get_origin(sig.return_annotation)
+                if origin is entrypoint.final:
+                    type_annotations = get_args(sig.return_annotation)
+                    if len(type_annotations) != 2:
+                        raise TypeError(
+                            "Please an annotation for both the return_ and "
+                            "the save values."
+                            "For example, `-> entrypoint.final[int, str]` would assign a "
+                            "return_ a type of `int` and save the type `str`."
+                        )
+                    output_type, save_type = get_args(sig.return_annotation)
+                else:
+                    output_type = save_type = sig.return_annotation
+
+        graph: Pregel[Any, ContextT, Any, Any] = Pregel(
+            nodes={
+                func.__name__: PregelNode(
+                    bound=bound,
+                    triggers=[START],
+                    channels=START,
+                    timeout=self.timeout,
+                    writers=[
+                        ChannelWrite(
+                            [
+                                ChannelWriteEntry(END, mapper=_pluck_return_value),
+                                ChannelWriteEntry(PREVIOUS, mapper=_pluck_save_value),
+                            ]
+                        )
+                    ],
+                )
+            },
+            channels={
+                START: EphemeralValue(input_type),
+                END: LastValue(output_type, END),
+                PREVIOUS: LastValue(save_type, PREVIOUS),
+            },
+            input_channels=START,
+            output_channels=END,
+            stream_channels=END,
+            stream_mode=stream_mode,
+            stream_eager=True,
+            checkpointer=self.checkpointer,
+            store=self.store,
+            cache=self.cache,
+            cache_policy=self.cache_policy,
+            retry_policy=self.retry_policy or (),
+            context_schema=self.context_schema,
+        )
+        if _serde.STRICT_MSGPACK_ENABLED:
+            serde_allowlist = _serde.build_serde_allowlist(
+                schemas=[input_type, output_type, save_type]
+                + ([self.context_schema] if self.context_schema is not None else []),
+                channels=graph.channels,
+            )
+            graph._serde_allowlist = serde_allowlist
+            graph.checkpointer = _serde.apply_checkpointer_allowlist(
+                graph.checkpointer, serde_allowlist
+            )
+        return graph

coco_graph/graph/__init__.py

@@ -0,0 +1,12 @@
+from coco_graph.constants import END, START
+from coco_graph.graph.message import MessageGraph, MessagesState, add_messages
+from coco_graph.graph.state import StateGraph
+
+__all__ = (
+    "END",
+    "START",
+    "StateGraph",
+    "add_messages",
+    "MessagesState",
+    "MessageGraph",
+)