diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java index fbc672d4a..2703255ea 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java @@ -42,8 +42,10 @@ public class ChatRequestEvent extends Event { /** * The fixed cross-language attribute schema for this event: {@code model} and {@code messages} - * are required, {@code prompt_args} and {@code output_schema} are optional, and no other - * attribute is allowed. {@link #fromEvent} enforces it at the JSON boundary. + * are required, {@code prompt_args}, {@code output_schema}, and {@code user_context} are + * optional, and no other attribute is allowed. {@link #fromEvent} enforces it at the JSON + * boundary. {@code user_context} is application-defined data propagated to the response; it is + * not part of the prompt. */ private static final List ATTRIBUTE_SCHEMA = List.of( @@ -54,13 +56,15 @@ public class ChatRequestEvent extends Event { ChatMessage.class, Map.class), BuiltInAttribute.optional("prompt_args", Map.class), - BuiltInAttribute.optionalUntyped("output_schema")); + BuiltInAttribute.optionalUntyped("output_schema"), + BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class)); public ChatRequestEvent( String model, List messages, @Nullable Map promptArgs, - @Nullable Object outputSchema) { + @Nullable Object outputSchema, + @Nullable Map userContext) { super(EVENT_TYPE); setAttr("model", model); setAttr("messages", new ArrayList<>(messages)); @@ -68,6 +72,15 @@ public ChatRequestEvent( if (outputSchema != null) { setAttr("output_schema", outputSchema); } + UserDefinedContext.set(this, userContext); + } + + public ChatRequestEvent( + String model, + List messages, + @Nullable Map promptArgs, + @Nullable Object outputSchema) { + this(model, messages, promptArgs, outputSchema, null); } public ChatRequestEvent( @@ -108,10 +121,10 @@ private static Map normalizeAttributes(Map attri * Reconstructs a typed ChatRequestEvent from a base Event, deserializing nested types. * *

Enforces the fixed cross-language attribute schema: {@code model} and {@code messages} are - * required, {@code prompt_args} and {@code output_schema} are optional, no other attribute is - * allowed, {@code model} must be a string, and every {@code messages} element must be a {@link - * ChatMessage} or its serialized map — an element such as a bare number is rejected rather than - * silently dropped. + * required, {@code prompt_args}, {@code output_schema}, and {@code user_context} are optional, + * no other attribute is allowed, {@code model} must be a string, and every {@code messages} + * element must be a {@link ChatMessage} or its serialized map — an element such as a bare + * number is rejected rather than silently dropped. * * @param event the base event containing chat request data in attributes * @return a typed ChatRequestEvent @@ -145,4 +158,15 @@ public Map getPromptArgs() { Map args = (Map) getAttr("prompt_args"); return args != null ? args : Collections.emptyMap(); } + + /** + * Returns application-defined data to propagate to the matching {@link ChatResponseEvent}, or + * {@code null} when this request omitted it. The map is not added to messages or prompt + * arguments. + */ + @JsonIgnore + @Nullable + public Map getUserContext() { + return UserDefinedContext.get(this); + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java index c4982ceaa..487756006 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java @@ -25,6 +25,8 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.chat.messages.ChatResult; +import javax.annotation.Nullable; + import java.util.List; import java.util.Map; import java.util.Objects; @@ -54,7 +56,8 @@ public class ChatResponseEvent extends Event { BuiltInAttribute.optional(RETRY_COUNT, Number.class), BuiltInAttribute.optional(TOTAL_RETRY_WAIT_SEC, Number.class), BuiltInAttribute.optionalUntyped(STRUCTURED_OUTPUT), - BuiltInAttribute.optional(MODEL_ROUTING, Map.class)); + BuiltInAttribute.optional(MODEL_ROUTING, Map.class), + BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class)); private static final ObjectMapper MAPPER = new ObjectMapper(); @@ -143,6 +146,25 @@ public void setRoutingMetadata(Map value) { setAttr(MODEL_ROUTING, value); } + /** + * Attaches application-defined data from the originating chat request. {@code null} omits it. + * The value is not derived from event attachments. + */ + public ChatResponseEvent withUserContext(@Nullable Map userContext) { + UserDefinedContext.set(this, userContext); + return this; + } + + /** + * Returns the originating request's application-defined data, or {@code null} when it was + * omitted. + */ + @JsonIgnore + @Nullable + public Map getUserContext() { + return UserDefinedContext.get(this); + } + @JsonIgnore public String getStatus() { return (String) getAttr(STATUS); diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java index 668075be1..2cec65ddc 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java @@ -23,6 +23,8 @@ import com.fasterxml.jackson.annotation.JsonProperty; import org.apache.flink.agents.api.Event; +import javax.annotation.Nullable; + import java.util.List; import java.util.Map; import java.util.UUID; @@ -38,7 +40,8 @@ public class ContextRetrievalRequestEvent extends Event { List.of( BuiltInAttribute.required("query", String.class), BuiltInAttribute.required("vector_store", String.class), - BuiltInAttribute.optional("max_results", Number.class)); + BuiltInAttribute.optional("max_results", Number.class), + BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class)); public ContextRetrievalRequestEvent(String query, String vectorStore) { this(query, vectorStore, DEFAULT_MAX_RESULTS); @@ -50,14 +53,32 @@ public ContextRetrievalRequestEvent(String query, String vectorStore) { * @param query the search query text * @param vectorStore the name of the vector store to query * @param maxResults the maximum number of results to return; must be positive + * @param userContext application-defined data propagated to the response; {@code null} omits it * @throws IllegalArgumentException if {@code maxResults} is non-positive */ - public ContextRetrievalRequestEvent(String query, String vectorStore, int maxResults) { + public ContextRetrievalRequestEvent( + String query, + String vectorStore, + int maxResults, + @Nullable Map userContext) { super(EVENT_TYPE); validateMaxResults(maxResults); setAttr("query", query); setAttr("vector_store", vectorStore); setAttr("max_results", maxResults); + UserDefinedContext.set(this, userContext); + } + + /** + * Creates a new context retrieval request event. + * + * @param query the search query text + * @param vectorStore the name of the vector store to query + * @param maxResults the maximum number of results to return; must be positive + * @throws IllegalArgumentException if {@code maxResults} is non-positive + */ + public ContextRetrievalRequestEvent(String query, String vectorStore, int maxResults) { + this(query, vectorStore, maxResults, null); } /** @@ -93,8 +114,8 @@ private static void validateMaxResults(int maxResults) { * Reconstructs a typed ContextRetrievalRequestEvent from a base Event. * *

Enforces the fixed cross-language schema: {@code query} and {@code vector_store} are - * required strings, {@code max_results} is an optional positive number, and no other attribute - * is allowed. + * required strings, {@code max_results} is an optional positive number, {@code user_context} is + * an optional map, and no other attribute is allowed. * * @param event the base event containing context retrieval request data in attributes * @return a typed ContextRetrievalRequestEvent @@ -125,6 +146,16 @@ public int getMaxResults() { return ((Number) getAttr("max_results")).intValue(); } + /** + * Returns application-defined data to propagate to the matching response, or {@code null} when + * this request omitted it. The map is not added to the vector-store query. + */ + @JsonIgnore + @Nullable + public Map getUserContext() { + return UserDefinedContext.get(this); + } + @Override public String toString() { return "ContextRetrievalRequestEvent{" diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java index f8481bc17..611c6ce6a 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java @@ -25,6 +25,8 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.vectorstores.Document; +import javax.annotation.Nullable; + import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -45,7 +47,8 @@ public class ContextRetrievalResponseEvent extends Event { "documents", "a Document or its serialized map", Document.class, - Map.class)); + Map.class), + BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class)); public ContextRetrievalResponseEvent(UUID requestId, String query, List documents) { super(EVENT_TYPE); @@ -88,8 +91,9 @@ private static Map normalizeAttributes(Map attri * types. * *

Enforces the fixed cross-language schema: {@code request_id}, {@code query}, and {@code - * documents} are required and no other attribute is allowed; {@code query} must be a string and - * every {@code documents} element must be a {@link Document} or its serialized map. + * documents} are required, {@code user_context} is an optional map, and no other attribute is + * allowed; {@code query} must be a string and every {@code documents} element must be a {@link + * Document} or its serialized map. * * @param event the base event containing context retrieval response data in attributes * @return a typed ContextRetrievalResponseEvent @@ -120,6 +124,26 @@ public List getDocuments() { return (List) getAttr("documents"); } + /** + * Attaches application-defined data from the originating retrieval request. {@code null} omits + * it. + */ + public ContextRetrievalResponseEvent withUserContext( + @Nullable Map userContext) { + UserDefinedContext.set(this, userContext); + return this; + } + + /** + * Returns the originating request's application-defined data, or {@code null} when it was + * omitted. + */ + @JsonIgnore + @Nullable + public Map getUserContext() { + return UserDefinedContext.get(this); + } + @Override public String toString() { return "ContextRetrievalResponseEvent{" diff --git a/api/src/main/java/org/apache/flink/agents/api/event/DecisionRequestEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/DecisionRequestEvent.java index b81b2baba..66be51b6f 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/DecisionRequestEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/DecisionRequestEvent.java @@ -43,7 +43,8 @@ public class DecisionRequestEvent extends Event { List.of( BuiltInAttribute.required("model", String.class), BuiltInAttribute.requiredUntyped("request"), - BuiltInAttribute.optionalUntyped("request_params")); + BuiltInAttribute.optionalUntyped("request_params"), + BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class)); /** The model is the registered resource name, not a provider model identifier. */ public DecisionRequestEvent(String model, DecisionRequest request) { @@ -53,10 +54,25 @@ public DecisionRequestEvent(String model, DecisionRequest request) { /** Requests a decision with provider-specific parameters for this call. */ public DecisionRequestEvent( String model, DecisionRequest request, @Nullable Map requestParams) { + this(model, request, requestParams, null); + } + + /** + * Requests a decision with provider-specific parameters and application-defined data. + * + *

{@code userContext} is propagated to the response and is not part of {@code request} or + * {@code requestParams}. {@code null} omits it. + */ + public DecisionRequestEvent( + String model, + DecisionRequest request, + @Nullable Map requestParams, + @Nullable Map userContext) { super(EVENT_TYPE); setAttr("model", Objects.requireNonNull(model, "model")); setAttr("request", Objects.requireNonNull(request, "request")); setAttr("request_params", normalizeRequestParams(requestParams)); + UserDefinedContext.set(this, userContext); } @JsonCreator @@ -111,4 +127,15 @@ public DecisionRequest getRequest() { public Map getRequestParams() { return (Map) getAttr("request_params"); } + + /** + * Returns application-defined data to propagate to the matching response, or {@code null} when + * this request omitted it. The map is not added to the decision request or {@code + * request_params}. + */ + @JsonIgnore + @Nullable + public Map getUserContext() { + return UserDefinedContext.get(this); + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/event/DecisionResponseEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/DecisionResponseEvent.java index 9e9f94021..841352f60 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/DecisionResponseEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/DecisionResponseEvent.java @@ -25,6 +25,8 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.decision.DecisionResult; +import javax.annotation.Nullable; + import java.util.List; import java.util.Map; import java.util.Objects; @@ -51,7 +53,8 @@ public class DecisionResponseEvent extends Event { BuiltInAttribute.optionalUntyped(RESPONSE), BuiltInAttribute.optionalUntyped(ERROR), BuiltInAttribute.optional(RETRY_COUNT, Number.class), - BuiltInAttribute.optional(TOTAL_RETRY_WAIT_SEC, Number.class)); + BuiltInAttribute.optional(TOTAL_RETRY_WAIT_SEC, Number.class), + BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class)); private static final ObjectMapper MAPPER = new ObjectMapper(); @@ -206,4 +209,23 @@ public int getRetryCount() { public int getTotalRetryWaitSec() { return ((Number) getAttr(TOTAL_RETRY_WAIT_SEC)).intValue(); } + + /** + * Attaches application-defined data from the originating decision request. {@code null} omits + * it. Present on both successful and failed responses. + */ + public DecisionResponseEvent withUserContext(@Nullable Map userContext) { + UserDefinedContext.set(this, userContext); + return this; + } + + /** + * Returns the originating request's application-defined data, or {@code null} when it was + * omitted. + */ + @JsonIgnore + @Nullable + public Map getUserContext() { + return UserDefinedContext.get(this); + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ToolRequestEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ToolRequestEvent.java index 24abf3ae3..3b4802166 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ToolRequestEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ToolRequestEvent.java @@ -25,6 +25,8 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.chat.messages.ToolCallBlock; +import javax.annotation.Nullable; + import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; @@ -41,14 +43,23 @@ public class ToolRequestEvent extends Event { List.of( BuiltInAttribute.required("model", String.class), BuiltInAttribute.requiredList( - "tool_calls", "a tool call", Map.class, ToolCallBlock.class)); + "tool_calls", "a tool call", Map.class, ToolCallBlock.class), + BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class)); private static final ObjectMapper MAPPER = new ObjectMapper(); - public ToolRequestEvent(String model, List toolCalls) { + public ToolRequestEvent( + String model, + List toolCalls, + @Nullable Map userContext) { super(EVENT_TYPE); setAttr("model", model); setAttr("tool_calls", validate(toolCalls)); + UserDefinedContext.set(this, userContext); + } + + public ToolRequestEvent(String model, List toolCalls) { + this(model, toolCalls, null); } @JsonCreator @@ -81,6 +92,16 @@ public List getToolCalls() { return Collections.unmodifiableList((List) getAttr("tool_calls")); } + /** + * Returns application-defined data to propagate to the matching {@link ToolResponseEvent}, or + * {@code null} when this request omitted it. The map is not added to tool-call arguments. + */ + @JsonIgnore + @Nullable + public Map getUserContext() { + return UserDefinedContext.get(this); + } + private static List restoreToolCalls(List values) { List calls = new ArrayList<>(); for (Object value : values) { diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ToolResponseEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ToolResponseEvent.java index bdd34e2e2..6d6f96664 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ToolResponseEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ToolResponseEvent.java @@ -25,6 +25,8 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.tools.ToolResponse; +import javax.annotation.Nullable; + import java.util.HashMap; import java.util.List; import java.util.Map; @@ -41,7 +43,8 @@ public class ToolResponseEvent extends Event { BuiltInAttribute.required("responses", Map.class), BuiltInAttribute.optional("success", Map.class), BuiltInAttribute.optional("error", Map.class), - BuiltInAttribute.optional("timestamp", Number.class)); + BuiltInAttribute.optional("timestamp", Number.class), + BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class)); private static final ObjectMapper MAPPER = new ObjectMapper(); @@ -145,6 +148,25 @@ public long getTimestamp() { return ((Number) getAttr("timestamp")).longValue(); } + /** + * Attaches application-defined data from the originating tool request. {@code null} omits it. + * Per-call success and error maps are independent of this value. + */ + public ToolResponseEvent withUserContext(@Nullable Map userContext) { + UserDefinedContext.set(this, userContext); + return this; + } + + /** + * Returns the originating request's application-defined data, or {@code null} when it was + * omitted. + */ + @JsonIgnore + @Nullable + public Map getUserContext() { + return UserDefinedContext.get(this); + } + @Override public String toString() { return "ToolResponseEvent{" diff --git a/api/src/main/java/org/apache/flink/agents/api/event/UserDefinedContext.java b/api/src/main/java/org/apache/flink/agents/api/event/UserDefinedContext.java new file mode 100644 index 000000000..31b566323 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/UserDefinedContext.java @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.api.event; + +import org.apache.flink.agents.api.Event; + +import javax.annotation.Nullable; + +import java.lang.reflect.Array; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Optional application-defined data carried on a built-in request and propagated to its response. + * + *

The map is not interpreted and is not injected into prompts, tool arguments, or queries. It is + * omitted when absent so existing events keep their serialized shape. The attribute follows the + * same JSON rules as other event attributes; this class only detaches mutable inputs. + * + *

Maps, lists, and arrays, including nested ones, are copied when attached. A later change to + * the caller's object does not change an event that already holds the value. Other values are + * stored as given. + */ +final class UserDefinedContext { + + static final String ATTRIBUTE = "user_context"; + + private UserDefinedContext() {} + + @SuppressWarnings("unchecked") + static void set(Event event, @Nullable Map userContext) { + if (userContext == null) { + event.getAttributes().remove(ATTRIBUTE); + } else { + event.setAttr(ATTRIBUTE, snapshot(userContext)); + } + } + + /** + * Returns a detached copy of {@code value}. Maps, lists, and arrays are copied recursively so + * the event does not share the caller's mutable structure. + */ + static Object snapshot(Object value) { + if (value instanceof Map) { + Map map = (Map) value; + Map copy = new LinkedHashMap<>(); + for (Map.Entry entry : map.entrySet()) { + copy.put(snapshot(entry.getKey()), snapshot(entry.getValue())); + } + return copy; + } + if (value instanceof List) { + List list = (List) value; + List copy = new ArrayList<>(list.size()); + for (Object item : list) { + copy.add(snapshot(item)); + } + return copy; + } + if (value != null && value.getClass().isArray()) { + int length = Array.getLength(value); + Object copy = Array.newInstance(value.getClass().getComponentType(), length); + for (int i = 0; i < length; i++) { + Array.set(copy, i, snapshot(Array.get(value, i))); + } + return copy; + } + return value; + } + + @Nullable + @SuppressWarnings("unchecked") + static Map get(Event event) { + return (Map) event.getAttr(ATTRIBUTE); + } +} diff --git a/api/src/test/java/org/apache/flink/agents/api/BuiltInEventsTest.java b/api/src/test/java/org/apache/flink/agents/api/BuiltInEventsTest.java index ec7bfe50f..da12e909e 100644 --- a/api/src/test/java/org/apache/flink/agents/api/BuiltInEventsTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/BuiltInEventsTest.java @@ -621,22 +621,28 @@ private static Stream mistypedAttributeCases() { Arguments.of(ChatRequestEvent.EVENT_TYPE, "model", 0), Arguments.of(ChatRequestEvent.EVENT_TYPE, "messages", List.of(0)), Arguments.of(ChatRequestEvent.EVENT_TYPE, "prompt_args", "not-a-map"), + Arguments.of(ChatRequestEvent.EVENT_TYPE, "user_context", "not-a-map"), // ChatResponseEvent Arguments.of(ChatResponseEvent.EVENT_TYPE, "retry_count", "not-a-number"), Arguments.of(ChatResponseEvent.EVENT_TYPE, "total_retry_wait_sec", "not-a-number"), + Arguments.of(ChatResponseEvent.EVENT_TYPE, "user_context", "not-a-map"), // Decision events Arguments.of(DecisionRequestEvent.EVENT_TYPE, "model", 0), + Arguments.of(DecisionRequestEvent.EVENT_TYPE, "user_context", "not-a-map"), Arguments.of(DecisionResponseEvent.EVENT_TYPE, "retry_count", "not-a-number"), Arguments.of( DecisionResponseEvent.EVENT_TYPE, "total_retry_wait_sec", "not-a-number"), + Arguments.of(DecisionResponseEvent.EVENT_TYPE, "user_context", "not-a-map"), // ToolRequestEvent Arguments.of(ToolRequestEvent.EVENT_TYPE, "model", 0), Arguments.of(ToolRequestEvent.EVENT_TYPE, "tool_calls", List.of(0)), + Arguments.of(ToolRequestEvent.EVENT_TYPE, "user_context", "not-a-map"), // ToolResponseEvent Arguments.of(ToolResponseEvent.EVENT_TYPE, "request_id", "not-a-uuid"), Arguments.of(ToolResponseEvent.EVENT_TYPE, "responses", "not-a-map"), Arguments.of(ToolResponseEvent.EVENT_TYPE, "success", "not-a-map"), Arguments.of(ToolResponseEvent.EVENT_TYPE, "error", "not-a-map"), + Arguments.of(ToolResponseEvent.EVENT_TYPE, "user_context", "not-a-map"), Arguments.of(ChatResponseEvent.EVENT_TYPE, "model_routing", "not-a-map"), Arguments.of(ToolResponseEvent.EVENT_TYPE, "timestamp", "not-a-number"), // ContextRetrievalRequestEvent @@ -644,10 +650,12 @@ private static Stream mistypedAttributeCases() { Arguments.of(ContextRetrievalRequestEvent.EVENT_TYPE, "vector_store", 0), Arguments.of( ContextRetrievalRequestEvent.EVENT_TYPE, "max_results", "not-a-number"), + Arguments.of(ContextRetrievalRequestEvent.EVENT_TYPE, "user_context", "not-a-map"), // ContextRetrievalResponseEvent Arguments.of(ContextRetrievalResponseEvent.EVENT_TYPE, "request_id", "not-a-uuid"), Arguments.of(ContextRetrievalResponseEvent.EVENT_TYPE, "query", 0), Arguments.of(ContextRetrievalResponseEvent.EVENT_TYPE, "documents", List.of(0)), + Arguments.of(ContextRetrievalResponseEvent.EVENT_TYPE, "user_context", "not-a-map"), // ModelRoutingEvent (Java-only) Arguments.of(ModelRoutingEvent.EVENT_TYPE, "request_id", "not-a-uuid"), Arguments.of(ModelRoutingEvent.EVENT_TYPE, "router", 0), diff --git a/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java b/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java index 742d610e1..df4236043 100644 --- a/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java @@ -374,6 +374,98 @@ void javaCanDeserializeChatResponseEventFromPythonSnapshot() throws Exception { assertEquals("hi there", response.getText()); } + // ── Request context ──────────────────────────────────────────────────── + + private static Map snapshotContext() { + Map context = new LinkedHashMap<>(); + context.put("task", "summary"); + context.put("item_id", "42"); + context.put("n", 1); + context.put("ok", true); + context.put("tags", List.of("a")); + return context; + } + + @SuppressWarnings("unchecked") + private static void assertSnapshotContext(Object context) { + assertInstanceOf(Map.class, context); + Map values = (Map) context; + assertEquals("summary", values.get("task")); + assertEquals("42", values.get("item_id")); + assertEquals(1, ((Number) values.get("n")).intValue()); + assertEquals(Boolean.TRUE, values.get("ok")); + assertEquals(List.of("a"), values.get("tags")); + } + + private static ChatRequestEvent buildChatRequestEventWithContext() { + ChatRequestEvent event = + new ChatRequestEvent( + "test-model", + List.of(new ChatMessage(MessageRole.USER, "hello world")), + null, + null, + snapshotContext()); + return new ChatRequestEvent(FIXED_EVENT_ID, event.getAttributes()); + } + + @Test + void regenerateChatRequestEventWithContextJavaSnapshot() throws Exception { + assumeTrue(regenerateRequested(), "Set -Dregenerate.snapshots=true to refresh."); + writeJavaSnapshot( + "chat_request_event_with_context.json", buildChatRequestEventWithContext()); + } + + @Test + void chatRequestEventWithContextJavaSnapshotIsStable() throws Exception { + assertJavaSnapshotStable( + "chat_request_event_with_context.json", buildChatRequestEventWithContext()); + } + + @Test + void javaCanDeserializeChatRequestEventWithContextFromPythonSnapshot() throws Exception { + ChatRequestEvent typed = + ChatRequestEvent.fromEvent( + readPythonSnapshot("chat_request_event_with_context.json")); + assertEquals(FIXED_EVENT_ID, typed.getId()); + assertEquals("test-model", typed.getModel()); + assertEquals("hello world", typed.getMessages().get(0).getText()); + assertSnapshotContext(typed.getUserContext()); + } + + private static ChatResponseEvent buildChatResponseEventWithContext() { + ChatResponseEvent event = + ChatResponseEvent.success( + FIXED_REQUEST_ID, + new ChatResult( + ChatMessage.assistant(List.of(new TextBlock("hi there"))))) + .withUserContext(snapshotContext()); + return new ChatResponseEvent(FIXED_EVENT_ID, event.getAttributes()); + } + + @Test + void regenerateChatResponseEventWithContextJavaSnapshot() throws Exception { + assumeTrue(regenerateRequested(), "Set -Dregenerate.snapshots=true to refresh."); + writeJavaSnapshot( + "chat_response_event_with_context.json", buildChatResponseEventWithContext()); + } + + @Test + void chatResponseEventWithContextJavaSnapshotIsStable() throws Exception { + assertJavaSnapshotStable( + "chat_response_event_with_context.json", buildChatResponseEventWithContext()); + } + + @Test + void javaCanDeserializeChatResponseEventWithContextFromPythonSnapshot() throws Exception { + ChatResponseEvent typed = + ChatResponseEvent.fromEvent( + readPythonSnapshot("chat_response_event_with_context.json")); + assertEquals(FIXED_EVENT_ID, typed.getId()); + assertEquals(FIXED_REQUEST_ID, typed.getRequestId()); + assertEquals("hi there", typed.getResponse().getText()); + assertSnapshotContext(typed.getUserContext()); + } + // ── ToolRequestEvent ─────────────────────────────────────────────────── private static ToolRequestEvent buildToolRequestEvent() { diff --git a/api/src/test/java/org/apache/flink/agents/api/event/RequestContextTest.java b/api/src/test/java/org/apache/flink/agents/api/event/RequestContextTest.java new file mode 100644 index 000000000..1f75c72ad --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/event/RequestContextTest.java @@ -0,0 +1,331 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.api.event; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ChatResult; +import org.apache.flink.agents.api.chat.messages.ToolCallBlock; +import org.apache.flink.agents.api.decision.DecisionRequest; +import org.apache.flink.agents.api.decision.DecisionResult; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.vectorstores.Document; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Optional application-defined user context on built-in request and response events. */ +class RequestContextTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private static Map context() { + Map context = new LinkedHashMap<>(); + context.put("task", "summary"); + context.put("item_id", "42"); + context.put("n", 1); + context.put("ok", true); + context.put("tags", List.of("a")); + return context; + } + + @Test + void omittedContextIsAbsentFromTheSerializedEvent() throws Exception { + ChatRequestEvent request = new ChatRequestEvent("model", List.of(ChatMessage.user("hi"))); + assertThat(request.getUserContext()).isNull(); + assertThat(request.getAttributes()).doesNotContainKey("user_context"); + assertThat(MAPPER.writeValueAsString(request)).doesNotContain("\"user_context\""); + + ChatResponseEvent response = + ChatResponseEvent.success( + UUID.randomUUID(), new ChatResult(ChatMessage.assistant("ok"))); + assertThat(response.getUserContext()).isNull(); + assertThat(MAPPER.writeValueAsString(response)).doesNotContain("\"user_context\""); + } + + @Test + void emptyUserContextIsKeptAndNullClearsAResponse() { + Map empty = new HashMap<>(); + ChatRequestEvent request = + new ChatRequestEvent("model", List.of(ChatMessage.user("hi")), null, null, empty); + assertThat(request.getAttributes()).containsKey("user_context"); + assertThat(request.getUserContext()).isEqualTo(empty); + + ChatResponseEvent response = + ChatResponseEvent.success( + UUID.randomUUID(), new ChatResult(ChatMessage.assistant("ok"))) + .withUserContext(empty); + assertThat(response.getUserContext()).isEqualTo(empty); + response.withUserContext(null); + assertThat(response.getAttributes()).doesNotContainKey("user_context"); + assertThat(response.getUserContext()).isNull(); + } + + @Test + void eachRequestAndResponseRoundTripsItsOwnContext() throws Exception { + Map context = context(); + UUID requestId = UUID.randomUUID(); + DecisionRequest decisionRequest = + new DecisionRequest( + "evidence", + Map.of("urgent", DecisionRequest.Question.predicate("Needs attention?"))); + DecisionResult decisionResult = + new DecisionResult( + Map.of("urgent", new DecisionResult.PredicateAnswer(0.8)), + null, + null, + Map.of()); + + List events = + List.of( + new ChatRequestEvent( + "model", + List.of(ChatMessage.user("hi")), + Map.of("tone", "brief"), + null, + context), + new ChatRequestEvent( + "model", + List.of(ChatMessage.user("hi")), + null, + null, + Map.of("label", "summary-42")), + ChatResponseEvent.success( + requestId, new ChatResult(ChatMessage.assistant("ok"))) + .withUserContext(context), + ChatResponseEvent.failed(requestId, "TimeoutException: timed out", 1, 2) + .withUserContext(context), + new ToolRequestEvent( + "model", + List.of(new ToolCallBlock("call-1", "lookup", Map.of("x", 1))), + context), + new ToolResponseEvent( + requestId, + Map.of("call-1", ToolResponse.text("ok")), + Map.of("call-1", true), + Map.of()) + .withUserContext(context), + new ContextRetrievalRequestEvent("what is flink", "store", 2, context), + new ContextRetrievalResponseEvent( + requestId, + "what is flink", + List.of(new Document("body", Map.of(), "d1"))) + .withUserContext(context), + new DecisionRequestEvent("decision", decisionRequest, Map.of(), context), + DecisionResponseEvent.success(requestId, decisionResult, 0, 0) + .withUserContext(context), + DecisionResponseEvent.failed(requestId, "provider down", 2, 3) + .withUserContext(context)); + + for (Event event : events) { + Event restored = Event.fromJson(MAPPER.writeValueAsString(event)); + assertThat(restored).isInstanceOf(event.getClass()); + assertThat(restored.getAttr("user_context")).isEqualTo(event.getAttr("user_context")); + assertThat(restored.getId()).isEqualTo(event.getId()); + } + } + + @Test + void foreignJsonContextSurvivesTheCrossLanguageBoundary() throws Exception { + String requestJson = + "{" + + "\"id\":\"00000000-0000-0000-0000-000000000001\"," + + "\"type\":\"_chat_request_event\"," + + "\"attributes\":{" + + "\"model\":\"test-model\"," + + "\"messages\":[{\"role\":\"user\",\"blocks\":[{\"type\":\"text\",\"text\":\"hi\"}],\"metadata\":{}}]," + + "\"prompt_args\":{}," + + "\"output_schema\":null," + + "\"user_context\":{\"task\":\"summary\",\"item_id\":\"42\",\"n\":1,\"ok\":true,\"tags\":[\"a\"]}" + + "}," + + "\"attachments\":{\"trace\":\"keep\"}" + + "}"; + ChatRequestEvent request = ChatRequestEvent.fromEvent(Event.fromJson(requestJson)); + assertThat(request.getModel()).isEqualTo("test-model"); + assertThat(request.getMessages().get(0).getText()).isEqualTo("hi"); + assertThat(request.getUserContext()).isEqualTo(context()); + assertThat(request.getAttachment("trace")).isEqualTo("keep"); + + String failedJson = + "{" + + "\"id\":\"00000000-0000-0000-0000-000000000001\"," + + "\"type\":\"_chat_response_event\"," + + "\"attributes\":{" + + "\"request_id\":\"00000000-0000-0000-0000-000000000002\"," + + "\"status\":\"FAILED\"," + + "\"response\":null," + + "\"error\":\"TimeoutError: timed out\"," + + "\"retry_count\":2," + + "\"total_retry_wait_sec\":3," + + "\"user_context\":{\"task\":\"summary\",\"item_id\":\"42\"}" + + "}," + + "\"attachments\":{}" + + "}"; + ChatResponseEvent failed = ChatResponseEvent.fromEvent(Event.fromJson(failedJson)); + assertThat(failed.isFailed()).isTrue(); + assertThat(failed.getError()).isEqualTo("TimeoutError: timed out"); + assertThat(failed.getUserContext()).isEqualTo(Map.of("task", "summary", "item_id", "42")); + } + + @Test + void unknownAttributeIsStillRejectedWhenContextIsPresent() { + Map attributes = new HashMap<>(); + attributes.put("model", "m"); + attributes.put("messages", List.of(ChatMessage.user("hi"))); + attributes.put("user_context", Map.of("task", "summary")); + attributes.put("extra", 1); + Event event = new Event(UUID.randomUUID(), ChatRequestEvent.EVENT_TYPE, attributes); + + assertThatThrownBy(() -> ChatRequestEvent.fromEvent(event)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("extra"); + } + + @Test + void reusedContextMapIsCopiedIncludingNestedMapsAndLists() { + Map shared = new LinkedHashMap<>(); + shared.put("task", "summary"); + List tags = new ArrayList<>(); + tags.add("keep"); + shared.put("tags", tags); + Map nested = new LinkedHashMap<>(); + nested.put("label", "keep"); + shared.put("meta", nested); + + ChatRequestEvent requestA = + new ChatRequestEvent("model", List.of(ChatMessage.user("a")), null, null, shared); + shared.put("task", "risk"); + tags.add("changed"); + nested.put("label", "changed"); + ChatRequestEvent requestB = + new ChatRequestEvent("model", List.of(ChatMessage.user("b")), null, null, shared); + + assertThat(requestA.getUserContext()) + .isEqualTo( + Map.of( + "task", + "summary", + "tags", + List.of("keep"), + "meta", + Map.of("label", "keep"))); + assertThat(requestB.getUserContext()) + .isEqualTo( + Map.of( + "task", + "risk", + "tags", + List.of("keep", "changed"), + "meta", + Map.of("label", "changed"))); + Map storedA = requestA.getUserContext(); + assertThat(storedA).isNotSameAs(shared); + assertThat(storedA.get("tags")).isNotSameAs(tags); + assertThat(storedA.get("meta")).isNotSameAs(nested); + } + + @Test + void userContextJsonObjectRoundTripsAsAMap() throws Exception { + String requestJson = + "{" + + "\"id\":\"00000000-0000-0000-0000-000000000001\"," + + "\"type\":\"_chat_request_event\"," + + "\"attributes\":{" + + "\"model\":\"model\"," + + "\"messages\":[{\"role\":\"user\",\"blocks\":[{\"type\":\"text\",\"text\":\"hi\"}],\"metadata\":{}}]," + + "\"user_context\":{\"type\":\"Row\",\"values\":[1],\"label\":\"keep\"}" + + "}}"; + ChatRequestEvent request = ChatRequestEvent.fromEvent(Event.fromJson(requestJson)); + assertThat(request.getUserContext()) + .isEqualTo(Map.of("type", "Row", "values", List.of(1), "label", "keep")); + + String responseJson = + "{" + + "\"id\":\"00000000-0000-0000-0000-000000000001\"," + + "\"type\":\"_chat_response_event\"," + + "\"attributes\":{" + + "\"request_id\":\"00000000-0000-0000-0000-000000000002\"," + + "\"status\":\"SUCCESS\"," + + "\"response\":{\"message\":{\"role\":\"assistant\",\"blocks\":[{\"type\":\"text\",\"text\":\"ok\"}],\"metadata\":{}}}," + + "\"error\":null," + + "\"retry_count\":0," + + "\"total_retry_wait_sec\":0," + + "\"user_context\":{\"type\":\"Row\",\"values\":[1],\"label\":\"keep\"}" + + "}}"; + ChatResponseEvent response = ChatResponseEvent.fromEvent(Event.fromJson(responseJson)); + assertThat(response.getUserContext()) + .isEqualTo(Map.of("type", "Row", "values", List.of(1), "label", "keep")); + } + + @Test + void nonMapUserContextIsRejected() { + Map attributes = new HashMap<>(); + attributes.put("model", "m"); + attributes.put("messages", List.of(ChatMessage.user("hi"))); + attributes.put("user_context", "summary-42"); + Event event = new Event(UUID.randomUUID(), ChatRequestEvent.EVENT_TYPE, attributes); + + assertThatThrownBy(() -> ChatRequestEvent.fromEvent(event)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("user_context") + .hasMessageContaining("Map"); + } + + @Test + void nestedArrayIsCopiedSoLaterMutationDoesNotChangeAnEarlierRequest() { + String[] labels = new String[] {"summary"}; + int[] counts = new int[] {1}; + List nested = new ArrayList<>(); + String[] inner = new String[] {"keep"}; + nested.add(inner); + Map shared = new LinkedHashMap<>(); + shared.put("labels", labels); + shared.put("counts", counts); + shared.put("nested", nested); + + ChatRequestEvent requestA = + new ChatRequestEvent("model", List.of(ChatMessage.user("a")), null, null, shared); + labels[0] = "risk"; + counts[0] = 9; + inner[0] = "changed"; + ChatRequestEvent requestB = + new ChatRequestEvent("model", List.of(ChatMessage.user("b")), null, null, shared); + + Map storedA = requestA.getUserContext(); + assertThat((String[]) storedA.get("labels")).containsExactly("summary"); + assertThat((int[]) storedA.get("counts")).containsExactly(1); + assertThat((String[]) ((List) storedA.get("nested")).get(0)).containsExactly("keep"); + assertThat(storedA.get("labels")).isNotSameAs(labels); + assertThat(storedA.get("counts")).isNotSameAs(counts); + + Map storedB = requestB.getUserContext(); + assertThat((String[]) storedB.get("labels")).containsExactly("risk"); + assertThat((int[]) storedB.get("counts")).containsExactly(9); + assertThat((String[]) ((List) storedB.get("nested")).get(0)).containsExactly("changed"); + } +} diff --git a/docs/content/docs/development/calling-llms/chat.md b/docs/content/docs/development/calling-llms/chat.md index 19beaf489..d44328424 100644 --- a/docs/content/docs/development/calling-llms/chat.md +++ b/docs/content/docs/development/calling-llms/chat.md @@ -107,6 +107,8 @@ InputEvent → request Action → ChatRequestEvent A `ChatRequestEvent` can name a model router instead of a chat model. The built-in action selects the model first and emits a `ModelRoutingEvent`; see [Model Routing]({{< ref "docs/development/calling-llms/model_routing" >}}). +Pass an optional `user_context` map to the `ChatRequestEvent` constructor (Python `user_context`, Java `userContext`) to carry application-defined data from a request to its corresponding response. Telling responses apart is one use, including when requests come from one Action or from several. The map is propagated to the corresponding `ChatResponseEvent`, including after tool-call rounds, retries, model fallback, and a failed response. It is not added to the prompt. Values follow the same rules as other Event attributes; see [Serialization and Typed Values]({{< ref "docs/development/create-agent/events#json-serialization" >}}) and [User context]({{< ref "docs/development/create-agent/events#user-context" >}}). + ### Example {#usage-example} Start with the Agent and model declarations in the diff --git a/docs/content/docs/development/create-agent/events.md b/docs/content/docs/development/create-agent/events.md index 22a58ef79..5351eee4d 100644 --- a/docs/content/docs/development/create-agent/events.md +++ b/docs/content/docs/development/create-agent/events.md @@ -278,3 +278,58 @@ Event that your Actions can handle to continue the workflow: | Call tools | `ToolRequestEvent`, `ToolResponseEvent` | [Tool Use]({{< ref "docs/development/calling-llms/tool_use#built-in-events-and-actions" >}}) | | Retrieve context | `ContextRetrievalRequestEvent`, `ContextRetrievalResponseEvent` | [Vector Stores]({{< ref "docs/development/vector-search/vector_stores#built-in-events-and-actions" >}}) | | Call a decision model | `DecisionRequestEvent`, `DecisionResponseEvent` | [Decision Models]({{< ref "docs/development/decision_models#event-driven-calls" >}}) | + +### User context + +Chat, tool, context-retrieval, and decision requests can carry an optional map of +application-defined data from a request to its corresponding response. Telling +responses apart is one use. Requests may come from one Action or from several. The +built-in Action propagates that map onto the matching response and does not add it to +the prompt, tool arguments, or vector-store query. Omitting it leaves the response +without a `user_context` attribute. + +Pass the map to the request constructor. The value is copied when the event is created, +including nested maps, lists, and arrays, so a later change to the object you passed in +does not change the event. `user_context` follows the same rules as other Event +attributes; see [Serialization and Typed Values](#json-serialization). JSON transfer or +replay may change runtime types, and callers restore application types when needed. +Copying the mutable input and restoring types are separate: construction copies the +structure and does not normalize values. + +A chat response keeps the user context from the original `ChatRequestEvent` after +tool-call rounds, retries, and model fallback. Tool responses keep it from the +`ToolRequestEvent` that asked for those calls, including when a call in the batch +fails. Decision responses keep it on both success and failure. Context retrieval has +no failure event; a query error propagates and the request still holds its user context. + +{{< tabs "User Context" >}} +{{< tab "Python" >}} +```python +ctx.send_event(ChatRequestEvent( + model="review_model", + messages=[ChatMessage.user(text)], + user_context={"task": "summary", "item_id": "42"}, +)) +``` + +```python +response = ChatResponseEvent.from_event(event) +task = response.user_context["task"] +``` +{{< /tab >}} +{{< tab "Java" >}} +```java +ctx.sendEvent(new ChatRequestEvent( + "reviewModel", + List.of(ChatMessage.user(text)), + null, + null, + Map.of("task", "summary", "item_id", "42"))); +``` + +```java +ChatResponseEvent response = ChatResponseEvent.fromEvent(event); +String task = String.valueOf(response.getUserContext().get("task")); +``` +{{< /tab >}} +{{< /tabs >}} diff --git a/e2e-test/cross-language-event-snapshots/java/chat_request_event_with_context.json b/e2e-test/cross-language-event-snapshots/java/chat_request_event_with_context.json new file mode 100644 index 000000000..213f796a2 --- /dev/null +++ b/e2e-test/cross-language-event-snapshots/java/chat_request_event_with_context.json @@ -0,0 +1,24 @@ +{ + "id" : "00000000-0000-0000-0000-000000000001", + "attributes" : { + "user_context" : { + "task" : "summary", + "item_id" : "42", + "n" : 1, + "ok" : true, + "tags" : [ "a" ] + }, + "messages" : [ { + "role" : "user", + "blocks" : [ { + "text" : "hello world", + "type" : "text" + } ], + "metadata" : { } + } ], + "model" : "test-model", + "prompt_args" : { } + }, + "attachments" : { }, + "type" : "_chat_request_event" +} diff --git a/e2e-test/cross-language-event-snapshots/java/chat_response_event_with_context.json b/e2e-test/cross-language-event-snapshots/java/chat_response_event_with_context.json new file mode 100644 index 000000000..a78a53e7a --- /dev/null +++ b/e2e-test/cross-language-event-snapshots/java/chat_response_event_with_context.json @@ -0,0 +1,34 @@ +{ + "id" : "00000000-0000-0000-0000-000000000001", + "attributes" : { + "response" : { + "message" : { + "role" : "assistant", + "blocks" : [ { + "text" : "hi there", + "type" : "text" + } ], + "metadata" : { } + }, + "model" : null, + "response_id" : null, + "usage" : null, + "finish_reason" : null, + "metadata" : { } + }, + "total_retry_wait_sec" : 0, + "retry_count" : 0, + "user_context" : { + "task" : "summary", + "item_id" : "42", + "n" : 1, + "ok" : true, + "tags" : [ "a" ] + }, + "error" : null, + "request_id" : "00000000-0000-0000-0000-000000000002", + "status" : "SUCCESS" + }, + "attachments" : { }, + "type" : "_chat_response_event" +} diff --git a/e2e-test/cross-language-event-snapshots/python/chat_request_event_with_context.json b/e2e-test/cross-language-event-snapshots/python/chat_request_event_with_context.json new file mode 100644 index 000000000..92971860f --- /dev/null +++ b/e2e-test/cross-language-event-snapshots/python/chat_request_event_with_context.json @@ -0,0 +1,31 @@ +{ + "id": "00000000-0000-0000-0000-000000000001", + "type": "_chat_request_event", + "attributes": { + "model": "test-model", + "messages": [ + { + "role": "user", + "blocks": [ + { + "type": "text", + "text": "hello world" + } + ], + "metadata": {} + } + ], + "prompt_args": {}, + "output_schema": null, + "user_context": { + "task": "summary", + "item_id": "42", + "n": 1, + "ok": true, + "tags": [ + "a" + ] + } + }, + "attachments": {} +} diff --git a/e2e-test/cross-language-event-snapshots/python/chat_response_event_with_context.json b/e2e-test/cross-language-event-snapshots/python/chat_response_event_with_context.json new file mode 100644 index 000000000..a24f30232 --- /dev/null +++ b/e2e-test/cross-language-event-snapshots/python/chat_response_event_with_context.json @@ -0,0 +1,38 @@ +{ + "id": "00000000-0000-0000-0000-000000000001", + "type": "_chat_response_event", + "attributes": { + "request_id": "00000000-0000-0000-0000-000000000002", + "status": "SUCCESS", + "response": { + "message": { + "role": "assistant", + "blocks": [ + { + "type": "text", + "text": "hi there" + } + ], + "metadata": {} + }, + "model": null, + "response_id": null, + "usage": null, + "finish_reason": null, + "metadata": {} + }, + "error": null, + "retry_count": 0, + "total_retry_wait_sec": 0, + "user_context": { + "task": "summary", + "item_id": "42", + "n": 1, + "ok": true, + "tags": [ + "a" + ] + } + }, + "attachments": {} +} diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java index 9f3cd78f4..37d44b9a2 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java @@ -98,6 +98,9 @@ public class ChatModelAction { private static final String TOOL_REQUEST_EVENT_CONTEXT = "_TOOL_REQUEST_EVENT_CONTEXT"; private static final String INITIAL_REQUEST_ID = "initialRequestId"; private static final String MODEL = "model"; + /** Sensory-memory key for the chat request's user context. Matches the event attribute. */ + private static final String USER_CONTEXT = "user_context"; + private static final String ROUTING_METADATA_CONTEXT = "_ROUTING_METADATA_CONTEXT"; private static final String OUTPUT_SCHEMA = "outputSchema"; private static final String PROMPT_ARGS = "prompt_args"; @@ -153,7 +156,8 @@ private static void saveToolRequestEventContext( UUID initialRequestId, String model, Map promptArgs, - Object outputSchema) + Object outputSchema, + @Nullable Map userContext) throws Exception { Map toolRequestEventContext; if (sensoryMem.isExist(TOOL_REQUEST_EVENT_CONTEXT)) { @@ -169,6 +173,9 @@ private static void saveToolRequestEventContext( if (outputSchema != null) { context.put(OUTPUT_SCHEMA, outputSchema); } + if (userContext != null) { + context.put(USER_CONTEXT, userContext); + } toolRequestEventContext.put(toolRequestEventId, context); sensoryMem.set(TOOL_REQUEST_EVENT_CONTEXT, toolRequestEventContext); } @@ -238,6 +245,7 @@ private static void handleToolCalls( List messages, Map promptArgs, Object outputSchema, + @Nullable Map userContext, RunnerContext ctx) throws Exception { updateToolCallContext( @@ -255,7 +263,8 @@ private static void handleToolCalls( initialRequestId, model, promptArgs, - outputSchema); + outputSchema, + userContext); ctx.sendEvent(toolRequestEvent); } @@ -342,6 +351,7 @@ public static void chat( messages, promptArgs, outputSchema, + null, ctx); } @@ -351,6 +361,7 @@ private static void chat( List messages, Map promptArgs, @Nullable Object outputSchema, + @Nullable Map userContext, RunnerContext ctx) throws Exception { int numRetries = ChatModelInvoker.configuredRetries(ctx); @@ -420,6 +431,7 @@ private static void chat( messages, promptArgs, outputSchema, + userContext, ctx); } else { if (routingMetadata == null) { @@ -433,10 +445,11 @@ private static void chat( clearRequestContext(ctx.getSensoryMemory(), initialRequestId); ChatResponseEvent completed = ChatResponseEvent.success( - initialRequestId, - result.response, - totalRetryCount, - totalRetryWaitSec); + initialRequestId, + result.response, + totalRetryCount, + totalRetryWaitSec) + .withUserContext(userContext); completed.setStructuredOutput(result.structuredOutput); if (routingMetadata != null) completed.setRoutingMetadata(routingMetadata); ctx.sendEvent(completed); @@ -472,20 +485,25 @@ private static void chat( // The reasoning loop is over; a routed loop that dies mid-way must not leak its // parked metadata before completing the failed logical request. takeRoutingMetadata(ctx.getSensoryMemory(), initialRequestId); - sendFailure(initialRequestId, Objects.requireNonNull(lastError), ctx); + sendFailure(initialRequestId, Objects.requireNonNull(lastError), userContext, ctx); } - private static void sendFailure(UUID requestId, Exception error, RunnerContext ctx) + private static void sendFailure( + UUID requestId, + Exception error, + @Nullable Map userContext, + RunnerContext ctx) throws Exception { LOG.warn("Chat request {} failed", requestId, error); Map stats = getRetryStats(ctx.getSensoryMemory(), requestId); clearRequestContext(ctx.getSensoryMemory(), requestId); ctx.sendEvent( ChatResponseEvent.failed( - requestId, - ChatModelInvoker.errorText(error), - stats.get(TOTAL_RETRY_COUNT).intValue(), - stats.get(TOTAL_RETRY_WAIT_SEC).intValue())); + requestId, + ChatModelInvoker.errorText(error), + stats.get(TOTAL_RETRY_COUNT).intValue(), + stats.get(TOTAL_RETRY_WAIT_SEC).intValue()) + .withUserContext(userContext)); } @SuppressWarnings("unchecked") @@ -618,7 +636,7 @@ private static void processChatRequest(ChatRequestEvent event, RunnerContext ctx event.getPromptArgs(), ctx); } catch (ModelRoutingResolver.RoutingFailure e) { - sendFailure(event.getId(), e, ctx); + sendFailure(event.getId(), e, event.getUserContext(), ctx); return; } chat( @@ -627,6 +645,7 @@ private static void processChatRequest(ChatRequestEvent event, RunnerContext ctx event.getMessages(), event.getPromptArgs(), event.getOutputSchema(), + event.getUserContext(), ctx); } @@ -643,6 +662,7 @@ private static void processToolResponse(ToolResponseEvent event, RunnerContext c Map promptArgs = (Map) context.getOrDefault(PROMPT_ARGS, Map.of()); Object outputSchema = context.get(OUTPUT_SCHEMA); + Map userContext = (Map) context.get(USER_CONTEXT); Map responses = event.getResponses(); Map success = event.getSuccess(); @@ -670,6 +690,7 @@ private static void processToolResponse(ToolResponseEvent event, RunnerContext c messages, promptArgs, outputSchema, + userContext, ctx); } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java index 461e3b8a0..54a35d7ff 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java @@ -97,9 +97,10 @@ public VectorStoreQueryResult call() throws Exception { ctx.sendEvent( new ContextRetrievalResponseEvent( - contextRetrievalRequestEvent.getId(), - contextRetrievalRequestEvent.getQuery(), - result.getDocuments())); + contextRetrievalRequestEvent.getId(), + contextRetrievalRequestEvent.getQuery(), + result.getDocuments()) + .withUserContext(contextRetrievalRequestEvent.getUserContext())); } } } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelAction.java index a9be4b80f..9fccb0725 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelAction.java @@ -66,15 +66,17 @@ public static void processDecisionRequest(Event event, RunnerContext ctx) throws DecisionResponseEvent terminal = result.isSuccess() ? DecisionResponseEvent.success( - request.getId(), - result.response, - result.retryCount, - result.totalWaitSec) + request.getId(), + result.response, + result.retryCount, + result.totalWaitSec) + .withUserContext(request.getUserContext()) : DecisionResponseEvent.failed( - request.getId(), - result.error, - result.retryCount, - result.totalWaitSec); + request.getId(), + result.error, + result.retryCount, + result.totalWaitSec) + .withUserContext(request.getUserContext()); ctx.sendEvent(terminal); } } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java index 3d33def7c..820a696bb 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java @@ -101,7 +101,9 @@ public static void processToolRequest(Event event, RunnerContext ctx) throws Exc executeSequentially(executions, toolCallAsync, ctx, success, error, responses); } - ctx.sendEvent(new ToolResponseEvent(toolRequest.getId(), responses, success, error)); + ctx.sendEvent( + new ToolResponseEvent(toolRequest.getId(), responses, success, error) + .withUserContext(toolRequest.getUserContext())); } @SuppressWarnings("unchecked") diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java index 8e4697361..6fea68348 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java @@ -17,6 +17,7 @@ */ package org.apache.flink.agents.plan.actions; +import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.agents.Agent; import org.apache.flink.agents.api.agents.AgentExecutionOptions; @@ -31,6 +32,7 @@ import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ChatRequestEvent; import org.apache.flink.agents.api.event.ChatResponseEvent; import org.apache.flink.agents.api.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; @@ -1279,6 +1281,236 @@ void truncatedConversionIsRejectedBeforeParsing() throws Exception { Map.of("promptTokens", 9L, "completionTokens", 0L)))); } + @Test + void omittedContextStaysOffTheResponse() throws Exception { + configureRetryStrategy(0, 0); + when(mockChatModel.chat(any(), any(), any())).thenReturn(ChatFixtures.response("hello")); + + ChatRequestEvent request = + new ChatRequestEvent( + "test-model", + List.of(ChatMessage.user("hi")), + Map.of("tone", "brief"), + null); + ChatModelAction.processChatRequestOrToolResponse(request, mockCtx); + + ChatResponseEvent response = ChatResponseEvent.fromEvent(sentEvents.get(0)); + assertThat(response.getUserContext()).isNull(); + assertThat(response.getAttributes()).doesNotContainKey("user_context"); + assertThat(response.getRequestId()).isEqualTo(request.getId()); + assertThat(response.getAttachments()).isEmpty(); + } + + @Test + @SuppressWarnings("unchecked") + void successFailureAndRetryEchoContextWithoutTouchingThePrompt() throws Exception { + Map context = new LinkedHashMap<>(); + context.put("task", "summary"); + context.put("item_id", "42"); + context.put("n", 1); + context.put("ok", true); + context.put("tags", List.of("a")); + Map promptArgs = Map.of("tone", "brief"); + List messages = List.of(ChatMessage.user("hi")); + + configureRetryStrategy(1, 0); + when(mockChatModel.chat(any(), any(), any())) + .thenThrow(new IllegalStateException("transient")) + .thenReturn(ChatFixtures.response("done")); + ChatRequestEvent request = + new ChatRequestEvent("test-model", messages, promptArgs, null, context); + request.setAttachment("trace", "keep-on-request"); + ChatModelAction.processChatRequestOrToolResponse(request, mockCtx); + + ChatResponseEvent response = ChatResponseEvent.fromEvent(sentEvents.get(0)); + assertThat(response.isSuccess()).isTrue(); + assertThat(response.getUserContext()).isEqualTo(context); + assertThat(response.getRetryCount()).isEqualTo(1); + assertThat(response.getAttachments()).isEmpty(); + assertThat(response.getRequestId()).isEqualTo(request.getId()); + + ArgumentCaptor> messagesCaptor = ArgumentCaptor.forClass(List.class); + ArgumentCaptor> promptCaptor = ArgumentCaptor.forClass(Map.class); + verify(mockChatModel, times(2)) + .chat(messagesCaptor.capture(), promptCaptor.capture(), any()); + assertThat(promptCaptor.getAllValues()).containsOnly(promptArgs); + assertThat(messagesCaptor.getAllValues()) + .allSatisfy(sent -> assertThat(sent).containsExactlyElementsOf(messages)); + + sentEvents.clear(); + when(mockChatModel.chat(any(), any(), any())) + .thenThrow(new IllegalStateException("persistent")); + ChatRequestEvent failing = + new ChatRequestEvent("test-model", messages, promptArgs, null, context); + ChatModelAction.processChatRequestOrToolResponse(failing, mockCtx); + ChatResponseEvent failed = ChatResponseEvent.fromEvent(sentEvents.get(0)); + assertThat(failed.isFailed()).isTrue(); + assertThat(failed.getError()).contains("persistent"); + assertThat(failed.getUserContext()).isEqualTo(context); + assertThat(failed.getRequestId()).isEqualTo(failing.getId()); + } + + @Test + void toolRoundsKeepTheInitialContextAndDoNotCopyItOntoToolRequests() throws Exception { + configureRetryStrategy(0, 0); + Map contextA = Map.of("task", "summary", "item_id", "42"); + Map contextB = Map.of("task", "risk_assessment", "item_id", "42"); + Map toolCall = + Map.of( + "id", + "call-1", + "type", + "function", + "function", + Map.of("name", "lookup", "arguments", Map.of("x", 1))); + when(mockChatModel.chat(any(), any(), any())) + .thenReturn(ChatFixtures.response("", List.of(toolCall))) + .thenReturn(ChatFixtures.response("", List.of(toolCall))) + .thenReturn(ChatFixtures.response("risk done")) + .thenReturn(ChatFixtures.response("summary done")); + + ChatRequestEvent requestA = + new ChatRequestEvent( + "test-model", + List.of(ChatMessage.user("a")), + Map.of("who", "a"), + null, + contextA); + ChatRequestEvent requestB = + new ChatRequestEvent( + "test-model", + List.of(ChatMessage.user("b")), + Map.of("who", "b"), + null, + contextB); + ChatModelAction.processChatRequestOrToolResponse(requestA, mockCtx); + ChatModelAction.processChatRequestOrToolResponse(requestB, mockCtx); + + ToolRequestEvent toolA = ToolRequestEvent.fromEvent(sentEvents.get(0)); + ToolRequestEvent toolB = ToolRequestEvent.fromEvent(sentEvents.get(1)); + assertThat(toolA.getUserContext()).isNull(); + assertThat(toolB.getUserContext()).isNull(); + assertThat(toolA.getToolCalls().get(0).getInput()).containsExactly(Map.entry("x", 1)); + assertThat(toolA.getAttachments()).isEmpty(); + + @SuppressWarnings("unchecked") + Map> parked = + (Map>) + sensoryMemory.get("_TOOL_REQUEST_EVENT_CONTEXT").getValue(); + Object checkpointed = + new ObjectMapper() + .readValue( + new ObjectMapper() + .writeValueAsBytes( + parked.get(toolB.getId()).get("user_context")), + Object.class); + parked.get(toolB.getId()).put("user_context", checkpointed); + + // Finish B before A. Each final response keeps its own request context. + ChatModelAction.processChatRequestOrToolResponse( + new ToolResponseEvent( + toolB.getId(), + Map.of("call-1", ToolResponse.text("b")), + Map.of("call-1", true), + Map.of()), + mockCtx); + ChatModelAction.processChatRequestOrToolResponse( + new ToolResponseEvent( + toolA.getId(), + Map.of("call-1", ToolResponse.text("a")), + Map.of("call-1", true), + Map.of()), + mockCtx); + + ChatResponseEvent responseB = ChatResponseEvent.fromEvent(sentEvents.get(2)); + ChatResponseEvent responseA = ChatResponseEvent.fromEvent(sentEvents.get(3)); + assertThat(responseB.getRequestId()).isEqualTo(requestB.getId()); + assertThat(responseB.getUserContext()).isEqualTo(contextB); + assertThat(responseB.getResponse().getText()).isEqualTo("risk done"); + assertThat(responseA.getRequestId()).isEqualTo(requestA.getId()); + assertThat(responseA.getUserContext()).isEqualTo(contextA); + assertThat(responseA.getResponse().getText()).isEqualTo("summary done"); + assertThat(responseA.getAttachments()).isEmpty(); + } + + @Test + void restoredRequestStillEchoesContext() throws Exception { + configureRetryStrategy(0, 0); + when(mockChatModel.chat(any(), any(), any())).thenReturn(ChatFixtures.response("ok")); + Map context = Map.of("task", "summary", "item_id", "42"); + ChatRequestEvent original = + new ChatRequestEvent( + "test-model", List.of(ChatMessage.user("hi")), null, null, context); + ChatRequestEvent restored = + ChatRequestEvent.fromEvent( + Event.fromJson(new ObjectMapper().writeValueAsString(original))); + + ChatModelAction.processChatRequestOrToolResponse(restored, mockCtx); + + ChatResponseEvent response = ChatResponseEvent.fromEvent(sentEvents.get(0)); + assertThat(response.getRequestId()).isEqualTo(original.getId()); + assertThat(response.getUserContext()).isEqualTo(context); + } + + @Test + @SuppressWarnings("unchecked") + void reusedContextMapKeepsEachResponseSnapshot() throws Exception { + configureRetryStrategy(0, 0); + when(mockChatModel.chat(any(), any(), any())) + .thenReturn(ChatFixtures.response("summary done")) + .thenReturn(ChatFixtures.response("risk done")); + + Map shared = new LinkedHashMap<>(); + shared.put("task", "summary"); + List tags = new ArrayList<>(); + tags.add("keep"); + shared.put("tags", tags); + Map nested = new LinkedHashMap<>(); + nested.put("label", "keep"); + shared.put("meta", nested); + + ChatRequestEvent requestA = + new ChatRequestEvent( + "test-model", List.of(ChatMessage.user("a")), null, null, shared); + shared.put("task", "risk"); + tags.add("changed"); + nested.put("label", "changed"); + ChatRequestEvent requestB = + new ChatRequestEvent( + "test-model", List.of(ChatMessage.user("b")), null, null, shared); + + ChatModelAction.processChatRequestOrToolResponse(requestA, mockCtx); + ChatModelAction.processChatRequestOrToolResponse(requestB, mockCtx); + + ChatResponseEvent responseA = ChatResponseEvent.fromEvent(sentEvents.get(0)); + ChatResponseEvent responseB = ChatResponseEvent.fromEvent(sentEvents.get(1)); + assertThat(responseA.getRequestId()).isEqualTo(requestA.getId()); + assertThat(responseA.getResponse().getText()).isEqualTo("summary done"); + assertThat(responseA.getUserContext()) + .isEqualTo( + Map.of( + "task", + "summary", + "tags", + List.of("keep"), + "meta", + Map.of("label", "keep"))); + assertThat(responseB.getRequestId()).isEqualTo(requestB.getId()); + assertThat(responseB.getResponse().getText()).isEqualTo("risk done"); + assertThat(responseB.getUserContext()) + .isEqualTo( + Map.of( + "task", + "risk", + "tags", + List.of("keep", "changed"), + "meta", + Map.of("label", "changed"))); + Map storedA = (Map) responseA.getUserContext(); + assertThat(storedA).isNotSameAs(shared); + assertThat(storedA.get("tags")).isNotSameAs(tags); + } + // --- Helper methods --- private static void assertInvocationIdentities(ExecutionReporter reporter, int expectedCalls) { diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java index ae59ea870..4857018b0 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java @@ -1135,6 +1135,94 @@ void fallbackExhaustedReturnsFailure() throws Exception { assertThat(ctx.chatResponse().isFailed()).isTrue(); } + @Test + void fallbackSuccessAndExhaustionRetainRequestContext() throws Exception { + Map context = Map.of("task", "summary", "item_id", "42"); + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .defaultModel("small") + .fallback(true) + .build(), + null); + FakeRunnerContext succeeded = + new FakeRunnerContext(router) + .register("big", new FakeChatModel(new RuntimeException("big is down"))) + .register( + "small", new FakeChatModel(ChatFixtures.response("ok from small"))); + ChatRequestEvent request = + new ChatRequestEvent( + "router", List.of(ChatMessage.user("write sql")), null, null, context); + request.setAttachment("trace", "stay-on-request"); + ChatModelAction.processChatRequestOrToolResponse(request, succeeded); + assertThat(succeeded.chatResponse().isSuccess()).isTrue(); + assertThat(succeeded.chatResponse().getUserContext()).isEqualTo(context); + assertThat(succeeded.chatResponse().getAttachments()).isEmpty(); + assertThat(succeeded.resolvedChatModels).containsExactly("big", "small"); + + FakeRunnerContext exhausted = + new FakeRunnerContext(router) + .register("big", new FakeChatModel(new RuntimeException("big-exploded"))) + .register( + "small", new FakeChatModel(new RuntimeException("small-exploded"))); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(ChatMessage.user("write sql")), null, null, context), + exhausted); + assertThat(exhausted.chatResponse().isFailed()).isTrue(); + assertThat(exhausted.chatResponse().getUserContext()).isEqualTo(context); + } + + @Test + void fallbackThenToolRoundRetainsInitialContext() throws Exception { + Map context = Map.of("task", "summary", "item_id", "42"); + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .defaultModel("small") + .fallback(true) + .build(), + null); + List> toolCall = + List.of( + Map.of( + "id", + "call-1", + "type", + "function", + "function", + Map.of("name", "lookup", "arguments", Map.of("q", "sql")))); + FakeRunnerContext ctx = + new FakeRunnerContext(router) + .register("big", new FakeChatModel(new RuntimeException("big is down"))) + .register( + "small", + new FakeChatModel( + ChatFixtures.response("", toolCall), + ChatFixtures.response("final from small"))); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(ChatMessage.user("write sql")), null, null, context), + ctx); + ToolRequestEvent toolRequest = ctx.toolRequestEvent(); + assertThat(toolRequest.getUserContext()).isNull(); + assertThat(toolRequest.getToolCalls().get(0).getInput()) + .containsExactly(Map.entry("q", "sql")); + + ChatModelAction.processChatRequestOrToolResponse( + new ToolResponseEvent( + toolRequest.getId(), + Map.of("call-1", ToolResponse.text("42")), + Map.of("call-1", true), + Map.of()), + ctx); + assertThat(ctx.chatResponse().getResponse().getText()).isEqualTo("final from small"); + assertThat(ctx.chatResponse().getUserContext()).isEqualTo(context); + assertThat(ctx.resolvedChatModels).containsExactly("big", "small", "small"); + } + @Test void unresolvableCandidateCountsAsFailedAttemptAndFallsBack() throws Exception { // "big" is selected by the rule but its resource lookup fails (typo'd candidate); the diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ContextRetrievalActionTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ContextRetrievalActionTest.java new file mode 100644 index 000000000..d12837808 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ContextRetrievalActionTest.java @@ -0,0 +1,130 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.plan.actions; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.configuration.ReadableConfiguration; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ContextRetrievalRequestEvent; +import org.apache.flink.agents.api.event.ContextRetrievalResponseEvent; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.vectorstores.BaseVectorStore; +import org.apache.flink.agents.api.vectorstores.Document; +import org.apache.flink.agents.api.vectorstores.VectorStoreQuery; +import org.apache.flink.agents.api.vectorstores.VectorStoreQueryResult; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Context retrieval copies optional request context onto the response only. */ +class ContextRetrievalActionTest { + + @Test + void successEchoesContextAndLeavesTheQueryUnchanged() throws Exception { + Harness harness = new Harness(); + when(harness.store.query(any())) + .thenReturn( + new VectorStoreQueryResult(List.of(new Document("body", Map.of(), "d1")))); + Map context = Map.of("task", "retrieve", "item_id", "42"); + ContextRetrievalRequestEvent request = + new ContextRetrievalRequestEvent("what is flink", "store", 2, context); + request.setAttachment("trace", "keep-on-request"); + + ContextRetrievalAction.processContextRetrievalRequest(request, harness.ctx); + + ContextRetrievalResponseEvent response = + ContextRetrievalResponseEvent.fromEvent(harness.sent.get(0)); + assertThat(response.getRequestId()).isEqualTo(request.getId()); + assertThat(response.getQuery()).isEqualTo("what is flink"); + assertThat(response.getDocuments()).hasSize(1); + assertThat(response.getUserContext()).isEqualTo(context); + assertThat(response.getAttachments()).isEmpty(); + verify(harness.store) + .query( + org.mockito.ArgumentMatchers.argThat( + (VectorStoreQuery query) -> + "what is flink".equals(query.getQueryText()) + && Integer.valueOf(2).equals(query.getLimit()))); + } + + @Test + void omittedContextStaysAbsent() throws Exception { + Harness harness = new Harness(); + when(harness.store.query(any())).thenReturn(new VectorStoreQueryResult(List.of())); + + ContextRetrievalAction.processContextRetrievalRequest( + new ContextRetrievalRequestEvent("q", "store"), harness.ctx); + + ContextRetrievalResponseEvent response = + ContextRetrievalResponseEvent.fromEvent(harness.sent.get(0)); + assertThat(response.getUserContext()).isNull(); + assertThat(response.getAttributes()).doesNotContainKey("user_context"); + } + + @Test + void queryFailurePropagatesWithoutAResponseAndKeepsTheRequestContext() throws Exception { + Harness harness = new Harness(); + when(harness.store.query(any())).thenThrow(new IllegalStateException("store down")); + Map context = Map.of("task", "retrieve"); + ContextRetrievalRequestEvent request = + new ContextRetrievalRequestEvent("q", "store", 3, context); + + assertThatThrownBy( + () -> + ContextRetrievalAction.processContextRetrievalRequest( + request, harness.ctx)) + .isInstanceOf(IllegalStateException.class) + .hasMessage("store down"); + assertThat(harness.sent).isEmpty(); + assertThat(request.getUserContext()).isEqualTo(context); + } + + private static final class Harness { + private final RunnerContext ctx = mock(RunnerContext.class); + private final BaseVectorStore store = mock(BaseVectorStore.class); + private final List sent = new ArrayList<>(); + + private Harness() throws Exception { + ReadableConfiguration config = mock(ReadableConfiguration.class); + when(config.get(AgentExecutionOptions.RAG_ASYNC)).thenReturn(false); + when(ctx.getConfig()).thenReturn(config); + when(ctx.getResource(eq("store"), eq(ResourceType.VECTOR_STORE))).thenReturn(store); + when(ctx.durableExecute(any())) + .thenAnswer( + invocation -> + invocation + .>getArgument(0) + .call()); + doAnswer(invocation -> sent.add(invocation.getArgument(0))).when(ctx).sendEvent(any()); + } + } +} diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/DecisionModelActionTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/DecisionModelActionTest.java index c5595d4f4..f5d4f7ddb 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/DecisionModelActionTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/DecisionModelActionTest.java @@ -47,6 +47,8 @@ import org.mockito.MockedStatic; import pemja.core.object.PyObject; +import javax.annotation.Nullable; + import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -66,6 +68,35 @@ void clearInterrupt() { Thread.interrupted(); } + @Test + void successAndFailureEchoContextAndLeaveTheProviderRequestUnchanged() throws Exception { + Map context = Map.of("task", "review", "item_id", "42"); + Harness ok = new Harness(false, 0, context); + ok.request.setAttachment("trace", "keep-on-request"); + when(ok.model.decide(any(), anyMap())).thenReturn(result("provider-model", null)); + + ok.run(); + + assertThat(ok.response().isSuccess()).isTrue(); + assertThat(ok.response().getUserContext()).isEqualTo(context); + assertThat(ok.response().getAttachments()).isEmpty(); + verify(ok.model).decide(ok.request.getRequest(), ok.request.getRequestParams()); + + Harness failed = new Harness(false, 0, context); + when(failed.model.decide(any(), anyMap())) + .thenThrow(new IllegalStateException("provider down")); + failed.run(); + assertThat(failed.response().isFailed()).isTrue(); + assertThat(failed.response().getError()).contains("provider down"); + assertThat(failed.response().getUserContext()).isEqualTo(context); + + Harness omitted = new Harness(false, 0); + when(omitted.model.decide(any(), anyMap())).thenReturn(result(null, null)); + omitted.run(); + assertThat(omitted.response().getUserContext()).isNull(); + assertThat(omitted.response().getAttributes()).doesNotContainKey("user_context"); + } + @ParameterizedTest @ValueSource(booleans = {false, true}) void succeedsThroughConfiguredDurablePath(boolean async) throws Exception { @@ -507,17 +538,24 @@ private static final class Harness { private final ObjectMapper mapper = new ObjectMapper(); private final List stored = new ArrayList<>(); private int slot; - private final DecisionRequestEvent request = - new DecisionRequestEvent( - MODEL, - new DecisionRequest( - "The delivery arrived damaged", - Map.of( - "damaged", - DecisionRequest.Question.predicate("Is it damaged?"))), - Map.of("safety_identifier", "review-user", "timeout", 4.0)); + private final DecisionRequestEvent request; private Harness(boolean async, int retries) throws Exception { + this(async, retries, null); + } + + private Harness(boolean async, int retries, @Nullable Map userContext) + throws Exception { + request = + new DecisionRequestEvent( + MODEL, + new DecisionRequest( + "The delivery arrived damaged", + Map.of( + "damaged", + DecisionRequest.Question.predicate("Is it damaged?"))), + Map.of("safety_identifier", "review-user", "timeout", 4.0), + userContext); ReadableConfiguration config = mock(ReadableConfiguration.class); when(config.get(AgentExecutionOptions.DECISION_ASYNC)).thenReturn(async); when(config.get(AgentExecutionOptions.MAX_RETRIES)).thenReturn(retries); diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java index 7591f49b8..51fe0bd6b 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java @@ -428,6 +428,47 @@ void processToolRequestReportsUninitializedMemory() throws Exception { "Cannot inject tool parameter from sensory_memory because memory is not initialized."); } + @Test + void processToolRequestEchoesContextOnSuccessAndFailureWithoutInjectingIt() throws Exception { + Map context = Map.of("task", "lookup", "item_id", "42"); + FakeRunnerContext success = new FakeRunnerContext(); + ToolRequestEvent request = toolRequest("queryOrder", context); + request.setAttachment("trace", "keep-on-request"); + + ToolCallAction.processToolRequest(request, success); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(success.sentEvents.get(0)); + assertThat(response.getUserContext()).isEqualTo(context); + assertThat(response.getRequestId()).isEqualTo(request.getId()); + assertThat(response.getResponses().get("call-1").getText()).isEqualTo("tenant-1:order-1"); + assertThat(response.getAttachments()).isEmpty(); + assertThat(response.getSuccess()).containsEntry("call-1", true); + + FakeRunnerContext failed = + new FakeRunnerContext() { + @Override + public Resource getResource(String name, ResourceType type) { + return new FailingTool(); + } + }; + ToolRequestEvent failing = toolRequest("queryOrder", context); + ToolCallAction.processToolRequest(failing, failed); + ToolResponseEvent failure = ToolResponseEvent.fromEvent(failed.sentEvents.get(0)); + assertThat(failure.getSuccess()).containsEntry("call-1", false); + assertThat(failure.getError()).containsEntry("call-1", "database timeout"); + assertThat(failure.getUserContext()).isEqualTo(context); + assertThat(failure.getAttachments()).isEmpty(); + } + + @Test + void processToolRequestOmitsContextWhenTheRequestDoes() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(); + ToolCallAction.processToolRequest(toolRequest("queryOrder"), ctx); + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getUserContext()).isNull(); + assertThat(response.getAttributes()).doesNotContainKey("user_context"); + } + @Test void processToolRequestKeepsDiagnosticErrorWhenToolExecutionFails() throws Exception { FakeRunnerContext ctx = @@ -751,6 +792,11 @@ private static ToolRequestEvent toolRequest(String toolName) { return ChatFixtures.request("model", List.of(toolCall(toolName, "call-1", "order-1"))); } + private static ToolRequestEvent toolRequest(String toolName, Map userContext) { + ToolRequestEvent base = toolRequest(toolName); + return new ToolRequestEvent(base.getModel(), base.getToolCalls(), userContext); + } + private static ToolRequestEvent toolRequest(String toolName, String... ids) { List> toolCalls = new ArrayList<>(); for (int i = 0; i < ids.length; i++) { diff --git a/python/flink_agents/api/events/built_in_events.py b/python/flink_agents/api/events/built_in_events.py index a618ffa7d..c39135eb3 100644 --- a/python/flink_agents/api/events/built_in_events.py +++ b/python/flink_agents/api/events/built_in_events.py @@ -92,7 +92,6 @@ LongTermSearchEvent.EVENT_TYPE: MemoryEvent, } - def restore(event: Event) -> Event: """Restore a registered built-in event; unknown types stay generic. diff --git a/python/flink_agents/api/events/chat_event.py b/python/flink_agents/api/events/chat_event.py index 5d17dfde4..4eb426d6f 100644 --- a/python/flink_agents/api/events/chat_event.py +++ b/python/flink_agents/api/events/chat_event.py @@ -26,7 +26,7 @@ from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage from flink_agents.api.chat_result import ChatResult -from flink_agents.api.events.event import BuiltInAttribute, Event +from flink_agents.api.events.event import BuiltInAttribute, Event, copy_user_context class ChatRequestEvent(Event): @@ -43,6 +43,9 @@ class ChatRequestEvent(Event): resource is configured on the chat model setup. Empty by default. output_schema: OutputSchema | None The expected output schema of the chat model final response. Optional. + user_context : dict[str, Any] | None + Application-defined data propagated to the matching response. Not added + to messages or prompt arguments. Omitted when unset. """ EVENT_TYPE: ClassVar[str] = "_chat_request_event" @@ -54,6 +57,7 @@ class ChatRequestEvent(Event): ), BuiltInAttribute.optional("prompt_args", dict), BuiltInAttribute.optional_untyped("output_schema"), + BuiltInAttribute.optional("user_context", dict), ) def __init__( @@ -62,16 +66,20 @@ def __init__( messages: List[ChatMessage], prompt_args: Dict[str, Any] | None = None, output_schema: OutputSchema | None = None, + user_context: dict[str, Any] | None = None, ) -> None: """Create a ChatRequestEvent.""" + attributes: Dict[str, Any] = { + "model": model, + "messages": messages, + "prompt_args": prompt_args if prompt_args is not None else {}, + "output_schema": output_schema, + } + if user_context is not None: + attributes["user_context"] = copy_user_context(user_context) super().__init__( type=ChatRequestEvent.EVENT_TYPE, - attributes={ - "model": model, - "messages": messages, - "prompt_args": prompt_args if prompt_args is not None else {}, - "output_schema": output_schema, - }, + attributes=attributes, ) @classmethod @@ -93,6 +101,7 @@ def from_event(cls, event: Event) -> "ChatRequestEvent": messages=messages, prompt_args=attributes.get("prompt_args"), output_schema=output_schema_raw, + user_context=attributes.get("user_context"), ) return result.reconstruct_from(event) @@ -117,6 +126,11 @@ def output_schema(self) -> OutputSchema | None: """Return the expected output schema, if any.""" return self.get_attr("output_schema") + @property + def user_context(self) -> dict[str, Any] | None: + """Return application-defined data to propagate to the response, if set.""" + return self.get_attr("user_context") + class ChatResponseError(RuntimeError): """A failed chat result, retaining the request ID and textual error.""" @@ -140,6 +154,8 @@ class ChatResponseEvent(Event): The total number of retries across all tool call rounds. total_retry_wait_sec : int The total time spent waiting during retries in seconds. + user_context : dict[str, Any] | None + Application-defined data from the originating request. Omitted when unset. """ EVENT_TYPE: ClassVar[str] = "_chat_response_event" @@ -164,6 +180,7 @@ class ChatResponseEvent(Event): BuiltInAttribute.optional(_TOTAL_RETRY_WAIT_SEC, int), BuiltInAttribute.optional_untyped(_STRUCTURED_OUTPUT), BuiltInAttribute.optional(_MODEL_ROUTING, dict), + BuiltInAttribute.optional("user_context", dict), ) def __init__( @@ -174,6 +191,7 @@ def __init__( error: str | None = None, retry_count: int = 0, total_retry_wait_sec: int = 0, + user_context: dict[str, Any] | None = None, ) -> None: """Create a ChatResponseEvent.""" if not ( @@ -195,16 +213,19 @@ def __init__( if not isinstance(request_id, UUID): msg = "request_id must be a UUID" raise TypeError(msg) + attributes: Dict[str, Any] = { + self._REQUEST_ID: request_id, + self._STATUS: status, + self._RESPONSE: response, + self._ERROR: error, + self._RETRY_COUNT: retry_count, + self._TOTAL_RETRY_WAIT_SEC: total_retry_wait_sec, + } + if user_context is not None: + attributes["user_context"] = copy_user_context(user_context) super().__init__( type=ChatResponseEvent.EVENT_TYPE, - attributes={ - self._REQUEST_ID: request_id, - self._STATUS: status, - self._RESPONSE: response, - self._ERROR: error, - self._RETRY_COUNT: retry_count, - self._TOTAL_RETRY_WAIT_SEC: total_retry_wait_sec, - }, + attributes=attributes, ) @classmethod @@ -214,6 +235,7 @@ def success( response: ChatResult, retry_count: int = 0, total_retry_wait_sec: int = 0, + user_context: dict[str, Any] | None = None, ) -> "ChatResponseEvent": """Create a successful terminal response.""" return cls( @@ -222,6 +244,7 @@ def success( response=response, retry_count=retry_count, total_retry_wait_sec=total_retry_wait_sec, + user_context=user_context, ) @classmethod @@ -231,6 +254,7 @@ def failed( error: str, retry_count: int = 0, total_retry_wait_sec: int = 0, + user_context: dict[str, Any] | None = None, ) -> "ChatResponseEvent": """Create a failed terminal response.""" return cls( @@ -239,6 +263,7 @@ def failed( error=error, retry_count=retry_count, total_retry_wait_sec=total_retry_wait_sec, + user_context=user_context, ) @classmethod @@ -261,6 +286,7 @@ def from_event(cls, event: Event) -> "ChatResponseEvent": error=attributes.get(cls._ERROR), retry_count=attributes.get(cls._RETRY_COUNT, 0), total_retry_wait_sec=attributes.get(cls._TOTAL_RETRY_WAIT_SEC, 0), + user_context=attributes.get("user_context"), ) result.attributes[cls._STRUCTURED_OUTPUT] = event.attributes.get( cls._STRUCTURED_OUTPUT @@ -319,3 +345,8 @@ def retry_count(self) -> int: def total_retry_wait_sec(self) -> int: """Return the total retry wait time in seconds.""" return self.get_attr(self._TOTAL_RETRY_WAIT_SEC) + + @property + def user_context(self) -> dict[str, Any] | None: + """Return the originating request's application-defined data, if set.""" + return self.get_attr("user_context") diff --git a/python/flink_agents/api/events/context_retrieval_event.py b/python/flink_agents/api/events/context_retrieval_event.py index 63f7e8d23..e4af90fd3 100644 --- a/python/flink_agents/api/events/context_retrieval_event.py +++ b/python/flink_agents/api/events/context_retrieval_event.py @@ -15,7 +15,7 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# -from typing import ClassVar, List, Tuple +from typing import Any, ClassVar, List, Tuple try: from typing import override @@ -23,7 +23,7 @@ from typing_extensions import override from uuid import UUID -from flink_agents.api.events.event import BuiltInAttribute, Event +from flink_agents.api.events.event import BuiltInAttribute, Event, copy_user_context from flink_agents.api.vector_stores.vector_store import Document @@ -38,6 +38,9 @@ class ContextRetrievalRequestEvent(Event): Name of the vector store setup resource to use max_results : int Maximum number of results to return (default: 3) + user_context : dict[str, Any] | None + Application-defined data propagated to the matching response. Not added + to the vector-store query. Omitted when unset. """ EVENT_TYPE: ClassVar[str] = "_context_retrieval_request_event" @@ -46,9 +49,16 @@ class ContextRetrievalRequestEvent(Event): BuiltInAttribute.required("query", str), BuiltInAttribute.required("vector_store", str), BuiltInAttribute.optional("max_results", int), + BuiltInAttribute.optional("user_context", dict), ) - def __init__(self, query: str, vector_store: str, max_results: int = 3) -> None: + def __init__( + self, + query: str, + vector_store: str, + max_results: int = 3, + user_context: dict[str, Any] | None = None, + ) -> None: """Create a ContextRetrievalRequestEvent. Args: @@ -56,6 +66,8 @@ def __init__(self, query: str, vector_store: str, max_results: int = 3) -> None: query: The search query text to find relevant context for. vector_store: Name of the vector store setup resource to use. max_results: Maximum number of results to return (default: 3). Must be positive. + user_context: Application-defined data propagated to the response. + Omitted when unset. Raises: ------ @@ -64,13 +76,16 @@ def __init__(self, query: str, vector_store: str, max_results: int = 3) -> None: if max_results <= 0: msg = f"`max_results` must be positive, but was {max_results}." raise ValueError(msg) + attributes: dict[str, Any] = { + "query": query, + "vector_store": vector_store, + "max_results": max_results, + } + if user_context is not None: + attributes["user_context"] = copy_user_context(user_context) super().__init__( type=ContextRetrievalRequestEvent.EVENT_TYPE, - attributes={ - "query": query, - "vector_store": vector_store, - "max_results": max_results, - }, + attributes=attributes, ) @classmethod @@ -84,6 +99,7 @@ def from_event(cls, event: Event) -> "ContextRetrievalRequestEvent": query=attributes["query"], vector_store=attributes["vector_store"], max_results=attributes.get("max_results", 3), + user_context=attributes.get("user_context"), ) return result.reconstruct_from(event) @@ -102,6 +118,11 @@ def max_results(self) -> int: """Return the maximum number of results.""" return self.get_attr("max_results") + @property + def user_context(self) -> dict[str, Any] | None: + """Return application-defined data to propagate to the response, if set.""" + return self.get_attr("user_context") + class ContextRetrievalResponseEvent(Event): """Event representing retrieved context results. @@ -114,6 +135,8 @@ class ContextRetrievalResponseEvent(Event): The original search query from the request documents : List[Document] List of retrieved documents from the vector store + user_context : dict[str, Any] | None + Application-defined data from the originating request. Omitted when unset. """ EVENT_TYPE: ClassVar[str] = "_context_retrieval_response_event" @@ -124,17 +147,27 @@ class ContextRetrievalResponseEvent(Event): BuiltInAttribute.required_list( "documents", "a Document or its serialized dict", (Document, dict) ), + BuiltInAttribute.optional("user_context", dict), ) - def __init__(self, request_id: UUID, query: str, documents: List[Document]) -> None: + def __init__( + self, + request_id: UUID, + query: str, + documents: List[Document], + user_context: dict[str, Any] | None = None, + ) -> None: """Create a ContextRetrievalResponseEvent.""" + attributes: dict[str, Any] = { + "request_id": request_id, + "query": query, + "documents": documents, + } + if user_context is not None: + attributes["user_context"] = copy_user_context(user_context) super().__init__( type=ContextRetrievalResponseEvent.EVENT_TYPE, - attributes={ - "request_id": request_id, - "query": query, - "documents": documents, - }, + attributes=attributes, ) @classmethod @@ -152,6 +185,7 @@ def from_event(cls, event: Event) -> "ContextRetrievalResponseEvent": request_id=attributes["request_id"], query=attributes["query"], documents=documents, + user_context=attributes.get("user_context"), ) return result.reconstruct_from(event) @@ -170,3 +204,8 @@ def query(self) -> str: def documents(self) -> List[Document]: """Return the retrieved documents.""" return self.get_attr("documents") + + @property + def user_context(self) -> dict[str, Any] | None: + """Return the originating request's application-defined data, if set.""" + return self.get_attr("user_context") diff --git a/python/flink_agents/api/events/decision_event.py b/python/flink_agents/api/events/decision_event.py index 1ff40e5cc..04f9af852 100644 --- a/python/flink_agents/api/events/decision_event.py +++ b/python/flink_agents/api/events/decision_event.py @@ -19,7 +19,7 @@ import json from collections.abc import Mapping -from typing import ClassVar, Tuple +from typing import Any, ClassVar, Tuple from uuid import UUID from pydantic import JsonValue @@ -27,7 +27,7 @@ from flink_agents.api.decision_models import DecisionRequest, DecisionResult from flink_agents.api.decision_models.decision_model import _validate_request_params -from flink_agents.api.events.event import BuiltInAttribute, Event +from flink_agents.api.events.event import BuiltInAttribute, Event, copy_user_context class DecisionRequestEvent(Event): @@ -36,6 +36,8 @@ class DecisionRequestEvent(Event): ``model`` is the registered resource name, not a provider model identifier. ``request`` supplies the shared input and named questions for this call. ``request_params`` supplies provider-specific parameters for this call. + ``user_context`` is application-defined data propagated to the response. It + is not part of ``request`` or ``request_params``, and is omitted when unset. """ EVENT_TYPE: ClassVar[str] = "_decision_request_event" @@ -43,6 +45,7 @@ class DecisionRequestEvent(Event): BuiltInAttribute.required("model", str), BuiltInAttribute.required_untyped("request"), BuiltInAttribute.optional_untyped("request_params"), + BuiltInAttribute.optional("user_context", dict), ) def __init__( @@ -51,18 +54,22 @@ def __init__( request: DecisionRequest, *, request_params: Mapping[str, JsonValue] | None = None, + user_context: dict[str, Any] | None = None, ) -> None: """Create a decision request event.""" if not isinstance(model, str) or not isinstance(request, DecisionRequest): msg = "Decision request requires a model resource name and DecisionRequest." raise TypeError(msg) + attributes: dict[str, Any] = { + "model": model, + "request": request, + "request_params": _validate_request_params(request_params), + } + if user_context is not None: + attributes["user_context"] = copy_user_context(user_context) super().__init__( type=self.EVENT_TYPE, - attributes={ - "model": model, - "request": request, - "request_params": _validate_request_params(request_params), - }, + attributes=attributes, ) @classmethod @@ -76,6 +83,7 @@ def from_event(cls, event: Event) -> "DecisionRequestEvent": model=attributes["model"], request=DecisionRequest.model_validate(attributes["request"]), request_params=attributes.get("request_params"), + user_context=attributes.get("user_context"), ) return result.reconstruct_from(event) @@ -94,6 +102,11 @@ def request_params(self) -> dict[str, JsonValue]: """Return provider-specific parameters for this request.""" return self.get_attr("request_params") + @property + def user_context(self) -> dict[str, Any] | None: + """Return application-defined data to propagate to the response, if set.""" + return self.get_attr("user_context") + class DecisionResponseError(RuntimeError): """A failed decision result, retaining the request ID and textual error.""" @@ -117,6 +130,8 @@ class DecisionResponseEvent(Event): The total number of retries for this request. total_retry_wait_sec : int The total time spent waiting during retries in seconds. + user_context : dict[str, Any] | None + Application-defined data from the originating request. Omitted when unset. """ EVENT_TYPE: ClassVar[str] = "_decision_response_event" @@ -137,6 +152,7 @@ class DecisionResponseEvent(Event): BuiltInAttribute.optional_untyped(_ERROR), BuiltInAttribute.optional(_RETRY_COUNT, int), BuiltInAttribute.optional(_TOTAL_RETRY_WAIT_SEC, int), + BuiltInAttribute.optional("user_context", dict), ) def __init__( @@ -147,6 +163,7 @@ def __init__( error: str | None = None, retry_count: int = 0, total_retry_wait_sec: int = 0, + user_context: dict[str, Any] | None = None, ) -> None: """Create a DecisionResponseEvent.""" if not ( @@ -170,16 +187,19 @@ def __init__( if not isinstance(request_id, UUID): msg = "request_id must be a UUID" raise TypeError(msg) + attributes: dict[str, Any] = { + self._REQUEST_ID: request_id, + self._STATUS: status, + self._RESPONSE: response, + self._ERROR: error, + self._RETRY_COUNT: retry_count, + self._TOTAL_RETRY_WAIT_SEC: total_retry_wait_sec, + } + if user_context is not None: + attributes["user_context"] = copy_user_context(user_context) super().__init__( type=DecisionResponseEvent.EVENT_TYPE, - attributes={ - self._REQUEST_ID: request_id, - self._STATUS: status, - self._RESPONSE: response, - self._ERROR: error, - self._RETRY_COUNT: retry_count, - self._TOTAL_RETRY_WAIT_SEC: total_retry_wait_sec, - }, + attributes=attributes, ) @classmethod @@ -189,6 +209,7 @@ def success( response: DecisionResult, retry_count: int = 0, total_retry_wait_sec: int = 0, + user_context: dict[str, Any] | None = None, ) -> "DecisionResponseEvent": """Create a successful terminal response.""" return cls( @@ -197,6 +218,7 @@ def success( response=response, retry_count=retry_count, total_retry_wait_sec=total_retry_wait_sec, + user_context=user_context, ) @classmethod @@ -206,6 +228,7 @@ def failed( error: str, retry_count: int = 0, total_retry_wait_sec: int = 0, + user_context: dict[str, Any] | None = None, ) -> "DecisionResponseEvent": """Create a failed terminal response.""" return cls( @@ -214,6 +237,7 @@ def failed( error=error, retry_count=retry_count, total_retry_wait_sec=total_retry_wait_sec, + user_context=user_context, ) @classmethod @@ -238,6 +262,7 @@ def from_event(cls, event: Event) -> "DecisionResponseEvent": error=attributes.get(cls._ERROR), retry_count=attributes.get(cls._RETRY_COUNT, 0), total_retry_wait_sec=attributes.get(cls._TOTAL_RETRY_WAIT_SEC, 0), + user_context=attributes.get("user_context"), ) return result.reconstruct_from(event) @@ -286,3 +311,8 @@ def retry_count(self) -> int: def total_retry_wait_sec(self) -> int: """Return the total retry wait time in seconds.""" return self.get_attr(self._TOTAL_RETRY_WAIT_SEC) + + @property + def user_context(self) -> dict[str, Any] | None: + """Return the originating request's application-defined data, if set.""" + return self.get_attr("user_context") diff --git a/python/flink_agents/api/events/event.py b/python/flink_agents/api/events/event.py index d54c295be..fba6c110f 100644 --- a/python/flink_agents/api/events/event.py +++ b/python/flink_agents/api/events/event.py @@ -59,6 +59,34 @@ def _reconstruct_row_if_needed(data: Any) -> Any: return data +def copy_user_context(user_context: dict[str, Any]) -> dict[str, Any]: + """Detach a caller-supplied ``user_context`` map from the event that stores it. + + Built-in request and response events store this copy so a later mutation of + the caller's object does not change an event that already holds the value. + Nested maps and lists are copied. Other values are returned as given. + Copying does not restore or rewrite application types. + + Raises: + TypeError: If ``user_context`` is not a dict. + """ + if not isinstance(user_context, dict): + msg = "user_context must be a dict" + raise TypeError(msg) + return { + key: _copy_user_context_value(value) for key, value in user_context.items() + } + + +def _copy_user_context_value(value: Any) -> Any: + """Copy nested maps and lists inside a user-context value.""" + if isinstance(value, dict): + return {key: _copy_user_context_value(item) for key, item in value.items()} + if isinstance(value, list): + return [_copy_user_context_value(item) for item in value] + return value + + class _AttrKind(Enum): """The shape a built-in attribute value takes at the JSON boundary.""" @@ -466,12 +494,12 @@ def from_json(cls, json_str: str) -> "Event": msg = "Event JSON must contain a non-empty 'type' field." raise ValueError(msg) event = cls.model_validate(data) - for key in list(event.attributes): - event.attributes[key] = _reconstruct_row_if_needed(event.attributes[key]) # Imported lazily: built_in_events imports the concrete subclasses, which # import this module, so a top-level import here would be circular. from flink_agents.api.events.built_in_events import restore + for key in list(event.attributes): + event.attributes[key] = _reconstruct_row_if_needed(event.attributes[key]) return restore(event) diff --git a/python/flink_agents/api/events/tool_event.py b/python/flink_agents/api/events/tool_event.py index ffbd8e7d1..8c28fc405 100644 --- a/python/flink_agents/api/events/tool_event.py +++ b/python/flink_agents/api/events/tool_event.py @@ -25,7 +25,7 @@ from uuid import UUID from flink_agents.api.chat_message import ToolCallBlock -from flink_agents.api.events.event import BuiltInAttribute, Event +from flink_agents.api.events.event import BuiltInAttribute, Event, copy_user_context from flink_agents.api.tools.tool_response import ToolResponse @@ -38,6 +38,9 @@ class ToolRequestEvent(Event): name of the model that generated the tool request. tool_calls : List[ToolCallBlock] tool calls that should be executed in batch. + user_context : dict[str, Any] | None + Application-defined data propagated to the matching response. Not added + to tool-call arguments. Omitted when unset. """ EVENT_TYPE: ClassVar[str] = "_tool_request_event" @@ -47,20 +50,29 @@ class ToolRequestEvent(Event): BuiltInAttribute.required_list( "tool_calls", "a tool call", (dict, ToolCallBlock) ), + BuiltInAttribute.optional("user_context", dict), ) - def __init__(self, model: str, tool_calls: List[ToolCallBlock]) -> None: + def __init__( + self, + model: str, + tool_calls: List[ToolCallBlock], + user_context: dict[str, Any] | None = None, + ) -> None: """Create a ToolRequestEvent.""" tool_calls = [ToolCallBlock.model_validate(c) for c in tool_calls] if len({c.call_id for c in tool_calls}) != len(tool_calls): msg = "Duplicate tool call ID in one request" raise ValueError(msg) + attributes: Dict[str, Any] = { + "model": model, + "tool_calls": tool_calls, + } + if user_context is not None: + attributes["user_context"] = copy_user_context(user_context) super().__init__( type=ToolRequestEvent.EVENT_TYPE, - attributes={ - "model": model, - "tool_calls": tool_calls, - }, + attributes=attributes, ) @classmethod @@ -72,6 +84,7 @@ def from_event(cls, event: Event) -> "ToolRequestEvent": result = ToolRequestEvent( model=event.attributes["model"], tool_calls=event.attributes["tool_calls"], + user_context=event.attributes.get("user_context"), ) return result.reconstruct_from(event) @@ -85,6 +98,11 @@ def tool_calls(self) -> List[ToolCallBlock]: """Return the list of tool calls.""" return self.get_attr("tool_calls") + @property + def user_context(self) -> dict[str, Any] | None: + """Return application-defined data to propagate to the response, if set.""" + return self.get_attr("user_context") + class ToolResponseEvent(Event): """Event representing a result from tool call. @@ -95,6 +113,8 @@ class ToolResponseEvent(Event): The id of the request event. responses : Dict[str, Any] The dict maps tool call id to result. + user_context : dict[str, Any] | None + Application-defined data from the originating request. Omitted when unset. """ EVENT_TYPE: ClassVar[str] = "_tool_response_event" @@ -108,6 +128,7 @@ class ToolResponseEvent(Event): # Java-produced events restore, then drop it (Python does not carry it) # to stay consistent with the Python snapshot shape. BuiltInAttribute.optional("timestamp", int), + BuiltInAttribute.optional("user_context", dict), ) def __init__( @@ -116,6 +137,7 @@ def __init__( responses: Dict[str, Any], success: Dict[str, bool] | None = None, error: Dict[str, str] | None = None, + user_context: dict[str, Any] | None = None, ) -> None: """Create a ToolResponseEvent.""" responses = { @@ -134,16 +156,17 @@ def __init__( } if success is None: success = {k: v.is_success() for k, v in responses.items()} + attributes: Dict[str, Any] = { + "request_id": request_id, + "responses": responses, + "success": success, + "error": error if error is not None else {}, + } + if user_context is not None: + attributes["user_context"] = copy_user_context(user_context) super().__init__( type=ToolResponseEvent.EVENT_TYPE, - attributes={ - "request_id": request_id, - "responses": responses, - "success": success - if success is not None - else dict.fromkeys(responses, True), - "error": error if error is not None else {}, - }, + attributes=attributes, ) @classmethod @@ -158,6 +181,7 @@ def from_event(cls, event: Event) -> "ToolResponseEvent": responses=responses, success=event.attributes.get("success", dict.fromkeys(responses, True)), error=event.attributes.get("error", {}), + user_context=event.attributes.get("user_context"), ) return result.reconstruct_from(event) @@ -181,3 +205,8 @@ def success(self) -> Dict[str, bool]: def error(self) -> Dict[str, str]: """Return diagnostic errors for failed tool calls.""" return self.get_attr("error") + + @property + def user_context(self) -> dict[str, Any] | None: + """Return the originating request's application-defined data, if set.""" + return self.get_attr("user_context") diff --git a/python/flink_agents/api/tests/test_built_in_events.py b/python/flink_agents/api/tests/test_built_in_events.py index c0b1994aa..717dab0c7 100644 --- a/python/flink_agents/api/tests/test_built_in_events.py +++ b/python/flink_agents/api/tests/test_built_in_events.py @@ -546,31 +546,39 @@ def _mistyped_attribute_cases() -> list[tuple[str, str, object]]: (ChatRequestEvent.EVENT_TYPE, "model", 0), (ChatRequestEvent.EVENT_TYPE, "messages", [0]), (ChatRequestEvent.EVENT_TYPE, "prompt_args", "not-a-dict"), + (ChatRequestEvent.EVENT_TYPE, "user_context", "not-a-dict"), # ChatResponseEvent (ChatResponseEvent.EVENT_TYPE, "retry_count", "not-a-number"), (ChatResponseEvent.EVENT_TYPE, "total_retry_wait_sec", "not-a-number"), + (ChatResponseEvent.EVENT_TYPE, "user_context", "not-a-dict"), # Decision events (DecisionRequestEvent.EVENT_TYPE, "model", 0), + (DecisionRequestEvent.EVENT_TYPE, "user_context", "not-a-dict"), (DecisionResponseEvent.EVENT_TYPE, "retry_count", "not-a-number"), (DecisionResponseEvent.EVENT_TYPE, "total_retry_wait_sec", "not-a-number"), + (DecisionResponseEvent.EVENT_TYPE, "user_context", "not-a-dict"), # ToolRequestEvent (ToolRequestEvent.EVENT_TYPE, "model", 0), (ToolRequestEvent.EVENT_TYPE, "tool_calls", [0]), + (ToolRequestEvent.EVENT_TYPE, "user_context", "not-a-dict"), # ToolResponseEvent (ToolResponseEvent.EVENT_TYPE, "request_id", "not-a-uuid"), (ToolResponseEvent.EVENT_TYPE, "responses", "not-a-dict"), (ToolResponseEvent.EVENT_TYPE, "success", "not-a-dict"), (ToolResponseEvent.EVENT_TYPE, "error", "not-a-dict"), + (ToolResponseEvent.EVENT_TYPE, "user_context", "not-a-dict"), (ChatResponseEvent.EVENT_TYPE, "model_routing", "not-a-dict"), (ToolResponseEvent.EVENT_TYPE, "timestamp", "not-a-number"), # ContextRetrievalRequestEvent (ContextRetrievalRequestEvent.EVENT_TYPE, "query", 0), (ContextRetrievalRequestEvent.EVENT_TYPE, "vector_store", 0), (ContextRetrievalRequestEvent.EVENT_TYPE, "max_results", "not-a-number"), + (ContextRetrievalRequestEvent.EVENT_TYPE, "user_context", "not-a-dict"), # ContextRetrievalResponseEvent (ContextRetrievalResponseEvent.EVENT_TYPE, "request_id", "not-a-uuid"), (ContextRetrievalResponseEvent.EVENT_TYPE, "query", 0), (ContextRetrievalResponseEvent.EVENT_TYPE, "documents", [0]), + (ContextRetrievalResponseEvent.EVENT_TYPE, "user_context", "not-a-dict"), ] diff --git a/python/flink_agents/api/tests/test_cross_language_event_snapshots.py b/python/flink_agents/api/tests/test_cross_language_event_snapshots.py index 6a636d633..a718eac1a 100644 --- a/python/flink_agents/api/tests/test_cross_language_event_snapshots.py +++ b/python/flink_agents/api/tests/test_cross_language_event_snapshots.py @@ -344,6 +344,93 @@ def test_python_can_deserialize_chat_response_event_from_java_snapshot() -> None assert typed.response.text == "hi there" +# ── Request context ───────────────────────────────────────────────────── + +_SNAPSHOT_CONTEXT = { + "task": "summary", + "item_id": "42", + "n": 1, + "ok": True, + "tags": ["a"], +} + + +def _build_chat_request_event_with_context() -> ChatRequestEvent: + event = ChatRequestEvent( + model="test-model", + messages=[ChatMessage.of(MessageRole.USER, "hello world")], + user_context=dict(_SNAPSHOT_CONTEXT), + ) + return _force_id(event, _FIXED_EVENT_ID) + + +def test_regenerate_chat_request_event_with_context_python_snapshot() -> None: + if not _regenerate_enabled(): + pytest.skip("Set REGENERATE_SNAPSHOTS=1 to refresh.") + _write_python_snapshot( + "chat_request_event_with_context.json", + _build_chat_request_event_with_context(), + ) + + +def test_chat_request_event_with_context_python_snapshot_is_stable() -> None: + _assert_python_snapshot_stable( + "chat_request_event_with_context.json", + _build_chat_request_event_with_context(), + ) + + +def test_python_can_deserialize_chat_request_event_with_context_from_java_snapshot() -> ( + None +): + typed = ChatRequestEvent.from_event( + _read_java_snapshot("chat_request_event_with_context.json") + ) + assert typed.id == _FIXED_EVENT_ID + assert typed.model == "test-model" + assert typed.messages[0].text == "hello world" + assert typed.user_context == _SNAPSHOT_CONTEXT + + +def _build_chat_response_event_with_context() -> ChatResponseEvent: + event = ChatResponseEvent.success( + request_id=_FIXED_REQUEST_ID, + response=ChatResult( + message=ChatMessage.assistant([TextBlock(text="hi there")]) + ), + user_context=dict(_SNAPSHOT_CONTEXT), + ) + return _force_id(event, _FIXED_EVENT_ID) + + +def test_regenerate_chat_response_event_with_context_python_snapshot() -> None: + if not _regenerate_enabled(): + pytest.skip("Set REGENERATE_SNAPSHOTS=1 to refresh.") + _write_python_snapshot( + "chat_response_event_with_context.json", + _build_chat_response_event_with_context(), + ) + + +def test_chat_response_event_with_context_python_snapshot_is_stable() -> None: + _assert_python_snapshot_stable( + "chat_response_event_with_context.json", + _build_chat_response_event_with_context(), + ) + + +def test_python_can_deserialize_chat_response_event_with_context_from_java_snapshot() -> ( + None +): + typed = ChatResponseEvent.from_event( + _read_java_snapshot("chat_response_event_with_context.json") + ) + assert typed.id == _FIXED_EVENT_ID + assert typed.request_id == _FIXED_REQUEST_ID + assert typed.response.text == "hi there" + assert typed.user_context == _SNAPSHOT_CONTEXT + + # ── ToolRequestEvent ──────────────────────────────────────────────────── diff --git a/python/flink_agents/api/tests/test_request_context.py b/python/flink_agents/api/tests/test_request_context.py new file mode 100644 index 000000000..95af63121 --- /dev/null +++ b/python/flink_agents/api/tests/test_request_context.py @@ -0,0 +1,282 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +"""Optional opaque context on built-in request and response events.""" + +import json +from uuid import UUID, uuid4 + +import pytest +from pyflink.common import Row + +from flink_agents.api.chat_message import ChatMessage, ToolCallBlock +from flink_agents.api.chat_result import ChatResult +from flink_agents.api.decision_models import ( + DecisionRequest, + DecisionResult, + PredicateAnswer, + PredicateQuestion, +) +from flink_agents.api.events.chat_event import ChatRequestEvent, ChatResponseEvent +from flink_agents.api.events.context_retrieval_event import ( + ContextRetrievalRequestEvent, + ContextRetrievalResponseEvent, +) +from flink_agents.api.events.decision_event import ( + DecisionRequestEvent, + DecisionResponseEvent, +) +from flink_agents.api.events.event import Event, InputEvent +from flink_agents.api.events.tool_event import ToolRequestEvent, ToolResponseEvent +from flink_agents.api.tools.tool_response import ToolResponse +from flink_agents.api.vector_stores.vector_store import Document + +_CONTEXT = { + "task": "summary", + "item_id": "42", + "n": 1, + "ok": True, + "tags": ["a"], +} + + +def test_omitted_context_is_absent_from_the_serialized_event() -> None: + request = ChatRequestEvent(model="model", messages=[ChatMessage.user("hi")]) + assert request.user_context is None + assert "user_context" not in request.attributes + assert "user_context" not in request.model_dump_json() + + response = ChatResponseEvent.success( + uuid4(), ChatResult(message=ChatMessage.assistant("ok")) + ) + assert response.user_context is None + assert "user_context" not in response.model_dump_json() + + +def test_empty_context_is_kept() -> None: + request = ChatRequestEvent( + model="model", messages=[ChatMessage.user("hi")], user_context={} + ) + assert "user_context" in request.attributes + assert request.user_context == {} + + +def test_each_request_and_response_round_trips_its_own_context() -> None: + request_id = uuid4() + decision_request = DecisionRequest( + input="evidence", + questions={"urgent": PredicateQuestion(instructions="Needs attention?")}, + ) + decision_result = DecisionResult( + answers={"urgent": PredicateAnswer(probability_true=0.8)} + ) + events = [ + ChatRequestEvent( + model="model", + messages=[ChatMessage.user("hi")], + prompt_args={"tone": "brief"}, + user_context=_CONTEXT, + ), + ChatRequestEvent( + model="model", + messages=[ChatMessage.user("hi")], + user_context={"label": "summary-42"}, + ), + ChatResponseEvent.success( + request_id, + ChatResult(message=ChatMessage.assistant("ok")), + user_context=_CONTEXT, + ), + ChatResponseEvent.failed( + request_id, "TimeoutError: timed out", 1, 2, user_context=_CONTEXT + ), + ToolRequestEvent( + model="model", + tool_calls=[ToolCallBlock(call_id="call-1", name="lookup", input={"x": 1})], + user_context=_CONTEXT, + ), + ToolResponseEvent( + request_id=request_id, + responses={"call-1": ToolResponse.text("ok")}, + user_context=_CONTEXT, + ), + ContextRetrievalRequestEvent( + query="what is flink", vector_store="store", max_results=2, user_context=_CONTEXT + ), + ContextRetrievalResponseEvent( + request_id=request_id, + query="what is flink", + documents=[Document(content="body", id="d1")], + user_context=_CONTEXT, + ), + DecisionRequestEvent( + model="decision", request=decision_request, user_context=_CONTEXT + ), + DecisionResponseEvent.success(request_id, decision_result, user_context=_CONTEXT), + DecisionResponseEvent.failed( + request_id, "provider down", 2, 3, user_context=_CONTEXT + ), + ] + for event in events: + restored = Event.from_json(event.model_dump_json()) + assert type(restored) is type(event) + assert restored.user_context == event.user_context + assert restored.id == event.id + + +def test_foreign_json_context_survives_the_cross_language_boundary() -> None: + request_json = """ + { + "id": "00000000-0000-0000-0000-000000000001", + "type": "_chat_request_event", + "attributes": { + "model": "test-model", + "messages": [ + {"role": "user", "blocks": [{"type": "text", "text": "hi"}], "metadata": {}} + ], + "user_context": {"task": "summary", "item_id": "42", "n": 1, "ok": true, "tags": ["a"]} + }, + "attachments": {"trace": "keep"} + } + """ + request = ChatRequestEvent.from_event(Event.from_json(request_json)) + assert request.model == "test-model" + assert request.messages[0].text == "hi" + assert request.user_context == _CONTEXT + assert request.get_attachment("trace") == "keep" + assert request.id == UUID("00000000-0000-0000-0000-000000000001") + + failed_json = """ + { + "id": "00000000-0000-0000-0000-000000000001", + "type": "_chat_response_event", + "attributes": { + "request_id": "00000000-0000-0000-0000-000000000002", + "status": "FAILED", + "response": null, + "error": "TimeoutError: timed out", + "retry_count": 2, + "total_retry_wait_sec": 3, + "user_context": {"task": "summary", "item_id": "42"} + }, + "attachments": {} + } + """ + failed = ChatResponseEvent.from_event(Event.from_json(failed_json)) + assert failed.is_failed + assert failed.error == "TimeoutError: timed out" + assert failed.user_context == {"task": "summary", "item_id": "42"} + + +def test_unknown_attribute_is_still_rejected_when_context_is_present() -> None: + event = Event( + type=ChatRequestEvent.EVENT_TYPE, + attributes={ + "model": "m", + "messages": [ChatMessage.user("hi")], + "user_context": {"task": "summary"}, + "extra": 1, + }, + ) + with pytest.raises(ValueError, match="extra"): + ChatRequestEvent.from_event(event) + + +def test_serialized_context_is_json_compatible() -> None: + event = ChatRequestEvent( + model="model", messages=[ChatMessage.user("hi")], user_context=_CONTEXT + ) + payload = json.loads(event.model_dump_json()) + assert payload["attributes"]["user_context"] == _CONTEXT + assert "user_context" not in payload["attributes"]["messages"][0] + + +def test_reused_context_mapping_is_copied_including_nested_values() -> None: + shared = { + "task": "summary", + "tags": ["keep"], + "meta": {"label": "keep"}, + } + request_a = ChatRequestEvent( + model="model", messages=[ChatMessage.user("a")], user_context=shared + ) + shared["task"] = "risk" + shared["tags"].append("changed") + shared["meta"]["label"] = "changed" + request_b = ChatRequestEvent( + model="model", messages=[ChatMessage.user("b")], user_context=shared + ) + + assert request_a.user_context == { + "task": "summary", + "tags": ["keep"], + "meta": {"label": "keep"}, + } + assert request_b.user_context == { + "task": "risk", + "tags": ["keep", "changed"], + "meta": {"label": "changed"}, + } + assert request_a.user_context is not shared + assert request_a.user_context["tags"] is not shared["tags"] + assert request_a.user_context["meta"] is not shared["meta"] + + +_ROW_SHAPED_CONTEXT = {"type": "Row", "values": [1], "label": "keep"} + + +def test_row_shaped_user_context_follows_ordinary_attribute_decoding() -> None: + """A Row-shaped value is reconstructed like any other event attribute. + + Construction leaves the supplied map alone. JSON decoding then applies the + shared row reconstruction, so a nested object becomes a Row and a top-level + Row-shaped map is no longer a dict. + """ + nested = {"note": dict(_ROW_SHAPED_CONTEXT)} + request = ChatRequestEvent( + model="model", messages=[ChatMessage.user("hi")], user_context=nested + ) + assert request.user_context == nested + assert type(request.user_context["note"]) is dict + + restored = Event.from_json(request.model_dump_json()) + assert isinstance(restored.user_context, dict) + assert restored.user_context["note"] == Row(1) + + raw = { + "type": ChatRequestEvent.EVENT_TYPE, + "attributes": { + "model": "model", + "messages": [ChatMessage.user("hi").model_dump(mode="json")], + "user_context": dict(_ROW_SHAPED_CONTEXT), + }, + } + with pytest.raises(ValueError, match="Malformed built-in event"): + Event.from_json(json.dumps(raw)) + + +def test_non_context_row_values_are_still_reconstructed() -> None: + event = Event.from_json( + json.dumps( + { + "type": InputEvent.EVENT_TYPE, + "attributes": {"input": {"type": "Row", "values": [1]}}, + } + ) + ) + assert isinstance(event, InputEvent) + assert event.input == Row(1) diff --git a/python/flink_agents/plan/actions/chat_model_action.py b/python/flink_agents/plan/actions/chat_model_action.py index 1496e6703..036c56170 100644 --- a/python/flink_agents/plan/actions/chat_model_action.py +++ b/python/flink_agents/plan/actions/chat_model_action.py @@ -66,6 +66,11 @@ # position in the action's durable call sequence. _CHAT_DURABLE_ID = "chat" _PROMPT_ARGS = "prompt_args" +# Sensory-memory key. Matches the event attribute name. +_USER_CONTEXT = "user_context" +# Not a built-in event type. A carrier of this type runs Event JSON encode and +# decode, including Row reconstruction, for the stored map. +_USER_CONTEXT_CARRIER = "_user_context_memory" _FINISH_REASON = "finish_reason" _TRUNCATED_FINISH_REASON = "length" _CONTENT_FILTERED_FINISH_REASON = "content_filter" @@ -189,7 +194,12 @@ async def _durable_invoke( return ctx.durable_execute(func, *args, **kwargs) -def _send_failure(request_id: UUID, error: Exception, ctx: RunnerContext) -> None: +def _send_failure( + request_id: UUID, + error: Exception, + ctx: RunnerContext, + user_context: dict[str, Any] | None = None, +) -> None: stats = _get_retry_stats(ctx.sensory_memory, request_id) _clear_request_context(ctx.sensory_memory, request_id) ctx.send_event( @@ -198,6 +208,7 @@ def _send_failure(request_id: UUID, error: Exception, ctx: RunnerContext) -> Non _error_text(error), stats["total_retry_count"], stats["total_retry_wait_sec"], + user_context=user_context, ) ) @@ -252,6 +263,41 @@ def _update_tool_call_context( return _deserialize_messages(tool_call_context[key]) +def _serialize_user_context(user_context: dict[str, Any]) -> dict[str, Any]: + """Encode ``user_context`` with the JSON rules used for Event attributes. + + Sensory memory accepts only checkpoint-stable values. An Event attribute + also accepts UUIDs, datetimes, Pydantic models, tuples, and ``Row``, so the + value written here is that Event JSON form. + """ + encoded = json.loads( + Event( + type=_USER_CONTEXT_CARRIER, + attributes={_USER_CONTEXT: user_context}, + ).model_dump_json() + )["attributes"][_USER_CONTEXT] + if not isinstance(encoded, dict): + msg = "user_context must serialize to a JSON object" + raise TypeError(msg) + return encoded + + +def _deserialize_user_context(stored: dict[str, Any]) -> dict[str, Any]: + """Restore stored ``user_context`` the way ``Event.from_json`` decodes it.""" + restored = Event.from_json( + json.dumps( + { + "type": _USER_CONTEXT_CARRIER, + "attributes": {_USER_CONTEXT: stored}, + } + ) + ).attributes[_USER_CONTEXT] + if not isinstance(restored, dict): + msg = "user_context must deserialize to a dict" + raise TypeError(msg) + return restored + + def _save_tool_request_event_context( sensory_memory: MemoryObject, tool_request_event_id: UUID, @@ -259,10 +305,11 @@ def _save_tool_request_event_context( model: str, prompt_args: Dict | None, output_schema: OutputSchema | None, + user_context: dict[str, Any] | None = None, ) -> None: """Save the context for a specific tool request event.""" - context = sensory_memory.get(_TOOL_REQUEST_EVENT_CONTEXT) or {} - context[str(tool_request_event_id)] = { + stored = sensory_memory.get(_TOOL_REQUEST_EVENT_CONTEXT) or {} + entry = { "initial_request_id": str(initial_request_id), "model": model, _PROMPT_ARGS: prompt_args if prompt_args is not None else {}, @@ -270,7 +317,13 @@ def _save_tool_request_event_context( if output_schema is not None else None, } - sensory_memory.set(_TOOL_REQUEST_EVENT_CONTEXT, context) + # Absent when the request omitted it, so a checkpoint from before this field + # still resumes with no user context. Present values are Event JSON: memory + # rejects objects an Event attribute can already carry. + if user_context is not None: + entry[_USER_CONTEXT] = _serialize_user_context(user_context) + stored[str(tool_request_event_id)] = entry + sensory_memory.set(_TOOL_REQUEST_EVENT_CONTEXT, stored) def _get_tool_request_event_context( @@ -290,6 +343,10 @@ def _get_tool_request_event_context( if output_schema is not None else None ) + if _USER_CONTEXT in removed_context: + removed_context[_USER_CONTEXT] = _deserialize_user_context( + removed_context[_USER_CONTEXT] + ) return removed_context @@ -391,6 +448,7 @@ def _handle_tool_calls( prompt_args: Dict | None, output_schema: OutputSchema | None, ctx: RunnerContext, + user_context: dict[str, Any] | None = None, ) -> None: """Handle tool calls in chat response.""" _update_tool_call_context( @@ -413,6 +471,7 @@ def _handle_tool_calls( model, prompt_args, output_schema, + user_context, ) ctx.send_event(tool_request_event) @@ -758,6 +817,8 @@ async def chat( prompt_args: Dict | None, output_schema: OutputSchema | None, ctx: RunnerContext, + *, + user_context: dict[str, Any] | None = None, ) -> None: """Chat with llm. @@ -776,7 +837,7 @@ async def chat( initial_request_id, model, messages, prompt_args, output_schema, ctx ) except _ChatAttemptFailed as failure: - _send_failure(initial_request_id, failure.error, ctx) + _send_failure(initial_request_id, failure.error, ctx, user_context) return response = attempt.response chat_model = attempt.chat_model @@ -794,6 +855,7 @@ async def chat( prompt_args, output_schema, ctx, + user_context, ) else: # if there is no tool call generated, return chat response directly retry_stats = _get_retry_stats(ctx.sensory_memory, initial_request_id) @@ -805,6 +867,7 @@ async def chat( response=response, retry_count=total_retry_count, total_retry_wait_sec=total_retry_wait_sec, + user_context=user_context, ) completed.attributes["structured_output"] = structured_output ctx.send_event(completed) @@ -819,6 +882,7 @@ async def _process_chat_request(event: ChatRequestEvent, ctx: RunnerContext) -> prompt_args=event.prompt_args, output_schema=event.output_schema, ctx=ctx, + user_context=event.user_context, ) @@ -851,6 +915,7 @@ async def _process_tool_response(event: ToolResponseEvent, ctx: RunnerContext) - prompt_args=tool_request_event_context.get(_PROMPT_ARGS, {}), output_schema=tool_request_event_context["output_schema"], ctx=ctx, + user_context=tool_request_event_context.get(_USER_CONTEXT), ) diff --git a/python/flink_agents/plan/actions/context_retrieval_action.py b/python/flink_agents/plan/actions/context_retrieval_action.py index 38dcdd918..2456b5487 100644 --- a/python/flink_agents/plan/actions/context_retrieval_action.py +++ b/python/flink_agents/plan/actions/context_retrieval_action.py @@ -59,7 +59,10 @@ async def process_context_retrieval_request(event: Event, ctx: RunnerContext) -> ctx.send_event( ContextRetrievalResponseEvent( - request_id=event.id, query=event.query, documents=result.documents + request_id=event.id, + query=event.query, + documents=result.documents, + user_context=event.user_context, ) ) diff --git a/python/flink_agents/plan/actions/decision_model_action.py b/python/flink_agents/plan/actions/decision_model_action.py index b8dab7f2e..111ccc77b 100644 --- a/python/flink_agents/plan/actions/decision_model_action.py +++ b/python/flink_agents/plan/actions/decision_model_action.py @@ -158,13 +158,21 @@ def invoke(request: DecisionRequest) -> tuple[DecisionResult | None, str | None] if failure is not None: ctx.send_event( DecisionResponseEvent.failed( - event.id, failure, retry_count, total_retry_wait_sec + event.id, + failure, + retry_count, + total_retry_wait_sec, + user_context=event.user_context, ) ) else: ctx.send_event( DecisionResponseEvent.success( - event.id, response, retry_count, total_retry_wait_sec + event.id, + response, + retry_count, + total_retry_wait_sec, + user_context=event.user_context, ) ) diff --git a/python/flink_agents/plan/actions/tool_call_action.py b/python/flink_agents/plan/actions/tool_call_action.py index 160e4c1e9..9a7d8bd0a 100644 --- a/python/flink_agents/plan/actions/tool_call_action.py +++ b/python/flink_agents/plan/actions/tool_call_action.py @@ -170,6 +170,7 @@ async def process_tool_request(event: Event, ctx: RunnerContext) -> None: responses=responses, success=success, error=error, + user_context=event.user_context, ) ) diff --git a/python/flink_agents/plan/tests/actions/test_decision_model_action.py b/python/flink_agents/plan/tests/actions/test_decision_model_action.py index 66b85f19c..a0549df93 100644 --- a/python/flink_agents/plan/tests/actions/test_decision_model_action.py +++ b/python/flink_agents/plan/tests/actions/test_decision_model_action.py @@ -141,6 +141,51 @@ def _response(ctx) -> DecisionResponseEvent: return ctx.send_event.call_args.args[0] +def test_success_and_failure_echo_context_without_changing_the_provider_request(): + context = {"task": "review", "item_id": "42"} + event = DecisionRequestEvent( + model="review-decision", + request=DecisionRequest( + input={"review": "The package arrived damaged."}, + questions={"urgent": PredicateQuestion(instructions="Needs attention?")}, + ), + request_params={"safety_identifier": "review-user", "timeout": 4.0}, + user_context=context, + ) + event.set_attachment("trace", "keep-on-request") + model = _Model([_result()]) + ctx, _ = _context(model, decision_async=False, interval=0) + asyncio.run(process_decision_request(event, ctx)) + response = _response(ctx) + assert response.is_success + assert response.user_context == context + assert response.attachments == {} + model.decide.assert_called_once_with( + event.request, request_params=event.request_params + ) + + failed_event = DecisionRequestEvent( + model="review-decision", + request=event.request, + request_params=event.request_params, + user_context=context, + ) + failed_model = _Model([RuntimeError("provider down")]) + failed_ctx, _ = _context(failed_model, decision_async=False, retries=0, interval=0) + asyncio.run(process_decision_request(failed_event, failed_ctx)) + failed = _response(failed_ctx) + assert failed.is_failed + assert "provider down" in failed.error + assert failed.user_context == context + + omitted = _request() + omitted_model = _Model([_result()]) + omitted_ctx, _ = _context(omitted_model, decision_async=False, interval=0) + asyncio.run(process_decision_request(omitted, omitted_ctx)) + assert _response(omitted_ctx).user_context is None + assert "user_context" not in _response(omitted_ctx).attributes + + @pytest.mark.parametrize("decision_async", [False, True]) def test_success_uses_selected_durable_path_and_action_owned_context(decision_async): event = _request() diff --git a/python/flink_agents/plan/tests/actions/test_request_context_propagation.py b/python/flink_agents/plan/tests/actions/test_request_context_propagation.py new file mode 100644 index 000000000..1054ca351 --- /dev/null +++ b/python/flink_agents/plan/tests/actions/test_request_context_propagation.py @@ -0,0 +1,427 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +"""Built-in actions echo optional request context onto the matching response.""" + +import asyncio +import json +from datetime import datetime, timezone +from typing import Any +from unittest.mock import MagicMock +from uuid import UUID + +import pytest +from pydantic import BaseModel +from pyflink.common import Row + +from flink_agents.api.chat_message import ChatMessage, ToolCallBlock +from flink_agents.api.chat_result import ChatResult +from flink_agents.api.core_options import AgentExecutionOptions +from flink_agents.api.events.chat_event import ChatRequestEvent, ChatResponseEvent +from flink_agents.api.events.context_retrieval_event import ( + ContextRetrievalRequestEvent, + ContextRetrievalResponseEvent, +) +from flink_agents.api.events.event import Event +from flink_agents.api.events.tool_event import ToolRequestEvent, ToolResponseEvent +from flink_agents.api.memory_object import validate_memory_value +from flink_agents.api.resource import ResourceType +from flink_agents.api.tools.tool_response import ToolResponse +from flink_agents.api.trace import ExecutionReporter +from flink_agents.api.vector_stores.vector_store import Document, VectorStoreQuery +from flink_agents.plan.actions.chat_model_action import ( + process_chat_request_or_tool_response, +) +from flink_agents.plan.actions.context_retrieval_action import ( + process_context_retrieval_request, +) + + +class _Memory: + def __init__(self) -> None: + self._store: dict[str, Any] = {} + + def get(self, path: str) -> Any: + return self._store.get(path) + + def set(self, path: str, value: Any) -> None: + validate_memory_value(path, value) + self._store[path] = value + + +def _runner(chat_model: Any, max_retries: int = 0) -> tuple[MagicMock, list, _Memory]: + sent: list = [] + memory = _Memory() + chat_model.model = "configured-model" + chat_model.will_apply_native_structured_output = MagicMock(return_value=False) + config = MagicMock() + options = { + id(AgentExecutionOptions.MAX_RETRIES): max_retries, + id(AgentExecutionOptions.RETRY_WAIT_INTERVAL): 0, + id(AgentExecutionOptions.CHAT_ASYNC): False, + } + config.get = MagicMock( + side_effect=lambda option: options.get(id(option), option.get_default_value()) + ) + ctx = MagicMock(spec=ExecutionReporter) + ctx.config = config + ctx.sensory_memory = memory + ctx.action_metric_group = None + ctx.send_event = MagicMock(side_effect=sent.append) + ctx.get_resource = MagicMock(return_value=chat_model) + ctx.durable_execute = MagicMock( + side_effect=lambda fn, *args, durable_id=None, **kwargs: fn(*args, **kwargs) + ) + + async def _async( + fn: Any, *args: Any, durable_id: str | None = None, **kwargs: Any + ) -> Any: + return fn(*args, **kwargs) + + ctx.durable_execute_async = MagicMock(side_effect=_async) + return ctx, sent, memory + + +def _text(text: str) -> ChatResult: + return ChatResult(message=ChatMessage.assistant(text)) + + +def _tool_result() -> ChatResult: + return ChatResult( + message=ChatMessage.assistant( + [ToolCallBlock(call_id="call-1", name="lookup", input={"x": 1})] + ) + ) + + +def _request(text: str, prompt_args: dict, context: Any | None) -> ChatRequestEvent: + event = ChatRequestEvent( + model="test-model", + messages=[ChatMessage.user(text)], + prompt_args=prompt_args, + user_context=context, + ) + event.set_attachment("trace", "keep-on-request") + return event + + +def test_omitted_context_stays_off_the_response() -> None: + chat_model = MagicMock() + chat_model.chat = MagicMock(return_value=_text("hello")) + ctx, sent, _ = _runner(chat_model) + request = ChatRequestEvent( + model="test-model", + messages=[ChatMessage.user("hi")], + prompt_args={"tone": "brief"}, + ) + + asyncio.run(process_chat_request_or_tool_response(request, ctx)) + + response = sent[0] + assert isinstance(response, ChatResponseEvent) + assert response.user_context is None + assert "user_context" not in response.attributes + assert response.request_id == request.id + assert response.attachments == {} + + +def test_success_failure_and_retry_echo_context_without_touching_the_prompt() -> None: + context = {"task": "summary", "item_id": "42", "n": 1, "ok": True, "tags": ["a"]} + prompt_args = {"tone": "brief"} + chat_model = MagicMock() + chat_model.chat = MagicMock(side_effect=[RuntimeError("transient"), _text("done")]) + ctx, sent, _ = _runner(chat_model, max_retries=1) + request = _request("hi", prompt_args, context) + + asyncio.run(process_chat_request_or_tool_response(request, ctx)) + + response = sent[0] + assert response.is_success + assert response.user_context == context + assert response.retry_count == 1 + assert response.attachments == {} + assert response.request_id == request.id + for call in chat_model.chat.call_args_list: + assert call.kwargs == {"prompt_args": prompt_args} + assert call.args[0][0].text == "hi" + assert "summary" not in call.args[0][0].text + + chat_model.chat = MagicMock(side_effect=RuntimeError("persistent")) + failing = _request("hi", prompt_args, context) + sent.clear() + asyncio.run(process_chat_request_or_tool_response(failing, ctx)) + failed = sent[0] + assert failed.is_failed + assert "persistent" in failed.error + assert failed.user_context == context + assert failed.request_id == failing.id + + +def test_tool_rounds_keep_each_request_context() -> None: + context_a = {"task": "summary", "item_id": "42"} + context_b = {"task": "risk_assessment", "item_id": "42"} + chat_model = MagicMock() + chat_model.chat = MagicMock( + side_effect=[ + _tool_result(), + _tool_result(), + _text("risk done"), + _text("summary done"), + ] + ) + ctx, sent, memory = _runner(chat_model) + request_a = _request("a", {"who": "a"}, context_a) + request_b = _request("b", {"who": "b"}, context_b) + + asyncio.run(process_chat_request_or_tool_response(request_a, ctx)) + asyncio.run(process_chat_request_or_tool_response(request_b, ctx)) + tool_a = sent[0] + tool_b = sent[1] + assert isinstance(tool_a, ToolRequestEvent) + assert tool_a.user_context is None + assert tool_b.user_context is None + assert tool_a.tool_calls[0].input == {"x": 1} + assert tool_a.attachments == {} + + parked = memory.get("_TOOL_REQUEST_EVENT_CONTEXT") + parked[str(tool_b.id)]["user_context"] = json.loads( + json.dumps(parked[str(tool_b.id)]["user_context"]) + ) + + def _tool_response(tool_request: ToolRequestEvent) -> ToolResponseEvent: + return ToolResponseEvent( + request_id=tool_request.id, + responses={"call-1": ToolResponse.text("ok")}, + success={"call-1": True}, + error={}, + ) + + asyncio.run(process_chat_request_or_tool_response(_tool_response(tool_b), ctx)) + asyncio.run(process_chat_request_or_tool_response(_tool_response(tool_a), ctx)) + + response_b = sent[2] + response_a = sent[3] + assert response_b.request_id == request_b.id + assert response_b.user_context == context_b + assert response_b.response.text == "risk done" + assert response_a.request_id == request_a.id + assert response_a.user_context == context_a + assert response_a.response.text == "summary done" + assert response_a.attachments == {} + for call in chat_model.chat.call_args_list: + assert "context" not in call.kwargs + assert set(call.kwargs) == {"prompt_args"} + + +class _Note(BaseModel): + label: str + n: int + + +def test_tool_round_restores_event_serializable_user_context() -> None: + """A tool round stores user context as Event JSON and restores it. + + UUIDs, datetimes, Pydantic models, tuples, and Row are valid Event + attributes and are rejected by sensory memory. The final response carries + the same value ``Event.from_json`` would restore. + """ + context = { + "request_id": UUID("12345678-1234-5678-1234-567812345678"), + "when": datetime(2026, 11, 6, 8, 46, 54, tzinfo=timezone.utc), + "note": _Note(label="keep", n=1), + "pair": ("left", 2), + "payload": {"row": Row(1, "x"), "tags": ("a", "b")}, + } + chat_model = MagicMock() + chat_model.chat = MagicMock(side_effect=[_tool_result(), _text("done")]) + ctx, sent, memory = _runner(chat_model) + request = _request("hi", {"tone": "brief"}, context) + + asyncio.run(process_chat_request_or_tool_response(request, ctx)) + + tool_request = sent[0] + assert isinstance(tool_request, ToolRequestEvent) + assert tool_request.user_context is None + stored = memory.get("_TOOL_REQUEST_EVENT_CONTEXT")[str(tool_request.id)][ + "user_context" + ] + validate_memory_value("user_context", stored) + assert type(stored["request_id"]) is str + assert type(stored["when"]) is str + assert stored["note"] == {"label": "keep", "n": 1} + assert stored["pair"] == ["left", 2] + assert stored["payload"]["row"]["type"] == "Row" + assert type(stored["payload"]["tags"]) is list + # Checkpoint restore returns the JSON form sensory memory persisted. + parked = memory.get("_TOOL_REQUEST_EVENT_CONTEXT") + parked[str(tool_request.id)]["user_context"] = json.loads(json.dumps(stored)) + + asyncio.run( + process_chat_request_or_tool_response( + ToolResponseEvent( + request_id=tool_request.id, + responses={"call-1": ToolResponse.text("ok")}, + success={"call-1": True}, + error={}, + ), + ctx, + ) + ) + + response = sent[1] + assert isinstance(response, ChatResponseEvent) + assert response.request_id == request.id + assert response.response.text == "done" + assert ( + response.user_context == Event.from_json(request.model_dump_json()).user_context + ) + assert isinstance(response.user_context["payload"]["row"], Row) + assert response.user_context["payload"]["row"] == Row(1, "x") + + +def test_restored_request_still_echoes_context() -> None: + chat_model = MagicMock() + chat_model.chat = MagicMock(return_value=_text("ok")) + ctx, sent, _ = _runner(chat_model) + original = ChatRequestEvent( + model="test-model", + messages=[ChatMessage.user("hi")], + user_context={"task": "summary", "item_id": "42"}, + ) + restored = Event.from_json(original.model_dump_json()) + + asyncio.run(process_chat_request_or_tool_response(restored, ctx)) + + assert sent[0].request_id == original.id + assert sent[0].user_context == original.user_context + + +def test_retrieval_echoes_context_and_leaves_the_query_unchanged() -> None: + seen: list[VectorStoreQuery] = [] + + def query(vector_query: VectorStoreQuery) -> Any: + seen.append(vector_query) + return MagicMock(documents=[Document(content="body", id="d1")]) + + ctx = MagicMock() + ctx.config.get = MagicMock(return_value=False) + ctx.get_resource = MagicMock(return_value=MagicMock(query=query)) + ctx.durable_execute = MagicMock( + side_effect=lambda fn, *args, durable_id=None, **kwargs: fn(*args, **kwargs) + ) + sent: list = [] + ctx.send_event = MagicMock(side_effect=sent.append) + context = {"task": "retrieve", "item_id": "42"} + request = ContextRetrievalRequestEvent( + query="what is flink", vector_store="store", max_results=2, user_context=context + ) + request.set_attachment("trace", "keep-on-request") + + asyncio.run(process_context_retrieval_request(request, ctx)) + + response = sent[0] + assert isinstance(response, ContextRetrievalResponseEvent) + assert response.request_id == request.id + assert response.query == "what is flink" + assert response.user_context == context + assert response.attachments == {} + assert seen[0].query_text == "what is flink" + assert seen[0].limit == 2 + ctx.get_resource.assert_called_once_with("store", ResourceType.VECTOR_STORE) + + +def test_retrieval_omits_context_when_the_request_does() -> None: + ctx = MagicMock() + ctx.config.get = MagicMock(return_value=False) + ctx.get_resource = MagicMock( + return_value=MagicMock(query=MagicMock(return_value=MagicMock(documents=[]))) + ) + ctx.durable_execute = MagicMock( + side_effect=lambda fn, *args, durable_id=None, **kwargs: fn(*args, **kwargs) + ) + sent: list = [] + ctx.send_event = MagicMock(side_effect=sent.append) + + asyncio.run( + process_context_retrieval_request( + ContextRetrievalRequestEvent(query="q", vector_store="store"), ctx + ) + ) + + assert sent[0].user_context is None + assert "user_context" not in sent[0].attributes + + +def test_retrieval_query_failure_propagates_without_a_response() -> None: + def query(vector_query: VectorStoreQuery) -> Any: + msg = "store down" + raise RuntimeError(msg) + + ctx = MagicMock() + ctx.config.get = MagicMock(return_value=False) + ctx.get_resource = MagicMock(return_value=MagicMock(query=query)) + ctx.durable_execute = MagicMock( + side_effect=lambda fn, *args, durable_id=None, **kwargs: fn(*args, **kwargs) + ) + sent: list = [] + ctx.send_event = MagicMock(side_effect=sent.append) + request = ContextRetrievalRequestEvent( + query="q", vector_store="store", user_context={"task": "retrieve"} + ) + + with pytest.raises(RuntimeError, match="store down"): + asyncio.run(process_context_retrieval_request(request, ctx)) + + assert sent == [] + assert request.user_context == {"task": "retrieve"} + + +def test_reused_context_mapping_keeps_each_response_snapshot() -> None: + chat_model = MagicMock() + chat_model.chat = MagicMock(side_effect=[_text("summary done"), _text("risk done")]) + ctx, sent, _ = _runner(chat_model) + shared = { + "task": "summary", + "tags": ["keep"], + "meta": {"label": "keep"}, + } + request_a = _request("a", {"who": "a"}, shared) + shared["task"] = "risk" + shared["tags"].append("changed") + shared["meta"]["label"] = "changed" + request_b = _request("b", {"who": "b"}, shared) + + asyncio.run(process_chat_request_or_tool_response(request_a, ctx)) + asyncio.run(process_chat_request_or_tool_response(request_b, ctx)) + + response_a, response_b = sent + assert response_a.request_id == request_a.id + assert response_a.response.text == "summary done" + assert response_a.user_context == { + "task": "summary", + "tags": ["keep"], + "meta": {"label": "keep"}, + } + assert response_b.request_id == request_b.id + assert response_b.response.text == "risk done" + assert response_b.user_context == { + "task": "risk", + "tags": ["keep", "changed"], + "meta": {"label": "changed"}, + } + assert response_a.user_context is not shared + assert response_a.user_context["tags"] is not shared["tags"] diff --git a/python/flink_agents/plan/tests/actions/test_tool_call_action.py b/python/flink_agents/plan/tests/actions/test_tool_call_action.py index 29b7017e1..2183b5148 100644 --- a/python/flink_agents/plan/tests/actions/test_tool_call_action.py +++ b/python/flink_agents/plan/tests/actions/test_tool_call_action.py @@ -249,6 +249,63 @@ def test_restored_tool_request_reaches_action_as_concrete_type() -> None: assert response.success["call-1"] is True +def test_tool_response_echoes_context_on_success_and_failure() -> None: + context = {"task": "lookup", "item_id": "42"} + ctx = _Context() + event = _request( + model="model", + tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": { + "name": "query_order", + "arguments": {"order_id": "order-1"}, + }, + } + ], + ) + event.attributes["user_context"] = context + event = ToolRequestEvent.from_event(event) + event.set_attachment("trace", "keep-on-request") + + asyncio.run(process_tool_request(event, ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.user_context == context + assert response.request_id == event.id + assert response.responses["call-1"].get_text() == "tenant-1:order-1" + assert response.attachments == {} + assert response.success["call-1"] is True + _func, _args, executed = ctx.durable_execute_calls[0] + assert "user_context" not in executed + assert executed["order_id"] == "order-1" + + missing = ToolRequestEvent( + model="model", + tool_calls=[ + ToolCallBlock( + call_id="call-1", name="missing_tool", input={"order_id": "order-1"} + ) + ], + user_context=context, + ) + failed_ctx = _Context() + asyncio.run(process_tool_request(missing, failed_ctx)) + failure = ToolResponseEvent.from_event(failed_ctx.sent_events[0]) + assert failure.success["call-1"] is False + assert failure.user_context == context + assert failure.attachments == {} + + +def test_tool_response_omits_context_when_the_request_does() -> None: + ctx = _Context() + asyncio.run(process_tool_request(tool_request(), ctx)) + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.user_context is None + assert "user_context" not in response.attributes + + def test_tool_call_action_injects_args_from_config_without_mutating_request() -> None: ctx = _Context() arguments = {"order_id": "order-1"} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/ChatModelActionContextRecoveryTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/ChatModelActionContextRecoveryTest.java new file mode 100644 index 000000000..589e54fe5 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/ChatModelActionContextRecoveryTest.java @@ -0,0 +1,165 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.runtime.context; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ChatResult; +import org.apache.flink.agents.api.chat.messages.ToolCallBlock; +import org.apache.flink.agents.api.chat.model.BaseChatModelSetup; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.MemoryObject; +import org.apache.flink.agents.api.context.MemoryUpdate; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ChatRequestEvent; +import org.apache.flink.agents.api.event.ChatResponseEvent; +import org.apache.flink.agents.api.event.ToolRequestEvent; +import org.apache.flink.agents.api.event.ToolResponseEvent; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.apache.flink.agents.plan.actions.ChatModelAction; +import org.apache.flink.agents.runtime.actionstate.ActionState; +import org.apache.flink.agents.runtime.actionstate.ActionStateSerde; +import org.apache.flink.agents.runtime.memory.CachedMemoryStore; +import org.apache.flink.agents.runtime.memory.ForTestMemoryMapState; +import org.apache.flink.agents.runtime.memory.MemoryObjectImpl; +import org.apache.flink.agents.runtime.memory.MemoryUpdateReplayer; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * The chat tool-round parks the request context in sensory memory. Restoring that {@link + * MemoryUpdate} through {@link ActionStateSerde} and replaying it must still echo the original + * context on the final chat response. + */ +class ChatModelActionContextRecoveryTest { + + private static final String TOOL_REQUEST_EVENT_CONTEXT = "_TOOL_REQUEST_EVENT_CONTEXT"; + + @Test + void restoredToolRequestMemoryKeepsTheOriginalRequestContext() throws Exception { + Map context = new LinkedHashMap<>(); + context.put("task", "summary"); + context.put("item_id", "42"); + List tags = new ArrayList<>(); + tags.add("a"); + context.put("tags", tags); + + List sentEvents = new ArrayList<>(); + List updates = new ArrayList<>(); + MemoryObject sensory = newMemory(updates); + RunnerContext ctx = context(sensory, sentEvents); + BaseChatModelSetup chatModel = mock(BaseChatModelSetup.class); + when(ctx.getResource(anyString(), eq(ResourceType.CHAT_MODEL))).thenReturn(chatModel); + when(chatModel.chat(any(), any(), any())) + .thenReturn( + new ChatResult( + ChatMessage.assistant( + List.of( + new ToolCallBlock( + "call-1", "lookup", Map.of("x", 1)))))) + .thenReturn(new ChatResult(ChatMessage.assistant("done"))); + + ChatRequestEvent request = + new ChatRequestEvent( + "test-model", List.of(ChatMessage.user("hi")), null, null, context); + ChatModelAction.processChatRequestOrToolResponse(request, ctx); + + ToolRequestEvent toolRequest = ToolRequestEvent.fromEvent(sentEvents.get(0)); + assertThat(toolRequest.getUserContext()).isNull(); + Object liveParked = sensory.get(TOOL_REQUEST_EVENT_CONTEXT).getValue(); + + ActionState state = new ActionState(request); + for (MemoryUpdate update : updates) { + state.addSensoryMemoryUpdate(update); + } + state.addEvent(toolRequest); + ActionState restored = ActionStateSerde.deserialize(ActionStateSerde.serialize(state)); + + MemoryUpdate parked = + restored.getSensoryMemoryUpdates().stream() + .filter(update -> TOOL_REQUEST_EVENT_CONTEXT.equals(update.getPath())) + .reduce((first, second) -> second) + .orElseThrow(); + assertThat(parked.getValue()).isEqualTo(liveParked); + assertThat(parked.getValue()).isNotSameAs(liveParked); + + MemoryObject recovered = newMemory(new ArrayList<>()); + MemoryUpdateReplayer.replay(recovered, restored.getSensoryMemoryUpdates()); + when(ctx.getSensoryMemory()).thenReturn(recovered); + sentEvents.clear(); + + ToolRequestEvent restoredTool = + ToolRequestEvent.fromEvent(restored.getOutputEvents().get(0)); + assertThat(restoredTool.getId()).isEqualTo(toolRequest.getId()); + ChatModelAction.processChatRequestOrToolResponse( + new ToolResponseEvent( + restoredTool.getId(), + Map.of("call-1", ToolResponse.text("ok")), + Map.of("call-1", true), + Map.of()), + ctx); + + ChatResponseEvent response = ChatResponseEvent.fromEvent(sentEvents.get(0)); + assertThat(response.getRequestId()).isEqualTo(request.getId()); + assertThat(response.getResponse().getText()).isEqualTo("done"); + assertThat(response.getUserContext()).isEqualTo(context); + assertThat(response.getUserContext()).isNotSameAs(context); + } + + private static MemoryObject newMemory(List updates) throws Exception { + return new MemoryObjectImpl( + MemoryObject.MemoryType.SENSORY, + new CachedMemoryStore(new ForTestMemoryMapState<>()), + MemoryObjectImpl.ROOT_KEY, + updates); + } + + private static RunnerContext context(MemoryObject sensory, List sentEvents) + throws Exception { + RunnerContext ctx = mock(RunnerContext.class); + AgentConfiguration config = new AgentConfiguration(); + config.set(AgentExecutionOptions.MAX_RETRIES, 0); + config.set(AgentExecutionOptions.RETRY_WAIT_INTERVAL, 0); + config.set(AgentExecutionOptions.CHAT_ASYNC, false); + when(ctx.getConfig()).thenReturn(config); + when(ctx.getSensoryMemory()).thenReturn(sensory); + when(ctx.durableExecute(any())) + .thenAnswer( + invocation -> invocation.>getArgument(0).call()); + doAnswer(invocation -> sentEvents.add(invocation.getArgument(0))) + .when(ctx) + .sendEvent(any()); + return ctx; + } +}