@@ -3,15 +3,16 @@ import { AiPluginOptions, CommandTool } from './types';
3
3
import CommandKit , { getCommandKit , Logger } from 'commandkit' ;
4
4
import { AiContext } from './context' ;
5
5
import { Collection , Events , Message } from 'discord.js' ;
6
- import { tool , Tool , generateText } from 'ai' ;
6
+ import { tool , Tool , generateText , stepCountIs , ModelMessage } from 'ai' ;
7
7
import { getAiWorkerContext , runInAiWorkerContext } from './ai-context-worker' ;
8
8
import { getAvailableCommands } from './tools/get-available-commands' ;
9
9
import { getChannelById } from './tools/get-channel-by-id' ;
10
10
import { getCurrentClientInfo } from './tools/get-current-client-info' ;
11
11
import { getGuildById } from './tools/get-guild-by-id' ;
12
12
import { getUserById } from './tools/get-user-by-id' ;
13
- import { AiMessage , getAIConfig } from './configure' ;
13
+ import { getAIConfig } from './configure' ;
14
14
import { augmentCommandKit } from './augmentation' ;
15
+ import { ToolParameterType } from './tools/common' ;
15
16
16
17
/**
17
18
* Represents the configuration options for the AI plugin scoped to a specific command.
@@ -24,7 +25,7 @@ export interface AiConfig {
24
25
/**
25
26
* A zod schema defining the parameters that the AI command accepts.
26
27
*/
27
- parameters : any ;
28
+ inputSchema : ToolParameterType ;
28
29
}
29
30
30
31
const defaultTools : Record < string , Tool > = {
@@ -99,11 +100,19 @@ export class AiPlugin extends RuntimePlugin<AiPluginOptions> {
99
100
await runInAiWorkerContext ( ctx , message , async ( ) => {
100
101
const systemPrompt = await prepareSystemPrompt ( ctx , message ) ;
101
102
const prompt = await preparePrompt ( ctx , message ) ;
102
- const { model, abortSignal, maxSteps, ...modelOptions } =
103
- await selectAiModel ( ctx , message ) ;
104
-
105
- const promptOrMessage =
106
- typeof prompt === 'string' ? { prompt } : { messages : prompt } ;
103
+ const {
104
+ model,
105
+ abortSignal,
106
+ stopWhen,
107
+ prompt : _prompt ,
108
+ ...modelOptions
109
+ } = await selectAiModel ( ctx , message ) ;
110
+
111
+ const promptOrMessage = (
112
+ typeof prompt === 'string' ? { prompt } : { messages : prompt }
113
+ ) as
114
+ | { prompt : string ; messages : never }
115
+ | { messages : ModelMessage [ ] ; prompt : never } ;
107
116
108
117
await onProcessingStart ( ctx , message ) ;
109
118
@@ -112,7 +121,7 @@ export class AiPlugin extends RuntimePlugin<AiPluginOptions> {
112
121
model,
113
122
abortSignal : abortSignal ?? AbortSignal . timeout ( 60_000 ) ,
114
123
system : systemPrompt ,
115
- maxSteps : maxSteps ?? 5 ,
124
+ stopWhen : stopWhen ?? stepCountIs ( 5 ) ,
116
125
...modelOptions ,
117
126
tools : {
118
127
// Include built-in least significant tools if not disabled
@@ -181,12 +190,12 @@ export class AiPlugin extends RuntimePlugin<AiPluginOptions> {
181
190
182
191
const cmdTool = tool ( {
183
192
description,
184
- parameters : cmd . data . aiConfig . parameters ,
193
+ inputSchema : cmd . data . aiConfig . inputSchema ,
185
194
async execute ( params ) {
186
195
const config = getAIConfig ( ) ;
187
196
const ctx = getAiWorkerContext ( ) ;
188
197
189
- ctx . ctx . setParams ( params ) ;
198
+ ctx . ctx . setParams ( params as Record < string , unknown > ) ;
190
199
191
200
try {
192
201
const target = await commandkit . commandHandler . prepareCommandRun (
0 commit comments