fix: enforce configured model input modalities

src/prerna/engine/impl/model/AbstractModelEngine.java

@@ -28,11 +28,15 @@
 package prerna.engine.impl.model;
 
 import java.time.ZonedDateTime;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Properties;
+import java.util.Set;
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -46,6 +50,10 @@
 import prerna.engine.impl.model.inferencetracking.ModelInferenceLogsUtils;
 import prerna.engine.impl.model.message.AbstractMessage;
 import prerna.engine.impl.model.message.InputMessage;
+import prerna.engine.impl.model.message.MediaMessagePart;
+import prerna.engine.impl.model.message.MessageInputMedia;
+import prerna.engine.impl.model.message.MessagePart;
+import prerna.engine.impl.model.message.MessagePartType;
 import prerna.engine.impl.model.message.MessageUtils;
 import prerna.engine.impl.model.message.ResponseMessage;
 import prerna.engine.impl.model.responses.AskErrorModelEngineResponse;
@@ -61,6 +69,8 @@
 public abstract class AbstractModelEngine extends AbstractEngine implements IModelEngine {
 
 	private static final Logger classLogger = LogManager.getLogger(AbstractModelEngine.class);
+	private static final String FILE_INPUT_MODALITY = "FILE";
+	private static final String PDF_INPUT_MODALITY = "PDF";
 
 	public static final String OPEN_AI_KEY = "OPEN_AI_KEY";
 	public static final String AWS_SECRET_KEY = "AWS_SECRET_KEY";
@@ -133,6 +143,12 @@ public abstract class AbstractModelEngine extends AbstractEngine implements IMod
 	 */
 	protected Boolean temperatureSupported = null;
 
+	/**
+	 * Input modalities configured in MODELMETADATA. Null means the metadata does
+	 * not restrict request content.
+	 */
+	protected Set<String> inputModalities = null;
+
 	@Override
 	public void open(Properties smssProp) throws Exception {
 		super.open(smssProp);
@@ -198,6 +214,7 @@ private void fillModelSettingsFromMetadata() {
 		fillIfMissing("context_window", metadata.get("contextWindow"));
 		fillIfMissing("max_tokens", metadata.get("maxOutputTokens"));
 		this.builtinTools = metadata.get("builtinTools");
+		this.inputModalities = toModalitySet(metadata.get("inputModalities"));
 
 		if (metadata.get("reasoning") instanceof Boolean) {
 			this.reasoning = (Boolean) metadata.get("reasoning");
@@ -212,6 +229,19 @@ private void fillModelSettingsFromMetadata() {
 		}
 	}
 
+	private static Set<String> toModalitySet(Object value) {
+		if (!(value instanceof Collection<?>)) {
+			return null;
+		}
+		Set<String> modalities = new LinkedHashSet<>();
+		for (Object modality : (Collection<?>) value) {
+			if (modality != null && !modality.toString().isBlank()) {
+				modalities.add(modality.toString().trim().toUpperCase(Locale.ROOT));
+			}
+		}
+		return modalities.isEmpty() ? null : modalities;
+	}
+
 	/**
 	 * Set the smss property to the metadata value only when the smss file does not
 	 * already define a non-empty value for the key.
@@ -467,6 +497,8 @@ public AskModelEngineResponse askRoom(String question, Room room, AbstractMessag
 				question = MessageUtils.toJsonArray(room.getMessages());
 			}
 
+			validateInputModalities(room.getMessages(), inputMessage);
+
 			ZonedDateTime inputTime = ZonedDateTime.now();
 			AskModelEngineResponse askModelResponse = askCall(question, null, context, room.getInsight(), room.getId(),
 					parameters);
@@ -573,6 +605,64 @@ public AskModelEngineResponse askRoom(String question, Room room, AbstractMessag
 		}
 	}
 
+	void validateInputModalities(List<AbstractMessage> messages, AbstractMessage inputMessage) {
+		if (this.inputModalities == null) {
+			return;
+		}
+		List<AbstractMessage> requestMessages = messages == null ? List.of() : messages;
+		if (inputMessage != null) {
+			requestMessages = MessageUtils.getMessageBranchWithNewMessage(requestMessages, inputMessage);
+		}
+		for (AbstractMessage message : requestMessages) {
+			validateInputModalities(message);
+		}
+	}
+
+	private void validateInputModalities(AbstractMessage message) {
+		if (message == null) {
+			return;
+		}
+		for (MessagePart part : message.getParts()) {
+			String modality = modalityFor(part);
+			if (modality != null && !this.inputModalities.contains(modality)) {
+				String model = this.engineName == null || this.engineName.isBlank() ? this.engineId : this.engineName;
+				throw new IllegalArgumentException("Model " + model + " does not allow " + modality
+						+ " input. Configured input modalities: " + this.inputModalities);
+			}
+		}
+	}
+
+	private static String modalityFor(MessagePart part) {
+		if (part == null) {
+			return null;
+		}
+		return switch (part.getType()) {
+		case TEXT, SYSTEM -> MessagePartType.TEXT.name();
+		case MEDIA -> part instanceof MediaMessagePart ? modalityFor((MediaMessagePart) part) : FILE_INPUT_MODALITY;
+		default -> null;
+		};
+	}
+
+	private static String modalityFor(MediaMessagePart part) {
+		MessageInputMedia media = part.getMediaInfo();
+		String mimeType = media == null ? null : media.getMimeType();
+		if (mimeType == null || mimeType.isBlank()) {
+			// URL media currently represents image input and does not carry a MIME type.
+			return AskModelEngineResponse.IMAGE;
+		}
+
+		String[] mimeParts = mimeType.split("/", 2);
+		String mimeFamily = mimeParts[0].trim().toUpperCase(Locale.ROOT);
+		if (mimeFamily.equals(AskModelEngineResponse.IMAGE) || mimeFamily.equals("AUDIO")
+				|| mimeFamily.equals("VIDEO")) {
+			return mimeFamily;
+		}
+		if (mimeParts.length == 2 && PDF_INPUT_MODALITY.equalsIgnoreCase(mimeParts[1].trim())) {
+			return PDF_INPUT_MODALITY;
+		}
+		return FILE_INPUT_MODALITY;
+	}
+
 	@Override
 	@Deprecated
 	public AskModelEngineResponse ask(String question, String context, Insight insight,

test/prerna/engine/impl/model/AbstractModelEngineInputModalityUnitTests.java

@@ -0,0 +1,158 @@
+/*******************************************************************************
+ * Copyright 2015 Defense Health Agency (DHA)
+ *
+ * If your use of this software does not include any GPLv2 components:
+ * 	Licensed 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.
+ * ----------------------------------------------------------------------------
+ * If your use of this software includes any GPLv2 components:
+ * 	This program is free software; you can redistribute it and/or
+ * 	modify it under the terms of the GNU General Public License
+ * 	as published by the Free Software Foundation; either version 2
+ * 	of the License, or (at your option) any later version.
+ *
+ * 	This program is distributed in the hope that it will be useful,
+ * 	but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * 	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * 	GNU General Public License for more details.
+ *******************************************************************************/
+package prerna.engine.impl.model;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import org.junit.jupiter.api.Test;
+
+import com.google.gson.Gson;
+
+import prerna.engine.api.ModelTypeEnum;
+import prerna.engine.impl.model.message.InputMessage;
+import prerna.engine.impl.model.message.MediaMessagePart;
+import prerna.engine.impl.model.message.MessageInputMedia;
+import prerna.engine.impl.model.message.MessagePartType;
+import prerna.engine.impl.model.message.TextMessagePart;
+import prerna.engine.impl.model.responses.AskModelEngineResponse;
+import prerna.engine.impl.model.responses.EmbeddingsModelEngineResponse;
+import prerna.om.Insight;
+
+class AbstractModelEngineInputModalityUnitTests {
+
+	@Test
+	void rejectsImageWhenModelOnlyAllowsTextInput() {
+		TestModelEngine engine = new TestModelEngine(Set.of(MessagePartType.TEXT.name()));
+		InputMessage message = newMessage();
+		message.addPart(new TextMessagePart("describe this"));
+		message.addPart(new MediaMessagePart(MessageInputMedia.fromUrl("https://example.com/image.png")));
+
+		IllegalArgumentException exception = assertThrows(IllegalArgumentException.class,
+				() -> engine.validateInputModalities(List.of(), message));
+
+		assertEquals("Model test-model does not allow IMAGE input. Configured input modalities: [TEXT]",
+				exception.getMessage());
+	}
+
+	@Test
+	void acceptsImageWhenModelAllowsImageInput() {
+		TestModelEngine engine = new TestModelEngine(
+				Set.of(MessagePartType.TEXT.name(), AskModelEngineResponse.IMAGE));
+		InputMessage message = newMessage();
+		message.addPart(new TextMessagePart("describe this"));
+		message.addPart(new MediaMessagePart(MessageInputMedia.fromUrl("https://example.com/image.png")));
+
+		assertDoesNotThrow(() -> engine.validateInputModalities(List.of(), message));
+	}
+
+	@Test
+	void rejectsPdfWhenModelDoesNotAllowPdfInput() {
+		TestModelEngine engine = new TestModelEngine(
+				Set.of(MessagePartType.TEXT.name(), AskModelEngineResponse.IMAGE));
+		InputMessage message = newMessage();
+		message.addPart(new MediaMessagePart(pdfMedia()));
+
+		IllegalArgumentException exception = assertThrows(IllegalArgumentException.class,
+				() -> engine.validateInputModalities(List.of(), message));
+
+		assertTrue(exception.getMessage().contains("does not allow PDF input"));
+	}
+
+	@Test
+	void doesNotRestrictInputWhenMetadataDoesNotConfigureModalities() {
+		TestModelEngine engine = new TestModelEngine(null);
+		InputMessage message = newMessage();
+		message.addPart(new MediaMessagePart(MessageInputMedia.fromUrl("https://example.com/image.png")));
+
+		assertDoesNotThrow(() -> engine.validateInputModalities(List.of(), message));
+	}
+
+	@Test
+	void ignoresUnsupportedPartsOnConversationBranchesNotSentToModel() {
+		TestModelEngine engine = new TestModelEngine(Set.of(MessagePartType.TEXT.name()));
+		InputMessage root = newMessage();
+		root.addPart(new TextMessagePart("root"));
+		InputMessage imageBranch = newMessage();
+		imageBranch.setParentMessageId(root.getMessageId());
+		imageBranch.addPart(new MediaMessagePart(MessageInputMedia.fromUrl("https://example.com/image.png")));
+		InputMessage textBranch = newMessage();
+		textBranch.setParentMessageId(root.getMessageId());
+		textBranch.addPart(new TextMessagePart("continue here"));
+
+		assertDoesNotThrow(
+				() -> engine.validateInputModalities(List.of(root, imageBranch), textBranch));
+	}
+
+	private static InputMessage newMessage() {
+		Room room = new Room();
+		room.setId("test-room");
+		return InputMessage.builder(room).build();
+	}
+
+	private static MessageInputMedia pdfMedia() {
+		return new Gson().fromJson("{\"mimeType\":\"application/pdf\"}", MessageInputMedia.class);
+	}
+
+	private static class TestModelEngine extends AbstractModelEngine {
+
+		private TestModelEngine(Set<String> inputModalities) {
+			this.inputModalities = inputModalities;
+			setEngineName("test-model");
+		}
+
+		@Override
+		protected AskModelEngineResponse askCall(String question, Object fullPrompt, String context, Insight insight,
+				String roomId, Map<String, Object> hyperParameters) {
+			return null;
+		}
+
+		@Override
+		protected EmbeddingsModelEngineResponse embeddingsCall(List<String> stringsToEmbed, Insight insight,
+				Map<String, Object> parameters) {
+			return null;
+		}
+
+		@Override
+		public ModelTypeEnum getModelType() {
+			return ModelTypeEnum.OPEN_AI;
+		}
+
+		@Override
+		public void close() throws IOException {
+			// Nothing to close in this test engine.
+		}
+	}
+}