Skip to content

Commit b6685c6

Browse files
jiapingzengjiapingzeng
committed
refactor events and constants
Signed-off-by: Jiaping Zeng <[email protected]>
1 parent 3cc5304 commit b6685c6

File tree

8 files changed

+411
-232
lines changed

8 files changed

+411
-232
lines changed

‎common/src/main/java/org/opensearch/ml/common/agui/AGUIConstants.java‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.agui;
7+
8+
/**
9+
* Constants for AG-UI (Agent UI) implementation.
10+
*
11+
* <p>Naming Conventions:
12+
* <ul>
13+
* <li>AGUI_ROLE_* - Message role identifiers</li>
14+
* <li>AGUI_PARAM_* - Internal parameter keys (snake_case format)</li>
15+
* <li>AGUI_FIELD_* - External API field names (camelCase format)</li>
16+
* <li>AGUI_EVENT_* - Event type identifiers</li>
17+
* <li>AGUI_PREFIX_* - ID prefixes for generated identifiers</li>
18+
* </ul>
19+
*/
20+
public final class AGUIConstants {
21+
22+
// ========== Message Roles ==========
23+
24+
/** Role identifier for assistant messages */
25+
public static final String AGUI_ROLE_ASSISTANT = "assistant";
26+
27+
/** Role identifier for user messages */
28+
public static final String AGUI_ROLE_USER = "user";
29+
30+
/** Role identifier for tool result messages */
31+
public static final String AGUI_ROLE_TOOL = "tool";
32+
33+
// ========== Parameter Keys (Internal) ==========
34+
35+
/** Parameter key for AG-UI thread identifier */
36+
public static final String AGUI_PARAM_THREAD_ID = "agui_thread_id";
37+
38+
/** Parameter key for AG-UI run identifier */
39+
public static final String AGUI_PARAM_RUN_ID = "agui_run_id";
40+
41+
/** Parameter key for AG-UI messages array */
42+
public static final String AGUI_PARAM_MESSAGES = "agui_messages";
43+
44+
/** Parameter key for AG-UI tools array */
45+
public static final String AGUI_PARAM_TOOLS = "agui_tools";
46+
47+
/** Parameter key for AG-UI context array */
48+
public static final String AGUI_PARAM_CONTEXT = "agui_context";
49+
50+
/** Parameter key for AG-UI state object */
51+
public static final String AGUI_PARAM_STATE = "agui_state";
52+
53+
/** Parameter key for AG-UI forwarded properties */
54+
public static final String AGUI_PARAM_FORWARDED_PROPS = "agui_forwarded_props";
55+
56+
/** Parameter key for AG-UI tool call results */
57+
public static final String AGUI_PARAM_TOOL_CALL_RESULTS = "agui_tool_call_results";
58+
59+
/** Parameter key for AG-UI assistant tool call messages */
60+
public static final String AGUI_PARAM_ASSISTANT_TOOL_CALL_MESSAGES = "agui_assistant_tool_call_messages";
61+
62+
/** Parameter key for backend tool names (used for filtering) */
63+
public static final String AGUI_PARAM_BACKEND_TOOL_NAMES = "backend_tool_names";
64+
65+
// ========== Field Names (External API) ==========
66+
67+
/** Field name for thread identifier in AG-UI input */
68+
public static final String AGUI_FIELD_THREAD_ID = "threadId";
69+
70+
/** Field name for run identifier in AG-UI input */
71+
public static final String AGUI_FIELD_RUN_ID = "runId";
72+
73+
/** Field name for messages array in AG-UI input */
74+
public static final String AGUI_FIELD_MESSAGES = "messages";
75+
76+
/** Field name for tools array in AG-UI input */
77+
public static final String AGUI_FIELD_TOOLS = "tools";
78+
79+
/** Field name for context array in AG-UI input */
80+
public static final String AGUI_FIELD_CONTEXT = "context";
81+
82+
/** Field name for state object in AG-UI input */
83+
public static final String AGUI_FIELD_STATE = "state";
84+
85+
/** Field name for forwarded properties in AG-UI input */
86+
public static final String AGUI_FIELD_FORWARDED_PROPS = "forwardedProps";
87+
88+
/** Field name for message role */
89+
public static final String AGUI_FIELD_ROLE = "role";
90+
91+
/** Field name for message content */
92+
public static final String AGUI_FIELD_CONTENT = "content";
93+
94+
/** Field name for tool call identifier */
95+
public static final String AGUI_FIELD_TOOL_CALL_ID = "toolCallId";
96+
97+
/** Field name for tool calls array */
98+
public static final String AGUI_FIELD_TOOL_CALLS = "toolCalls";
99+
100+
/** Field name for message identifier */
101+
public static final String AGUI_FIELD_ID = "id";
102+
103+
/** Field name for tool call type */
104+
public static final String AGUI_FIELD_TYPE = "type";
105+
106+
/** Field name for function object in tool calls */
107+
public static final String AGUI_FIELD_FUNCTION = "function";
108+
109+
/** Field name for function name */
110+
public static final String AGUI_FIELD_NAME = "name";
111+
112+
/** Field name for function arguments */
113+
public static final String AGUI_FIELD_ARGUMENTS = "arguments";
114+
}

‎common/src/main/java/org/opensearch/ml/common/agui/AGUIInputConverter.java‎

Lines changed: 61 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,26 @@
55

66
package org.opensearch.ml.common.agui;
77

8+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_CONTENT;
9+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_CONTEXT;
10+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_FORWARDED_PROPS;
11+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_MESSAGES;
12+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_ROLE;
13+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_RUN_ID;
14+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_STATE;
15+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_THREAD_ID;
16+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_TOOLS;
17+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_FIELD_TOOL_CALL_ID;
18+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_CONTEXT;
19+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_FORWARDED_PROPS;
20+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_MESSAGES;
21+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_RUN_ID;
22+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_STATE;
23+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_THREAD_ID;
24+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_TOOLS;
25+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_TOOL_CALL_RESULTS;
26+
import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_ROLE_USER;
27+
828
import java.util.HashMap;
929
import java.util.List;
1030
import java.util.Map;
@@ -27,7 +47,28 @@ public class AGUIInputConverter {
2747
public static boolean isAGUIInput(String inputJson) {
2848
try {
2949
JsonObject jsonObj = JsonParser.parseString(inputJson).getAsJsonObject();
30-
return jsonObj.has("threadId") && jsonObj.has("runId") && jsonObj.has("messages") && jsonObj.has("tools");
50+
51+
// Check required fields exist
52+
if (!jsonObj.has(AGUI_FIELD_THREAD_ID)
53+
|| !jsonObj.has(AGUI_FIELD_RUN_ID)
54+
|| !jsonObj.has(AGUI_FIELD_MESSAGES)
55+
|| !jsonObj.has(AGUI_FIELD_TOOLS)) {
56+
return false;
57+
}
58+
59+
// Validate messages is an array
60+
JsonElement messages = jsonObj.get(AGUI_FIELD_MESSAGES);
61+
if (!messages.isJsonArray()) {
62+
return false;
63+
}
64+
65+
// Validate tools is an array
66+
JsonElement tools = jsonObj.get(AGUI_FIELD_TOOLS);
67+
if (!tools.isJsonArray()) {
68+
return false;
69+
}
70+
71+
return true;
3172
} catch (Exception e) {
3273
log.debug("Failed to parse input as JSON for AG-UI detection", e);
3374
return false;
@@ -38,37 +79,37 @@ public static AgentMLInput convertFromAGUIInput(String aguiInputJson, String age
3879
try {
3980
JsonObject aguiInput = JsonParser.parseString(aguiInputJson).getAsJsonObject();
4081

41-
String threadId = getStringField(aguiInput, "threadId");
42-
String runId = getStringField(aguiInput, "runId");
43-
JsonElement state = aguiInput.get("state");
44-
JsonElement messages = aguiInput.get("messages");
45-
JsonElement tools = aguiInput.get("tools");
46-
JsonElement context = aguiInput.get("context");
47-
JsonElement forwardedProps = aguiInput.get("forwardedProps");
82+
String threadId = getStringField(aguiInput, AGUI_FIELD_THREAD_ID);
83+
String runId = getStringField(aguiInput, AGUI_FIELD_RUN_ID);
84+
JsonElement state = aguiInput.get(AGUI_FIELD_STATE);
85+
JsonElement messages = aguiInput.get(AGUI_FIELD_MESSAGES);
86+
JsonElement tools = aguiInput.get(AGUI_FIELD_TOOLS);
87+
JsonElement context = aguiInput.get(AGUI_FIELD_CONTEXT);
88+
JsonElement forwardedProps = aguiInput.get(AGUI_FIELD_FORWARDED_PROPS);
4889

4990
Map<String, String> parameters = new HashMap<>();
50-
parameters.put("agui_thread_id", threadId);
51-
parameters.put("agui_run_id", runId);
91+
parameters.put(AGUI_PARAM_THREAD_ID, threadId);
92+
parameters.put(AGUI_PARAM_RUN_ID, runId);
5293

5394
if (state != null) {
54-
parameters.put("agui_state", gson.toJson(state));
95+
parameters.put(AGUI_PARAM_STATE, gson.toJson(state));
5596
}
5697

5798
if (messages != null) {
58-
parameters.put("agui_messages", gson.toJson(messages));
99+
parameters.put(AGUI_PARAM_MESSAGES, gson.toJson(messages));
59100
extractUserQuestion(messages, parameters);
60101
}
61102

62103
if (tools != null) {
63-
parameters.put("agui_tools", gson.toJson(tools));
104+
parameters.put(AGUI_PARAM_TOOLS, gson.toJson(tools));
64105
}
65106

66107
if (context != null) {
67-
parameters.put("agui_context", gson.toJson(context));
108+
parameters.put(AGUI_PARAM_CONTEXT, gson.toJson(context));
68109
}
69110

70111
if (forwardedProps != null) {
71-
parameters.put("agui_forwarded_props", gson.toJson(forwardedProps));
112+
parameters.put(AGUI_PARAM_FORWARDED_PROPS, gson.toJson(forwardedProps));
72113
}
73114
RemoteInferenceInputDataSet inputDataSet = RemoteInferenceInputDataSet.builder().parameters(parameters).build();
74115
AgentMLInput agentMLInput = new AgentMLInput(
@@ -93,40 +134,6 @@ private static String getStringField(JsonObject obj, String fieldName) {
93134
return element != null && !element.isJsonNull() ? element.getAsString() : null;
94135
}
95136

96-
public static JsonObject reconstructAGUIInput(Map<String, String> parameters) {
97-
JsonObject aguiInput = new JsonObject();
98-
99-
try {
100-
String threadId = parameters.get("agui_thread_id");
101-
String runId = parameters.get("agui_run_id");
102-
String stateJson = parameters.get("agui_state");
103-
String messagesJson = parameters.get("agui_messages");
104-
String toolsJson = parameters.get("agui_tools");
105-
String contextJson = parameters.get("agui_context");
106-
String forwardedPropsJson = parameters.get("agui_forwarded_props");
107-
108-
if (threadId != null)
109-
aguiInput.addProperty("threadId", threadId);
110-
if (runId != null)
111-
aguiInput.addProperty("runId", runId);
112-
if (stateJson != null)
113-
aguiInput.add("state", JsonParser.parseString(stateJson));
114-
if (messagesJson != null)
115-
aguiInput.add("messages", JsonParser.parseString(messagesJson));
116-
if (toolsJson != null)
117-
aguiInput.add("tools", JsonParser.parseString(toolsJson));
118-
if (contextJson != null)
119-
aguiInput.add("context", JsonParser.parseString(contextJson));
120-
if (forwardedPropsJson != null)
121-
aguiInput.add("forwardedProps", JsonParser.parseString(forwardedPropsJson));
122-
123-
} catch (Exception e) {
124-
log.error("Failed to reconstruct AG-UI input from parameters", e);
125-
}
126-
127-
return aguiInput;
128-
}
129-
130137
private static void extractUserQuestion(JsonElement messages, Map<String, String> parameters) {
131138
if (messages == null || !messages.isJsonArray()) {
132139
throw new IllegalArgumentException("Invalid AG-UI messages");
@@ -140,12 +147,12 @@ private static void extractUserQuestion(JsonElement messages, Map<String, String
140147
for (JsonElement messageElement : messages.getAsJsonArray()) {
141148
if (messageElement.isJsonObject()) {
142149
JsonObject message = messageElement.getAsJsonObject();
143-
JsonElement roleElement = message.get("role");
144-
JsonElement contentElement = message.get("content");
145-
JsonElement toolCallIdElement = message.get("toolCallId");
150+
JsonElement roleElement = message.get(AGUI_FIELD_ROLE);
151+
JsonElement contentElement = message.get(AGUI_FIELD_CONTENT);
152+
JsonElement toolCallIdElement = message.get(AGUI_FIELD_TOOL_CALL_ID);
146153

147154
if (roleElement != null
148-
&& "user".equals(roleElement.getAsString())
155+
&& AGUI_ROLE_USER.equals(roleElement.getAsString())
149156
&& contentElement != null
150157
&& !contentElement.isJsonNull()) {
151158

@@ -173,7 +180,7 @@ private static void extractUserQuestion(JsonElement messages, Map<String, String
173180

174181
// Set appropriate parameters based on what was found
175182
if (toolCallResults != null) {
176-
parameters.put("agui_tool_call_results", toolCallResults);
183+
parameters.put(AGUI_PARAM_TOOL_CALL_RESULTS, toolCallResults);
177184
log.debug("Detected AG-UI tool call results: {}", toolCallResults);
178185
} else if (lastUserMessage != null) {
179186
parameters.put("question", lastUserMessage);

‎common/src/main/java/org/opensearch/ml/common/agui/BaseEvent.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,11 @@
88
import java.io.IOException;
99
import java.util.Map;
1010

11+
import org.opensearch.common.xcontent.XContentFactory;
1112
import org.opensearch.core.common.io.stream.StreamInput;
1213
import org.opensearch.core.common.io.stream.StreamOutput;
1314
import org.opensearch.core.common.io.stream.Writeable;
15+
import org.opensearch.core.xcontent.ToXContent;
1416
import org.opensearch.core.xcontent.ToXContentFragment;
1517
import org.opensearch.core.xcontent.XContentBuilder;
1618

@@ -63,4 +65,12 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
6365
}
6466

6567
protected abstract void addEventSpecificFields(XContentBuilder builder, Params params) throws IOException;
68+
69+
public String toJsonString() {
70+
try {
71+
return toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString();
72+
} catch (IOException e) {
73+
throw new RuntimeException("Failed to serialize event to JSON", e);
74+
}
75+
}
6676
}

0 commit comments

Comments
 (0)