@cogitator-ai/self-modifying 17.0.17 → 18.0.0

package/src/tool-generation/tool-generator.ts

@@ -1,5 +1,6 @@
 import type {
   Tool,
+  ToolSchema,
   LLMBackend,
   CapabilityGap,
   GeneratedTool,
@@ -8,7 +9,8 @@ import type {
 } from '@cogitator-ai/types';
 import { z, type ZodType } from 'zod';
 import { ToolValidator } from './tool-validator';
-import { ToolSandbox } from './tool-sandbox';
+import { ToolSandbox, type SandboxTestCase } from './tool-sandbox';
+import { llmChat } from '../utils/llm-helper';
 import {
   TOOL_GENERATION_SYSTEM_PROMPT,
   buildToolGenerationPrompt,
@@ -22,6 +24,10 @@ export interface ToolGeneratorOptions {
   model?: string;
 }
 
+export interface GenerateOptions {
+  parameters?: Record<string, unknown>;
+}
+
 export interface GenerationResult {
   tool: GeneratedTool | null;
   validationResult: ToolValidationResult | null;
@@ -52,7 +58,8 @@ export class ToolGenerator {
   async generate(
     gap: CapabilityGap,
     existingTools: Tool[],
-    testCases?: Array<{ input: unknown; expectedOutput?: unknown }>
+    testCases?: SandboxTestCase[],
+    options: GenerateOptions = {}
   ): Promise<GenerationResult> {
     let currentTool: GeneratedTool | null = null;
     let validationResult: ToolValidationResult | null = null;
@@ -65,15 +72,23 @@ export class ToolGenerator {
 
       try {
         if (currentTool === null) {
-          currentTool = await this.generateInitial(gap, existingTools);
+          currentTool = await this.generateInitial(gap, existingTools, options.parameters);
         } else if (validationResult) {
-          currentTool = await this.improve(currentTool, validationResult, iterations);
+          const improved = await this.improve(currentTool, validationResult, iterations);
+          if (!improved) {
+            continue;
+          }
+          currentTool = improved;
         }
 
         if (!currentTool) {
           continue;
         }
 
+        if (options.parameters) {
+          currentTool.parameters = options.parameters;
+        }
+
         validationResult = await this.validator.validate(currentTool, testCases);
 
         if (validationResult.isValid) {
@@ -120,7 +135,7 @@ export class ToolGenerator {
   async generateQuick(
     description: string,
     name: string,
-    _parameters: Record<string, unknown>
+    parameters: Record<string, unknown>
   ): Promise<GeneratedTool | null> {
     const gap: CapabilityGap = {
       id: `quick_${Date.now()}`,
@@ -132,13 +147,22 @@ export class ToolGenerator {
       reasoning: 'User-requested quick generation',
     };
 
-    const result = await this.generate(gap, []);
-    return result.tool;
+    const result = await this.generate(
+      gap,
+      [],
+      undefined,
+      Object.keys(parameters).length > 0
+        ? { parameters: { type: 'object', properties: parameters } }
+        : {}
+    );
+
+    return result.success ? result.tool : null;
   }
 
   private async generateInitial(
     gap: CapabilityGap,
-    existingTools: Tool[]
+    existingTools: Tool[],
+    parameters?: Record<string, unknown>
   ): Promise<GeneratedTool | null> {
     const toolSummaries = existingTools.map((t) => ({
       name: t.name,
@@ -149,17 +173,19 @@ export class ToolGenerator {
       maxLines: 100,
       securityLevel: 'strict',
       allowedModules: this.config.sandboxConfig?.allowedModules,
+      parameters,
     });
 
-    const response = await this.callLLM(
+    const content = await llmChat(
+      this.llm,
       [
         { role: 'system', content: TOOL_GENERATION_SYSTEM_PROMPT },
         { role: 'user', content: prompt },
       ],
-      0.4
+      { model: this.model, temperature: 0.4 }
     );
 
-    const tool = parseToolGenerationResponse(response.content);
+    const tool = parseToolGenerationResponse(content);
 
     if (tool) {
       tool.metadata = {
@@ -179,18 +205,21 @@ export class ToolGenerator {
   ): Promise<GeneratedTool | null> {
     const prompt = buildToolImprovementPrompt(tool, validationResult, iteration);
 
-    const response = await this.callLLM(
+    const content = await llmChat(
+      this.llm,
       [
         { role: 'system', content: TOOL_GENERATION_SYSTEM_PROMPT },
         { role: 'user', content: prompt },
       ],
-      0.3
+      { model: this.model, temperature: 0.3 }
     );
 
-    const improved = parseToolGenerationResponse(response.content);
+    const improved = parseToolGenerationResponse(content);
 
     if (improved) {
       improved.id = tool.id;
+      improved.name = tool.name;
+      improved.createdAt = tool.createdAt;
       improved.version = tool.version + 1;
       improved.metadata = {
         ...tool.metadata,
@@ -206,9 +235,18 @@ export class ToolGenerator {
   createExecutableTool(generated: GeneratedTool): Tool {
     const sandbox = this.sandbox;
     const tool = generated;
+    const paramSchema = buildZodSchema(generated.parameters);
+    const jsonSchema = toToolJsonSchema(generated.parameters);
 
     const execute = async (params: unknown): Promise<unknown> => {
-      const result = await sandbox.execute(tool, params);
+      const parsed = paramSchema.safeParse(params);
+      if (!parsed.success) {
+        throw new Error(
+          `Invalid parameters: ${parsed.error.issues.map((i) => `${i.path.join('.')}: ${i.message}`).join('; ')}`
+        );
+      }
+
+      const result = await sandbox.execute(tool, parsed.data);
       if (!result.success) {
         throw new Error(result.error ?? 'Tool execution failed');
       }
@@ -218,26 +256,82 @@ export class ToolGenerator {
     return {
       name: generated.name,
       description: generated.description,
-      parameters: z.record(z.string(), z.unknown()) as ZodType<unknown>,
+      parameters: paramSchema as ZodType<unknown>,
       execute,
       toJSON: () => ({
         name: generated.name,
         description: generated.description,
-        parameters: {
-          type: 'object' as const,
-          properties: generated.parameters as Record<string, unknown>,
-        },
+        parameters: jsonSchema,
       }),
     };
   }
+}
+
+function toToolJsonSchema(parameters: Record<string, unknown>): ToolSchema['parameters'] {
+  const properties =
+    parameters.type === 'object' &&
+    parameters.properties &&
+    typeof parameters.properties === 'object' &&
+    !Array.isArray(parameters.properties)
+      ? (parameters.properties as Record<string, unknown>)
+      : {};
+  const required = Array.isArray(parameters.required)
+    ? parameters.required.filter((key): key is string => typeof key === 'string')
+    : [];
+
+  return required.length > 0
+    ? { type: 'object', properties, required }
+    : { type: 'object', properties };
+}
 
-  private async callLLM(
-    messages: Array<{ role: 'system' | 'user' | 'assistant'; content: string }>,
-    temperature: number
+function buildZodSchema(parameters: Record<string, unknown>): ZodType<unknown> {
+  if (
+    parameters.type !== 'object' ||
+    !parameters.properties ||
+    typeof parameters.properties !== 'object'
   ) {
-    if (this.llm.complete) {
-      return this.llm.complete({ messages, temperature });
+    return z.record(z.string(), z.unknown());
+  }
+
+  const properties = parameters.properties as Record<string, Record<string, unknown>>;
+  const required = Array.isArray(parameters.required)
+    ? new Set(parameters.required as string[])
+    : new Set<string>();
+
+  const shape: Record<string, ZodType<unknown>> = {};
+
+  for (const [key, propSchema] of Object.entries(properties)) {
+    let fieldSchema: ZodType<unknown> = jsonTypeToZod(propSchema.type as string | undefined);
+
+    if (!required.has(key)) {
+      fieldSchema = fieldSchema.optional() as ZodType<unknown>;
     }
-    return this.llm.chat({ model: this.model, messages, temperature });
+
+    shape[key] = fieldSchema;
+  }
+
+  if (Object.keys(shape).length === 0) {
+    return z.record(z.string(), z.unknown());
+  }
+
+  return z.object(shape).passthrough() as ZodType<unknown>;
+}
+
+function jsonTypeToZod(type?: string): ZodType<unknown> {
+  switch (type) {
+    case 'string':
+      return z.string();
+    case 'number':
+      return z.number();
+    case 'integer':
+      return z.number().int();
+    case 'boolean':
+      return z.boolean();
+    case 'array':
+      return z.array(z.unknown());
+    case 'object':
+      return z.record(z.string(), z.unknown());
+    default:
+      return z.unknown();
   }
 }