langgraph-celery 0.1.0__tar.gz

langgraph_celery-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,22 @@
+name: CI
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - run: uv sync --frozen
+
+      - run: uv run ruff check src/
+
+      - run: uv run pytest

langgraph_celery-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,31 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - "v*"
+
+jobs:
+  release:
+    runs-on: ubuntu-latest
+    environment: release
+    permissions:
+      id-token: write
+      contents: read
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - run: uv sync --frozen
+
+      - run: uv run ruff check src/
+
+      - run: uv run pytest
+
+      - run: uv build
+
+      - uses: pypa/gh-action-pypi-publish@release/v1

langgraph_celery-0.1.0/.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+.venv/
+.env

langgraph_celery-0.1.0/.python-version

@@ -0,0 +1 @@
+3.13

langgraph_celery-0.1.0/PKG-INFO

@@ -0,0 +1,8 @@
+Metadata-Version: 2.4
+Name: langgraph-celery
+Version: 0.1.0
+Summary: Bridge between Celery workers and LangGraph agents
+Author-email: Charef <ccherrad@gmail.com>
+Requires-Python: >=3.11
+Requires-Dist: celery>=5.6.3
+Requires-Dist: langgraph>=1.1.10

langgraph_celery-0.1.0/pyproject.toml

@@ -0,0 +1,35 @@
+[project]
+name = "langgraph-celery"
+version = "0.1.0"
+description = "Bridge between Celery workers and LangGraph agents"
+readme = "README.md"
+authors = [
+    { name = "Charef", email = "ccherrad@gmail.com" }
+]
+requires-python = ">=3.11"
+dependencies = [
+    "celery>=5.6.3",
+    "langgraph>=1.1.10",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.0.3",
+    "pytest-asyncio>=1.3.0",
+    "ruff>=0.15.12",
+]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP"]

langgraph_celery-0.1.0/src/langgraph_celery/__init__.py

@@ -0,0 +1,5 @@
+from langgraph_celery.events import GraphResult, NodeEvent
+from langgraph_celery.interrupt import resume
+from langgraph_celery.task import task
+
+__all__ = ["task", "resume", "GraphResult", "NodeEvent"]

langgraph_celery-0.1.0/src/langgraph_celery/bridge.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+
+async def invoke(graph: Any, input: dict, config: dict | None = None) -> dict:
+    return await graph.ainvoke(input, config=config or {})
+
+
+async def stream_events(
+    graph: Any,
+    input: dict,
+    config: dict | None = None,
+) -> AsyncIterator[tuple[str, str, dict]]:
+    async for event in graph.astream_events(input, config=config or {}, version="v2"):
+        kind = event.get("event", "")
+        name = event.get("name", "")
+        data = event.get("data", {})
+        yield kind, name, data
+
+
+def run_sync(coro) -> Any:
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return asyncio.run(coro)
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        future = pool.submit(asyncio.run, coro)
+        return future.result()

langgraph_celery-0.1.0/src/langgraph_celery/checkpointing.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def build_checkpointer(checkpointer: str) -> Any:
+    if checkpointer == "memory":
+        from langgraph.checkpoint.memory import MemorySaver
+        return MemorySaver()
+    if checkpointer == "postgres":
+        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+        return AsyncPostgresSaver
+    raise ValueError(f"Unknown checkpointer: {checkpointer!r}. Use 'memory' or 'postgres'.")
+
+
+def resolve_thread_id(thread_id_from: str, kwargs: dict) -> str:
+    value = kwargs.get(thread_id_from)
+    if value is None:
+        raise KeyError(
+            f"thread_id_from={thread_id_from!r} not found in task kwargs: {list(kwargs)}"
+        )
+    return f"task:{value}"
+
+
+def inject_thread_id(config: dict, thread_id: str) -> dict:
+    config = dict(config)
+    configurable = dict(config.get("configurable", {}))
+    configurable["thread_id"] = thread_id
+    config["configurable"] = configurable
+    return config

langgraph_celery-0.1.0/src/langgraph_celery/emitter.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from langgraph_celery.bridge import stream_events
+from langgraph_celery.events import GraphResult, NodeEvent
+
+
+class NodeEventEmitter:
+    def __init__(self, callback: Callable[[NodeEvent], None]) -> None:
+        self._callback = callback
+
+    async def run(self, graph: Any, input: dict, config: dict | None = None) -> GraphResult:
+        result = await GraphResult.from_stream(stream_events(graph, input, config))
+        for event in result.node_events:
+            self._callback(event)
+        return result

langgraph_celery-0.1.0/src/langgraph_celery/events.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class NodeEvent:
+    kind: str
+    node: str
+    data: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class GraphResult:
+    output: dict[str, Any]
+    full_answer: str = ""
+    tool_calls: list[dict] = field(default_factory=list)
+    node_events: list[NodeEvent] = field(default_factory=list)
+    interrupted: bool = False
+
+    @classmethod
+    def from_output(cls, output: dict) -> GraphResult:
+        from langchain_core.messages import AIMessage
+
+        messages = output.get("messages", [])
+        full_answer = ""
+        tool_calls = []
+
+        for msg in reversed(messages):
+            if isinstance(msg, AIMessage):
+                if msg.content and not getattr(msg, "tool_calls", None):
+                    full_answer = msg.content
+                    break
+                if getattr(msg, "tool_calls", None):
+                    tool_calls.extend(msg.tool_calls)
+
+        return cls(output=output, full_answer=full_answer, tool_calls=tool_calls)
+
+    @classmethod
+    async def from_stream(
+        cls,
+        event_iter: AsyncIterator[tuple[str, str, dict]],
+    ) -> GraphResult:
+        tokens: list[str] = []
+        tool_calls: list[dict] = []
+        node_events: list[NodeEvent] = []
+        output: dict[str, Any] = {}
+
+        async for kind, name, data in event_iter:
+            if kind == "on_chat_model_stream":
+                chunk = data.get("chunk")
+                if chunk is not None:
+                    content = getattr(chunk, "content", None)
+                    if isinstance(content, str) and content:
+                        tokens.append(content)
+                    elif isinstance(content, list):
+                        for part in content:
+                            if isinstance(part, dict) and part.get("type") == "text":
+                                tokens.append(part.get("text", ""))
+
+            elif kind == "on_tool_start":
+                node_events.append(NodeEvent(kind=kind, node=name, data=data))
+
+            elif kind == "on_tool_end":
+                node_events.append(NodeEvent(kind=kind, node=name, data=data))
+                tool_input = data.get("input") or {}
+                tool_output = data.get("output")
+                tool_calls.append({"name": name, "input": tool_input, "output": tool_output})
+
+            elif kind == "on_chain_end" and name == "LangGraph":
+                output = data.get("output", {})
+
+        if not output and tokens:
+            output = {}
+
+        result = cls.from_output(output)
+        if tokens and not result.full_answer:
+            result.full_answer = "".join(tokens)
+        result.node_events = node_events
+        if not result.tool_calls:
+            result.tool_calls = tool_calls
+        return result

langgraph_celery-0.1.0/src/langgraph_celery/interrupt.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from typing import Any
+
+from langgraph_celery.bridge import invoke, run_sync
+from langgraph_celery.events import GraphResult
+
+
+async def _run_with_interrupt(graph: Any, graph_input: Any, config: dict) -> GraphResult:
+    output = await invoke(graph, graph_input, config)
+    interrupts = output.get("__interrupt__")
+    if interrupts:
+        interrupt_values = [i.value for i in interrupts]
+        result = GraphResult(
+            output=output,
+            interrupted=True,
+        )
+        result.output["__interrupt_values__"] = interrupt_values
+        return result
+    return GraphResult.from_output(output)
+
+
+async def _resume(graph: Any, resume_value: Any, config: dict) -> GraphResult:
+    from langgraph.types import Command
+
+    output = await invoke(graph, Command(resume=resume_value), config)
+    interrupts = output.get("__interrupt__")
+    if interrupts:
+        interrupt_values = [i.value for i in interrupts]
+        result = GraphResult(output=output, interrupted=True)
+        result.output["__interrupt_values__"] = interrupt_values
+        return result
+    return GraphResult.from_output(output)
+
+
+def run_with_interrupt(graph: Any, graph_input: Any, config: dict) -> GraphResult:
+    return run_sync(_run_with_interrupt(graph, graph_input, config))
+
+
+def resume(graph: Any, resume_value: Any, config: dict) -> GraphResult:
+    return run_sync(_resume(graph, resume_value, config))

langgraph_celery-0.1.0/src/langgraph_celery/streaming/redis.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from langgraph_celery.bridge import stream_events
+from langgraph_celery.events import GraphResult
+
+
+class RedisStreamer:
+    def __init__(self, redis_url: str, channel: str) -> None:
+        self._redis_url = redis_url
+        self._channel = channel
+
+    def _resolve_channel(self, kwargs: dict) -> str:
+        try:
+            return self._channel.format(**kwargs)
+        except KeyError:
+            return self._channel
+
+    async def run(
+        self,
+        graph: Any,
+        input: dict,
+        config: dict | None = None,
+        *,
+        task_kwargs: dict | None = None,
+    ) -> GraphResult:
+        import redis.asyncio as aioredis
+
+        channel = self._resolve_channel(task_kwargs or {})
+        client = aioredis.from_url(self._redis_url)
+
+        try:
+            result = await GraphResult.from_stream(
+                self._publish_stream(client, channel, graph, input, config)
+            )
+        finally:
+            await client.aclose()
+
+        return result
+
+    async def _publish_stream(
+        self, client: Any, channel: str, graph: Any, input: dict, config: dict | None
+    ):
+        async for kind, name, data in stream_events(graph, input, config):
+            if kind == "on_chat_model_stream":
+                chunk = data.get("chunk")
+                if chunk is not None:
+                    content = getattr(chunk, "content", None)
+                    if isinstance(content, str) and content:
+                        msg = json.dumps({"type": "token", "content": content})
+                        await client.publish(channel, msg)
+                    elif isinstance(content, list):
+                        for part in content:
+                            if isinstance(part, dict) and part.get("type") == "text":
+                                text = part.get("text", "")
+                                if text:
+                                    msg = json.dumps({"type": "token", "content": text})
+                                    await client.publish(channel, msg)
+            elif kind == "on_tool_start":
+                await client.publish(channel, json.dumps({"type": "tool_start", "name": name}))
+            elif kind == "on_tool_end":
+                await client.publish(channel, json.dumps({"type": "tool_end", "name": name}))
+            elif kind == "on_chain_end" and name == "LangGraph":
+                await client.publish(channel, json.dumps({"type": "done"}))
+            yield kind, name, data

langgraph_celery-0.1.0/src/langgraph_celery/task.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import functools
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+from langgraph_celery.bridge import invoke, run_sync
+
+
+def task(
+    graph_factory: Callable[..., Any] | None = None,
+    *,
+    streaming: str | None = None,
+    channel: str | None = None,
+    checkpointer: str | None = None,
+    thread_id_from: str | None = None,
+    emit_node_events: bool = False,
+    on_node_event: Callable | None = None,
+    interrupt_before: list[str] | None = None,
+):
+    def decorator(fn: Callable) -> Callable:
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            graph = _resolve_graph(graph_factory, fn, args, kwargs)
+            graph_input = _build_input(fn, args, kwargs)
+            config = _build_config(checkpointer, thread_id_from, kwargs, graph)
+
+            if streaming == "redis":
+                return run_sync(_run_redis(graph, graph_input, config, channel, kwargs))
+
+            if emit_node_events:
+                return run_sync(_run_emit(graph, graph_input, config, on_node_event))
+
+            if interrupt_before:
+                from langgraph_celery.interrupt import run_with_interrupt
+
+                return run_with_interrupt(graph, graph_input, config)
+
+            return run_sync(invoke(graph, graph_input, config))
+
+        wrapper._langgraph_celery = {
+            "graph_factory": graph_factory,
+            "streaming": streaming,
+            "channel": channel,
+            "checkpointer": checkpointer,
+            "thread_id_from": thread_id_from,
+            "emit_node_events": emit_node_events,
+            "interrupt_before": interrupt_before,
+        }
+        return wrapper
+
+    if callable(graph_factory):
+        fn = graph_factory
+        graph_factory = None
+        return decorator(fn)
+
+    return decorator
+
+
+def _resolve_graph(graph_factory, fn, args, kwargs):
+    if graph_factory is not None:
+        sig = inspect.signature(graph_factory)
+        params = list(sig.parameters.keys())
+        bound_kwargs = {k: v for k, v in kwargs.items() if k in params}
+        return graph_factory(**bound_kwargs)
+    raise ValueError("graph_factory must be provided via task(graph=...)")
+
+
+def _build_input(fn, args, kwargs) -> dict:
+    sig = inspect.signature(fn)
+    bound = sig.bind(*args, **kwargs)
+    bound.apply_defaults()
+    return dict(bound.arguments)
+
+
+def _build_config(checkpointer_name, thread_id_from, kwargs, graph) -> dict:
+    config: dict = {}
+
+    if checkpointer_name is not None:
+        from langgraph_celery.checkpointing import (
+            build_checkpointer,
+            inject_thread_id,
+            resolve_thread_id,
+        )
+
+        cp = build_checkpointer(checkpointer_name)
+        try:
+            graph.checkpointer = cp
+        except Exception:
+            pass
+
+        if thread_id_from:
+            thread_id = resolve_thread_id(thread_id_from, kwargs)
+            config = inject_thread_id(config, thread_id)
+
+    return config
+
+
+async def _run_redis(graph, graph_input, config, channel, task_kwargs) -> Any:
+    from langgraph_celery.streaming.redis import RedisStreamer
+
+    if channel is None:
+        raise ValueError("channel must be set when streaming='redis'")
+
+    redis_url = task_kwargs.get("redis_url", "redis://localhost:6379")
+    streamer = RedisStreamer(redis_url=redis_url, channel=channel)
+    return await streamer.run(graph, graph_input, config, task_kwargs=task_kwargs)
+
+
+async def _run_emit(graph, graph_input, config, on_node_event) -> Any:
+    from langgraph_celery.emitter import NodeEventEmitter
+
+    callback = on_node_event if on_node_event is not None else lambda e: None
+    emitter = NodeEventEmitter(callback=callback)
+    return await emitter.run(graph, graph_input, config)

langgraph_celery-0.1.0/tests/test_bridge.py

@@ -0,0 +1,22 @@
+import asyncio
+
+import pytest
+
+from langgraph_celery.bridge import run_sync
+
+
+async def _add(x, y):
+    return x + y
+
+
+def test_run_sync_no_loop():
+    result = run_sync(_add(1, 2))
+    assert result == 3
+
+
+def test_run_sync_inside_loop():
+    async def inner():
+        return run_sync(_add(3, 4))
+
+    result = asyncio.run(inner())
+    assert result == 7

langgraph_celery-0.1.0/tests/test_events.py

@@ -0,0 +1,22 @@
+from langchain_core.messages import AIMessage
+
+from langgraph_celery.events import GraphResult
+
+
+def test_graph_result_extracts_final_answer():
+    msg = AIMessage(content="Hello world")
+    result = GraphResult.from_output({"messages": [msg]})
+    assert result.full_answer == "Hello world"
+    assert not result.interrupted
+
+
+def test_graph_result_skips_tool_call_messages():
+    tool_msg = AIMessage(content="", tool_calls=[{"name": "search", "args": {}, "id": "1"}])
+    final_msg = AIMessage(content="Final answer")
+    result = GraphResult.from_output({"messages": [tool_msg, final_msg]})
+    assert result.full_answer == "Final answer"
+
+
+def test_graph_result_empty_messages():
+    result = GraphResult.from_output({"messages": []})
+    assert result.full_answer == ""

langgraph_celery-0.1.0/tests/test_graph_result_stream.py

@@ -0,0 +1,50 @@
+import pytest
+from unittest.mock import MagicMock
+
+from langgraph_celery.events import GraphResult, NodeEvent
+
+
+async def _make_iter(events):
+    for e in events:
+        yield e
+
+
+def _chunk(text):
+    m = MagicMock()
+    m.content = text
+    return m
+
+
+async def test_from_stream_tokens():
+    events = [
+        ("on_chat_model_stream", "model", {"chunk": _chunk("Hello ")}),
+        ("on_chat_model_stream", "model", {"chunk": _chunk("world")}),
+        ("on_chain_end", "LangGraph", {"output": {}}),
+    ]
+    result = await GraphResult.from_stream(_make_iter(events))
+    assert result.full_answer == "Hello world"
+
+
+async def test_from_stream_tool_events():
+    events = [
+        ("on_tool_start", "search", {"input": {"query": "foo"}}),
+        ("on_tool_end", "search", {"output": "bar", "input": {"query": "foo"}}),
+        ("on_chain_end", "LangGraph", {"output": {}}),
+    ]
+    result = await GraphResult.from_stream(_make_iter(events))
+    assert len(result.node_events) == 2
+    assert result.node_events[0].kind == "on_tool_start"
+    assert result.node_events[1].kind == "on_tool_end"
+    assert result.tool_calls[0]["name"] == "search"
+
+
+async def test_from_stream_prefers_chain_end_output():
+    from langchain_core.messages import AIMessage
+
+    msg = AIMessage(content="From output")
+    events = [
+        ("on_chat_model_stream", "model", {"chunk": _chunk("From stream")}),
+        ("on_chain_end", "LangGraph", {"output": {"messages": [msg]}}),
+    ]
+    result = await GraphResult.from_stream(_make_iter(events))
+    assert result.full_answer == "From output"