langgraph-celery 0.1.0__py3-none-any.whl

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/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/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/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/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/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/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/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.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,13 @@
+langgraph_celery/__init__.py,sha256=bL5p51zqAw7C7mqazOunH-H4nFru_l6JYiQVJZJ-7nM,202
+langgraph_celery/bridge.py,sha256=YfEYV5sIGHj-WCkGcVz7Bd4rZzqEazwkx3DPym6yv88,933
+langgraph_celery/checkpointing.py,sha256=5XLrrjsa9umViY6FFT6CtCrT2NBC11A13EgvMbu0Hu0,993
+langgraph_celery/emitter.py,sha256=jIzlZBvK9G9-egOFKm6895ibmMCfzyzPAM5HotwMMr8,612
+langgraph_celery/events.py,sha256=8tE1eDNa_x3us8hImypNXBPoNjxrMCwUVGPvF1IY2NE,2901
+langgraph_celery/interrupt.py,sha256=AkrvPl0P94_SgTXxQAz-hxO9okK41-_YJajg7-R3CP4,1447
+langgraph_celery/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langgraph_celery/task.py,sha256=XUf70IbQKwynGOFh7ExlOcHlqX8XSiX6JauLo1iXFYc,3783
+langgraph_celery/streaming/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langgraph_celery/streaming/redis.py,sha256=4PVn2bsnPiX-D4E73-wGn0fr1wlMI3zuJk9SomUeQbk,2547
+langgraph_celery-0.1.0.dist-info/METADATA,sha256=jrZcQQQUaJr8tp0bgBERT8b-4u_epqzqq91A_G68qwA,248
+langgraph_celery-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+langgraph_celery-0.1.0.dist-info/RECORD,,

langgraph_celery-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any