smartagent/ts/smartagent.classes.dualagent.ts

import * as plugins from './plugins.js';
import * as interfaces from './smartagent.interfaces.js';
import { BaseToolWrapper } from './smartagent.tools.base.js';
import { DriverAgent } from './smartagent.classes.driveragent.js';
import { GuardianAgent } from './smartagent.classes.guardianagent.js';
import { FilesystemTool } from './smartagent.tools.filesystem.js';
import { HttpTool } from './smartagent.tools.http.js';
import { ShellTool } from './smartagent.tools.shell.js';
import { BrowserTool } from './smartagent.tools.browser.js';

/**
 * DualAgentOrchestrator - Coordinates Driver and Guardian agents
 * Manages the complete lifecycle of task execution with tool approval
 */
export class DualAgentOrchestrator {
  private options: interfaces.IDualAgentOptions;
  private smartai: plugins.smartai.SmartAi;
  private driverProvider: plugins.smartai.MultiModalModel;
  private guardianProvider: plugins.smartai.MultiModalModel;
  private driver: DriverAgent;
  private guardian: GuardianAgent;
  private tools: Map<string, BaseToolWrapper> = new Map();
  private isRunning = false;
  private conversationHistory: interfaces.IAgentMessage[] = [];

  constructor(options: interfaces.IDualAgentOptions) {
    this.options = {
      maxIterations: 20,
      maxConsecutiveRejections: 3,
      defaultProvider: 'openai',
      ...options,
    };

    // Create SmartAi instance
    this.smartai = new plugins.smartai.SmartAi(options);

    // Get providers
    this.driverProvider = this.getProviderByName(this.options.defaultProvider!);
    this.guardianProvider = this.options.guardianProvider
      ? this.getProviderByName(this.options.guardianProvider)
      : this.driverProvider;

    // Create agents
    this.driver = new DriverAgent(this.driverProvider, options.driverSystemMessage);
    this.guardian = new GuardianAgent(this.guardianProvider, options.guardianPolicyPrompt);
  }

  /**
   * Get provider by name
   */
  private getProviderByName(providerName: plugins.smartai.TProvider): plugins.smartai.MultiModalModel {
    switch (providerName) {
      case 'openai':
        return this.smartai.openaiProvider;
      case 'anthropic':
        return this.smartai.anthropicProvider;
      case 'perplexity':
        return this.smartai.perplexityProvider;
      case 'ollama':
        return this.smartai.ollamaProvider;
      case 'groq':
        return this.smartai.groqProvider;
      case 'xai':
        return this.smartai.xaiProvider;
      case 'exo':
        return this.smartai.exoProvider;
      default:
        return this.smartai.openaiProvider;
    }
  }

  /**
   * Register a custom tool
   */
  public registerTool(tool: BaseToolWrapper): void {
    this.tools.set(tool.name, tool);
    this.driver.registerTool(tool);
    this.guardian.registerTool(tool);
  }

  /**
   * Register all standard tools
   */
  public registerStandardTools(): void {
    const standardTools = [
      new FilesystemTool(),
      new HttpTool(),
      new ShellTool(),
      new BrowserTool(),
    ];

    for (const tool of standardTools) {
      this.registerTool(tool);
    }
  }

  /**
   * Initialize all tools (eager loading)
   */
  public async start(): Promise<void> {
    // Start smartai
    await this.smartai.start();

    // Initialize all tools
    const initPromises: Promise<void>[] = [];
    for (const tool of this.tools.values()) {
      initPromises.push(tool.initialize());
    }

    await Promise.all(initPromises);
    this.isRunning = true;
  }

  /**
   * Cleanup all tools
   */
  public async stop(): Promise<void> {
    const cleanupPromises: Promise<void>[] = [];

    for (const tool of this.tools.values()) {
      cleanupPromises.push(tool.cleanup());
    }

    await Promise.all(cleanupPromises);
    await this.smartai.stop();
    this.isRunning = false;
    this.driver.reset();
  }

  /**
   * Run a task through the dual-agent system
   */
  public async run(task: string): Promise<interfaces.IDualAgentRunResult> {
    if (!this.isRunning) {
      throw new Error('Orchestrator not started. Call start() first.');
    }

    this.conversationHistory = [];
    let iterations = 0;
    let consecutiveRejections = 0;
    let completed = false;
    let finalResult: string | null = null;

    // Add initial task to history
    this.conversationHistory.push({
      role: 'user',
      content: task,
    });

    // Start the driver with the task
    let driverResponse = await this.driver.startTask(task);
    this.conversationHistory.push(driverResponse);

    while (
      iterations < this.options.maxIterations! &&
      consecutiveRejections < this.options.maxConsecutiveRejections! &&
      !completed
    ) {
      iterations++;

      // Check if task is complete
      if (this.driver.isTaskComplete(driverResponse.content)) {
        completed = true;
        finalResult = this.driver.extractTaskResult(driverResponse.content) || driverResponse.content;
        break;
      }

      // Check if driver needs clarification
      if (this.driver.needsClarification(driverResponse.content)) {
        // Return with clarification needed status
        return {
          success: false,
          completed: false,
          result: driverResponse.content,
          iterations,
          history: this.conversationHistory,
          status: 'clarification_needed',
        };
      }

      // Parse tool call proposals
      const proposals = this.driver.parseToolCallProposals(driverResponse.content);

      if (proposals.length === 0) {
        // No tool calls, continue the conversation
        driverResponse = await this.driver.continueWithMessage(
          'Please either use a tool to make progress on the task, or indicate that the task is complete with <task_complete>summary</task_complete>.'
        );
        this.conversationHistory.push(driverResponse);
        continue;
      }

      // Process the first proposal (one at a time)
      const proposal = proposals[0];

      // Quick validation first
      const quickDecision = this.guardian.quickValidate(proposal);
      let decision: interfaces.IGuardianDecision;

      if (quickDecision) {
        decision = quickDecision;
      } else {
        // Full AI evaluation
        decision = await this.guardian.evaluate(proposal, task);
      }

      if (decision.decision === 'approve') {
        consecutiveRejections = 0;

        // Execute the tool
        const tool = this.tools.get(proposal.toolName);
        if (!tool) {
          const errorMessage = `Tool "${proposal.toolName}" not found.`;
          driverResponse = await this.driver.continueWithMessage(
            `TOOL ERROR: ${errorMessage}\n\nPlease try a different approach.`
          );
          this.conversationHistory.push(driverResponse);
          continue;
        }

        try {
          const result = await tool.execute(proposal.action, proposal.params);

          // Send result to driver
          const resultMessage = result.success
            ? `TOOL RESULT (${proposal.toolName}.${proposal.action}):\n${JSON.stringify(result.result, null, 2)}`
            : `TOOL ERROR (${proposal.toolName}.${proposal.action}):\n${result.error}`;

          this.conversationHistory.push({
            role: 'system',
            content: resultMessage,
            toolCall: proposal,
            toolResult: result,
          });

          driverResponse = await this.driver.continueWithMessage(resultMessage);
          this.conversationHistory.push(driverResponse);
        } catch (error) {
          const errorMessage = `Tool execution failed: ${error instanceof Error ? error.message : String(error)}`;
          driverResponse = await this.driver.continueWithMessage(
            `TOOL ERROR: ${errorMessage}\n\nPlease try a different approach.`
          );
          this.conversationHistory.push(driverResponse);
        }
      } else {
        // Rejected
        consecutiveRejections++;

        // Build rejection feedback
        let feedback = `TOOL CALL REJECTED by Guardian:\n`;
        feedback += `- Reason: ${decision.reason}\n`;

        if (decision.concerns && decision.concerns.length > 0) {
          feedback += `- Concerns:\n${decision.concerns.map(c => `  - ${c}`).join('\n')}\n`;
        }

        if (decision.suggestions) {
          feedback += `- Suggestions: ${decision.suggestions}\n`;
        }

        feedback += `\nPlease adapt your approach based on this feedback.`;

        this.conversationHistory.push({
          role: 'system',
          content: feedback,
          toolCall: proposal,
          guardianDecision: decision,
        });

        driverResponse = await this.driver.continueWithMessage(feedback);
        this.conversationHistory.push(driverResponse);
      }
    }

    // Determine final status
    let status: interfaces.TDualAgentRunStatus = 'completed';
    if (!completed) {
      if (iterations >= this.options.maxIterations!) {
        status = 'max_iterations_reached';
      } else if (consecutiveRejections >= this.options.maxConsecutiveRejections!) {
        status = 'max_rejections_reached';
      }
    }

    return {
      success: completed,
      completed,
      result: finalResult || driverResponse.content,
      iterations,
      history: this.conversationHistory,
      status,
    };
  }

  /**
   * Continue an existing task with user input
   */
  public async continueTask(userInput: string): Promise<interfaces.IDualAgentRunResult> {
    if (!this.isRunning) {
      throw new Error('Orchestrator not started. Call start() first.');
    }

    this.conversationHistory.push({
      role: 'user',
      content: userInput,
    });

    const driverResponse = await this.driver.continueWithMessage(userInput);
    this.conversationHistory.push(driverResponse);

    // Continue the run loop
    // For simplicity, we return the current state - full continuation would need refactoring
    return {
      success: false,
      completed: false,
      result: driverResponse.content,
      iterations: 1,
      history: this.conversationHistory,
      status: 'in_progress',
    };
  }

  /**
   * Get the conversation history
   */
  public getHistory(): interfaces.IAgentMessage[] {
    return [...this.conversationHistory];
  }

  /**
   * Update the guardian policy
   */
  public setGuardianPolicy(policyPrompt: string): void {
    this.guardian.setPolicy(policyPrompt);
  }

  /**
   * Check if orchestrator is running
   */
  public isActive(): boolean {
    return this.isRunning;
  }

  /**
   * Get registered tool names
   */
  public getToolNames(): string[] {
    return Array.from(this.tools.keys());
  }
}