WIP: Model with dependency injection support

apps/dbagent/package.json

@@ -55,6 +55,7 @@
     "@opentelemetry/semantic-conventions": "^1.32.0",
     "@tailwindcss/postcss": "^4.1.6",
     "@tanstack/react-query": "^5.75.7",
+    "@types/ndjson": "^2.0.4",
     "@vercel/functions": "^2.0.3",
     "ai": "^4.3.15",
     "bytes": "^3.1.2",
@@ -72,9 +73,11 @@
     "geist": "^1.4.2",
     "import-in-the-middle": "^1.13.1",
     "kysely": "^0.28.2",
+    "langfuse": "^3.37.2",
     "langfuse-vercel": "^3.37.2",
     "litellm-api": "^0.0.3",
     "lucide-react": "^0.509.0",
+    "ndjson": "^2.0.0",
     "next": "^15.3.2",
     "next-auth": "5.0.0-beta.25",
     "next-themes": "^0.4.6",
@@ -118,6 +121,7 @@
     "@types/react-dom": "^19.1.3",
     "@types/react-syntax-highlighter": "^15.5.13",
     "autoprefixer": "^10.4.21",
+    "cmd-ts": "^0.13.0",
     "dockerode": "^4.0.6",
     "dotenv": "^16.5.0",
     "filenamify": "^6.0.0",

apps/dbagent/playground/default/index.ts

@@ -1,29 +1,38 @@
 import { Agent } from '@mastra/core/agent';
-import { AnswerRelevancyMetric, PromptAlignmentMetric } from '@mastra/evals/llm';
+import { AnswerRelevancyMetric } from '@mastra/evals/llm';
 import { CompletenessMetric } from '@mastra/evals/nlp';
 import { CloudProvider } from '~/lib/db/schema';
 
 import { openai } from '@ai-sdk/openai';
+import { RuntimeContext } from '@mastra/core/runtime-context';
+import { z } from 'zod';
 import { getChatSystemPrompt, getMonitoringSystemPrompt } from '~/lib/ai/agent';
-import { getBuiltinProviderRegistry } from '~/lib/ai/providers';
+import { getProviderRegistry } from '~/lib/ai/providers';
 import { buildPlaygroundTools } from '../tools';
 
-/* eslint-disable no-process-env */
-const defaultModel = getBuiltinProviderRegistry()
-  .languageModel(process.env.MASTRA_MODEL ?? 'chat')
-  .instance();
-const cloudProvider = (process.env.MASTRA_CLOUD_PROVIDER ?? 'aws') as CloudProvider;
-const defaultTools = buildPlaygroundTools({
-  projectConnection: process.env.MASTRA_PROJECT_CONNECTION ?? undefined,
-  dbUrl: process.env.MASTRA_DB_URL ?? undefined,
-  userId: process.env.MASTRA_USER_ID ?? undefined
+export const runtimeContextSchema = z.object({
+  cloudProvider: z.enum(['other', 'aws', 'gcp']).default('other'),
+  model_id: z.string().default('chat'),
+  projectConnection: z.string().optional(),
+  dbUrl: z.string().optional(),
+  userId: z.string().optional()
 });
-/* eslint-enable no-process-env */
-
-const chatPrompt = getChatSystemPrompt({ cloudProvider });
-const monitoringPrompt = getMonitoringSystemPrompt({ cloudProvider });
 
 export function createAgents() {
+  /* eslint-disable no-process-env */
+  const runtimeContext = new RuntimeContext(
+    Object.entries({
+      cloudProvider: (process.env.MASTRA_CLOUD_PROVIDER ?? 'aws') as CloudProvider,
+      model_id: process.env.MASTRA_MODEL ?? 'chat',
+      projectConnection: process.env.MASTRA_PROJECT_CONNECTION ?? undefined,
+      dbUrl: process.env.MASTRA_DB_URL ?? undefined,
+      userId: process.env.MASTRA_USER_ID ?? undefined
+    })
+  );
+  // verify runtimeContextSchema
+  console.log(runtimeContextSchema.parse(runtimeContext));
+  /* eslint-enable no-process-env */
+
   const evals = {
     completeness: new CompletenessMetric(),
     relevancy: new AnswerRelevancyMetric(openai('gpt-4o-mini'), {
@@ -32,33 +41,44 @@ export function createAgents() {
     })
   };
 
+  const defaultOptions = {
+    defaultGenerateOptions: { runtimeContext },
+    defaultStreamOptions: { runtimeContext },
+    model: getModel,
+    tools: getTools,
+    evals: {
+      ...evals
+    }
+  };
+
   // TODO: create agent per supported model (assuming we have the keys)?
   return {
     monitoringAgent: new Agent({
       name: 'monitoring-agent',
-      instructions: monitoringPrompt,
-      model: defaultModel,
-      tools: defaultTools,
-      evals: {
-        ...evals,
-        'prompt-alignment': new PromptAlignmentMetric(openai('gpt-4o-mini'), {
-          instructions: [monitoringPrompt],
-          scale: 1 // Scale for the final score
-        })
+      ...defaultOptions,
+      instructions: ({ runtimeContext }) => {
+        return getMonitoringSystemPrompt({ cloudProvider: runtimeContext.get('cloudProvider') });
       }
     }),
     chatAgent: new Agent({
       name: 'chat-agent',
-      instructions: chatPrompt,
-      model: defaultModel,
-      tools: defaultTools,
-      evals: {
-        ...evals,
-        'prompt-alignment': new PromptAlignmentMetric(openai('gpt-4o-mini'), {
-          instructions: [chatPrompt],
-          scale: 1 // Scale for the final score
-        })
+      ...defaultOptions,
+      instructions: ({ runtimeContext }) => {
+        return getChatSystemPrompt({ cloudProvider: runtimeContext.get('cloudProvider') });
       }
     })
   };
 }
+
+const getModel = async ({ runtimeContext }: { runtimeContext: RuntimeContext }) => {
+  const { model_id } = runtimeContextSchema.parse(runtimeContext);
+  console.log('model_id', model_id);
+  const registry = await getProviderRegistry();
+  return registry.languageModel(model_id || 'chat', true).instance();
+};
+
+const getTools = async ({ runtimeContext }: { runtimeContext: RuntimeContext }) => {
+  const entries = runtimeContextSchema.parse(runtimeContext);
+  console.log('entries', entries);
+  return await buildPlaygroundTools(entries);
+};

apps/dbagent/playground/index.ts

@@ -1,5 +1,80 @@
+import { Agent } from '@mastra/core/agent';
+import { AnswerRelevancyMetric } from '@mastra/evals/llm';
+import { CompletenessMetric } from '@mastra/evals/nlp';
+import { CloudProvider } from '~/lib/db/schema';
+
+import { openai } from '@ai-sdk/openai';
 import { Mastra } from '@mastra/core';
-import { createAgents } from './default';
+import { RuntimeContext } from '@mastra/core/runtime-context';
+import { z } from 'zod';
+import { getChatSystemPrompt, getMonitoringSystemPrompt } from '~/lib/ai/agent';
+import { getProviderRegistry } from '~/lib/ai/providers';
+import { buildPlaygroundTools } from './tools';
+
+/* eslint-disable no-process-env */
+const cloudProvider = process.env.MASTRA_CLOUD_PROVIDER ?? ('aws' as CloudProvider);
+
+export const runtimeContextSchema = z.object({
+  cloudProvider: z.enum(['other', 'aws', 'gcp']).default(cloudProvider as any as CloudProvider),
+  model_id: z.string().default(process.env.MASTRA_MODEL ?? 'chat'),
+  projectConnection: process.env.MASTRA_PROJECT_CONNECTION
+    ? z.string().default(process.env.MASTRA_PROJECT_CONNECTION)
+    : z.string().optional(),
+  dbUrl: process.env.MASTRA_DB_URL ? z.string().default(process.env.MASTRA_DB_URL) : z.string().optional(),
+  userId: process.env.MASTRA_USER_ID ? z.string().default(process.env.MASTRA_USER_ID) : z.string().optional()
+});
+/* eslint-enable no-process-env */
+
+// verify runtimeContextSchema
+
+export function createAgents() {
+  const evals = {
+    completeness: new CompletenessMetric(),
+    relevancy: new AnswerRelevancyMetric(openai('gpt-4o-mini'), {
+      uncertaintyWeight: 0.3, // Weight for 'unsure' verdicts
+      scale: 1 // Scale for the final score
+    })
+  };
+
+  const defaultOptions = {
+    model: getModel,
+    tools: getTools,
+    evals: {
+      ...evals
+    }
+  };
+
+  // TODO: create agent per supported model (assuming we have the keys)?
+  return {
+    monitoringAgent: new Agent({
+      name: 'monitoring-agent',
+      ...defaultOptions,
+      instructions: ({ runtimeContext }) => {
+        return getMonitoringSystemPrompt({ cloudProvider: runtimeContext.get('cloudProvider') });
+      }
+    }),
+    chatAgent: new Agent({
+      name: 'chat-agent',
+      ...defaultOptions,
+      instructions: ({ runtimeContext }) => {
+        return getChatSystemPrompt({ cloudProvider: runtimeContext.get('cloudProvider') });
+      }
+    })
+  };
+}
+
+const getModel = async ({ runtimeContext }: { runtimeContext: RuntimeContext }) => {
+  const { model_id } = runtimeContextSchema.parse(runtimeContext);
+  console.log('model_id', model_id);
+  const registry = await getProviderRegistry();
+  return registry.languageModel(model_id || 'chat', true).instance();
+};
+
+const getTools = async ({ runtimeContext }: { runtimeContext: RuntimeContext }) => {
+  const entries = runtimeContextSchema.parse(runtimeContext);
+  console.log('entries', entries);
+  return await buildPlaygroundTools(entries);
+};
 
 export const mastra = new Mastra({
   agents: createAgents()

apps/dbagent/playground/tools/index.ts

@@ -2,17 +2,18 @@ import { Tool } from 'ai';
 import {
   builtinPlaybookToolset,
   commonToolset,
-  DBSQLTools,
   getDBClusterTools,
-  mergeToolsets,
-  playbookTools,
-  Toolset
+  getDBSQLTools,
+  getPlaybookToolset,
+  getProjectClusterService,
+  projectPlaybookService,
+  targetDBService
 } from '~/lib/ai/tools';
 import { getConnectionByName, getDefaultConnection } from '~/lib/db/connections';
 import { DBUserAccess } from '~/lib/db/db';
 import { getProjectByName } from '~/lib/db/projects';
 import { CloudProvider, Connection } from '~/lib/db/schema';
-import { getTargetDbPool, Pool } from '~/lib/targetdb/db';
+import { getTargetDbPool } from '~/lib/targetdb/db';
 
 type PlaygroundToolsConfig = {
   projectConnection?: string;
@@ -21,8 +22,8 @@ type PlaygroundToolsConfig = {
   cloudProvider?: CloudProvider;
 };
 
-export function buildPlaygroundTools(config: PlaygroundToolsConfig): Record<string, Tool> {
-  const toolsets: Toolset[] = [commonToolset];
+export async function buildPlaygroundTools(config: PlaygroundToolsConfig): Promise<Record<string, Tool>> {
+  const toolsets: Record<string, Tool>[] = [commonToolset];
   const connString: string | null = config.dbUrl ?? null;
 
   // Set up user ID provider if configured
@@ -37,79 +38,56 @@ export function buildPlaygroundTools(config: PlaygroundToolsConfig): Record<stri
       }
 
       const db = new DBUserAccess(config.userId);
-      const connGetter = createConnectionGetter(config.userId, projectName, connectionName);
+      const conn = await getConnection(config.userId, projectName, connectionName);
       if (config.cloudProvider) {
-        toolsets.push(getDBClusterTools(db, connGetter, config.cloudProvider));
+        const clusterService = getProjectClusterService(db, conn, config.cloudProvider);
+        toolsets.push(getDBClusterTools(clusterService));
       }
 
-      const playbookToolset = new playbookTools(db, async () => {
-        const conn = await connGetter();
-        return { projectId: conn.projectId };
-      });
+      const playbookToolset = getPlaybookToolset(projectPlaybookService(db, conn.projectId));
       toolsets.push(playbookToolset);
       hasPlaybookToolset = true;
 
       if (!connString) {
-        const targetDBPoolGetter = createTargetDBPoolGetter(connGetter);
-        toolsets.push(new DBSQLTools(targetDBPoolGetter));
+        const targetDBPool = getTargetDbPool(conn.connectionString);
+        toolsets.push(getDBSQLTools({ db: targetDBService(targetDBPool) }));
         hasDBTools = true;
       }
     }
   }
 
   if (connString && !hasDBTools) {
-    toolsets.push(new DBSQLTools(getTargetDbPool(connString)));
+    toolsets.push(getDBSQLTools({ db: targetDBService(getTargetDbPool(connString)) }));
   }
   if (!hasPlaybookToolset) {
     toolsets.push(builtinPlaybookToolset);
   }
 
-  return mergeToolsets(...toolsets);
+  return toolsets.reduce((acc, toolset) => {
+    return {
+      ...acc,
+      ...toolset
+    };
+  }, {});
 }
 
-function createConnectionGetter(
-  userId: string,
-  projectName: string,
-  connectionName?: string
-): () => Promise<Connection> {
-  let cachedConnection: Connection | null = null;
-  return async () => {
-    if (cachedConnection) {
-      return cachedConnection;
-    }
-
-    const db = new DBUserAccess(userId);
-
-    // First find the project by name where user has access
-    const project = await getProjectByName(db, projectName);
-    if (!project) {
-      throw new Error(`Project "${projectName}" not found or access denied`);
-    }
-
-    const connection = connectionName
-      ? await getConnectionByName(db, project.id, connectionName)
-      : await getDefaultConnection(db, project.id);
-    if (!connection) {
-      throw new Error(
-        connectionName
-          ? `Connection "${connectionName}" not found in project "${projectName}"`
-          : `No default connection found in project "${projectName}"`
-      );
-    }
+export async function getConnection(userId: string, projectName: string, connectionName?: string): Promise<Connection> {
+  const db = new DBUserAccess(userId);
+  const project = await getProjectByName(db, projectName);
+  if (!project) {
+    throw new Error(`Project "${projectName}" not found or access denied`);
+  }
 
-    cachedConnection = connection;
-    return connection;
-  };
-}
+  const connection = connectionName
+    ? await getConnectionByName(db, project.id, connectionName)
+    : await getDefaultConnection(db, project.id);
+  if (!connection) {
+    throw new Error(
+      connectionName
+        ? `Connection "${connectionName}" not found in project "${projectName}"`
+        : `No default connection found in project "${projectName}"`
+    );
+  }
 
-function createTargetDBPoolGetter(connGetter: () => Promise<Connection>) {
-  let cachedPool: Pool | null = null;
-  return async () => {
-    if (cachedPool) {
-      return cachedPool;
-    }
-    const connection = await connGetter();
-    cachedPool = getTargetDbPool(connection.connectionString);
-    return cachedPool;
-  };
+  return connection;
 }

apps/dbagent/scripts/eval/simple.ts

@@ -0,0 +1,14 @@
+'use server';
+
+import Langfuse from 'langfuse';
+import { evalAllTurn, LLMTurn } from '~/evals/lib/llmcheck';
+import { conciseAnswerMetric, finalAnswerMetric, Measure } from '~/evals/lib/llmcheck/metrics';
+
+const metrics = [conciseAnswerMetric(), finalAnswerMetric()];
+
+export async function evalFunction(langfuse: Langfuse, turn: LLMTurn): Promise<Record<string, Measure>> {
+  console.log(turn);
+
+  const evals = await evalAllTurn(turn, metrics);
+  return evals;
+}

apps/dbagent/scripts/lftool/env.ts

@@ -0,0 +1,8 @@
+import { config } from '@dotenvx/dotenvx';
+
+// Load environment variables first
+config({
+  path: ['.env', '.env.local'],
+  strict: false,
+  ignore: ['MISSING_ENV_FILE']
+});