Merge pull request #3383 from quantified-uncertainty/load-prompt-fix

Fixed local ts code generation, added README to TS files

packages/ai/package.json

@@ -13,7 +13,7 @@
     "lint": "prettier --check .",
     "format": "prettier --write .",
     "test": "NODE_OPTIONS=--experimental-vm-modules jest",
-    "getCode": "tsx src/scripts/fetchSquiggleLibraries.ts",
+    "prepareSquiggleFiles": "tsx src/scripts/prepareSquiggleFiles.ts",
     "createSquiggle": "tsx src/scripts/createSquiggle.ts",
     "editSquiggle": "tsx src/scripts/editSquiggle.ts"
   },

packages/ai/src/Code.ts

@@ -11,10 +11,10 @@ import {
 } from "@quri/squiggle-lang";
 
 import { formatSquiggleCode } from "./squiggle/formatSquiggleCode.js";
-import { libraryContents } from "./squiggle/squiggleLibraryContents.js";
+import { LIBRARY_CONTENTS } from "./squiggle/squiggleLibraryContents.js";
 
 export const llmLinker = makeSelfContainedLinker(
-  Object.fromEntries(libraryContents)
+  Object.fromEntries(LIBRARY_CONTENTS)
 );
 
 // Evaluated code; source + result or error.

packages/ai/src/prompts.ts

@@ -1,27 +1,5 @@
-import fs from "fs";
-import path from "path";
-import { fileURLToPath } from "url";
-
-import { libraryContents } from "./squiggle/squiggleLibraryContents.js";
-
-const SQUIGGLE_DOCS_PATH = path.join(
-  path.dirname(fileURLToPath(import.meta.url)),
-  "..",
-  "files",
-  "squiggleDocs.md"
-);
-
-// Utility functions
-function readTxtFileSync(filePath: string) {
-  try {
-    return fs.readFileSync(filePath, "utf8");
-  } catch (err) {
-    console.error(`Error reading file: ${err}`);
-    throw err;
-  }
-}
-// Load Squiggle docs
-export const squiggleDocs = readTxtFileSync(SQUIGGLE_DOCS_PATH);
+import { README } from "./squiggle/README.js";
+import { LIBRARY_CONTENTS } from "./squiggle/squiggleLibraryContents.js";
 
 // Used as context for Claude, and as first message for other LLMs.
 export const squiggleSystemContent: string = `You are an AI assistant specialized in generating Squiggle code. Squiggle is a probabilistic programming language designed for estimation. Always respond with valid Squiggle code enclosed in triple backticks (\`\`\`). Do not give any more explanation, just provide the code and nothing else. Think through things, step by step.
@@ -30,11 +8,11 @@ Write the entire code, don't truncate it. So don't ever use "...", just write ou
 
 Here's the full Squiggle Documentation. It's important that all of the functions you use are contained here. Check this before finishing your work.
 
-${squiggleDocs}
+${README}
 
 ## Available libraries:
 
-${[...libraryContents.entries()]
+${[...LIBRARY_CONTENTS.entries()]
   .map(([name, content]) => `### Library ${name} \n\n ${content}`)
   .join("\n\n")}`;
 

packages/ai/src/scripts/fetchSquiggleLibraries.ts → packages/ai/src/scripts/prepareSquiggleFiles.ts

@@ -1,6 +1,8 @@
 import axios from "axios";
 import fs from "fs/promises";
-import path from "path";
+import path, { dirname } from "path";
+import prettier from "prettier";
+import { fileURLToPath } from "url";
 
 export const librariesToImport = ["ozziegooen/sTest", "ozziegooen/helpers"];
 
@@ -26,7 +28,7 @@ function getQuery(owner: string, slug: string) {
   `;
 }
 
-export function getLibraryPath(libName: string): string {
+function getLibraryPath(libName: string): string {
   const [author, name] = libName.split("/");
   return path.join(
     "squiggleLibraries",
@@ -35,6 +37,11 @@ export function getLibraryPath(libName: string): string {
   );
 }
 
+function getScriptPath() {
+  const __filename = fileURLToPath(import.meta.url);
+  return dirname(__filename);
+}
+
 async function fetchSquiggleCode(owner: string, slug: string) {
   try {
     const query = getQuery(owner, slug);
@@ -61,19 +68,26 @@ async function fetchAllLibraries() {
   }
 }
 
-// We save the files into a TS file, so that Vercel could read them.
-async function saveToTsFile(contents: (string | null)[], fileName: string) {
+async function saveToTsFile(
+  contents: string,
+  variableName: string,
+  fileName: string
+) {
   const fileContent = `// This file is auto-generated. Do not edit manually.
-export const squiggleLibraryContents = new Map([
-${contents.filter(Boolean).join(",\n")}
-]);
+export const ${variableName} = ${contents};
 `;
 
-  const outputPath = path.join(__dirname, "..", fileName);
-  await fs.writeFile(outputPath, fileContent);
-  console.log(`${fileName} has been generated successfully.`);
+  const prettierConfig = await prettier.resolveConfig(process.cwd());
+  const formattedContent = await prettier.format(fileContent, {
+    ...prettierConfig,
+    parser: "typescript",
+  });
+
+  const outputPath = path.join(getScriptPath(), "..", fileName);
+  await fs.writeFile(outputPath, formattedContent);
+  console.log(`${fileName} has been generated and formatted successfully.`);
 }
-// New function to generate squiggleLibraryContents.ts.
+
 async function generateLibraryContents() {
   const contents = await Promise.all(
     librariesToImport.map(async (lib) => {
@@ -88,13 +102,33 @@ async function generateLibraryContents() {
     })
   );
 
-  await saveToTsFile(contents, "squiggleLibraryContents.ts");
+  await saveToTsFile(
+    `new Map([${contents.filter(Boolean).join(",\n")}]);`,
+    "LIBRARY_CONTENTS",
+    "squiggle/squiggleLibraryContents.ts"
+  );
+}
+
+function getDocsPath(): string {
+  return path.join(getScriptPath(), "..", "..", "files", "squiggleDocs.md");
+}
+
+async function generateReadme() {
+  const docsPath = getDocsPath();
+  const readmeContent = await fs.readFile(docsPath, "utf8");
+
+  await saveToTsFile(
+    `${JSON.stringify(readmeContent)}`,
+    "README",
+    "squiggle/README.ts"
+  );
 }
 
-// Run the fetching process and generate LibraryContents.ts
+// We save the files into TS files, so that Vercel could read them.
 async function main() {
   await fetchAllLibraries();
   await generateLibraryContents();
+  await generateReadme();
 }
 
 main().catch((error) => {

packages/ai/src/squiggle/squiggleLibraryContents.ts

@@ -1,5 +1,5 @@
 // This file is auto-generated. Do not edit manually.
-export const libraryContents = new Map([
+export const LIBRARY_CONTENTS = new Map([
   [
     "hub:ozziegooen/sTest",
     '@startOpen\n@name("Documentation")\ndocumentation = "\n# SquiggleJest Testing Library\n\nSquiggleJest is a simple testing library for Squiggle, inspired by Jest for JavaScript. It provides a way to write and run tests for your Squiggle models and functions.\n\n## How to Use\n\n1. Import the library (assuming it\'s in a file named \'squiggleJest.squiggle\'):\n   ```squiggle\n   import \'squiggleJest.squiggle\' as SJ\n   ```\n\n2. Create your tests using the `test` function:\n   ```squiggle\n   test = SJ.test\n   expect = SJ.expect\n\n   myTest = test(\\"My test description\\", {|| \n     expect(2 + 2).toBe(4)\n   })\n   ```\n\n3. Group related tests using the `describe` function:\n   ```squiggle\n   describe = SJ.describe\n\n   myTestSuite = describe(\\"My Test Suite\\", [\n     test(\\"First test\\", {|| expect(true).toBeTrue()}),\n     test(\\"Second test\\", {|| expect(5).toBeGreaterThan(3)})\n   ])\n   ```\n\n4. Run your test suite and view the results.\n\n## Available Matchers\n\n- `toBe(expected)`: Checks for exact equality\n- `toBeGreaterThan(expected)`: Checks if the actual value is greater than the expected\n- `toBeGreaterThanOrEqual(expected)`: Checks if the actual value is greater or equal than the expected\n- `toBeLessThan(expected)`: Checks if the actual value is less than the expected\n- `toBeLessThanOrEqual(expected)`: Checks if the actual value is less than or equal than the expected\n- `toBeTrue()`: Checks if the value is true\n- `toBeFalse()`: Checks if the value is false\n- `toBeCloseTo(expected, epsilon)`: Checks if the actual value is close to the expected value within a given epsilon.\n- `toBeBetween(low, high)`: Checks if the actual value is between the given low and high values (inclusive).\n\n## Examples\n\n### Testing a Simple Function\n\n```squiggle\nadd(a, b) = a + b\n\ndescribe(\\"Add function\\", [\n  test(\\"adds two positive numbers\\", {|| \n    expect(add(2, 3)).toBe(5)\n  }),\n  test(\\"adds a positive and a negative number\\", {|| \n    expect(add(5, -3)).toBe(2)\n  })\n])\n```\n\n### Testing a Distribution\n\n```squiggle\nmyDist = normal(10, 2)\n\ndescribe(\\"My Distribution\\", [\n  test(\\"has correct mean\\", {|| \n    expect(mean(myDist)).toBe(10)\n  }),\n  test(\\"has correct standard deviation\\", {|| \n    expect(stdev(myDist)).toBe(2)\n  }),\n  test(\\"90% of values are within 2 standard deviations\\", {||\n    lower = 10 - 2 * 2\n    upper = 10 + 2 * 2\n    expect(cdf(myDist, upper) - cdf(myDist, lower)).toBeGreaterThan(0.9)\n  })\n])\n```\n\nThese examples demonstrate how to use SquiggleJest to test various aspects of your Squiggle models, from simple functions to complex models with distributions.\n"\n\n@startClosed\ncreateExpectation(actual) = {\n  toBe: {\n    |expected|\n    if actual != expected then "Expected " + expected + " but got " +\n      actual else true\n  },\n  toBeGreaterThan: {\n    |expected|\n    if actual > expected then true else "Expected " + actual +\n      " to be greater than " +\n      expected\n  },\n  toBeGreaterThanOrEqual: {\n    |expected|\n    if actual >= expected then true else "Expected " + actual +\n      " to be less than or equal" +\n      expected\n  },\n  toBeLessThan: {\n    |expected|\n    if actual < expected then true else "Expected " + actual +\n      " to be less than " +\n      expected\n  },\n  toBeLessThanOrEqual: {\n    |expected|\n    if actual <= expected then true else "Expected " + actual +\n      " to be less than or equal" +\n      expected\n  },\n  toBeBetween: {\n    |low, high|\n    if actual < low || actual > high then "Expected " + actual +\n      " to be between " +\n      low +\n      " and " +\n      high else true\n  },\n  toBeCloseTo: {\n    |expected, epsilon|\n    if abs(actual - expected) > epsilon then "Expected " + actual +\n      " to be close to " +\n      expected +\n      " (within " +\n      epsilon +\n      ")" else true\n  },\n  toBeTrue: {|| if !actual then "Expected true but got " + actual else true},\n  toBeFalse: {|| if actual then "Expected false but got " + actual else true},\n}\n\nrunTest(test) = {\n  fnResult = test.fn()\n  {\n    name: test.name,\n    passed: fnResult == true,\n    error: if fnResult != true then fnResult else "",\n  }\n}\n\n@startClosed\ngenerateTestReport(name, testResults) = {\n  passedTests = List.filter(testResults, {|t| t.passed})\n  failedTests = List.filter(testResults, {|t| !t.passed})\n\n  [\n    "## Test Suite: " + name + "  ",\n    "**Total tests**: " + List.length(testResults) + "  ",\n    "**Passed**: " + List.length(passedTests) + "  ",\n    "**Failed**: " + List.length(failedTests),\n    "",\n    "**Results:**  ",\n  ]\n}\n\n@startClosed\nformatTestResult(testResult) = (if testResult.passed then "✅" else "❌") +\n  "  " +\n  testResult.name +\n  (if testResult.error != "" then "\n    --- Error: *" + testResult.error +\n    "*" else "") +\n  "  "\n\n// Main squiggleJest framework\n@startClosed\nsquiggleJest = {\n  expect: createExpectation,\n  test: {|name, fn| { name: name, fn: fn }},\n  describe: {\n    |name, tests|\n    testResults = List.map(tests, runTest)\n    report = generateTestReport(name, testResults)\n    testDetails = List.map(testResults, formatTestResult)\n    List.concat(report, testDetails) -> List.join("\n")\n  },\n}\n\nexport test = squiggleJest.test\nexport describe = squiggleJest.describe\nexport expect = squiggleJest.expect\n\n/// Testing ---\n@name("Example Model")\nmodel = { items: [1, 2, 3] }\n\ntestResults = describe(\n  "Model Tests",\n  [\n    test(\n      "has items with length 3",\n      {|| expect(List.length(model.items)).toBe(3)}\n    ),\n    test("first item is 1", {|| expect(model.items[0]).toBe(1)}),\n    test(\n      "last item is greater than 2",\n      {|| expect(model.items[2]).toBeGreaterThan(1)}\n    ),\n    test(\n      "second item is less than 3",\n      {|| expect(model.items[1]).toBeLessThan(8)}\n    ),\n    test(\n      "second item is between 1 and 5",\n      {|| expect(model.items[1]).toBeBetween(1, 3)}\n    ),\n    test(\n      "contains truthy value",\n      {|| expect(List.some(model.items, {|i| i > 0})).toBeTrue()}\n    ),\n    test(\n      "doesn\'t contain 4",\n      {|| expect(List.some(model.items, {|i| i == 4})).toBeFalse()}\n    ),\n    test("this test should fail", {|| expect(1).toBe(2)}),\n  ]\n)\n\ncomparisonTests = describe(\n  "Number Comparisons",\n  [\n    test("5 is greater than 3", {|| expect(5).toBeGreaterThan(3)}),\n    test(\n      "5 is greater than or equal to 5",\n      {|| expect(5).toBeGreaterThanOrEqual(5)}\n    ),\n    test("3 is less than 5", {|| expect(3).toBeLessThan(5)}),\n    test("5 is less than or equal to 5", {|| expect(5).toBeLessThanOrEqual(5)}),\n    test("7 is between 5 and 10", {|| expect(7).toBeBetween(5, 10)}),\n    test(\n      "5 is close to 5.0001 within 0.001",\n      {|| expect(5).toBeCloseTo(5.0001, 0.001)}\n    ),\n    test("0 is not greater than 0", {|| expect(0).toBeLessThanOrEqual(0)}),\n    test("-1 is less than 0", {|| expect(-1).toBeLessThan(0)}),\n    test(\n      "1000000 is greater than 999999",\n      {|| expect(1000000).toBeGreaterThan(999999)}\n    ),\n    test(\n      "0.1 + 0.2 is close to 0.3",\n      {|| expect(0.1 + 0.2).toBeCloseTo(0.3, 0.0000001)}\n    ),\n    test(\n      "PI is approximately 3.14159",\n      {|| expect(3.14159).toBeCloseTo(Math.pi, 0.00001)}\n    ),\n    test(\n      "e is approximately 2.71828",\n      {|| expect(2.71828).toBeCloseTo(Math.e, 0.00001)}\n    ),\n    test(\n      "10 is between 9.99999 and 10.00001",\n      {|| expect(10).toBeBetween(9.99999, 10.00001)}\n    ),\n    test(\n      "5 is not between 5.00001 and 10",\n      {|| expect(5).toBeLessThan(5.00001)}\n    ),\n    test("1e-10 is greater than 0", {|| expect(1e-10).toBeGreaterThan(0)}),\n    test(\n      "The absolute difference between 1/3 and 0.333333 is less than 1e-5",\n      {|| expect(abs(1 / 3 - 0.333333)).toBeLessThan(1e-5)}\n    ),\n    test(\n      "2^53 - 1 is the largest integer precisely representable in IEEE 754",\n      {|| expect(2 ^ 53 - 1).toBe(9007199254740991)}\n    ),\n  ]\n)\n',