obsidian_ollama/src/workflow-engine/workflow-engine.ts

// src/workflow-engine/workflow-engine.ts

import { Vault, App } from 'obsidian';
import {
  WorkflowStep,
  WorkflowStepType,
  WorkflowStepResult,
  WorkflowDefinition,
  WorkflowExecutionContext,
  WorkflowExecutionResult,
  LlmStepConfig,
  VaultSearchStepConfig,
  ToolStepConfig,
  FormatStepConfig,
  OllamaMessage,
  OllamaTool,
} from '../types';
import { VaultIndexer } from '../vault-indexer';
import { ToolExecutor } from '../tool-executor';
import { OllamaClient } from '../ollama-client';
import { ConversationStateManager } from '../conversation-state';
import { Logger } from '../utils';
import { safeParseJson } from '../utils';

/**
 * Pattern to match {{variable}} or {{variable.property}} syntax in strings.
 */
const VARIABLE_PATTERN = /\{\{([\w.]+)\}\}/g;

/**
 * WorkflowEngine orchestrates multi-step/chained reasoning workflows.
 *
 * Each workflow is a sequence of steps that can:
 * - Call the LLM with prompts
 * - Search the vault for relevant notes
 * - Execute tools (file operations, vault queries)
 * - Format/transform data
 *
 * Steps can reference outputs from previous steps via {{stepId.output}} syntax.
 */
export class WorkflowEngine {
  private vaultIndexer: VaultIndexer;
  private toolExecutor: ToolExecutor;
  private ollamaClient: OllamaClient;
  private conversationStateManager: ConversationStateManager;
  private maxSteps: number;
  private maxWorkflowDuration: number;

  constructor(
    vault: Vault,
    app: App,
    ollamaUrl: string,
    model: string,
    options?: {
      maxSteps?: number;
      maxWorkflowDuration?: number;
      cacheConfig?: import('../types').CacheConfig;
    }
  ) {
    this.vaultIndexer = new VaultIndexer(vault);
    this.toolExecutor = new ToolExecutor(vault, app);
    this.ollamaClient = new OllamaClient(ollamaUrl, model, undefined, options?.cacheConfig);
    this.conversationStateManager = new ConversationStateManager();
    this.maxSteps = options?.maxSteps ?? 20;
    this.maxWorkflowDuration = options?.maxWorkflowDuration ?? 300_000; // 5 minutes
  }

  /**
   * Execute a workflow definition.
   * @param definition The workflow to execute
   * @param initialVariables Optional initial variables to seed the context
   * @returns The execution result with all step outputs
   */
  async executeWorkflow(
    definition: WorkflowDefinition,
    initialVariables?: Record<string, unknown>
  ): Promise<WorkflowExecutionResult> {
    Logger.info(`Starting workflow: ${definition.name} (${definition.id})`, 'workflow-engine');

    const startTime = Date.now();
    const context = this.createExecutionContext(initialVariables);

    // Validate the workflow before execution
    const validationError = this.validateWorkflow(definition);
    if (validationError) {
      Logger.error(`Workflow validation failed: ${validationError}`, 'workflow-engine');
      return {
        workflowId: definition.id,
        workflowName: definition.name,
        success: false,
        stepResults: [],
        finalOutput: null,
        error: validationError,
      };
    }

    try {
      // Execute steps in topological order
      const orderedSteps = this.topologicalSort(definition.steps);

      let stepCount = 0;
      for (const step of orderedSteps) {
        // Check timeout
        const elapsed = Date.now() - startTime;
        if (elapsed > this.maxWorkflowDuration) {
          throw new Error(`Workflow exceeded maximum duration of ${this.maxWorkflowDuration}ms`);
        }

        // Check max steps
        stepCount++;
        if (stepCount > this.maxSteps) {
          throw new Error(`Workflow exceeded maximum step count of ${this.maxSteps}`);
        }

        // Wait for dependency if exists
        if (step.dependsOn) {
          const depResult = context.stepResults.find((r) => r.stepId === step.dependsOn);
          if (!depResult) {
            throw new Error(`Step ${step.id}: dependency '${step.dependsOn}' not found`);
          }
          if (!depResult.success) {
            Logger.warn(
              `Step ${step.id}: dependency '${step.dependsOn}' failed, skipping`,
              'workflow-engine'
            );
            // Record a failure for this step due to failed dependency
            context.stepResults.push({
              stepId: step.id,
              stepName: step.name,
              success: false,
              data: null,
              error: `Dependency '${step.dependsOn}' failed: ${depResult.error}`,
              timestamp: Date.now(),
            });
            continue;
          }
        }

        // Execute the step
        Logger.debug(`Executing step: ${step.name} (${step.id})`, 'workflow-engine');
        const result = await this.executeStep(step, context);
        context.stepResults.push(result);

        // Store the result data in variables for interpolation
        context.variables.set(step.id, result.data);

        // Update conversation history for LLM steps
        if (step.type === 'llm' && result.success) {
          this.conversationStateManager.updateShortTermContext({
            role: 'assistant',
            content: String(result.data),
          });
        }
      }

      // Determine final output (last successful step's data)
      const lastSuccessfulResult = [...context.stepResults].reverse().find((r) => r.success);
      const finalOutput = lastSuccessfulResult?.data ?? null;

      const success = context.stepResults.every((r) => r.success);

      Logger.info(
        `Workflow completed: ${definition.name} (${success ? 'success' : 'partial'})`,
        'workflow-engine'
      );

      return {
        workflowId: definition.id,
        workflowName: definition.name,
        success,
        stepResults: context.stepResults,
        finalOutput,
      };
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : String(error);
      Logger.error(`Workflow failed: ${errorMessage}`, 'workflow-engine');

      return {
        workflowId: definition.id,
        workflowName: definition.name,
        success: false,
        stepResults: context.stepResults,
        finalOutput: null,
        error: errorMessage,
      };
    }
  }

  /**
   * Execute a workflow from a natural language description.
   * The LLM will generate the workflow steps, then we execute them.
   *
   * @param userQuery The user's natural language request
   * @param availableTools Optional list of available tools to inform the LLM
   * @returns The execution result
   */
  async executeWorkflowFromQuery(
    userQuery: string,
    availableTools?: OllamaTool[]
  ): Promise<WorkflowExecutionResult> {
    Logger.info(`Generating workflow from query: ${userQuery}`, 'workflow-engine');

    // First, ask the LLM to generate a workflow plan
    const workflow = await this.generateWorkflowFromQuery(userQuery, availableTools);
    if (!workflow) {
      return {
        workflowId: 'auto-generated',
        workflowName: 'Auto-generated workflow',
        success: false,
        stepResults: [],
        finalOutput: null,
        error: 'Failed to generate workflow from query',
      };
    }

    // Add a system variable with the original query
    const initialVariables: Record<string, unknown> = {
      original_query: userQuery,
    };

    return this.executeWorkflow(workflow, initialVariables);
  }

  /**
   * Generate a workflow definition from a natural language query using the LLM.
   */
  private async generateWorkflowFromQuery(
    query: string,
    availableTools?: OllamaTool[]
  ): Promise<WorkflowDefinition | null> {
    const toolDescriptions =
      availableTools?.map((t) => `- ${t.function.name}: ${t.function.description}`).join('\n') ??
      '';

    const systemPrompt = `You are a workflow planner. Given a user query, break it down into a sequence of workflow steps.

Available step types:
- vault_search: Search the vault for notes. Config: { type: 'vault_search', query: string, limit?: number, tagFilter?: string }
- llm: Call an LLM. Config: { type: 'llm', systemPrompt?: string, userPrompt: string }
- tool: Execute a tool. Config: { type: 'tool', toolName: string, args: object }
- format: Format output. Config: { type: 'format', template: string, outputFormat?: 'markdown' | 'text' | 'json' }

Available tools:
${toolDescriptions}

Variable interpolation syntax:
- Use {{stepId.output}} to reference a previous step's output
- Use {{stepId.output.property}} to reference a property of a step's output
- Use {{original_query}} to reference the original user query

Return a JSON object with this structure:
{
  "id": "workflow-uuid",
  "name": "workflow name",
  "description": "description",
  "steps": [
    {
      "id": "step_1",
      "type": "vault_search" | "llm" | "tool" | "format",
      "name": "step name",
      "description": "optional description",
      "config": { /* step-specific config */ },
      "dependsOn": "optional_step_id"
    }
  ]
}

Rules:
1. Number steps sequentially (step_1, step_2, etc.)
2. Use dependsOn to specify ordering when needed
3. Use variable interpolation to pass data between steps
4. Keep the workflow minimal but effective
5. End with a format step if the user wants structured output`;

    const userPrompt = `User query: ${query}`;

    try {
      const messages: OllamaMessage[] = [
        { role: 'system', content: systemPrompt },
        { role: 'user', content: userPrompt },
      ];

      const response = await this.ollamaClient.chat(messages, []);
      const content = response.content?.trim();

      if (!content) {
        Logger.error('Empty response from LLM when generating workflow', 'workflow-engine');
        return null;
      }

      // Extract JSON from the response (handle markdown code blocks)
      const jsonMatch =
        content.match(/```(?:json)?\s*([\s\S]*?)```/) ?? content.match(/\{[\s\S]*\}/);
      const jsonString = jsonMatch ? jsonMatch[1] : content;

      const parsed = safeParseJson(jsonString);
      if (!parsed || typeof parsed !== 'object') {
        Logger.error('Invalid workflow JSON from LLM', 'workflow-engine');
        return null;
      }

      return parsed as WorkflowDefinition;
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : String(error);
      Logger.error(`Failed to generate workflow: ${errorMessage}`, 'workflow-engine');
      return null;
    }
  }

  /**
   * Create a fresh execution context.
   */
  private createExecutionContext(
    initialVariables?: Record<string, unknown>
  ): WorkflowExecutionContext {
    const variables = new Map<string, unknown>();
    if (initialVariables) {
      for (const [key, value] of Object.entries(initialVariables)) {
        variables.set(key, value);
      }
    }

    return {
      variables,
      stepResults: [],
      conversationHistory: [],
    };
  }

  /**
   * Execute a single workflow step.
   */
  private async executeStep(
    step: WorkflowStep,
    context: WorkflowExecutionContext
  ): Promise<WorkflowStepResult> {
    const timestamp = Date.now();
    try {
      // Interpolate variables in the step config
      const interpolatedConfig = this.interpolateVariables(step.config, context.variables);

      let data: unknown;
      switch (step.type) {
        case 'llm':
          data = await this.executeLlmStep(interpolatedConfig as LlmStepConfig, context);
          break;
        case 'vault_search':
          data = await this.executeVaultSearchStep(interpolatedConfig as VaultSearchStepConfig);
          break;
        case 'tool':
          data = await this.executeToolStep(interpolatedConfig as ToolStepConfig);
          break;
        case 'format':
          data = this.executeFormatStep(interpolatedConfig as FormatStepConfig, context);
          break;
        default:
          throw new Error(`Unknown step type: ${String(step.type)}`);
      }

      return {
        stepId: step.id,
        stepName: step.name,
        success: true,
        data,
        timestamp,
      };
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : String(error);
      Logger.error(`Step ${step.id} failed: ${errorMessage}`, 'workflow-engine');

      return {
        stepId: step.id,
        stepName: step.name,
        success: false,
        data: null,
        error: errorMessage,
        timestamp,
      };
    }
  }

  /**
   * Execute an LLM step.
   */
  private async executeLlmStep(
    config: LlmStepConfig,
    context: WorkflowExecutionContext
  ): Promise<string> {
    const messages: OllamaMessage[] = [];

    // Add system prompt if provided
    if (config.systemPrompt) {
      messages.push({
        role: 'system',
        content: config.systemPrompt,
      });
    }

    // Add conversation history
    messages.push(...context.conversationHistory);

    // Add user prompt
    messages.push({
      role: 'user',
      content: config.userPrompt,
    });

    // Determine if we should include tools
    const tools = config.includeToolCalls ? [] : [];

    const response = await this.ollamaClient.chat(messages, tools);
    return response.content ?? '';
  }

  /**
   * Execute a vault search step.
   */
  private async executeVaultSearchStep(config: VaultSearchStepConfig): Promise<unknown> {
    const limit = config.limit ?? 5;
    const entries = await this.vaultIndexer.searchVault(config.query, limit);

    // Apply tag filter if specified
    const filtered = config.tagFilter
      ? entries.filter((entry) => {
          const tags = entry.tags ?? '';
          return tags.toLowerCase().includes(config.tagFilter!.toLowerCase());
        })
      : entries;

    return filtered.map((entry) => ({
      path: entry.path,
      title: entry.title,
      content: entry.content,
      score: entry.score,
      tags: entry.tags,
    }));
  }

  /**
   * Execute a tool step.
   */
  private async executeToolStep(config: ToolStepConfig): Promise<unknown> {
    const args: Record<string, unknown> = config.args ?? {};
    const result = await this.toolExecutor.executeTool(config.toolName, args);

    return {
      success: result.success,
      message: result.message,
      data: result.data,
    };
  }

  /**
   * Execute a format step (template rendering).
   */
  private executeFormatStep(config: FormatStepConfig, context: WorkflowExecutionContext): string {
    let output = config.template;
    output = String(this.interpolateVariables(output, context.variables));

    // Apply output formatting
    if (config.outputFormat === 'json') {
      try {
        const parsed = safeParseJson(output);
        return JSON.stringify(parsed, null, 2);
      } catch {
        // Return as-is if not valid JSON
        return output;
      }
    }

    return output;
  }

  /**
   * Interpolate {{variables}} in a string or object.
   * Supports:
   * - {{variableName}} -> value from context
   * - {{stepId.output}} -> data from a step result
   * - {{stepId.output.property}} -> nested property access
   */
  private interpolateVariables(input: unknown, variables: Map<string, unknown>): unknown {
    if (typeof input === 'string') {
      return this.interpolateString(input, variables);
    }

    if (Array.isArray(input)) {
      return input.map((item) => this.interpolateVariables(item, variables));
    }

    if (input !== null && typeof input === 'object') {
      const result: Record<string, unknown> = {};
      for (const [key, value] of Object.entries(input as Record<string, unknown>)) {
        result[key] = this.interpolateVariables(value, variables);
      }
      return result;
    }

    return input;
  }

  /**
   * Interpolate variables in a string.
   */
  private interpolateString(input: string, variables: Map<string, unknown>): string {
    return input.replace(VARIABLE_PATTERN, (_match: string, variablePath: string) => {
      const value = this.resolveVariable(variablePath, variables);
      if (value === undefined) {
        // Keep the original placeholder if variable not found
        Logger.warn(`Variable '${variablePath}' not found during interpolation`, 'workflow-engine');
        return _match;
      }

      // Handle different value types
      if (typeof value === 'object' && value !== null) {
        return JSON.stringify(value);
      }

      if (
        typeof value === 'string' ||
        typeof value === 'number' ||
        typeof value === 'boolean' ||
        value === null ||
        value === undefined
      ) {
        return String(value);
      }

      // Fallback for symbols, functions, etc.
      // eslint-disable-next-line @typescript-eslint/no-base-to-string
      return String(value);
    });
  }

  /**
   * Resolve a variable path like "step_1.output" or "step_1.output.property".
   */
  private resolveVariable(path: string, variables: Map<string, unknown>): unknown {
    const parts = path.split('.');

    // Check if this is a step output reference (stepId.output or stepId.output.property)
    if (parts.length >= 2 && parts[1] === 'output') {
      const stepId = parts[0];
      const stepData = variables.get(stepId);

      if (parts.length === 2) {
        return stepData;
      }

      // Navigate into nested properties
      let current = stepData;
      for (let i = 2; i < parts.length; i++) {
        if (current === null || current === undefined || typeof current !== 'object') {
          return undefined;
        }
        current = (current as Record<string, unknown>)[parts[i]];
      }
      return current;
    }

    // Direct variable lookup
    return variables.get(path);
  }

  /**
   * Perform topological sort on steps to determine execution order.
   * This respects the dependsOn field and ensures steps run in the correct order.
   */
  private topologicalSort(steps: WorkflowStep[]): WorkflowStep[] {
    const stepMap = new Map<string, WorkflowStep>();
    for (const step of steps) {
      stepMap.set(step.id, step);
    }

    const result: WorkflowStep[] = [];
    const visited = new Set<string>();
    const visiting = new Set<string>();

    const visit = (stepId: string): void => {
      if (visited.has(stepId)) return;
      if (visiting.has(stepId)) {
        throw new Error(`Circular dependency detected involving step '${stepId}'`);
      }

      const step = stepMap.get(stepId);
      if (!step) {
        throw new Error(`Step '${stepId}' not found`);
      }

      visiting.add(stepId);

      // Visit dependencies first
      if (step.dependsOn) {
        visit(step.dependsOn);
      }

      visiting.delete(stepId);
      visited.add(stepId);
      result.push(step);
    };

    // Visit all steps
    for (const step of steps) {
      visit(step.id);
    }

    return result;
  }

  /**
   * Validate a workflow definition before execution.
   * Returns null if valid, or an error message string if invalid.
   */
  private validateWorkflow(definition: WorkflowDefinition): string | null {
    if (!definition.id) {
      return 'Workflow must have an id';
    }

    if (!definition.name) {
      return 'Workflow must have a name';
    }

    if (!definition.steps || !Array.isArray(definition.steps)) {
      return 'Workflow must have a steps array';
    }

    if (definition.steps.length === 0) {
      return 'Workflow must have at least one step';
    }

    const stepIds = new Set<string>();
    for (const step of definition.steps) {
      if (!step.id) {
        return 'Each step must have an id';
      }

      if (!step.type) {
        return `Step '${step.id}' must have a type`;
      }

      if (!step.name) {
        return `Step '${step.id}' must have a name`;
      }

      if (!step.config) {
        return `Step '${step.id}' must have a config`;
      }

      // Check for duplicate step ids
      if (stepIds.has(step.id)) {
        return `Duplicate step id: '${step.id}'`;
      }
      stepIds.add(step.id);

      // Validate step type
      const validTypes: WorkflowStepType[] = ['llm', 'vault_search', 'tool', 'format'];
      if (!validTypes.includes(step.type)) {
        return `Step '${step.id}' has invalid type: '${step.type}'`;
      }

      // Validate dependsOn references a valid step
      if (step.dependsOn && !definition.steps.some((s) => s.id === step.dependsOn)) {
        return `Step '${step.id}' depends on unknown step: '${step.dependsOn}'`;
      }

      // Validate step-specific config
      const configError = this.validateStepConfig(step);
      if (configError) {
        return configError;
      }
    }

    return null;
  }

  /**
   * Validate a step's configuration based on its type.
   */
  private validateStepConfig(step: WorkflowStep): string | null {
    switch (step.type) {
      case 'llm': {
        const config = step.config as LlmStepConfig;
        if (!config.userPrompt) {
          return `LLM step '${step.id}' requires a userPrompt`;
        }
        break;
      }

      case 'vault_search': {
        const config = step.config as VaultSearchStepConfig;
        if (!config.query) {
          return `Vault search step '${step.id}' requires a query`;
        }
        break;
      }

      case 'tool': {
        const config = step.config as ToolStepConfig;
        if (!config.toolName) {
          return `Tool step '${step.id}' requires a toolName`;
        }
        break;
      }

      case 'format': {
        const config = step.config as FormatStepConfig;
        if (!config.template) {
          return `Format step '${step.id}' requires a template`;
        }
        break;
      }
    }

    return null;
  }

  /**
   * Get all built-in workflow definitions (presets).
   */
  static getBuiltInWorkflows(): WorkflowDefinition[] {
    return [
      WorkflowEngine.createMeetingSummaryWorkflow(),
      WorkflowEngine.createNoteAnalyzerWorkflow(),
    ];
  }

  /**
   * Create a workflow that summarizes meeting notes from the last week.
   */
  private static createMeetingSummaryWorkflow(): WorkflowDefinition {
    return {
      id: 'meeting-summary',
      name: 'Meeting Notes Summary',
      description:
        'Analyzes meeting notes from the last week, extracts decisions, and creates a summary.',
      steps: [
        {
          id: 'step_1',
          type: 'vault_search',
          name: 'Find Meeting Notes',
          description: 'Search for notes tagged with #meeting',
          config: {
            type: 'vault_search',
            query: 'meeting',
            limit: 10,
            tagFilter: 'meeting',
          },
        },
        {
          id: 'step_2',
          type: 'llm',
          name: 'Extract Decisions',
          description: 'Extract key decisions and assigned owners from meeting notes',
          config: {
            type: 'llm',
            systemPrompt:
              'You are a meeting analyst. Extract key decisions, action items, and assigned owners from meeting notes.',
            userPrompt: `Here are meeting notes from recent meetings:\n\n{{step_1.output}}\n\nPlease extract:\n1. Key decisions made\n2. Action items with assigned owners\n3. Deadlines if mentioned\n\nFormat as a structured list.`,
          },
          dependsOn: 'step_1',
        },
        {
          id: 'step_3',
          type: 'format',
          name: 'Format Summary',
          description: 'Format the results into a markdown table',
          config: {
            type: 'format',
            template:
              '# Meeting Summary\n\n## Decisions and Action Items\n\n{{step_2.output}}\n\n---\n*Generated by Workflow Engine*',
            outputFormat: 'markdown',
          },

          dependsOn: 'step_2',
        },
      ],
    };
  }

  /**
   * Create a workflow that analyzes notes and generates insights.
   */
  private static createNoteAnalyzerWorkflow(): WorkflowDefinition {
    return {
      id: 'note-analyzer',
      name: 'Note Analyzer',
      description: 'Analyzes notes and generates insights, summaries, and suggestions.',
      steps: [
        {
          id: 'step_1',
          type: 'vault_search',
          name: 'Search Notes',
          description: 'Search for relevant notes based on query',
          config: {
            type: 'vault_search',
            query: '{{original_query}}',
            limit: 5,
          },
        },
        {
          id: 'step_2',
          type: 'llm',
          name: 'Analyze Content',
          description: 'Analyze the found notes for insights',
          config: {
            type: 'llm',
            systemPrompt:
              'You are an analytical assistant. Analyze the provided notes and identify key themes, insights, and connections.',
            userPrompt: `Original query: {{original_query}}\n\nFound notes:\n{{step_1.output}}\n\nPlease provide:\n1. Key themes identified\n2. Important insights\n3. Potential connections between notes\n4. Suggestions for further exploration`,
          },
          dependsOn: 'step_1',
        },
        {
          id: 'step_3',
          type: 'format',
          name: 'Format Results',
          description: 'Format the analysis into a readable report',
          config: {
            type: 'format',
            template:
              '# Analysis Report\n\n## Query: {{original_query}}\n\n## Findings\n\n{{step_2.output}}\n\n---\n*Generated by Workflow Engine*',
            outputFormat: 'markdown',
          },

          dependsOn: 'step_2',
        },
      ],
    };
  }
}