import { Mistral } from "@mistralai/mistralai";
import { z } from "zod";

const apiKey = process.env["MISTRAL_API_KEY"];
if (!apiKey) {
  throw new Error("missing MISTRAL_API_KEY environment variable");
}

const Explanation = z.object({
    explanation: z.string(),
    output: z.string(),
  });

const MathDemonstration = z.object({
    steps: z.array(Explanation),
    final_answer: z.string(),
  });

const client = new Mistral({ apiKey: apiKey });

// Structured Output using the parseStream method
// This method will return a stream of events containing the JSON response. It will not be parsed back to the initial Zod object.
// If you need to parse the stream, see the "Parse the accumulated stream back to the initial zod object" comment below.

const stream = await client.chat.parseStream({
  model: "mistral-tiny-latest",
  messages: [
    { role: "system", content: "You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning." },
    { role: "user", content: "How can I solve 8x + 7 = -23" }
  ],
  responseFormat: MathDemonstration,
});

let accumulatedStream = '';
for await (const event of stream) {
  const content = event.data?.choices[0]?.delta.content;
  if (!content) {
    continue;
  }
  process.stdout.write(content);
  accumulatedStream += content;
}

// Parse the accumulated stream back to the initial zod object.
let parsedAccumulatedStream = MathDemonstration.safeParse(JSON.parse(accumulatedStream)).data;
console.log('\n', parsedAccumulatedStream);

// Structured Output using the parse method
// This method will return the response parsed back to the initial Zod object.

const chatResponse = await client.chat.parse({
  model: "mistral-tiny-latest",
  messages: [
    { role: "system", content: "You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning." },
    { role: "user", content: "How can I solve 8x + 7 = -23" }
  ],
  responseFormat: MathDemonstration,
});

console.log('\n', chatResponse.choices[0].message.parsed);
