Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Project search assets #5449

Merged
merged 26 commits into from
Nov 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
81ac626
WIP
kbirk Oct 31, 2024
e8f01b3
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 4, 2024
823f00c
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 5, 2024
850f217
Add method
kbirk Nov 6, 2024
096dd72
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 6, 2024
eb90e72
Add project assets into the project document.
kbirk Nov 6, 2024
b492104
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 7, 2024
5efee6f
More fixes
kbirk Nov 11, 2024
720cf6c
Missing commit
kbirk Nov 11, 2024
5efd27e
Get inner hits working
kbirk Nov 11, 2024
588718b
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 11, 2024
fc1d131
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 12, 2024
7768dc3
chore: update generated types
kbirk Nov 12, 2024
e167790
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 12, 2024
48378cf
chore: update generated types
kbirk Nov 12, 2024
173bcb9
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 12, 2024
b452667
chore: update generated types
kbirk Nov 12, 2024
7273b13
Merge branch 'main' into kbirk/project-search-assets
dgauldie Nov 13, 2024
4f5cabf
chore: update generated types
dgauldie Nov 13, 2024
b3d8d9b
Merge branch 'main' into kbirk/project-search-assets
dgauldie Nov 13, 2024
8302d98
Fix bad import.
kbirk Nov 14, 2024
30bcfa6
Merge branch 'kbirk/project-search-assets' of github.com:DARPA-ASKEM/…
kbirk Nov 14, 2024
c6f7220
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 14, 2024
4f7bdd5
Merge branch 'main' into kbirk/project-search-assets
kbirk Nov 14, 2024
2730f3e
Add docs
kbirk Nov 14, 2024
93a611a
Merge branch 'kbirk/project-search-assets' of github.com:DARPA-ASKEM/…
kbirk Nov 14, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/client/hmi-client/src/types/Types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1024,6 +1024,7 @@ export enum AssetType {
ModelConfiguration = "model-configuration",
Artifact = "artifact",
InterventionPolicy = "intervention-policy",
Project = "project",
}

export enum EvaluationScenarioStatus {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@
import software.uncharted.terarium.hmiserver.service.data.ProjectAssetService;
import software.uncharted.terarium.hmiserver.service.data.ProjectPermissionsService;
import software.uncharted.terarium.hmiserver.service.data.ProjectSearchService;
import software.uncharted.terarium.hmiserver.service.data.ProjectSearchService.ProjectSearchResponse;
import software.uncharted.terarium.hmiserver.service.data.ProjectService;
import software.uncharted.terarium.hmiserver.service.data.TerariumAssetServices;
import software.uncharted.terarium.hmiserver.service.data.WorkflowService;
Expand Down Expand Up @@ -929,7 +930,7 @@ public ResponseEntity<PermissionRelationships> getProjectPermissions(@PathVariab
try {
for (final RebacPermissionRelationship permissionRelationship : rebacProject.getPermissionRelationships()) {
if (permissionRelationship.getSubjectType().equals(Schema.Type.USER)) {
PermissionUser user = reBACService.getUser(permissionRelationship.getSubjectId());
final PermissionUser user = reBACService.getUser(permissionRelationship.getSubjectId());
if (user != null) {
permissions.addUser(user, permissionRelationship.getRelationship());
}
Expand All @@ -948,6 +949,55 @@ public ResponseEntity<PermissionRelationships> getProjectPermissions(@PathVariab
return ResponseEntity.ok(permissions);
}

@GetMapping("/knn")
@Secured(Roles.USER)
@Operation(summary = "Executes a knn search against the provided asset type")
@ApiResponses(
value = {
@ApiResponse(
responseCode = "200",
description = "Query results",
content = @Content(
mediaType = "application/json",
schema = @io.swagger.v3.oas.annotations.media.Schema(implementation = JsonNode.class)
)
),
@ApiResponse(responseCode = "204", description = "There was no concept found", content = @Content),
@ApiResponse(
responseCode = "500",
description = "There was an issue retrieving the concept from the data store",
content = @Content
)
}
)
public ResponseEntity<List<ProjectSearchResponse>> projectKnnSearch(
@RequestParam(value = "page-size", defaultValue = "100", required = false) final Integer pageSize,
@RequestParam(value = "page", defaultValue = "0", required = false) final Integer page,
@RequestParam(value = "text", defaultValue = "") final String text,
@RequestParam(value = "k", defaultValue = "100") final int k,
@RequestParam(value = "num-candidates", defaultValue = "1000") final int numCandidates
) {
try {
final String userId = currentUserService.get().getId();

final List<ProjectSearchResponse> res = projectSearchService.searchProjectsKNN(
userId,
pageSize,
page,
text,
k,
numCandidates,
null
);

return ResponseEntity.ok(res);
} catch (final Exception e) {
final String error = "Unable to get execute knn search";
log.error(error, e);
throw new ResponseStatusException(org.springframework.http.HttpStatus.INTERNAL_SERVER_ERROR, error);
}
}

// --------------------------------------------------------------------------
// Project Permissions
// --------------------------------------------------------------------------
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package software.uncharted.terarium.hmiserver.models;

import com.fasterxml.jackson.annotation.JsonAlias;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.MapperFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
Expand All @@ -19,6 +20,7 @@
import lombok.EqualsAndHashCode;
import lombok.experimental.Accessors;
import org.hibernate.annotations.Type;
import software.uncharted.terarium.hmiserver.annotations.TSIgnore;
import software.uncharted.terarium.hmiserver.annotations.TSModel;
import software.uncharted.terarium.hmiserver.annotations.TSOptional;

Expand Down Expand Up @@ -80,14 +82,14 @@ protected TerariumAsset cloneSuperFields(final TerariumAsset asset) {
return asset;
}

public static void removeFieldsWithKeys(ObjectNode objectNode, List<String> keys) {
Iterator<String> keysIterator = objectNode.fieldNames();
public static void removeFieldsWithKeys(final ObjectNode objectNode, final List<String> keys) {
final Iterator<String> keysIterator = objectNode.fieldNames();
while (keysIterator.hasNext()) {
String key = keysIterator.next();
final String key = keysIterator.next();
if (keys.contains(key)) {
keysIterator.remove();
} else {
JsonNode node = objectNode.get(key);
final JsonNode node = objectNode.get(key);
if (node.isObject()) {
removeFieldsWithKeys((ObjectNode) node, keys);
}
Expand All @@ -96,22 +98,23 @@ public static void removeFieldsWithKeys(ObjectNode objectNode, List<String> keys
}

/**
* Serialize the asset to a JSON string, removing the fields that are not needed.
* Serialize the asset to a JSON string, removing the fields that are not
* needed.
*
* @param keepFields A list of fields that should not be removed.
* @param additionalDeleteFields Additional fields to remove.
* @param keepFields A list of fields that should not be removed.
* @param additionalDeleteFields Additional fields to remove.
* @return The JSON string.
*/
public String serializeWithoutTerariumFields(
@Nullable String[] keepFields,
@Nullable String[] additionalDeleteFields
@Nullable final String[] keepFields,
@Nullable final String[] additionalDeleteFields
) {
final ObjectMapper mapper = new ObjectMapper();
mapper.setConfig(mapper.getSerializationConfig().with(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY));
final ObjectNode objectNode = mapper.convertValue(this, ObjectNode.class);

// Fields to delete
String[] deleteFields = new String[] {
final String[] deleteFields = new String[] {
"id",
"createdOn",
"updatedOn",
Expand All @@ -123,15 +126,15 @@ public String serializeWithoutTerariumFields(
"fileNames",
"userId"
};
List<String> deleteFieldsList = new ArrayList<>(Arrays.asList(deleteFields));
final List<String> deleteFieldsList = new ArrayList<>(Arrays.asList(deleteFields));
if (keepFields != null) {
for (String field : keepFields) {
for (final String field : keepFields) {
deleteFieldsList.removeIf(key -> key.equals(field));
}
}

// Remove the fields that are not needed
for (String key : deleteFieldsList) {
for (final String key : deleteFieldsList) {
objectNode.remove(key);
}

Expand All @@ -149,4 +152,10 @@ public String serializeWithoutTerariumFields(

return objectNode.toString();
}

@JsonIgnore
@TSIgnore
public String getEmbeddingSourceText() {
return name + " " + description;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ public static class Embeddings {

private String embeddingId;
private double[] vector;
private long[] spans;
private long[] span;
}

private List<Embeddings> embeddings = new ArrayList<>();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.util.List;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
Expand All @@ -13,6 +14,7 @@
import software.uncharted.terarium.hmiserver.models.dataservice.document.DocumentAsset;
import software.uncharted.terarium.hmiserver.models.dataservice.model.Model;
import software.uncharted.terarium.hmiserver.models.dataservice.model.configurations.ModelConfiguration;
import software.uncharted.terarium.hmiserver.models.dataservice.project.Project;
import software.uncharted.terarium.hmiserver.models.dataservice.simulation.Simulation;
import software.uncharted.terarium.hmiserver.models.dataservice.workflow.Workflow;
import software.uncharted.terarium.hmiserver.models.simulationservice.interventions.InterventionPolicy;
Expand Down Expand Up @@ -46,7 +48,10 @@ public enum AssetType {
ARTIFACT,

@JsonProperty("intervention-policy")
INTERVENTION_POLICY;
INTERVENTION_POLICY,

@JsonProperty("project")
PROJECT;

public static AssetType getAssetType(final String assetTypeName, final ObjectMapper objectMapper)
throws ResponseStatusException {
Expand All @@ -58,6 +63,32 @@ public static AssetType getAssetType(final String assetTypeName, final ObjectMap
}
}

public static AssetType getAssetType(final Class<? extends TerariumAsset> clazz) {
if (clazz == Artifact.class) {
return ARTIFACT;
} else if (clazz == Code.class) {
return CODE;
} else if (clazz == Dataset.class) {
return DATASET;
} else if (clazz == DocumentAsset.class) {
return DOCUMENT;
} else if (clazz == Model.class) {
return MODEL;
} else if (clazz == ModelConfiguration.class) {
return MODEL_CONFIGURATION;
} else if (clazz == Simulation.class) {
return SIMULATION;
} else if (clazz == Workflow.class) {
return WORKFLOW;
} else if (clazz == InterventionPolicy.class) {
return INTERVENTION_POLICY;
} else if (clazz == Project.class) {
return PROJECT;
} else {
throw new IllegalArgumentException("Unrecognized asset type: " + clazz);
}
}

public Class<? extends TerariumAsset> getAssetClass() {
switch (this) {
case ARTIFACT:
Expand All @@ -78,8 +109,25 @@ public Class<? extends TerariumAsset> getAssetClass() {
return Workflow.class;
case INTERVENTION_POLICY:
return InterventionPolicy.class;
case PROJECT:
return Project.class;
default:
throw new IllegalArgumentException("Unrecognized asset type: " + this);
}
}

public static List<AssetType> getAllAssetTypes() {
return List.of(
ARTIFACT,
CODE,
DATASET,
DOCUMENT,
MODEL,
MODEL_CONFIGURATION,
SIMULATION,
WORKFLOW,
PROJECT,
INTERVENTION_POLICY
);
}
}
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
package software.uncharted.terarium.hmiserver.models.dataservice.document;

import com.fasterxml.jackson.annotation.JsonAlias;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.hypersistence.utils.hibernate.type.json.JsonType;
import jakarta.persistence.CascadeType;
import jakarta.persistence.Column;
Expand All @@ -21,6 +23,7 @@
import lombok.experimental.Accessors;
import org.hibernate.annotations.JdbcTypeCode;
import org.hibernate.annotations.Type;
import software.uncharted.terarium.hmiserver.annotations.TSIgnore;
import software.uncharted.terarium.hmiserver.annotations.TSModel;
import software.uncharted.terarium.hmiserver.annotations.TSOptional;
import software.uncharted.terarium.hmiserver.models.TerariumAsset;
Expand Down Expand Up @@ -135,4 +138,20 @@ public DocumentAsset clone() {

return clone;
}

@JsonIgnore
@TSIgnore
public String getEmbeddingSourceText() {
try {
if (getMetadata() != null && getMetadata().containsKey("gollmCard")) {
// update embeddings
final JsonNode card = getMetadata().get("gollmCard");
final ObjectMapper objectMapper = new ObjectMapper();
return objectMapper.writeValueAsString(card);
}
return null;
} catch (final Exception e) {
throw new RuntimeException("Failed to serialize model embedding text into JSON", e);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -354,4 +354,18 @@ public boolean isPetrinet() {
}
return this.getHeader().getSchemaName().equalsIgnoreCase("petrinet");
}

@JsonIgnore
@TSIgnore
public String getEmbeddingSourceText() {
try {
final ObjectMapper objectMapper = new ObjectMapper();
if (getMetadata() != null && getMetadata().getGollmCard() != null) {
return objectMapper.writeValueAsString(getMetadata().getGollmCard());
}
return objectMapper.writeValueAsString(this);
} catch (final Exception e) {
throw new RuntimeException("Failed to serialize model embedding text into JSON", e);
}
}
}
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package software.uncharted.terarium.hmiserver.models.dataservice.project;

import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonManagedReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.swagger.v3.oas.annotations.media.Schema;
import jakarta.persistence.CascadeType;
import jakarta.persistence.Entity;
Expand All @@ -17,8 +19,10 @@
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import org.hibernate.annotations.JdbcTypeCode;
import org.hibernate.annotations.Where;
import software.uncharted.terarium.hmiserver.annotations.TSIgnore;
import software.uncharted.terarium.hmiserver.annotations.TSModel;
import software.uncharted.terarium.hmiserver.annotations.TSOptional;
import software.uncharted.terarium.hmiserver.models.TerariumAsset;
Expand All @@ -28,6 +32,7 @@
@Accessors(chain = true)
@TSModel
@Entity
@Slf4j
public class Project extends TerariumAsset {

@Serial
Expand Down Expand Up @@ -71,7 +76,10 @@ public class Project extends TerariumAsset {
@Schema(accessMode = Schema.AccessMode.READ_ONLY)
private Boolean publicProject;

/** Information for the front-end to enable/disable features based on user permissions (Read/Write). */
/**
* Information for the front-end to enable/disable features based on user
* permissions (Read/Write).
*/
@TSOptional
@Transient
@Schema(accessMode = Schema.AccessMode.READ_ONLY)
Expand Down Expand Up @@ -111,4 +119,19 @@ public Project clone() {
cloned.userPermission = userPermission;
return cloned;
}

@JsonIgnore
@TSIgnore
public String getEmbeddingSourceText() {
try {
if (overviewContent != null) {
log.info(new String(overviewContent));
return new String(overviewContent);
}
final ObjectMapper objectMapper = new ObjectMapper();
return objectMapper.writeValueAsString(this);
} catch (final Exception e) {
throw new RuntimeException("Failed to serialize model embedding text into JSON", e);
}
}
}
Loading
Loading