mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-29 18:46:59 +00:00
fix(relay): isolate streaming callback contexts
Signed-off-by: Alex Fournier <afournier@nvidia.com>
This commit is contained in:
Alex Fournier
parent
14bed44c8c
commit
eda54775e2
3 changed files with 195 additions and 9 deletions
|
|
@ -344,6 +344,11 @@ class ManagedLlmStream(Iterator[Any]):
|
|||
self.output_modified = False
|
||||
callback_context = contextvars.copy_context()
|
||||
|
||||
def run_callback(callback: Callable[..., Any], *args: Any) -> Any:
|
||||
# Relay can invoke stream surfaces while another callback still
|
||||
# owns the captured Context. A fresh copy is safe to enter.
|
||||
return callback_context.copy().run(callback, *args)
|
||||
|
||||
runtime, session, parent = relay_runtime.resolve_execution_context(session_id)
|
||||
if (
|
||||
runtime is None
|
||||
|
|
@ -378,7 +383,7 @@ class ManagedLlmStream(Iterator[Any]):
|
|||
async def provider_stream(next_request: Any):
|
||||
raw_stream = None
|
||||
try:
|
||||
raw_stream = callback_context.run(
|
||||
raw_stream = run_callback(
|
||||
stream_factory,
|
||||
_provider_request(
|
||||
request,
|
||||
|
|
@ -390,7 +395,7 @@ class ManagedLlmStream(Iterator[Any]):
|
|||
)
|
||||
if (
|
||||
completed_response_predicate is not None
|
||||
and callback_context.run(
|
||||
and run_callback(
|
||||
completed_response_predicate,
|
||||
raw_stream,
|
||||
)
|
||||
|
|
@ -399,14 +404,14 @@ class ManagedLlmStream(Iterator[Any]):
|
|||
self._provider_completed = True
|
||||
return
|
||||
if on_stream_created is not None:
|
||||
callback_context.run(on_stream_created, raw_stream)
|
||||
raw_iterator = callback_context.run(iter, raw_stream)
|
||||
run_callback(on_stream_created, raw_stream)
|
||||
raw_iterator = run_callback(iter, raw_stream)
|
||||
while True:
|
||||
try:
|
||||
chunk = callback_context.run(next, raw_iterator)
|
||||
chunk = run_callback(next, raw_iterator)
|
||||
except StopIteration:
|
||||
break
|
||||
if self._accept_chunk is not None and not callback_context.run(
|
||||
if self._accept_chunk is not None and not run_callback(
|
||||
self._accept_chunk,
|
||||
chunk,
|
||||
):
|
||||
|
|
@ -421,17 +426,17 @@ class ManagedLlmStream(Iterator[Any]):
|
|||
finally:
|
||||
close = getattr(raw_stream, "close", None)
|
||||
if callable(close):
|
||||
callback_context.run(close)
|
||||
run_callback(close)
|
||||
|
||||
def observe_chunk(chunk: Any) -> None:
|
||||
if self._on_chunk is not None:
|
||||
callback_context.run(self._on_chunk, _jsonable(chunk))
|
||||
run_callback(self._on_chunk, _jsonable(chunk))
|
||||
|
||||
def relay_finalizer() -> Any:
|
||||
try:
|
||||
if self.final_response is not None:
|
||||
return _jsonable(self.final_response)
|
||||
return _jsonable(callback_context.run(finalizer))
|
||||
return _jsonable(run_callback(finalizer))
|
||||
except BaseException as exc:
|
||||
self._callback_error = exc
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import contextvars
|
||||
import json
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
|
@ -436,6 +437,97 @@ def test_stream_provider_callbacks_preserve_caller_context(relay_turn):
|
|||
]
|
||||
|
||||
|
||||
def test_anthropic_stream_callbacks_do_not_reenter_captured_context(
|
||||
relay_turn,
|
||||
monkeypatch,
|
||||
):
|
||||
del relay_turn
|
||||
caller_value = contextvars.ContextVar(
|
||||
"anthropic_stream_caller_value",
|
||||
default="default",
|
||||
)
|
||||
caller_value.set("caller")
|
||||
callback_context = contextvars.copy_context()
|
||||
real_copy_context = contextvars.copy_context
|
||||
copy_count = 0
|
||||
|
||||
def capture_callback_context():
|
||||
nonlocal copy_count
|
||||
copy_count += 1
|
||||
if copy_count == 1:
|
||||
return callback_context
|
||||
return real_copy_context()
|
||||
|
||||
monkeypatch.setattr(
|
||||
relay_llm.contextvars,
|
||||
"copy_context",
|
||||
capture_callback_context,
|
||||
)
|
||||
observed = []
|
||||
accumulator = relay_llm.AnthropicStreamAccumulator()
|
||||
|
||||
def observe_chunk(chunk):
|
||||
observed.append(caller_value.get())
|
||||
accumulator.observe(chunk)
|
||||
|
||||
chunks = [
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "message-1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-test",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 1},
|
||||
},
|
||||
]
|
||||
stream = relay_llm.stream(
|
||||
{
|
||||
"model": "claude-test",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
},
|
||||
lambda _request: iter(chunks),
|
||||
session_id="session-1",
|
||||
name="anthropic",
|
||||
model_name="claude-test",
|
||||
finalizer=accumulator.finalize,
|
||||
on_chunk=observe_chunk,
|
||||
metadata={
|
||||
"api_mode": "anthropic_messages",
|
||||
"api_request_id": "request-anthropic-context-reentry",
|
||||
},
|
||||
)
|
||||
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def hold_callback_context() -> None:
|
||||
def wait() -> None:
|
||||
entered.set()
|
||||
assert release.wait(timeout=5)
|
||||
|
||||
callback_context.run(wait)
|
||||
|
||||
holder = threading.Thread(target=hold_callback_context)
|
||||
holder.start()
|
||||
assert entered.wait(timeout=1)
|
||||
try:
|
||||
assert list(stream) == chunks
|
||||
finally:
|
||||
release.set()
|
||||
holder.join(timeout=1)
|
||||
|
||||
assert holder.is_alive() is False
|
||||
assert observed == ["caller", "caller"]
|
||||
|
||||
|
||||
def test_non_stream_does_not_forward_relay_session_headers(relay_turn):
|
||||
_relay, _turn = relay_turn
|
||||
captured_requests = []
|
||||
|
|
|
|||
|
|
@ -1313,6 +1313,95 @@ class TestAnthropicStreamCallbacks:
|
|||
assert touch_calls.count("receiving stream response") == len(events)
|
||||
mock_stream.close.assert_called_once()
|
||||
|
||||
@pytest.mark.filterwarnings("ignore:Pydantic serializer warnings:UserWarning")
|
||||
def test_anthropic_sdk_stream_runs_through_relay_managed_execution(
|
||||
self,
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
anthropic = pytest.importorskip("anthropic")
|
||||
httpx = pytest.importorskip("httpx")
|
||||
pytest.importorskip("nemo_relay")
|
||||
from agent import relay_runtime
|
||||
from run_agent import AIAgent
|
||||
|
||||
monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes-home"))
|
||||
response_body = b"""event: message_start
|
||||
data: {"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","content":[],"model":"claude-test","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":0}}}
|
||||
|
||||
event: content_block_start
|
||||
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}
|
||||
|
||||
event: content_block_stop
|
||||
data: {"type":"content_block_stop","index":0}
|
||||
|
||||
event: message_delta
|
||||
data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":1}}
|
||||
|
||||
event: message_stop
|
||||
data: {"type":"message_stop"}
|
||||
|
||||
"""
|
||||
|
||||
def respond(request):
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
content=response_body,
|
||||
request=request,
|
||||
)
|
||||
|
||||
client = anthropic.Anthropic(
|
||||
api_key="test-key",
|
||||
http_client=httpx.Client(transport=httpx.MockTransport(respond)),
|
||||
)
|
||||
relay_runtime._reset_for_tests()
|
||||
agent = AIAgent(
|
||||
api_key="test-key",
|
||||
base_url="https://api.anthropic.com",
|
||||
provider="anthropic",
|
||||
model="claude-test",
|
||||
quiet_mode=True,
|
||||
skip_context_files=True,
|
||||
skip_memory=True,
|
||||
)
|
||||
agent.api_mode = "anthropic_messages"
|
||||
agent.session_id = "anthropic-relay-session"
|
||||
agent._interrupt_requested = False
|
||||
agent._create_request_anthropic_client = lambda *args, **kwargs: client
|
||||
lease = relay_runtime.SESSION_COORDINATOR.acquire_conversation(
|
||||
profile_key=relay_runtime.current_profile_key(),
|
||||
session_id=agent.session_id,
|
||||
platform="cli",
|
||||
)
|
||||
turn = relay_runtime.SESSION_COORDINATOR.begin_turn(
|
||||
lease,
|
||||
turn_id="anthropic-relay-turn",
|
||||
task_id="anthropic-relay-task",
|
||||
)
|
||||
lease.host.retain_managed_execution("test.anthropic_relay")
|
||||
|
||||
try:
|
||||
result = agent._interruptible_streaming_api_call(
|
||||
{
|
||||
"model": "claude-test",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
)
|
||||
finally:
|
||||
lease.host.release_managed_execution("test.anthropic_relay")
|
||||
relay_runtime.SESSION_COORDINATOR.end_turn(turn, outcome="success")
|
||||
relay_runtime.SESSION_COORDINATOR.release_conversation(lease)
|
||||
relay_runtime._reset_for_tests()
|
||||
client.close()
|
||||
|
||||
assert result.content[0].text == "hello"
|
||||
assert result.stop_reason == "end_turn"
|
||||
|
||||
@patch("run_agent.AIAgent._rebuild_anthropic_client")
|
||||
@patch("run_agent.AIAgent._replace_primary_openai_client")
|
||||
def test_anthropic_stream_parser_valueerror_retries_before_delivery(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue