@finos/legend-lego 2.0.202 → 2.0.203

package/src/legend-ai/LegendAIServiceRetrieval.ts

@@ -0,0 +1,228 @@
+/**
+ * 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 { guaranteeNonNullable } from '@finos/legend-shared';
+import type { TDSServiceSchema } from './LegendAITypes.js';
+import {
+  splitIdentifierTokens,
+  tokenizeText,
+} from './LegendAIDocEnrichment.js';
+
+const BM25_K1 = 1.2;
+const BM25_B = 0.75;
+const FIELD_WEIGHT_TITLE = 3;
+const FIELD_WEIGHT_DESCRIPTION = 2;
+const MIN_FUZZY_TOKEN_LENGTH = 4;
+
+interface BM25Doc {
+  termFreq: Map<string, number>;
+  length: number;
+}
+
+interface BM25Index {
+  docs: BM25Doc[];
+  docFreq: Map<string, number>;
+  avgLength: number;
+}
+
+function tokensFromIdentifier(name: string): string[] {
+  return [...tokenizeText(name), ...splitIdentifierTokens(name)];
+}
+
+function addTokensToDoc(doc: BM25Doc, tokens: string[], weight: number): void {
+  for (const token of tokens) {
+    doc.termFreq.set(token, (doc.termFreq.get(token) ?? 0) + weight);
+    doc.length += weight;
+  }
+}
+
+function buildServiceDoc(svc: TDSServiceSchema): BM25Doc {
+  const doc: BM25Doc = { termFreq: new Map(), length: 0 };
+
+  addTokensToDoc(doc, tokensFromIdentifier(svc.title), FIELD_WEIGHT_TITLE);
+  if (svc.description) {
+    addTokensToDoc(
+      doc,
+      tokenizeText(svc.description),
+      FIELD_WEIGHT_DESCRIPTION,
+    );
+  }
+  for (const col of svc.columns) {
+    addTokensToDoc(doc, tokensFromIdentifier(col.name), 1);
+    if (col.documentation) {
+      addTokensToDoc(doc, tokenizeText(col.documentation), 1);
+    }
+  }
+  for (const param of svc.parameters) {
+    addTokensToDoc(doc, tokensFromIdentifier(param), 1);
+  }
+  for (const pf of svc.preFilters ?? []) {
+    addTokensToDoc(doc, tokenizeText(pf.property), 1);
+    if (pf.value !== undefined) {
+      addTokensToDoc(doc, tokenizeText(String(pf.value)), 1);
+    }
+  }
+  return doc;
+}
+
+function buildBM25Index(services: readonly TDSServiceSchema[]): BM25Index {
+  const docs = services.map(buildServiceDoc);
+  const docFreq = new Map<string, number>();
+  let totalLength = 0;
+  for (const doc of docs) {
+    totalLength += doc.length;
+    for (const term of doc.termFreq.keys()) {
+      docFreq.set(term, (docFreq.get(term) ?? 0) + 1);
+    }
+  }
+  return {
+    docs,
+    docFreq,
+    avgLength: docs.length === 0 ? 0 : totalLength / docs.length,
+  };
+}
+
+/**
+ * Standard Okapi BM25:
+ *   IDF(q) = ln((N - df + 0.5) / (df + 0.5) + 1)
+ *   tfNorm = tf * (k1 + 1) / (tf + k1 * (1 - b + b * |D| / avgdl))
+ *   score  = Σ IDF(q) * tfNorm
+ */
+function scoreBM25(
+  queryTerms: string[],
+  doc: BM25Doc,
+  index: BM25Index,
+): number {
+  const n = index.docs.length;
+  const avgdl = index.avgLength === 0 ? 1 : index.avgLength;
+  let score = 0;
+  for (const term of queryTerms) {
+    const tf = doc.termFreq.get(term);
+    if (!tf) {
+      continue;
+    }
+    const df = index.docFreq.get(term) ?? 0;
+    const idf = Math.log((n - df + 0.5) / (df + 0.5) + 1);
+    const tfNorm =
+      (tf * (BM25_K1 + 1)) /
+      (tf + BM25_K1 * (1 - BM25_B + (BM25_B * doc.length) / avgdl));
+    score += idf * tfNorm;
+  }
+  return score;
+}
+
+/** Rewrites a typoed query term to its closest indexed vocabulary entry. */
+function resolveQueryTermViaFuzzy(token: string, index: BM25Index): string {
+  if (token.length < MIN_FUZZY_TOKEN_LENGTH || index.docFreq.has(token)) {
+    return token;
+  }
+  for (const term of index.docFreq.keys()) {
+    if (term.length >= MIN_FUZZY_TOKEN_LENGTH && isFuzzyMatch(token, term)) {
+      return term;
+    }
+  }
+  return token;
+}
+
+/**
+ * Lexical pre-filter that ranks services against the user question using
+ * Okapi BM25. IDF down-weights corpus-wide common terms and length
+ * normalization demotes denormalized aggregate tables — no stopword list
+ * or specificity penalty needed. Query typos are recovered via
+ * {@link resolveQueryTermViaFuzzy} before scoring.
+ */
+export function preFilterServicesByRelevance(
+  question: string,
+  services: TDSServiceSchema[],
+  limit: number,
+): TDSServiceSchema[] {
+  const rawTerms = tokenizeText(question);
+  if (rawTerms.length === 0) {
+    return services.slice(0, limit);
+  }
+  const index = buildBM25Index(services);
+  const queryTerms = rawTerms.map((t) => resolveQueryTermViaFuzzy(t, index));
+  const scored = services.map((svc, i) => ({
+    svc,
+    score: scoreBM25(
+      queryTerms,
+      guaranteeNonNullable(index.docs[i], `BM25 doc missing for service ${i}`),
+      index,
+    ),
+  }));
+  scored.sort((a, b) => b.score - a.score);
+  return scored.slice(0, limit).map((s) => s.svc);
+}
+
+/**
+ * Returns `true` when two tokens are within Levenshtein edit distance 1 (for
+ * short tokens) or 2 (for longer tokens). Used to recover typoed query terms
+ * during retrieval and coverage checks.
+ */
+export function isFuzzyMatch(a: string, b: string): boolean {
+  const lenDiff = Math.abs(a.length - b.length);
+  if (lenDiff > 2) {
+    return false;
+  }
+  const maxDist = Math.min(a.length, b.length) <= 4 ? 1 : 2;
+  return levenshteinDistance(a, b) <= maxDist;
+}
+
+/**
+ * Computes the Levenshtein edit distance between two strings.
+ * Uses a single-row DP approach for space efficiency.
+ */
+export function levenshteinDistance(a: string, b: string): number {
+  if (a === b) {
+    return 0;
+  }
+  if (a.length === 0) {
+    return b.length;
+  }
+  if (b.length === 0) {
+    return a.length;
+  }
+  if (a.length > b.length) {
+    const tmp = a;
+    a = b;
+    b = tmp;
+  }
+
+  const aLen = a.length;
+  const bLen = b.length;
+  const row = new Array<number>(aLen + 1);
+
+  for (let i = 0; i <= aLen; i++) {
+    row[i] = i;
+  }
+
+  for (let j = 1; j <= bLen; j++) {
+    let prev = guaranteeNonNullable(row[0]);
+    row[0] = j;
+    for (let i = 1; i <= aLen; i++) {
+      const cost = a[i - 1] === b[j - 1] ? 0 : 1;
+      const current = guaranteeNonNullable(row[i]);
+      row[i] = Math.min(
+        current + 1, // deletion
+        guaranteeNonNullable(row[i - 1]) + 1, // insertion
+        prev + cost, // substitution
+      );
+      prev = current;
+    }
+  }
+
+  return guaranteeNonNullable(row[aLen]);
+}

package/src/legend-ai/LegendAITypes.ts

@@ -209,6 +209,14 @@ export class LegendAIServiceError extends Error {
   }
 }
 
+export class LegendAIUnsupportedEngineShapeError extends LegendAIServiceError {
+  override name = 'LegendAIUnsupportedEngineShapeError';
+
+  constructor(hint: string) {
+    super(hint, LegendAIErrorType.EXECUTION);
+  }
+}
+
 export enum LegendAIMessageRole {
   USER = 'user',
   ASSISTANT = 'assistant',
@@ -256,6 +264,7 @@ export class LegendAIAssistantMessage {
   isExecuting!: boolean;
   suggestedQueries!: string[];
   fallbackAction!: LegendAIFallbackAction | null;
+  queriedAccessPointGroups!: string[];
 }
 
 export type LegendAIMessage = LegendAIUserMessage | LegendAIAssistantMessage;
@@ -366,6 +375,70 @@ export class LegendAIProductMetadata {
   domainContext?: string;
 }
 
+// ────────────────────────────────────────────────────────────────────────────
+// Model context for local entity resolution (DataSpaces only)
+// ────────────────────────────────────────────────────────────────────────────
+
+export interface LegendAIModelProperty {
+  name: string;
+  type: string;
+  isCollection: boolean;
+  isOptional: boolean;
+}
+
+export interface LegendAIModelEntity {
+  path: string;
+  name: string;
+  description?: string;
+  properties: LegendAIModelProperty[];
+  isRootMapped?: boolean;
+  isQueryable?: boolean;
+  superTypes?: string[];
+}
+
+export interface LegendAIModelAssociation {
+  name: string;
+  leftEntity: string;
+  leftProperty: string;
+  rightEntity: string;
+  rightProperty: string;
+}
+
+export interface LegendAIColumnPropertyMapping {
+  columnName: string;
+  propertyPath: string;
+}
+
+export interface LegendAIParameterInfo {
+  name: string;
+  type: string;
+}
+
+export interface LegendAIEnumerationInfo {
+  path: string;
+  name: string;
+  values: string[];
+}
+
+export interface LegendAIExecutableInfo {
+  title: string;
+  description?: string;
+  rootEntityPath: string;
+  referencedEntityPaths?: string[];
+  columns?: string[];
+  queryTemplate?: string;
+  requiredParameters?: LegendAIParameterInfo[];
+  columnPropertyMappings?: LegendAIColumnPropertyMapping[];
+}
+
+export interface LegendAIModelContext {
+  entities: LegendAIModelEntity[];
+  associations: LegendAIModelAssociation[];
+  enumerations?: LegendAIEnumerationInfo[];
+  executables?: LegendAIExecutableInfo[];
+  dataspaceDescription?: string;
+}
+
 export enum LegendAIQuestionIntent {
   DATA_QUERY = 'data_query',
   METADATA = 'metadata',
@@ -474,6 +547,58 @@ export class QuestionIntentClassification {
   ambiguous!: boolean;
 }
 
+function maskEntityNames(question: string, entityNames: string[]): string {
+  const escaped = entityNames
+    .filter((name) => name.length > 0)
+    .map((name) =>
+      name.toLowerCase().replaceAll(/[.*+?^${}()|[\]\\]/g, String.raw`\$&`),
+    );
+  if (escaped.length === 0) {
+    return question;
+  }
+  return question.replaceAll(new RegExp(escaped.join('|'), 'g'), '___');
+}
+
+function classifyMixedSignal(
+  q: string,
+  metaScore: number,
+  dataScore: number,
+): QuestionIntentClassification {
+  if (dataScore >= metaScore * 2) {
+    return {
+      intent: LegendAIQuestionIntent.DATA_QUERY,
+      metaScore,
+      dataScore,
+      ambiguous: false,
+    };
+  }
+  if (PRODUCT_REFERENCE_PATTERN.test(q) || STRUCTURAL_KEYWORD_PATTERN.test(q)) {
+    return {
+      intent: LegendAIQuestionIntent.METADATA,
+      metaScore,
+      dataScore,
+      ambiguous: false,
+    };
+  }
+  if (CAPABILITY_DISCOVERY_PATTERNS.some((p) => p.test(q))) {
+    return {
+      intent: LegendAIQuestionIntent.METADATA,
+      metaScore,
+      dataScore,
+      ambiguous: false,
+    };
+  }
+  return {
+    intent:
+      metaScore > dataScore
+        ? LegendAIQuestionIntent.METADATA
+        : LegendAIQuestionIntent.DATA_QUERY,
+    metaScore,
+    dataScore,
+    ambiguous: true,
+  };
+}
+
 /**
  * Fast deterministic regex classifier (sync, < 1ms).
  * Returns both the resolved intent AND whether the result is ambiguous,
@@ -482,6 +607,7 @@ export class QuestionIntentClassification {
 export function classifyQuestionIntentFast(
   question: string,
   hasServices: boolean,
+  entityNames?: string[],
 ): QuestionIntentClassification {
   const q = question.toLowerCase().trim();
 
@@ -495,7 +621,12 @@ export function classifyQuestionIntentFast(
   }
 
   const metaScore = countPatternMatches(q, METADATA_SIGNAL_PATTERNS);
-  const dataScore = countPatternMatches(q, DATA_QUERY_SIGNAL_PATTERNS);
+  const qForDataScore =
+    entityNames && entityNames.length > 0 ? maskEntityNames(q, entityNames) : q;
+  const dataScore = countPatternMatches(
+    qForDataScore,
+    DATA_QUERY_SIGNAL_PATTERNS,
+  );
 
   if (metaScore > 0 && dataScore === 0) {
     const isStructural =
@@ -518,37 +649,8 @@ export function classifyQuestionIntentFast(
       ambiguous: false,
     };
   }
-
   if (metaScore > 0 && dataScore > 0) {
-    if (dataScore >= metaScore * 2) {
-      return {
-        intent: LegendAIQuestionIntent.DATA_QUERY,
-        metaScore,
-        dataScore,
-        ambiguous: false,
-      };
-    }
-    if (
-      PRODUCT_REFERENCE_PATTERN.test(q) ||
-      STRUCTURAL_KEYWORD_PATTERN.test(q)
-    ) {
-      return {
-        intent: LegendAIQuestionIntent.METADATA,
-        metaScore,
-        dataScore,
-        ambiguous: false,
-      };
-    }
-    const tentative =
-      metaScore > dataScore
-        ? LegendAIQuestionIntent.METADATA
-        : LegendAIQuestionIntent.DATA_QUERY;
-    return {
-      intent: tentative,
-      metaScore,
-      dataScore,
-      ambiguous: true,
-    };
+    return classifyMixedSignal(q, metaScore, dataScore);
   }
 
   return {
@@ -568,8 +670,9 @@ export function classifyQuestionIntentFast(
 export function classifyQuestionIntent(
   question: string,
   hasServices: boolean,
+  entityNames?: string[],
 ): LegendAIQuestionIntent {
-  return classifyQuestionIntentFast(question, hasServices).intent;
+  return classifyQuestionIntentFast(question, hasServices, entityNames).intent;
 }
 
 export async function extractParameterSchemas(
@@ -629,6 +732,7 @@ export interface LegendAIChatProps {
    * returned by the orchestrator. Required for the execute-after-generate flow.
    */
   pureExecutionContext?: QueryExplicitExecutionContextInfo;
+  modelContext?: LegendAIModelContext;
   availableScopes?: LegendAIScopeItem[];
   /** Called when the user clicks the close button in the chat header. */
   onClose?: () => void;
@@ -641,4 +745,6 @@ export interface LegendAIChatProps {
   onMessageFeedback?: (
     feedback: LegendAIMessageFeedback,
   ) => Promise<void> | void;
+  onRequestAccess?: (accessPointGroupTitle: string) => void;
+  contextBannerMessage?: string;
 }

package/src/legend-ai/LegendAI_LegendApplicationPlugin_Extension.ts

@@ -28,6 +28,7 @@ import type {
   LegendAIConversationTurn,
   LegendAIProductMetadata,
   LegendAIQuestionIntent,
+  LegendAIModelContext,
 } from './LegendAITypes.js';
 
 export class LegendAISqlExtractionResult {
@@ -63,6 +64,30 @@ export interface LegendAISemanticSearchResolutionDetails {
   data_product_coordinates: LegendAIOrchestratorDataProductCoordinates;
   root_entity: string;
   related_entities: string[];
+  enriched_business_context?: LegendAIEnrichedBusinessContext;
+}
+
+export interface LegendAIEnrichedBusinessContext {
+  naturalLanguageQuery?: string;
+  businessContextMatch?: LegendAIBusinessContextMatch;
+}
+
+export interface LegendAIBusinessContextMatch {
+  properties?: LegendAIBusinessContextProperty[];
+  additionalNlModelContext?: LegendAIAdditionalNlModelContext[];
+}
+
+export interface LegendAIBusinessContextProperty {
+  propertyName: string;
+  matchType?: string[];
+  propertyValues?: string[];
+  probablePropertyValues?: string[];
+}
+
+export interface LegendAIAdditionalNlModelContext {
+  id: string;
+  description: string;
+  category: string;
 }
 
 export class LegendAIOrchestratorRequest {
@@ -161,6 +186,8 @@ export abstract class LegendAI_LegendApplicationPlugin_Extension extends LegendA
     question: string,
     metadata: LegendAIProductMetadata,
     history?: LegendAIConversationTurn[],
+    services?: TDSServiceSchema[],
+    modelContextEnrichment?: string,
   ): string;
 
   /**
@@ -172,6 +199,7 @@ export abstract class LegendAI_LegendApplicationPlugin_Extension extends LegendA
     coordinates: string,
     history?: LegendAIConversationTurn[],
     metadata?: LegendAIProductMetadata,
+    modelContextEnrichment?: string,
   ): string;
 
   /**
@@ -194,6 +222,7 @@ export abstract class LegendAI_LegendApplicationPlugin_Extension extends LegendA
     question: string,
     accessPoints: TDSServiceSchema[],
     history?: LegendAIConversationTurn[],
+    modelContextEnrichment?: string,
   ): string;
 
   /**
@@ -205,8 +234,40 @@ export abstract class LegendAI_LegendApplicationPlugin_Extension extends LegendA
     question: string,
     accessPoints: TDSServiceSchema[],
     history?: LegendAIConversationTurn[],
+    modelContextEnrichment?: string,
   ): string;
 
+  buildAccessPointErrorCorrectionPrompt(
+    failedSql: string,
+    question: string,
+    executionError: string,
+    accessPoints: TDSServiceSchema[],
+    history?: LegendAIConversationTurn[],
+    modelContextEnrichment?: string,
+  ): string {
+    return this.buildAccessPointJudgePrompt(
+      failedSql,
+      `${question}\n\nEXECUTION ERROR — the previous SQL failed. Fix it using the SAME access points:\n${executionError}`,
+      accessPoints,
+      history,
+      modelContextEnrichment,
+    );
+  }
+
+  normalizeQuestion(
+    question: string,
+    _config: LegendAIConfig,
+  ): Promise<string> {
+    return Promise.resolve(question);
+  }
+
+  preWarmSchemaAnalysis(
+    _services: TDSServiceSchema[],
+    _config: LegendAIConfig,
+  ): void {
+    /* no-op by default — plugins override */
+  }
+
   /**
    * Send a prompt to the LLM service and return the raw response text.
    * The plugin manages conversation lifecycle internally — callers
@@ -252,6 +313,7 @@ export abstract class LegendAI_LegendApplicationPlugin_Extension extends LegendA
     dataProductCoordinates: LegendAIOrchestratorDataProductCoordinates,
     config: LegendAIConfig,
     pureExecutionContext?: QueryExplicitExecutionContextInfo,
+    modelContext?: LegendAIModelContext,
   ): Promise<LegendAIResolvedEntities>;
 
   /**

package/src/legend-ai/__test-utils__/LegendAITestUtils.ts

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import { guaranteeNonNullable } from '@finos/legend-shared';
 import { createMock } from '@finos/legend-shared/test';
 import type {
   MessageSetter,
@@ -72,6 +73,7 @@ export const TEST__makeAssistantMessage = (
   isExecuting: false,
   suggestedQueries: [],
   fallbackAction: null,
+  queriedAccessPointGroups: [],
   ...overrides,
 });
 export const TEST__seedAssistant = (setter: MessageSetter): void => {
@@ -80,6 +82,27 @@ export const TEST__seedAssistant = (setter: MessageSetter): void => {
     TEST__makeAssistantMessage(),
   ]);
 };
+
+/**
+ * Retrieve a message by index and assert it is an assistant message,
+ * narrowing the type via the discriminated `role` field without a cast.
+ */
+export const TEST__getAssistantMessage = (
+  messages: LegendAIMessage[],
+  index: number,
+): LegendAIAssistantMessage => {
+  const msg = guaranteeNonNullable(
+    messages[index],
+    `Expected message at index ${index}`,
+  );
+  if (msg.role !== LegendAIMessageRole.ASSISTANT) {
+    throw new Error(
+      `Expected assistant message at index ${index}, got role '${msg.role}'`,
+    );
+  }
+  return msg;
+};
+
 export const TEST_DATA__legendAIConfig: LegendAIConfig = {
   enabled: true,
   llmServiceUrl: 'http://localhost/llm',
@@ -140,13 +163,18 @@ export const TEST__createMockLegendAIPlugin = (
       _q: string,
       _s: TDSServiceSchema[],
       _h?: LegendAIConversationTurn[],
+      _m?: string,
     ) => 'ap generator prompt',
     buildAccessPointJudgePrompt: (
       _sql: string,
       _q: string,
       _s: TDSServiceSchema[],
       _h?: LegendAIConversationTurn[],
+      _m?: string,
     ) => 'ap judge prompt',
+    preWarmSchemaAnalysis: () => {
+      /* no-op in test */
+    },
     callLLM: createMock(),
     executeSql: createMock(),
     extractSqlFromResponse: (_a: string): LegendAISqlExtractionResult => ({