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>; 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 { class AgentService extends GadgetService {
private toolbox = new AiToolbox(aiEnv);
get name(): string { get name(): string {
return "AgentService"; return "AgentService";
} }
@ -49,6 +63,15 @@ class AgentService extends GadgetService {
} }
async start(): Promise<void> { 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"); this.log.info("started");
} }
@ -70,13 +93,13 @@ class AgentService extends GadgetService {
this.log.debug("stream chunk received", { chunk }); this.log.debug("stream chunk received", { chunk });
switch (chunk.type) { switch (chunk.type) {
case 'thinking': case "thinking":
socket.emit("thinking", chunk.data); socket.emit("thinking", chunk.data);
break; break;
case 'response': case "response":
socket.emit("response", chunk.data); socket.emit("response", chunk.data);
break; break;
case 'toolCall': case "toolCall":
socket.emit( socket.emit(
"toolCall", "toolCall",
chunk.toolCallId!, chunk.toolCallId!,
@ -124,8 +147,13 @@ class AgentService extends GadgetService {
); );
// Check for model loading failure // Check for model loading failure
if (response.doneReason === 'load' && !response.response && !response.thinking && (!response.toolCalls || response.toolCalls.length === 0)) { if (
throw new Error('Model failed to respond (still loading or error)'); 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 // Emit thinking content if present
@ -146,7 +174,8 @@ class AgentService extends GadgetService {
toolCall.callId, toolCall.callId,
toolCall.function.name, toolCall.function.name,
toolCall.function.arguments, toolCall.function.arguments,
response.toolCallResults?.find(r => r.callId === toolCall.callId)?.result || "", response.toolCallResults?.find((r) => r.callId === toolCall.callId)
?.result || "",
); );
} }
} }
@ -213,9 +242,12 @@ class AgentService extends GadgetService {
// Extract thinking and response from blocks // Extract thinking and response from blocks
for (const block of turn.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>`; 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) { if (content && content.length) {
content += "\n"; content += "\n";
} }
@ -245,27 +277,7 @@ class AgentService extends GadgetService {
} }
private getToolsForMode(mode: ChatSessionMode): any[] { private getToolsForMode(mode: ChatSessionMode): any[] {
const aiEnv: IAiEnvironment = { return Array.from(this.toolbox.getModeSet(mode) || []);
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) || []);
} }
} }