mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-10 13:27:36 +00:00
format
This commit is contained in:
parent
f365865fb5
commit
cd14cc5014
@ -238,10 +238,10 @@ def shielded(func: Func) -> Func:
|
||||
async def wrapped(*args: Any, **kwargs: Any) -> Any:
|
||||
# Capture the current context to preserve context variables
|
||||
ctx = copy_context()
|
||||
|
||||
|
||||
# Create the coroutine
|
||||
coro = func(*args, **kwargs)
|
||||
|
||||
|
||||
# For Python 3.11+, create task with explicit context
|
||||
# For older versions, fallback to original behavior
|
||||
try:
|
||||
|
@ -153,56 +153,56 @@ async def test_inline_handlers_share_parent_context_multiple() -> None:
|
||||
|
||||
async def test_shielded_callback_context_preservation() -> None:
|
||||
"""Verify that shielded callbacks preserve context variables.
|
||||
|
||||
|
||||
This test specifically addresses the issue where async callbacks decorated
|
||||
with @shielded do not properly preserve context variables, breaking
|
||||
instrumentation and other context-dependent functionality.
|
||||
|
||||
|
||||
The issue manifests in callbacks that use the @shielded decorator:
|
||||
* on_llm_end
|
||||
* on_llm_error
|
||||
* on_llm_error
|
||||
* on_chain_end
|
||||
* on_chain_error
|
||||
* And other shielded callback methods
|
||||
"""
|
||||
context_var: contextvars.ContextVar[str] = contextvars.ContextVar("test_context")
|
||||
|
||||
|
||||
class ContextTestHandler(AsyncCallbackHandler):
|
||||
"""Handler that reads context variables in shielded callbacks."""
|
||||
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.run_inline = False
|
||||
self.context_values = []
|
||||
|
||||
@override
|
||||
|
||||
@override
|
||||
async def on_llm_end(self, response: Any, **kwargs: Any) -> None:
|
||||
"""This method is decorated with @shielded in the run manager."""
|
||||
# This should preserve the context variable value
|
||||
self.context_values.append(context_var.get("not_found"))
|
||||
|
||||
|
||||
@override
|
||||
async def on_chain_end(self, outputs: Any, **kwargs: Any) -> None:
|
||||
"""This method is decorated with @shielded in the run manager."""
|
||||
# This should preserve the context variable value
|
||||
self.context_values.append(context_var.get("not_found"))
|
||||
|
||||
|
||||
# Set up the test context
|
||||
context_var.set("test_value")
|
||||
handler = ContextTestHandler()
|
||||
manager = AsyncCallbackManager(handlers=[handler])
|
||||
|
||||
|
||||
# Create run managers that have the shielded methods
|
||||
llm_managers = await manager.on_llm_start({}, ["test prompt"])
|
||||
llm_run_manager = llm_managers[0]
|
||||
|
||||
|
||||
chain_run_manager = await manager.on_chain_start({}, {"test": "input"})
|
||||
|
||||
|
||||
# Test LLM end callback (which is shielded)
|
||||
await llm_run_manager.on_llm_end({"response": "test"})
|
||||
|
||||
# Test Chain end callback (which is shielded)
|
||||
|
||||
# Test Chain end callback (which is shielded)
|
||||
await chain_run_manager.on_chain_end({"output": "test"})
|
||||
|
||||
|
||||
# The context should be preserved in shielded callbacks
|
||||
# This was the main issue - shielded decorators were not preserving context
|
||||
assert handler.context_values == ["test_value", "test_value"], (
|
||||
|
Loading…
Reference in New Issue
Block a user