Skip to content

Commit

Permalink
task: separate enrichement alignement and extractions (#2120)
Browse files Browse the repository at this point in the history
  • Loading branch information
YohannParis authored Oct 31, 2023
1 parent c37abf8 commit 6ce0935
Show file tree
Hide file tree
Showing 2 changed files with 128 additions and 53 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -8,28 +8,19 @@
<tera-asset-link
:label="document.name!"
:asset-route="{ assetId: document.id!, pageType: AssetType.Documents }"
></tera-asset-link>
/>
</li>
</ul>
<Button text label="Enrich description" :loading="isLoading" @click="dialogForEnrichment" />
<Button text label="Extract variables" :loading="isLoading" @click="dialogForExtraction" />
<Button
label="Enrich Description"
text
:loading="enriching"
@click="
dialogType = 'enrich';
visible = true;
"
/>
<Button
v-if="props.assetType === ResourceType.MODEL"
label="Align Model"
text
:loading="aligning"
@click="
dialogType = 'align';
visible = true;
"
:disabled="props.assetType != ResourceType.MODEL"
:label="`Align extractions to ${assetType}`"
:loading="isLoading"
@click="dialogForAlignment"
/>

<Dialog
v-model:visible="visible"
modal
Expand Down Expand Up @@ -61,37 +52,46 @@
</div>
</div>
</div>
<template #footer>
<Button severity="secondary" outlined label="Cancel" @click="visible = false" />
<Button
:label="
dialogType === 'enrich'
? 'Use these resources to enrich descriptions'
: 'Use these resources to align the model'
"
@click="
dialogType === 'enrich' ? sendForEnrichments() : sendToAlignModel();
visible = false;
"
<aside v-if="dialogType === DialogType.EXTRACT && assetType === ResourceType.MODEL">
<p>Which extraction service would you like to use?</p>
<RadioButton
v-model="extractionService"
inputId="extractionServiceSkema"
name="skema"
:value="Extractor.SKEMA"
/>
<label for="extractionServiceSkema">SKEMA</label>
<RadioButton
v-model="extractionService"
inputId="extractionServiceMit"
name="mit"
:value="Extractor.MIT"
/>
<label for="extractionServiceMit">MIT</label>
</aside>
<template #footer>
<Button severity="secondary" outlined label="Cancel" @click="closeDialog" />
<Button :label="dialogActionCopy" :disabled="isDialogDisabled" @click="acceptDialog" />
</template>
</Dialog>
</main>
</template>

<script setup lang="ts">
import { onMounted, ref, watch } from 'vue';
import { computed, onMounted, ref, watch } from 'vue';
import Button from 'primevue/button';
import Dialog from 'primevue/dialog';
import DataTable from 'primevue/datatable';
import Column from 'primevue/column';
import RadioButton from 'primevue/radiobutton';
import { ResourceType } from '@/types/common';
import {
alignModel,
fetchExtraction,
pdfExtractions,
profileDataset,
profileModel,
fetchExtraction,
alignModel,
pdfExtractions
Extractor
} from '@/services/knowledge';
import { PollerResult } from '@/api/api';
import { isEmpty } from 'lodash';
Expand All @@ -106,20 +106,70 @@ const props = defineProps<{
assetId: string;
}>();
enum DialogType {
ENRICH,
EXTRACT,
ALIGN
}
const emit = defineEmits(['enriched']);
const visible = ref(false);
const selectedResources = ref();
const dialogType = ref<'enrich' | 'align'>('enrich');
const aligning = ref(false);
const enriching = ref(false);
const dialogType = ref<DialogType>(DialogType.ENRICH);
const extractionService = ref<Extractor>(Extractor.SKEMA);
const isLoading = ref(false);
const relatedDocuments = ref<Array<{ name: string | undefined; id: string | undefined }>>([]);
const sendForEnrichments = async (/* _selectedResources */) => {
const dialogActionCopy = ref('');
function openDialog() {
visible.value = true;
}
function closeDialog() {
visible.value = false;
}
function dialogForEnrichment() {
dialogType.value = DialogType.ENRICH;
dialogActionCopy.value = 'Use this resource to enrich descriptions';
openDialog();
}
function dialogForExtraction() {
dialogType.value = DialogType.EXTRACT;
dialogActionCopy.value = 'Use this resource to extract variables';
openDialog();
}
function dialogForAlignment() {
dialogType.value = DialogType.ALIGN;
dialogActionCopy.value = `Use this resource to align the ${props.assetType}`;
openDialog();
}
const acceptDialog = () => {
if (dialogType.value === DialogType.ENRICH) {
sendForEnrichment();
} else if (dialogType.value === DialogType.EXTRACT) {
sendForExtractions();
} else if (dialogType.value === DialogType.ALIGN) {
sendToAlignModel();
}
closeDialog();
};
// Disable the dialog action button if no resources are selected
// and the dialog type is not enrichment
const isDialogDisabled = computed(() => {
if (dialogType.value === DialogType.ENRICH) return false;
return !selectedResources.value;
});
const sendForEnrichment = async () => {
const jobIds: (string | null)[] = [];
const selectedResourceId = selectedResources.value?.id ?? null;
const extractionList: Promise<PollerResult<any>>[] = [];
enriching.value = true;
isLoading.value = true;
// Build enrichment job ids list (profile asset, align model, etc...)
if (props.assetType === ResourceType.MODEL) {
const profileModelJobId = await profileModel(props.assetId, selectedResourceId);
Expand All @@ -141,27 +191,36 @@ const sendForEnrichments = async (/* _selectedResources */) => {
// Poll all extractions
await Promise.all(extractionList);
enriching.value = false;
isLoading.value = false;
emit('enriched');
getRelatedDocuments();
await getRelatedDocuments();
};
const sendForExtractions = async () => {
const selectedResourceId = selectedResources.value?.id ?? null;
isLoading.value = true;
const pdfExtractionsJobId = await pdfExtractions(selectedResourceId, extractionService.value);
if (!pdfExtractionsJobId) return;
await fetchExtraction(pdfExtractionsJobId);
isLoading.value = false;
emit('enriched');
await getRelatedDocuments();
};
const sendToAlignModel = async () => {
const selectedResourceId = selectedResources.value?.id ?? null;
if (props.assetType === ResourceType.MODEL && selectedResourceId) {
// fetch pdf extractions and link amr synchronously
aligning.value = true;
const pdfExtractionsJobId = await pdfExtractions(selectedResourceId);
if (!pdfExtractionsJobId) return;
await fetchExtraction(pdfExtractionsJobId);
isLoading.value = true;
const linkAmrJobId = await alignModel(props.assetId, selectedResourceId);
if (!linkAmrJobId) return;
await fetchExtraction(linkAmrJobId);
aligning.value = false;
isLoading.value = false;
emit('enriched');
getRelatedDocuments();
await getRelatedDocuments();
}
};
Expand All @@ -181,7 +240,6 @@ async function getRelatedDocuments() {
const provenanceType = mapResourceTypeToProvenanceType(props.assetType);
if (!provenanceType) return;
const provenanceNodes = await getRelatedArtifacts(props.assetId, provenanceType, [
ProvenanceType.Document
]);
Expand Down Expand Up @@ -214,10 +272,17 @@ ul {
.no-documents-text {
padding: 5px;
font-size: var(--font-body);
font-size: var(--font-body-medium);
font-family: var(--font-family);
font-weight: 500;
color: var(--text-color-secondary);
text-align: left;
}
.p-dialog aside > * {
margin-top: 1rem;
}
.p-dialog aside label {
margin: 0 1rem 0 0.5rem;
}
</style>
18 changes: 14 additions & 4 deletions packages/client/hmi-client/src/services/knowledge.ts
Original file line number Diff line number Diff line change
Expand Up @@ -131,15 +131,24 @@ export const extractTextFromPDFDocument = async (documentId: string): Promise<st
return null;
};

export enum Extractor {
SKEMA = 'SKEMA',
MIT = 'MIT'
}
export const pdfExtractions = async (
documentId: string,
extractor: Extractor,
pdfName?: string,
description?: string
): Promise<string | null> => {
// I've purposefully excluded the MIT and SKEMA options here, so they're always
// defaulted to true.

let url = `/knowledge/pdf-extractions?document_id=${documentId}`;

if (extractor === Extractor.SKEMA) {
url += '&annotate_skema=true&annotate_mit=false';
} else if (extractor === Extractor.MIT) {
url += '&annotate_skema=false&annotate_mit=true';
}

if (pdfName) {
url += `&name=${pdfName}`;
}
Expand Down Expand Up @@ -175,7 +184,8 @@ export const extractPDF = async (documentId: string) => {
if (resp) {
const pollResult = await fetchExtraction(resp);
if (pollResult?.state === PollerState.Done) {
await pdfExtractions(documentId); // we don't care now. fire and forget.
pdfExtractions(documentId, Extractor.SKEMA);
pdfExtractions(documentId, Extractor.MIT);
}
}
}
Expand Down

0 comments on commit 6ce0935

Please sign in to comment.