@finos/legend-lego 2.0.211 → 2.0.213

package/src/legend-ai/stores/LegendAIChatState.ts

@@ -15,7 +15,7 @@
  */
 
 import { useState, useRef, useCallback, useEffect, useMemo } from 'react';
-import { noop } from '@finos/legend-shared';
+import { guaranteeNonNullable, noop } from '@finos/legend-shared';
 import {
   type TDSServiceSchema,
   type LegendAIConfig,
@@ -30,10 +30,13 @@ import {
   type LegendAIAssistantMessage,
   LegendAIMessageRole,
   LegendAIChatTelemetryEventType,
+  toRowCountBucket,
+  LEGEND_AI_DEFAULT_MODEL_LABEL,
 } from '../LegendAITypes.js';
 import {
   type LegendAI_LegendApplicationPlugin_Extension,
   type LegendAIOrchestratorDataProductCoordinates,
+  type LegendAIResolvedEntities,
 } from '../LegendAI_LegendApplicationPlugin_Extension.js';
 import type { QueryExplicitExecutionContextInfo } from '@finos/legend-graph';
 import {
@@ -45,6 +48,8 @@ import {
   processQuestionViaOrchestrator,
 } from './LegendAIChatProcessors.js';
 
+const LEGEND_AI_MCP_SCOPE_ID = 'legend-ai-mcp';
+
 export const useLegendAIChatState = (
   services: TDSServiceSchema[],
   coordinates: string,
@@ -56,7 +61,6 @@ export const useLegendAIChatState = (
   modelContext?: LegendAIModelContext,
   onLogTelemetryEvent?: (event: LegendAIChatTelemetryEvent) => void,
 ): LegendAIChatState => {
-  const LEGEND_AI_MCP_SCOPE_ID = 'legend-ai-mcp';
   const [questionText, setQuestionText] = useState('');
   const [isSending, setIsSending] = useState(false);
   const [messages, setMessages] = useState<LegendAIMessage[]>([]);
@@ -118,17 +122,47 @@ export const useLegendAIChatState = (
     });
   }, []);
 
-  const toggleScope = useCallback((scope: LegendAIScopeItem) => {
-    setSelectedScopes((prev) =>
-      prev.some((s) => s.id === scope.id)
-        ? prev.filter((s) => s.id !== scope.id)
-        : [...prev, scope],
-    );
-  }, []);
+  const logScopeChanged = useCallback(
+    (next: LegendAIScopeItem[]): void => {
+      onLogTelemetryEvent?.({
+        type: LegendAIChatTelemetryEventType.SCOPE_CHANGED,
+        selectedCount: next.length,
+        includesMcp: next.some((s) => s.id === LEGEND_AI_MCP_SCOPE_ID),
+      });
+    },
+    [onLogTelemetryEvent],
+  );
 
-  const removeScope = useCallback((scopeId: string) => {
-    setSelectedScopes((prev) => prev.filter((s) => s.id !== scopeId));
-  }, []);
+  const toggleScope = useCallback(
+    (scope: LegendAIScopeItem) => {
+      const next = selectedScopes.some((s) => s.id === scope.id)
+        ? selectedScopes.filter((s) => s.id !== scope.id)
+        : [...selectedScopes, scope];
+      setSelectedScopes(next);
+      logScopeChanged(next);
+    },
+    [selectedScopes, logScopeChanged],
+  );
+
+  const removeScope = useCallback(
+    (scopeId: string) => {
+      const next = selectedScopes.filter((s) => s.id !== scopeId);
+      setSelectedScopes(next);
+      logScopeChanged(next);
+    },
+    [selectedScopes, logScopeChanged],
+  );
+
+  const handleSelectModel = useCallback(
+    (model: string | undefined): void => {
+      setSelectedModelName(model);
+      onLogTelemetryEvent?.({
+        type: LegendAIChatTelemetryEventType.MODEL_CHANGED,
+        model: model ?? LEGEND_AI_DEFAULT_MODEL_LABEL,
+      });
+    },
+    [onLogTelemetryEvent],
+  );
 
   const configForRequest = useMemo(
     () => ({
@@ -214,6 +248,15 @@ export const useLegendAIChatState = (
                 type: LegendAIChatTelemetryEventType.RESPONSE_RECEIVED,
                 outcome: classifyResponseOutcome(lastAssistant),
                 durationMs: Date.now() - startTime,
+                hasSql: Boolean(lastAssistant?.sql),
+                hasResultGrid: Boolean(lastAssistant?.gridData),
+                rowCountBucket: toRowCountBucket(
+                  lastAssistant?.gridData?.rowData.length ?? 0,
+                ),
+                suggestionCount: lastAssistant?.suggestedQueries.length ?? 0,
+                ...(lastAssistant?.errorType
+                  ? { errorType: lastAssistant.errorType }
+                  : {}),
               });
             }
             sendTimeoutRef.current = undefined;
@@ -223,6 +266,14 @@ export const useLegendAIChatState = (
     [isSending, messages, onLogTelemetryEvent],
   );
 
+  const shouldRouteToOrchestrator = useCallback(
+    (): boolean =>
+      selectedScopes.some((scope) => scope.id === LEGEND_AI_MCP_SCOPE_ID) &&
+      Boolean(configForRequest.orchestratorUrl) &&
+      dataProductCoordinates !== undefined,
+    [selectedScopes, configForRequest.orchestratorUrl, dataProductCoordinates],
+  );
+
   const askQuestion = useCallback(
     (): void =>
       dispatchQuestion(questionText, (trimmed, history) => {
@@ -232,14 +283,10 @@ export const useLegendAIChatState = (
           history,
           setMessages,
         };
-        if (
-          selectedScopes.some((scope) => scope.id === LEGEND_AI_MCP_SCOPE_ID) &&
-          configForRequest.orchestratorUrl &&
-          dataProductCoordinates
-        ) {
+        if (shouldRouteToOrchestrator()) {
           return processQuestionViaOrchestrator(
             trimmed,
-            dataProductCoordinates,
+            guaranteeNonNullable(dataProductCoordinates),
             metadata,
             operationContext,
             pureExecutionContext,
@@ -268,7 +315,7 @@ export const useLegendAIChatState = (
       plugin,
       dataProductCoordinates,
       pureExecutionContext,
-      selectedScopes,
+      shouldRouteToOrchestrator,
       modelContext,
     ],
   );
@@ -282,14 +329,10 @@ export const useLegendAIChatState = (
           history,
           setMessages,
         };
-        if (
-          selectedScopes.some((scope) => scope.id === LEGEND_AI_MCP_SCOPE_ID) &&
-          configForRequest.orchestratorUrl &&
-          dataProductCoordinates
-        ) {
+        if (shouldRouteToOrchestrator()) {
           return processQuestionViaOrchestrator(
             trimmed,
-            dataProductCoordinates,
+            guaranteeNonNullable(dataProductCoordinates),
             metadata,
             operationContext,
             pureExecutionContext,
@@ -324,7 +367,7 @@ export const useLegendAIChatState = (
       plugin,
       dataProductCoordinates,
       pureExecutionContext,
-      selectedScopes,
+      shouldRouteToOrchestrator,
       modelContext,
     ],
   );
@@ -335,6 +378,7 @@ export const useLegendAIChatState = (
         return;
       }
       let question: string | undefined;
+      let preResolvedEntities: LegendAIResolvedEntities | undefined;
       for (let i = 0; i < messages.length; i++) {
         const msg = messages[i];
         if (
@@ -342,6 +386,7 @@ export const useLegendAIChatState = (
           msg.id === messageId &&
           i > 0
         ) {
+          preResolvedEntities = msg.fallbackAction?.resolvedEntities;
           const userMsg = messages[i - 1];
           if (userMsg?.role === LegendAIMessageRole.USER) {
             question = userMsg.text;
@@ -373,7 +418,7 @@ export const useLegendAIChatState = (
           setMessages,
         },
         pureExecutionContext,
-        undefined,
+        preResolvedEntities,
         modelContext,
       )
         .catch(noop)
@@ -401,7 +446,7 @@ export const useLegendAIChatState = (
     messages,
     selectedModelName,
     availableModelNames,
-    setSelectedModelName,
+    setSelectedModelName: handleSelectModel,
     askQuestion,
     askQuestionWithIntent,
     runFallbackAction,

package/src/legend-ai/stores/LegendAIJoinAnalysis.ts

@@ -0,0 +1,165 @@
+/**
+ * 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 {
+  type TDSServiceSchema,
+  TDSServiceSourceType,
+} from '../LegendAITypes.js';
+import {
+  AP_CALL_PATTERN,
+  buildColumnByNameIndex,
+  buildServiceByPIdIndex,
+  sharedColumnNames,
+} from './LegendAISqlHelpers.js';
+
+const MAX_JOINABLE_SUGGESTIONS = 3;
+const SAMPLE_VALUE_PREVIEW_COUNT = 3;
+const SUGGESTED_ROW_COUNT = 20;
+
+/** Parses a comma-separated `sampleValues` string into a lowercased value set. */
+export function parseSampleValueSet(
+  sampleValues: string | undefined,
+): Set<string> {
+  if (sampleValues === undefined) {
+    return new Set();
+  }
+  return new Set(
+    sampleValues
+      .split(',')
+      .map((v) => v.trim().toLowerCase())
+      .filter((v) => v.length > 0),
+  );
+}
+
+// True when the two access points share a required (non-nullable) column — a
+// candidate cross-access-point join key.
+function hasSharedRequiredKey(
+  apA: TDSServiceSchema,
+  apB: TDSServiceSchema,
+): boolean {
+  const colsA = buildColumnByNameIndex(apA.columns);
+  const colsB = buildColumnByNameIndex(apB.columns);
+  return sharedColumnNames(apA, apB).some((name) => {
+    const key = name.toLowerCase();
+    return (
+      colsA.get(key)?.nullable === false && colsB.get(key)?.nullable === false
+    );
+  });
+}
+
+// Resolves the access points referenced by a cross-`p()` join SQL to their
+// schemas via `findServiceForPId`, de-duplicated by pId.
+function resolveJoinedAccessPoints(
+  sql: string,
+  services: TDSServiceSchema[],
+): TDSServiceSchema[] {
+  const serviceByPId = buildServiceByPIdIndex(services);
+  const resolved: TDSServiceSchema[] = [];
+  const seenPIds = new Set<string>();
+  for (const match of sql.matchAll(AP_CALL_PATTERN)) {
+    const pId = match.groups?.pId;
+    if (pId === undefined || seenPIds.has(pId)) {
+      continue;
+    }
+    seenPIds.add(pId);
+    const svc = serviceByPId.get(pId);
+    if (svc?.sourceType === TDSServiceSourceType.ACCESS_POINT) {
+      resolved.push(svc);
+    }
+  }
+  return resolved;
+}
+
+/**
+ * Explains why a cross-access-point join returned 0 rows from grounded sample
+ * values. Returns undefined when this is not a cross-access-point join.
+ */
+export function buildCrossJoinZeroRowExplanation(
+  sql: string,
+  services: TDSServiceSchema[],
+): string | undefined {
+  const involved = resolveJoinedAccessPoints(sql, services);
+  const apA = involved[0];
+  const apB = involved[1];
+  if (apA === undefined || apB === undefined) {
+    return undefined;
+  }
+  const colsA = buildColumnByNameIndex(apA.columns);
+  const colsB = buildColumnByNameIndex(apB.columns);
+  for (const name of sharedColumnNames(apA, apB)) {
+    const key = name.toLowerCase();
+    const colA = colsA.get(key);
+    const colB = colsB.get(key);
+    const setA = parseSampleValueSet(colA?.sampleValues);
+    const setB = parseSampleValueSet(colB?.sampleValues);
+    if (setA.size === 0 || setB.size === 0) {
+      continue;
+    }
+    if (![...setA].some((v) => setB.has(v))) {
+      const previewA = colA?.sampleValues
+        ?.split(',')
+        .slice(0, SAMPLE_VALUE_PREVIEW_COUNT)
+        .map((v) => v.trim())
+        .join(', ');
+      const previewB = colB?.sampleValues
+        ?.split(',')
+        .slice(0, SAMPLE_VALUE_PREVIEW_COUNT)
+        .map((v) => v.trim())
+        .join(', ');
+      return `The join executed successfully but returned **0 rows**: **${apA.title}** and **${apB.title}** cover different data. Their shared column "${name}" has no values in common (e.g. ${previewA} vs ${previewB}), so there are no rows to join.`;
+    }
+  }
+  return `The join executed successfully but returned **0 rows**: no rows share the join key between **${apA.title}** and **${apB.title}**. These access points may cover different products or time periods.`;
+}
+
+/**
+ * Suggests alternative joins likely to return rows: same-group access points
+ * sharing a required key, plus a guaranteed single-access-point query.
+ */
+export function buildJoinablePairSuggestions(
+  selectedAPs: TDSServiceSchema[],
+  allAccessPoints: TDSServiceSchema[],
+): string[] {
+  const involved = selectedAPs.filter(
+    (s) => s.sourceType === TDSServiceSourceType.ACCESS_POINT,
+  );
+  const involvedPatterns = new Set(involved.map((s) => s.pattern));
+  const others = allAccessPoints.filter(
+    (s) =>
+      s.sourceType === TDSServiceSourceType.ACCESS_POINT &&
+      !involvedPatterns.has(s.pattern),
+  );
+  const candidates: string[] = [];
+  for (const base of involved) {
+    for (const other of others) {
+      if (
+        base.accessPointGroupTitle !== undefined &&
+        base.accessPointGroupTitle === other.accessPointGroupTitle &&
+        hasSharedRequiredKey(base, other)
+      ) {
+        candidates.push(
+          `Join "${base.title}" and "${other.title}" and show ${SUGGESTED_ROW_COUNT} rows`,
+        );
+      }
+    }
+  }
+  const suggestions = candidates.slice(0, MAX_JOINABLE_SUGGESTIONS - 1);
+  const first = involved[0];
+  if (first !== undefined) {
+    suggestions.push(`Show ${SUGGESTED_ROW_COUNT} rows from "${first.title}"`);
+  }
+  return suggestions.slice(0, MAX_JOINABLE_SUGGESTIONS);
+}

package/src/legend-ai/stores/LegendAISqlHelpers.ts

@@ -0,0 +1,82 @@
+/**
+ * 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 type { TDSColumnSchema, TDSServiceSchema } from '../LegendAITypes.js';
+
+export const AP_CALL_PATTERN = /(?<!\w)p\(\s*'(?<pId>[^']+)'\s*\)/g;
+export const HAS_LIMIT_PATTERN = /\bLIMIT\s+\d+/i;
+
+export function servicePId(service: TDSServiceSchema): string | undefined {
+  return service.dataProductPath === undefined
+    ? undefined
+    : `${service.dataProductPath}.${service.pattern.replace(/^\//, '')}`;
+}
+
+export function pureRelationColumnRef(column: string): string {
+  if (/^[A-Za-z_]\w*$/u.test(column)) {
+    return column;
+  }
+  const escaped = column.replaceAll("'", String.raw`\'`);
+  return `'${escaped}'`;
+}
+
+// Resolves the access-point schema whose `p('...')` identifier matches `pId`
+// (`<dataProductPath>.<pattern>`). Shared by the join sanitizers and analysis.
+export function findServiceForPId(
+  pId: string,
+  services: TDSServiceSchema[] | undefined,
+): TDSServiceSchema | undefined {
+  return services?.find((s) => servicePId(s) === pId);
+}
+
+// Builds a `pId → service` index for callers that resolve many `p('...')`
+// accessors in a loop, avoiding a linear scan (and per-call path strip) per
+// accessor.
+export function buildServiceByPIdIndex(
+  services: TDSServiceSchema[] | undefined,
+): Map<string, TDSServiceSchema> {
+  const index = new Map<string, TDSServiceSchema>();
+  for (const service of services ?? []) {
+    const pId = servicePId(service);
+    if (pId !== undefined) {
+      index.set(pId, service);
+    }
+  }
+  return index;
+}
+
+// Case-insensitive intersection of two access points' column names, returning
+// the names as cased in `apA`.
+export function sharedColumnNames(
+  apA: TDSServiceSchema,
+  apB: TDSServiceSchema,
+): string[] {
+  const bColumns = new Set(apB.columns.map((c) => c.name.toLowerCase()));
+  return apA.columns
+    .map((c) => c.name)
+    .filter((name) => bColumns.has(name.toLowerCase()));
+}
+
+// Indexes a column list by lowercased name for O(1) case-insensitive lookup.
+export function buildColumnByNameIndex(
+  columns: readonly TDSColumnSchema[],
+): Map<string, TDSColumnSchema> {
+  const index = new Map<string, TDSColumnSchema>();
+  for (const column of columns) {
+    index.set(column.name.toLowerCase(), column);
+  }
+  return index;
+}