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

/**
 * Reporting Tools
 *
 * Tools for generating reports and exporting data.
 */

import { execSync } from "child_process";
import { writeFile } from "fs/promises";
import { resolve } from "path";
import Database from "better-sqlite3";
import { AtomizerTool } from "../index.js";
import { PYTHON_PATH, ATOMIZER_ROOT, getStudyDir, getStudyDbPath } from "../utils/paths.js";

export const reportingTools: AtomizerTool[] = [
  {
    definition: {
      name: "generate_report",
      description:
        "Generate a markdown report for an optimization study, including configuration summary, results analysis, and recommendations.",
      inputSchema: {
        type: "object" as const,
        properties: {
          study_name: {
            type: "string",
            description: "Name of the study to report on",
          },
          report_type: {
            type: "string",
            enum: ["summary", "detailed", "executive"],
            description:
              "Type of report: summary (quick overview), detailed (full analysis), executive (high-level for stakeholders)",
          },
          include_plots: {
            type: "boolean",
            description: "Whether to generate convergence plots (requires matplotlib)",
          },
        },
        required: ["study_name"],
      },
    },
    handler: async (args) => {
      const studyName = args.study_name as string;
      const reportType = (args.report_type as string) || "summary";
      const includePlots = args.include_plots as boolean;
      const studyDir = getStudyDir(studyName);

      // Use Python report generator
      const script = `
import sys
sys.path.insert(0, r"${ATOMIZER_ROOT}")
from optimization_engine.reporting.markdown_report import generate_study_report

try:
    report = generate_study_report(
        study_name="${studyName}",
        report_type="${reportType}",
        include_plots=${includePlots ? "True" : "False"}
    )
    print(report)
except Exception as e:
    print(f"ERROR: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)
`;

      try {
        const output = execSync(`"${PYTHON_PATH}" -c "${script}"`, {
          encoding: "utf-8",
          cwd: studyDir,
          timeout: 120000, // 2 minute timeout for plots
          maxBuffer: 10 * 1024 * 1024, // 10MB buffer for large reports
        });

        // Save report to file
        const reportPath = resolve(studyDir, "3_results", `STUDY_REPORT_${reportType}.md`);
        try {
          await writeFile(reportPath, output);
        } catch {
          // Ignore save errors - still return the report
        }

        return {
          content: [
            {
              type: "text",
              text: output,
            },
          ],
        };
      } catch (error) {
        // If Python report generator fails, generate a basic report from DB
        const message = error instanceof Error ? error.message : String(error);

        // Fall back to basic report
        try {
          const basicReport = await generateBasicReport(studyName, reportType);
          return {
            content: [
              {
                type: "text",
                text: basicReport,
              },
            ],
          };
        } catch (fallbackError) {
          return {
            content: [
              {
                type: "text",
                text: `Error generating report: ${message}`,
              },
            ],
            isError: true,
          };
        }
      }
    },
  },

  {
    definition: {
      name: "export_data",
      description:
        "Export optimization results to CSV or JSON format for external analysis.",
      inputSchema: {
        type: "object" as const,
        properties: {
          study_name: {
            type: "string",
            description: "Name of the study to export",
          },
          format: {
            type: "string",
            enum: ["csv", "json"],
            description: "Export format",
          },
          include_failed: {
            type: "boolean",
            description: "Whether to include failed trials (default: false)",
          },
        },
        required: ["study_name", "format"],
      },
    },
    handler: async (args) => {
      const studyName = args.study_name as string;
      const format = args.format as string;
      const includeFailed = args.include_failed as boolean;
      const studyDir = getStudyDir(studyName);
      const dbPath = getStudyDbPath(studyName);

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

        // Get all trial data
        const stateFilter = includeFailed
          ? ""
          : "WHERE t.state = 'COMPLETE'";

        const trials = db
          .prepare(
            `
            SELECT t.number, t.state, t.datetime_start, t.datetime_complete
            FROM trials t
            ${stateFilter}
            ORDER BY t.number ASC
          `
          )
          .all() as {
          number: number;
          state: string;
          datetime_start: string;
          datetime_complete: string | null;
        }[];

        // Get params and values for each trial
        const data = trials.map((trial) => {
          const trialId = db
            .prepare("SELECT trial_id FROM trials WHERE number = ?")
            .get(trial.number) as { trial_id: number };

          const params = db
            .prepare("SELECT param_name, param_value FROM trial_params WHERE trial_id = ?")
            .all(trialId.trial_id) as { param_name: string; param_value: string }[];

          const values = db
            .prepare("SELECT objective_id, value FROM trial_values WHERE trial_id = ?")
            .all(trialId.trial_id) as { objective_id: number; value: number }[];

          return {
            trial_number: trial.number,
            state: trial.state,
            datetime_start: trial.datetime_start,
            datetime_complete: trial.datetime_complete,
            ...Object.fromEntries(
              params.map((p) => [p.param_name, parseFloat(p.param_value) || p.param_value])
            ),
            ...Object.fromEntries(values.map((v) => [`objective_${v.objective_id}`, v.value])),
          };
        });

        db.close();

        let output: string;
        let filename: string;

        if (format === "csv") {
          // Generate CSV
          if (data.length === 0) {
            output = "No data to export";
          } else {
            const headers = Object.keys(data[0]);
            const rows = data.map((row) => headers.map((h) => row[h as keyof typeof row] ?? "").join(","));
            output = [headers.join(","), ...rows].join("\n");
          }
          filename = `${studyName}_export.csv`;
        } else {
          // Generate JSON
          output = JSON.stringify(data, null, 2);
          filename = `${studyName}_export.json`;
        }

        // Save to file
        const exportPath = resolve(studyDir, "3_results", filename);
        await writeFile(exportPath, output);

        return {
          content: [
            {
              type: "text",
              text: `Exported ${data.length} trials to ${exportPath}\n\nPreview (first 5 rows):\n${format === "csv" ? output.split("\n").slice(0, 6).join("\n") : JSON.stringify(data.slice(0, 5), null, 2)}`,
            },
          ],
        };
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        return {
          content: [
            {
              type: "text",
              text: `Error exporting data: ${message}`,
            },
          ],
          isError: true,
        };
      }
    },
  },
];

// Helper function for basic report generation
async function generateBasicReport(
  studyName: string,
  reportType: string
): Promise<string> {
  const dbPath = getStudyDbPath(studyName);
  const db = new Database(dbPath, { readonly: true });

  // Get trial counts
  const total = (db.prepare("SELECT COUNT(*) as c FROM trials").get() as { c: number }).c;
  const completed = (
    db.prepare("SELECT COUNT(*) as c FROM trials WHERE state = 'COMPLETE'").get() as { c: number }
  ).c;
  const failed = (
    db.prepare("SELECT COUNT(*) as c FROM trials WHERE state = 'FAIL'").get() as { c: number }
  ).c;

  // Get best result
  const best = db
    .prepare(
      `
      SELECT t.number, tv.value
      FROM trials t
      JOIN trial_values tv ON t.trial_id = tv.trial_id
      WHERE t.state = 'COMPLETE'
      ORDER BY tv.value ASC
      LIMIT 1
    `
    )
    .get() as { number: number; value: number } | undefined;

  db.close();

  let report = `# Optimization Report: ${studyName}\n\n`;
  report += `**Generated:** ${new Date().toISOString()}\n\n`;
  report += `## Summary\n\n`;
  report += `- Total trials: ${total}\n`;
  report += `- Completed: ${completed}\n`;
  report += `- Failed: ${failed}\n`;

  if (best) {
    report += `\n## Best Result\n\n`;
    report += `- Trial #${best.number}\n`;
    report += `- Objective value: ${best.value.toFixed(6)}\n`;
  }

  return report;
}