@finos/legend-extension-dsl-data-space
Version:
Legend extension for Data Space DSL
415 lines • 18.3 kB
JavaScript
import { jsx as _jsx } from "react/jsx-runtime";
/**
* Copyright (c) 2026-present, Goldman Sachs
*
* 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.
*/
import { useState, useEffect, useMemo } from 'react';
import { observer } from 'mobx-react-lite';
import { Enumeration, getAllSuperclasses, PRIMITIVE_TYPE, } from '@finos/legend-graph';
import { LegendAIChat, LegendAIErrorBoundary, findLegendAIPlugin, extractParameterSchemas, buildPropertyDocIndex, enrichColumnsFromElementDocs, inferServiceRelationshipsFromAssociations, extractLambdaPreFilters, extractModelContext, } from '@finos/legend-lego/legend-ai';
import {} from '@finos/legend-extension-dsl-diagram';
import { assertErrorThrown, LogEvent } from '@finos/legend-shared';
import { DSL_DATASPACE_EVENT } from '../__lib__/DSL_DataSpace_Event.js';
import { DataSpaceServiceExecutableInfo, DataSpaceMultiExecutionServiceExecutableInfo, DataSpaceExecutableTDSResult, } from '../graph-manager/action/analytics/DataSpaceAnalysis.js';
import { DataSpaceSupportEmail, DataSpaceSupportCombinedInfo, } from '../graph/metamodel/pure/model/packageableElements/dataSpace/DSL_DataSpace_DataSpace.js';
const MAX_QUERY_SCAN_LENGTH = 5_000;
const MAX_QUERY_TEMPLATE_LENGTH = 1000;
const ENTITY_PATH_PATTERN = /(?<entityPath>[A-Za-z_]\w{0,63}(?:::[A-Za-z_]\w{0,63}){0,10})\.all\(\)/g;
// ────────────────────────────────────────────────────────────────────────────
// Diagram-based model context extraction (fallback when elementDocs is sparse)
// ────────────────────────────────────────────────────────────────────────────
/**
* Extracts a {@link LegendAIModelContext} from diagram metamodel objects.
* Diagrams always contain the actual Class metamodel references so this
* works even when elementDocs documentation is empty or incomplete.
*/
export function extractModelContextFromDiagrams(diagrams) {
const entityMap = new Map();
const associations = [];
const enumPaths = new Map();
for (const { diagram } of diagrams) {
for (const classView of diagram.classViews) {
const cls = classView.class.value;
if (entityMap.has(cls.path)) {
continue;
}
const properties = cls.properties.map((prop) => {
const rawType = prop.genericType.value.rawType;
const typePath = rawType.path;
if (rawType instanceof Enumeration && !enumPaths.has(typePath)) {
enumPaths.set(typePath, {
path: typePath,
name: rawType.name,
values: rawType.values.map((v) => v.name),
});
}
return {
name: prop.name,
type: typePath,
isCollection: prop.multiplicity.upperBound === undefined ||
prop.multiplicity.upperBound > 1,
isOptional: prop.multiplicity.lowerBound === 0,
};
});
const entity = {
path: cls.path,
name: cls.name,
properties,
};
const superTypes = getAllSuperclasses(cls).map((c) => c.path);
if (superTypes.length > 0) {
entity.superTypes = superTypes;
}
entityMap.set(cls.path, entity);
}
for (const assocView of diagram.associationViews) {
const assoc = assocView.association.value;
const [propA, propB] = assoc.properties;
associations.push({
name: assoc.name,
leftEntity: assocView.from.classView.value.class.value.path,
leftProperty: propA.name,
rightEntity: assocView.to.classView.value.class.value.path,
rightProperty: propB.name,
});
}
}
const result = {
entities: Array.from(entityMap.values()),
associations,
};
if (enumPaths.size > 0) {
result.enumerations = Array.from(enumPaths.values());
}
return result;
}
/**
* Enriches a {@link LegendAIModelContext} with mapping coverage data.
* Marks entities that are root-mapped (directly queryable by the orchestrator)
* so entity resolution can prefer them.
*/
export function enrichModelContextWithMappingCoverage(modelContext, mappingCoverage) {
if (!mappingCoverage) {
return;
}
for (const entity of modelContext.entities) {
const mapped = mappingCoverage.mappedEntities.find((me) => me.path === entity.path);
if (mapped?.info?.isRootEntity) {
entity.isRootMapped = true;
}
}
}
/**
* Builds column-to-property mappings by matching TDS result column names
* to model property names on the root entity.
* Uses case-insensitive, whitespace/underscore-normalized comparison.
*/
function buildColumnPropertyMappings(columns, rootEntity) {
if (!rootEntity || rootEntity.properties.length === 0) {
return undefined;
}
const propLookup = new Map();
for (const prop of rootEntity.properties) {
const normalized = prop.name.toLowerCase().replaceAll('_', '');
propLookup.set(normalized, prop.name);
}
const mappings = [];
for (const colName of columns) {
const normalized = colName.toLowerCase().replaceAll(/[\s_]/g, '');
const propName = propLookup.get(normalized);
if (propName && propName !== colName) {
mappings.push({ columnName: colName, propertyPath: propName });
}
}
return mappings.length > 0 ? mappings : undefined;
}
function collectReferencedEntityPaths(queryStr) {
const safeQuery = queryStr.slice(0, MAX_QUERY_SCAN_LENGTH);
// Reset lastIndex defensively — the regex is module-scoped and reused.
ENTITY_PATH_PATTERN.lastIndex = 0;
const paths = [];
const seen = new Set();
let match;
while ((match = ENTITY_PATH_PATTERN.exec(safeQuery)) !== null) {
const path = match.groups?.entityPath;
if (path && !seen.has(path)) {
seen.add(path);
paths.push(path);
}
}
return paths;
}
function buildExecutableInfo(exec, entityMap, serviceParameters) {
if (!(exec.info instanceof DataSpaceServiceExecutableInfo ||
exec.info instanceof DataSpaceMultiExecutionServiceExecutableInfo)) {
return undefined;
}
const queryStr = exec.info.query;
const referencedEntityPaths = collectReferencedEntityPaths(queryStr);
const rootEntityPath = referencedEntityPaths[0];
if (!rootEntityPath) {
return undefined;
}
const info = {
title: exec.title,
rootEntityPath,
queryTemplate: queryStr.length > MAX_QUERY_TEMPLATE_LENGTH
? `${queryStr.slice(0, MAX_QUERY_TEMPLATE_LENGTH)}...`
: queryStr,
};
if (referencedEntityPaths.length > 1) {
info.referencedEntityPaths = referencedEntityPaths;
}
if (exec.description !== undefined) {
info.description = exec.description;
}
const reqParams = serviceParameters.get(exec.title);
if (reqParams && reqParams.length > 0) {
info.requiredParameters = reqParams;
}
if (exec.result instanceof DataSpaceExecutableTDSResult) {
info.columns = exec.result.columns.map((c) => c.name);
const colMappings = buildColumnPropertyMappings(info.columns, entityMap.get(rootEntityPath));
if (colMappings) {
info.columnPropertyMappings = colMappings;
}
}
return info;
}
export function extractExecutableInfo(executables, modelContext, services) {
const entityMap = new Map(modelContext.entities.map((e) => [e.path, e]));
const serviceParameters = new Map();
for (const svc of services) {
if (svc.parameterSchemas && svc.parameterSchemas.length > 0) {
serviceParameters.set(svc.title, svc.parameterSchemas.map((p) => ({
name: p.name,
type: p.type ?? PRIMITIVE_TYPE.STRING,
})));
}
}
const result = [];
for (const exec of executables) {
const info = buildExecutableInfo(exec, entityMap, serviceParameters);
if (!info) {
continue;
}
result.push(info);
const queryableEntityPaths = info.referencedEntityPaths ?? [
info.rootEntityPath,
];
for (const entityPath of queryableEntityPaths) {
const entity = entityMap.get(entityPath);
if (entity) {
entity.isQueryable = true;
}
}
}
return result;
}
export async function extractTDSServicesFromDataSpace(viewerState) {
const executables = viewerState.dataSpaceAnalysisResult.executables;
const graphManager = viewerState.graphManagerState.graphManager;
const nested = await Promise.all(executables.map(async (exec) => {
if (!(exec.info instanceof DataSpaceServiceExecutableInfo ||
exec.info instanceof DataSpaceMultiExecutionServiceExecutableInfo) ||
!(exec.result instanceof DataSpaceExecutableTDSResult)) {
return [];
}
const info = exec.info;
const tdsResult = exec.result;
const { parameters, parameterSchemas, parameterExtractionFailed } = await extractParameterSchemas(info.query, graphManager, viewerState.graphManagerState);
// Extract hardcoded pre-filter constraints from the lambda
let preFilters;
try {
const rawLambda = await graphManager.pureCodeToLambda(info.query);
const extracted = extractLambdaPreFilters(rawLambda.body);
if (extracted.length > 0) {
preFilters = extracted;
}
}
catch (error) {
assertErrorThrown(error);
// pre-filter extraction is best-effort — a service without pre-filters
// still works, the AI just won't know about hardcoded constraints.
}
const entry = {
title: exec.title,
pattern: info.pattern,
columns: tdsResult.columns.map((col) => {
const column = { name: col.name };
if (col.type !== undefined) {
column.type = col.type;
}
if (col.relationalType !== undefined) {
column.relationalType = col.relationalType;
}
if (col.documentation !== undefined) {
column.documentation = col.documentation;
}
if (col.sampleValues !== undefined) {
column.sampleValues = col.sampleValues;
}
return column;
}),
parameters,
...(parameterSchemas.length > 0 ? { parameterSchemas } : {}),
...(parameterExtractionFailed
? { parameterExtractionFailed: true }
: {}),
...(preFilters ? { preFilters } : {}),
};
if (exec.description !== undefined) {
entry.description = exec.description;
}
return [entry];
}));
const services = nested.flat();
// Enrich columns with documentation and nullability from elementDocs
const propIndex = buildPropertyDocIndex(viewerState.dataSpaceAnalysisResult.elementDocs);
if (propIndex.size > 0) {
for (const svc of services) {
enrichColumnsFromElementDocs(svc.columns, propIndex);
}
}
return services;
}
function extractMetadataFromDataSpace(viewerState, coordinates, services) {
const result = viewerState.dataSpaceAnalysisResult;
let supportInfoText;
if (result.supportInfo instanceof DataSpaceSupportEmail) {
supportInfoText = result.supportInfo.address;
}
else if (result.supportInfo instanceof DataSpaceSupportCombinedInfo) {
supportInfoText = (result.supportInfo.emails ?? []).join(', ');
}
const metadata = {
name: result.title ?? result.name,
coordinates,
serviceSummaries: result.executables.map((exec) => {
const summary = {
title: exec.title,
};
if (exec.description !== undefined) {
summary.description = exec.description;
}
if (exec.result instanceof DataSpaceExecutableTDSResult) {
summary.columnNames = exec.result.columns.map((col) => col.name);
}
return summary;
}),
tags: result.taggedValues.map((tv) => ({
profile: tv.profile,
value: tv.value,
})),
};
if (result.description !== undefined) {
metadata.description = result.description;
}
if (supportInfoText !== undefined) {
metadata.supportInfo = supportInfoText;
}
if (services.length >= 2) {
const serviceRels = inferServiceRelationshipsFromAssociations(services, result.elementDocs);
if (serviceRels.length > 0) {
metadata.serviceRelationships = serviceRels;
}
}
return metadata;
}
const DataSpaceLegendAIIntegrationInner = observer((props) => {
const { dataSpaceViewerState, config, onClose, onMinimize } = props;
const legendAIPlugin = useMemo(() => findLegendAIPlugin(dataSpaceViewerState.applicationStore.pluginManager.getApplicationPlugins()), [dataSpaceViewerState]);
const [services, setServices] = useState([]);
useEffect(() => {
let cancelled = false;
extractTDSServicesFromDataSpace(dataSpaceViewerState)
.then((result) => {
if (!cancelled) {
setServices(result);
}
})
.catch((error) => {
assertErrorThrown(error);
dataSpaceViewerState.applicationStore.logService.warn(LogEvent.create(DSL_DATASPACE_EVENT.ERROR_EXTRACT_LEGEND_AI_SERVICES), error);
});
return () => {
cancelled = true;
};
}, [dataSpaceViewerState]);
const coordinates = `${dataSpaceViewerState.groupId}:${dataSpaceViewerState.artifactId}:${dataSpaceViewerState.versionId}`;
const dataProductCoordinates = useMemo(() => ({
data_product: dataSpaceViewerState.dataSpaceAnalysisResult.path,
group_id: dataSpaceViewerState.groupId,
artifact_id: dataSpaceViewerState.artifactId,
version: dataSpaceViewerState.versionId,
}), [
dataSpaceViewerState.dataSpaceAnalysisResult.path,
dataSpaceViewerState.groupId,
dataSpaceViewerState.artifactId,
dataSpaceViewerState.versionId,
]);
const pureExecutionContext = useMemo(() => ({
mapping: dataSpaceViewerState.currentExecutionContext.mapping.path,
runtime: dataSpaceViewerState.currentRuntime.path,
}), [
dataSpaceViewerState.currentExecutionContext.mapping,
dataSpaceViewerState.currentRuntime,
]);
const metadata = useMemo(() => extractMetadataFromDataSpace(dataSpaceViewerState, coordinates, services), [dataSpaceViewerState, coordinates, services]);
const modelContext = useMemo(() => {
const result = dataSpaceViewerState.dataSpaceAnalysisResult;
let ctx;
if (result.elementDocs.length > 0) {
const docsCtx = extractModelContext(result.elementDocs);
if (docsCtx.entities.length > 0) {
ctx = docsCtx;
}
}
// Fallback: extract from diagram metamodel objects (always available)
if (!ctx && result.diagrams.length > 0) {
const diagramCtx = extractModelContextFromDiagrams(result.diagrams);
if (diagramCtx.entities.length > 0) {
ctx = diagramCtx;
}
}
if (!ctx) {
return undefined;
}
const mappingPath = dataSpaceViewerState.currentExecutionContext.mapping.path;
const mappingCoverage = result.mappingToMappingCoverageResult?.get(mappingPath);
enrichModelContextWithMappingCoverage(ctx, mappingCoverage);
if (result.executables.length > 0) {
const execInfos = extractExecutableInfo(result.executables, ctx, services);
if (execInfos.length > 0) {
ctx.executables = execInfos;
}
}
if (result.description) {
ctx.dataspaceDescription = result.description;
}
return ctx;
}, [
dataSpaceViewerState.dataSpaceAnalysisResult,
dataSpaceViewerState.currentExecutionContext.mapping,
services,
]);
if (!config.enabled || !legendAIPlugin) {
return null;
}
const dsTitle = dataSpaceViewerState.dataSpaceAnalysisResult.title ??
dataSpaceViewerState.dataSpaceAnalysisResult.name;
return (_jsx(LegendAIChat, { services: services, coordinates: coordinates, config: config, metadata: metadata, title: `Ask ${dsTitle}`, plugin: legendAIPlugin, contextBannerMessage: "You can query available TDS Executables within Data Space, or use Legend AI MCP for Pure queries on models.", dataProductCoordinates: dataProductCoordinates, pureExecutionContext: pureExecutionContext, ...(modelContext ? { modelContext } : {}), ...(onClose ? { onClose } : {}), ...(onMinimize ? { onMinimize } : {}) }));
});
export const DataSpaceLegendAIIntegration = observer((props) => (_jsx(LegendAIErrorBoundary, { children: _jsx(DataSpaceLegendAIIntegrationInner, { ...props }) })));
//# sourceMappingURL=DataSpaceLegendAIIntegration.js.map