Skip to content

Commit

Permalink
generate json schema (json) files as part of run
Browse files Browse the repository at this point in the history
  • Loading branch information
pelikhan committed May 17, 2024
1 parent 6717499 commit eea7235
Show file tree
Hide file tree
Showing 10 changed files with 44 additions and 21 deletions.
2 changes: 2 additions & 0 deletions packages/cli/src/llamaindexretrieval.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import {
JSON_MIME_TYPE,
PDF_MIME_TYPE,
DOCX_MIME_TYPE,
JSON_SCHEMA_MIME_TYPE,
} from "genaiscript-core"
import { type BaseReader, OllamaEmbedding } from "llamaindex"
import type { GenericFileSystem } from "@llamaindex/env"
Expand Down Expand Up @@ -90,6 +91,7 @@ export class LlamaIndexRetrievalService
"text/plain": new this.module.TextFileReader(),
[JAVASCRIPT_MIME_TYPE]: new this.module.TextFileReader(),
[JSON_MIME_TYPE]: new this.module.TextFileReader(),
[JSON_SCHEMA_MIME_TYPE]: new this.module.TextFileReader(),
[PDF_MIME_TYPE]: new this.module.PDFReader(),
"text/markdown": new this.module.MarkdownReader(),
[DOCX_MIME_TYPE]: new this.module.DocxReader(),
Expand Down
16 changes: 9 additions & 7 deletions packages/cli/src/run.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,9 @@ import {
resolveModelConnectionInfo,
CONFIGURATION_ERROR_CODE,
parseKeyValuePairs,
stringifySchemaToTypeScript,
JSONSchemaStringifyToTypeScript,
filePathOrUrlToWorkspaceFile,
JSONSchemaStringify,
} from "genaiscript-core"
import { capitalize } from "inflection"
import { basename, resolve, join } from "node:path"
Expand Down Expand Up @@ -304,16 +305,17 @@ ${Array.from(files)
}
if (res.schemas) {
for (const [sname, schema] of Object.entries(res.schemas)) {
const schemaFile = join(
out,
`${sname.toLocaleLowerCase()}.schema.d.ts`
)
await writeText(
schemaFile,
stringifySchemaToTypeScript(schema, {
join(out, `${sname.toLocaleLowerCase()}.schema.ts`),
JSONSchemaStringifyToTypeScript(schema, {
typeName: capitalize(sname),
export: true,
})
)
await writeText(
join(out, `${sname.toLocaleLowerCase()}.schema.json`),
JSONSchemaStringify(schema)
)
}
}
if (annotationf) {
Expand Down
5 changes: 4 additions & 1 deletion packages/core/src/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -113,4 +113,7 @@ export const PDF_MIME_TYPE = "application/pdf"
export const DOCX_MIME_TYPE = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
export const XLSX_MIME_TYPE = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
export const JSON_MIME_TYPE = "application/json"
export const JAVASCRIPT_MIME_TYPE = "application/javascript"
export const JSON_SCHEMA_MIME_TYPE = "application/schema+json"
export const JAVASCRIPT_MIME_TYPE = "application/javascript"

export const JSON_META_SCHEMA_URI = "https://json-schema.org/draft/2020-12/schema"
7 changes: 3 additions & 4 deletions packages/core/src/expander.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { Fragment, Project, PromptScript } from "./ast"
import { assert, normalizeFloat, normalizeInt, normalizeString } from "./util"
import { assert, normalizeFloat, normalizeInt } from "./util"
import { MarkdownTrace } from "./trace"
import { errorMessage, isCancelError } from "./error"
import { estimateTokens } from "./tokens"
Expand All @@ -21,9 +21,8 @@ import {
toChatCompletionUserMessage,
} from "./chat"
import { importPrompt } from "./importprompt"
import { lookupMime } from "./mime"
import { parseModelIdentifier } from "./models"
import { stringifySchemaToTypeScript } from "./schema"
import { JSONSchemaStringifyToTypeScript } from "./schema"

const defaultTopP: number = undefined
const defaultSeed: number = undefined
Expand Down Expand Up @@ -387,7 +386,7 @@ export async function expandTemplate(
if (responseSchema) {
responseType = "json_object"
const typeName = "Output"
const schemaTs = stringifySchemaToTypeScript(responseSchema, {
const schemaTs = JSONSchemaStringifyToTypeScript(responseSchema, {
typeName,
})
messages.unshift({
Expand Down
1 change: 0 additions & 1 deletion packages/core/src/parameters.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import { assert } from "node:console"
import { Project } from "./ast"
import { NotSupportedError } from "./error"
import { resolveSystems } from "./expander"
Expand Down
4 changes: 2 additions & 2 deletions packages/core/src/promptdom.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { CSVToMarkdown, CSVTryParse } from "./csv"
import { resolveFileContent } from "./file"
import { addLineNumbers } from "./liner"
import { stringifySchemaToTypeScript } from "./schema"
import { JSONSchemaStringifyToTypeScript } from "./schema"
import { estimateTokens } from "./tokens"
import { MarkdownTrace, TraceOptions } from "./trace"
import { assert, toStringList, trimNewlines } from "./util"
Expand Down Expand Up @@ -548,7 +548,7 @@ export async function renderPromptNode(
schemaText = YAMLStringify(schema)
break
default:
schemaText = stringifySchemaToTypeScript(schema, {
schemaText = JSONSchemaStringifyToTypeScript(schema, {
typeName: schemaName,
})
break
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/retrieval.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import {
DOCX_MIME_TYPE,
JAVASCRIPT_MIME_TYPE,
JSON_MIME_TYPE,
JSON_SCHEMA_MIME_TYPE,
PDF_MIME_TYPE,
} from "./constants"
import {
Expand All @@ -16,6 +17,7 @@ const UPSERTFILE_MIME_TYPES = [
PDF_MIME_TYPE,
DOCX_MIME_TYPE,
JSON_MIME_TYPE,
JSON_SCHEMA_MIME_TYPE,
JAVASCRIPT_MIME_TYPE,
]

Expand Down
6 changes: 3 additions & 3 deletions packages/core/src/schema.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { describe, test } from "node:test"
import assert from "node:assert/strict"
import { stringifySchemaToTypeScript } from "./schema"
import { JSONSchemaStringifyToTypeScript } from "./schema"

describe("schema", () => {
test("cities", () => {
Expand Down Expand Up @@ -30,7 +30,7 @@ describe("schema", () => {
},
}

const ts = stringifySchemaToTypeScript(source, { typeName: "Foo" })
const ts = JSONSchemaStringifyToTypeScript(source, { typeName: "Foo" })
// console.log(ts)
assert.equal(
ts,
Expand Down Expand Up @@ -68,7 +68,7 @@ of the city.`,
required: ["name", "url"],
}

const ts = stringifySchemaToTypeScript(source)
const ts = JSONSchemaStringifyToTypeScript(source)
// console.log(ts)
assert.equal(
ts,
Expand Down
20 changes: 17 additions & 3 deletions packages/core/src/schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,18 @@ export function isJSONSchema(obj: any) {
return false
}

export function stringifySchemaToTypeScript(
export function JSONSchemaStringifyToTypeScript(
schema: JSONSchema,
options?: { typeName?: string }
options?: { typeName?: string; export?: boolean }
) {
const { typeName = "Response" } = options || {}
let lines: string[] = []
let indent = 0

appendJsDoc(schema.description)
append(`type ${typeName.replace(/\s+/g, "_")} =`)
append(
`${options?.export ? "export " : ""}type ${typeName.replace(/\s+/g, "_")} =`
)
stringifyNode(schema)
return lines.join("\n")

Expand Down Expand Up @@ -171,3 +173,15 @@ export function validateFencesWithSchema(
}
return frames
}

export function JSONSchemaStringify(schema: JSONSchema) {
return JSON.stringify(
{
$schema:
schema.$schema ?? "http://json-schema.org/draft-07/schema#",
...schema,
},
null,
2
)
}
2 changes: 2 additions & 0 deletions packages/core/src/types/prompt_template.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -617,6 +617,7 @@ interface JSONSchemaBoolean {
}

interface JSONSchemaObject {
$schema?: string
type: "object"
description?: string
properties?: {
Expand All @@ -627,6 +628,7 @@ interface JSONSchemaObject {
}

interface JSONSchemaArray {
$schema?: string
type: "array"
description?: string
items?: JSONSchemaType
Expand Down

0 comments on commit eea7235

Please sign in to comment.