rpc.ts

// Run `pnpm build:rpc` before running this example

import assert from 'node:assert'
import {
  type FileHandle,
  mkdir,
  open,
  readFile,
  writeFile,
} from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { inspect } from 'node:util'
import { WASI } from 'node:wasi'

const PROMPTL_WASM_PATH = './dist-rpc/promptl.wasm'

// --- Example 1: Basic chain without references ---

const prompt = `
<step>
  <system>
    You are a helpful assistant.
  </system>
  <user>
    Say hello.
  </user>
</step>
<step>
  <user>
    Now say goodbye.
  </user>
</step>
`
let chain, conversation
chain = await createChain(prompt);
({ chain, ...conversation } = await stepChain(chain));
({ chain, ...conversation } = await stepChain(chain, 'Hello!'));
({ chain, ...conversation } = await stepChain(chain, 'Goodbye!'));

assert(chain.completed)
assert(conversation.completed)

console.log('--- Basic chain ---')
console.log(inspect(conversation.messages, { depth: null }))

// --- Example 2: Chain with prompt references ---

const mainPrompt = `
<step>
  <prompt path="instructions" />
  <user>
    Tell me about {{ topic }}.
  </user>
</step>
`

const references = {
  instructions: `
<system>
  You are an expert assistant. Always be concise.
</system>
  `,
}

let refChain, refConversation
refChain = await createChain(mainPrompt, {
  parameters: { topic: 'TypeScript' },
  references: references,
});
({ chain: refChain, ...refConversation } = await stepChain(refChain));
({ chain: refChain, ...refConversation } = await stepChain(refChain, 'It is good!'));

assert(refChain.completed)
assert(refConversation.completed)

console.log('\n--- Chain with references ---')
console.log(inspect(refConversation.messages, { depth: null }))

// Utility functions

async function createChain(
  prompt: string,
  opts?: {
    parameters?: Record<string, unknown>
    references?: Record<string, string>
  },
): Promise<any> {
  return await execute([
    {
      procedure: 'createChain',
      parameters: {
        prompt,
        ...opts,
      },
    },
  ]).then((result) => {
    if (result[0]!.error) {
      console.log('Error dump: ', inspect(result[0]!.error, { depth: null }))
      throw new Error(result[0]!.error.message)
    }
    return result[0]!.value
  })
}

async function stepChain(chain: any, response?: any): Promise<any> {
  return await execute([
    {
      procedure: 'stepChain',
      parameters: {
        chain: chain,
        response: response,
      },
    },
  ]).then((result) => {
    if (result[0]!.error) {
      console.log('Error dump: ', inspect(result[0]!.error, { depth: null }))
      throw new Error(result[0]!.error.message)
    }
    return result[0]!.value
  })
}

async function execute(data: any): Promise<any> {
  const dir = join(tmpdir(), 'promptl')
  const stdin_path = join(dir, 'stdin')
  const stdout_path = join(dir, 'stdout')
  const stderr_path = join(dir, 'stderr')

  await mkdir(dir, { recursive: true })
  await writeFile(stdin_path, '')
  await writeFile(stdout_path, '')
  await writeFile(stderr_path, '')

  let stdin: FileHandle | undefined
  let stdout: FileHandle | undefined
  let stderr: FileHandle | undefined

  let wasmStdin: FileHandle | undefined
  let wasmStdout: FileHandle | undefined
  let wasmStderr: FileHandle | undefined

  try {
    stdin = await open(stdin_path, 'w')
    stdout = await open(stdout_path, 'r')
    stderr = await open(stderr_path, 'r')

    wasmStdin = await open(stdin_path, 'r')
    wasmStdout = await open(stdout_path, 'w')
    wasmStderr = await open(stderr_path, 'w')

    const wasi = new WASI({
      version: 'preview1',
      args: [],
      env: {},
      stdin: wasmStdin.fd,
      stdout: wasmStdout.fd,
      stderr: wasmStderr.fd,
      returnOnExit: true,
    })

    const bytes = await readFile(PROMPTL_WASM_PATH)
    WebAssembly.validate(bytes)
    const module = await WebAssembly.compile(bytes)
    const instance = await WebAssembly.instantiate(
      module,
      wasi.getImportObject(),
    )

    await send(stdin, data)

    wasi.start(instance)

    const [out, err] = await Promise.all([receive(stdout), receive(stderr)])
    if (err) throw new Error(err)

    return out
  } finally {
    await Promise.all([
      stdin?.close(),
      stdout?.close(),
      stderr?.close(),
      wasmStdin?.close(),
      wasmStdout?.close(),
      wasmStderr?.close(),
    ])
  }
}

async function send(file: FileHandle, data: any) {
  await writeFile(file, JSON.stringify(data) + '\n', {
    encoding: 'utf8',
    flush: true,
  })
}

async function receive(file: FileHandle): Promise<any> {
  const data = await readFile(file, {
    encoding: 'utf8',
  }).then((data) => data.trim())

  try {
    return JSON.parse(data)
  } catch {
    return data
  }
}