refactor the toolbox and tool registration

This commit is contained in:
Rob Colbert 2026-05-08 02:27:01 -04:00
parent 2dd3d55ecd
commit 2e8c4c4ae9

View File

@ -40,7 +40,21 @@ interface IAgentWorkflow {
type DroneSocket = Socket<ServerToClientEvents, ClientToServerEvents>;
const aiEnv: IAiEnvironment = {
NODE_ENV: env.NODE_ENV || "develop",
services: {
google: {
cse: {
apiKey: env.google.cse.apiKey,
engineId: env.google.cse.engineId,
},
},
},
};
class AgentService extends GadgetService {
private toolbox = new AiToolbox(aiEnv);
get name(): string {
return "AgentService";
}
@ -49,6 +63,15 @@ class AgentService extends GadgetService {
}
async start(): Promise<void> {
const googleSearchTool = new GoogleSearchTool(this.toolbox);
this.toolbox.register(googleSearchTool, [
ChatSessionMode.Plan,
ChatSessionMode.Build,
ChatSessionMode.Test,
ChatSessionMode.Ship,
ChatSessionMode.Develop,
]);
this.log.info("started");
}
@ -68,15 +91,15 @@ class AgentService extends GadgetService {
const onStreamChunk = async (chunk: IAiStreamChunk): Promise<void> => {
this.log.debug("stream chunk received", { chunk });
switch (chunk.type) {
case 'thinking':
case "thinking":
socket.emit("thinking", chunk.data);
break;
case 'response':
case "response":
socket.emit("response", chunk.data);
break;
case 'toolCall':
case "toolCall":
socket.emit(
"toolCall",
chunk.toolCallId!,
@ -124,8 +147,13 @@ class AgentService extends GadgetService {
);
// Check for model loading failure
if (response.doneReason === 'load' && !response.response && !response.thinking && (!response.toolCalls || response.toolCalls.length === 0)) {
throw new Error('Model failed to respond (still loading or error)');
if (
response.doneReason === "load" &&
!response.response &&
!response.thinking &&
(!response.toolCalls || response.toolCalls.length === 0)
) {
throw new Error("Model failed to respond (still loading or error)");
}
// Emit thinking content if present
@ -146,7 +174,8 @@ class AgentService extends GadgetService {
toolCall.callId,
toolCall.function.name,
toolCall.function.arguments,
response.toolCallResults?.find(r => r.callId === toolCall.callId)?.result || "",
response.toolCallResults?.find((r) => r.callId === toolCall.callId)
?.result || "",
);
}
}
@ -210,12 +239,15 @@ class AgentService extends GadgetService {
* (reasoning) output (if any).
*/
let content = "";
// Extract thinking and response from blocks
for (const block of turn.blocks) {
if (block.mode === 'thinking' && typeof block.content === 'string') {
if (block.mode === "thinking" && typeof block.content === "string") {
content += `<thinking>${block.content}</thinking>`;
} else if (block.mode === 'responding' && typeof block.content === 'string') {
} else if (
block.mode === "responding" &&
typeof block.content === "string"
) {
if (content && content.length) {
content += "\n";
}
@ -245,27 +277,7 @@ class AgentService extends GadgetService {
}
private getToolsForMode(mode: ChatSessionMode): any[] {
const aiEnv: IAiEnvironment = {
NODE_ENV: env.NODE_ENV || "develop",
services: {
google: {
cse: {
apiKey: env.google.cse.apiKey,
engineId: env.google.cse.engineId,
},
},
},
};
const toolbox = new AiToolbox(aiEnv);
const googleSearchTool = new GoogleSearchTool(toolbox);
toolbox.register(googleSearchTool, [
ChatSessionMode.Plan,
ChatSessionMode.Build,
ChatSessionMode.Test,
ChatSessionMode.Ship,
ChatSessionMode.Develop,
]);
return Array.from(toolbox.getModeSet(mode) || []);
return Array.from(this.toolbox.getModeSet(mode) || []);
}
}