WIP: Model with dependency injection support

[urso]

Note: this change is build on top of #168.

Tools, prompts, and even telemetry support feels a little involved to construct.

While vercel AI sdk supports LanguageModel middleware to add some customization to models, that I found that middleware to be lacking to some extend. This is because most functionality is implemented in streamText, generateText commands.

This change experiments with a way of adding dependency injection to the language model. Plus it sets a few defaults and wraps streamText/generateText instead of the more low level. This way all models automatically have telemetry and some other defaults enabled without us repeating the settings.

The alternative language model is implemented by the AugmentedLanguageModel (for lack of a better name) type. The model accepts a providerRegistry and a model or model-id. In case a ID is given we will lookup the actual LanguageModel right before calling into the vercel AI SDK.
The augmented model can also be configured with system prompts, tools, and toolsets. All these can be plain values, but optionally support dependency injection.

With these changes the full model can be defined as a const for easier sharing, for example with test code.

For example to define the agents chat model I didn't reuse the builders we normally use, but instead use the lower level functions and prompts:

type ChatModelDeps = { ... };

const chatModel = new AugmentedLanguageModel<ChatModelDeps>({
  providerRegistry: getProviderRegistry,
  baseModel: 'chat',
  metadata: {
    tags: ['chat']
  },
  systemPrompt: [commonSystemPrompt, chatSystemPrompt],
  toolsSets: [
    { tools: mcpToolset.listMCPTools },
    { tools: commonToolset },
    { tools: async ({ targetDb }) => getDBSQLTools(targetDb).toolset() },

    // Playbook support
    { tools: async ({ dbAccess, project }) => getPlaybookToolset(dbAccess, project.id) },

    // Common cloud provider DB support
    {
      tools: async ({ dbAccess, connection }) =>
        new CommonDBClusterTools(dbAccess, () => Promise.resolve(connection)).toolset()
    }
  ]
});

By adding all features to the model we can group prompts and tools when creating the model:

// AWS cloud provider support
chatModel.addSystemPrompts(({ project }) => (project.cloudProvider === 'aws' ? awsCloudProviderPrompt : ''));
chatModel.addToolSet({
  active: (deps?: ChatModelDeps) => deps?.project.cloudProvider === 'aws',
  tools: async ({ dbAccess, connection }) => {
    return new AWSDBClusterTools(dbAccess, () => Promise.resolve(connection)).toolset();
  }
});

Due to the model dependencies being typed Typescript will tell us if we miss a dependency when calling streamText/generateText:

chatModel.streamText({
  deps: { ... },
  messages,
})

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Updated (UTC)
agent	✅ Ready (Inspect)	Visit Preview	May 16, 2025 3:29pm