fix: preserve annotation reply metadata

This commit is contained in:
minzifang 2026-06-03 16:28:11 +08:00
parent 86af36429d
commit 38118aff07
5 changed files with 9 additions and 13 deletions

View File

@ -1078,9 +1078,6 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
"""
extras = self._task_state.metadata.model_dump()
if self._task_state.metadata.annotation_reply:
del extras["annotation_reply"]
return MessageEndStreamResponse(
task_id=self._application_generate_entity.task_id,
id=self._message_id,

View File

@ -99,10 +99,6 @@ class AppGenerateResponseConverter[TBlockingResponse: AppBlockingResponse](ABC):
)
metadata["retriever_resources"] = updated_resources
# show annotation reply
if "annotation_reply" in metadata:
del metadata["annotation_reply"]
# show usage
if "usage" in metadata:
del metadata["usage"]

View File

@ -264,7 +264,7 @@ class TestAdvancedChatGenerateTaskPipeline:
assert "UPDATE messages" in stmt_str
assert "WHERE messages.id" in stmt_str
def test_message_end_to_stream_response_strips_annotation_reply(self):
def test_message_end_to_stream_response_preserves_annotation_reply(self):
pipeline = _make_pipeline()
pipeline._task_state.metadata.annotation_reply = AnnotationReply(
id="ann",
@ -273,7 +273,10 @@ class TestAdvancedChatGenerateTaskPipeline:
response = pipeline._message_end_to_stream_response()
assert "annotation_reply" not in response.metadata
assert response.metadata["annotation_reply"] == {
"id": "ann",
"account": {"id": "acc", "name": "acc"},
}
def test_handle_output_moderation_chunk_publishes_stop(self):
pipeline = _make_pipeline()

View File

@ -71,7 +71,7 @@ class TestAgentChatAppGenerateResponseConverterBlocking:
result = AgentChatAppGenerateResponseConverter.convert_blocking_simple_response(blocking)
assert "annotation_reply" not in result["metadata"]
assert result["metadata"]["annotation_reply"] == {"id": "a"}
assert "usage" not in result["metadata"]
def test_convert_blocking_simple_response_with_non_dict_metadata(self):
@ -169,7 +169,7 @@ class TestAgentChatAppGenerateResponseConverterStream:
assert items[2]["event"] == "message_end"
assert "metadata" in items[2]
metadata = items[2]["metadata"]
assert "annotation_reply" not in metadata
assert metadata["annotation_reply"] == {"id": "a"}
assert "usage" not in metadata
assert metadata["retriever_resources"] == [
{

View File

@ -75,7 +75,7 @@ class TestCompletionAppGenerateResponseConverter:
result = CompletionAppGenerateResponseConverter.convert_blocking_simple_response(blocking)
assert "annotation_reply" not in result["metadata"]
assert result["metadata"]["annotation_reply"] == {"a": 1}
assert "usage" not in result["metadata"]
assert result["metadata"]["retriever_resources"][0]["dataset_id"] == "dataset-1"
assert result["metadata"]["retriever_resources"][0]["document_id"] == "document-1"
@ -164,6 +164,6 @@ class TestCompletionAppGenerateResponseConverter:
assert result[0] == "ping"
assert result[1]["event"] == "message_end"
assert "annotation_reply" not in result[1]["metadata"]
assert result[1]["metadata"]["annotation_reply"] == {"a": 1}
assert "usage" not in result[1]["metadata"]
assert result[2]["event"] == "error"