Skip to content

Commit 6e3656d

Browse files
pyek-botpyek-bot
authored
Add support for date time injection for agents (opensearch-project#4008)
* feat: add support for date and date format injection into chat and per agent Signed-off-by: Pavan Yekbote <pybot@amazon.com> * chore: clean up code Signed-off-by: Pavan Yekbote <pybot@amazon.com> * code quality: minor changes to improve code Signed-off-by: Pavan Yekbote <pybot@amazon.com> * minor: use isblank Signed-off-by: Pavan Yekbote <pybot@amazon.com> * fix: change default to iso format Signed-off-by: Pavan Yekbote <pybot@amazon.com> * fix: test Signed-off-by: Pavan Yekbote <pybot@amazon.com> --------- Signed-off-by: Pavan Yekbote <pybot@amazon.com>
1 parent 0db22c9 commit 6e3656d

File tree

6 files changed

+235
-5
lines changed

6 files changed

+235
-5
lines changed

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
package org.opensearch.ml.engine.algorithms.agent;
77

8+
import static org.apache.commons.lang3.StringUtils.isBlank;
89
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
910
import static org.opensearch.ml.common.CommonValue.MCP_CONNECTORS_FIELD;
1011
import static org.opensearch.ml.common.CommonValue.MCP_CONNECTOR_ID_FIELD;
@@ -34,6 +35,9 @@
3435
import java.security.AccessController;
3536
import java.security.PrivilegedActionException;
3637
import java.security.PrivilegedExceptionAction;
38+
import java.time.Instant;
39+
import java.time.ZoneId;
40+
import java.time.format.DateTimeFormatter;
3741
import java.util.ArrayList;
3842
import java.util.Collection;
3943
import java.util.Collections;
@@ -133,6 +137,10 @@ public class AgentUtils {
133137
// For function calling, do not escape the below params in connector by default
134138
public static final String DEFAULT_NO_ESCAPE_PARAMS = "_chat_history,_tools,_interactions,tool_configs";
135139

140+
public static final String DEFAULT_DATETIME_FORMAT = "yyyy-MM-dd'T'HH:mm:ss'Z'";
141+
public static final String DEFAULT_DATETIME_PREFIX = "Current date and time: ";
142+
private static final ZoneId UTC_ZONE = ZoneId.of("UTC");
143+
136144
public static String addExamplesToPrompt(Map<String, String> parameters, String prompt) {
137145
Map<String, String> examplesMap = new HashMap<>();
138146
if (parameters.containsKey(EXAMPLES)) {
@@ -946,4 +954,45 @@ public static void cleanUpResource(Map<String, Tool> tools) {
946954
}
947955
}
948956
}
957+
958+
/**
959+
* Generates a formatted current date and time string in UTC timezone.
960+
*
961+
* <p>This method returns the current date and time formatted according to the provided pattern.
962+
* If no format is provided or the format is invalid, it uses the default format:
963+
* "yyyy-MM-dd'T'HH:mm:ss'Z'" (e.g., "2024-01-15T14:30:00Z").
964+
*
965+
* <p>The method always returns the time in UTC timezone regardless of the system's local timezone.
966+
*
967+
* @param dateFormat The date format pattern to use. Can be null or empty to use the default format.
968+
* Must be a valid {@link java.time.format.DateTimeFormatter} pattern.
969+
* Examples:
970+
* <ul>
971+
* <li>"yyyy-MM-dd HH:mm:ss" → "2024-01-15 14:30:00"</li>
972+
* <li>"EEEE, MMMM d, yyyy 'at' h:mm a z" → "Monday, January 15, 2024 at 2:30 PM UTC"</li>
973+
* <li>"MM/dd/yyyy h:mm a" → "01/15/2024 2:30 PM"</li>
974+
* <li>"yyyy-MM-dd'T'HH:mm:ss'Z'" → "2024-01-15T14:30:00Z"</li>
975+
* </ul>
976+
* @return A string containing the current date and time prefixed with "Current date and time: ".
977+
* If the provided format is invalid, a warning is logged and the default format is used.
978+
* @see java.time.format.DateTimeFormatter
979+
* @see java.time.ZoneId
980+
*/
981+
public static String getCurrentDateTime(String dateFormat) {
982+
Instant now = Instant.now();
983+
DateTimeFormatter formatter;
984+
985+
if (!isBlank(dateFormat)) {
986+
try {
987+
formatter = DateTimeFormatter.ofPattern(dateFormat).withZone(UTC_ZONE);
988+
} catch (IllegalArgumentException e) {
989+
log.warn("Invalid date format provided: {}. Using default format.", dateFormat);
990+
formatter = DateTimeFormatter.ofPattern(DEFAULT_DATETIME_FORMAT).withZone(UTC_ZONE);
991+
}
992+
} else {
993+
formatter = DateTimeFormatter.ofPattern(DEFAULT_DATETIME_FORMAT).withZone(UTC_ZONE);
994+
}
995+
996+
return DEFAULT_DATETIME_PREFIX + formatter.format(now);
997+
}
949998
}

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.cleanUpResource;
2323
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.constructToolParams;
2424
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.createTools;
25+
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getCurrentDateTime;
2526
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getMcpToolSpecs;
2627
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getMessageHistoryLimit;
2728
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getMlToolSpecs;
@@ -82,6 +83,8 @@
8283
import org.opensearch.remote.metadata.client.SdkClient;
8384
import org.opensearch.transport.client.Client;
8485

86+
import com.google.common.annotations.VisibleForTesting;
87+
8588
import lombok.Data;
8689
import lombok.NoArgsConstructor;
8790
import lombok.extern.log4j.Log4j2;
@@ -117,6 +120,9 @@ public class MLChatAgentRunner implements MLAgentRunner {
117120
public static final String CHAT_HISTORY_RESPONSE_TEMPLATE = "chat_history_template.ai_response";
118121
public static final String CHAT_HISTORY_MESSAGE_PREFIX = "${_chat_history.message.";
119122
public static final String LLM_INTERFACE = "_llm_interface";
123+
public static final String INJECT_DATETIME_FIELD = "inject_datetime";
124+
public static final String DATETIME_FORMAT_FIELD = "datetime_format";
125+
public static final String SYSTEM_PROMPT_FIELD = "system_prompt";
120126

121127
private static final String DEFAULT_MAX_ITERATIONS = "10";
122128

@@ -748,7 +754,8 @@ private static String constructLLMPrompt(Map<String, Tool> tools, Map<String, St
748754
return prompt;
749755
}
750756

751-
private static Map<String, String> constructLLMParams(LLMSpec llm, Map<String, String> parameters) {
757+
@VisibleForTesting
758+
static Map<String, String> constructLLMParams(LLMSpec llm, Map<String, String> parameters) {
752759
Map<String, String> tmpParameters = new HashMap<>();
753760
if (llm.getParameters() != null) {
754761
tmpParameters.putAll(llm.getParameters());
@@ -774,6 +781,23 @@ private static Map<String, String> constructLLMParams(LLMSpec llm, Map<String, S
774781
);
775782
}
776783

784+
boolean injectDate = Boolean.parseBoolean(tmpParameters.getOrDefault(INJECT_DATETIME_FIELD, "false"));
785+
if (injectDate) {
786+
String dateFormat = tmpParameters.get(DATETIME_FORMAT_FIELD);
787+
String currentDateTime = getCurrentDateTime(dateFormat);
788+
// If system_prompt exists, inject datetime into it
789+
if (tmpParameters.containsKey(SYSTEM_PROMPT_FIELD)) {
790+
String systemPrompt = tmpParameters.get(SYSTEM_PROMPT_FIELD);
791+
systemPrompt = systemPrompt + "\n\n" + currentDateTime;
792+
tmpParameters.put(SYSTEM_PROMPT_FIELD, systemPrompt);
793+
} else {
794+
// Otherwise inject datetime into prompt_prefix
795+
String promptPrefix = tmpParameters.getOrDefault(PROMPT_PREFIX, PromptTemplate.PROMPT_TEMPLATE_PREFIX);
796+
promptPrefix = promptPrefix + "\n\n" + currentDateTime;
797+
tmpParameters.put(PROMPT_PREFIX, promptPrefix);
798+
}
799+
}
800+
777801
tmpParameters.putIfAbsent(PROMPT_PREFIX, PromptTemplate.PROMPT_TEMPLATE_PREFIX);
778802
tmpParameters.putIfAbsent(PROMPT_SUFFIX, PromptTemplate.PROMPT_TEMPLATE_SUFFIX);
779803
tmpParameters.putIfAbsent(RESPONSE_FORMAT_INSTRUCTION, PromptTemplate.PROMPT_FORMAT_INSTRUCTION);

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.LLM_RESPONSE_FILTER;
1818
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.cleanUpResource;
1919
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.createTools;
20+
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getCurrentDateTime;
2021
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getMcpToolSpecs;
2122
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getMlToolSpecs;
2223
import static org.opensearch.ml.engine.algorithms.agent.MLChatAgentRunner.LLM_INTERFACE;
@@ -98,10 +99,11 @@ public class MLPlanExecuteAndReflectAgentRunner implements MLAgentRunner {
9899
private String reflectPromptTemplate;
99100
private String plannerWithHistoryPromptTemplate;
100101

101-
// defaults
102-
private static final String DEFAULT_PLANNER_SYSTEM_PROMPT =
102+
@VisibleForTesting
103+
static final String DEFAULT_PLANNER_SYSTEM_PROMPT =
103104
"You are part of an OpenSearch cluster. When you deliver your final result, include a comprehensive report. This report MUST:\\n1. List every analysis or step you performed.\\n2. Summarize the inputs, methods, tools, and data used at each step.\\n3. Include key findings from all intermediate steps — do NOT omit them.\\n4. Clearly explain how the steps led to your final conclusion.\\n5. Return the full analysis and conclusion in the 'result' field, even if some of this was mentioned earlier.\\n\\nThe final response should be fully self-contained and detailed, allowing a user to understand the full investigation without needing to reference prior messages. Always respond in JSON format.";
104-
private static final String DEFAULT_EXECUTOR_SYSTEM_PROMPT =
105+
@VisibleForTesting
106+
static final String DEFAULT_EXECUTOR_SYSTEM_PROMPT =
105107
"You are a dedicated helper agent working as part of a plan‑execute‑reflect framework. Your role is to receive a discrete task, execute all necessary internal reasoning or tool calls, and return a single, final response that fully addresses the task. You must never return an empty response. If you are unable to complete the task or retrieve meaningful information, you must respond with a clear explanation of the issue or what was missing. Under no circumstances should you end your reply with a question or ask for more information. If you search any index, always include the raw documents in the final result instead of summarizing the content. This is critical to give visibility into what the query retrieved.";
106108
private static final String DEFAULT_NO_ESCAPE_PARAMS = "tool_configs,_tools";
107109
private static final String DEFAULT_MAX_STEPS_EXECUTED = "20";
@@ -136,6 +138,8 @@ public class MLPlanExecuteAndReflectAgentRunner implements MLAgentRunner {
136138
public static final String REFLECT_PROMPT_TEMPLATE_FIELD = "reflect_prompt_template";
137139
public static final String PLANNER_WITH_HISTORY_TEMPLATE_FIELD = "planner_with_history_template";
138140
public static final String EXECUTOR_MAX_ITERATIONS_FIELD = "executor_max_iterations";
141+
public static final String INJECT_DATETIME_FIELD = "inject_datetime";
142+
public static final String DATETIME_FORMAT_FIELD = "datetime_format";
139143

140144
public MLPlanExecuteAndReflectAgentRunner(
141145
Client client,
@@ -170,7 +174,22 @@ void setupPromptParameters(Map<String, String> params) {
170174

171175
String userPrompt = params.get(QUESTION_FIELD);
172176
params.put(USER_PROMPT_FIELD, userPrompt);
173-
params.put(SYSTEM_PROMPT_FIELD, params.getOrDefault(SYSTEM_PROMPT_FIELD, DEFAULT_PLANNER_SYSTEM_PROMPT));
177+
178+
boolean injectDate = Boolean.parseBoolean(params.getOrDefault(INJECT_DATETIME_FIELD, "false"));
179+
String dateFormat = params.get(DATETIME_FORMAT_FIELD);
180+
String currentDateTime = injectDate ? getCurrentDateTime(dateFormat) : "";
181+
182+
String plannerSystemPrompt = params.getOrDefault(SYSTEM_PROMPT_FIELD, DEFAULT_PLANNER_SYSTEM_PROMPT);
183+
if (injectDate) {
184+
plannerSystemPrompt = String.format("%s\n\n%s", plannerSystemPrompt, currentDateTime);
185+
}
186+
params.put(SYSTEM_PROMPT_FIELD, plannerSystemPrompt);
187+
188+
String executorSystemPrompt = params.getOrDefault(EXECUTOR_SYSTEM_PROMPT_FIELD, DEFAULT_EXECUTOR_SYSTEM_PROMPT);
189+
if (injectDate) {
190+
executorSystemPrompt = String.format("%s\n\n%s", executorSystemPrompt, currentDateTime);
191+
}
192+
params.put(EXECUTOR_SYSTEM_PROMPT_FIELD, executorSystemPrompt);
174193

175194
if (params.get(PLANNER_PROMPT_FIELD) != null) {
176195
this.plannerPrompt = params.get(PLANNER_PROMPT_FIELD);

‎ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/AgentUtilsTest.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import static org.opensearch.ml.common.CommonValue.MCP_CONNECTORS_FIELD;
1818
import static org.opensearch.ml.common.CommonValue.MCP_CONNECTOR_ID_FIELD;
1919
import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
20+
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.DEFAULT_DATETIME_PREFIX;
2021
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.LLM_FINISH_REASON_PATH;
2122
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.LLM_FINISH_REASON_TOOL_USE;
2223
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.LLM_GEN_INPUT;
@@ -1665,4 +1666,30 @@ public void testCreateTool_ToolNotFound() {
16651666

16661667
assertThrows(IllegalArgumentException.class, () -> AgentUtils.createTool(toolFactories, new HashMap<>(), toolSpec, "test_tenant"));
16671668
}
1669+
1670+
@Test
1671+
public void testGetCurrentDateTime_WithInvalidFormats() {
1672+
// null
1673+
String result = AgentUtils.getCurrentDateTime(null);
1674+
Assert.assertNotNull(result);
1675+
Assert.assertTrue(result.startsWith(DEFAULT_DATETIME_PREFIX));
1676+
1677+
// empty
1678+
result = AgentUtils.getCurrentDateTime("");
1679+
Assert.assertNotNull(result);
1680+
Assert.assertTrue(result.startsWith(DEFAULT_DATETIME_PREFIX));
1681+
1682+
// invalid
1683+
result = AgentUtils.getCurrentDateTime("invalid-format");
1684+
Assert.assertNotNull(result);
1685+
Assert.assertTrue(result.startsWith(DEFAULT_DATETIME_PREFIX));
1686+
}
1687+
1688+
@Test
1689+
public void testGetCurrentDateTime_WithValidFormat() {
1690+
String result = AgentUtils.getCurrentDateTime("EEEE, MMMM d, yyyy 'at' h:mm a z");
1691+
Assert.assertNotNull(result);
1692+
Assert.assertTrue(result.startsWith(DEFAULT_DATETIME_PREFIX));
1693+
Assert.assertTrue(result.contains("UTC"));
1694+
}
16681695
}

‎ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import static org.mockito.Mockito.never;
1717
import static org.mockito.Mockito.verify;
1818
import static org.mockito.Mockito.when;
19+
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.DEFAULT_DATETIME_PREFIX;
1920
import static org.opensearch.ml.engine.algorithms.agent.MLAgentExecutor.MESSAGE_HISTORY_LIMIT;
2021
import static org.opensearch.ml.engine.memory.ConversationIndexMemory.LAST_N_INTERACTIONS;
2122

@@ -977,4 +978,63 @@ private Answer generateToolFailure(Exception e) {
977978
};
978979
}
979980

981+
@Test
982+
public void testConstructLLMParams_WithSystemPromptAndDateTimeInjection() {
983+
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
984+
Map<String, String> parameters = new HashMap<>();
985+
parameters.put(MLChatAgentRunner.SYSTEM_PROMPT_FIELD, "You are a helpful assistant.");
986+
parameters.put(MLChatAgentRunner.INJECT_DATETIME_FIELD, "true");
987+
988+
Map<String, String> result = MLChatAgentRunner.constructLLMParams(llmSpec, parameters);
989+
990+
Assert.assertNotNull(result);
991+
Assert.assertTrue(result.containsKey(MLChatAgentRunner.SYSTEM_PROMPT_FIELD));
992+
String systemPrompt = result.get(MLChatAgentRunner.SYSTEM_PROMPT_FIELD);
993+
Assert.assertTrue(systemPrompt.startsWith("You are a helpful assistant."));
994+
Assert.assertTrue(systemPrompt.contains(DEFAULT_DATETIME_PREFIX));
995+
}
996+
997+
@Test
998+
public void testConstructLLMParams_WithoutSystemPromptAndDateTimeInjection() {
999+
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
1000+
Map<String, String> parameters = new HashMap<>();
1001+
parameters.put(MLChatAgentRunner.INJECT_DATETIME_FIELD, "true");
1002+
1003+
Map<String, String> result = MLChatAgentRunner.constructLLMParams(llmSpec, parameters);
1004+
1005+
Assert.assertNotNull(result);
1006+
Assert.assertTrue(result.containsKey(AgentUtils.PROMPT_PREFIX));
1007+
String promptPrefix = result.get(AgentUtils.PROMPT_PREFIX);
1008+
Assert.assertTrue(promptPrefix.contains(DEFAULT_DATETIME_PREFIX));
1009+
}
1010+
1011+
@Test
1012+
public void testConstructLLMParams_DateTimeInjectionDisabled() {
1013+
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
1014+
Map<String, String> parameters = new HashMap<>();
1015+
parameters.put(MLChatAgentRunner.INJECT_DATETIME_FIELD, "false");
1016+
parameters.put(MLChatAgentRunner.SYSTEM_PROMPT_FIELD, "You are a helpful assistant.");
1017+
1018+
Map<String, String> result = MLChatAgentRunner.constructLLMParams(llmSpec, parameters);
1019+
1020+
Assert.assertNotNull(result);
1021+
Assert.assertTrue(result.containsKey(MLChatAgentRunner.SYSTEM_PROMPT_FIELD));
1022+
String systemPrompt = result.get(MLChatAgentRunner.SYSTEM_PROMPT_FIELD);
1023+
Assert.assertEquals("You are a helpful assistant.", systemPrompt);
1024+
Assert.assertFalse(systemPrompt.contains(DEFAULT_DATETIME_PREFIX));
1025+
}
1026+
1027+
@Test
1028+
public void testConstructLLMParams_DefaultValues() {
1029+
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
1030+
Map<String, String> parameters = new HashMap<>();
1031+
1032+
Map<String, String> result = MLChatAgentRunner.constructLLMParams(llmSpec, parameters);
1033+
1034+
Assert.assertNotNull(result);
1035+
Assert.assertTrue(result.containsKey(AgentUtils.PROMPT_PREFIX));
1036+
Assert.assertTrue(result.containsKey(AgentUtils.PROMPT_SUFFIX));
1037+
Assert.assertTrue(result.containsKey(AgentUtils.RESPONSE_FORMAT_INSTRUCTION));
1038+
Assert.assertTrue(result.containsKey(AgentUtils.TOOL_RESPONSE));
1039+
}
9801040
}

0 commit comments

Comments
 (0)