core[minor]: Add dispatching for custom events (#24080)

This PR allows dispatching adhoc events for a given run.

# Context

This PR allows users to send arbitrary data to the callback system and
to the astream events API from within a given runnable. This can be
extremely useful to surface custom information to end users about
progress etc.

Integration with langsmith tracer will be done separately since the data
cannot be currently visualized. It'll be accommodated using the events
attribute of the Run

# Examples with astream events

```python
from langchain_core.callbacks import adispatch_custom_event
from langchain_core.tools import tool

@tool
async def foo(x: int) -> int:
    """Foo"""
    await adispatch_custom_event("event1", {"x": x})
    await adispatch_custom_event("event2", {"x": x})
    return x + 1

async for event in foo.astream_events({'x': 1}, version='v2'):
    print(event)
```

```python
{'event': 'on_tool_start', 'data': {'input': {'x': 1}}, 'name': 'foo', 'tags': [], 'run_id': 'fd6fb7a7-dd37-4191-962c-e43e245909f6', 'metadata': {}, 'parent_ids': []}
{'event': 'on_custom_event', 'run_id': 'fd6fb7a7-dd37-4191-962c-e43e245909f6', 'name': 'event1', 'tags': [], 'metadata': {}, 'data': {'x': 1}, 'parent_ids': []}
{'event': 'on_custom_event', 'run_id': 'fd6fb7a7-dd37-4191-962c-e43e245909f6', 'name': 'event2', 'tags': [], 'metadata': {}, 'data': {'x': 1}, 'parent_ids': []}
{'event': 'on_tool_end', 'data': {'output': 2}, 'run_id': 'fd6fb7a7-dd37-4191-962c-e43e245909f6', 'name': 'foo', 'tags': [], 'metadata': {}, 'parent_ids': []}
```

```python
from langchain_core.callbacks import adispatch_custom_event
from langchain_core.runnables import RunnableLambda

@RunnableLambda
async def foo(x: int) -> int:
    """Foo"""
    await adispatch_custom_event("event1", {"x": x})
    await adispatch_custom_event("event2", {"x": x})
    return x + 1

async for event in foo.astream_events(1, version='v2'):
    print(event)
```

```python
{'event': 'on_chain_start', 'data': {'input': 1}, 'name': 'foo', 'tags': [], 'run_id': 'ce2beef2-8608-49ea-8eba-537bdaafb8ec', 'metadata': {}, 'parent_ids': []}
{'event': 'on_custom_event', 'run_id': 'ce2beef2-8608-49ea-8eba-537bdaafb8ec', 'name': 'event1', 'tags': [], 'metadata': {}, 'data': {'x': 1}, 'parent_ids': []}
{'event': 'on_custom_event', 'run_id': 'ce2beef2-8608-49ea-8eba-537bdaafb8ec', 'name': 'event2', 'tags': [], 'metadata': {}, 'data': {'x': 1}, 'parent_ids': []}
{'event': 'on_chain_stream', 'run_id': 'ce2beef2-8608-49ea-8eba-537bdaafb8ec', 'name': 'foo', 'tags': [], 'metadata': {}, 'data': {'chunk': 2}, 'parent_ids': []}
{'event': 'on_chain_end', 'data': {'output': 2}, 'run_id': 'ce2beef2-8608-49ea-8eba-537bdaafb8ec', 'name': 'foo', 'tags': [], 'metadata': {}, 'parent_ids': []}
```

# Examples with handlers 

This is copy pasted from unit tests

```python
    class CustomCallbackManager(BaseCallbackHandler):
        def __init__(self) -> None:
            self.events: List[Any] = []

        def on_custom_event(
            self,
            name: str,
            data: Any,
            *,
            run_id: UUID,
            tags: Optional[List[str]] = None,
            metadata: Optional[Dict[str, Any]] = None,
            **kwargs: Any,
        ) -> None:
            assert kwargs == {}
            self.events.append(
                (
                    name,
                    data,
                    run_id,
                    tags,
                    metadata,
                )
            )

    callback = CustomCallbackManager()

    run_id = uuid.UUID(int=7)

    @RunnableLambda
    def foo(x: int, config: RunnableConfig) -> int:
        dispatch_custom_event("event1", {"x": x})
        dispatch_custom_event("event2", {"x": x}, config=config)
        return x

    foo.invoke(1, {"callbacks": [callback], "run_id": run_id})

    assert callback.events == [
        ("event1", {"x": 1}, UUID("00000000-0000-0000-0000-000000000007"), [], {}),
        ("event2", {"x": 1}, UUID("00000000-0000-0000-0000-000000000007"), [], {}),
    ]
```
This commit is contained in:
Eugene Yurtsev
2024-07-10 22:25:12 -04:00
committed by GitHub
parent 14a8bbc21a
commit dc131ac42a
8 changed files with 800 additions and 18 deletions

View File

@@ -0,0 +1,161 @@
import sys
import uuid
from typing import Any, Dict, List, Optional
from uuid import UUID
import pytest
from langchain_core.callbacks import AsyncCallbackHandler, BaseCallbackHandler
from langchain_core.callbacks.manager import (
adispatch_custom_event,
dispatch_custom_event,
)
from langchain_core.runnables import RunnableLambda
from langchain_core.runnables.config import RunnableConfig
class AsyncCustomCallbackHandler(AsyncCallbackHandler):
def __init__(self) -> None:
self.events: List[Any] = []
async def on_custom_event(
self,
name: str,
data: Any,
*,
run_id: UUID,
tags: Optional[List[str]] = None,
metadata: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> None:
assert kwargs == {}
self.events.append(
(
name,
data,
run_id,
tags,
metadata,
)
)
def test_custom_event_root_dispatch() -> None:
"""Test adhoc event in a nested chain."""
# This just tests that nothing breaks on the path.
# It shouldn't do anything at the moment, since the tracer isn't configured
# to handle adhoc events.
# Expected behavior is that the event cannot be dispatched
with pytest.raises(RuntimeError):
dispatch_custom_event("event1", {"x": 1})
async def test_async_custom_event_root_dispatch() -> None:
"""Test adhoc event in a nested chain."""
# This just tests that nothing breaks on the path.
# It shouldn't do anything at the moment, since the tracer isn't configured
# to handle adhoc events.
# Expected behavior is that the event cannot be dispatched
with pytest.raises(RuntimeError):
await adispatch_custom_event("event1", {"x": 1})
IS_GTE_3_11 = sys.version_info >= (3, 11)
@pytest.mark.skipif(not IS_GTE_3_11, reason="Requires Python >=3.11")
async def test_async_custom_event_implicit_config() -> None:
"""Test dispatch without passing config explicitly."""
callback = AsyncCustomCallbackHandler()
run_id = uuid.UUID(int=7)
# Typing not working well with RunnableLambda when used as
# a decorator for async functions
@RunnableLambda # type: ignore[arg-type]
async def foo(x: int, config: RunnableConfig) -> int:
await adispatch_custom_event("event1", {"x": x})
await adispatch_custom_event("event2", {"x": x})
return x
await foo.ainvoke(
1, # type: ignore[arg-type]
{"callbacks": [callback], "run_id": run_id},
)
assert callback.events == [
("event1", {"x": 1}, UUID("00000000-0000-0000-0000-000000000007"), [], {}),
("event2", {"x": 1}, UUID("00000000-0000-0000-0000-000000000007"), [], {}),
]
async def test_async_callback_manager() -> None:
"""Test async callback manager."""
callback = AsyncCustomCallbackHandler()
run_id = uuid.UUID(int=7)
# Typing not working well with RunnableLambda when used as
# a decorator for async functions
@RunnableLambda # type: ignore[arg-type]
async def foo(x: int, config: RunnableConfig) -> int:
await adispatch_custom_event("event1", {"x": x}, config=config)
await adispatch_custom_event("event2", {"x": x}, config=config)
return x
await foo.ainvoke(
1, # type: ignore[arg-type]
{"callbacks": [callback], "run_id": run_id},
)
assert callback.events == [
("event1", {"x": 1}, UUID("00000000-0000-0000-0000-000000000007"), [], {}),
("event2", {"x": 1}, UUID("00000000-0000-0000-0000-000000000007"), [], {}),
]
def test_sync_callback_manager() -> None:
"""Test async callback manager."""
class CustomCallbackManager(BaseCallbackHandler):
def __init__(self) -> None:
self.events: List[Any] = []
def on_custom_event(
self,
name: str,
data: Any,
*,
run_id: UUID,
tags: Optional[List[str]] = None,
metadata: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> None:
assert kwargs == {}
self.events.append(
(
name,
data,
run_id,
tags,
metadata,
)
)
callback = CustomCallbackManager()
run_id = uuid.UUID(int=7)
@RunnableLambda
def foo(x: int, config: RunnableConfig) -> int:
dispatch_custom_event("event1", {"x": x})
dispatch_custom_event("event2", {"x": x}, config=config)
return x
foo.invoke(1, {"callbacks": [callback], "run_id": run_id})
assert callback.events == [
("event1", {"x": 1}, UUID("00000000-0000-0000-0000-000000000007"), [], {}),
("event2", {"x": 1}, UUID("00000000-0000-0000-0000-000000000007"), [], {}),
]

View File

@@ -31,6 +31,8 @@ EXPECTED_ALL = [
"StdOutCallbackHandler",
"StreamingStdOutCallbackHandler",
"FileCallbackHandler",
"adispatch_custom_event",
"dispatch_custom_event",
]

View File

@@ -2353,3 +2353,245 @@ async def test_cancel_astream_events() -> None:
# node "anotherwhile" should never start
assert anotherwhile.started is False
async def test_custom_event() -> None:
"""Test adhoc event."""
from langchain_core.callbacks.manager import adispatch_custom_event
# Ignoring type due to RunnableLamdba being dynamic when it comes to being
# applied as a decorator to async functions.
@RunnableLambda # type: ignore[arg-type]
async def foo(x: int, config: RunnableConfig) -> int:
"""Simple function that emits some adhoc events."""
await adispatch_custom_event("event1", {"x": x}, config=config)
await adispatch_custom_event("event2", "foo", config=config)
return x + 1
uuid1 = uuid.UUID(int=7)
events = await _collect_events(
foo.astream_events(
1,
version="v2",
config={"run_id": uuid1},
),
with_nulled_ids=False,
)
run_id = str(uuid1)
assert events == [
{
"data": {"input": 1},
"event": "on_chain_start",
"metadata": {},
"name": "foo",
"parent_ids": [],
"run_id": run_id,
"tags": [],
},
{
"data": {"x": 1},
"event": "on_custom_event",
"metadata": {},
"name": "event1",
"parent_ids": [],
"run_id": run_id,
"tags": [],
},
{
"data": "foo",
"event": "on_custom_event",
"metadata": {},
"name": "event2",
"parent_ids": [],
"run_id": run_id,
"tags": [],
},
{
"data": {"chunk": 2},
"event": "on_chain_stream",
"metadata": {},
"name": "foo",
"parent_ids": [],
"run_id": run_id,
"tags": [],
},
{
"data": {"output": 2},
"event": "on_chain_end",
"metadata": {},
"name": "foo",
"parent_ids": [],
"run_id": run_id,
"tags": [],
},
]
async def test_custom_event_nested() -> None:
"""Test adhoc event in a nested chain."""
from langchain_core.callbacks.manager import adispatch_custom_event
# Ignoring type due to RunnableLamdba being dynamic when it comes to being
# applied as a decorator to async functions.
@RunnableLambda # type: ignore[arg-type]
async def foo(x: int, config: RunnableConfig) -> int:
"""Simple function that emits some adhoc events."""
await adispatch_custom_event("event1", {"x": x}, config=config)
await adispatch_custom_event("event2", "foo", config=config)
return x + 1
run_id = uuid.UUID(int=7)
child_run_id = uuid.UUID(int=8)
# Ignoring type due to RunnableLamdba being dynamic when it comes to being
# applied as a decorator to async functions.
@RunnableLambda # type: ignore[arg-type]
async def bar(x: int, config: RunnableConfig) -> int:
"""Simple function that emits some adhoc events."""
return await foo.ainvoke(
x, # type: ignore[arg-type]
{"run_id": child_run_id, **config},
)
events = await _collect_events(
bar.astream_events(
1,
version="v2",
config={"run_id": run_id},
),
with_nulled_ids=False,
)
run_id = str(run_id) # type: ignore[assignment]
child_run_id = str(child_run_id) # type: ignore[assignment]
assert events == [
{
"data": {"input": 1},
"event": "on_chain_start",
"metadata": {},
"name": "bar",
"parent_ids": [],
"run_id": "00000000-0000-0000-0000-000000000007",
"tags": [],
},
{
"data": {"input": 1},
"event": "on_chain_start",
"metadata": {},
"name": "foo",
"parent_ids": ["00000000-0000-0000-0000-000000000007"],
"run_id": "00000000-0000-0000-0000-000000000008",
"tags": [],
},
{
"data": {"x": 1},
"event": "on_custom_event",
"metadata": {},
"name": "event1",
"parent_ids": ["00000000-0000-0000-0000-000000000007"],
"run_id": "00000000-0000-0000-0000-000000000008",
"tags": [],
},
{
"data": "foo",
"event": "on_custom_event",
"metadata": {},
"name": "event2",
"parent_ids": ["00000000-0000-0000-0000-000000000007"],
"run_id": "00000000-0000-0000-0000-000000000008",
"tags": [],
},
{
"data": {"input": 1, "output": 2},
"event": "on_chain_end",
"metadata": {},
"name": "foo",
"parent_ids": ["00000000-0000-0000-0000-000000000007"],
"run_id": "00000000-0000-0000-0000-000000000008",
"tags": [],
},
{
"data": {"chunk": 2},
"event": "on_chain_stream",
"metadata": {},
"name": "bar",
"parent_ids": [],
"run_id": "00000000-0000-0000-0000-000000000007",
"tags": [],
},
{
"data": {"output": 2},
"event": "on_chain_end",
"metadata": {},
"name": "bar",
"parent_ids": [],
"run_id": "00000000-0000-0000-0000-000000000007",
"tags": [],
},
]
async def test_custom_event_root_dispatch() -> None:
"""Test adhoc event in a nested chain."""
# This just tests that nothing breaks on the path.
# It shouldn't do anything at the moment, since the tracer isn't configured
# to handle adhoc events.
from langchain_core.callbacks.manager import adispatch_custom_event
# Expected behavior is that the event cannot be dispatched
with pytest.raises(RuntimeError):
await adispatch_custom_event("event1", {"x": 1})
IS_GTE_3_11 = sys.version_info >= (3, 11)
# Test relies on automatically picking up RunnableConfig from contextvars
@pytest.mark.skipif(not IS_GTE_3_11, reason="Requires Python >=3.11")
async def test_custom_event_root_dispatch_with_in_tool() -> None:
"""Test adhoc event in a nested chain."""
from langchain_core.callbacks.manager import adispatch_custom_event
from langchain_core.tools import tool
@tool
async def foo(x: int) -> int:
"""Foo"""
await adispatch_custom_event("event1", {"x": x})
return x + 1
# Ignoring type due to @tool not returning correct type annotations
events = await _collect_events(
foo.astream_events({"x": 2}, version="v2") # type: ignore[attr-defined]
)
assert events == [
{
"data": {"input": {"x": 2}},
"event": "on_tool_start",
"metadata": {},
"name": "foo",
"parent_ids": [],
"run_id": "",
"tags": [],
},
{
"data": {"x": 2},
"event": "on_custom_event",
"metadata": {},
"name": "event1",
"parent_ids": [],
"run_id": "",
"tags": [],
},
{
"data": {"output": 3},
"event": "on_tool_end",
"metadata": {},
"name": "foo",
"parent_ids": [],
"run_id": "",
"tags": [],
},
]