Build a CSV data analysis app with the code interpreter

const generateCodePrompt = ({
  csvFileUrl,
  csvHeaders,
  csvRows,
}: {
  csvFileUrl?: string;
  csvHeaders?: string[];
  csvRows?: { [key: string]: string }[];
}) => {
  // Prepare sample rows as a markdown table if available
  let sampleRowsSection = "";
  if (csvRows && csvRows.length > 0 && csvHeaders && csvHeaders.length > 0) {
    const sampleRows = csvRows.slice(0, 3);
    const headerRow = `| ${csvHeaders.join(" | ")} |`;
    const separatorRow = `|${csvHeaders.map(() => "---").join("|")}|`;
    const dataRows = sampleRows
      .map((row) => `| ${csvHeaders.map((h) => row[h] ?? "").join(" | ")} |`)
      .join("\n");
    sampleRowsSection = `\n\nHere are a few sample rows from the dataset:\n\n${headerRow}\n${separatorRow}\n${dataRows}`;
  }

  return `
You are an expert data scientist assistant that writes python code to answer questions about a dataset.

You are given a dataset and a question.

The dataset is available at the following S3 URL: ${
    csvFileUrl || "[NO FILE URL PROVIDED]"
  }
The dataset has the following columns: ${
    csvHeaders?.join(", ") || "[NO HEADERS PROVIDED]"
  }
${sampleRowsSection}

You must always write python code that:
- Downloads the CSV from the provided S3 URL (using requests or pandas.read_csv).
- Uses the provided columns for analysis.
- Never outputs more than one graph per code response.
- Limits graph complexity to maintain readability.
- Never generates HTML output. Only use Python print statements or graphs/plots for output.

Always return the python code in a single unique code block.

Python sessions come pre-installed with pandas, matplotlib, seaborn, numpy, and other essential data science libraries.
`;
};

import { streamText } from "ai";
import { createTogetherAI } from "@ai-sdk/togetherai";

const togetherai = createTogetherAI({
  apiKey: process.env.TOGETHER_API_KEY ?? "",
});

export const generateAnalysisCode = (
  userQuestion: string,
  csvFileUrl: string,
  csvHeaders: string[],
  csvRows: { [key: string]: string }[],
  model: string = "moonshotai/Kimi-K2.7-Code",
) => {
  const result = streamText({
    model: togetherai(model),
    system: generateCodePrompt({ csvFileUrl, csvHeaders, csvRows }),
    messages: [{ role: "user", content: userQuestion }],
  });

  // Stream the response to the client as UI messages
  return result.toUIMessageStreamResponse({ sendReasoning: true });
};

import Together from "together-ai";

const together = new Together({
  apiKey: process.env.TOGETHER_API_KEY,
});

/**
 * Executes Python code with the Together code interpreter.
 */
export async function runPython(code: string, sessionId?: string) {
  const response = await together.codeInterpreter.execute({
    code,
    language: "python",
    ...(sessionId ? { session_id: sessionId } : {}),
  });

  if (response.errors) {
    return { status: "error" as const, errors: response.errors };
  }

  return {
    status: "success" as const,
    // Pass this back as sessionId on the next call to reuse the same
    // sandbox, so variables and installed packages persist across questions
    sessionId: response.data.session_id,
    // stdout/stderr text plus rich outputs such as image/png for charts
    outputs: response.data.outputs,
  };
}

const extractPythonCode = (text: string) => {
  const match = text.match(/```python\s*([\s\S]*?)```/);
  return match ? match[1].trim() : null;
};

export const analyzeCsv = async (
  question: string,
  csvFileUrl: string,
  csvHeaders: string[],
  csvRows: { [key: string]: string }[],
) => {
  // 1. Generate analysis code
  const result = streamText({
    model: togetherai("moonshotai/Kimi-K2.7-Code"),
    system: generateCodePrompt({ csvFileUrl, csvHeaders, csvRows }),
    messages: [{ role: "user", content: question }],
  });

  // 2. Wait for the full response and pull out the code block
  const code = extractPythonCode(await result.text);
  if (!code) return null;

  // 3. Run the code in the sandbox
  return runPython(code);
};