ZON Vs TOON Langchain example script

langchain-demo.ts

import * as ZON from 'zon-format';
// @ts-ignore
import * as TOON from '@toon-format/toon';

async function main() {
  try {
    const dotenv = await import("dotenv");
    dotenv.config();

    const { ChatOpenAI, AzureChatOpenAI } = await import("@langchain/openai");
    const { PromptTemplate } = await import("@langchain/core/prompts");
    const { RunnableSequence, RunnableLambda } = await import("@langchain/core/runnables");
    const { StringOutputParser } = await import("@langchain/core/output_parsers");

    const hasOpenAI = !!process.env.OPENAI_API_KEY;
    const hasAzure = !!process.env.AZURE_OPENAI_API_KEY;

    if (!hasOpenAI && !hasAzure) {
      console.warn("⚠️  No API keys found.");
      return;
    }

    let model;
    if (hasAzure) {
      const deploymentName = process.env.AZURE_OPENAI_API_DEPLOYMENT_NAME 
        || process.env.AZURE_GPT5_DEPLOYMENT;

      let instanceName = process.env.AZURE_OPENAI_API_INSTANCE_NAME;
      const endpoint = process.env.AZURE_OPENAI_ENDPOINT;
      
      if (!instanceName && endpoint) {
        try {
          const url = new URL(endpoint);
          const hostParts = url.hostname.split('.');
          if (hostParts.length > 0) instanceName = hostParts[0];
        } catch (e) { }
      }

      model = new AzureChatOpenAI({
        azureOpenAIApiDeploymentName: deploymentName,
        azureOpenAIApiInstanceName: instanceName,
        azureOpenAIApiKey: process.env.AZURE_OPENAI_API_KEY,
        azureOpenAIApiVersion: process.env.AZURE_OPENAI_API_VERSION || "2024-02-01"
      });
    } else {
      model = new ChatOpenAI({ modelName: "gpt-4o", temperature: 0 });
    }

    const stringParser = new StringOutputParser();

    const createParser = (name: string, decodeFn: (text: string) => any) => 
      new RunnableLambda({
        func: async (text: string) => {
          console.log(`\n[${name}] Raw LLM Response:`);
          console.log(text);
          try {
            return decodeFn(text.trim().replace(/^```(zon|zonf|yaml)?/, "").replace(/```$/, "").trim());
          } catch (e) {
            console.error(`Error parsing ${name}:`, e);
            return null;
          }
        }
      });

    const exampleObject = {
      portfolios: [
        { 
          owner: "Alice Smith", 
          role: "Designer",
          skills: [
             { level: "Expert", name: "Figma" }
          ]
        },
        { 
          owner: "Bob Jones", 
          role: "Developer",
          skills: [
             { level: "Advanced", name: "CSS" }
          ]
        }
      ]
    };
    
    console.log("\n[1] Input Data (JSON):");
    console.dir(exampleObject, { depth: null, colors: true });

    console.log("\n" + "=".repeat(40));
    console.log("   RUNNING ZON DEMO");
    console.log("=".repeat(40));

    const exampleZon = ZON.encode(exampleObject);
    console.log("\n[ZON] Encoded Example (Input to LLM):");
    console.log(exampleZon);

    const zonPrompt = PromptTemplate.fromTemplate(`
     You are a data assistant.
      I will show you an example data structure.
      Your task is to generate output for the new request following exactly the SAME format and structure.

      Example Data:
      {example_data}
      
      New User Request: {query}
    `);

    const zonChain = RunnableSequence.from([
      zonPrompt,
      model,
      stringParser,
      createParser("ZON", ZON.decode)
    ]);

    const query = "Create a portfolio for two different people, Roni Bhakta & Rakesh Bhakta, a software engineer who likes TypeScript (Expert) and Python (Advanced).";

    const zonResult = await zonChain.invoke({
      query: query,
      example_data: exampleZon
    });
    
    console.log("\n[ZON] Final Result:");
    console.dir(zonResult, { depth: null, colors: true });

    console.log("\n" + "=".repeat(40));
    console.log("   RUNNING TOON DEMO");
    console.log("=".repeat(40));

    const exampleToon = TOON.encode(exampleObject);
    console.log("\n[TOON] Encoded Example (Input to LLM):");
    console.log(exampleToon);

    const toonPrompt = PromptTemplate.fromTemplate(`
      You are a data assistant.
      I will show you an example data structure.
      Your task is to generate output for the new request following exactly the SAME format and structure.

      Example Data:
      {example_data}
      
      New User Request: {query}
    `);

    const toonChain = RunnableSequence.from([
      toonPrompt,
      model,
      stringParser,
      createParser("TOON", TOON.decode)
    ]);

    const toonResult = await toonChain.invoke({
      query: query,
      example_data: exampleToon
    });

    console.log("\n[TOON] Final Result:");
    console.dir(toonResult, { depth: null, colors: true });

    console.log("\n" + "=".repeat(40));
    console.log("   RUNNING ZON GUARDRAILS DEMO");
    console.log("=".repeat(40));

    const PortfolioSchema = ZON.zon.object({
      portfolios: ZON.zon.array(ZON.zon.object({
        owner: ZON.zon.string().describe("Full Name"),
        role: ZON.zon.enum(["Software Engineer", "Designer", "Manager"]).describe("Job Title"),
        skills: ZON.zon.array(ZON.zon.object({
          name: ZON.zon.string(),
          level: ZON.zon.enum(["Expert", "Advanced", "Intermediate", "Beginner"])
        }))
      }))
    });

    console.log("\n[Guardrails] Defined Schema:");
    // @ts-ignore - access internal structure for display if simple stringify fails
    console.log(JSON.stringify(PortfolioSchema, null, 2)); 

    const guardrailsPrompt = PromptTemplate.fromTemplate(`
      You are a strict API. 
      Generate a response for the user request matching this exact ZON schema:

      ${PortfolioSchema.toPrompt()}

      Format Reference (Use this syntax style):
      {example_code}

      User Request: {query}
    `);

    const guardrailsChain = RunnableSequence.from([
      guardrailsPrompt,
      model,
      stringParser,
      createParser("GUARDRAILS", ZON.decode)
    ]);

    const grQuery = "Create a portfolio for Roni (Engineer), Alice (Designer) and Rakesh (Manager). Roni knows Rust (Beginner), Rakesh knows TypeScript (Expert).";
    
    const grResult = await guardrailsChain.invoke({
      query: grQuery,
      example_code: exampleZon
    });

    console.log("\n[Guardrails] Raw Result:");
    console.dir(grResult, { depth: null, colors: true });

    const validation = ZON.validate(grResult, PortfolioSchema);
    
    console.log("\n[Guardrails] Validation Report:");
    if (validation.success) {
      console.log("✅ SUCCESS: Output matches strict schema!");
    } else {
      console.log("❌ FAILURE: Schema violation detected!");
      console.dir(validation.error, { depth: null, colors: true });
    }


  } catch (e: any) {
    console.error("Error:", e);
  }
}

main().catch(console.error);