mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-02 11:39:18 +00:00
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:
@@ -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"), [], {}),
|
||||
]
|
@@ -31,6 +31,8 @@ EXPECTED_ALL = [
|
||||
"StdOutCallbackHandler",
|
||||
"StreamingStdOutCallbackHandler",
|
||||
"FileCallbackHandler",
|
||||
"adispatch_custom_event",
|
||||
"dispatch_custom_event",
|
||||
]
|
||||
|
||||
|
||||
|
@@ -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": [],
|
||||
},
|
||||
]
|
||||
|
Reference in New Issue
Block a user