ag-ui-protocol / ag-ui-protocol/ag-ui

[Feature]: Scoped message snapshots

未關閉
#1,308 0 則留言 2 個 reaction 已指派 0 人 在 GitHub 檢視
enhancement
主要語言
Python
星號
15.9k
分支
1.4k
平均合併
1 天 17 小時
30 天內合併 PR
163

描述

### Pre-flight Checklist

- [x] I have searched existing issues and this hasn't been requested yet.

### Problem or Motivation

The documentation for [compacting event streams](https://github.com/ag-ui-protocol/ag-ui/blob/d5d9573cef5b76b1407aa478107be6fa64d8be2b/docs/concepts/serialization.mdx?plain=1#L107-L135) currently shows a trivial example of compacting text message events of a single message into a messages snapshot event. That works perfectly fine even for multiple messages as long as they are contained within the scope of a single, "flat" run.

However, using a messages snapshot event with the [current definition](https://github.com/ag-ui-protocol/ag-ui/blob/d5d9573cef5b76b1407aa478107be6fa64d8be2b/docs/concepts/events.mdx?plain=1#L461-L475) loses important information in at least two use cases:

1. If the event stream contains more than one run, e.g.

```python
from ag_ui.core import *

events = [
RunStartedEvent(thread_id="thread", run_id="run1"),
TextMessageStartEvent(message_id="msg1"),
TextMessageContentEvent(message_id="msg1", delta="msg1 delta1"),
TextMessageContentEvent(message_id="msg1", delta="delta2"),
TextMessageEndEvent(message_id="msg1"),
RunFinishedEvent(thread_id="thread", run_id="run1"),
RunStartedEvent(thread_id="thread", run_id="run2"),
TextMessageStartEvent(message_id="msg2"),
TextMessageContentEvent(message_id="msg2", delta="msg2 delta1"),
TextMessageContentEvent(message_id="msg2", delta="delta2"),
TextMessageEndEvent(message_id="msg2"),
RunFinishedEvent(thread_id="thread", run_id="run2"),
]
```

the compacted event stream could look something like this:

```python
compacted = [
RunStartedEvent(thread_id="thread", run_id="run1"),
RunFinishedEvent(thread_id="thread", run_id="run1"),
RunStartedEvent(thread_id="thread", run_id="run2"),
MessagesSnapshotEvent(
messages=[
AssistantMessage(id="msg1", content="msg1 delta1delta2"),
AssistantMessage(id="msg2", content="msg2 delta1delta2"),
]
),
RunStartedEvent(thread_id="thread", run_id="run2"),
]
```

This correctly snapshots all messages, but it loses the association of `"msg1"` to `"run1"`. This association is important however for [branching](https://github.com/ag-ui-protocol/ag-ui/blob/d5d9573cef5b76b1407aa478107be6fa64d8be2b/docs/concepts/serialization.mdx?plain=1#L66-L89): since branching happens on a per-run level, compacting the event stream as shown above makes it now impossible for the user to branch after `"msg1"`.
2. The exact same issue of message disassociation also happens for non-"flat" runs, e.g. ones that include at least one [step](https://github.com/ag-ui-protocol/ag-ui/blob/d5d9573cef5b76b1407aa478107be6fa64d8be2b/docs/concepts/events.mdx?plain=1#L126-L154):

```python
from ag_ui.core import *

events = [
RunStartedEvent(thread_id="thread", run_id="run"),
StepStartedEvent(step_name="step"),
TextMessageStartEvent(message_id="msg1"),
TextMessageContentEvent(message_id="msg1", delta="msg1 delta1"),
TextMessageContentEvent(message_id="msg1", delta="delta2"),
TextMessageEndEvent(message_id="msg1"),
StepFinishedEvent(step_name="step"),
TextMessageStartEvent(message_id="msg2"),
TextMessageContentEvent(message_id="msg2", delta="msg2 delta1"),
TextMessageContentEvent(message_id="msg2", delta="delta2"),
TextMessageEndEvent(message_id="msg2"),
RunFinishedEvent(thread_id="thread", run_id="run"),
]

compacted = [
RunStartedEvent(thread_id="thread", run_id="run"),
StepStartedEvent(step_name="step"),
StepFinishedEvent(step_name="step"),
MessagesSnapshotEvent(
messages=[
AssistantMessage(id="msg1", content="msg1 delta1delta2"),
AssistantMessage(id="msg2", content="msg2 delta1delta2"),
]
),
RunFinishedEvent(thread_id="thread", run_id="run"),
]
```

The association of `"msg1"` to `"step"` is lost.

### Proposed Solution

Instead of requiring the messages snapshot to

https://github.com/ag-ui-protocol/ag-ui/blob/d5d9573cef5b76b1407aa478107be6fa64d8be2b/docs/concepts/events.mdx?plain=1#L463

we could add a scope, e.g. thread, run, step, etc., as well as a scope ID to it:

- The current behavior would be `scope="thread", scope_id=thread_id`,
- scoping it to a run would be `scope="run", scope_id=run_id`, and
- scoping it to a step would be `scope="step", scope_id=step_id`.

Both the `thread_id` and `run_id` are already required input for the agent so no additional changes are needed other than the agent setting the proper `scope_id`. For steps we currently [don't have an ID at all](https://github.com/ag-ui-protocol/ag-ui/blob/d5d9573cef5b76b1407aa478107be6fa64d8be2b/docs/concepts/events.mdx?plain=1#L126-L154) that would have to be added.

Example:

```python
from ag_ui.core import *

events = [
RunStartedEvent(thread_id="thread-id", run_id="run1-id"),
StepStartedEvent(step_id="step-id"),
TextMessageStartEvent(message_id="msg1"),
TextMessageContentEvent(message_id="msg1", delta="msg1 delta1"),
TextMessageContentEvent(message_id="msg1", delta="delta2"),
TextMessageEndEvent(message_id="msg1"),
StepFinishedEvent(step_id="step-id"),
TextMessageStartEvent(message_id="msg2"),
TextMessageContentEvent(message_id="msg2", delta="msg2 delta1"),
TextMessageContentEvent(message_id="msg2", delta="delta2"),
TextMessageEndEvent(message_id="msg2"),
RunFinishedEvent(thread_id="thread-id", run_id="run1-id"),
RunStartedEvent(thread_id="thread-id", run_id="run2-id"),
TextMessageStartEvent(message_id="msg3"),
TextMessageContentEvent(message_id="msg3", delta="msg3 delta1"),
TextMessageContentEvent(message_id="msg3", delta="delta2"),
TextMessageEndEvent(message_id="msg3"),
RunFinishedEvent(thread_id="thread-id", run_id="run2-id"),
]

compact = [
RunStartedEvent(thread_id="thread-id", run_id="run1-id"),
StepStartedEvent(step_id="step-id"),
MessagesSnapshotEvent(
scope="step",
scope_id="step-id",
messages=[AssistantMessage(id="msg1", content="msg1 delta1delta2")],
),
StepFinishedEvent(step_id="step-id"),
MessagesSnapshotEvent(
scope="run",
scope_id="run1-id",
messages=[AssistantMessage(id="msg2", content="msg2 delta1delta2")],
),
RunFinishedEvent(thread_id="thread-id", run_id="run1-id"),
RunStartedEvent(thread_id="thread-id", run_id="run2-id"),
MessagesSnapshotEvent(
scope="run",
scope_id="run2-id",
messages=[AssistantMessage(id="msg3", content="msg3 delta1delta2")],
),
RunFinishedEvent(thread_id="thread-id", run_id="run2-id"),
]
```

The compaction algorithm itself could also take a `compact_scope: Literal["thread", "run", "step"]` parameter that controls if steps and runs should be preserved or if they should be removed if they are empty after compaction. The example above would be for `compact_scope="step"`. For `compact_scope="run"` it could be

```python
compact = [
RunStartedEvent(thread_id="thread-id", run_id="run1-id"),
MessagesSnapshotEvent(
scope="run",
scope_id="run1-id",
messages=[
AssistantMessage(id="msg1", content="msg1 delta1delta2"),
AssistantMessage(id="msg2", content="msg2 delta1delta2")
],
),
RunFinishedEvent(thread_id="thread-id", run_id="run1-id"),
RunStartedEvent(thread_id="thread-id", run_id="run2-id"),
MessagesSnapshotEvent(
scope="run",
scope_id="run2-id",
messages=[AssistantMessage(id="msg3", content="msg3 delta1delta2")],
),
RunFinishedEvent(thread_id="thread-id", run_id="run2-id"),
]
```

and for `compact_scope="thread"`

```python
compact = [
MessagesSnapshotEvent(
scope="thread",
scope_id="thread-id",
messages=[
AssistantMessage(id="msg1", content="msg1 delta1delta2"),
AssistantMessage(id="msg2", content="msg2 delta1delta2"),
AssistantMessage(id="msg3", content="msg3 delta1delta2")
],
),
]
```

The latter is basically again how compaction is currently described in the documentation.

### Alternatives Considered

- Instead of just sending one messages snapshot event, compaction could result one multiple ones, each for a given "scope":

```python
from ag_ui.core import *

compacted = [
RunStartedEvent(thread_id="thread", run_id="run1"),
MessagesSnapshotEvent(
messages=[
AssistantMessage(id="msg1", content="msg1 delta1delta2"),
]
),
RunFinishedEvent(thread_id="thread", run_id="run1"),
RunStartedEvent(thread_id="thread", run_id="run2"),
MessagesSnapshotEvent(
messages=[
AssistantMessage(id="msg1", content="msg1 delta1delta2"),
AssistantMessage(id="msg2", content="msg2 delta1delta2"),
]
),
RunFinishedEvent(thread_id="thread", run_id="run2"),
]
```

This will indeed not lose the scope association of the messages, but has two downsides:

1. The client now needs to perform the normalization to avoid showing the same message multiple times.
2. The number of messages in the compacted event stream grows with `O(n**2)` where `n` is the number of messages in the conversation. This completely defeats the compaction.
- Instead of using the messages snapshot event, message compaction could be built on the `*MessageChunk` events, e.g. [`TextMessageChunk`](https://github.com/ag-ui-protocol/ag-ui/blob/d5d9573cef5b76b1407aa478107be6fa64d8be2b/docs/concepts/events.mdx?plain=1#L238-L257) as they at least let's us fold all the corresponding start, content, and end events into one:

```python
from ag_ui.core import *

compacted = [
RunStartedEvent(thread_id="thread", run_id="run1"),
TextMessageChunkEvent(message_id="msg1", role="assistant", delta="msg1 delta1delta2"),
RunFinishedEvent(thread_id="thread", run_id="run1"),
RunStartedEvent(thread_id="thread", run_id="run2"),
TextMessageChunkEvent(message_id="msg2", role="assistant", delta="msg2 delta1delta2"),
RunFinishedEvent(thread_id="thread", run_id="run2"),
]
```

This should work, but it means that the client still has to perform the `*MessageChunkEvent` to `Message` conversion, e.g. assembling an `AssistantMessage` from text messages and tool calls.

### Additional Context

This was initially discussed on [Discord](https://discord.com/channels/1379082175625953370/1482026446775124160/1482026446775124160)

貢獻指南

開啟貢獻指南

評估

這個 Issue 還沒有評估資料。

把新 issue 寄到你的電子郵件信箱

精選適合新手參與的 GitHub issue 摘要。