diff --git a/AGENTS.md b/AGENTS.md index 4b3c547b508e..2a376464ee3e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -462,6 +462,17 @@ pnpm exec tsx examples/claudeAgents.ts - Keep streaming and batch paths aligned: `Turn._build_attrs()` must apply content gating and PII redaction to both lists before passing them to `invoke_agent_attributes()`, and `log_turn()` must accept both fields. +- The Python Claude Agent SDK integration taps async prompt iterables, queues + user inputs in FIFO order, and closes one `Turn` for each `ResultMessage`. +- Keep its output adapters split by SDK contract: string `query()` and each + `ClaudeSDKClient.receive_response()` use the linear single-turn tracer, while + only standalone `query(AsyncIterable)` uses multi-result queue/lookahead logic. +- Name the corresponding test cases "string prompt" and "async-iterable + prompt" rather than sync/async: both SDK entry points are asynchronous. Keep + one-input success and pre/post-result failure assertions paired across them. +- SDK output can fail before the first message or between completed turns. Keep + the submitted/pending input observable by ending its `Turn` through the + exception path so the span records the error before the exception propagates. ### Claude Agent SDK token accounting diff --git a/tests/integrations/claude_agent_sdk/claude_agent_sdk_otel_test.py b/tests/integrations/claude_agent_sdk/claude_agent_sdk_otel_test.py index 723cd3e2f43a..089911756a37 100644 --- a/tests/integrations/claude_agent_sdk/claude_agent_sdk_otel_test.py +++ b/tests/integrations/claude_agent_sdk/claude_agent_sdk_otel_test.py @@ -10,7 +10,7 @@ import asyncio import json -from collections.abc import AsyncIterator, Generator +from collections.abc import AsyncIterable, AsyncIterator, Generator from typing import Any import pytest @@ -21,6 +21,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import SpanKind, StatusCode +import weave.integrations.claude_agent_sdk.otel_integration as claude_agent_sdk_otel_integration from tests.integrations.claude_agent_sdk.conftest import ReplayTransport, load_cassette from weave.conversation import agent_name_override from weave.integrations.claude_agent_sdk.otel_integration import ( @@ -47,9 +48,7 @@ def otel_spans(monkeypatch: pytest.MonkeyPatch) -> Generator[InMemorySpanExporte @pytest.fixture(autouse=True) def patch_claude_agent_sdk_otel() -> Generator[None]: - import weave.integrations.claude_agent_sdk.otel_integration as mod - - mod._claude_agent_sdk_otel_patcher = None + claude_agent_sdk_otel_integration._claude_agent_sdk_otel_patcher = None patcher = get_claude_agent_sdk_otel_patcher() patcher.attempt_patch() yield @@ -111,7 +110,36 @@ def get_part_types(messages: list[dict[str, Any]]) -> set[str]: } -async def run_query(cassette: str, prompt: str) -> None: +def user_prompt(text: str, *, should_query: bool | None = None) -> dict[str, Any]: + prompt: dict[str, Any] = { + "type": "user", + "message": {"role": "user", "content": text}, + "parent_tool_use_id": None, + } + if should_query is not None: + prompt["shouldQuery"] = should_query + return prompt + + +async def async_prompt_messages( + *messages: dict[str, Any], +) -> AsyncIterator[dict[str, Any]]: + for message in messages: + yield message + + +def text_query_prompt( + text: str, *, as_async_iterable: bool +) -> str | AsyncIterable[dict[str, Any]]: + if as_async_iterable: + return async_prompt_messages(user_prompt(text)) + return text + + +async def run_query( + cassette: str, + prompt: str | AsyncIterable[dict[str, Any]], +) -> None: async for _ in query( prompt=prompt, options=ClaudeAgentOptions(), @@ -120,12 +148,23 @@ async def run_query(cassette: str, prompt: str) -> None: pass -# --- query(): simple text --------------------------------------------------- +# --- query(): string and single-input async-iterable parity ----------------- @pytest.mark.asyncio -async def test_simple_text_query_otel(otel_spans: InMemorySpanExporter) -> None: - await run_query("simple_text_response", "What is 2+2?") +@pytest.mark.parametrize( + "as_async_iterable", + [False, True], + ids=["string-prompt", "async-iterable-prompt"], +) +async def test_single_text_query_otel( + otel_spans: InMemorySpanExporter, + as_async_iterable: bool, +) -> None: + await run_query( + "simple_text_response", + text_query_prompt("What is 2+2?", as_async_iterable=as_async_iterable), + ) spans = otel_spans.get_finished_spans() assert {span.instrumentation_scope.name for span in spans} == {"weave.conversation"} @@ -319,11 +358,13 @@ async def test_ambient_trace_nesting_otel(otel_spans: InMemorySpanExporter) -> N assert agent_span.context.trace_id == outer_context.trace_id -# --- ClaudeSDKClient: multi-turn -------------------------------------------- +# --- ClaudeSDKClient: one turn per receive_response() ------------------------ @pytest.mark.asyncio -async def test_multi_turn_client_otel(otel_spans: InMemorySpanExporter) -> None: +async def test_string_client_queries_use_one_turn_per_response_otel( + otel_spans: InMemorySpanExporter, +) -> None: sdk_client = ClaudeSDKClient( options=ClaudeAgentOptions(), transport=ReplayTransport(load_cassette("multi_turn_response")), @@ -357,11 +398,51 @@ async def test_multi_turn_client_otel(otel_spans: InMemorySpanExporter) -> None: assert prompts == {"Hello", "What is the capital of France?"} -# --- query(): streamed image prompt ----------------------------------------- +@pytest.mark.asyncio +async def test_async_iterable_client_query_uses_one_turn_per_response_otel( + otel_spans: InMemorySpanExporter, +) -> None: + """Each receive_response() remains a single turn for async client input.""" + sdk_client = ClaudeSDKClient( + options=ClaudeAgentOptions(), + transport=ReplayTransport(load_cassette("multi_turn_response")), + ) + await sdk_client.connect() + await sdk_client.query( + async_prompt_messages( + user_prompt("Hello"), + user_prompt("What is the capital of France?"), + ) + ) + + responses = [ + [type(message).__name__ async for message in sdk_client.receive_response()], + [type(message).__name__ async for message in sdk_client.receive_response()], + ] + await sdk_client.disconnect() + + assert responses == [ + ["SystemMessage", "AssistantMessage", "ResultMessage"], + ["AssistantMessage", "ResultMessage"], + ] + agent_spans = sorted( + get_spans_by_op(otel_spans.get_finished_spans(), "invoke_agent"), + key=lambda span: span.start_time, + ) + assert [ + get_all_text(get_messages(agent_span, "gen_ai.input.messages")) + for agent_span in agent_spans + ] == [ + "Hello", + "What is the capital of France?", + ] + + +# --- query(): async-iterable image prompt ----------------------------------- @pytest.mark.asyncio -async def test_streamed_image_prompt_otel( +async def test_async_iterable_image_prompt_otel( otel_spans: InMemorySpanExporter, ) -> None: image_base64 = ( @@ -424,6 +505,252 @@ async def image_prompt() -> AsyncIterator[dict[str, Any]]: ] +# --- query(): async-iterable multi-turn prompts ------------------------------ + + +@pytest.mark.asyncio +async def test_async_iterable_query_multi_turn_otel( + otel_spans: InMemorySpanExporter, +) -> None: + messages = [ + message + async for message in query( + prompt=async_prompt_messages( + user_prompt("Hello"), + user_prompt("What is the capital of France?"), + ), + options=ClaudeAgentOptions(), + transport=ReplayTransport(load_cassette("multi_turn_response")), + ) + ] + + assert [type(message).__name__ for message in messages] == [ + "SystemMessage", + "AssistantMessage", + "ResultMessage", + "AssistantMessage", + "ResultMessage", + ] + agent_spans = sorted( + get_spans_by_op(otel_spans.get_finished_spans(), "invoke_agent"), + key=lambda span: span.start_time, + ) + assert len(agent_spans) == 2 + assert [ + get_messages(agent_span, "gen_ai.input.messages") for agent_span in agent_spans + ] == [ + [{"role": "user", "parts": [{"type": "text", "content": "Hello"}]}], + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "What is the capital of France?", + } + ], + } + ], + ] + assert [ + get_messages(agent_span, "gen_ai.output.messages") for agent_span in agent_spans + ] == [ + [ + { + "role": "assistant", + "parts": [{"type": "text", "content": "Hello! How can I help you?"}], + } + ], + [ + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "The capital of France is Paris.", + } + ], + } + ], + ] + assert { + get_attrs(agent_span)["gen_ai.conversation.id"] for agent_span in agent_spans + } == {"s-mt001"} + assert len({agent_span.context.trace_id for agent_span in agent_spans}) == 2 + + +@pytest.mark.asyncio +async def test_async_iterable_query_buffers_non_query_inputs_otel( + otel_spans: InMemorySpanExporter, +) -> None: + messages = [ + message + async for message in query( + prompt=async_prompt_messages( + user_prompt("Background context", should_query=False), + user_prompt("Answer using that context"), + ), + options=ClaudeAgentOptions(), + transport=ReplayTransport(load_cassette("simple_text_response")), + ) + ] + + assert [type(message).__name__ for message in messages] == [ + "SystemMessage", + "AssistantMessage", + "ResultMessage", + ] + agent_spans = get_spans_by_op(otel_spans.get_finished_spans(), "invoke_agent") + assert len(agent_spans) == 1 + assert get_messages(agent_spans[0], "gen_ai.input.messages") == [ + { + "role": "user", + "parts": [{"type": "text", "content": "Background context"}], + }, + { + "role": "user", + "parts": [{"type": "text", "content": "Answer using that context"}], + }, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "as_async_iterable", + [False, True], + ids=["string-prompt", "async-iterable-prompt"], +) +async def test_query_failure_before_first_response_records_error_turn_otel( + otel_spans: InMemorySpanExporter, + as_async_iterable: bool, +) -> None: + transport = ReplayTransport( + load_cassette("simple_text_response"), + fail_after_messages=0, + ) + + with pytest.raises(Exception, match="replay transport failed"): + _ = [ + message + async for message in query( + prompt=text_query_prompt("Hello", as_async_iterable=as_async_iterable), + options=ClaudeAgentOptions(), + transport=transport, + ) + ] + + agent_spans = get_spans_by_op(otel_spans.get_finished_spans(), "invoke_agent") + assert len(agent_spans) == 1 + assert agent_spans[0].status.status_code == StatusCode.ERROR + assert get_messages(agent_spans[0], "gen_ai.input.messages") == [ + {"role": "user", "parts": [{"type": "text", "content": "Hello"}]} + ] + assert get_messages(agent_spans[0], "gen_ai.output.messages") == [] + + +@pytest.mark.asyncio +async def test_async_iterable_query_failure_between_turns_records_error_turn_otel( + otel_spans: InMemorySpanExporter, +) -> None: + transport = ReplayTransport( + load_cassette("multi_turn_response"), + fail_after_messages=3, + ) + seen_messages: list[str] = [] + + async def consume_query() -> None: + async for message in query( + prompt=async_prompt_messages( + user_prompt("Hello"), + user_prompt("Second prompt"), + ), + options=ClaudeAgentOptions(), + transport=transport, + ): + seen_messages.append(type(message).__name__) + + with pytest.raises(Exception, match="replay transport failed"): + await consume_query() + + assert seen_messages == ["SystemMessage", "AssistantMessage", "ResultMessage"] + agent_spans = sorted( + get_spans_by_op(otel_spans.get_finished_spans(), "invoke_agent"), + key=lambda span: span.start_time, + ) + assert len(agent_spans) == 2 + assert [span.status.status_code for span in agent_spans] == [ + StatusCode.UNSET, + StatusCode.ERROR, + ] + assert [ + get_messages(agent_span, "gen_ai.input.messages") for agent_span in agent_spans + ] == [ + [{"role": "user", "parts": [{"type": "text", "content": "Hello"}]}], + [ + { + "role": "user", + "parts": [{"type": "text", "content": "Second prompt"}], + } + ], + ] + assert [ + get_messages(agent_span, "gen_ai.output.messages") for agent_span in agent_spans + ] == [ + [ + { + "role": "assistant", + "parts": [{"type": "text", "content": "Hello! How can I help you?"}], + } + ], + [], + ] + assert { + get_attrs(agent_span)["gen_ai.conversation.id"] for agent_span in agent_spans + } == {"s-mt001"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "as_async_iterable", + [False, True], + ids=["string-prompt", "async-iterable-prompt"], +) +async def test_query_failure_after_final_result_does_not_add_turn_otel( + otel_spans: InMemorySpanExporter, + as_async_iterable: bool, +) -> None: + transport = ReplayTransport( + load_cassette("simple_text_response"), + fail_after_messages=3, + ) + seen_messages: list[str] = [] + + async def consume_query() -> None: + async for message in query( + prompt=text_query_prompt("Hello", as_async_iterable=as_async_iterable), + options=ClaudeAgentOptions(), + transport=transport, + ): + seen_messages.append(type(message).__name__) + + with pytest.raises(Exception, match="replay transport failed"): + await consume_query() + + assert seen_messages == ["SystemMessage", "AssistantMessage", "ResultMessage"] + agent_spans = get_spans_by_op(otel_spans.get_finished_spans(), "invoke_agent") + assert len(agent_spans) == 1 + assert agent_spans[0].status.status_code == StatusCode.UNSET + assert get_messages(agent_spans[0], "gen_ai.input.messages") == [ + {"role": "user", "parts": [{"type": "text", "content": "Hello"}]} + ] + assert get_messages(agent_spans[0], "gen_ai.output.messages") == [ + { + "role": "assistant", + "parts": [{"type": "text", "content": "The answer is 4."}], + } + ] + + # --- agent name override ---------------------------------------------------- diff --git a/tests/integrations/claude_agent_sdk/conftest.py b/tests/integrations/claude_agent_sdk/conftest.py index d7017c3dcf30..1f6082ad4f32 100644 --- a/tests/integrations/claude_agent_sdk/conftest.py +++ b/tests/integrations/claude_agent_sdk/conftest.py @@ -26,8 +26,14 @@ class ReplayTransport(Transport): then yields recorded messages from the cassette. """ - def __init__(self, messages: list[dict[str, Any]]) -> None: + def __init__( + self, + messages: list[dict[str, Any]], + *, + fail_after_messages: int | None = None, + ) -> None: self._messages = messages + self._fail_after_messages = fail_after_messages self._connected = False self._pending_control_ids: list[str] = [] self._control_event = anyio.Event() @@ -66,8 +72,12 @@ async def _read_impl(self) -> AsyncIterator[dict[str, Any]]: # Wait briefly for the user message write + end_input await anyio.sleep(0.01) - for msg in self._messages: + for index, msg in enumerate(self._messages): + if self._fail_after_messages == index: + raise RuntimeError("replay transport failed") yield msg + if self._fail_after_messages == len(self._messages): + raise RuntimeError("replay transport failed") async def close(self) -> None: self._connected = False diff --git a/weave/integrations/claude_agent_sdk/otel_integration.py b/weave/integrations/claude_agent_sdk/otel_integration.py index a6e2b390abbd..712cc62d63cf 100644 --- a/weave/integrations/claude_agent_sdk/otel_integration.py +++ b/weave/integrations/claude_agent_sdk/otel_integration.py @@ -4,10 +4,12 @@ sibling ``claude_agent_sdk_integration.py`` emits legacy Weave calls instead. The dispatcher selects this variant when ``WEAVE_USE_OTEL_V2`` is set. -Each ``query()`` call / ``ClaudeSDKClient`` turn becomes an ``invoke_agent`` -span, with a child ``chat`` span per model response and an ``execute_tool`` -span per tool call. The SDK reports token usage only on the final -``ResultMessage``, so it is attached to the last ``chat`` span. +String queries and ``ClaudeSDKClient.receive_response()`` calls each emit one +``invoke_agent`` turn. A standalone ``query()`` driven by an async prompt +iterable can return several results, so that path emits one turn per +``ResultMessage``. Each turn has a child ``chat`` span per model response and an +``execute_tool`` span per tool call. The SDK reports token usage only on a +turn's final ``ResultMessage``, so it is attached to the last ``chat`` span. The Claude Agent SDK has no agent-name concept, so ``invoke_agent`` spans default to ``claude_agent_sdk``. Users relabel them for a block of work with the @@ -21,6 +23,7 @@ import importlib import json import logging +from collections import deque from collections.abc import AsyncIterable, AsyncIterator from dataclasses import dataclass, field from datetime import datetime, timezone @@ -77,6 +80,74 @@ class _AssistantOutput: reasoning: Reasoning +@dataclass(frozen=True, slots=True) +class _PendingTurnInput: + messages: list[Message] + started_at: datetime + + +def _text_turn_input(prompt: str) -> _PendingTurnInput: + return _PendingTurnInput( + messages=[Message.user(prompt)], + started_at=datetime.now(timezone.utc), + ) + + +@dataclass(slots=True) +class _InputTracker: + """FIFO queue of user inputs, one entry per turn the SDK will answer.""" + + pending: deque[_PendingTurnInput] = field(default_factory=deque) + buffered: list[Message] = field(default_factory=list) + buffered_started_at: datetime | None = None + + @classmethod + def from_prompt(cls, prompt: Any) -> _InputTracker: + tracker = cls() + if isinstance(prompt, str): + tracker.pending.append(_text_turn_input(prompt)) + return tracker + + def record(self, prompt: dict[str, Any]) -> None: + message = _message_from_prompt(prompt) + if message is None: + return + if self.buffered_started_at is None: + self.buffered_started_at = datetime.now(timezone.utc) + self.buffered.append(message) + # ``shouldQuery: False`` means the SDK keeps collecting input instead of + # answering, so those messages belong to the next turn's input. + if prompt.get("shouldQuery") is not False: + self.pending.append( + _PendingTurnInput( + messages=self.buffered, + started_at=self.buffered_started_at, + ) + ) + self.buffered = [] + self.buffered_started_at = None + + def take(self) -> _PendingTurnInput: + if self.pending: + return self.pending.popleft() + if self.buffered: + pending = _PendingTurnInput( + messages=self.buffered, + started_at=self.buffered_started_at or datetime.now(timezone.utc), + ) + self.buffered = [] + self.buffered_started_at = None + return pending + return _PendingTurnInput( + messages=[], + started_at=datetime.now(timezone.utc), + ) + + def has_pending_turn(self) -> bool: + """Return whether a submitted input still expects an SDK result.""" + return bool(self.pending) + + def _message_from_prompt(prompt: dict[str, Any]) -> Message | None: if prompt.get("type") != "user": return None @@ -125,15 +196,13 @@ def _message_from_prompt(prompt: dict[str, Any]) -> Message | None: return Message(role="user", parts=parts) if parts else None -async def _track_prompt( +async def _track_async_prompt( prompt: AsyncIterable[dict[str, Any]], - input_messages: list[Message], + tracker: _InputTracker, ) -> AsyncIterator[dict[str, Any]]: async for message in prompt: try: - traced_message = _message_from_prompt(message) - if traced_message is not None: - input_messages.append(traced_message) + tracker.record(message) except Exception: logger.exception("claude_agent_sdk input tracing failed") yield message @@ -183,8 +252,8 @@ def _assistant_output_message( class _TurnState: """Mutable per-turn accumulator (one invoke_agent span and its children). - Not frozen: it accumulates as the message stream is consumed. Scope is one - turn, owned by ``_trace_turn``. + Not frozen: it accumulates as the message stream is consumed. Its scope is + one turn, regardless of which query adapter owns that turn. """ conversation: Conversation @@ -312,49 +381,123 @@ def _finalize_turn(state: _TurnState) -> None: ) -async def _trace_turn( +def _start_tracked_turn( + conversation: Conversation, turn_input: _PendingTurnInput +) -> Turn: + # Resolve per turn so an ``agent_name_override(...)`` block that spans only + # part of the stream labels only the turns inside it. + turn = conversation.start_turn(agent_name=resolve_agent_name(_DEFAULT_AGENT_NAME)) + turn.started_at = turn_input.started_at + turn.record(messages=turn_input.messages) + return turn + + +def _process_message_for_turn( + msg: Any, + state: _TurnState, + conversation_id_holder: list[str] | None = None, +) -> None: + """Record one SDK message without allowing tracing to break the stream.""" + try: + session_id = _process_message(msg, state) + if conversation_id_holder is not None and session_id: + conversation_id_holder[0] = session_id + except Exception: + logger.exception("claude_agent_sdk GenAI span processing failed") + + +async def _trace_single_turn( messages: AsyncIterator[Any], *, - input_messages: list[Message], + turn_input: _PendingTurnInput, conversation_id_holder: list[str] | None = None, ) -> AsyncIterator[Any]: - """Wrap a message stream, emitting the span tree for one turn. + """Trace a stream whose public API contract permits exactly one result. ``conversation_id_holder`` carries the SDK ``session_id`` across turns of a single ``ClaudeSDKClient``: the ``system/init`` message (which holds the session_id) is only sent on the first turn, so later turns must inherit it. """ - agent_name = resolve_agent_name(_DEFAULT_AGENT_NAME) + message_iterator = aiter(messages) conversation_id = ( conversation_id_holder[0] if conversation_id_holder is not None else "" ) with Conversation( conversation_id=conversation_id, - agent_name=agent_name, + agent_name=_DEFAULT_AGENT_NAME, continue_parent_trace=True, attributes=_INTEGRATION_OTEL_ATTRS, ) as conversation: - with conversation.start_turn() as turn: - state = _TurnState( - conversation=conversation, - turn=turn, - ) + turn = _start_tracked_turn(conversation, turn_input) + with turn: + state = _TurnState(conversation=conversation, turn=turn) try: - async for msg in messages: - try: - session_id = _process_message(msg, state) - if conversation_id_holder is not None and session_id: - conversation_id_holder[0] = session_id - except Exception: - # Never let span bookkeeping break the user's stream. - logger.exception( - "claude_agent_sdk GenAI span processing failed" - ) + async for msg in message_iterator: + _process_message_for_turn( + msg, + state, + conversation_id_holder, + ) yield msg + if isinstance(msg, ResultMessage): + break finally: - state.turn.record(messages=list(input_messages)) _finalize_turn(state) + # A ResultMessage completes the one allowed turn. Continue forwarding any + # teardown messages or errors without retroactively changing that turn. + async for msg in message_iterator: + yield msg + + +async def _trace_async_iterable_query( + messages: AsyncIterator[Any], + *, + inputs: _InputTracker, +) -> AsyncIterator[Any]: + """Trace every result produced by a standalone async-iterable query. + + Unlike the single-turn APIs, this output stream may contain several + ``ResultMessage`` boundaries. Looking ahead between turns keeps transport + failures after a completed result out of that completed turn while still + assigning failures to an already-submitted next input. + """ + with Conversation( + conversation_id="", + agent_name=_DEFAULT_AGENT_NAME, + continue_parent_trace=True, + attributes=_INTEGRATION_OTEL_ATTRS, + ) as conversation: + message_iterator = aiter(messages) + while True: + try: + next_message = await anext(message_iterator) + except StopAsyncIteration: + return + except BaseException: + if not inputs.has_pending_turn(): + raise + turn = _start_tracked_turn(conversation, inputs.take()) + with turn: + raise + + turn = _start_tracked_turn(conversation, inputs.take()) + with turn: + state = _TurnState(conversation=conversation, turn=turn) + try: + while True: + msg = next_message + _process_message_for_turn(msg, state) + yield msg + if isinstance(msg, ResultMessage): + break + try: + next_message = await anext(message_iterator) + except StopAsyncIteration: + return + finally: + _finalize_turn(state) + def _patched_process_query_wrapper(settings: IntegrationSettings) -> Any: """Wrap ``InternalClient.process_query`` (the async-gen behind ``query()``).""" @@ -375,19 +518,28 @@ async def wrapped_process_query( yield msg return - input_messages = [Message.user(prompt)] if isinstance(prompt, str) else [] - traced_prompt = ( - prompt - if isinstance(prompt, str) - else _track_prompt(prompt, input_messages) - ) + if isinstance(prompt, str): + inner = original_process_query( + self_client, + prompt=prompt, + options=options, + transport=transport, + ) + async for msg in _trace_single_turn( + inner, + turn_input=_text_turn_input(prompt), + ): + yield msg + return + + inputs = _InputTracker() inner = original_process_query( self_client, - prompt=traced_prompt, + prompt=_track_async_prompt(prompt, inputs), options=options, transport=transport, ) - async for msg in _trace_turn(inner, input_messages=input_messages): + async for msg in _trace_async_iterable_query(inner, inputs=inputs): yield msg return wrapped_process_query @@ -412,21 +564,19 @@ def patched_init(self: Any, *args: Any, **kwargs: Any) -> None: original_receive_response = self.receive_response # One-element holder so wrapped_query can hand the prompt to the # next receive_response() turn. - input_messages_holder: list[list[Message]] = [[]] + input_tracker_holder: list[_InputTracker] = [_InputTracker()] # Persists the SDK session_id across turns: system/init is only sent # on the first turn, so later turns inherit the conversation id here. conversation_id_holder: list[str] = [""] @wraps(original_query) async def wrapped_query(prompt: Any, session_id: str = "default") -> None: - input_messages = ( - [Message.user(prompt)] if isinstance(prompt, str) else [] - ) - input_messages_holder[0] = input_messages + inputs = _InputTracker.from_prompt(prompt) + input_tracker_holder[0] = inputs traced_prompt = ( prompt if isinstance(prompt, str) - else _track_prompt(prompt, input_messages) + else _track_async_prompt(prompt, inputs) ) return await original_query(traced_prompt, session_id=session_id) @@ -437,9 +587,9 @@ async def wrapped_receive_response() -> AsyncIterator[Any]: async for msg in inner: yield msg return - async for msg in _trace_turn( + async for msg in _trace_single_turn( inner, - input_messages=input_messages_holder[0], + turn_input=input_tracker_holder[0].take(), conversation_id_holder=conversation_id_holder, ): yield msg