Skip to content

Commit

Permalink
Refactor tracing logic and add summarize-concurrent script (#718)
Browse files Browse the repository at this point in the history 
* Refactor tracing logic and add summarize-concurrent script for concurrent file summarization.

* Add trace details handling and refactor trace methods in core and CLI packages

* speed up build
  • Loading branch information
@pelikhan
pelikhan authored Sep 19, 2024
1 parent 052a523 commit a8460db
Show file tree
Hide file tree
Showing 7 changed files with 105 additions and 54 deletions.
4 changes: 0 additions & 4 deletions .github/workflows/build.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,3 @@ jobs:
run: yarn test:core
- name: unit tests
run: yarn test:samples
- name: slides
run: yarn build:slides
- name: docs
run: yarn build:docs
6 changes: 6 additions & 0 deletions packages/cli/src/run.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import {
RUNS_DIR_NAME,
CONSOLE_COLOR_DEBUG,
DOCS_CONFIGURATION_URL,
TRACE_DETAILS,
} from "../../core/src/constants"
import { isCancelError, errorMessage } from "../../core/src/error"
import { Fragment, GenerationResult } from "../../core/src/generation"
Expand Down Expand Up @@ -65,6 +66,7 @@ import {
} from "../../core/src/azuredevops"
import { resolveTokenEncoder } from "../../core/src/encoders"
import { writeFile } from "fs/promises"
import { writeFileSync } from "node:fs"

async function setupTraceWriting(trace: MarkdownTrace, filename: string) {
logVerbose(`trace: ${filename}`)
Expand All @@ -78,6 +80,10 @@ async function setupTraceWriting(trace: MarkdownTrace, filename: string) {
},
false
)
trace.addEventListener(TRACE_DETAILS, (ev) => {
const content = trace.content
writeFileSync(filename, content, { encoding: "utf-8" })
})
}

export async function runScriptWithExitCode(
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/constants.ts
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
export const CHANGE = "change"
export const TRACE_CHUNK = "traceChunk"
export const TRACE_DETAILS = "traceDetails"
export const RECONNECT = "reconnect"
export const OPEN = "open"
export const MAX_TOOL_CALLS = 10000
Expand Down
13 changes: 6 additions & 7 deletions packages/core/src/openai.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,13 +144,14 @@ export const OpenAIChatCompletion: ChatCompletionHandler = async (
throw e
}

trace.itemValue(`response`, `${r.status} ${r.statusText}`)
trace.itemValue(`status`, `${r.status} ${r.statusText}`)
if (r.status !== 200) {
let body: string
let responseBody: string
try {
body = await r.text()
responseBody = await r.text()
} catch (e) {}
const { error, message } = JSON5TryParse(body, {}) as {
trace.detailsFenced(`response`, responseBody, "json")
const { error, message } = JSON5TryParse(responseBody, {}) as {
error: any
message: string
}
Expand All @@ -165,13 +166,11 @@ export const OpenAIChatCompletion: ChatCompletionHandler = async (
r.status,
message ?? error?.message ?? r.statusText,
error,
body,
responseBody,
normalizeInt(r.headers.get("retry-after"))
)
}

trace.appendContent("\n\n")

let done = false
let finishReason: ChatCompletionResponse["finishReason"] = undefined
let chatResp = ""
Expand Down
72 changes: 39 additions & 33 deletions packages/core/src/promptcontext.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,18 +71,18 @@ export async function createPromptContext(
return res
},
grep: async (query, globs, options) => {
trace.startDetails(
const grepTrace = trace.startTraceDetails(
`🌐 grep <code>${HTMLEscape(typeof query === "string" ? query : query.source)}</code>`
)
try {
const { files } = await grepSearch(query, arrayify(globs), {
trace,
trace: grepTrace,
...options,
})
trace.files(files, { model, secrets: env.secrets })
grepTrace.files(files, { model, secrets: env.secrets })
return { files }
} finally {
trace.endDetails()
grepTrace.endDetails()
}
},
}
Expand Down Expand Up @@ -110,35 +110,38 @@ export async function createPromptContext(
fuzzSearch: async (q, files_, searchOptions) => {
const files = arrayify(files_)
searchOptions = searchOptions || {}
const fuzzTrace = trace.startTraceDetails(
`🧐 fuzz search <code>${HTMLEscape(q)}</code>`
)
try {
trace.startDetails(
`🧐 fuzz search <code>${HTMLEscape(q)}</code>`
)
if (!files?.length) {
trace.error("no files provided")
fuzzTrace.error("no files provided")
return []
} else {
const res = await fuzzSearch(q, files, searchOptions)
trace.files(res, {
const res = await fuzzSearch(q, files, {
...searchOptions,
trace: fuzzTrace,
})
fuzzTrace.files(res, {
model,
secrets: env.secrets,
skipIfEmpty: true,
})
return res
}
} finally {
trace.endDetails()
fuzzTrace.endDetails()
}
},
vectorSearch: async (q, files_, searchOptions) => {
const files = arrayify(files_).map(toWorkspaceFile)
searchOptions = { ...(searchOptions || {}) }
const vecTrace = trace.startTraceDetails(
`🔍 vector search <code>${HTMLEscape(q)}</code>`
)
try {
trace.startDetails(
`🔍 vector search <code>${HTMLEscape(q)}</code>`
)
if (!files?.length) {
trace.error("no files provided")
vecTrace.error("no files provided")
return []
}

Expand All @@ -154,17 +157,17 @@ export async function createPromptContext(
const res = await vectorSearch(q, files, {
...searchOptions,
folderPath,
trace,
trace: vecTrace,
})
// search
trace.files(res, {
vecTrace.files(res, {
model,
secrets: env.secrets,
skipIfEmpty: true,
})
return res
} finally {
trace.endDetails()
vecTrace.endDetails()
}
},
}
Expand Down Expand Up @@ -256,14 +259,17 @@ export async function createPromptContext(
return p
},
runPrompt: async (generator, runOptions): Promise<RunPromptResult> => {
const { label } = runOptions || {}
const runTrace = trace.startTraceDetails(
`🎁 run prompt ${label || ""}`
)
try {
const { label } = runOptions || {}
trace.startDetails(`🎁 run prompt ${label || ""}`)
infoCb?.({ text: `run prompt ${label || ""}` })

const genOptions = mergeGenerationOptions(options, runOptions)
genOptions.inner = true
const ctx = createChatGenerationContext(genOptions, trace)
genOptions.trace = runTrace
const ctx = createChatGenerationContext(genOptions, runTrace)
if (typeof generator === "string")
ctx.node.children.push(createTextNode(generator))
else await generator(ctx)
Expand Down Expand Up @@ -291,7 +297,7 @@ export async function createPromptContext(
chatParticipants: cps,
} = await renderPromptNode(genOptions.model, node, {
flexTokens: genOptions.flexTokens,
trace,
trace: runTrace,
})

schemas = scs
Expand All @@ -313,12 +319,12 @@ export async function createPromptContext(
const system = prj.getTemplate(systemId)
if (!system)
throw new Error(`system template ${systemId} not found`)
trace.startDetails(`👾 ${system.id}`)
runTrace.startDetails(`👾 ${system.id}`)
const sysr = await callExpander(
prj,
system,
env,
trace,
runTrace,
genOptions
)
if (sysr.images?.length)
Expand All @@ -334,18 +340,18 @@ export async function createPromptContext(
if (sysr.fileOutputs?.length)
throw new NotSupportedError("fileOutputs")
if (sysr.logs?.length)
trace.details("📝 console.log", sysr.logs)
runTrace.details("📝 console.log", sysr.logs)
if (sysr.text) {
systemMessage.content +=
SYSTEM_FENCE + "\n" + sysr.text + "\n"
trace.fence(sysr.text, "markdown")
runTrace.fence(sysr.text, "markdown")
}
if (sysr.aici) {
trace.fence(sysr.aici, "yaml")
runTrace.fence(sysr.aici, "yaml")
messages.push(sysr.aici)
}
trace.detailsFenced("js", system.jsSource, "js")
trace.endDetails()
runTrace.detailsFenced("js", system.jsSource, "js")
runTrace.endDetails()
if (sysr.status !== "success")
throw new Error(
`system ${system.id} failed ${sysr.status} ${sysr.statusText}`
Expand All @@ -355,7 +361,7 @@ export async function createPromptContext(

const connection = await resolveModelConnectionInfo(
genOptions,
{ trace, token: true }
{ trace: runTrace, token: true }
)
checkCancelled(cancellationToken)
if (!connection.configuration)
Expand All @@ -381,17 +387,17 @@ export async function createPromptContext(
chatParticipants,
genOptions
)
tracePromptResult(trace, resp)
tracePromptResult(runTrace, resp)
return resp
} catch (e) {
trace.error(e)
runTrace.error(e)
return {
text: "",
finishReason: isCancelError(e) ? "cancel" : "fail",
error: serializeError(e),
}
} finally {
trace.endDetails()
runTrace.endDetails()
}
},
}
Expand Down
42 changes: 32 additions & 10 deletions packages/core/src/trace.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import {
EMOJI_UNDEFINED,
TOOL_ID,
TRACE_CHUNK,
TRACE_DETAILS,
} from "./constants"
import { fenceMD, parseTraceTree, TraceTree } from "./markdown"
import { stringify as yamlStringify } from "yaml"
Expand All @@ -22,9 +23,9 @@ export class TraceChunkEvent extends Event {
}

export class MarkdownTrace extends EventTarget implements ToolCallTrace {
readonly errors: { message: string; error: SerializedError }[] = []
readonly _errors: { message: string; error: SerializedError }[] = []
private detailsDepth = 0
private _content: string = ""
private _content: (string | MarkdownTrace)[] = []
private _tree: TraceTree

constructor(
Expand All @@ -42,17 +43,34 @@ export class MarkdownTrace extends EventTarget implements ToolCallTrace {
}

get tree() {
if (!this._tree) this._tree = parseTraceTree(this._content)
if (!this._tree) this._tree = parseTraceTree(this.content)
return this._tree
}

get content() {
get content(): string {
return this._content
.map((c) => (typeof c === "string" ? c : c.content))
.join("")
}

startTraceDetails(title: string) {
const trace = new MarkdownTrace({ ...this.options })
trace.addEventListener(TRACE_CHUNK, (ev) =>
this.dispatchEvent(
new TraceChunkEvent((ev as TraceChunkEvent).chunk)
)
)
trace.addEventListener(TRACE_DETAILS, () =>
this.dispatchEvent(new Event(TRACE_DETAILS))
)
trace.startDetails(title)
this._content.push(trace)
return trace
}

appendContent(value: string) {
if (value !== undefined && value !== null && value !== "") {
this._content = this._content + value
this._content.push(value)
this._tree = undefined
this.dispatchChange()
this.dispatchEvent(new TraceChunkEvent(value))
Expand Down Expand Up @@ -83,6 +101,7 @@ ${this.toResultIcon(success, "")}${title}
if (this.detailsDepth > 0) {
this.detailsDepth--
this.appendContent(`\n</details>\n\n`)
this.dispatchEvent(new Event(TRACE_DETAILS))
}
}

Expand Down Expand Up @@ -161,10 +180,6 @@ ${this.toResultIcon(success, "")}${title}
this.appendContent(fenceMD(res, contentType))
}

append(trace: MarkdownTrace) {
this.appendContent("\n" + trace.content)
}

tip(message: string) {
this.appendContent(`> ${message}\n`)
}
Expand Down Expand Up @@ -195,11 +210,18 @@ ${this.toResultIcon(success, "")}${title}
message,
error: serializeError(error),
}
this.errors.push(err)
this._errors.push(err)
this.renderError(err, { details: true })
})
}

get errors(): { message: string; error: SerializedError }[] {
const traces = this._content.filter(
(c) => typeof c !== "string"
) as MarkdownTrace[]
return this._errors.concat(...traces.map((t) => t.errors))
}

renderErrors(): void {
while (this.detailsDepth > 0) this.endDetails()
const errors = this.errors || []
Expand Down
21 changes: 21 additions & 0 deletions packages/sample/genaisrc/summarize-concurrent.genai.mjs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
script({
title: "summarize concurrently",
model: "openai:gpt-3.5-turbo",
files: "src/rag/*",
})

const summaries = await Promise.all(
env.files.map((file) =>
runPrompt(
(_) => {
_.def("FILE", file)
_.$`Summarize FILE with one paragraph.`
},
{ label: file.filename }
)
)
)

summaries.forEach((s) => def("FILE", s.text))

$`Summarize FILE.`

0 comments on commit a8460db

Please sign in to comment.