Atomizer/mcp-server/atomizer-tools/src/tools/analysis.ts

/**
 * Analysis Tools
 *
 * Tools for querying and analyzing optimization results.
 */

import { execSync } from "child_process";
import Database from "better-sqlite3";
import { AtomizerTool } from "../index.js";
import { PYTHON_PATH, ATOMIZER_ROOT, getStudyDbPath } from "../utils/paths.js";

export const analysisTools: AtomizerTool[] = [
  {
    definition: {
      name: "get_trial_data",
      description:
        "Query trial data from a study's database. Supports various query types for different analysis needs.",
      inputSchema: {
        type: "object" as const,
        properties: {
          study_name: {
            type: "string",
            description: "Name of the study to query",
          },
          query: {
            type: "string",
            enum: ["all", "best", "pareto", "recent", "failed"],
            description:
              "Type of query: all (all trials), best (top N by objective), pareto (Pareto-optimal), recent (last N), failed (failed trials)",
          },
          limit: {
            type: "number",
            description: "Maximum number of results to return (default: 10)",
          },
          objective: {
            type: "string",
            description:
              "Objective name to sort by (for 'best' query). Default is first objective.",
          },
        },
        required: ["study_name", "query"],
      },
    },
    handler: async (args) => {
      const studyName = args.study_name as string;
      const query = args.query as string;
      const limit = (args.limit as number) || 10;
      const objective = args.objective as string | undefined;
      const dbPath = getStudyDbPath(studyName);

      try {
        const db = new Database(dbPath, { readonly: true });

        let sql: string;
        let params: unknown[] = [];

        switch (query) {
          case "best":
            sql = `
              SELECT t.number as trial_number,
                     tv.value as objective_value,
                     t.datetime_start as timestamp,
                     GROUP_CONCAT(tp.param_name || '=' || tp.param_value, '; ') as params
              FROM trials t
              JOIN trial_values tv ON t.trial_id = tv.trial_id
              LEFT JOIN trial_params tp ON t.trial_id = tp.trial_id
              WHERE t.state = 'COMPLETE'
              GROUP BY t.trial_id
              ORDER BY tv.value ASC
              LIMIT ?
            `;
            params = [limit];
            break;

          case "recent":
            sql = `
              SELECT t.number as trial_number,
                     t.state,
                     tv.value as objective_value,
                     t.datetime_start as timestamp
              FROM trials t
              LEFT JOIN trial_values tv ON t.trial_id = tv.trial_id
              ORDER BY t.number DESC
              LIMIT ?
            `;
            params = [limit];
            break;

          case "failed":
            sql = `
              SELECT t.number as trial_number,
                     t.state,
                     t.datetime_start as timestamp,
                     tua.value as error_message
              FROM trials t
              LEFT JOIN trial_user_attributes tua ON t.trial_id = tua.trial_id
                AND tua.key = 'error'
              WHERE t.state = 'FAIL'
              ORDER BY t.number DESC
              LIMIT ?
            `;
            params = [limit];
            break;

          case "pareto":
            // For multi-objective, get Pareto-optimal solutions
            // This is a simplified version - full Pareto requires Python
            sql = `
              SELECT t.number as trial_number,
                     tv.objective_id,
                     tv.value as objective_value,
                     t.datetime_start as timestamp
              FROM trials t
              JOIN trial_values tv ON t.trial_id = tv.trial_id
              WHERE t.state = 'COMPLETE'
              ORDER BY tv.value ASC
              LIMIT ?
            `;
            params = [limit * 2]; // Get more to filter
            break;

          default: // "all"
            sql = `
              SELECT t.number as trial_number,
                     t.state,
                     tv.value as objective_value,
                     t.datetime_start as timestamp
              FROM trials t
              LEFT JOIN trial_values tv ON t.trial_id = tv.trial_id
              ORDER BY t.number ASC
              LIMIT ?
            `;
            params = [limit];
        }

        const rows = db.prepare(sql).all(...params);
        db.close();

        return {
          content: [
            {
              type: "text",
              text: JSON.stringify(
                {
                  study: studyName,
                  query_type: query,
                  count: rows.length,
                  results: rows,
                },
                null,
                2
              ),
            },
          ],
        };
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        return {
          content: [
            {
              type: "text",
              text: `Error querying trials: ${message}`,
            },
          ],
          isError: true,
        };
      }
    },
  },

  {
    definition: {
      name: "analyze_convergence",
      description:
        "Analyze optimization convergence for a study. Returns trend metrics, improvement rate, and recommendations.",
      inputSchema: {
        type: "object" as const,
        properties: {
          study_name: {
            type: "string",
            description: "Name of the study to analyze",
          },
          window_size: {
            type: "number",
            description:
              "Rolling window size for convergence analysis (default: 10)",
          },
        },
        required: ["study_name"],
      },
    },
    handler: async (args) => {
      const studyName = args.study_name as string;
      const windowSize = (args.window_size as number) || 10;
      const dbPath = getStudyDbPath(studyName);

      try {
        const db = new Database(dbPath, { readonly: true });

        // Get all completed trials ordered by number
        const rows = db
          .prepare(
            `
            SELECT t.number, tv.value as objective_value
            FROM trials t
            JOIN trial_values tv ON t.trial_id = tv.trial_id
            WHERE t.state = 'COMPLETE'
            ORDER BY t.number ASC
          `
          )
          .all() as { number: number; objective_value: number }[];

        db.close();

        if (rows.length < 3) {
          return {
            content: [
              {
                type: "text",
                text: JSON.stringify(
                  {
                    study: studyName,
                    status: "insufficient_data",
                    message: `Only ${rows.length} completed trials. Need at least 3 for convergence analysis.`,
                  },
                  null,
                  2
                ),
              },
            ],
          };
        }

        // Compute metrics
        const values = rows.map((r) => r.objective_value);
        const n = values.length;

        // Best value so far at each point (cumulative minimum)
        const bestSoFar: number[] = [];
        let runningBest = Infinity;
        for (const v of values) {
          runningBest = Math.min(runningBest, v);
          bestSoFar.push(runningBest);
        }

        // Improvement rate (how often we find a new best)
        let improvements = 0;
        for (let i = 1; i < n; i++) {
          if (bestSoFar[i] < bestSoFar[i - 1]) {
            improvements++;
          }
        }
        const improvementRate = improvements / (n - 1);

        // Recent variance (last window_size trials)
        const recentValues = values.slice(-windowSize);
        const recentMean =
          recentValues.reduce((a, b) => a + b, 0) / recentValues.length;
        const recentVariance =
          recentValues.reduce((a, b) => a + Math.pow(b - recentMean, 2), 0) /
          recentValues.length;
        const recentStdDev = Math.sqrt(recentVariance);
        const coeffOfVariation = recentStdDev / Math.abs(recentMean);

        // Recent improvement rate
        const recentBestSoFar = bestSoFar.slice(-windowSize);
        let recentImprovements = 0;
        for (let i = 1; i < recentBestSoFar.length; i++) {
          if (recentBestSoFar[i] < recentBestSoFar[i - 1]) {
            recentImprovements++;
          }
        }

        // Convergence status
        let status: string;
        let recommendation: string;

        if (recentImprovements === 0 && coeffOfVariation < 0.05) {
          status = "converged";
          recommendation =
            "Optimization has converged. Consider stopping or switching to a local optimizer for refinement.";
        } else if (improvementRate > 0.2) {
          status = "exploring";
          recommendation =
            "Still finding improvements regularly. Continue optimization.";
        } else if (improvementRate > 0.05) {
          status = "improving";
          recommendation =
            "Making steady progress. Consider running more trials.";
        } else {
          status = "plateaued";
          recommendation =
            "Improvement rate is low. Consider changing sampler or expanding search space.";
        }

        const result = {
          study: studyName,
          total_trials: n,
          best_value: bestSoFar[n - 1],
          best_trial: rows.findIndex((r) => r.objective_value === bestSoFar[n - 1]) + 1,
          convergence: {
            status,
            improvement_rate: improvementRate,
            recent_improvements: recentImprovements,
            recent_coefficient_of_variation: coeffOfVariation,
          },
          recommendation,
          history: {
            first_value: values[0],
            last_value: values[n - 1],
            improvement_from_start:
              ((values[0] - bestSoFar[n - 1]) / Math.abs(values[0])) * 100,
          },
        };

        return {
          content: [
            {
              type: "text",
              text: JSON.stringify(result, null, 2),
            },
          ],
        };
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        return {
          content: [
            {
              type: "text",
              text: `Error analyzing convergence: ${message}`,
            },
          ],
          isError: true,
        };
      }
    },
  },

  {
    definition: {
      name: "compare_trials",
      description:
        "Compare multiple trials side by side, showing parameters and objectives for each.",
      inputSchema: {
        type: "object" as const,
        properties: {
          study_name: {
            type: "string",
            description: "Name of the study",
          },
          trial_numbers: {
            type: "array",
            items: { type: "number" },
            description: "List of trial numbers to compare",
          },
        },
        required: ["study_name", "trial_numbers"],
      },
    },
    handler: async (args) => {
      const studyName = args.study_name as string;
      const trialNumbers = args.trial_numbers as number[];
      const dbPath = getStudyDbPath(studyName);

      try {
        const db = new Database(dbPath, { readonly: true });

        const trials = trialNumbers.map((num) => {
          // Get trial info
          const trial = db
            .prepare(
              `
              SELECT trial_id, number, state, datetime_start
              FROM trials WHERE number = ?
            `
            )
            .get(num) as {
            trial_id: number;
            number: number;
            state: string;
            datetime_start: string;
          } | undefined;

          if (!trial) {
            return { number: num, error: "Trial not found" };
          }

          // Get parameters
          const params = db
            .prepare(
              `
              SELECT param_name, param_value FROM trial_params
              WHERE trial_id = ?
            `
            )
            .all(trial.trial_id) as { param_name: string; param_value: string }[];

          // Get objectives
          const objectives = db
            .prepare(
              `
              SELECT objective_id, value FROM trial_values
              WHERE trial_id = ?
            `
            )
            .all(trial.trial_id) as { objective_id: number; value: number }[];

          return {
            number: trial.number,
            state: trial.state,
            timestamp: trial.datetime_start,
            parameters: Object.fromEntries(
              params.map((p) => [p.param_name, parseFloat(p.param_value) || p.param_value])
            ),
            objectives: objectives.map((o) => o.value),
          };
        });

        db.close();

        return {
          content: [
            {
              type: "text",
              text: JSON.stringify(
                {
                  study: studyName,
                  comparison: trials,
                },
                null,
                2
              ),
            },
          ],
        };
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        return {
          content: [
            {
              type: "text",
              text: `Error comparing trials: ${message}`,
            },
          ],
          isError: true,
        };
      }
    },
  },

  {
    definition: {
      name: "get_best_design",
      description:
        "Get the best design found in an optimization study, including full parameter values and objectives.",
      inputSchema: {
        type: "object" as const,
        properties: {
          study_name: {
            type: "string",
            description: "Name of the study",
          },
          objective_index: {
            type: "number",
            description:
              "Index of objective to optimize for (0-based). For multi-objective studies.",
          },
        },
        required: ["study_name"],
      },
    },
    handler: async (args) => {
      const studyName = args.study_name as string;
      const objectiveIndex = (args.objective_index as number) || 0;
      const dbPath = getStudyDbPath(studyName);

      try {
        const db = new Database(dbPath, { readonly: true });

        // Find best trial
        const bestTrial = db
          .prepare(
            `
            SELECT t.trial_id, t.number, tv.value as objective_value
            FROM trials t
            JOIN trial_values tv ON t.trial_id = tv.trial_id
            WHERE t.state = 'COMPLETE' AND tv.objective_id = ?
            ORDER BY tv.value ASC
            LIMIT 1
          `
          )
          .get(objectiveIndex) as {
          trial_id: number;
          number: number;
          objective_value: number;
        } | undefined;

        if (!bestTrial) {
          db.close();
          return {
            content: [
              {
                type: "text",
                text: `No completed trials found for study "${studyName}"`,
              },
            ],
            isError: true,
          };
        }

        // Get all parameters
        const params = db
          .prepare(
            `
            SELECT param_name, param_value FROM trial_params
            WHERE trial_id = ?
          `
          )
          .all(bestTrial.trial_id) as { param_name: string; param_value: string }[];

        // Get all objectives
        const objectives = db
          .prepare(
            `
            SELECT objective_id, value FROM trial_values
            WHERE trial_id = ?
          `
          )
          .all(bestTrial.trial_id) as { objective_id: number; value: number }[];

        // Get total trials for context
        const totalTrials = (
          db.prepare("SELECT COUNT(*) as count FROM trials WHERE state = 'COMPLETE'").get() as {
            count: number;
          }
        ).count;

        db.close();

        const result = {
          study: studyName,
          best_trial: bestTrial.number,
          total_trials_evaluated: totalTrials,
          parameters: Object.fromEntries(
            params.map((p) => [p.param_name, parseFloat(p.param_value) || p.param_value])
          ),
          objectives: Object.fromEntries(
            objectives.map((o, i) => [`objective_${i}`, o.value])
          ),
          primary_objective_value: bestTrial.objective_value,
        };

        return {
          content: [
            {
              type: "text",
              text: JSON.stringify(result, null, 2),
            },
          ],
        };
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        return {
          content: [
            {
              type: "text",
              text: `Error getting best design: ${message}`,
            },
          ],
          isError: true,
        };
      }
    },
  },
];