feat(migration): migrate lb4-llm-chat-extension from LangGraph to Mastra [WIP]

package.json

@@ -127,33 +127,28 @@
     "@loopback/core": "^6.1.11"
   },
   "dependencies": {
-    "@langchain/community": "^0.3.50",
-    "@langchain/core": "^0.3.80",
-    "@langchain/langgraph": "^0.4.4",
     "@loopback/repository": "^7.0.14",
     "@sourceloop/chat-service": "15.0.3",
     "@sourceloop/core": "17.0.3",
     "@sourceloop/file-utils": "0.3.7",
-    "langchain": "^0.3.37",
+    "ai": "^6.0.168",
     "loopback4-authentication": "^12.2.0",
     "loopback4-authorization": "^7.0.3",
     "tslib": "^2.8.0",
     "winston": "^3.15.0"
   },
   "devDependencies": {
+    "@ai-sdk/amazon-bedrock": "^4.0.97",
+    "@ai-sdk/anthropic": "^3.0.72",
+    "@ai-sdk/cerebras": "^2.0.46",
+    "@ai-sdk/google": "^3.0.65",
+    "@ai-sdk/groq": "^3.0.36",
+    "@ai-sdk/openai": "^3.0.54",
     "@commitlint/cli": "^19.8.1",
     "@commitlint/config-conventional": "^16.2.1",
     "@google/generative-ai": "^0.24.1",
     "@istanbuljs/nyc-config-typescript": "^1.0.2",
-    "@langchain/anthropic": "^0.3.26",
-    "@langchain/aws": "^0.1.13",
-    "@langchain/cerebras": "^0.0.2",
-    "@langchain/google-genai": "^0.2.16",
-    "@langchain/groq": "^0.2.3",
-    "@langchain/ollama": "^0.2.3",
-    "@langchain/openai": "^0.6.11",
     "@langfuse/core": "^4.4.2",
-    "@langfuse/langchain": "^4.4.2",
     "@loopback/build": "^11.0.9",
     "@loopback/eslint-config": "^15.0.5",
     "@loopback/testlab": "^7.0.9",
@@ -175,6 +170,7 @@
     "jsonwebtoken": "^9.0.2",
     "loopback-connector-sqlite3": "^3.0.0",
     "mochawesome": "^7.1.3",
+    "ollama-ai-provider": "^1.2.0",
     "pg": "^8.16.3",
     "semantic-release": "^25.0.1",
     "source-map-support": "^0.5.21",

src/__tests__/acceptance/generation.controllers.acceptance.ts

@@ -1,11 +1,11 @@
-import {BindingScope} from '@loopback/core';
-import {Client, expect} from '@loopback/testlab';
-import {DbQueryAIExtensionBindings, IDataSetStore} from '../../components';
-import {LLMStreamEvent, LLMStreamEventType} from '../../graphs';
-import {AiIntegrationBindings} from '../../keys';
-import {PermissionKey} from '../../permissions';
-import {HttpTransport} from '../../transports';
-import {TestApp} from '../fixtures/test-app';
+import { BindingScope } from '@loopback/core';
+import { Client, expect } from '@loopback/testlab';
+import { DbQueryAIExtensionBindings, IDataSetStore } from '../../components';
+import { LLMStreamEvent, LLMStreamEventType } from '../../types/events';
+import { AiIntegrationBindings } from '../../keys';
+import { PermissionKey } from '../../permissions';
+import { HttpTransport } from '../../transports';
+import { TestApp } from '../fixtures/test-app';
 import {
   buildToken,
   seedCurrencies,
@@ -30,7 +30,7 @@ describe('GenerationController', () => {
   });
 
   before('setupApplication', async () => {
-    ({app, client} = await setupApplication({}));
+    ({ app, client } = await setupApplication({}));
     app
       .bind(AiIntegrationBindings.Transport)
       .toClass(HttpTransport)

src/__tests__/db-query/unit/db-knowledge-graph.service.unit.ts

@@ -1,20 +1,44 @@
-import {expect, sinon} from '@loopback/testlab';
-import {DbKnowledgeGraphService} from '../../../components';
-import {EmbeddingProvider, LLMProvider} from '../../../types';
+import { expect } from '@loopback/testlab';
+import { DbKnowledgeGraphService } from '../../../components';
+import { EmbeddingProvider, LLMProvider } from '../../../types';
+import { EmbeddingModelV2 } from '@ai-sdk/provider';
+import { createFakeStreamingLanguageModel } from '../../fixtures/fake-ai-models';
+
+function createTableEmbeddingModel(): EmbeddingModelV2<string> {
+  return {
+    specificationVersion: 'v2',
+    provider: 'test',
+    modelId: 'test-embedding',
+    maxEmbeddingsPerCall: null,
+    supportsParallelCalls: true,
+    doEmbed: async ({ values }: { values: string[] }) => {
+      const embeddings = values.map((v: string) => {
+        if (v.startsWith('employee_salaries')) return [0.1, 0.2, 0.3];
+        if (v.startsWith('employees')) return [0.1, 0.2, 0.3];
+        if (v.startsWith('orders')) return [0.9, 0.8, 0.7];
+        return [0.1, 0.2, 0.6];
+      });
+      return { embeddings, warnings: [] };
+    },
+  } as unknown as EmbeddingModelV2<string>;
+}
 
 describe(`DbKnowledgeGraphService Unit`, function () {
   let service: DbKnowledgeGraphService;
-  let llmStub: sinon.SinonStub;
-  let embedStub: sinon.SinonStub;
+  let fakeModel: LLMProvider;
 
   beforeEach(() => {
-    llmStub = sinon.stub();
-    embedStub = sinon.stub();
+    fakeModel = createFakeStreamingLanguageModel(
+      JSON.stringify({
+        concept: 'employees',
+        description: 'test description',
+        domain: 'test domain',
+        confidence: 0.9,
+      }),
+    ) as unknown as LLMProvider;
     service = new DbKnowledgeGraphService(
-      llmStub as unknown as LLMProvider,
-      {
-        embedDocuments: embedStub,
-      } as unknown as EmbeddingProvider,
+      fakeModel,
+      createTableEmbeddingModel() as unknown as EmbeddingProvider,
       {
         models: [],
         knowledgeGraph: {
@@ -28,26 +52,6 @@ describe(`DbKnowledgeGraphService Unit`, function () {
   });
 
   it('should generate a knowledge graph for a schema and should be able to find from it', async () => {
-    embedStub.callsFake(async doc => {
-      if (doc[0].startsWith('employee_salaries')) {
-        return [[0.1, 0.2, 0.3]];
-      }
-      if (doc[0].startsWith('employees')) {
-        return [[0.1, 0.2, 0.3]];
-      }
-      if (doc[0].startsWith('orders')) {
-        return [[0.9, 0.8, 0.7]];
-      }
-      return [[0.1, 0.2, 0.6]];
-    });
-    llmStub.resolves({
-      content: JSON.stringify({
-        concept: 'employees',
-        description: 'test description',
-        domain: 'test domain',
-        confidence: 0.9,
-      }),
-    });
     const schema = {
       tables: {
         // eslint-disable-next-line @typescript-eslint/naming-convention