diff --git a/api/src/main/java/org/apache/flink/agents/api/agents/ReActAgent.java b/api/src/main/java/org/apache/flink/agents/api/agents/ReActAgent.java index 10ed457e0..798b453f3 100644 --- a/api/src/main/java/org/apache/flink/agents/api/agents/ReActAgent.java +++ b/api/src/main/java/org/apache/flink/agents/api/agents/ReActAgent.java @@ -178,7 +178,7 @@ public static void stopAction(Event event, RunnerContext ctx) { if (response.getExtraArgs().containsKey(STRUCTURED_OUTPUT)) { output = response.getExtraArgs().get(STRUCTURED_OUTPUT); } else { - output = String.valueOf(response.getContent()); + output = response.getText(); } ctx.sendEvent(new OutputEvent(output)); diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/messages/AudioBlock.java b/api/src/main/java/org/apache/flink/agents/api/chat/messages/AudioBlock.java new file mode 100644 index 000000000..ea82c269e --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/messages/AudioBlock.java @@ -0,0 +1,54 @@ +/* + * 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.chat.messages; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import javax.annotation.Nullable; + +/** The audio content of a {@link ChatMessage} — see {@link MediaBlock} for the media shape. */ +public final class AudioBlock extends MediaBlock { + + @JsonCreator + public AudioBlock( + @JsonProperty("media_type") String mediaType, + @JsonProperty("data") @Nullable String data, + @JsonProperty("url") @Nullable String url, + @JsonProperty("name") @Nullable String name, + @JsonProperty("size_bytes") @Nullable Long sizeBytes, + @JsonProperty("sha256") @Nullable String sha256) { + super(mediaType, data, url, name, sizeBytes, sha256); + } + + /** Creates an audio block carrying an inline base64 payload. */ + public static AudioBlock fromBase64(String mediaType, String data) { + return new AudioBlock(mediaType, data, null, null, null, null); + } + + /** Creates an audio block referencing an externally managed URL or provider file URI. */ + public static AudioBlock fromUrl(String mediaType, String url) { + return new AudioBlock(mediaType, null, url, null, null, null); + } + + @Override + public String getType() { + return "audio"; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/messages/ChatMessage.java b/api/src/main/java/org/apache/flink/agents/api/chat/messages/ChatMessage.java index a15f220ce..24e009323 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/messages/ChatMessage.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/messages/ChatMessage.java @@ -20,21 +20,36 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.stream.Collectors; /** * Chat message class that represents all message types (user, system, assistant, tool) with - * different roles + * different roles. + * + *

Message content is an ordered list of typed {@link ContentBlock}s ({@link TextBlock} plus the + * media blocks); a text-only message simply carries one {@link TextBlock}. The string convenience + * constructors and factories preserve the text-message experience, and {@link #getText()} is the + * ordered concatenation of the text blocks. + * + *

Blocks are immutable and the message snapshots every block list it is handed, so {@link + * #getBlocks()} is an unmodifiable view: the content changes only by replacing it through {@link + * #setBlocks(List)} or {@link #setText(String)}. */ public class ChatMessage { + private static final ObjectMapper MAPPER = new ObjectMapper(); + private MessageRole role; - private String content; + private List blocks; @JsonProperty("tool_calls") private List> toolCalls; @@ -44,34 +59,61 @@ public class ChatMessage { /** Default constructor with SYSTEM role */ public ChatMessage() { - this(MessageRole.SYSTEM, null, null, null); + this(MessageRole.SYSTEM, (List) null, null, null); + } + + /** Constructor with role and text content */ + public ChatMessage(MessageRole role, String text) { + this(role, blocksOf(text), null, null); + } + + /** Constructor with role and content blocks */ + public ChatMessage(MessageRole role, List blocks) { + this(role, blocks, null, null); } - /** Constructor with role and content */ - public ChatMessage(MessageRole role, String content) { - this(role, content, null, null); + public ChatMessage(MessageRole role, String text, Map extraArgs) { + this(role, blocksOf(text), null, extraArgs); } - public ChatMessage(MessageRole role, String content, Map extraArgs) { - this(role, content, null, extraArgs); + public ChatMessage(MessageRole role, String text, List> toolCalls) { + this(role, blocksOf(text), toolCalls, null); } - public ChatMessage(MessageRole role, String content, List> toolCalls) { - this(role, content, toolCalls, null); + public ChatMessage( + MessageRole role, + String text, + List> toolCalls, + Map extraArgs) { + this(role, blocksOf(text), toolCalls, extraArgs); } /** Full constructor */ public ChatMessage( MessageRole role, - String content, + List blocks, List> toolCalls, Map extraArgs) { this.role = role != null ? role : MessageRole.SYSTEM; - this.content = content != null ? content : ""; + this.blocks = snapshotOf(blocks); this.toolCalls = toolCalls != null ? toolCalls : new ArrayList<>(); this.extraArgs = extraArgs != null ? new HashMap<>(extraArgs) : new HashMap<>(); } + /** An empty or null text becomes an empty block list rather than an empty text block. */ + private static List blocksOf(String text) { + return text == null || text.isEmpty() + ? Collections.emptyList() + : Collections.singletonList(new TextBlock(text)); + } + + /** An unmodifiable copy — since blocks are immutable, this freezes the content. */ + private static List snapshotOf(List blocks) { + return blocks == null || blocks.isEmpty() + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(blocks)); + } + public MessageRole getRole() { return role; } @@ -80,12 +122,19 @@ public void setRole(MessageRole role) { this.role = role; } - public String getContent() { - return content; + /** The content as an unmodifiable list of immutable blocks. */ + public List getBlocks() { + return blocks; + } + + public void setBlocks(List blocks) { + this.blocks = snapshotOf(blocks); } - public void setContent(String content) { - this.content = content; + /** Replaces the content with a single text block (empty text clears the content). */ + @JsonIgnore + public void setText(String text) { + this.blocks = blocksOf(text); } @JsonProperty("tool_calls") @@ -108,9 +157,40 @@ public void setExtraArgs(Map extraArgs) { this.extraArgs = extraArgs != null ? extraArgs : new HashMap<>(); } + /** + * The content blocks as plain maps in the serialized (snake_case, discriminated) shape — the + * same representation {@code tool_calls} uses. This is how blocks cross the Python bridge, + * which exchanges JSON-friendly lists and maps rather than typed Java objects. + */ + @JsonIgnore + public List> getBlocksAsMaps() { + return blocks.stream() + .map( + block -> + MAPPER.>convertValue( + block, new TypeReference>() {})) + .collect(Collectors.toList()); + } + + /** Replaces the content with blocks given as plain maps — see {@link #getBlocksAsMaps()}. */ + @JsonIgnore + public void setBlocksFromMaps(List> blockMaps) { + this.blocks = + blockMaps == null + ? Collections.emptyList() + : Collections.unmodifiableList( + blockMaps.stream() + .map(map -> MAPPER.convertValue(map, ContentBlock.class)) + .collect(Collectors.toList())); + } + + /** The text projection: the ordered concatenation of this message's {@link TextBlock}s. */ @JsonIgnore public String getText() { - return this.content; + return blocks.stream() + .filter(block -> block instanceof TextBlock) + .map(block -> ((TextBlock) block).getText()) + .collect(Collectors.joining()); } @JsonIgnore @@ -124,24 +204,32 @@ public MessageRole getMessageType() { } // Static factory methods for convenience - public static ChatMessage user(String content) { - return new ChatMessage(MessageRole.USER, content); + public static ChatMessage user(String text) { + return new ChatMessage(MessageRole.USER, text); + } + + public static ChatMessage user(List blocks) { + return new ChatMessage(MessageRole.USER, blocks); + } + + public static ChatMessage system(String text) { + return new ChatMessage(MessageRole.SYSTEM, text); } - public static ChatMessage system(String content) { - return new ChatMessage(MessageRole.SYSTEM, content); + public static ChatMessage assistant(String text) { + return new ChatMessage(MessageRole.ASSISTANT, text); } - public static ChatMessage assistant(String content) { - return new ChatMessage(MessageRole.ASSISTANT, content); + public static ChatMessage assistant(String text, List> toolCalls) { + return new ChatMessage(MessageRole.ASSISTANT, text, toolCalls, new HashMap<>()); } - public static ChatMessage assistant(String content, List> toolCalls) { - return new ChatMessage(MessageRole.ASSISTANT, content, toolCalls, new HashMap<>()); + public static ChatMessage tool(String text) { + return new ChatMessage(MessageRole.TOOL, text); } - public static ChatMessage tool(String content) { - return new ChatMessage(MessageRole.TOOL, content); + public static ChatMessage tool(List blocks) { + return new ChatMessage(MessageRole.TOOL, blocks); } @Override @@ -150,19 +238,19 @@ public boolean equals(Object o) { if (!(o instanceof ChatMessage)) return false; ChatMessage that = (ChatMessage) o; return Objects.equals(role, that.role) - && Objects.equals(content, that.content) + && Objects.equals(blocks, that.blocks) && Objects.equals(toolCalls, that.toolCalls) && Objects.equals(extraArgs, that.extraArgs); } @Override public int hashCode() { - return Objects.hash(role, content, toolCalls, extraArgs); + return Objects.hash(role, blocks, toolCalls, extraArgs); } @Override public String toString() { - return role.getValue() + ": " + content; + return role.getValue() + ": " + getText(); } /** Return the index of the first system message in the list, or -1 if none. */ diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/messages/ContentBlock.java b/api/src/main/java/org/apache/flink/agents/api/chat/messages/ContentBlock.java new file mode 100644 index 000000000..c1cc626ed --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/messages/ContentBlock.java @@ -0,0 +1,66 @@ +/* + * 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.chat.messages; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonSubTypes; +import com.fasterxml.jackson.annotation.JsonTypeInfo; + +import java.util.Map; + +/** + * A single, typed part of a {@link ChatMessage}'s content. + * + *

Blocks are ordered within a message and are immutable value objects: every construction path, + * including Jackson deserialization, runs the same validation, so sharing a block instance never + * shares mutable state. The concrete type answers how providers route the content ({@link + * TextBlock}, {@link ImageBlock}, {@link AudioBlock}, {@link VideoBlock}, {@link DocumentBlock}), + * while media encoding is carried by the media type on {@link MediaBlock}. + * + *

The serialized form carries a {@code type} discriminator with fixed values ({@code text}, + * {@code image}, {@code audio}, {@code video}, {@code document}) shared with the Python API, so + * blocks cross the Java/Python boundary as plain JSON. + */ +@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type") +@JsonSubTypes({ + @JsonSubTypes.Type(value = TextBlock.class, name = "text"), + @JsonSubTypes.Type(value = ImageBlock.class, name = "image"), + @JsonSubTypes.Type(value = AudioBlock.class, name = "audio"), + @JsonSubTypes.Type(value = VideoBlock.class, name = "video"), + @JsonSubTypes.Type(value = DocumentBlock.class, name = "document") +}) +public abstract class ContentBlock { + + /** The wire discriminator of this block: {@code text}, {@code image}, {@code audio}, ... */ + @JsonIgnore + public abstract String getType(); + + /** + * The log-safe projection of this block, as a plain map in the wire's snake_case shape. Each + * block type defines its own logging policy: text passes through unchanged (the Event Log's + * level-dependent truncation still applies downstream), while media blocks whitelist their + * metadata, omit inline payload bytes, and sanitize URLs. Normal Jackson serialization — the + * Java/Python bridge, event serialization, state recovery — is unaffected and preserves the + * complete payload. + * + *

The result is intentionally not a valid wire block (media payloads are gone for good), so + * it must never be deserialized back into a {@link ContentBlock}. + */ + public abstract Map sanitize(); +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/messages/DocumentBlock.java b/api/src/main/java/org/apache/flink/agents/api/chat/messages/DocumentBlock.java new file mode 100644 index 000000000..5955e9bae --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/messages/DocumentBlock.java @@ -0,0 +1,54 @@ +/* + * 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.chat.messages; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import javax.annotation.Nullable; + +/** The document content of a {@link ChatMessage} — see {@link MediaBlock} for the media shape. */ +public final class DocumentBlock extends MediaBlock { + + @JsonCreator + public DocumentBlock( + @JsonProperty("media_type") String mediaType, + @JsonProperty("data") @Nullable String data, + @JsonProperty("url") @Nullable String url, + @JsonProperty("name") @Nullable String name, + @JsonProperty("size_bytes") @Nullable Long sizeBytes, + @JsonProperty("sha256") @Nullable String sha256) { + super(mediaType, data, url, name, sizeBytes, sha256); + } + + /** Creates a document block carrying an inline base64 payload. */ + public static DocumentBlock fromBase64(String mediaType, String data) { + return new DocumentBlock(mediaType, data, null, null, null, null); + } + + /** Creates a document block referencing an externally managed URL or provider file URI. */ + public static DocumentBlock fromUrl(String mediaType, String url) { + return new DocumentBlock(mediaType, null, url, null, null, null); + } + + @Override + public String getType() { + return "document"; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/messages/ImageBlock.java b/api/src/main/java/org/apache/flink/agents/api/chat/messages/ImageBlock.java new file mode 100644 index 000000000..1f3dcd061 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/messages/ImageBlock.java @@ -0,0 +1,54 @@ +/* + * 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.chat.messages; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import javax.annotation.Nullable; + +/** The image content of a {@link ChatMessage} — see {@link MediaBlock} for the media shape. */ +public final class ImageBlock extends MediaBlock { + + @JsonCreator + public ImageBlock( + @JsonProperty("media_type") String mediaType, + @JsonProperty("data") @Nullable String data, + @JsonProperty("url") @Nullable String url, + @JsonProperty("name") @Nullable String name, + @JsonProperty("size_bytes") @Nullable Long sizeBytes, + @JsonProperty("sha256") @Nullable String sha256) { + super(mediaType, data, url, name, sizeBytes, sha256); + } + + /** Creates an image block carrying an inline base64 payload. */ + public static ImageBlock fromBase64(String mediaType, String data) { + return new ImageBlock(mediaType, data, null, null, null, null); + } + + /** Creates an image block referencing an externally managed URL or provider file URI. */ + public static ImageBlock fromUrl(String mediaType, String url) { + return new ImageBlock(mediaType, null, url, null, null, null); + } + + @Override + public String getType() { + return "image"; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/messages/MediaBlock.java b/api/src/main/java/org/apache/flink/agents/api/chat/messages/MediaBlock.java new file mode 100644 index 000000000..811eb6ee3 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/messages/MediaBlock.java @@ -0,0 +1,207 @@ +/* + * 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.chat.messages; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +import javax.annotation.Nullable; + +import java.net.URI; +import java.net.URISyntaxException; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; + +/** + * Shared shape for binary media blocks: modality is the concrete type, encoding is the media type + * (RFC 6838; historically called a MIME type). + * + *

Media blocks are immutable, and every construction path — the {@code fromBase64}/{@code + * fromUrl} factories, the full constructors, and Jackson deserialization — runs the same + * validation, so a block that exists carries exactly one of base64 {@code data} or an externally + * managed {@code url}. URL-backed content is externally managed: URLs may expire, may not be + * reachable by the model provider, and may be invalid after recovery from a checkpoint. + * + *

The optional {@code name}/{@code sizeBytes}/{@code sha256} metadata also serves the Event Log, + * which records media metadata instead of payload bytes — see {@link #sanitize()}. + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public abstract class MediaBlock extends ContentBlock { + + @JsonProperty("media_type") + private final String mediaType; + + @Nullable private final String data; + + @Nullable private final String url; + + @Nullable private final String name; + + @JsonProperty("size_bytes") + @Nullable + private final Long sizeBytes; + + @Nullable private final String sha256; + + protected MediaBlock( + String mediaType, + @Nullable String data, + @Nullable String url, + @Nullable String name, + @Nullable Long sizeBytes, + @Nullable String sha256) { + if (mediaType == null || mediaType.isEmpty()) { + throw new IllegalArgumentException("A media block requires a media type."); + } + if ((data == null) == (url == null)) { + throw new IllegalArgumentException( + "A media block carries exactly one of base64 data or a URL."); + } + this.mediaType = mediaType; + this.data = data; + this.url = url; + this.name = name; + this.sizeBytes = sizeBytes; + this.sha256 = sha256; + } + + @JsonProperty("media_type") + public String getMediaType() { + return mediaType; + } + + @Nullable + public String getData() { + return data; + } + + @Nullable + public String getUrl() { + return url; + } + + @Nullable + public String getName() { + return name; + } + + @JsonProperty("size_bytes") + @Nullable + public Long getSizeBytes() { + return sizeBytes; + } + + @Nullable + public String getSha256() { + return sha256; + } + + /** + * The metadata-only Event Log projection: type, media type, and the optional {@code + * name}/{@code size_bytes}/{@code sha256}. Inline {@code data} is dropped entirely (with {@code + * size_bytes} derived from the base64 length when not stored), and URLs are stripped of + * credentials, query, and fragment. Dropping {@code data} rather than masking it means an + * attempt to reconstruct a block from logged output fails loudly instead of yielding a block + * with a fake payload. + */ + @Override + public final Map sanitize() { + Map safe = new LinkedHashMap<>(); + safe.put("type", getType()); + safe.put("media_type", mediaType); + if (name != null) { + safe.put("name", name); + } + Long size = sizeBytes != null ? sizeBytes : decodedSizeOf(data); + if (size != null) { + safe.put("size_bytes", size); + } + if (sha256 != null) { + safe.put("sha256", sha256); + } + if (url != null) { + safe.put("url", sanitizeUrl(url)); + } + return safe; + } + + /** The decoded byte count implied by a base64 payload, or null when there is none. */ + @Nullable + private static Long decodedSizeOf(@Nullable String data) { + if (data == null || data.isEmpty()) { + return null; + } + long padding = data.endsWith("==") ? 2 : data.endsWith("=") ? 1 : 0; + return data.length() * 3L / 4 - padding; + } + + /** + * Reduces a URL to scheme, host, port, and path. Userinfo (credentials), query strings (signed + * URLs carry their tokens there), and fragments never reach the log; a URL that cannot be + * parsed is replaced entirely rather than logged raw. + */ + private static String sanitizeUrl(String url) { + try { + URI uri = new URI(url); + if (uri.isOpaque()) { + return uri.getScheme() + ":"; + } + return new URI( + uri.getScheme(), + null, + uri.getHost(), + uri.getPort(), + uri.getPath(), + null, + null) + .toString(); + } catch (URISyntaxException e) { + return ""; + } + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + MediaBlock that = (MediaBlock) o; + return Objects.equals(mediaType, that.mediaType) + && Objects.equals(data, that.data) + && Objects.equals(url, that.url) + && Objects.equals(name, that.name) + && Objects.equals(sizeBytes, that.sizeBytes) + && Objects.equals(sha256, that.sha256); + } + + @Override + public int hashCode() { + return Objects.hash(mediaType, data, url, name, sizeBytes, sha256); + } + + @Override + public String toString() { + return getClass().getSimpleName() + + "(" + + mediaType + + ", " + + (data != null ? "inline" : "url=" + url) + + ")"; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/messages/TextBlock.java b/api/src/main/java/org/apache/flink/agents/api/chat/messages/TextBlock.java new file mode 100644 index 000000000..b0eb78994 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/messages/TextBlock.java @@ -0,0 +1,75 @@ +/* + * 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.chat.messages; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; + +/** A plain-text, immutable part of a {@link ChatMessage}. */ +public final class TextBlock extends ContentBlock { + + private final String text; + + @JsonCreator + public TextBlock(@JsonProperty("text") String text) { + this.text = text != null ? text : ""; + } + + public static TextBlock of(String text) { + return new TextBlock(text); + } + + public String getText() { + return text; + } + + @Override + public String getType() { + return "text"; + } + + @Override + public Map sanitize() { + Map safe = new LinkedHashMap<>(); + safe.put("type", getType()); + safe.put("text", text); + return safe; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof TextBlock)) return false; + return Objects.equals(text, ((TextBlock) o).text); + } + + @Override + public int hashCode() { + return Objects.hash(text); + } + + @Override + public String toString() { + return text; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/messages/VideoBlock.java b/api/src/main/java/org/apache/flink/agents/api/chat/messages/VideoBlock.java new file mode 100644 index 000000000..0a87e12dd --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/messages/VideoBlock.java @@ -0,0 +1,54 @@ +/* + * 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.chat.messages; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import javax.annotation.Nullable; + +/** The video content of a {@link ChatMessage} — see {@link MediaBlock} for the media shape. */ +public final class VideoBlock extends MediaBlock { + + @JsonCreator + public VideoBlock( + @JsonProperty("media_type") String mediaType, + @JsonProperty("data") @Nullable String data, + @JsonProperty("url") @Nullable String url, + @JsonProperty("name") @Nullable String name, + @JsonProperty("size_bytes") @Nullable Long sizeBytes, + @JsonProperty("sha256") @Nullable String sha256) { + super(mediaType, data, url, name, sizeBytes, sha256); + } + + /** Creates a video block carrying an inline base64 payload. */ + public static VideoBlock fromBase64(String mediaType, String data) { + return new VideoBlock(mediaType, data, null, null, null, null); + } + + /** Creates a video block referencing an externally managed URL or provider file URI. */ + public static VideoBlock fromUrl(String mediaType, String url) { + return new VideoBlock(mediaType, null, url, null, null, null); + } + + @Override + public String getType() { + return "video"; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java index 3cb2e655b..1cbf05b4a 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java @@ -160,11 +160,10 @@ public ChatMessage chat( } } - // append meaningful messages + // append meaningful messages; any block counts, so image-only messages survive List promptMessages = prompt.formatMessages(MessageRole.USER, stringified); for (ChatMessage message : messages) { - if ((message.getContent() != null && !message.getContent().isEmpty()) - || message.getRole() == MessageRole.ASSISTANT) { + if (!message.getBlocks().isEmpty() || message.getRole() == MessageRole.ASSISTANT) { promptMessages.add(message); } } diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java index 67a1688a3..326516c54 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java @@ -38,10 +38,11 @@ * *

The isolation boundary is deliberate and one level deep: the message list, each message's * tool-call maps and extra args, and the prompt-args map are defensive copies, but values - * nested inside those maps are shared with the request that is actually sent. Strategies - * must treat the context as read-only; the copies exist to make accidental top-level mutation - * harmless, not to sandbox a hostile strategy (arbitrary-depth copies on every routing decision - * would tax the common case to guard a case the SPI already forbids). + * nested inside those maps are shared with the request that is actually sent. Content + * blocks are shared too, which is safe because they are immutable value objects. Strategies must + * treat the context as read-only; the copies exist to make accidental top-level mutation harmless, + * not to sandbox a hostile strategy (arbitrary-depth copies on every routing decision would tax the + * common case to guard a case the SPI already forbids). */ public final class RoutingContext { @@ -59,9 +60,10 @@ public RoutingContext( List candidates) { this.requestId = requestId; this.router = router; - // Deep copy: the wrapping list is unmodifiable, but ChatMessage is mutable and the - // caller passes the same instances that go to the model — a strategy calling - // setContent(...) on a shallow copy would silently rewrite the prompt actually sent. + // Deep copy: ChatMessage is mutable and the caller passes the same instances that go to + // the model — a strategy calling setText(...) on a shallow copy would silently rewrite + // the prompt actually sent. Only the mutable message/list structure needs copying; + // ContentBlocks are immutable and safe to share. this.messages = messages == null ? Collections.emptyList() @@ -97,7 +99,9 @@ private static List deepCopy(List messages) { toolCalls.add(call == null ? null : new HashMap<>(call)); } } - copy.add(new ChatMessage(m.getRole(), m.getContent(), toolCalls, m.getExtraArgs())); + // The full constructor snapshots the block list; block instances are shared, which + // is safe because ContentBlocks are immutable. + copy.add(new ChatMessage(m.getRole(), m.getBlocks(), toolCalls, m.getExtraArgs())); } return copy; } @@ -136,7 +140,7 @@ public List getCandidates() { public String firstUserMessage() { for (ChatMessage message : messages) { if (message.getRole() == MessageRole.USER) { - return message.getContent() == null ? "" : message.getContent(); + return message.getText(); } } return ""; @@ -151,7 +155,7 @@ public String lastUserMessage() { for (int i = messages.size() - 1; i >= 0; i--) { ChatMessage message = messages.get(i); if (message.getRole() == MessageRole.USER) { - return message.getContent() == null ? "" : message.getContent(); + return message.getText(); } } return ""; diff --git a/api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java b/api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java index 0dad24238..7d2514df4 100644 --- a/api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java +++ b/api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java @@ -24,7 +24,9 @@ import com.fasterxml.jackson.annotation.JsonSubTypes; import com.fasterxml.jackson.annotation.JsonTypeInfo; import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ContentBlock; import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.messages.TextBlock; import org.apache.flink.agents.api.resource.ResourceType; import org.apache.flink.agents.api.resource.SerializableResource; @@ -161,7 +163,7 @@ public String formatString(Map kwargs) { messages -> { List formattedMessages = new ArrayList<>(); for (ChatMessage message : messages) { - String formattedContent = format(message.getContent(), kwargs); + String formattedContent = format(message.getText(), kwargs); String formatted = message.getRole().getValue() + ": " + formattedContent; formattedMessages.add(formatted); @@ -186,10 +188,24 @@ public List formatMessages( message -> new ChatMessage( message.getRole(), - format(message.getContent(), kwargs))) + formatBlocks( + message.getBlocks(), kwargs))) .collect(Collectors.toList())); } + /** Placeholder substitution applies to text blocks; media blocks pass through as-is. */ + private List formatBlocks( + List blocks, Map kwargs) { + List formatted = new ArrayList<>(blocks.size()); + for (ContentBlock block : blocks) { + formatted.add( + block instanceof TextBlock + ? new TextBlock(format(((TextBlock) block).getText(), kwargs)) + : block); + } + return formatted; + } + @Override public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) return false; 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 b09e6a13a..f80a09117 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 @@ -217,7 +217,7 @@ void javaCanDeserializeChatRequestEventFromPythonSnapshot() throws Exception { assertEquals(1, typed.getMessages().size(), "Expected one message."); ChatMessage msg = typed.getMessages().get(0); assertEquals(MessageRole.USER, msg.getRole(), "Role mismatch on Python-produced message."); - assertEquals("hello world", msg.getContent()); + assertEquals("hello world", msg.getText()); } /** @@ -281,7 +281,7 @@ void javaCanDeserializeChatResponseEventFromPythonSnapshot() throws Exception { ChatMessage response = typed.getResponse(); assertNotNull(response, "response field is null."); assertEquals(MessageRole.ASSISTANT, response.getRole(), "Role mismatch on response."); - assertEquals("hi there", response.getContent()); + assertEquals("hi there", response.getText()); } // ── ToolRequestEvent ─────────────────────────────────────────────────── diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/messages/ChatMessageSerializationTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/messages/ChatMessageSerializationTest.java new file mode 100644 index 000000000..e3d4714a6 --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/chat/messages/ChatMessageSerializationTest.java @@ -0,0 +1,211 @@ +/* + * 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.chat.messages; + +import com.fasterxml.jackson.databind.JsonMappingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Jackson round-trip tests for {@link ChatMessage} content blocks — the wire contract shared with + * the Python API (see the cross-language snapshot tests for the full event-level contract) — plus + * the block-level immutability and {@link ContentBlock#sanitize()} contracts. + */ +class ChatMessageSerializationTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @Test + @DisplayName("A text-only message serializes to a single typed text block") + void testTextOnlyWireShape() throws Exception { + ChatMessage message = ChatMessage.user("hello world"); + + JsonNode json = MAPPER.valueToTree(message); + + assertThat(json.get("role").asText()).isEqualTo("user"); + assertThat(json.get("blocks")).hasSize(1); + assertThat(json.get("blocks").get(0).get("type").asText()).isEqualTo("text"); + assertThat(json.get("blocks").get(0).get("text").asText()).isEqualTo("hello world"); + assertThat(json.get("tool_calls")).isEmpty(); + assertThat(json.get("extra_args")).isEmpty(); + assertThat(json.has("content")).isFalse(); + } + + @Test + @DisplayName("Media blocks carry the snake_case discriminated shape and omit absent fields") + void testMediaBlockWireShape() throws Exception { + ChatMessage message = + ChatMessage.user( + List.of( + TextBlock.of("What's in this picture?"), + ImageBlock.fromBase64("image/png", "aGk="))); + + JsonNode image = MAPPER.valueToTree(message).get("blocks").get(1); + + // The discriminator is written once: getType() is @JsonIgnore, the resolver owns the key. + assertThat(MAPPER.writeValueAsString(message.getBlocks().get(1))) + .containsOnlyOnce("\"type\""); + assertThat(image.get("type").asText()).isEqualTo("image"); + assertThat(image.get("media_type").asText()).isEqualTo("image/png"); + assertThat(image.get("data").asText()).isEqualTo("aGk="); + // Absent optional fields are omitted, not serialized as nulls. + assertThat(image.has("url")).isFalse(); + assertThat(image.has("name")).isFalse(); + assertThat(image.has("size_bytes")).isFalse(); + assertThat(image.has("sha256")).isFalse(); + } + + @Test + @DisplayName("A mixed-block message round-trips through Jackson preserving order and types") + void testMixedBlocksRoundTrip() throws Exception { + ImageBlock image = + new ImageBlock( + "image/jpeg", null, "https://example.org/cat.jpg", "cat.jpg", 123L, null); + ChatMessage original = + new ChatMessage( + MessageRole.TOOL, + List.of( + TextBlock.of("before"), + image, + DocumentBlock.fromBase64("application/pdf", "cGRm"), + TextBlock.of("after"))); + + ChatMessage restored = + MAPPER.readValue(MAPPER.writeValueAsString(original), ChatMessage.class); + + assertThat(restored).isEqualTo(original); + assertThat(restored.getBlocks()) + .extracting(block -> block.getClass().getSimpleName()) + .containsExactly("TextBlock", "ImageBlock", "DocumentBlock", "TextBlock"); + assertThat(restored.getText()).isEqualTo("beforeafter"); + } + + @Test + @DisplayName("Audio and video blocks round-trip through the same discriminator") + void testAudioAndVideoRoundTrip() throws Exception { + ChatMessage original = + ChatMessage.user( + List.of( + AudioBlock.fromBase64("audio/wav", "d2F2"), + VideoBlock.fromUrl("video/mp4", "https://example.org/v.mp4"))); + + ChatMessage restored = + MAPPER.readValue(MAPPER.writeValueAsString(original), ChatMessage.class); + + assertThat(restored).isEqualTo(original); + } + + @Test + @DisplayName("Media factories enforce exactly one of data and url") + void testMediaSourceValidation() { + assertThatThrownBy(() -> ImageBlock.fromBase64("image/png", null)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> ImageBlock.fromUrl("image/png", null)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> ImageBlock.fromBase64(null, "aGk=")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + @DisplayName("The Jackson path runs the same validation as the factories") + void testJacksonPathValidation() { + // Both sources present. + assertThatThrownBy( + () -> + MAPPER.readValue( + "{\"type\":\"image\",\"media_type\":\"image/png\"," + + "\"data\":\"aGk=\",\"url\":\"https://example.org/x\"}", + ContentBlock.class)) + .isInstanceOf(JsonMappingException.class) + .hasRootCauseInstanceOf(IllegalArgumentException.class); + // Neither source present. + assertThatThrownBy( + () -> + MAPPER.readValue( + "{\"type\":\"image\",\"media_type\":\"image/png\"}", + ContentBlock.class)) + .isInstanceOf(JsonMappingException.class) + .hasRootCauseInstanceOf(IllegalArgumentException.class); + // Missing media type. + assertThatThrownBy( + () -> + MAPPER.readValue( + "{\"type\":\"image\",\"data\":\"aGk=\"}", + ContentBlock.class)) + .isInstanceOf(JsonMappingException.class) + .hasRootCauseInstanceOf(IllegalArgumentException.class); + } + + @Test + @DisplayName("A message's block list is an unmodifiable snapshot") + void testBlockListIsSnapshot() { + ChatMessage message = ChatMessage.user("hello"); + assertThatThrownBy(() -> message.getBlocks().add(TextBlock.of("injected"))) + .isInstanceOf(UnsupportedOperationException.class); + } + + @Test + @DisplayName("sanitize() keeps text and media metadata but never payload bytes") + void testSanitizeInlineMedia() { + Map text = TextBlock.of("hello").sanitize(); + assertThat(text).containsExactly(Map.entry("type", "text"), Map.entry("text", "hello")); + + // "aGVsbG8=" is 8 base64 chars with one padding char: 5 decoded bytes. + Map image = ImageBlock.fromBase64("image/png", "aGVsbG8=").sanitize(); + assertThat(image.get("type")).isEqualTo("image"); + assertThat(image.get("media_type")).isEqualTo("image/png"); + // The inline payload is dropped entirely, so a reconstruction attempt fails loudly + // instead of producing a block with a fake payload. + assertThat(image).doesNotContainKey("data"); + assertThat(image.get("size_bytes")).isEqualTo(5L); + assertThat(image).doesNotContainKey("url"); + + Map document = + new DocumentBlock("application/pdf", "cGRm", null, "a.pdf", 999L, "abc123") + .sanitize(); + // Stored metadata wins over the derived size, and the whitelist keeps name and sha256. + assertThat(document.get("size_bytes")).isEqualTo(999L); + assertThat(document.get("name")).isEqualTo("a.pdf"); + assertThat(document.get("sha256")).isEqualTo("abc123"); + assertThat(document).doesNotContainKey("data"); + } + + @Test + @DisplayName("sanitize() strips credentials, query, and fragment from URLs") + void testSanitizeUrls() { + Map signed = + ImageBlock.fromUrl( + "image/png", + "https://user:secret@example.org:8443/media/cat.png" + + "?X-Amz-Signature=abc&token=t#frag") + .sanitize(); + assertThat(signed.get("url")).isEqualTo("https://example.org:8443/media/cat.png"); + + Map unparseable = + ImageBlock.fromUrl("image/png", "http://exa mple.org/x").sanitize(); + assertThat(unparseable.get("url")).isEqualTo(""); + } +} diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/messages/ChatMessageTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/messages/ChatMessageTest.java index d58353ba8..531fda38f 100644 --- a/api/src/test/java/org/apache/flink/agents/api/chat/messages/ChatMessageTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/chat/messages/ChatMessageTest.java @@ -49,13 +49,13 @@ void setUp() { @DisplayName("Test ChatMessage creation with role and content") void testChatMessageCreation() { assertEquals(MessageRole.USER, userMessage.getRole()); - assertEquals("Hello, how are you?", userMessage.getContent()); + assertEquals("Hello, how are you?", userMessage.getText()); assertEquals(MessageRole.SYSTEM, systemMessage.getRole()); - assertEquals("You are a helpful assistant.", systemMessage.getContent()); + assertEquals("You are a helpful assistant.", systemMessage.getText()); assertEquals(MessageRole.ASSISTANT, assistantMessage.getRole()); - assertEquals("I'm doing well, thank you!", assistantMessage.getContent()); + assertEquals("I'm doing well, thank you!", assistantMessage.getText()); } @Test @@ -63,7 +63,7 @@ void testChatMessageCreation() { void testDefaultConstructor() { ChatMessage defaultMessage = new ChatMessage(); assertEquals(MessageRole.SYSTEM, defaultMessage.getRole()); - assertEquals("", defaultMessage.getContent()); + assertEquals("", defaultMessage.getText()); assertNotNull(defaultMessage.getToolCalls()); assertNotNull(defaultMessage.getExtraArgs()); } @@ -135,10 +135,10 @@ void testMessageRoles() { @DisplayName("Test ChatMessage content modification") void testContentModification() { ChatMessage message = new ChatMessage(MessageRole.USER, "Original content"); - assertEquals("Original content", message.getContent()); + assertEquals("Original content", message.getText()); - message.setContent("Modified content"); - assertEquals("Modified content", message.getContent()); + message.setText("Modified content"); + assertEquals("Modified content", message.getText()); } @Test diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSkillsTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSkillsTest.java index 0077ef7d5..487165748 100644 --- a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSkillsTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSkillsTest.java @@ -171,9 +171,9 @@ public List getSkillDirs(List skillNames) { // Expected: SYSTEM, SYSTEM(skill_prompt), USER assertEquals(3, connection.capturedMessages.size()); assertEquals(MessageRole.SYSTEM, connection.capturedMessages.get(0).getRole()); - assertEquals("you are an agent", connection.capturedMessages.get(0).getContent()); + assertEquals("you are an agent", connection.capturedMessages.get(0).getText()); assertEquals(MessageRole.SYSTEM, connection.capturedMessages.get(1).getRole()); - assertTrue(connection.capturedMessages.get(1).getContent().contains("")); + assertTrue(connection.capturedMessages.get(1).getText().contains("")); assertEquals(MessageRole.USER, connection.capturedMessages.get(2).getRole()); } } diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java index d27d79cb8..11604b7c8 100644 --- a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java @@ -71,7 +71,7 @@ public ChatMessage chat( String lastUserContent = ""; for (ChatMessage message : messages) { if (message.getRole() == MessageRole.USER) { - lastUserContent = message.getContent(); + lastUserContent = message.getText(); } } @@ -127,7 +127,7 @@ void testBasicChat() { assertNotNull(response); assertEquals(MessageRole.ASSISTANT, response.getRole()); - assertTrue(response.getContent().contains("Test Response:")); + assertTrue(response.getText().contains("Test Response:")); } @Test @@ -145,7 +145,7 @@ void testChatWithConversationPrompt() { assertNotNull(response); assertEquals(MessageRole.ASSISTANT, response.getRole()); - assertTrue(response.getContent().contains("What's the weather like?")); + assertTrue(response.getText().contains("What's the weather like?")); } @Test @@ -158,7 +158,7 @@ void testChatWithEmptyPrompt() { assertNotNull(response); assertEquals(MessageRole.ASSISTANT, response.getRole()); - assertTrue(response.getContent().contains("No user message found")); + assertTrue(response.getText().contains("No user message found")); } @Test @@ -178,7 +178,7 @@ void testChatWithMultipleUserMessages() { chatModel.chat(multiPrompt.formatMessages(MessageRole.USER, new HashMap<>())); assertNotNull(response); - assertTrue(response.getContent().contains("Second message - this should be the response")); + assertTrue(response.getText().contains("Second message - this should be the response")); } @Test @@ -195,7 +195,7 @@ void testChatModelConfiguration() { ChatMessage response = chatModel.chat(formattedPrompt.formatMessages(MessageRole.USER, new HashMap<>())); - assertTrue(response.getContent().startsWith("Custom Response:")); + assertTrue(response.getText().startsWith("Custom Response:")); } @Test @@ -211,7 +211,7 @@ void testChatWithSystemOnlyPrompt() { assertNotNull(response); assertEquals(MessageRole.ASSISTANT, response.getRole()); - assertTrue(response.getContent().contains("No user message found")); + assertTrue(response.getText().contains("No user message found")); } @Test @@ -228,10 +228,10 @@ void testChatResponseFormat() { // Verify response structure assertNotNull(response.getRole()); - assertNotNull(response.getContent()); + assertNotNull(response.getText()); assertNotNull(response.getToolCalls()); assertNotNull(response.getExtraArgs()); - assertTrue(response.getContent().length() > 0); + assertTrue(response.getText().length() > 0); } /** Connection that captures the messages passed to it for assertions. */ @@ -283,7 +283,7 @@ void testChatFillsTemplateFromPromptArgsParameter() { assertNotNull(connection.capturedMessages); assertEquals(1, connection.capturedMessages.size()); - assertEquals("Task: value", connection.capturedMessages.get(0).getContent()); + assertEquals("Task: value", connection.capturedMessages.get(0).getText()); } @Test @@ -299,8 +299,8 @@ void testChatDoesNotReadTemplateVarsFromExtraArgs() { assertNotNull(connection.capturedMessages); assertEquals(2, connection.capturedMessages.size()); - assertEquals("Task: {key}", connection.capturedMessages.get(0).getContent()); - assertEquals("hello", connection.capturedMessages.get(1).getContent()); + assertEquals("Task: {key}", connection.capturedMessages.get(0).getText()); + assertEquals("hello", connection.capturedMessages.get(1).getText()); } @Test @@ -313,13 +313,13 @@ void testChatRefillsTemplateOnSubsequentInvocations() { setup.chat(Collections.emptyList(), Map.of("key", "v1"), Map.of()); assertNotNull(connection.capturedMessages); assertEquals(1, connection.capturedMessages.size()); - assertEquals("Task: v1", connection.capturedMessages.get(0).getContent()); + assertEquals("Task: v1", connection.capturedMessages.get(0).getText()); ChatMessage toolResponse = new ChatMessage(MessageRole.TOOL, "tool result"); setup.chat(List.of(toolResponse), Map.of("key", "v1"), Map.of()); assertEquals(2, connection.capturedMessages.size()); - assertEquals("Task: v1", connection.capturedMessages.get(0).getContent()); - assertEquals("tool result", connection.capturedMessages.get(1).getContent()); + assertEquals("Task: v1", connection.capturedMessages.get(0).getText()); + assertEquals("tool result", connection.capturedMessages.get(1).getText()); } @Test @@ -359,7 +359,7 @@ void testDefaultChatOverloadDelegatesForNullOutputSchema() { // The 3-arg chat() ran (it is what produces "ok") and the overload added nothing // to modelParams that could travel on to a provider SDK request. - assertEquals("ok", response.getContent()); + assertEquals("ok", response.getText()); assertEquals(Map.of("temperature", 0.5), connection.capturedModelParams); } @@ -458,6 +458,6 @@ void testChatWithLongInput() { chatModel.chat(formattedPrompt.formatMessages(MessageRole.USER, new HashMap<>())); assertNotNull(response); - assertTrue(response.getContent().length() > 0); + assertTrue(response.getText().length() > 0); } } diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingTest.java index df4a60f07..67be6f687 100644 --- a/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingTest.java @@ -20,7 +20,9 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ImageBlock; import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.messages.TextBlock; import org.apache.flink.agents.api.event.ModelRoutingEvent; import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.junit.jupiter.api.Test; @@ -328,8 +330,34 @@ void routingContextMessagesAreDeepCopied() { new RoutingContext( UUID.randomUUID(), "router", List.of(original), Map.of(), List.of()); // A strategy mutating what it sees must not rewrite the message actually sent. - ctx.getMessages().get(0).setContent("REWRITTEN BY STRATEGY"); - assertEquals("original prompt", original.getContent()); + ctx.getMessages().get(0).setText("REWRITTEN BY STRATEGY"); + assertEquals("original prompt", original.getText()); + } + + @Test + void routingContextBlocksCannotBeMutatedThroughAliasing() { + // Blocks are shared between the context copy and the original message; the isolation + // relies on ContentBlock immutability plus the message's unmodifiable list snapshot, + // so a strategy can neither edit a shared block nor splice its own into the request. + ChatMessage original = + new ChatMessage( + MessageRole.USER, + List.of( + TextBlock.of("look at this"), + ImageBlock.fromUrl("image/png", "https://example.org/cat.png"))); + RoutingContext ctx = + new RoutingContext( + UUID.randomUUID(), "router", List.of(original), Map.of(), List.of()); + + List seenBlocks = ctx.getMessages().get(0).getBlocks(); + assertThrows( + UnsupportedOperationException.class, + () -> ctx.getMessages().get(0).getBlocks().add(TextBlock.of("INJECTED"))); + assertThrows(UnsupportedOperationException.class, () -> seenBlocks.remove(0)); + // Replacing the copied message's whole content stays local to the copy. + ctx.getMessages().get(0).setBlocks(List.of(TextBlock.of("REWRITTEN"))); + assertEquals("look at this", original.getText()); + assertEquals(2, original.getBlocks().size()); } @Test @@ -359,7 +387,7 @@ void routingContextToleratesNullToolCallsFromJsonSetter() { new RoutingContext( UUID.randomUUID(), "router", List.of(fromJson), Map.of(), List.of()); assertEquals(1, ctx.getMessages().size()); - assertEquals("hello", ctx.getMessages().get(0).getContent()); + assertEquals("hello", ctx.getMessages().get(0).getText()); // The copy re-normalizes through the constructor, so strategies see an empty list. assertEquals(0, ctx.getMessages().get(0).getToolCalls().size()); } diff --git a/api/src/test/java/org/apache/flink/agents/api/prompt/PromptTest.java b/api/src/test/java/org/apache/flink/agents/api/prompt/PromptTest.java index 08c465b40..b330e4b65 100644 --- a/api/src/test/java/org/apache/flink/agents/api/prompt/PromptTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/prompt/PromptTest.java @@ -95,9 +95,9 @@ void testTextPromptToMessages() { assertEquals(1, messages.size()); assertEquals(MessageRole.SYSTEM, messages.get(0).getRole()); - assertTrue(messages.get(0).getContent().contains("12345")); - assertTrue(messages.get(0).getContent().contains("wireless noise-canceling headphones")); - assertTrue(messages.get(0).getContent().contains("The headphones broke after one week")); + assertTrue(messages.get(0).getText().contains("12345")); + assertTrue(messages.get(0).getText().contains("wireless noise-canceling headphones")); + assertTrue(messages.get(0).getText().contains("The headphones broke after one week")); } @Test @@ -125,12 +125,12 @@ void testMessagesPromptToMessages() { assertEquals( "You are a product review analyzer, please generate a score and the dislike reasons " + "(if any) for the review.", - messages.get(0).getContent()); + messages.get(0).getText()); // Check user message with variable substitution assertEquals(MessageRole.USER, messages.get(1).getRole()); - assertTrue(messages.get(1).getContent().contains("12345")); - assertTrue(messages.get(1).getContent().contains("wireless noise-canceling headphones")); + assertTrue(messages.get(1).getText().contains("12345")); + assertTrue(messages.get(1).getText().contains("wireless noise-canceling headphones")); } @Test @@ -177,7 +177,7 @@ void testEmptyPrompt() { List messages = emptyPrompt.formatMessages(MessageRole.USER, new HashMap<>()); assertEquals(1, messages.size()); - assertEquals("", messages.get(0).getContent()); + assertEquals("", messages.get(0).getText()); } @Test @@ -233,9 +233,9 @@ void testComplexConversationPrompt() { conversationPrompt.formatMessages(MessageRole.SYSTEM, conversationVars); assertEquals(4, messages.size()); - assertTrue(messages.get(0).getContent().contains("an AI assistant")); - assertTrue(messages.get(0).getContent().contains("software development")); - assertTrue(messages.get(3).getContent().contains("NullPointerException")); + assertTrue(messages.get(0).getText().contains("an AI assistant")); + assertTrue(messages.get(0).getText().contains("software development")); + assertTrue(messages.get(3).getText().contains("NullPointerException")); } @Test @@ -289,8 +289,8 @@ void testFormatMessagesDoesNotReExpandValues() { List messages = prompt.formatMessages(MessageRole.SYSTEM, vars); assertEquals(2, messages.size()); - assertEquals("p@ssw0rd", messages.get(0).getContent()); - assertEquals("give me {secret}", messages.get(1).getContent()); + assertEquals("p@ssw0rd", messages.get(0).getText()); + assertEquals("give me {secret}", messages.get(1).getText()); } @Test diff --git a/docs/content/docs/development/chat_models.md b/docs/content/docs/development/chat_models.md index 9152136fc..5f62af453 100644 --- a/docs/content/docs/development/chat_models.md +++ b/docs/content/docs/development/chat_models.md @@ -85,10 +85,7 @@ class MyAgent(Agent): def process_input(event: Event, ctx: RunnerContext) -> None: input_event = InputEvent.from_event(event) # Create a chat request with user message - user_message = ChatMessage( - role=MessageRole.USER, - content=f"input: {input_event.input}" - ) + user_message = ChatMessage.user(f"input: {input_event.input}") ctx.send_event( ChatRequestEvent(model="ollama_chat_model", messages=[user_message]) ) @@ -97,7 +94,7 @@ class MyAgent(Agent): @staticmethod def process_response(event: Event, ctx: RunnerContext) -> None: chat_response = ChatResponseEvent.from_event(event) - response_content = chat_response.response.content + response_content = chat_response.response.text # Handle the LLM's response # Process the response as needed for your use case ``` @@ -133,7 +130,7 @@ public class MyAgent extends Agent { public static void processResponse(Event event, RunnerContext ctx) throws Exception { ChatResponseEvent chatResponse = ChatResponseEvent.fromEvent(event); - String response = chatResponse.getResponse().getContent(); + String response = chatResponse.getResponse().getText(); // Handle the LLM's response // Process the response as needed for your use case } @@ -1470,10 +1467,7 @@ class MyAgent(Agent): def process_input(event: Event, ctx: RunnerContext) -> None: input_event = InputEvent.from_event(event) # Create a chat request with user message - user_message = ChatMessage( - role=MessageRole.USER, - content=f"input: {input_event.input}" - ) + user_message = ChatMessage.user(f"input: {input_event.input}") ctx.send_event( ChatRequestEvent(model="java_chat_model", messages=[user_message]) ) @@ -1482,7 +1476,7 @@ class MyAgent(Agent): @staticmethod def process_response(event: Event, ctx: RunnerContext) -> None: chat_response = ChatResponseEvent.from_event(event) - response_content = chat_response.response.content + response_content = chat_response.response.text # Handle the LLM's response # Process the response as needed for your use case ``` @@ -1536,7 +1530,7 @@ public class MyAgent extends Agent { public static void processResponse(Event event, RunnerContext ctx) throws Exception { ChatResponseEvent chatResponse = ChatResponseEvent.fromEvent(event); - String response = chatResponse.getResponse().getContent(); + String response = chatResponse.getResponse().getText(); // Handle the LLM's response // Process the response as needed for your use case } diff --git a/docs/content/docs/development/prompts.md b/docs/content/docs/development/prompts.md index 8c57b5b0e..870a2d67e 100644 --- a/docs/content/docs/development/prompts.md +++ b/docs/content/docs/development/prompts.md @@ -121,9 +121,8 @@ For more control, create prompts from a sequence of `ChatMessage` objects using ```python review_analysis_prompt = Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content=""" + ChatMessage.system( + """ Analyze the user review and product information to determine a satisfaction score (1-5) and potential reasons for dissatisfaction. @@ -139,14 +138,13 @@ review_analysis_prompt = Prompt.from_messages( "score": 1, "reasons": ["poor quality"] } - """, + """ ), - ChatMessage( - role=MessageRole.USER, - content=""" + ChatMessage.user( + """ "input": {input} - """, + """ ), ], ) @@ -201,9 +199,8 @@ class ReviewAnalysisAgent(Agent): """Prompt for review analysis.""" return Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content=""" + ChatMessage.system( + """ Analyze the user review and product information to determine a satisfaction score (1-5) and potential reasons for dissatisfaction. @@ -219,14 +216,13 @@ class ReviewAnalysisAgent(Agent): "score": 1, "reasons": ["poor quality"] } - """, + """ ), - ChatMessage( - role=MessageRole.USER, - content=""" + ChatMessage.user( + """ "input": {input} - """, + """ ), ], ) diff --git a/docs/content/docs/development/react_agent.md b/docs/content/docs/development/react_agent.md index f660e8767..8bbd893c1 100644 --- a/docs/content/docs/development/react_agent.md +++ b/docs/content/docs/development/react_agent.md @@ -189,19 +189,15 @@ system_prompt_str = """ # Prompt for review analysis react agent. my_prompt = Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content=system_prompt_str, - ), + ChatMessage.system(system_prompt_str), # For react agent, if the input element is not primitive types, # framework will deserialize input element to dict and fill the prompt. # Note, the input element should be primitive types, BaseModel or Row. - ChatMessage( - role=MessageRole.USER, - content=""" + ChatMessage.user( + """ "id": {id}, "review": {review} - """, + """ ), ], ) @@ -244,10 +240,7 @@ If the input element is primitive types, like `int`, `str` and so on, the second {{< tab "Python" >}} ```python -ChatMessage( - role=MessageRole.USER, - content="{input}" -) +ChatMessage.user("{input}") ``` {{< /tab >}} diff --git a/docs/content/docs/development/workflow_agent.md b/docs/content/docs/development/workflow_agent.md index b8166cd0e..d95c37fdf 100644 --- a/docs/content/docs/development/workflow_agent.md +++ b/docs/content/docs/development/workflow_agent.md @@ -117,7 +117,7 @@ class ReviewAnalysisAgent(Agent): """Process chat response event and send output event.""" chat_response = ChatResponseEvent.from_event(event) try: - json_content = json.loads(chat_response.response.content) + json_content = json.loads(chat_response.response.text) ctx.send_event( OutputEvent( output=ProductReviewAnalysisRes( @@ -129,7 +129,7 @@ class ReviewAnalysisAgent(Agent): ) except Exception: logging.exception( - f"Error processing chat response {chat_response.response.content}" + f"Error processing chat response {chat_response.response.text}" ) # To fail the agent, you can raise an exception here. @@ -206,7 +206,7 @@ public class ReviewAnalysisAgent extends Agent { public static void processChatResponse(Event event, RunnerContext ctx) throws Exception { ChatResponseEvent chatResponse = ChatResponseEvent.fromEvent(event); - JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getContent()); + JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getText()); JsonNode scoreNode = jsonNode.findValue("score"); JsonNode reasonsNode = jsonNode.findValue("reasons"); if (scoreNode == null || reasonsNode == null) { @@ -426,7 +426,7 @@ access these framework variables: `attributes.score > 80` refer to the same field. - Nested values are not flattened. For `{input: {status: "ok"}}`, use `input.status` or `attributes.input.status`; bare `status` does not refer to the nested value. Other event payloads - keep their top-level envelope, for example `response.content`. + keep their top-level envelope, for example `response.blocks`. - Framework variables take precedence over attributes with the same names. Use `attributes["type"]` or `attributes["id"]` to access a colliding attribute. - For a top-level key containing dots, use a literal index such as `attributes["a.b.c"]`. Test its diff --git a/docs/content/docs/development/yaml.md b/docs/content/docs/development/yaml.md index a053873b4..7773a1a70 100644 --- a/docs/content/docs/development/yaml.md +++ b/docs/content/docs/development/yaml.md @@ -309,7 +309,7 @@ actions: - name: action2 function: my_pkg.actions:action2 trigger_conditions: - - "type == EventType.ChatResponseEvent && response.content != ''" + - "type == EventType.ChatResponseEvent && size(response.blocks) > 0" type: python - action3 # shared action reference (declared at file level) ``` diff --git a/docs/content/docs/get-started/quickstart/parallel_llm.md b/docs/content/docs/get-started/quickstart/parallel_llm.md index b42b61eb5..0f941a350 100644 --- a/docs/content/docs/get-started/quickstart/parallel_llm.md +++ b/docs/content/docs/get-started/quickstart/parallel_llm.md @@ -112,11 +112,8 @@ def _build_aspect_request(text: str, aspect: str) -> ChatRequestEvent: return ChatRequestEvent( model="sentiment_model", messages=[ - ChatMessage(role=MessageRole.SYSTEM, content=PARALLEL_SYSTEM_PROMPT), - ChatMessage( - role=MessageRole.USER, - content=f'Judge the "{aspect}" dimension: {text}', - ), + ChatMessage.system(PARALLEL_SYSTEM_PROMPT), + ChatMessage.user(f'Judge the "{aspect}" dimension: {text}'), ], output_schema=OutputSchema(output_schema=AspectResponse), ) @@ -132,8 +129,8 @@ def _build_summarize_request(text: str, sentiments: Dict[str, str]) -> ChatReque return ChatRequestEvent( model="sentiment_model", messages=[ - ChatMessage(role=MessageRole.SYSTEM, content=AGGREGATE_SYSTEM_PROMPT), - ChatMessage(role=MessageRole.USER, content=body), + ChatMessage.system(AGGREGATE_SYSTEM_PROMPT), + ChatMessage.user(body), ], output_schema=OutputSchema(output_schema=SummaryResponse), ) diff --git a/docs/content/docs/get-started/quickstart/skills_agent.md b/docs/content/docs/get-started/quickstart/skills_agent.md index a95a276a9..add62265b 100644 --- a/docs/content/docs/get-started/quickstart/skills_agent.md +++ b/docs/content/docs/get-started/quickstart/skills_agent.md @@ -82,12 +82,11 @@ class MathAgent(Agent): """System prompt instructing the agent to use the skill.""" return Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content="You are a helpful math assistant. Use the " + ChatMessage.system( + "You are a helpful math assistant. Use the " "math-calculator skill when asked to evaluate an expression. " "You must load the skill first and strictly follow its " - "instructions. Reply with only the final numeric result.", + "instructions. Reply with only the final numeric result." ) ], ) @@ -115,7 +114,7 @@ class MathAgent(Agent): ctx.send_event( ChatRequestEvent( model="math_model", - messages=[ChatMessage(role=MessageRole.USER, content=question)], + messages=[ChatMessage.user(question)], ) ) @@ -124,7 +123,7 @@ class MathAgent(Agent): def process_chat_response(event: Event, ctx: RunnerContext) -> None: """Process chat response event and send the answer as output.""" chat_response = ChatResponseEvent.from_event(event) - ctx.send_event(OutputEvent(output=chat_response.response.content)) + ctx.send_event(OutputEvent(output=chat_response.response.text)) ``` {{< /tab >}} @@ -179,7 +178,7 @@ public class MathAgent extends Agent { /** Process chat response event and send the answer as output. */ @Action(EventType.ChatResponseEvent) public static void processChatResponse(ChatResponseEvent event, RunnerContext ctx) { - ctx.sendEvent(new OutputEvent(event.getResponse().getContent())); + ctx.sendEvent(new OutputEvent(event.getResponse().getText())); } } ``` diff --git a/docs/content/docs/get-started/quickstart/workflow_agent.md b/docs/content/docs/get-started/quickstart/workflow_agent.md index 3e238475a..4e844fa88 100644 --- a/docs/content/docs/get-started/quickstart/workflow_agent.md +++ b/docs/content/docs/get-started/quickstart/workflow_agent.md @@ -154,7 +154,7 @@ class ReviewAnalysisAgent(Agent): """Process chat response event and send output event.""" chat_response = ChatResponseEvent.from_event(event) try: - json_content = json.loads(chat_response.response.content) + json_content = json.loads(chat_response.response.text) ctx.send_event( OutputEvent( output=ProductReviewAnalysisRes( @@ -166,7 +166,7 @@ class ReviewAnalysisAgent(Agent): ) except Exception: logging.exception( - f"Error processing chat response {chat_response.response.content}" + f"Error processing chat response {chat_response.response.text}" ) # To fail the agent, you can raise an exception here. @@ -243,7 +243,7 @@ public class ReviewAnalysisAgent extends Agent { public static void processChatResponse(Event event, RunnerContext ctx) throws Exception { ChatResponseEvent chatResponse = ChatResponseEvent.fromEvent(event); - JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getContent()); + JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getText()); JsonNode scoreNode = jsonNode.findValue("score"); JsonNode reasonsNode = jsonNode.findValue("reasons"); if (scoreNode == null || reasonsNode == null) { diff --git a/e2e-test/cross-language-event-snapshots/java/chat_request_event.json b/e2e-test/cross-language-event-snapshots/java/chat_request_event.json index 347c47e71..27cdb215c 100644 --- a/e2e-test/cross-language-event-snapshots/java/chat_request_event.json +++ b/e2e-test/cross-language-event-snapshots/java/chat_request_event.json @@ -4,7 +4,10 @@ "model" : "test-model", "messages" : [ { "role" : "user", - "content" : "hello world", + "blocks" : [ { + "type" : "text", + "text" : "hello world" + } ], "tool_calls" : [ ], "extra_args" : { } } ] diff --git a/e2e-test/cross-language-event-snapshots/java/chat_response_event.json b/e2e-test/cross-language-event-snapshots/java/chat_response_event.json index 3d5b4793c..a58b79dd8 100644 --- a/e2e-test/cross-language-event-snapshots/java/chat_response_event.json +++ b/e2e-test/cross-language-event-snapshots/java/chat_response_event.json @@ -4,7 +4,10 @@ "request_id" : "00000000-0000-0000-0000-000000000002", "response" : { "role" : "assistant", - "content" : "hi there", + "blocks" : [ { + "type" : "text", + "text" : "hi there" + } ], "tool_calls" : [ ], "extra_args" : { } }, diff --git a/e2e-test/cross-language-event-snapshots/python/chat_request_event.json b/e2e-test/cross-language-event-snapshots/python/chat_request_event.json index ac8808231..b9b1fb5a8 100644 --- a/e2e-test/cross-language-event-snapshots/python/chat_request_event.json +++ b/e2e-test/cross-language-event-snapshots/python/chat_request_event.json @@ -6,7 +6,12 @@ "messages": [ { "role": "user", - "content": "hello world", + "blocks": [ + { + "type": "text", + "text": "hello world" + } + ], "tool_calls": [], "extra_args": {} } diff --git a/e2e-test/cross-language-event-snapshots/python/chat_response_event.json b/e2e-test/cross-language-event-snapshots/python/chat_response_event.json index bafb28116..8f44219e4 100644 --- a/e2e-test/cross-language-event-snapshots/python/chat_response_event.json +++ b/e2e-test/cross-language-event-snapshots/python/chat_response_event.json @@ -5,7 +5,12 @@ "request_id": "00000000-0000-0000-0000-000000000002", "response": { "role": "assistant", - "content": "hi there", + "blocks": [ + { + "type": "text", + "text": "hi there" + } + ], "tool_calls": [], "extra_args": {} }, diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ChatModelIntegrationAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ChatModelIntegrationAgent.java index 2ad6bacfa..8b6553eb8 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ChatModelIntegrationAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ChatModelIntegrationAgent.java @@ -217,6 +217,6 @@ public static void process(Event event, RunnerContext ctx) throws Exception { @Action(EventType.ChatResponseEvent) public static void processChatResponse(Event event, RunnerContext ctx) { ChatResponseEvent chatResponse = ChatResponseEvent.fromEvent(event); - ctx.sendEvent(new OutputEvent(chatResponse.getResponse().getContent())); + ctx.sendEvent(new OutputEvent(chatResponse.getResponse().getText())); } } diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationAgent.java index 858f58d1a..824dd4832 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationAgent.java @@ -115,6 +115,6 @@ public static void process(InputEvent event, RunnerContext ctx) throws Exception @Action(EventType.ChatResponseEvent) public static void processChatResponse(ChatResponseEvent event, RunnerContext ctx) { - ctx.sendEvent(new OutputEvent(event.getResponse().getContent())); + ctx.sendEvent(new OutputEvent(event.getResponse().getText())); } } diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/TokenMetricsE2EAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/TokenMetricsE2EAgent.java index 22608a2c0..49559a7f2 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/TokenMetricsE2EAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/TokenMetricsE2EAgent.java @@ -72,6 +72,6 @@ public static void process(InputEvent event, RunnerContext ctx) throws Exception @Action(EventType.ChatResponseEvent) public static void processChatResponse(ChatResponseEvent event, RunnerContext ctx) { - ctx.sendEvent(new OutputEvent(event.getResponse().getContent())); + ctx.sendEvent(new OutputEvent(event.getResponse().getText())); } } diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ToolParameterInjectionAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ToolParameterInjectionAgent.java index 88d489ded..3088a48c2 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ToolParameterInjectionAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ToolParameterInjectionAgent.java @@ -68,7 +68,7 @@ public ChatMessage chat( Map modelParams) { ChatMessage lastMessage = messages.get(messages.size() - 1); if (lastMessage.getRole() == MessageRole.TOOL) { - return new ChatMessage(MessageRole.ASSISTANT, lastMessage.getContent()); + return new ChatMessage(MessageRole.ASSISTANT, lastMessage.getText()); } for (org.apache.flink.agents.api.tools.Tool tool : tools) { @@ -78,7 +78,7 @@ public ChatMessage chat( "Injected argument leaked into tool schema: " + inputSchema); } } - String orderId = lastMessage.getContent(); + String orderId = lastMessage.getText(); return new ChatMessage( MessageRole.ASSISTANT, "", @@ -143,6 +143,6 @@ public static void requestTool(Event event, RunnerContext ctx) { @Action(EventType.ChatResponseEvent) public static void emitResult(Event event, RunnerContext ctx) { ChatResponseEvent responseEvent = ChatResponseEvent.fromEvent(event); - ctx.sendEvent(new OutputEvent(responseEvent.getResponse().getContent())); + ctx.sendEvent(new OutputEvent(responseEvent.getResponse().getText())); } } diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/yaml/YamlChatActions.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/yaml/YamlChatActions.java index 524e0e0d6..574fbbbbe 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/yaml/YamlChatActions.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/yaml/YamlChatActions.java @@ -129,11 +129,11 @@ public static void chatRequest(Event event, RunnerContext ctx) throws Exception public static void processChatResponse(Event event, RunnerContext ctx) throws Exception { ChatResponseEvent chatResponse = ChatResponseEvent.fromEvent(event); ChatMessage response = chatResponse.getResponse(); - if (response == null || response.getContent() == null) { + if (response == null || response.getText() == null) { return; } Integer inputId = (Integer) ctx.getShortTermMemory().get("input_id").getValue(); - ctx.sendEvent(new OutputEvent(new YamlChatOutput(inputId, response.getContent()))); + ctx.sendEvent(new OutputEvent(new YamlChatOutput(inputId, response.getText()))); } /** diff --git a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ChatModelCrossLanguageAgent.java b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ChatModelCrossLanguageAgent.java index a2d78ecfb..b4f980223 100644 --- a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ChatModelCrossLanguageAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ChatModelCrossLanguageAgent.java @@ -170,6 +170,6 @@ public static void process(Event event, RunnerContext ctx) throws Exception { @Action(EventType.ChatResponseEvent) public static void processChatResponse(Event event, RunnerContext ctx) { ChatResponseEvent chatResponse = ChatResponseEvent.fromEvent(event); - ctx.sendEvent(new OutputEvent(chatResponse.getResponse().getContent())); + ctx.sendEvent(new OutputEvent(chatResponse.getResponse().getText())); } } diff --git a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/MCPCrossLanguageAgent.java b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/MCPCrossLanguageAgent.java index 30b921e8c..5505c07b1 100644 --- a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/MCPCrossLanguageAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/MCPCrossLanguageAgent.java @@ -66,8 +66,7 @@ public static void process(Event event, RunnerContext ctx) throws Exception { askSum.formatMessages(MessageRole.USER, Map.of("a", "1", "b", "2")); Assertions.assertEquals(1, chatMessages.size()); Assertions.assertEquals( - "Can you please calculate the sum of 1 and 2?", - chatMessages.get(0).getContent()); + "Can you please calculate the sum of 1 and 2?", chatMessages.get(0).getText()); Assertions.assertEquals(MessageRole.USER, chatMessages.get(0).getRole()); String content = askSum.formatString(Map.of("a", "3", "b", "4")); diff --git a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/YamlCrossLanguageActions.java b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/YamlCrossLanguageActions.java index 1dfd5109b..530622673 100644 --- a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/YamlCrossLanguageActions.java +++ b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/YamlCrossLanguageActions.java @@ -58,9 +58,9 @@ public static void processInput(Event event, RunnerContext ctx) throws Exception public static void processChatResponse(Event event, RunnerContext ctx) { ChatResponseEvent chatResponse = ChatResponseEvent.fromEvent(event); ChatMessage response = chatResponse.getResponse(); - if (response == null || response.getContent() == null) { + if (response == null || response.getText() == null) { return; } - ctx.sendEvent(new OutputEvent(response.getContent())); + ctx.sendEvent(new OutputEvent(response.getText())); } } diff --git a/examples/src/main/java/org/apache/flink/agents/examples/agents/MathAgent.java b/examples/src/main/java/org/apache/flink/agents/examples/agents/MathAgent.java index ba1ece089..c42816fc3 100644 --- a/examples/src/main/java/org/apache/flink/agents/examples/agents/MathAgent.java +++ b/examples/src/main/java/org/apache/flink/agents/examples/agents/MathAgent.java @@ -97,6 +97,6 @@ public static void processInput(InputEvent event, RunnerContext ctx) { /** Process chat response event and send the answer as output. */ @Action(EventType.ChatResponseEvent) public static void processChatResponse(ChatResponseEvent event, RunnerContext ctx) { - ctx.sendEvent(new OutputEvent(event.getResponse().getContent())); + ctx.sendEvent(new OutputEvent(event.getResponse().getText())); } } diff --git a/examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java b/examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java index 701715f33..77a0f113a 100644 --- a/examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java +++ b/examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java @@ -54,6 +54,6 @@ public static void processInput(InputEvent event, RunnerContext ctx) { /** Emit the model's answer as output. */ @Action(EventType.ChatResponseEvent) public static void processChatResponse(ChatResponseEvent event, RunnerContext ctx) { - ctx.sendEvent(new OutputEvent(event.getResponse().getContent())); + ctx.sendEvent(new OutputEvent(event.getResponse().getText())); } } diff --git a/examples/src/main/java/org/apache/flink/agents/examples/agents/ProductSuggestionAgent.java b/examples/src/main/java/org/apache/flink/agents/examples/agents/ProductSuggestionAgent.java index 8d5d0b90f..9279c7695 100644 --- a/examples/src/main/java/org/apache/flink/agents/examples/agents/ProductSuggestionAgent.java +++ b/examples/src/main/java/org/apache/flink/agents/examples/agents/ProductSuggestionAgent.java @@ -104,7 +104,7 @@ public static void processChatResponse(Event event, RunnerContext ctx) throws Ex // the handling that fits your pipeline: raise to fail the input (as below), emit an // OutputEvent carrying an error sentinel, or send a custom error event so downstream // operators can detect the failure. - JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getContent()); + JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getText()); JsonNode suggestionsNode = jsonNode.findValue("suggestion_list"); List suggestions = new ArrayList<>(); if (suggestionsNode.isArray()) { diff --git a/examples/src/main/java/org/apache/flink/agents/examples/agents/ReviewAnalysisAgent.java b/examples/src/main/java/org/apache/flink/agents/examples/agents/ReviewAnalysisAgent.java index 2862e0ae4..190aed4c1 100644 --- a/examples/src/main/java/org/apache/flink/agents/examples/agents/ReviewAnalysisAgent.java +++ b/examples/src/main/java/org/apache/flink/agents/examples/agents/ReviewAnalysisAgent.java @@ -117,7 +117,7 @@ public static void processChatResponse(Event event, RunnerContext ctx) throws Ex // the handling that fits your pipeline: raise to fail the input (as below), emit an // OutputEvent carrying an error sentinel, or send a custom error event so downstream // operators can detect the failure. - JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getContent()); + JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getText()); JsonNode scoreNode = jsonNode.findValue("score"); JsonNode reasonsNode = jsonNode.findValue("reasons"); if (scoreNode == null || reasonsNode == null) { diff --git a/examples/src/main/java/org/apache/flink/agents/examples/agents/TableReviewAnalysisAgent.java b/examples/src/main/java/org/apache/flink/agents/examples/agents/TableReviewAnalysisAgent.java index a692b3e9b..caf228b2b 100644 --- a/examples/src/main/java/org/apache/flink/agents/examples/agents/TableReviewAnalysisAgent.java +++ b/examples/src/main/java/org/apache/flink/agents/examples/agents/TableReviewAnalysisAgent.java @@ -137,7 +137,7 @@ public static void processChatResponse(Event event, RunnerContext ctx) throws Ex // the handling that fits your pipeline: raise to fail the input (as below), emit an // OutputEvent carrying an error sentinel, or send a custom error event so downstream // operators can detect the failure. - JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getContent()); + JsonNode jsonNode = MAPPER.readTree(chatResponse.getResponse().getText()); JsonNode scoreNode = jsonNode.findValue("score"); JsonNode reasonsNode = jsonNode.findValue("reasons"); if (scoreNode == null || reasonsNode == null) { diff --git a/integrations/chat-models/anthropic/src/main/java/org/apache/flink/agents/integrations/chatmodels/anthropic/AnthropicChatModelConnection.java b/integrations/chat-models/anthropic/src/main/java/org/apache/flink/agents/integrations/chatmodels/anthropic/AnthropicChatModelConnection.java index bf8e0afe0..f1a681691 100644 --- a/integrations/chat-models/anthropic/src/main/java/org/apache/flink/agents/integrations/chatmodels/anthropic/AnthropicChatModelConnection.java +++ b/integrations/chat-models/anthropic/src/main/java/org/apache/flink/agents/integrations/chatmodels/anthropic/AnthropicChatModelConnection.java @@ -448,13 +448,13 @@ static final class BuiltRequest { private List extractSystemMessages(List messages) { return messages.stream() .filter(m -> m.getRole() == MessageRole.SYSTEM) - .map(m -> TextBlockParam.builder().text(m.getContent()).build()) + .map(m -> TextBlockParam.builder().text(m.getText()).build()) .collect(Collectors.toList()); } private MessageParam convertToAnthropicMessage(ChatMessage message) { MessageRole role = message.getRole(); - String content = Optional.ofNullable(message.getContent()).orElse(""); + String content = Optional.ofNullable(message.getText()).orElse(""); switch (role) { case USER: diff --git a/integrations/chat-models/anthropic/src/test/java/org/apache/flink/agents/integrations/chatmodels/anthropic/AnthropicChatModelConnectionTest.java b/integrations/chat-models/anthropic/src/test/java/org/apache/flink/agents/integrations/chatmodels/anthropic/AnthropicChatModelConnectionTest.java index 091b900b5..1ef31b311 100644 --- a/integrations/chat-models/anthropic/src/test/java/org/apache/flink/agents/integrations/chatmodels/anthropic/AnthropicChatModelConnectionTest.java +++ b/integrations/chat-models/anthropic/src/test/java/org/apache/flink/agents/integrations/chatmodels/anthropic/AnthropicChatModelConnectionTest.java @@ -144,7 +144,7 @@ private static void assertPrefillDecision( assertThat(built.jsonPrefillApplied).isEqualTo(expectedApplied); assertThat(requestCarriesPrefill(built)).isEqualTo(expectedApplied); - assertThat(connection.convertResponse(built, textResponse(CONTINUATION)).getContent()) + assertThat(connection.convertResponse(built, textResponse(CONTINUATION)).getText()) .isEqualTo(expectedApplied ? COMPLETED : CONTINUATION); } @@ -475,7 +475,7 @@ void testJsonPrefillSuppressedWhenNativeApplies() { assertThat(built.jsonPrefillApplied).isFalse(); assertThat(requestCarriesPrefill(built)).isFalse(); // The provider returns a complete document, so nothing may be prepended to it. - assertThat(connection.convertResponse(built, textResponse(COMPLETED)).getContent()) + assertThat(connection.convertResponse(built, textResponse(COMPLETED)).getText()) .isEqualTo(COMPLETED); } @@ -494,7 +494,7 @@ void testJsonPrefillAppliedWhenSchemaFallsBack() { assertThat(built.jsonPrefillApplied).isTrue(); assertThat(requestCarriesPrefill(built)).isTrue(); - assertThat(connection.convertResponse(built, textResponse(CONTINUATION)).getContent()) + assertThat(connection.convertResponse(built, textResponse(CONTINUATION)).getText()) .isEqualTo(COMPLETED); } @@ -553,7 +553,7 @@ private static void assertPrefillDecisionForModel(String model, boolean expected assertThat(built.jsonPrefillApplied).isEqualTo(expectedApplied); assertThat(requestCarriesPrefill(built)).isEqualTo(expectedApplied); - assertThat(connection.convertResponse(built, textResponse(CONTINUATION)).getContent()) + assertThat(connection.convertResponse(built, textResponse(CONTINUATION)).getText()) .isEqualTo(expectedApplied ? COMPLETED : CONTINUATION); } diff --git a/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java b/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java index 86105572b..9b09787b1 100644 --- a/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java +++ b/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java @@ -139,7 +139,7 @@ public ChatMessage chat( if (!systemMsgs.isEmpty()) { requestBuilder.system( systemMsgs.stream() - .map(m -> SystemContentBlock.builder().text(m.getContent()).build()) + .map(m -> SystemContentBlock.builder().text(m.getText()).build()) .collect(Collectors.toList())); } @@ -233,7 +233,7 @@ private List mergeMessages(List msgs) { .toolUseId(toolCallId) .content( ToolResultContentBlock.builder() - .text(toolMsg.getContent()) + .text(toolMsg.getText()) .build()) .build())); i++; @@ -256,12 +256,12 @@ private Message toBedrockMessage(ChatMessage msg) { case USER: return Message.builder() .role(ConversationRole.USER) - .content(ContentBlock.fromText(msg.getContent())) + .content(ContentBlock.fromText(msg.getText())) .build(); case ASSISTANT: List blocks = new ArrayList<>(); - if (msg.getContent() != null && !msg.getContent().isEmpty()) { - blocks.add(ContentBlock.fromText(msg.getContent())); + if (msg.getText() != null && !msg.getText().isEmpty()) { + blocks.add(ContentBlock.fromText(msg.getText())); } if (msg.getToolCalls() != null && !msg.getToolCalls().isEmpty()) { for (Map call : msg.getToolCalls()) { @@ -290,7 +290,7 @@ private Message toBedrockMessage(ChatMessage msg) { .toolUseId(toolCallId) .content( ToolResultContentBlock.builder() - .text(msg.getContent()) + .text(msg.getText()) .build()) .build())) .build(); diff --git a/integrations/chat-models/gemini/src/main/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnection.java b/integrations/chat-models/gemini/src/main/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnection.java index e445fa276..4548ba26c 100644 --- a/integrations/chat-models/gemini/src/main/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnection.java +++ b/integrations/chat-models/gemini/src/main/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnection.java @@ -355,7 +355,7 @@ private Content extractSystemInstruction(List messages) { Part[] parts = messages.stream() .filter(m -> m.getRole() == MessageRole.SYSTEM) - .map(m -> Part.fromText(Optional.ofNullable(m.getContent()).orElse(""))) + .map(m -> Part.fromText(Optional.ofNullable(m.getText()).orElse(""))) .toArray(Part[]::new); return parts.length == 0 ? null : Content.fromParts(parts); } @@ -363,7 +363,7 @@ private Content extractSystemInstruction(List messages) { // Package-visible for unit testing of the message conversion. Content convertToContent(ChatMessage message, Map toolCallIdToName) { MessageRole role = message.getRole(); - String content = Optional.ofNullable(message.getContent()).orElse(""); + String content = Optional.ofNullable(message.getText()).orElse(""); switch (role) { case USER: diff --git a/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java b/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java index e86f12779..cd28f6a88 100644 --- a/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java +++ b/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java @@ -171,7 +171,7 @@ private OllamaChatMessage convertToOllamaChatMessages(ChatMessage message) { try { final OllamaChatMessageRole ollamaRole = OllamaChatMessageRole.getRole(role.name().toLowerCase()); - return new OllamaChatMessage(ollamaRole, message.getContent()); + return new OllamaChatMessage(ollamaRole, message.getText()); } catch (RoleNotFoundException e) { throw new RuntimeException(e); } diff --git a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAIChatCompletionsUtils.java b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAIChatCompletionsUtils.java index 35eddaf1c..60b1f2f0d 100644 --- a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAIChatCompletionsUtils.java +++ b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAIChatCompletionsUtils.java @@ -170,7 +170,7 @@ public static List convertToOpenAIMessages( /** Convert a single Flink Agents ChatMessage to an OpenAI ChatCompletionMessageParam. */ public static ChatCompletionMessageParam convertToOpenAIMessage(ChatMessage message) { MessageRole role = message.getRole(); - String content = Optional.ofNullable(message.getContent()).orElse(""); + String content = Optional.ofNullable(message.getText()).orElse(""); switch (role) { case SYSTEM: diff --git a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAIResponsesModelConnection.java b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAIResponsesModelConnection.java index 5dca0b6a4..9e9ce0f1f 100644 --- a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAIResponsesModelConnection.java +++ b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAIResponsesModelConnection.java @@ -223,7 +223,7 @@ private List convertInputItems(List messages) { private List convertSingleMessage(ChatMessage message) { List items = new ArrayList<>(); MessageRole role = message.getRole(); - String content = Optional.ofNullable(message.getContent()).orElse(""); + String content = Optional.ofNullable(message.getText()).orElse(""); switch (role) { case SYSTEM: diff --git a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java index 627ccd969..1ce286661 100644 --- a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java +++ b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java @@ -273,8 +273,8 @@ public ChatMessage chat( final ChatMessage chatMessage = parseResponse(MAPPER.readTree(response.body()), modelName); if (extractReasoning) { - final String[] parts = extractReasoning(chatMessage.getContent()); - chatMessage.setContent(parts[0]); + final String[] parts = extractReasoning(chatMessage.getText()); + chatMessage.setText(parts[0]); if (parts[1] != null) { chatMessage.getExtraArgs().put("reasoning", parts[1]); } @@ -469,12 +469,12 @@ static ArrayNode convertMessages(List messages) { case SYSTEM: case USER: node.put("role", role.name().toLowerCase()); - node.put("content", message.getContent()); + node.put("content", message.getText()); break; case ASSISTANT: node.put("role", "assistant"); - if (message.getContent() != null && !message.getContent().isEmpty()) { - node.put("content", message.getContent()); + if (message.getText() != null && !message.getText().isEmpty()) { + node.put("content", message.getText()); } final List> toolCalls = message.getToolCalls(); if (toolCalls != null && !toolCalls.isEmpty()) { @@ -488,7 +488,7 @@ static ArrayNode convertMessages(List messages) { "Tool message must have 'externalId' in extra args."); } node.put("role", "tool"); - node.put("content", message.getContent()); + node.put("content", message.getText()); node.put("tool_call_id", externalId.toString()); break; default: diff --git a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java index 9b1cd5b6d..7902f8d7c 100644 --- a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java +++ b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java @@ -375,7 +375,7 @@ void testParseResponse() throws Exception { WatsonxChatModelConnection.parseResponse(response, "ibm/granite-3-3-8b-instruct"); assertThat(message.getRole()).isEqualTo(MessageRole.ASSISTANT); - assertThat(message.getContent()).isEqualTo("Hello there!"); + assertThat(message.getText()).isEqualTo("Hello there!"); assertThat(message.getExtraArgs().get("model_name")) .isEqualTo("ibm/granite-3-3-8b-instruct"); assertThat(message.getExtraArgs().get("promptTokens")).isEqualTo(100L); @@ -405,8 +405,8 @@ void testIamTokenIsCached() throws Exception { WatsonxChatModelConnection connection = new WatsonxChatModelConnection( stubDescriptor(baseUrl(server), true, 0), NOOP, NO_ENVIRONMENT); - assertThat(chat(connection).getContent()).isEqualTo("Hello!"); - assertThat(chat(connection).getContent()).isEqualTo("Hello!"); + assertThat(chat(connection).getText()).isEqualTo("Hello!"); + assertThat(chat(connection).getText()).isEqualTo("Hello!"); assertThat(iamRequests).hasValue(1); assertThat(chatRequests).hasValue(2); } finally { @@ -473,7 +473,7 @@ void testRejectedIamTokenIsRefreshed(int rejectedStatus) throws Exception { WatsonxChatModelConnection connection = new WatsonxChatModelConnection( stubDescriptor(baseUrl(server), true, 0), NOOP, NO_ENVIRONMENT); - assertThat(chat(connection).getContent()).isEqualTo("Hello!"); + assertThat(chat(connection).getText()).isEqualTo("Hello!"); assertThat(iamRequests).hasValue(2); assertThat(chatRequests).hasValue(2); } finally { @@ -500,7 +500,7 @@ void testRetryLoop() throws Exception { WatsonxChatModelConnection retryingConnection = new WatsonxChatModelConnection( stubDescriptor(baseUrl(server), false, 1), NOOP, NO_ENVIRONMENT); - assertThat(chat(retryingConnection).getContent()).isEqualTo("Hello!"); + assertThat(chat(retryingConnection).getText()).isEqualTo("Hello!"); assertThat(chatRequests).hasValue(2); server.removeContext("/ml/v1/text/chat"); diff --git a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java index 50197e18d..4bac0e816 100644 --- a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java +++ b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java @@ -79,7 +79,7 @@ void testBasicChat() { Map.of("model", model(), "max_tokens", 100)); assertThat(response.getRole()).isEqualTo(MessageRole.ASSISTANT); - assertThat(response.getContent()).isNotBlank(); + assertThat(response.getText()).isNotBlank(); assertThat(response.getExtraArgs().get("promptTokens")).isNotNull(); assertThat(response.getExtraArgs().get("completionTokens")).isNotNull(); } diff --git a/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPPrompt.java b/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPPrompt.java index e539245a7..64201ea1c 100644 --- a/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPPrompt.java +++ b/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPPrompt.java @@ -154,7 +154,7 @@ public MCPServer getMcpServer() { public String formatString(Map arguments) { List messages = formatMessages(MessageRole.SYSTEM, arguments); return messages.stream() - .map(msg -> msg.getRole().getValue() + ": " + msg.getContent()) + .map(msg -> msg.getRole().getValue() + ": " + msg.getText()) .collect(Collectors.joining("\n")); } 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 57c191c75..b8e7bc3d5 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 @@ -321,7 +321,7 @@ static String cleanLlmResponse(String rawResponse) { @SuppressWarnings("unchecked") static ChatMessage generateStructuredOutput(ChatMessage response, Object outputSchema) throws JsonProcessingException { - String output = response.getContent(); + String output = response.getText(); output = cleanLlmResponse(output); Object structuredOutput; if (outputSchema instanceof Class) { diff --git a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonPrompt.java b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonPrompt.java index 2e7965377..ed05870f6 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonPrompt.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonPrompt.java @@ -18,6 +18,7 @@ package org.apache.flink.agents.plan.resource.python; import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ContentBlock; import org.apache.flink.agents.api.chat.messages.MessageRole; import org.apache.flink.agents.api.prompt.Prompt; @@ -69,20 +70,25 @@ public static PythonPrompt fromSerializedMap(Map serialized) { "Python prompt parsing failed. Template is not a string or list."); } - /** Parse a single ChatMessage from a Map representation. */ + /** + * Parse a single ChatMessage from the Python {@code model_dump()} representation, where the + * content arrives as a list of discriminated block maps under {@code blocks}. + */ @SuppressWarnings("unchecked") private static ChatMessage parseChatMessage(Map messageMap) { String roleValue = messageMap.get("role").toString(); MessageRole role = MessageRole.fromValue(roleValue); - Object contentObj = messageMap.get("content"); - String content = contentObj != null ? contentObj.toString() : ""; + List> blocks = (List>) messageMap.get("blocks"); List> toolCalls = (List>) messageMap.get("tool_calls"); Map extraArgs = (Map) messageMap.get("extra_args"); - return new ChatMessage(role, content, toolCalls, extraArgs); + ChatMessage message = + new ChatMessage(role, (List) null, toolCalls, extraArgs); + message.setBlocksFromMaps(blocks); + return message; } } diff --git a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareChatModelTest.java b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareChatModelTest.java index 561b2b5c6..bb0e3b516 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareChatModelTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareChatModelTest.java @@ -69,7 +69,7 @@ public Map getParameters() { @Override public ChatMessage chat(List messages) { // Return a deterministic response based on prompt name to assert on. - return new ChatMessage(MessageRole.ASSISTANT, "ok:" + messages.get(0).getContent()); + return new ChatMessage(MessageRole.ASSISTANT, "ok:" + messages.get(0).getText()); } } @@ -131,7 +131,7 @@ void retrieveAndChat() throws Exception { ChatMessage reply = model.chat(prompt.formatMessages(MessageRole.USER, new HashMap<>())); assertEquals(MessageRole.ASSISTANT, reply.getRole()); - assertEquals("ok:Hello world", reply.getContent()); + assertEquals("ok:Hello world", reply.getText()); } @Test @@ -154,7 +154,7 @@ void jsonRoundTrip() throws Exception { })); ChatMessage reply = model.chat(Prompt.fromText("Hi").formatMessages(MessageRole.USER, new HashMap<>())); - assertEquals("ok:Hi", reply.getContent()); + assertEquals("ok:Hi", reply.getText()); } @Test 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 46c9cc4b2..6ed47260e 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 @@ -403,7 +403,7 @@ void retryBudgetRunsBeforeFallback() throws Exception { // the selected model's retry budget is consumed BEFORE fallback: big's retry // succeeds and small is never resolved — the ordering the class javadoc guarantees - assertThat(ctx.chatResponse().getResponse().getContent()).isEqualTo("recovered on retry"); + assertThat(ctx.chatResponse().getResponse().getText()).isEqualTo("recovered on retry"); assertThat(ctx.resolvedChatModels).containsExactly("big"); assertThat(ctx.routingEventCount()).isEqualTo(1L); } @@ -446,7 +446,7 @@ void fallsBackToRemainingCandidateWhenSelectedModelFails() throws Exception { .containsExactly("route:router", "chat:router:big", "chat:router:small"); ChatResponseEvent response = ctx.chatResponse(); assertThat(response).isNotNull(); - assertThat(response.getResponse().getContent()).isEqualTo("ok from small"); + assertThat(response.getResponse().getText()).isEqualTo("ok from small"); Map routing = (Map) response.getResponse().getExtraArgs().get("model_routing"); assertThat(routing.get("final_model")).isEqualTo("small"); @@ -647,7 +647,7 @@ void routesOnceThenReusesSelectedModelAcrossToolRound() throws Exception { assertThat(ctx.routingEventCount()).isEqualTo(1L); assertThat(ctx.resolvedChatModels).containsExactly("big", "big"); assertThat(ctx.chatResponse()).isNotNull(); - assertThat(ctx.chatResponse().getResponse().getContent()).isEqualTo("final answer"); + assertThat(ctx.chatResponse().getResponse().getText()).isEqualTo("final answer"); // the routing metadata from the initial decision is carried onto the final response @SuppressWarnings("unchecked") diff --git a/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonPromptTest.java b/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonPromptTest.java index aa65bb9ef..cee1b27ab 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonPromptTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonPromptTest.java @@ -19,6 +19,7 @@ package org.apache.flink.agents.plan.resource.python; import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ImageBlock; import org.apache.flink.agents.api.chat.messages.MessageRole; import org.junit.jupiter.api.Test; @@ -50,14 +51,17 @@ public void testFromSerializedMapWithStringTemplate() { @Test public void testFromSerializedMapWithMessageListTemplate() { - // Create message map - Map systemMessage = new HashMap<>(); - systemMessage.put("role", "system"); - systemMessage.put("content", "You are a helpful assistant."); - - Map userMessage = new HashMap<>(); - userMessage.put("role", "user"); - userMessage.put("content", "Hello!"); + // The exact shape the Python side produces: LocalPrompt.model_dump() serializes each + // template ChatMessage with a `blocks` list (no `content` field) and keeps absent + // optional media fields as explicit nulls. + Map systemMessage = + messageDump("system", List.of(textBlockDump("You are a helpful assistant."))); + Map userMessage = + messageDump( + "user", + List.of( + textBlockDump("Hello! What's in {subject}?"), + imageBlockDump("image/png", "https://example.org/cat.png"))); List> messageList = new ArrayList<>(); messageList.add(systemMessage); @@ -70,14 +74,53 @@ public void testFromSerializedMapWithMessageListTemplate() { assertThat(prompt).isNotNull(); - // Test that the prompt formats messages correctly - List formattedMessages = - prompt.formatMessages(MessageRole.SYSTEM, new HashMap<>()); + // The restored prompt formats text blocks and passes media blocks through untouched. + Map kwargs = new HashMap<>(); + kwargs.put("subject", "this picture"); + List formattedMessages = prompt.formatMessages(MessageRole.SYSTEM, kwargs); assertThat(formattedMessages).hasSize(2); assertThat(formattedMessages.get(0).getRole()).isEqualTo(MessageRole.SYSTEM); - assertThat(formattedMessages.get(0).getContent()).isEqualTo("You are a helpful assistant."); + assertThat(formattedMessages.get(0).getText()).isEqualTo("You are a helpful assistant."); assertThat(formattedMessages.get(1).getRole()).isEqualTo(MessageRole.USER); - assertThat(formattedMessages.get(1).getContent()).isEqualTo("Hello!"); + assertThat(formattedMessages.get(1).getText()).isEqualTo("Hello! What's in this picture?"); + assertThat(formattedMessages.get(1).getBlocks()).hasSize(2); + assertThat(formattedMessages.get(1).getBlocks().get(1)) + .isInstanceOf(ImageBlock.class) + .satisfies( + block -> { + ImageBlock image = (ImageBlock) block; + assertThat(image.getMediaType()).isEqualTo("image/png"); + assertThat(image.getUrl()).isEqualTo("https://example.org/cat.png"); + assertThat(image.getData()).isNull(); + }); + } + + private static Map messageDump(String role, List> blocks) { + Map message = new HashMap<>(); + message.put("role", role); + message.put("blocks", blocks); + message.put("tool_calls", new ArrayList<>()); + message.put("extra_args", new HashMap<>()); + return message; + } + + private static Map textBlockDump(String text) { + Map block = new HashMap<>(); + block.put("type", "text"); + block.put("text", text); + return block; + } + + private static Map imageBlockDump(String mediaType, String url) { + Map block = new HashMap<>(); + block.put("type", "image"); + block.put("media_type", mediaType); + block.put("data", null); + block.put("url", url); + block.put("name", null); + block.put("size_bytes", null); + block.put("sha256", null); + return block; } @Test diff --git a/python/flink_agents/api/agents/react_agent.py b/python/flink_agents/api/agents/react_agent.py index 9f5f72df4..cdf37fa9d 100644 --- a/python/flink_agents/api/agents/react_agent.py +++ b/python/flink_agents/api/agents/react_agent.py @@ -78,13 +78,8 @@ class OutputData(BaseModel): # prepare prompt prompt = Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content='An example of output is {"result": 30.32}.', - ), - ChatMessage( - role=MessageRole.USER, content="What is ({a} + {b}) * {c}" - ), + ChatMessage.system('An example of output is {"result": 30.32}.'), + ChatMessage.user("What is ({a} + {b}) * {c}"), ], ) @@ -167,7 +162,7 @@ def start_action(event: Event, ctx: RunnerContext) -> None: role=MessageRole.USER, input=usr_input ) else: - usr_msgs = [ChatMessage(role=MessageRole.USER, content=usr_input)] + usr_msgs = [ChatMessage.user(usr_input)] else: if not prompt: err_msg = ( @@ -216,6 +211,6 @@ def stop_action(event: Event, ctx: RunnerContext) -> None: if STRUCTURED_OUTPUT in response.extra_args: output = response.extra_args[STRUCTURED_OUTPUT] else: - output = response.content + output = response.text ctx.send_event(OutputEvent(output=output)) diff --git a/python/flink_agents/api/chat_message.py b/python/flink_agents/api/chat_message.py index adb67e089..916789733 100644 --- a/python/flink_agents/api/chat_message.py +++ b/python/flink_agents/api/chat_message.py @@ -16,9 +16,10 @@ # limitations under the License. ################################################################################# from enum import Enum -from typing import Any, Dict, List +from typing import Any, Dict, List, Literal, Sequence -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator +from typing_extensions import Annotated class MessageRole(str, Enum): @@ -43,6 +44,88 @@ class MessageRole(str, Enum): TOOL = "tool" +class TextBlock(BaseModel): + """A plain-text, immutable part of a ChatMessage.""" + + model_config = ConfigDict(frozen=True) + + type: Literal["text"] = "text" + text: str = "" + + def __str__(self) -> str: + return self.text + + +class MediaBlock(BaseModel): + """Shared shape for binary media blocks: modality is the concrete type, + encoding is the media type (RFC 6838; historically called a MIME type). + + Media blocks are immutable, and the payload is carried by exactly one of + base64 ``data`` or an externally managed ``url``. URL-backed content is + externally managed: URLs may expire, may not be reachable by the model + provider, and may be invalid after recovery from a checkpoint. The optional + ``name``/``size_bytes``/``sha256`` metadata also serves the Event Log, + which records media metadata instead of payload bytes. + """ + + # Frozen keeps sharing a block (e.g. across a routing context copy) safe, + # and matches the validated immutable construction on the Java side. + model_config = ConfigDict(frozen=True) + + media_type: str + data: str | None = None # base64; exactly one of data / url set + url: str | None = None + name: str | None = None + size_bytes: int | None = None + sha256: str | None = None + + @model_validator(mode="after") + def _exactly_one_source(self) -> "MediaBlock": + if (self.data is None) == (self.url is None): + msg = "A media block carries exactly one of base64 data or a URL." + raise ValueError(msg) + return self + + def __str__(self) -> str: + source = "inline" if self.data is not None else f"url={self.url}" + return f"{type(self).__name__}({self.media_type}, {source})" + + +class ImageBlock(MediaBlock): + """The image content of a ChatMessage — see MediaBlock for the media shape.""" + + type: Literal["image"] = "image" + + +class AudioBlock(MediaBlock): + """The audio content of a ChatMessage — see MediaBlock for the media shape.""" + + type: Literal["audio"] = "audio" + + +class VideoBlock(MediaBlock): + """The video content of a ChatMessage — see MediaBlock for the media shape.""" + + type: Literal["video"] = "video" + + +class DocumentBlock(MediaBlock): + """The document content of a ChatMessage — see MediaBlock for the media shape.""" + + type: Literal["document"] = "document" + + +ContentBlock = Annotated[ + TextBlock | ImageBlock | AudioBlock | VideoBlock | DocumentBlock, + Field(discriminator="type"), +] + + +def _blocks_of(text: str) -> List[ContentBlock]: + """An empty text becomes an empty block list rather than an empty text block.""" + return [TextBlock(text=text)] if text else [] + + class ChatMessage(BaseModel): """Chat message. @@ -52,21 +135,76 @@ class ChatMessage(BaseModel): ---------- role : MessageRole The message productor or purpose. - content : str - The content of the message. + blocks : List[ContentBlock] + The ordered, typed content of the message; a text-only message carries + a single TextBlock. tool_calls: List[Dict[str, Any]] The tools call information. extra_args : dict[str, Any] Additional information about the message. """ + # Unknown keys fail loudly: the replaced `content` field would otherwise be + # silently ignored, producing an empty message instead of an error. + model_config = ConfigDict(extra="forbid") + role: MessageRole = MessageRole.USER - content: str = Field(default_factory=str) + blocks: List[ContentBlock] = Field(default_factory=list) tool_calls: List[Dict[str, Any]] = Field(default_factory=list) extra_args: Dict[str, Any] = Field(default_factory=dict) + @property + def text(self) -> str: + """The text projection: the ordered concatenation of the TextBlocks.""" + return "".join( + block.text for block in self.blocks if isinstance(block, TextBlock) + ) + + def set_text(self, text: str) -> None: + """Replace the content with a single text block (empty text clears it).""" + self.blocks = _blocks_of(text) + + @classmethod + def user( + cls, content: str | Sequence[ContentBlock], **kwargs: Any + ) -> "ChatMessage": + """Create a USER message from text or content blocks.""" + return cls.of(MessageRole.USER, content, **kwargs) + + @classmethod + def system( + cls, content: str | Sequence[ContentBlock], **kwargs: Any + ) -> "ChatMessage": + """Create a SYSTEM message from text or content blocks.""" + return cls.of(MessageRole.SYSTEM, content, **kwargs) + + @classmethod + def assistant( + cls, content: str | Sequence[ContentBlock], **kwargs: Any + ) -> "ChatMessage": + """Create an ASSISTANT message from text or content blocks.""" + return cls.of(MessageRole.ASSISTANT, content, **kwargs) + + @classmethod + def tool( + cls, content: str | Sequence[ContentBlock], **kwargs: Any + ) -> "ChatMessage": + """Create a TOOL message from text or content blocks.""" + return cls.of(MessageRole.TOOL, content, **kwargs) + + @classmethod + def of( + cls, + role: MessageRole, + content: str | Sequence[ContentBlock], + **kwargs: Any, + ) -> "ChatMessage": + """Create a message with the given role from text or content blocks.""" + blocks = _blocks_of(content) if isinstance(content, str) else list(content) + return cls(role=role, blocks=blocks, **kwargs) + def __str__(self) -> str: - return f"{self.role.value}: {self.content}" + return f"{self.role.value}: {self.text}" def find_first_system_message(messages: List[ChatMessage]) -> int: diff --git a/python/flink_agents/api/chat_models/chat_model.py b/python/flink_agents/api/chat_models/chat_model.py index 528ae347c..6c4acdd7f 100644 --- a/python/flink_agents/api/chat_models/chat_model.py +++ b/python/flink_agents/api/chat_models/chat_model.py @@ -405,10 +405,9 @@ def chat( prompt_messages = self._get_prompt().format_messages(**str_prompt_args) # append meaningful messages + # any block counts, so image-only messages survive for msg in messages: - if ( - msg.content is not None and msg.content != "" - ) or msg.role == MessageRole.ASSISTANT: + if len(msg.blocks) > 0 or msg.role == MessageRole.ASSISTANT: prompt_messages.append(msg) messages = prompt_messages @@ -417,9 +416,7 @@ def chat( messages = ( messages[: index + 1] + [ - ChatMessage( - role=MessageRole.SYSTEM, content=self.skill_discovery_prompt - ) + ChatMessage.system(self.skill_discovery_prompt) ] + messages[index + 1 :] ) diff --git a/python/flink_agents/api/chat_models/tests/test_chat_model_base.py b/python/flink_agents/api/chat_models/tests/test_chat_model_base.py index a7524973b..77a44f5da 100644 --- a/python/flink_agents/api/chat_models/tests/test_chat_model_base.py +++ b/python/flink_agents/api/chat_models/tests/test_chat_model_base.py @@ -66,7 +66,7 @@ def chat( self.captured_messages = list(messages) self.captured_kwargs = dict(kwargs) self.captured_output_schema = output_schema - return ChatMessage(role=MessageRole.ASSISTANT, content="ok") + return ChatMessage.of(MessageRole.ASSISTANT, "ok") class _RecordingChatModelSetup(BaseChatModelSetup): @@ -106,7 +106,7 @@ def test_chat_fills_template_from_prompt_args_parameter() -> None: setup.chat([], prompt_args={"key": "value"}) assert len(connection.captured_messages) == 1 - assert connection.captured_messages[0].content == "Task: value" + assert connection.captured_messages[0].text == "Task: value" def test_chat_does_not_read_template_vars_from_extra_args() -> None: @@ -114,14 +114,13 @@ def test_chat_does_not_read_template_vars_from_extra_args() -> None: prompt = Prompt.from_text(text="Task: {key}") setup, connection = _build_setup(prompt) - user_message = ChatMessage( - role=MessageRole.USER, content="hello", extra_args={"key": "value"} + user_message = ChatMessage.of(MessageRole.USER, "hello", extra_args={"key": "value"} ) setup.chat([user_message], prompt_args={}) assert len(connection.captured_messages) == 2 - assert connection.captured_messages[0].content == "Task: {key}" - assert connection.captured_messages[1].content == "hello" + assert connection.captured_messages[0].text == "Task: {key}" + assert connection.captured_messages[1].text == "hello" def test_chat_refills_template_on_subsequent_invocations() -> None: @@ -131,13 +130,13 @@ def test_chat_refills_template_on_subsequent_invocations() -> None: setup.chat([], prompt_args={"key": "v1"}) assert len(connection.captured_messages) == 1 - assert connection.captured_messages[0].content == "Task: v1" + assert connection.captured_messages[0].text == "Task: v1" - tool_response = ChatMessage(role=MessageRole.TOOL, content="tool result") + tool_response = ChatMessage.of(MessageRole.TOOL, "tool result") setup.chat([tool_response], prompt_args={"key": "v1"}) assert len(connection.captured_messages) == 2 - assert connection.captured_messages[0].content == "Task: v1" - assert connection.captured_messages[1].content == "tool result" + assert connection.captured_messages[0].text == "Task: v1" + assert connection.captured_messages[1].text == "tool result" def test_default_capability_predicate_is_false() -> None: diff --git a/python/flink_agents/api/chat_models/tests/test_token_metrics.py b/python/flink_agents/api/chat_models/tests/test_token_metrics.py index 15455836c..8ad0ed767 100644 --- a/python/flink_agents/api/chat_models/tests/test_token_metrics.py +++ b/python/flink_agents/api/chat_models/tests/test_token_metrics.py @@ -43,7 +43,7 @@ def resource_type(cls) -> ResourceType: def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: """Simple test implementation.""" - return ChatMessage(role=MessageRole.ASSISTANT, content="Test response") + return ChatMessage.of(MessageRole.ASSISTANT, "Test response") def test_record_token_metrics( self, diff --git a/python/flink_agents/api/prompts/prompt.py b/python/flink_agents/api/prompts/prompt.py index 37205d276..2e74d104a 100644 --- a/python/flink_agents/api/prompts/prompt.py +++ b/python/flink_agents/api/prompts/prompt.py @@ -20,7 +20,7 @@ from typing_extensions import override -from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.api.chat_message import ChatMessage, MessageRole, TextBlock from flink_agents.api.prompts.utils import format_string from flink_agents.api.resource import ResourceType, SerializableResource @@ -73,7 +73,7 @@ def format_string(self, **kwargs: str) -> str: else: msgs = [] for m in self.template: - msg = f"{m.role.value}: {format_string(m.content, **kwargs)}" + msg = f"{m.role.value}: {format_string(m.text, **kwargs)}" if m.extra_args is not None and len(m.extra_args) > 0: msg += f"{m.extra_args}" msgs.append(msg) @@ -84,14 +84,18 @@ def format_messages( ) -> List[ChatMessage]: """Generate list of ChatMessage from template with input arguments.""" if isinstance(self.template, str): - return [ - ChatMessage(role=role, content=format_string(self.template, **kwargs)) - ] + return [ChatMessage.of(role, format_string(self.template, **kwargs))] else: msgs = [] for m in self.template: msg = ChatMessage( - role=m.role, content=format_string(m.content, **kwargs) + role=m.role, + blocks=[ + TextBlock(text=format_string(b.text, **kwargs)) + if isinstance(b, TextBlock) + else b + for b in m.blocks + ], ) msgs.append(msg) return msgs diff --git a/python/flink_agents/api/tests/test_chat_message.py b/python/flink_agents/api/tests/test_chat_message.py new file mode 100644 index 000000000..96e8cd580 --- /dev/null +++ b/python/flink_agents/api/tests/test_chat_message.py @@ -0,0 +1,146 @@ +################################################################################ +# 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. +################################################################################# +import pytest +from pydantic import ValidationError + +from flink_agents.api.chat_message import ( + AudioBlock, + ChatMessage, + DocumentBlock, + ImageBlock, + MessageRole, + TextBlock, + VideoBlock, +) + + +def test_text_only_wire_shape() -> None: + """A text-only message serializes to a single typed text block.""" + message = ChatMessage.user("hello world") + dumped = message.model_dump(mode="json", exclude_none=True) + assert dumped == { + "role": "user", + "blocks": [{"type": "text", "text": "hello world"}], + "tool_calls": [], + "extra_args": {}, + } + + +def test_media_block_wire_shape_omits_absent_fields() -> None: + message = ChatMessage.user( + [ + TextBlock(text="What's in this picture?"), + ImageBlock(media_type="image/png", data="aGk="), + ] + ) + image = message.model_dump(mode="json", exclude_none=True)["blocks"][1] + assert image == {"type": "image", "media_type": "image/png", "data": "aGk="} + + +def test_mixed_blocks_round_trip_preserves_order_and_types() -> None: + original = ChatMessage( + role=MessageRole.TOOL, + blocks=[ + TextBlock(text="before"), + ImageBlock( + media_type="image/jpeg", + url="https://example.org/cat.jpg", + name="cat.jpg", + size_bytes=123, + ), + DocumentBlock(media_type="application/pdf", data="cGRm"), + TextBlock(text="after"), + ], + ) + restored = ChatMessage.model_validate_json(original.model_dump_json()) + assert restored == original + assert [type(b).__name__ for b in restored.blocks] == [ + "TextBlock", + "ImageBlock", + "DocumentBlock", + "TextBlock", + ] + assert restored.text == "beforeafter" + + +def test_audio_and_video_round_trip() -> None: + original = ChatMessage.user( + [ + AudioBlock(media_type="audio/wav", data="d2F2"), + VideoBlock(media_type="video/mp4", url="https://example.org/v.mp4"), + ] + ) + restored = ChatMessage.model_validate_json(original.model_dump_json()) + assert restored == original + + +def test_java_wire_shape_deserializes() -> None: + """The exact JSON the Java API emits validates into typed blocks.""" + payload = { + "role": "user", + "blocks": [ + {"type": "text", "text": "hi"}, + {"type": "image", "media_type": "image/png", "data": "aGk="}, + ], + "tool_calls": [], + "extra_args": {}, + } + message = ChatMessage.model_validate(payload) + assert isinstance(message.blocks[0], TextBlock) + assert isinstance(message.blocks[1], ImageBlock) + assert message.text == "hi" + + +def test_legacy_content_kwarg_fails_loudly() -> None: + """The replaced `content` field is rejected, never silently dropped.""" + with pytest.raises(ValidationError): + ChatMessage(role=MessageRole.USER, content="hi") + + +def test_media_requires_exactly_one_source() -> None: + with pytest.raises(ValidationError): + ImageBlock(media_type="image/png") + with pytest.raises(ValidationError): + ImageBlock(media_type="image/png", data="aGk=", url="https://example.org/x") + + +def test_blocks_are_frozen() -> None: + """Blocks are immutable value objects, so sharing them never shares state.""" + text = TextBlock(text="hi") + with pytest.raises(ValidationError): + text.text = "mutated" + image = ImageBlock(media_type="image/png", data="aGk=") + with pytest.raises(ValidationError): + image.data = "bXV0YXRlZA==" + with pytest.raises(ValidationError): + image.url = "https://example.org/x" + assert text.text == "hi" + assert image.data == "aGk=" + + +def test_factories_and_text_projection() -> None: + assert ChatMessage.system("be nice").role == MessageRole.SYSTEM + assert ChatMessage.assistant("ok").text == "ok" + assert ChatMessage.tool("result").blocks == [TextBlock(text="result")] + # Empty text becomes an empty block list rather than an empty text block. + empty = ChatMessage.user("") + assert empty.blocks == [] + assert empty.text == "" + empty.set_text("replaced") + assert empty.text == "replaced" + assert str(ChatMessage.user("hi")) == "user: hi" 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 545cbfdf4..76bb580f2 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 @@ -139,7 +139,7 @@ def test_python_can_deserialize_output_event_from_java_snapshot() -> None: def _build_chat_request_event() -> ChatRequestEvent: event = ChatRequestEvent( model="test-model", - messages=[ChatMessage(role=MessageRole.USER, content="hello world")], + messages=[ChatMessage.of(MessageRole.USER, "hello world")], ) return _force_id(event, _FIXED_EVENT_ID) @@ -163,7 +163,7 @@ def test_python_can_deserialize_chat_request_event_from_java_snapshot() -> None: assert len(typed.messages) == 1 msg = typed.messages[0] assert msg.role == MessageRole.USER, f"Role mismatch: got {msg.role!r}" - assert msg.content == "hello world" + assert msg.text == "hello world" def test_chat_request_row_type_info_output_schema_is_not_portable_across_languages_known_gap() -> ( @@ -188,7 +188,7 @@ def test_chat_request_row_type_info_output_schema_is_not_portable_across_languag ) event = ChatRequestEvent( model="test-model", - messages=[ChatMessage(role=MessageRole.USER, content="hi")], + messages=[ChatMessage.of(MessageRole.USER, "hi")], output_schema=schema, ) payload = event.model_dump_json() @@ -204,7 +204,7 @@ def test_chat_request_row_type_info_output_schema_is_not_portable_across_languag def _build_chat_response_event() -> ChatResponseEvent: event = ChatResponseEvent( request_id=_FIXED_REQUEST_ID, - response=ChatMessage(role=MessageRole.ASSISTANT, content="hi there"), + response=ChatMessage.of(MessageRole.ASSISTANT, "hi there"), ) return _force_id(event, _FIXED_EVENT_ID) @@ -235,7 +235,7 @@ def test_python_can_deserialize_chat_response_event_from_java_snapshot() -> None assert typed.response.role == MessageRole.ASSISTANT, ( f"Response role mismatch: got {typed.response.role!r}" ) - assert typed.response.content == "hi there" + assert typed.response.text == "hi there" # ── ToolRequestEvent ──────────────────────────────────────────────────── diff --git a/python/flink_agents/api/tests/test_prompt.py b/python/flink_agents/api/tests/test_prompt.py index 0da76661b..2f17f8401 100644 --- a/python/flink_agents/api/tests/test_prompt.py +++ b/python/flink_agents/api/tests/test_prompt.py @@ -51,9 +51,9 @@ def test_prompt_from_text_to_messages(text_prompt: LocalPrompt) -> None: description="wireless noise-canceling headphones with 20-hour battery life", review="The headphones broke after one week of use. Very poor quality", ) == [ - ChatMessage( - role=MessageRole.SYSTEM, - content="You ara a product review analyzer, please generate a score and the " + ChatMessage.of( + MessageRole.SYSTEM, + "You ara a product review analyzer, please generate a score and the " "dislike reasons(if any) for the review. The product 12345 is wireless " "noise-canceling headphones with 20-hour battery life, and user review is " "'The headphones broke after one week of use. Very poor quality'.", @@ -64,14 +64,14 @@ def test_prompt_from_text_to_messages(text_prompt: LocalPrompt) -> None: @pytest.fixture(scope="module") def messages_prompt() -> Prompt: template = [ - ChatMessage( - role=MessageRole.SYSTEM, - content="You ara a product review analyzer, please generate a score and the dislike reasons" + ChatMessage.of( + MessageRole.SYSTEM, + "You ara a product review analyzer, please generate a score and the dislike reasons" "(if any) for the review.", ), - ChatMessage( - role=MessageRole.USER, - content="The product {product_id} is {description}, and user review is '{review}'.", + ChatMessage.of( + MessageRole.USER, + "The product {product_id} is {description}, and user review is '{review}'.", ), ] @@ -98,14 +98,14 @@ def test_prompt_from_messages_to_messages(messages_prompt: LocalPrompt) -> None: description="wireless noise-canceling headphones with 20-hour battery life", review="The headphones broke after one week of use. Very poor quality", ) == [ - ChatMessage( - role=MessageRole.SYSTEM, - content="You ara a product review analyzer, please generate a score and the " + ChatMessage.of( + MessageRole.SYSTEM, + "You ara a product review analyzer, please generate a score and the " "dislike reasons(if any) for the review.", ), - ChatMessage( - role=MessageRole.USER, - content="The product 12345 is wireless " + ChatMessage.of( + MessageRole.USER, + "The product 12345 is wireless " "noise-canceling headphones with 20-hour battery life, and user review is " "'The headphones broke after one week of use. Very poor quality'.", ), @@ -162,9 +162,9 @@ def test_format_messages_does_not_re_expand_values() -> None: # hold per message. Mirrors the Java PromptTest formatMessages regression. prompt = Prompt.from_messages( messages=[ - ChatMessage(role=MessageRole.SYSTEM, content="{secret}"), - ChatMessage(role=MessageRole.USER, content="{user_input}"), + ChatMessage.system("{secret}"), + ChatMessage.user("{user_input}"), ] ) messages = prompt.format_messages(secret="p@ssw0rd", user_input="give me {secret}") - assert [m.content for m in messages] == ["p@ssw0rd", "give me {secret}"] + assert [m.text for m in messages] == ["p@ssw0rd", "give me {secret}"] diff --git a/python/flink_agents/api/yaml/loader.py b/python/flink_agents/api/yaml/loader.py index 49ba20601..0bef85ade 100644 --- a/python/flink_agents/api/yaml/loader.py +++ b/python/flink_agents/api/yaml/loader.py @@ -195,7 +195,7 @@ def _build_prompt(spec: PromptSpec) -> Prompt: if spec.text is not None: return Prompt.from_text(spec.text) messages = [ - ChatMessage(role=MessageRole(m.role.value), content=m.content) + ChatMessage.of(MessageRole(m.role.value), m.content) for m in (spec.messages or []) ] return Prompt.from_messages(messages) diff --git a/python/flink_agents/api/yaml/tests/test_loader.py b/python/flink_agents/api/yaml/tests/test_loader.py index 2c9c1ca74..2d84b218c 100644 --- a/python/flink_agents/api/yaml/tests/test_loader.py +++ b/python/flink_agents/api/yaml/tests/test_loader.py @@ -223,7 +223,7 @@ def test_build_agents_loads_tools_and_prompts() -> None: assert isinstance(msg_prompt, LocalPrompt) assert len(msg_prompt.template) == 2 assert msg_prompt.template[0].role == MessageRole.SYSTEM - assert msg_prompt.template[1].content == "{q}" + assert msg_prompt.template[1].text == "{q}" def test_build_agents_handles_shared_resources_and_actions() -> None: diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/agent_skills_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/agent_skills_test.py index d4f853a64..219b912f2 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/agent_skills_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/agent_skills_test.py @@ -99,9 +99,7 @@ def my_skills() -> Skills: def system_prompt() -> Prompt: return Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content="You are a help assistant. Use the math-calculator skill when asked to evaluate " + ChatMessage.of(MessageRole.SYSTEM, "You are a help assistant. Use the math-calculator skill when asked to evaluate " "an expression. You **must load the skill first** and strictly follow the instructions " "of the skill.", ) @@ -118,9 +116,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ChatRequestEvent( model="openai_setup", messages=[ - ChatMessage( - role=MessageRole.USER, - content=f"Please evaluate the expression: ({input.a} ^ {input.b})", + ChatMessage.of(MessageRole.USER, f"Please evaluate the expression: ({input.a} ^ {input.b})", ) ], ) @@ -131,9 +127,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ChatRequestEvent( model="openai_setup", messages=[ - ChatMessage( - role=MessageRole.USER, - content=input, + ChatMessage.of(MessageRole.USER, input, ) ], ) @@ -143,7 +137,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: @staticmethod def process_chat_response(event: Event, ctx: RunnerContext) -> None: chat_response = ChatResponseEvent.from_event(event) - ctx.send_event(OutputEvent(output=chat_response.response.content)) + ctx.send_event(OutputEvent(output=chat_response.response.text)) @pytest.mark.skipif(not API_KEY, reason="openai api key is required.") @@ -279,15 +273,11 @@ def test_react_agent_with_skills(tmp_path: Path) -> None: # prepare prompt prompt = Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content="You are a math calculate assistant. Use the math-calculator skill when asked to evaluate " + ChatMessage.of(MessageRole.SYSTEM, "You are a math calculate assistant. Use the math-calculator skill when asked to evaluate " "an expression. You **must load the skill first** and strictly follow the instructions " "of the skill.", ), - ChatMessage( - role=MessageRole.USER, - content="Please evaluate the expression: {a} ^ {b}", + ChatMessage.of(MessageRole.USER, "Please evaluate the expression: {a} ^ {b}", ), ], ) diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/built_in_action_async_execution_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/built_in_action_async_execution_test.py index cb0de86ee..e818e281a 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/built_in_action_async_execution_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/built_in_action_async_execution_test.py @@ -48,20 +48,19 @@ def model_kwargs(self) -> Dict[str, Any]: @override def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: time.sleep(5) # Simulate network delay - if "sum" in messages[-1].content: - input = messages[-1].content + if "sum" in messages[-1].text: + input = messages[-1].text function = {"name": "add", "arguments": {"a": 1, "b": 2}} tool_call = { "id": uuid.uuid4(), "type": ToolType.FUNCTION, "function": function, } - return ChatMessage( - role=MessageRole.ASSISTANT, content=input, tool_calls=[tool_call] + return ChatMessage.of(MessageRole.ASSISTANT, input, tool_calls=[tool_call] ) else: - content = "\n".join([message.content for message in messages]) - return ChatMessage(role=MessageRole.ASSISTANT, content=content) + content = "\n".join([message.text for message in messages]) + return ChatMessage.of(MessageRole.ASSISTANT, content) class AsyncTestAgent(Agent): @@ -92,7 +91,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ctx.send_event( ChatRequestEvent( model="slow_chat_model", - messages=[ChatMessage(role=MessageRole.USER, content=input)], + messages=[ChatMessage.of(MessageRole.USER, input)], prompt_args={"task": input}, ) ) @@ -101,7 +100,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: @staticmethod def process_chat_response(event: Event, ctx: RunnerContext) -> None: input = ChatResponseEvent.from_event(event).response - ctx.send_event(OutputEvent(output=input.content)) + ctx.send_event(OutputEvent(output=input.text)) def test_built_in_actions_async_execution() -> None: diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_integration_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_integration_agent.py index c7158dbd1..8934c77e5 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_integration_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_integration_agent.py @@ -180,7 +180,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ctx.send_event( ChatRequestEvent( model=model_name, - messages=[ChatMessage(role=MessageRole.USER, content=input_event.input)], + messages=[ChatMessage.of(MessageRole.USER, input_event.input)], ) ) @@ -190,5 +190,5 @@ def process_chat_response(event: Event, ctx: RunnerContext) -> None: """User defined action for processing chat model response.""" chat_response = ChatResponseEvent.from_event(event) input = chat_response.response - if chat_response.response and input.content: - ctx.send_event(OutputEvent(output=input.content)) + if chat_response.response and input.text: + ctx.send_event(OutputEvent(output=input.text)) diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/checkpoint_recovery_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/checkpoint_recovery_agent.py index 75fdf9836..49d0d797c 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/checkpoint_recovery_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/checkpoint_recovery_agent.py @@ -330,10 +330,8 @@ def chat( # emitted content, which is where the assertion can reach it. It is # wrapped as a JSON object because the caller parses this round against # the output schema; the joined text survives verbatim inside the field. - content = "\n".join(message.content for message in messages) - return ChatMessage( - role=MessageRole.ASSISTANT, - content=json.dumps({_STRUCTURED_TRANSCRIPT_FIELD: content}), + content = "\n".join(message.text for message in messages) + return ChatMessage.of(MessageRole.ASSISTANT, json.dumps({_STRUCTURED_TRANSCRIPT_FIELD: content}), ) # Validate the tool was bound before the model was invoked. @@ -344,9 +342,7 @@ def chat( "type": ToolType.FUNCTION, "function": {"name": BLOCKING_TOOL_NAME, "arguments": {}}, } - return ChatMessage( - role=MessageRole.ASSISTANT, - content=_ROUND_ONE_MARKER, + return ChatMessage.of(MessageRole.ASSISTANT, _ROUND_ONE_MARKER, tool_calls=[tool_call], ) @@ -444,7 +440,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ChatRequestEvent( model="recovery_chat_model", messages=[ - ChatMessage(role=MessageRole.USER, content=input_data.content) + ChatMessage.of(MessageRole.USER, input_data.content) ], prompt_args={"task": input_data.content}, # Set once, on the only request this agent issues. The framework @@ -507,7 +503,7 @@ def process_chat_response(event: Event, ctx: RunnerContext) -> None: "transcript": transcript, # The unparsed response alongside the unpacked transcript, so a payload # that failed to unpack is diagnosable from this file alone. - "response_content": chat_response.content, + "response_content": chat_response.text, } verdict = json.dumps(record, sort_keys=True) _atomic_write( diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/e2e_tests_mcp/mcp_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/e2e_tests_mcp/mcp_test.py index 2c625c566..fb8d0e4fe 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/e2e_tests_mcp/mcp_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/e2e_tests_mcp/mcp_test.py @@ -142,9 +142,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ) else: # Send chat request asking to use the add tool - msg = ChatMessage( - role=MessageRole.USER, - content=f"Please use the add tool to calculate the sum of {input_data.a} and {input_data.b}.", + msg = ChatMessage.of(MessageRole.USER, f"Please use the add tool to calculate the sum of {input_data.a} and {input_data.b}.", ) ctx.send_event(ChatRequestEvent(model="math_chat_model", messages=[msg])) @@ -153,8 +151,8 @@ def process_input(event: Event, ctx: RunnerContext) -> None: def process_chat_response(event: Event, ctx: RunnerContext) -> None: """Process chat response and output result.""" response = ChatResponseEvent.from_event(event).response - if response and response.content: - ctx.send_event(OutputEvent(output=response.content)) + if response and response.text: + ctx.send_event(OutputEvent(output=response.text)) current_dir = Path(__file__).parent diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py index 2aa03c264..450656d4e 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py @@ -106,8 +106,8 @@ def chat( schema out of ``**kwargs``. """ self._reject_unsupported_output_schema(output_schema) - if "sum" in messages[-1].content: - input = messages[-1].content + if "sum" in messages[-1].text: + input = messages[-1].text # Validate the tool was bound before the model was invoked. assert tools[0].name == "add" function = {"name": "add", "arguments": {"a": 1, "b": 2}} @@ -116,11 +116,10 @@ def chat( "type": ToolType.FUNCTION, "function": function, } - return ChatMessage( - role=MessageRole.ASSISTANT, content=input, tool_calls=[tool_call] + return ChatMessage.of(MessageRole.ASSISTANT, input, tool_calls=[tool_call] ) - content = "\n".join([message.content for message in messages]) - return ChatMessage(role=MessageRole.ASSISTANT, content=content) + content = "\n".join([message.text for message in messages]) + return ChatMessage.of(MessageRole.ASSISTANT, content) class MockChatModel(BaseChatModelSetup): @@ -153,7 +152,7 @@ def chat( else: prompt = self.prompt - if "sum" in messages[-1].content: + if "sum" in messages[-1].text: str_prompt_args = ( {k: str(v) for k, v in prompt_args.items()} if prompt_args else {} ) @@ -234,7 +233,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ChatRequestEvent( model="mock_chat_model", messages=[ - ChatMessage(role=MessageRole.USER, content=input_data.content) + ChatMessage.of(MessageRole.USER, input_data.content) ], prompt_args={"task": input_data.content}, ) @@ -248,7 +247,7 @@ def process_chat_response(event: Event, ctx: RunnerContext) -> None: input_id = ctx.short_term_memory.get("input_id") ctx.send_event( OutputEvent( - output=MockChatModelOutput(id=input_id, result=response.content) + output=MockChatModelOutput(id=input_id, result=response.text) ) ) @@ -269,9 +268,7 @@ def model_kwargs(self) -> Dict[str, Any]: def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: """Echo the input alongside the custom descriptor fields.""" - return ChatMessage( - role=MessageRole.ASSISTANT, - content=f"{messages[0].content} {self.host} {self.desc}", + return ChatMessage.of(MessageRole.ASSISTANT, f"{messages[0].text} {self.host} {self.desc}", ) @@ -302,8 +299,8 @@ def mock_action(event: Event, ctx: RunnerContext) -> None: name="mock_chat_model", type=ResourceType.CHAT_MODEL ) content = mock_chat_model.chat( - messages=[ChatMessage(role=MessageRole.USER, content=input_data.content)] - ).content + messages=[ChatMessage.of(MessageRole.USER, input_data.content)] + ).text ctx.send_event( OutputEvent( output=MockChatModelOutput(id=input_data.id, result=content) diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_test.py index 261378f87..6e427215e 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_test.py @@ -121,11 +121,9 @@ def test_react_agent_on_remote_runner( # prepare prompt prompt = Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content='An example of output is {"result": 30.32}.', + ChatMessage.of(MessageRole.SYSTEM, 'An example of output is {"result": 30.32}.', ), - ChatMessage(role=MessageRole.USER, content="What is ({a} + {b}) * {c}"), + ChatMessage.of(MessageRole.USER, "What is ({a} + {b}) * {c}"), ], ) @@ -245,7 +243,7 @@ def test_react_agent_no_output_schema_on_remote_runner( # prepare prompt prompt = Prompt.from_messages( messages=[ - ChatMessage(role=MessageRole.USER, content="What is ({a} + {b}) * {c}"), + ChatMessage.of(MessageRole.USER, "What is ({a} + {b}) * {c}"), ], ) diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py index 9358fec9d..b8d1b3ae8 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py @@ -95,7 +95,7 @@ def request_chat(event: Event, ctx: RunnerContext) -> None: ctx.send_event( ChatRequestEvent( model="mock_model", - messages=[ChatMessage(role=MessageRole.USER, content=order_id)], + messages=[ChatMessage.of(MessageRole.USER, order_id)], ) ) @@ -104,7 +104,7 @@ def request_chat(event: Event, ctx: RunnerContext) -> None: def emit_result(event: Event, ctx: RunnerContext) -> None: """Emit the final assistant response as output.""" response_event = ChatResponseEvent.from_event(event) - ctx.send_event(OutputEvent(output=response_event.response.content)) + ctx.send_event(OutputEvent(output=response_event.response.text)) class MockToolChatConnection(BaseChatModelConnection): @@ -126,17 +126,15 @@ def chat( self._reject_unsupported_output_schema(output_schema) last_message = messages[-1] if last_message.role == MessageRole.TOOL: - return ChatMessage(role=MessageRole.ASSISTANT, content=last_message.content) + return ChatMessage.of(MessageRole.ASSISTANT, last_message.text) for candidate_tool in tools or []: if "tenant_id" in str(candidate_tool.metadata.get_parameters_dict()): msg = "Injected argument leaked into tool schema." raise RuntimeError(msg) - order_id = str(last_message.content) - return ChatMessage( - role=MessageRole.ASSISTANT, - content="", + order_id = str(last_message.text) + return ChatMessage.of(MessageRole.ASSISTANT, "", tool_calls=[ { "id": f"call-{order_id}", diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/yaml_test_actions.py b/python/flink_agents/e2e_tests/e2e_tests_integration/yaml_test_actions.py index 44054abfb..2870dcee1 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/yaml_test_actions.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/yaml_test_actions.py @@ -112,7 +112,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ctx.send_event( ChatRequestEvent( model=model_name, - messages=[ChatMessage(role=MessageRole.USER, content=data.text)], + messages=[ChatMessage.of(MessageRole.USER, data.text)], ) ) @@ -129,7 +129,7 @@ def chat_request(event: Event, ctx: RunnerContext) -> None: ctx.send_event( ChatRequestEvent( model="chat_model", - messages=[ChatMessage(role=MessageRole.USER, content=data.text)], + messages=[ChatMessage.of(MessageRole.USER, data.text)], ) ) @@ -138,11 +138,11 @@ def process_chat_response(event: Event, ctx: RunnerContext) -> None: """Emit the model's text response, tagged with the original input id.""" chat_response = ChatResponseEvent.from_event(event) response = chat_response.response - if not response or not response.content: + if not response or not response.text: return input_id = ctx.short_term_memory.get("input_id") ctx.send_event( - OutputEvent(output=YamlChatOutput(id=input_id, answer=response.content)) + OutputEvent(output=YamlChatOutput(id=input_id, answer=response.text)) ) @@ -162,9 +162,7 @@ def commentary_request(event: Event, ctx: RunnerContext) -> None: ChatRequestEvent( model="chat_model", messages=[ - ChatMessage( - role=MessageRole.USER, - content=( + ChatMessage.of(MessageRole.USER, ( "Here is a math answer from another assistant: " f"{data.answer!r}. Reply with the numeric result only." ), diff --git a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/chat_model_cross_language_agent.py b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/chat_model_cross_language_agent.py index 3920f2389..d8c6daa64 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/chat_model_cross_language_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/chat_model_cross_language_agent.py @@ -58,9 +58,7 @@ def from_messages_prompt() -> Prompt: """Prompt for instruction.""" return Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content="Please answer the user's question.", + ChatMessage.of(MessageRole.SYSTEM, "Please answer the user's question.", ), ], ) @@ -147,7 +145,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ctx.send_event( ChatRequestEvent( model=model_name, - messages=[ChatMessage(role=MessageRole.USER, content=input_event.input)], + messages=[ChatMessage.of(MessageRole.USER, input_event.input)], ) ) @@ -157,5 +155,5 @@ def process_chat_response(event: Event, ctx: RunnerContext) -> None: """User defined action for processing chat model response.""" chat_response = ChatResponseEvent.from_event(event) input = chat_response.response - if chat_response.response and input.content: - ctx.send_event(OutputEvent(output=input.content)) + if chat_response.response and input.text: + ctx.send_event(OutputEvent(output=input.text)) diff --git a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/yaml_cross_language_actions.py b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/yaml_cross_language_actions.py index d1fea1e1f..752d8e823 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/yaml_cross_language_actions.py +++ b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/yaml_cross_language_actions.py @@ -41,7 +41,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ctx.send_event( ChatRequestEvent( model=model_name, - messages=[ChatMessage(role=MessageRole.USER, content=text)], + messages=[ChatMessage.of(MessageRole.USER, text)], ) ) @@ -50,8 +50,8 @@ def process_chat_response(event: Event, ctx: RunnerContext) -> None: """Emit the model's textual response.""" chat_response = ChatResponseEvent.from_event(event) response = chat_response.response - if response and response.content: - ctx.send_event(OutputEvent(output=response.content)) + if response and response.text: + ctx.send_event(OutputEvent(output=response.text)) def calculate_bmi(weight_kg: float, height_m: float) -> float: diff --git a/python/flink_agents/examples/quickstart/agents/custom_types_and_resources.py b/python/flink_agents/examples/quickstart/agents/custom_types_and_resources.py index b12087a02..ca9d22744 100644 --- a/python/flink_agents/examples/quickstart/agents/custom_types_and_resources.py +++ b/python/flink_agents/examples/quickstart/agents/custom_types_and_resources.py @@ -51,15 +51,11 @@ review_analysis_prompt = Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content=review_analysis_system_prompt_str, + ChatMessage.of(MessageRole.SYSTEM, review_analysis_system_prompt_str, ), # Here we just fill the prompt with input, user should deserialize # input element to input text self in action. - ChatMessage( - role=MessageRole.USER, - content=""" + ChatMessage.of(MessageRole.USER, """ "input": {input} """, @@ -70,16 +66,12 @@ # Prompt for review analysis react agent. review_analysis_react_prompt = Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content=review_analysis_system_prompt_str, + ChatMessage.of(MessageRole.SYSTEM, review_analysis_system_prompt_str, ), # For react agent, if the input element is not primitive types, # framework will deserialize input element to dict and fill the prompt. # Note, the input element should be primitive types, BaseModel or Row. - ChatMessage( - role=MessageRole.USER, - content=""" + ChatMessage.of(MessageRole.USER, """ "id": {id}, "review": {review} """, diff --git a/python/flink_agents/examples/quickstart/agents/math_agent.py b/python/flink_agents/examples/quickstart/agents/math_agent.py index 77f177607..b21140a13 100644 --- a/python/flink_agents/examples/quickstart/agents/math_agent.py +++ b/python/flink_agents/examples/quickstart/agents/math_agent.py @@ -55,9 +55,7 @@ def system_prompt() -> Prompt: """System prompt instructing the model to use the skill.""" return Prompt.from_messages( messages=[ - ChatMessage( - role=MessageRole.SYSTEM, - content="You are a helpful math assistant. Use the " + ChatMessage.of(MessageRole.SYSTEM, "You are a helpful math assistant. Use the " "math-calculator skill when asked to evaluate an expression. " "You must load the skill first and strictly follow its " "instructions. Reply with only the final numeric result.", @@ -88,7 +86,7 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ctx.send_event( ChatRequestEvent( model="math_model", - messages=[ChatMessage(role=MessageRole.USER, content=question)], + messages=[ChatMessage.of(MessageRole.USER, question)], ) ) @@ -97,4 +95,4 @@ def process_input(event: Event, ctx: RunnerContext) -> None: def process_chat_response(event: Event, ctx: RunnerContext) -> None: """Process chat response event and send the answer as output.""" chat_response = ChatResponseEvent.from_event(event) - ctx.send_event(OutputEvent(output=chat_response.response.content)) + ctx.send_event(OutputEvent(output=chat_response.response.text)) diff --git a/python/flink_agents/examples/quickstart/agents/parallel_chat_agent.py b/python/flink_agents/examples/quickstart/agents/parallel_chat_agent.py index f2e206522..3492f73c6 100644 --- a/python/flink_agents/examples/quickstart/agents/parallel_chat_agent.py +++ b/python/flink_agents/examples/quickstart/agents/parallel_chat_agent.py @@ -56,10 +56,8 @@ def _build_aspect_request(text: str, aspect: str) -> ChatRequestEvent: return ChatRequestEvent( model="sentiment_model", messages=[ - ChatMessage(role=MessageRole.SYSTEM, content=PARALLEL_SYSTEM_PROMPT), - ChatMessage( - role=MessageRole.USER, - content=f'Judge the "{aspect}" dimension: {text}', + ChatMessage.of(MessageRole.SYSTEM, PARALLEL_SYSTEM_PROMPT), + ChatMessage.of(MessageRole.USER, f'Judge the "{aspect}" dimension: {text}', ), ], output_schema=OutputSchema(output_schema=AspectResponse), @@ -76,8 +74,8 @@ def _build_summarize_request(text: str, sentiments: Dict[str, str]) -> ChatReque return ChatRequestEvent( model="sentiment_model", messages=[ - ChatMessage(role=MessageRole.SYSTEM, content=AGGREGATE_SYSTEM_PROMPT), - ChatMessage(role=MessageRole.USER, content=body), + ChatMessage.of(MessageRole.SYSTEM, AGGREGATE_SYSTEM_PROMPT), + ChatMessage.of(MessageRole.USER, body), ], output_schema=OutputSchema(output_schema=SummaryResponse), ) diff --git a/python/flink_agents/examples/quickstart/agents/product_suggestion_agent.py b/python/flink_agents/examples/quickstart/agents/product_suggestion_agent.py index b0979fa76..658326e27 100644 --- a/python/flink_agents/examples/quickstart/agents/product_suggestion_agent.py +++ b/python/flink_agents/examples/quickstart/agents/product_suggestion_agent.py @@ -99,7 +99,7 @@ def process_chat_response(event: Event, ctx: RunnerContext) -> None: # pipeline: raise to fail the input (as below), emit an OutputEvent # carrying an error sentinel, or send a custom error event so # downstream operators can detect the failure. - json_content = json.loads(chat_response.response.content) + json_content = json.loads(chat_response.response.text) ctx.send_event( OutputEvent( output=ProductSuggestion( diff --git a/python/flink_agents/examples/quickstart/agents/review_analysis_agent.py b/python/flink_agents/examples/quickstart/agents/review_analysis_agent.py index 8e7f9a105..d494e422c 100644 --- a/python/flink_agents/examples/quickstart/agents/review_analysis_agent.py +++ b/python/flink_agents/examples/quickstart/agents/review_analysis_agent.py @@ -115,7 +115,7 @@ def process_chat_response(event: Event, ctx: RunnerContext) -> None: # pipeline: raise to fail the input (as below), emit an OutputEvent # carrying an error sentinel, or send a custom error event so # downstream operators can detect the failure. - json_content = json.loads(chat_response.response.content) + json_content = json.loads(chat_response.response.text) ctx.send_event( OutputEvent( output=ProductReviewAnalysisRes( diff --git a/python/flink_agents/examples/quickstart/agents/table_review_analysis_agent.py b/python/flink_agents/examples/quickstart/agents/table_review_analysis_agent.py index 7ddf94390..c95d7dfbf 100644 --- a/python/flink_agents/examples/quickstart/agents/table_review_analysis_agent.py +++ b/python/flink_agents/examples/quickstart/agents/table_review_analysis_agent.py @@ -147,7 +147,7 @@ def process_chat_response(event: Event, ctx: RunnerContext) -> None: # pipeline: raise to fail the input (as below), emit an OutputEvent # carrying an error sentinel, or send a custom error event so # downstream operators can detect the failure. - json_content = json.loads(chat_response.response.content) + json_content = json.loads(chat_response.response.text) ctx.send_event( OutputEvent( output=ProductReviewAnalysisRes( diff --git a/python/flink_agents/examples/rag/agents/rag_agent.py b/python/flink_agents/examples/rag/agents/rag_agent.py index 461080d19..88a6282e4 100644 --- a/python/flink_agents/examples/rag/agents/rag_agent.py +++ b/python/flink_agents/examples/rag/agents/rag_agent.py @@ -145,7 +145,7 @@ def process_retrieved_context(event: Event, ctx: RunnerContext) -> None: ctx.send_event( ChatRequestEvent( model="chat_model", - messages=[ChatMessage(role=MessageRole.USER, content=enhanced_prompt)], + messages=[ChatMessage.of(MessageRole.USER, enhanced_prompt)], ) ) @@ -154,5 +154,5 @@ def process_retrieved_context(event: Event, ctx: RunnerContext) -> None: def process_chat_response(event: Event, ctx: RunnerContext) -> None: """Process chat model response and generate output.""" chat_response = ChatResponseEvent.from_event(event) - if chat_response.response and chat_response.response.content: - ctx.send_event(OutputEvent(output=chat_response.response.content)) + if chat_response.response and chat_response.response.text: + ctx.send_event(OutputEvent(output=chat_response.response.text)) diff --git a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py index 2cf1d8da1..b35fcc9ec 100644 --- a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py +++ b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py @@ -59,7 +59,7 @@ def convert_to_anthropic_message(message: ChatMessage) -> MessageParam: { "type": "tool_result", "tool_use_id": message.extra_args.get("external_id"), - "content": message.content, + "content": message.text, } ], } @@ -69,7 +69,7 @@ def convert_to_anthropic_message(message: ChatMessage) -> MessageParam: content = ( anthropic_content_blocks if anthropic_content_blocks is not None - else message.content + else message.text ) return { "role": message.role.value, @@ -78,7 +78,7 @@ def convert_to_anthropic_message(message: ChatMessage) -> MessageParam: else: return { "role": message.role.value, - "content": message.content, + "content": message.text, } @@ -107,7 +107,7 @@ def convert_to_anthropic_system_prompts( message for message in messages if message.role == MessageRole.SYSTEM ] return [ - TextBlockParam(type="text", text=message.content) for message in system_messages + TextBlockParam(type="text", text=message.text) for message in system_messages ] @@ -426,18 +426,14 @@ def chat( ] extra_args["anthropic_content_blocks"] = message.content - return ChatMessage( - role=MessageRole(message.role), - content=text, + return ChatMessage.of(MessageRole(message.role), text, tool_calls=tool_calls, extra_args=extra_args, ) else: # TODO: handle other stop_reason values according to Anthropic API: # https://docs.anthropic.com/en/api/messages#response-stop-reason - return ChatMessage( - role=MessageRole(message.role), - content=text, + return ChatMessage.of(MessageRole(message.role), text, extra_args=extra_args, ) diff --git a/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_chat_model.py b/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_chat_model.py index c6a6e8f5c..9cf046ab4 100644 --- a/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_chat_model.py +++ b/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_chat_model.py @@ -54,7 +54,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: connection="anthropic_server", resource_context=mock_ctx, ) - response = chat_model.chat([ChatMessage(role=MessageRole.USER, content="Hello!")]) + response = chat_model.chat([ChatMessage.of(MessageRole.USER, "Hello!")]) assert response is not None assert str(response).strip() != "" @@ -97,7 +97,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: resource_context=mock_ctx, ) response = chat_model.chat( - [ChatMessage(role=MessageRole.USER, content="What is 1 + 1?")] + [ChatMessage.of(MessageRole.USER, "What is 1 + 1?")] ) tool_calls = response.tool_calls assert len(tool_calls) == 1 diff --git a/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_response_parsing.py b/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_response_parsing.py index 421ac3cd8..5455ace1b 100644 --- a/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_response_parsing.py +++ b/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_response_parsing.py @@ -64,9 +64,9 @@ def test_tool_use_response_without_leading_text() -> None: usage=_usage(), ) response = _connection_returning(message).chat( - [ChatMessage(role=MessageRole.USER, content="add 1 and 2")] + [ChatMessage.of(MessageRole.USER, "add 1 and 2")] ) - assert response.content == "" + assert response.text == "" assert len(response.tool_calls) == 1 assert response.tool_calls[0]["function"]["name"] == "add" @@ -86,9 +86,9 @@ def test_tool_use_response_keeps_leading_text() -> None: usage=_usage(), ) response = _connection_returning(message).chat( - [ChatMessage(role=MessageRole.USER, content="add 1 and 2")] + [ChatMessage.of(MessageRole.USER, "add 1 and 2")] ) - assert response.content == "Let me add those." + assert response.text == "Let me add those." assert len(response.tool_calls) == 1 @@ -103,9 +103,9 @@ def test_plain_text_response() -> None: usage=_usage(), ) response = _connection_returning(message).chat( - [ChatMessage(role=MessageRole.USER, content="hi")] + [ChatMessage.of(MessageRole.USER, "hi")] ) - assert response.content == "Hello!" + assert response.text == "Hello!" def test_plain_text_response_keeps_token_usage() -> None: @@ -122,7 +122,7 @@ def test_plain_text_response_keeps_token_usage() -> None: usage=Usage(input_tokens=7, output_tokens=3), ) response = _connection_returning(message).chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model="claude-sonnet-4-5", ) assert response.extra_args["model_name"] == "claude-sonnet-4-5" @@ -144,7 +144,7 @@ def test_tool_use_response_keeps_token_usage() -> None: usage=Usage(input_tokens=7, output_tokens=3), ) response = _connection_returning(message).chat( - [ChatMessage(role=MessageRole.USER, content="add 1 and 2")], + [ChatMessage.of(MessageRole.USER, "add 1 and 2")], model="claude-sonnet-4-5", ) assert response.extra_args["promptTokens"] == 7 @@ -219,7 +219,7 @@ def _request_kwargs(**chat_kwargs: Any) -> Dict[str, Any]: usage=_usage(), ) connection = _connection_returning(message) - connection.chat([ChatMessage(role=MessageRole.USER, content="hi")], **chat_kwargs) + connection.chat([ChatMessage.of(MessageRole.USER, "hi")], **chat_kwargs) return connection.client.messages.create.call_args.kwargs @@ -386,10 +386,10 @@ def _prefill_outcome(**chat_kwargs: Any) -> tuple: ) connection = _connection_returning(message) response = connection.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], **chat_kwargs + [ChatMessage.of(MessageRole.USER, "hi")], **chat_kwargs ) sent = connection.client.messages.create.call_args.kwargs["messages"] - return sent[-1] == {"role": "assistant", "content": "{"}, response.content + return sent[-1] == {"role": "assistant", "content": "{"}, response.text def test_json_prefill_not_applied_by_default() -> None: diff --git a/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_chat_model.py b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_chat_model.py index 3febba044..8516914be 100644 --- a/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_chat_model.py +++ b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_chat_model.py @@ -64,7 +64,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: connection="azure_openai", resource_context=mock_ctx, ) - response = chat_model.chat([ChatMessage(role=MessageRole.USER, content="Hello!")]) + response = chat_model.chat([ChatMessage.of(MessageRole.USER, "Hello!")]) assert response is not None assert str(response).strip() != "" @@ -112,9 +112,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: ) response = chat_model.chat( [ - ChatMessage( - role=MessageRole.USER, - content="You MUST use the add tool to calculate: What is 377 + 688?", + ChatMessage.of(MessageRole.USER, "You MUST use the add tool to calculate: What is 377 + 688?", ) ] ) @@ -212,7 +210,7 @@ def test_chat_rejects_reserved_key_in_additional_kwargs() -> None: ) with pytest.raises(ValueError, match="additional_kwargs"): connection.chat( - messages=[ChatMessage(role=MessageRole.USER, content="hi")], + messages=[ChatMessage.of(MessageRole.USER, "hi")], model="my-deployment", temperature=0.3, additional_kwargs={"temperature": 5.0}, diff --git a/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_native_structured_output.py b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_native_structured_output.py index cbaf084d8..6f1748369 100644 --- a/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_native_structured_output.py +++ b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_native_structured_output.py @@ -113,7 +113,7 @@ def _chat_with_caller_response_format( else {"response_format": CALLER_RESPONSE_FORMAT} ) conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment=model_of_azure_deployment, output_schema=None if schema is None else OutputSchema(output_schema=schema), @@ -125,7 +125,7 @@ def test_native_applied_for_capable_deployment_model() -> None: """response_format json_schema strict applied for a BaseModel on a capable model.""" conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o-mini", output_schema=OutputSchema(output_schema=Person), @@ -146,7 +146,7 @@ def test_capable_native_request_still_targets_the_deployment() -> None: """ conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o-mini", output_schema=OutputSchema(output_schema=Person), @@ -158,7 +158,7 @@ def test_native_not_applied_when_deployment_model_absent() -> None: """Native NOT applied when the backing model of the deployment is unknown.""" conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, output_schema=OutputSchema(output_schema=Person), ) @@ -169,7 +169,7 @@ def test_native_not_applied_for_unknown_deployment_model() -> None: """Native NOT applied for a backing model outside the allowlist.""" conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="some-unknown-model", output_schema=OutputSchema(output_schema=Person), @@ -185,7 +185,7 @@ def test_native_not_applied_for_bare_gpt_4o() -> None: """ conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o", output_schema=OutputSchema(output_schema=Person), @@ -203,7 +203,7 @@ def test_native_applied_for_ga_date_at_or_above_floor(api_version: str) -> None: """ conn = _connection(api_version=api_version) conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o-mini", output_schema=OutputSchema(output_schema=Person), @@ -222,7 +222,7 @@ def test_native_not_applied_for_non_date_api_version(api_version: str) -> None: """ conn = _connection(api_version=api_version) conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o-mini", output_schema=OutputSchema(output_schema=Person), @@ -234,7 +234,7 @@ def test_native_not_applied_when_api_version_below_floor() -> None: """Native NOT applied when the configured api-version predates the floor.""" conn = _connection(api_version=BELOW_FLOOR_API_VERSION) conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o-mini", output_schema=OutputSchema(output_schema=Person), @@ -250,7 +250,7 @@ def test_native_not_applied_when_api_version_empty() -> None: """ conn = _connection(api_version="") conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o-mini", output_schema=OutputSchema(output_schema=Person), @@ -262,7 +262,7 @@ def test_native_not_applied_when_schema_none() -> None: """Native NOT applied when no output schema is supplied.""" conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o-mini", output_schema=None, @@ -274,7 +274,7 @@ def test_native_not_applied_for_row_type_info() -> None: """Native NOT applied for a RowTypeInfo schema (BaseModel-only scope).""" conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o-mini", output_schema=OutputSchema(output_schema=ROW_TYPE), @@ -292,7 +292,7 @@ def test_native_applied_even_when_tools_bound() -> None: conn = _connection() tool = FunctionTool(func=PythonFunction.from_callable(_add)) conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], tools=[tool], model=DEPLOYMENT, model_of_azure_deployment="gpt-4o-mini", diff --git a/python/flink_agents/integrations/chat_models/ollama_chat_model.py b/python/flink_agents/integrations/chat_models/ollama_chat_model.py index bd16905a1..ef926bbc8 100644 --- a/python/flink_agents/integrations/chat_models/ollama_chat_model.py +++ b/python/flink_agents/integrations/chat_models/ollama_chat_model.py @@ -149,9 +149,7 @@ def chat( extra_args["promptTokens"] = response.prompt_eval_count extra_args["completionTokens"] = response.eval_count - return ChatMessage( - role=MessageRole(response.message.role), - content=content, + return ChatMessage.of(MessageRole(response.message.role), content, tool_calls=tool_calls, extra_args=extra_args, ) @@ -160,7 +158,7 @@ def chat( def __convert_to_ollama_messages(messages: Sequence[ChatMessage]) -> List[Message]: ollama_messages = [] for message in messages: - ollama_message = Message(role=message.role.value, content=message.content) + ollama_message = Message(role=message.role.value, content=message.text) if len(message.tool_calls) > 0: ollama_tool_calls = [] for tool_call in message.tool_calls: diff --git a/python/flink_agents/integrations/chat_models/openai/openai_utils.py b/python/flink_agents/integrations/chat_models/openai/openai_utils.py index 9601d9ebc..41802c7a1 100644 --- a/python/flink_agents/integrations/chat_models/openai/openai_utils.py +++ b/python/flink_agents/integrations/chat_models/openai/openai_utils.py @@ -143,7 +143,7 @@ def convert_to_openai_message(message: ChatMessage) -> ChatCompletionMessagePara if role == MessageRole.SYSTEM: system_message: ChatCompletionSystemMessageParam = { "role": "system", - "content": message.content, + "content": message.text, } system_message.update(message.extra_args) return system_message @@ -152,7 +152,7 @@ def convert_to_openai_message(message: ChatMessage) -> ChatCompletionMessagePara elif role == MessageRole.USER: user_message: ChatCompletionUserMessageParam = { "role": "user", - "content": message.content, + "content": message.text, } user_message.update(message.extra_args) return user_message @@ -160,7 +160,7 @@ def convert_to_openai_message(message: ChatMessage) -> ChatCompletionMessagePara elif role == MessageRole.ASSISTANT: # Assistant messages may have empty content when tool_calls are present - content = message.content if message.content or not message.tool_calls else None + content = message.text if message.text or not message.tool_calls else None assistant_message: ChatCompletionAssistantMessageParam = { "role": "assistant", "content": content, @@ -183,7 +183,7 @@ def convert_to_openai_message(message: ChatMessage) -> ChatCompletionMessagePara raise ValueError(msg) tool_message: ChatCompletionToolMessageParam = { "role": "tool", - "content": message.content, + "content": message.text, "tool_call_id": tool_call_id, } return tool_message @@ -220,9 +220,7 @@ def convert_from_openai_message( ] if message.refusal is not None: extra_args = {**extra_args, "refusal": message.refusal} - return ChatMessage( - role=MessageRole(message.role), - content=message.content or "", + return ChatMessage.of(MessageRole(message.role), message.content or "", tool_calls=tool_calls, extra_args=extra_args, ) diff --git a/python/flink_agents/integrations/chat_models/openai/tests/test_openai_chat_model.py b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_chat_model.py index 5a30004c5..89acbff26 100644 --- a/python/flink_agents/integrations/chat_models/openai/tests/test_openai_chat_model.py +++ b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_chat_model.py @@ -59,7 +59,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: chat_model = OpenAIChatModelSetup( model=test_model, connection="openai", resource_context=mock_ctx ) - response = chat_model.chat([ChatMessage(role=MessageRole.USER, content="Hello!")]) + response = chat_model.chat([ChatMessage.of(MessageRole.USER, "Hello!")]) assert response is not None assert str(response).strip() != "" @@ -102,7 +102,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: resource_context=mock_ctx, ) response = chat_model.chat( - [ChatMessage(role=MessageRole.USER, content="What is 377 + 688?")] + [ChatMessage.of(MessageRole.USER, "What is 377 + 688?")] ) tool_calls = response.tool_calls assert len(tool_calls) == 1 diff --git a/python/flink_agents/integrations/chat_models/openai/tests/test_openai_native_structured_output.py b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_native_structured_output.py index 7132d8040..60d45f146 100644 --- a/python/flink_agents/integrations/chat_models/openai/tests/test_openai_native_structured_output.py +++ b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_native_structured_output.py @@ -82,7 +82,7 @@ def test_native_applied_for_basemodel_capable_model() -> None: """response_format json_schema strict applied for a BaseModel on a capable model.""" conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model="gpt-4o", output_schema=OutputSchema(output_schema=Person), ) @@ -96,7 +96,7 @@ def test_native_not_applied_for_incapable_model() -> None: """Native NOT applied for a BaseModel on an incapable model (prompt fallback).""" conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model="gpt-3.5-turbo", output_schema=OutputSchema(output_schema=Person), ) @@ -111,7 +111,7 @@ def test_native_not_applied_for_pre_cutoff_snapshot() -> None: """ conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model="gpt-4o-2024-05-13", output_schema=OutputSchema(output_schema=Person), ) @@ -122,7 +122,7 @@ def test_native_not_applied_when_schema_none() -> None: """Native NOT applied when no output schema is supplied.""" conn = _connection() conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model="gpt-4o", output_schema=None, ) @@ -134,7 +134,7 @@ def test_native_not_applied_for_row_type_info() -> None: conn = _connection() row_type = Types.ROW_NAMED(["name"], [Types.STRING()]) conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model="gpt-4o", output_schema=OutputSchema(output_schema=row_type), ) @@ -146,7 +146,7 @@ def test_native_applied_even_when_tools_bound() -> None: conn = _connection() tool = FunctionTool(func=PythonFunction.from_callable(_add)) conn.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], tools=[tool], model="gpt-4o", output_schema=OutputSchema(output_schema=Person), diff --git a/python/flink_agents/integrations/chat_models/openai/tests/test_openai_response_parsing.py b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_response_parsing.py index 07c5e3d91..2ffd8b8cd 100644 --- a/python/flink_agents/integrations/chat_models/openai/tests/test_openai_response_parsing.py +++ b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_response_parsing.py @@ -38,7 +38,7 @@ def test_refusal_is_preserved_in_extra_args(refusal: str) -> None: assert result.extra_args["refusal"] == refusal assert result.extra_args["promptTokens"] == 3 - assert result.content == "" + assert result.text == "" def test_no_refusal_key_when_refusal_absent() -> None: diff --git a/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py b/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py index cceedcd2e..9ba102d0d 100644 --- a/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py +++ b/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py @@ -44,7 +44,7 @@ def test_ollama_chat() -> None: server = OllamaChatModelConnection(request_timeout=120.0) response = server.chat( - [ChatMessage(role=MessageRole.USER, content="Hello!")], model=test_model + [ChatMessage.of(MessageRole.USER, "Hello!")], model=test_model ) assert response is not None assert str(response).strip() != "" @@ -98,9 +98,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: response = llm.chat( [ - ChatMessage( - role=MessageRole.USER, - content="Could you help me calculate the sum of 1 and 2?", + ChatMessage.of(MessageRole.USER, "Could you help me calculate the sum of 1 and 2?", ) ] ) @@ -177,9 +175,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: # Call the chat method response = llm.chat( [ - ChatMessage( - role=MessageRole.USER, - content="What's the meaning of life?", + ChatMessage.of(MessageRole.USER, "What's the meaning of life?", ) ] ) @@ -189,7 +185,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: # Check that the response content has been cleaned assert ( - response.content + response.text == "The meaning of life is often considered to be 42, according to the Hitchhiker's Guide to the Galaxy." ) # Check that the reasoning has been extracted and stored diff --git a/python/flink_agents/integrations/chat_models/tests/test_tongyi_chat_model.py b/python/flink_agents/integrations/chat_models/tests/test_tongyi_chat_model.py index f46d7bc8b..4831f887d 100644 --- a/python/flink_agents/integrations/chat_models/tests/test_tongyi_chat_model.py +++ b/python/flink_agents/integrations/chat_models/tests/test_tongyi_chat_model.py @@ -43,11 +43,11 @@ def test_tongyi_chat() -> None: """Test basic chat functionality of TongyiChatModelConnection.""" connection = TongyiChatModelConnection() response = connection.chat( - [ChatMessage(role=MessageRole.USER, content="Hello!")], model=test_model + [ChatMessage.of(MessageRole.USER, "Hello!")], model=test_model ) assert response is not None - assert response.content is not None - assert response.content.strip() != "" + assert response.text is not None + assert response.text.strip() != "" assert response.role == MessageRole.ASSISTANT @@ -99,9 +99,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: response = llm.chat( [ - ChatMessage( - role=MessageRole.USER, - content="Could you help me calculate the sum of 1 and 2?", + ChatMessage.of(MessageRole.USER, "Could you help me calculate the sum of 1 and 2?", ) ] ) @@ -164,13 +162,13 @@ def get_resource(name: str, type: ResourceType) -> Resource: llm.open() response = llm.chat( - [ChatMessage(role=MessageRole.USER, content="What's the meaning of life?")] + [ChatMessage.of(MessageRole.USER, "What's the meaning of life?")] ) mock_call.assert_called_once() assert ( - response.content + response.text == "The meaning of life is often considered to be 42, according to the Hitchhiker's Guide to the Galaxy." ) assert "reasoning" in response.extra_args diff --git a/python/flink_agents/integrations/chat_models/tongyi_chat_model.py b/python/flink_agents/integrations/chat_models/tongyi_chat_model.py index 5e7b52a1c..a619b4ecc 100644 --- a/python/flink_agents/integrations/chat_models/tongyi_chat_model.py +++ b/python/flink_agents/integrations/chat_models/tongyi_chat_model.py @@ -173,9 +173,9 @@ def chat( if extract_reasoning and reasoning_content: extra_args["reasoning"] = reasoning_content - return ChatMessage( - role=MessageRole(response_message.get("role", "assistant")), - content=content, + return ChatMessage.of( + MessageRole(response_message.get("role", "assistant")), + content, tool_calls=tool_calls, extra_args=extra_args, ) @@ -188,7 +188,7 @@ def __convert_to_tongyi_messages( for message in messages: msg_dict: Dict[str, Any] = { "role": message.role.value, - "content": message.content, + "content": message.text, } if message.tool_calls: diff --git a/python/flink_agents/integrations/chat_models/vllm/tests/test_vllm_chat_model.py b/python/flink_agents/integrations/chat_models/vllm/tests/test_vllm_chat_model.py index 77192e10a..8e9352950 100644 --- a/python/flink_agents/integrations/chat_models/vllm/tests/test_vllm_chat_model.py +++ b/python/flink_agents/integrations/chat_models/vllm/tests/test_vllm_chat_model.py @@ -106,7 +106,7 @@ def test_native_response_format_applied_for_qwen_model() -> None: connection._client = mock_client connection.chat( - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], model="Qwen/Qwen2.5-7B-Instruct", output_schema=OutputSchema(output_schema=_Person), ) diff --git a/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py b/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py index ab79674c6..806256933 100644 --- a/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py +++ b/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py @@ -56,11 +56,11 @@ def test_watsonx_chat() -> None: """Test basic chat functionality of WatsonxChatModelConnection.""" connection = WatsonxChatModelConnection() response = connection.chat( - [ChatMessage(role=MessageRole.USER, content="Hello!")], model=test_model + [ChatMessage.of(MessageRole.USER, "Hello!")], model=test_model ) assert response is not None - assert response.content is not None - assert response.content.strip() != "" + assert response.text is not None + assert response.text.strip() != "" assert response.role == MessageRole.ASSISTANT @@ -113,7 +113,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: llm.open() response = llm.chat( - [ChatMessage(role=MessageRole.USER, content="Hello!")], top_p=0.5 + [ChatMessage.of(MessageRole.USER, "Hello!")], top_p=0.5 ) mock_model.chat.assert_called_once() @@ -126,7 +126,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: } assert response.role == MessageRole.ASSISTANT - assert response.content == "Hello there!" + assert response.text == "Hello there!" assert response.extra_args["model_name"] == test_model assert response.extra_args["promptTokens"] == 100 assert response.extra_args["completionTokens"] == 50 @@ -160,7 +160,7 @@ def test_watsonx_tool_call_response_mocked(monkeypatch: pytest.MonkeyPatch) -> N connection = _fake_connection() response = connection.chat( - [ChatMessage(role=MessageRole.USER, content="What is 1 + 2?")], + [ChatMessage.of(MessageRole.USER, "What is 1 + 2?")], model=test_model, ) @@ -203,10 +203,10 @@ def test_chat_retries_transient_failures(monkeypatch: pytest.MonkeyPatch) -> Non connection = _fake_connection(max_retries=3) response = connection.chat( - [ChatMessage(role=MessageRole.USER, content="Hello!")], model=test_model + [ChatMessage.of(MessageRole.USER, "Hello!")], model=test_model ) - assert response.content == "Recovered!" + assert response.text == "Recovered!" assert mock_model.chat.call_count == 3 assert [call.args[0] for call in sleep.call_args_list] == [5, 5] @@ -218,7 +218,7 @@ def test_chat_retries_transient_failures(monkeypatch: pytest.MonkeyPatch) -> Non mock_model.chat.reset_mock() with pytest.raises(ApiRequestFailure): connection.chat( - [ChatMessage(role=MessageRole.USER, content="Hello!")], model=test_model + [ChatMessage.of(MessageRole.USER, "Hello!")], model=test_model ) assert mock_model.chat.call_count == 1 @@ -260,8 +260,8 @@ def test_parse_tool_arguments_messy_formats() -> None: def test_convert_to_watsonx_messages_round_trip() -> None: """Test conversion of assistant tool calls and tool results to watsonx format.""" messages = [ - ChatMessage(role=MessageRole.SYSTEM, content="You are helpful."), - ChatMessage(role=MessageRole.USER, content="What is 1 + 2?"), + ChatMessage.of(MessageRole.SYSTEM, "You are helpful."), + ChatMessage.of(MessageRole.USER, "What is 1 + 2?"), ChatMessage( role=MessageRole.ASSISTANT, tool_calls=[ @@ -273,9 +273,7 @@ def test_convert_to_watsonx_messages_round_trip() -> None: } ], ), - ChatMessage( - role=MessageRole.TOOL, - content="3", + ChatMessage.of(MessageRole.TOOL, "3", extra_args={"external_id": "call_abc123"}, ), ] @@ -359,7 +357,7 @@ def test_configuration_contract(monkeypatch: pytest.MonkeyPatch) -> None: with pytest.raises(ValueError, match="additional_kwargs"): _fake_connection().chat( - [ChatMessage(role=MessageRole.USER, content="Hello!")], + [ChatMessage.of(MessageRole.USER, "Hello!")], model=test_model, additional_kwargs={"temperature": 5.0}, ) @@ -372,14 +370,14 @@ def test_additional_kwargs_reject_request_owned_fields(reserved_key: str) -> Non """Framework-owned request fields cannot be replaced by static configuration.""" with pytest.raises(ValueError, match=reserved_key): _fake_connection().chat( - [ChatMessage(role=MessageRole.USER, content="Hello!")], + [ChatMessage.of(MessageRole.USER, "Hello!")], model=test_model, additional_kwargs={reserved_key: "override"}, ) if reserved_key not in {"messages", "tools"}: with pytest.raises(ValueError, match=reserved_key): _fake_connection().chat( - [ChatMessage(role=MessageRole.USER, content="Hello!")], + [ChatMessage.of(MessageRole.USER, "Hello!")], model=test_model, **{reserved_key: "override"}, ) diff --git a/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py index 6544d2218..0bb16a6c0 100644 --- a/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py +++ b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py @@ -88,8 +88,8 @@ def convert_to_watsonx_messages( if role == MessageRole.ASSISTANT: assistant_message: Dict[str, Any] = {"role": "assistant"} - if message.content: - assistant_message["content"] = message.content + if message.text: + assistant_message["content"] = message.text if message.tool_calls: assistant_message["tool_calls"] = [ _convert_to_watsonx_tool_call(tool_call) @@ -104,12 +104,12 @@ def convert_to_watsonx_messages( watsonx_messages.append( { "role": "tool", - "content": message.content, + "content": message.text, "tool_call_id": tool_call_id, } ) else: - watsonx_messages.append({"role": role.value, "content": message.content}) + watsonx_messages.append({"role": role.value, "content": message.text}) return watsonx_messages @@ -427,9 +427,9 @@ def chat( if reasoning: extra_args["reasoning"] = reasoning - return ChatMessage( - role=MessageRole(response_message.get("role", "assistant")), - content=content, + return ChatMessage.of( + MessageRole(response_message.get("role", "assistant")), + content, tool_calls=tool_calls, extra_args=extra_args, ) diff --git a/python/flink_agents/integrations/mcp/mcp.py b/python/flink_agents/integrations/mcp/mcp.py index aa80073f2..ccccfac26 100644 --- a/python/flink_agents/integrations/mcp/mcp.py +++ b/python/flink_agents/integrations/mcp/mcp.py @@ -123,7 +123,7 @@ def format_string(self, **arguments: str) -> str: Returns a text representation of the prompt. """ text = "\n".join( - message.content for message in self.format_messages(**arguments) + message.text for message in self.format_messages(**arguments) ) return text @@ -335,7 +335,7 @@ async def _get_prompt_async( for message in prompt.messages: if isinstance(message.content, TextContent): chat_messages.append( - ChatMessage(role=message.role, content=message.content.text) + ChatMessage.of(message.role, message.content.text) ) else: err_msg = f"Unsupported content type: {type(message.content)}" diff --git a/python/flink_agents/integrations/mcp/tests/test_mcp.py b/python/flink_agents/integrations/mcp/tests/test_mcp.py index 1b013ebb3..155fa9b05 100644 --- a/python/flink_agents/integrations/mcp/tests/test_mcp.py +++ b/python/flink_agents/integrations/mcp/tests/test_mcp.py @@ -50,9 +50,7 @@ def test_mcp() -> None: assert prompt.name == "ask_sum" message = prompt.format_messages(role=MessageRole.SYSTEM, a="1", b="2") assert [ - ChatMessage( - role=MessageRole.USER, - content="Can you please calculate the sum of 1 and 2?", + ChatMessage.of(MessageRole.USER, "Can you please calculate the sum of 1 and 2?", ) ] == message tools = mcp_server.list_tools() diff --git a/python/flink_agents/plan/actions/chat_model_action.py b/python/flink_agents/plan/actions/chat_model_action.py index 0b6e7f4ae..73e5972cc 100644 --- a/python/flink_agents/plan/actions/chat_model_action.py +++ b/python/flink_agents/plan/actions/chat_model_action.py @@ -257,7 +257,7 @@ def _generate_structured_output( ) -> ChatMessage: """Deserialize output to expected output schema.""" output_schema = output_schema.output_schema - output = json.loads(_clean_llm_response(response.content)) + output = json.loads(_clean_llm_response(response.text)) if isinstance(output_schema, type) and issubclass(output_schema, BaseModel): output = output_schema.model_validate(output) @@ -494,9 +494,7 @@ async def _process_tool_response(event: ToolResponseEvent, ctx: RunnerContext) - initial_request_id, None, [ - ChatMessage( - role=MessageRole.TOOL, - content=str(response), + ChatMessage.of(MessageRole.TOOL, str(response), extra_args={"external_id": event.external_ids.get(tool_id)} if event.external_ids and event.external_ids.get(tool_id) else {}, diff --git a/python/flink_agents/plan/tests/actions/test_chat_model_action.py b/python/flink_agents/plan/tests/actions/test_chat_model_action.py index f0bba58e4..a96efb2ae 100644 --- a/python/flink_agents/plan/tests/actions/test_chat_model_action.py +++ b/python/flink_agents/plan/tests/actions/test_chat_model_action.py @@ -98,19 +98,19 @@ def test_clean_llm_response_with_multiple_lines_in_block(): def test_update_tool_call_context_stores_primitive_only(): mem = _memory() - initial = [ChatMessage(role=MessageRole.USER, content="hi")] - added = [ChatMessage(role=MessageRole.ASSISTANT, content="hello")] + initial = [ChatMessage.of(MessageRole.USER, "hi")] + added = [ChatMessage.of(MessageRole.ASSISTANT, "hello")] _update_tool_call_context(mem, uuid4(), initial, added) _assert_primitive(mem.get(_TOOL_CALL_CONTEXT)) def test_update_tool_call_context_returns_chat_messages(): mem = _memory() - initial = [ChatMessage(role=MessageRole.USER, content="hi")] - added = [ChatMessage(role=MessageRole.ASSISTANT, content="hello")] + initial = [ChatMessage.of(MessageRole.USER, "hi")] + added = [ChatMessage.of(MessageRole.ASSISTANT, "hello")] result = _update_tool_call_context(mem, uuid4(), initial, added) assert all(isinstance(message, ChatMessage) for message in result) - assert [(m.role, m.content) for m in result] == [ + assert [(m.role, m.text) for m in result] == [ (MessageRole.USER, "hi"), (MessageRole.ASSISTANT, "hello"), ] @@ -172,9 +172,9 @@ def test_request_event_key_match_after_normalization(): def test_tool_call_context_key_match_after_normalization(): mem = _memory() request_id = uuid4() - initial = [ChatMessage(role=MessageRole.USER, content="hi")] + initial = [ChatMessage.of(MessageRole.USER, "hi")] _update_tool_call_context(mem, request_id, initial, []) - extra = ChatMessage(role=MessageRole.TOOL, content="result") + extra = ChatMessage.of(MessageRole.TOOL, "result") result = _update_tool_call_context(mem, request_id, None, [extra]) assert len(result) == 2 assert len(mem.get(_TOOL_CALL_CONTEXT)[str(request_id)]) == 2 diff --git a/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py b/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py index f73648aa8..50f96af9b 100644 --- a/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py +++ b/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py @@ -167,7 +167,7 @@ def test_chat_succeeds_without_retry(self) -> None: """No retry needed: retry_count=0, total_retry_wait_sec=0, no metrics.""" chat_model = MagicMock() chat_model.chat = MagicMock( - return_value=ChatMessage(role=MessageRole.ASSISTANT, content="hello") + return_value=ChatMessage.of(MessageRole.ASSISTANT, "hello") ) ctx, sent_events, metric_group, _ = _create_mock_runner_context(chat_model) @@ -177,7 +177,7 @@ def test_chat_succeeds_without_retry(self) -> None: chat( request_id, chat_model.connection, - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], {}, None, ctx, @@ -214,7 +214,7 @@ def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: if call_count <= 1: err_msg = "transient error" raise RuntimeError(err_msg) - return ChatMessage(role=MessageRole.ASSISTANT, content="success") + return ChatMessage.of(MessageRole.ASSISTANT, "success") chat_model = MagicMock() chat_model.chat = mock_chat @@ -229,7 +229,7 @@ def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: chat( request_id, "test-model", - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], {}, None, ctx, @@ -276,7 +276,7 @@ def test_chat_exhausts_retries_and_raises(self) -> None: chat( request_id, "test-model", - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], {}, None, ctx, @@ -298,8 +298,8 @@ def test_structured_output_parse_error_retries_without_failing_llm( chat_model = MagicMock() chat_model.chat = MagicMock( side_effect=[ - ChatMessage(role=MessageRole.ASSISTANT, content="not-json"), - ChatMessage(role=MessageRole.ASSISTANT, content='{"result": 42}'), + ChatMessage.of(MessageRole.ASSISTANT, "not-json"), + ChatMessage.of(MessageRole.ASSISTANT, '{"result": 42}'), ] ) @@ -311,7 +311,7 @@ def test_structured_output_parse_error_retries_without_failing_llm( chat( uuid4(), "test-model", - [ChatMessage(role=MessageRole.USER, content="hi")], + [ChatMessage.of(MessageRole.USER, "hi")], {}, OutputSchema(output_schema=_StructuredResult), ctx, @@ -350,7 +350,7 @@ def test_default_retry_fields(self) -> None: """Default construction has retry_count=0, total_retry_wait_sec=0.""" event = ChatResponseEvent( request_id=uuid4(), - response=ChatMessage(role=MessageRole.ASSISTANT, content="test"), + response=ChatMessage.of(MessageRole.ASSISTANT, "test"), ) assert event.retry_count == 0 assert event.total_retry_wait_sec == 0 @@ -359,7 +359,7 @@ def test_with_retry_fields(self) -> None: """Full construction carries retry info.""" event = ChatResponseEvent( request_id=uuid4(), - response=ChatMessage(role=MessageRole.ASSISTANT, content="test"), + response=ChatMessage.of(MessageRole.ASSISTANT, "test"), retry_count=5, total_retry_wait_sec=31, ) @@ -391,7 +391,7 @@ def test_forwards_saved_prompt_args_to_chat(self) -> None: def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: captured_prompt_args.append(kwargs.get("prompt_args")) - return ChatMessage(role=MessageRole.ASSISTANT, content="done") + return ChatMessage.of(MessageRole.ASSISTANT, "done") chat_model = MagicMock() chat_model.chat = mock_chat @@ -420,7 +420,7 @@ def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: "_TOOL_CALL_CONTEXT", { str(initial_request_id): [ - ChatMessage(role=MessageRole.USER, content="hi").model_dump( + ChatMessage.of(MessageRole.USER, "hi").model_dump( mode="json" ) ] @@ -449,7 +449,7 @@ def test_failed_tool_response_uses_generic_response_message(self) -> None: def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: captured_messages.append(messages) - return ChatMessage(role=MessageRole.ASSISTANT, content="done") + return ChatMessage.of(MessageRole.ASSISTANT, "done") chat_model = MagicMock() chat_model.chat = mock_chat @@ -472,7 +472,7 @@ def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: "_TOOL_CALL_CONTEXT", { str(initial_request_id): [ - ChatMessage(role=MessageRole.USER, content="hi").model_dump( + ChatMessage.of(MessageRole.USER, "hi").model_dump( mode="json" ) ] @@ -494,4 +494,4 @@ def mock_chat(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: assert captured_messages tool_message = captured_messages[0][-1] assert tool_message.role == MessageRole.TOOL - assert tool_message.content == "Tool `query_order` execute failed." + assert tool_message.text == "Tool `query_order` execute failed." diff --git a/python/flink_agents/plan/tests/test_agent_plan.py b/python/flink_agents/plan/tests/test_agent_plan.py index 35226da00..1278b1858 100644 --- a/python/flink_agents/plan/tests/test_agent_plan.py +++ b/python/flink_agents/plan/tests/test_agent_plan.py @@ -231,8 +231,7 @@ def resource_type(cls) -> ResourceType: def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage: """Testing Implementation.""" - return ChatMessage( - role=MessageRole.ASSISTANT, content=self.host + " " + self.desc + return ChatMessage.of(MessageRole.ASSISTANT, self.host + " " + self.desc ) @@ -551,7 +550,7 @@ def test_get_resource() -> None: cache = ResourceCache(agent_plan.resource_providers, agent_plan.config) mock = cache.get_resource("mock", ResourceType.CHAT_MODEL) assert ( - mock.chat(ChatMessage(role=MessageRole.USER, content="")).content + mock.chat(ChatMessage.of(MessageRole.USER, "")).text == "8.8.8.8 mock resource just for testing." ) diff --git a/python/flink_agents/runtime/java/java_resource_wrapper.py b/python/flink_agents/runtime/java/java_resource_wrapper.py index 87094b4a6..c12db259a 100644 --- a/python/flink_agents/runtime/java/java_resource_wrapper.py +++ b/python/flink_agents/runtime/java/java_resource_wrapper.py @@ -90,11 +90,13 @@ def format_messages( j_MessageRole.fromValue(role.value), kwargs ) chatMessages = [ - ChatMessage( - role=MessageRole(j_chat_message.getRole().getValue()), - content=j_chat_message.getContent(), - tool_calls=j_chat_message.getToolCalls(), - extra_args=j_chat_message.getExtraArgs(), + ChatMessage.model_validate( + { + "role": MessageRole(j_chat_message.getRole().getValue()), + "blocks": j_chat_message.getBlocksAsMaps(), + "tool_calls": j_chat_message.getToolCalls(), + "extra_args": j_chat_message.getExtraArgs(), + } ) for j_chat_message in j_chat_messages ] diff --git a/python/flink_agents/runtime/memory/mem0/flink_agents_mem0_adapters.py b/python/flink_agents/runtime/memory/mem0/flink_agents_mem0_adapters.py index 2d08694bf..e84acd3d0 100644 --- a/python/flink_agents/runtime/memory/mem0/flink_agents_mem0_adapters.py +++ b/python/flink_agents/runtime/memory/mem0/flink_agents_mem0_adapters.py @@ -132,9 +132,7 @@ def generate_response( The generated response content as a string. """ chat_messages = [ - ChatMessage( - role=MessageRole(msg["role"]), - content=msg["content"], + ChatMessage.of(MessageRole(msg["role"]), msg["content"], ) for msg in messages ] @@ -146,7 +144,7 @@ def generate_response( # Mem0 expects a plain string response from generate_response. # It handles JSON parsing internally via remove_code_blocks/json.loads. - return response.content + return response.text class _OutputData(BaseModel): diff --git a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py index 3f3faf4e5..1d152f759 100644 --- a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py +++ b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py @@ -77,7 +77,7 @@ def chat(self, messages, **kwargs: Any): ), None, ) - user_text = user_msg.content if user_msg else "" + user_text = user_msg.text if user_msg else "" if "Input:" in user_text: user_text = user_text.split("Input:", 1)[1].strip() # Mem0 formats input as "role: content" — strip the role prefix. @@ -85,16 +85,12 @@ def chat(self, messages, **kwargs: Any): user_text = user_text.split(": ", 1)[1] self._last_facts = [user_text] if user_text else [] - return ChatMessage( - role=MessageRole.ASSISTANT, - content=json.dumps({"facts": self._last_facts}), + return ChatMessage.of(MessageRole.ASSISTANT, json.dumps({"facts": self._last_facts}), ) # Call 2: Memory update — return ADD for each extracted fact. memory_ops = [{"text": fact, "event": "ADD"} for fact in self._last_facts] - return ChatMessage( - role=MessageRole.ASSISTANT, - content=json.dumps({"memory": memory_ops}), + return ChatMessage.of(MessageRole.ASSISTANT, json.dumps({"memory": memory_ops}), ) @@ -372,23 +368,19 @@ def chat(self, messages, **kwargs: Any): ), None, ) - user_text = user_msg.content if user_msg else "" + user_text = user_msg.text if user_msg else "" if "Input:" in user_text: user_text = user_text.split("Input:", 1)[1].strip() if ": " in user_text: user_text = user_text.split(": ", 1)[1] self._last_facts = [user_text] if user_text else [] - return ChatMessage( - role=MessageRole.ASSISTANT, - content=json.dumps({"facts": self._last_facts}), + return ChatMessage.of(MessageRole.ASSISTANT, json.dumps({"facts": self._last_facts}), extra_args=extra_args, ) memory_ops = [{"text": fact, "event": "ADD"} for fact in self._last_facts] - return ChatMessage( - role=MessageRole.ASSISTANT, - content=json.dumps({"memory": memory_ops}), + return ChatMessage.of(MessageRole.ASSISTANT, json.dumps({"memory": memory_ops}), extra_args=extra_args, ) diff --git a/python/flink_agents/runtime/python_java_utils.py b/python/flink_agents/runtime/python_java_utils.py index fc1b06662..7ba26250b 100644 --- a/python/flink_agents/runtime/python_java_utils.py +++ b/python/flink_agents/runtime/python_java_utils.py @@ -18,7 +18,7 @@ import importlib import json import typing -from typing import Any, Dict +from typing import Any, Dict, List import cloudpickle @@ -281,16 +281,27 @@ def normalize_tool_call_id(tool_call: Dict[str, Any]) -> Dict[str, Any]: return normalized_call +def _dump_blocks(chat_message: ChatMessage) -> List[Dict[str, Any]]: + """Content blocks as plain dicts in the serialized shape, for the Java bridge.""" + return [ + block.model_dump(mode="json", exclude_none=True) + for block in chat_message.blocks + ] + + def from_java_chat_message(j_chat_message: Any) -> ChatMessage: """Convert a chat message to a python chat message.""" - return ChatMessage( - role=MessageRole(j_chat_message.getRole().getValue()), - content=j_chat_message.getContent(), - tool_calls=[ - normalize_tool_call_id(tool_call) - for tool_call in j_chat_message.getToolCalls() - ], - extra_args=j_chat_message.getExtraArgs(), + return ChatMessage.model_validate( + { + "role": MessageRole(j_chat_message.getRole().getValue()), + # Blocks cross the bridge as plain dicts in the serialized shape. + "blocks": j_chat_message.getBlocksAsMaps(), + "tool_calls": [ + normalize_tool_call_id(tool_call) + for tool_call in j_chat_message.getToolCalls() + ], + "extra_args": j_chat_message.getExtraArgs(), + } ) @@ -303,7 +314,7 @@ def to_java_chat_message(chat_message: ChatMessage) -> Any: j_MessageRole = findClass("org.apache.flink.agents.api.chat.messages.MessageRole") j_chat_message.setRole(j_MessageRole.fromValue(chat_message.role.value)) - j_chat_message.setContent(chat_message.content) + j_chat_message.setBlocksFromMaps(_dump_blocks(chat_message)) j_chat_message.setExtraArgs(chat_message.extra_args) if chat_message.tool_calls: tool_calls = [ @@ -317,7 +328,7 @@ def to_java_chat_message(chat_message: ChatMessage) -> Any: # TODO: Replace this with `to_java_chat_message()` when the `find_class` bug is fixed. def update_java_chat_message(chat_message: ChatMessage, j_chat_message: Any) -> str: """Update a Java chat message using Python chat message.""" - j_chat_message.setContent(chat_message.content) + j_chat_message.setBlocksFromMaps(_dump_blocks(chat_message)) j_chat_message.setExtraArgs(chat_message.extra_args) if chat_message.tool_calls: tool_calls = [ diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/ChatMessageEventLogSerializer.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/ChatMessageEventLogSerializer.java new file mode 100644 index 000000000..2a90ebebf --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/ChatMessageEventLogSerializer.java @@ -0,0 +1,67 @@ +/* + * 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.eventlog; + +import com.fasterxml.jackson.core.JsonGenerator; +import com.fasterxml.jackson.databind.JsonSerializer; +import com.fasterxml.jackson.databind.Module; +import com.fasterxml.jackson.databind.SerializerProvider; +import com.fasterxml.jackson.databind.module.SimpleModule; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ContentBlock; + +import java.io.IOException; + +/** + * The Event Log's {@link ChatMessage} serializer: each content block is written through its own + * {@link ContentBlock#sanitize()} projection, so media payload bytes and unsanitized URLs never + * reach the log — at any log level, VERBOSE included. Everything else (role, tool calls, extra + * args) keeps the regular shape, and the level-dependent {@link JsonTruncator} still applies to the + * result afterwards at STANDARD. + * + *

This serializer is registered only on the Event Log mappers via {@link #module()}; the global + * {@link ChatMessage} wire format — the Java/Python bridge, event serialization, state recovery — + * is untouched and preserves the complete payload. Logged output is not a faithful {@link + * ChatMessage} and must never be reconstructed into one: an inline-backed media block drops its + * {@code data} (so reconstruction fails loudly), and a URL-backed one carries only the stripped + * URL. + */ +public class ChatMessageEventLogSerializer extends JsonSerializer { + + /** The module Event Log mappers register to apply the sanitized {@link ChatMessage} shape. */ + public static Module module() { + return new SimpleModule("flink-agents-event-log-chat-messages") + .addSerializer(ChatMessage.class, new ChatMessageEventLogSerializer()); + } + + @Override + public void serialize(ChatMessage message, JsonGenerator gen, SerializerProvider serializers) + throws IOException { + gen.writeStartObject(); + gen.writeStringField("role", message.getRole().getValue()); + gen.writeArrayFieldStart("blocks"); + for (ContentBlock block : message.getBlocks()) { + gen.writeObject(block.sanitize()); + } + gen.writeEndArray(); + gen.writeObjectField("tool_calls", message.getToolCalls()); + gen.writeObjectField("extra_args", message.getExtraArgs()); + gen.writeEndObject(); + } +} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/FileEventLogger.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/FileEventLogger.java index 6e9988bfb..1023f79ca 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/FileEventLogger.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/FileEventLogger.java @@ -88,7 +88,10 @@ public class FileEventLogger implements EventLogger { private static final String DEFAULT_BASE_LOG_DIR = Paths.get(System.getProperty("java.io.tmpdir"), "flink-agents").toString(); - private static final ObjectMapper MAPPER = new ObjectMapper(); + // Event-Log-only mapper: chat messages are logged through their sanitized projection + // (media metadata instead of payload bytes) at every log level. + private static final ObjectMapper MAPPER = + new ObjectMapper().registerModule(ChatMessageEventLogSerializer.module()); private final EventLoggerConfig config; private boolean prettyPrint; diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLogger.java b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLogger.java index 769c9fd28..1f98e092f 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLogger.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/eventlog/Slf4jEventLogger.java @@ -77,7 +77,11 @@ public class Slf4jEventLogger implements EventLogger { private static final String EVENT_LOG_APPENDER_NAME = "FlinkAgentsEventLogAppender"; private static final Logger EVENT_LOG = LoggerFactory.getLogger(EVENT_LOGGER_NAME); - private static final ObjectMapper MAPPER = new ObjectMapper(); + + // Event-Log-only mapper: chat messages are logged through their sanitized projection + // (media metadata instead of payload bytes) at every log level. + private static final ObjectMapper MAPPER = + new ObjectMapper().registerModule(ChatMessageEventLogSerializer.module()); private final EventLoggerConfig config; private boolean prettyPrint; diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java index c1f235040..6eb74390e 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java @@ -24,7 +24,9 @@ import org.apache.flink.agents.api.InputEvent; import org.apache.flink.agents.api.OutputEvent; import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ImageBlock; import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.messages.TextBlock; import org.apache.flink.agents.api.context.MemoryUpdate; import org.apache.flink.agents.api.event.AgentRunBeginEvent; import org.apache.flink.agents.api.event.ChatRequestEvent; @@ -603,12 +605,12 @@ public void testBuiltInEventSerDeRoundTrip() throws Exception { // Typed getters must return typed values directly on the deserialized events. ChatRequestEvent chatRequest = (ChatRequestEvent) outputEvents.get(0); - assertEquals("hello", chatRequest.getMessages().get(0).getContent()); + assertEquals("hello", chatRequest.getMessages().get(0).getText()); assertEquals(MessageRole.USER, chatRequest.getMessages().get(0).getRole()); ChatResponseEvent chatResponse = (ChatResponseEvent) outputEvents.get(1); assertEquals(requestId, chatResponse.getRequestId()); - assertEquals("hello", chatResponse.getResponse().getContent()); + assertEquals("hello", chatResponse.getResponse().getText()); ToolResponseEvent toolResponse = (ToolResponseEvent) outputEvents.get(3); assertEquals(requestId, toolResponse.getRequestId()); @@ -620,4 +622,25 @@ public void testBuiltInEventSerDeRoundTrip() throws Exception { assertEquals("doc content", retrievalResponse.getDocuments().get(0).getContent()); assertEquals("doc-1", retrievalResponse.getDocuments().get(0).getId()); } + + @Test + void testChatMessageWithMediaBlocksRoundTrip() throws Exception { + // Content blocks are polymorphic Jackson types; the durable-execution path must + // round-trip them without interference from the serde's own type handling. + ChatMessage mixed = + new ChatMessage( + MessageRole.USER, + java.util.List.of( + TextBlock.of("look at this"), + ImageBlock.fromBase64("image/png", "aGk="))); + ActionState state = new ActionState(new ChatRequestEvent("m", java.util.List.of(mixed))); + + ActionState restored = ActionStateSerde.deserialize(ActionStateSerde.serialize(state)); + + ChatMessage restoredMessage = + ((ChatRequestEvent) restored.getTaskEvent()).getMessages().get(0); + assertEquals(mixed, restoredMessage); + assertEquals("look at this", restoredMessage.getText()); + assertEquals(ImageBlock.class, restoredMessage.getBlocks().get(1).getClass()); + } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/condition/ConditionEvaluatorTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/condition/ConditionEvaluatorTest.java index ea1f48ffd..283235fc0 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/condition/ConditionEvaluatorTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/condition/ConditionEvaluatorTest.java @@ -177,7 +177,7 @@ void typedEventAttributesMatchJsonRoundTrip() throws Exception { List sources = List.of( "request_id == '550e8400-e29b-41d4-a716-446655440000'", - "response.content == 'hello'", + "response.blocks[0].text == 'hello'", "response.role == 'assistant'"); EvaluatorHarness testEvaluator = new EvaluatorHarness(sources, ConditionEvaluationFailureStrategy.FAIL); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/FileEventLoggerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/FileEventLoggerTest.java index dd2d24a6c..985493c57 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/FileEventLoggerTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/eventlog/FileEventLoggerTest.java @@ -25,7 +25,12 @@ import org.apache.flink.agents.api.EventContext; import org.apache.flink.agents.api.InputEvent; import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.DocumentBlock; +import org.apache.flink.agents.api.chat.messages.ImageBlock; +import org.apache.flink.agents.api.chat.messages.TextBlock; import org.apache.flink.agents.api.configuration.AgentConfigOptions; +import org.apache.flink.agents.api.event.ChatRequestEvent; import org.apache.flink.agents.api.logger.EventLogger; import org.apache.flink.agents.api.logger.EventLoggerConfig; import org.apache.flink.agents.api.logger.EventLoggerOpenParams; @@ -419,6 +424,57 @@ void testVerboseLevelNoTruncation() throws Exception { "this is a very long string that exceeds 10", attrsNode.get("customData").asText()); } + @Test + void testMediaPayloadsNeverReachTheLogAtStandard() throws Exception { + assertMediaPayloadsSanitizedAt("STANDARD"); + } + + @Test + void testMediaPayloadsNeverReachTheLogAtVerbose() throws Exception { + // VERBOSE lifts truncation, not sanitization: payload bytes and raw URLs stay out. + assertMediaPayloadsSanitizedAt("VERBOSE"); + } + + private void assertMediaPayloadsSanitizedAt(String level) throws Exception { + Map agentConfig = new HashMap<>(); + agentConfig.put("event-log.level", level); + + config = buildConfig(agentConfig); + logger = new FileEventLogger(config); + logger.open(openParams); + + String payload = "aW5saW5lLXBheWxvYWQtYnl0ZXM="; + String signedUrl = "https://user:secret@example.org/media/cat.png?X-Amz-Signature=abc123"; + ChatMessage message = + ChatMessage.user( + List.of( + TextBlock.of("what is in this picture?"), + ImageBlock.fromBase64("image/png", payload), + new DocumentBlock( + "application/pdf", null, signedUrl, "cat.pdf", 42L, null))); + ChatRequestEvent event = new ChatRequestEvent("test-model", List.of(message)); + + append(logger, event, null); + logger.flush(); + + Path logFile = getExpectedLogFilePath(); + String line = Files.readAllLines(logFile).get(0); + assertFalse(line.contains(payload), "Inline payload bytes must never be logged"); + assertFalse(line.contains("secret"), "URL credentials must never be logged"); + assertFalse(line.contains("X-Amz-Signature"), "URL query strings must never be logged"); + + JsonNode logged = objectMapper.readTree(line).get("eventAttributes").get("messages").get(0); + assertEquals("what is in this picture?", logged.get("blocks").get(0).get("text").asText()); + JsonNode image = logged.get("blocks").get(1); + assertEquals("image/png", image.get("media_type").asText()); + assertFalse(image.has("data"), "Inline data is dropped entirely, not masked"); + assertTrue(image.get("size_bytes").isNumber(), "Derived size metadata should be logged"); + JsonNode document = logged.get("blocks").get(2); + assertEquals("https://example.org/media/cat.png", document.get("url").asText()); + assertEquals("cat.pdf", document.get("name").asText()); + assertEquals(42L, document.get("size_bytes").asLong()); + } + @Test void testOffLevelSkipsEvent() throws Exception { // Given - config with OFF level