All Downloads are FREE. Search and download functionalities are using the official Maven repository.

prerna.engine.impl.model.LLMInstructReactor Maven / Gradle / Ivy

The newest version!
package prerna.engine.impl.model;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import prerna.auth.utils.SecurityEngineUtils;
import prerna.auth.utils.SecurityProjectUtils;
import prerna.engine.api.IModelEngine;
import prerna.engine.api.IRawSelectWrapper;
import prerna.reactor.AbstractReactor;
import prerna.sablecc2.om.GenRowStruct;
import prerna.sablecc2.om.PixelDataType;
import prerna.sablecc2.om.PixelOperationType;
import prerna.sablecc2.om.ReactorKeysEnum;
import prerna.sablecc2.om.nounmeta.NounMetadata;
import prerna.util.Constants;
import prerna.util.Utility;

public class LLMInstructReactor extends AbstractReactor {

	private static final Logger classLogger = LogManager.getLogger(LLMInstructReactor.class);

	public LLMInstructReactor() {
		this.keysToGet = new String[] {ReactorKeysEnum.ENGINE.getKey(), 
				ReactorKeysEnum.COMMAND.getKey(), ReactorKeysEnum.CONTEXT.getKey(), 
				ReactorKeysEnum.PARAM_VALUES_MAP.getKey()};
		this.keyRequired = new int[] {1, 1, 0, 0};
	}

	@Override
	public NounMetadata execute() {
		organizeKeys();

		List> projectData = this.getProjectData();

		String engineId = this.keyValue.get(this.keysToGet[0]);
		if(!SecurityEngineUtils.userCanViewEngine(this.insight.getUser(), engineId)) {
			throw new IllegalArgumentException("Model " + engineId + " does not exist or user does not have access to this model");
		}

		String task = Utility.decodeURIComponent(this.keyValue.get(this.keysToGet[1]));
		String context = this.keyValue.get(this.keysToGet[2]);
		if (context != null) {
			context = Utility.decodeURIComponent(context);
		}

		Map paramMap = getMap();
		IModelEngine modelEngine = Utility.getModel(engineId);
		if(paramMap == null) {
			paramMap = new HashMap();
		}

		Map output = modelEngine.instruct(task, context, projectData, this.insight, paramMap).toMap();
		return new NounMetadata(output, PixelDataType.MAP, PixelOperationType.OPERATION);
	}

	private List> getProjectData() {
		try {
			List> projectInfo = SecurityProjectUtils.getUserProjectList(this.insight.getUser(), null, 
					false, true, null, null, null, null, null);
			// Create the keySet of project ids
			if(!projectInfo.isEmpty()) {
				Map index = new HashMap<>(projectInfo.size());
				int size = projectInfo.size();
				for(int i = 0; i < size; i++) {
					Map project = projectInfo.get(i);
					String projectId = project.get("project_id").toString();
					index.put(projectId, Integer.valueOf(i));
				}

				IRawSelectWrapper wrapper = null;
				List metaKeys = Arrays.asList("description");

				try {
					wrapper = SecurityProjectUtils.getProjectMetadataWrapper(index.keySet(), metaKeys, true);
					while(wrapper.hasNext()) {
						Object[] data = wrapper.next().getValues();
						String projectId = (String) data[0];
						String metaKey = (String) data[1];
						String metaValue = (String) data[2];
						if (metaValue == null) {
							continue;
						}
						int indexToFind = index.get(projectId);
						Map res = projectInfo.get(indexToFind);
						// whatever it is, if it is single send a single value, if it is multi send as array
						if(res.containsKey(metaKey)) {
							Object obj = res.get(metaKey);
							if(obj instanceof List) {
								((List) obj).add(metaValue);
							} else {
								List newList = new ArrayList<>();
								newList.add(obj);
								newList.add(metaValue);
								res.put(metaKey, newList);
							}
						} else {
							res.put(metaKey, metaValue);
						}
					}
				} catch (Exception e) {
					classLogger.error(Constants.STACKTRACE, e);
				} finally {
					if(wrapper != null) {
						try {
							wrapper.close();
						} catch (IOException e) {
							classLogger.error(Constants.STACKTRACE, e);

						}
					}
				}
			}

			return projectInfo;
		} catch (Exception e) {
			classLogger.error("There was a problem fetching the user project descriptions", e);
			return Collections.emptyList();
		}
	}

	/**
	 * 
	 * @return
	 */
	private Map getMap() {
		GenRowStruct mapGrs = this.store.getNoun(keysToGet[3]);
		if(mapGrs != null && !mapGrs.isEmpty()) {
			List mapInputs = mapGrs.getNounsOfType(PixelDataType.MAP);
			if(mapInputs != null && !mapInputs.isEmpty()) {
				return (Map) mapInputs.get(0).getValue();
			}
		}
		List mapInputs = this.curRow.getNounsOfType(PixelDataType.MAP);
		if(mapInputs != null && !mapInputs.isEmpty()) {
			return (Map) mapInputs.get(0).getValue();
		}
		return null;
	}

}