Skip to content

Commit 90b7434

Browse files
authored
Add typescript collector (#1573)
1 parent 2a3c865 commit 90b7434

36 files changed

Lines changed: 7747 additions & 4258 deletions

engine/language_client_codegen/src/typescript/templates/async_client.ts.j2

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { BamlRuntime, FunctionResult, BamlCtxManager, ClientRegistry, Image, Audio } from "@boundaryml/baml"
1+
import type { BamlRuntime, FunctionResult, BamlCtxManager, ClientRegistry, Image, Audio, Collector } from "@boundaryml/baml"
22
import { toBamlError, BamlStream } from "@boundaryml/baml"
33
import type { Checked, Check, RecursivePartialNull as MovedRecursivePartialNull } from "./types"
44
import type { partial_types } from "./partial_types"
@@ -34,9 +34,10 @@ export class BamlAsyncClient {
3434
{% for (name, optional, type) in fn.args -%}
3535
{{name}}{% if optional %}?{% endif %}: {{type}},
3636
{%- endfor %}
37-
__baml_options__?: { tb?: TypeBuilder, clientRegistry?: ClientRegistry }
37+
__baml_options__?: { tb?: TypeBuilder, clientRegistry?: ClientRegistry, collector?: Collector | Collector[] }
3838
): Promise<{{fn.return_type}}> {
3939
try {
40+
const collector = __baml_options__?.collector ? (Array.isArray(__baml_options__.collector) ? __baml_options__.collector : [__baml_options__.collector]) : [];
4041
const raw = await this.runtime.callFunction(
4142
"{{fn.name}}",
4243
{
@@ -47,6 +48,7 @@ export class BamlAsyncClient {
4748
this.ctx_manager.cloneContext(),
4849
__baml_options__?.tb?.__tb(),
4950
__baml_options__?.clientRegistry,
51+
collector,
5052
)
5153
return raw.parsed(false) as {{fn.return_type}}
5254
} catch (error) {
@@ -64,9 +66,10 @@ class BamlStreamClient {
6466
{% for (name, optional, type) in fn.args -%}
6567
{{name}}{% if optional %}?{% endif %}: {{type}},
6668
{%- endfor %}
67-
__baml_options__?: { tb?: TypeBuilder, clientRegistry?: ClientRegistry }
69+
__baml_options__?: { tb?: TypeBuilder, clientRegistry?: ClientRegistry, collector?: Collector | Collector[] }
6870
): BamlStream<{{ fn.partial_return_type }}, {{ fn.return_type }}> {
6971
try {
72+
const collector = __baml_options__?.collector ? (Array.isArray(__baml_options__.collector) ? __baml_options__.collector : [__baml_options__.collector]) : [];
7073
const raw = this.runtime.streamFunction(
7174
"{{fn.name}}",
7275
{
@@ -78,6 +81,7 @@ class BamlStreamClient {
7881
this.ctx_manager.cloneContext(),
7982
__baml_options__?.tb?.__tb(),
8083
__baml_options__?.clientRegistry,
84+
collector,
8185
)
8286
return new BamlStream<{{ fn.partial_return_type }}, {{ fn.return_type }}>(
8387
raw,

engine/language_client_codegen/src/typescript/templates/sync_client.ts.j2

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { BamlRuntime, FunctionResult, BamlCtxManager, Image, Audio, ClientRegistry } from "@boundaryml/baml"
1+
import type { BamlRuntime, FunctionResult, BamlCtxManager, Image, Audio, ClientRegistry, Collector } from "@boundaryml/baml"
22
import { toBamlError } from "@boundaryml/baml"
33
import type { Checked, Check, RecursivePartialNull as MovedRecursivePartialNull } from "./types"
44
import type * as types from "./types"
@@ -35,21 +35,23 @@ export class BamlSyncClient {
3535
{% for (name, optional, type) in fn.args -%}
3636
{{name}}{% if optional %}?{% endif %}: {{type}},
3737
{%- endfor %}
38-
__baml_options__?: { tb?: TypeBuilder, clientRegistry?: ClientRegistry }
38+
__baml_options__?: { tb?: TypeBuilder, clientRegistry?: ClientRegistry, collector?: Collector | Collector[] }
3939
): {{fn.return_type}} {
4040
try {
41-
const raw = this.runtime.callFunctionSync(
42-
"{{fn.name}}",
43-
{
44-
{% for (name, optional, type) in fn.args -%}
45-
"{{name}}": {{name}}{% if optional %}?? null{% endif %}{% if !loop.last %},{% endif %}
46-
{%- endfor %}
47-
},
48-
this.ctx_manager.cloneContext(),
49-
__baml_options__?.tb?.__tb(),
50-
__baml_options__?.clientRegistry,
51-
)
52-
return raw.parsed(false) as {{fn.return_type}}
41+
const collector = __baml_options__?.collector ? (Array.isArray(__baml_options__.collector) ? __baml_options__.collector : [__baml_options__.collector]) : [];
42+
const raw = this.runtime.callFunctionSync(
43+
"{{fn.name}}",
44+
{
45+
{% for (name, optional, type) in fn.args -%}
46+
"{{name}}": {{name}}{% if optional %}?? null{% endif %}{% if !loop.last %},{% endif %}
47+
{%- endfor %}
48+
},
49+
this.ctx_manager.cloneContext(),
50+
__baml_options__?.tb?.__tb(),
51+
__baml_options__?.clientRegistry,
52+
collector,
53+
)
54+
return raw.parsed(false) as {{fn.return_type}}
5355
} catch (error: any) {
5456
throw toBamlError(error);
5557
}
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
export * from "./errors";
2-
export { BamlRuntime, FunctionResult, FunctionResultStream, BamlImage as Image, ClientBuilder, BamlAudio as Audio, invoke_runtime_cli, ClientRegistry, BamlLogEvent, } from "./native";
2+
export { BamlRuntime, FunctionResult, FunctionResultStream, BamlImage as Image, ClientBuilder, BamlAudio as Audio, invoke_runtime_cli, ClientRegistry, BamlLogEvent, Collector, FunctionLog, Usage, } from "./native";
33
export { BamlStream } from "./stream";
44
export { BamlCtxManager } from "./async_context_vars";
55
//# sourceMappingURL=index.d.ts.map

engine/language_client_typescript/index.d.ts.map

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

engine/language_client_typescript/index.js

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ var __exportStar = (this && this.__exportStar) || function(m, exports) {
1414
for (var p in m) if (p !== "default" && !Object.prototype.hasOwnProperty.call(exports, p)) __createBinding(exports, m, p);
1515
};
1616
Object.defineProperty(exports, "__esModule", { value: true });
17-
exports.BamlCtxManager = exports.BamlStream = exports.BamlLogEvent = exports.ClientRegistry = exports.invoke_runtime_cli = exports.Audio = exports.ClientBuilder = exports.Image = exports.FunctionResultStream = exports.FunctionResult = exports.BamlRuntime = void 0;
17+
exports.BamlCtxManager = exports.BamlStream = exports.Usage = exports.FunctionLog = exports.Collector = exports.BamlLogEvent = exports.ClientRegistry = exports.invoke_runtime_cli = exports.Audio = exports.ClientBuilder = exports.Image = exports.FunctionResultStream = exports.FunctionResult = exports.BamlRuntime = void 0;
1818
__exportStar(require("./errors"), exports);
1919
var native_1 = require("./native");
2020
Object.defineProperty(exports, "BamlRuntime", { enumerable: true, get: function () { return native_1.BamlRuntime; } });
@@ -26,6 +26,9 @@ Object.defineProperty(exports, "Audio", { enumerable: true, get: function () { r
2626
Object.defineProperty(exports, "invoke_runtime_cli", { enumerable: true, get: function () { return native_1.invoke_runtime_cli; } });
2727
Object.defineProperty(exports, "ClientRegistry", { enumerable: true, get: function () { return native_1.ClientRegistry; } });
2828
Object.defineProperty(exports, "BamlLogEvent", { enumerable: true, get: function () { return native_1.BamlLogEvent; } });
29+
Object.defineProperty(exports, "Collector", { enumerable: true, get: function () { return native_1.Collector; } });
30+
Object.defineProperty(exports, "FunctionLog", { enumerable: true, get: function () { return native_1.FunctionLog; } });
31+
Object.defineProperty(exports, "Usage", { enumerable: true, get: function () { return native_1.Usage; } });
2932
var stream_1 = require("./stream");
3033
Object.defineProperty(exports, "BamlStream", { enumerable: true, get: function () { return stream_1.BamlStream; } });
3134
var async_context_vars_1 = require("./async_context_vars");

engine/language_client_typescript/native.d.ts

Lines changed: 92 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,10 @@ export declare class BamlRuntime {
2323
static fromFiles(rootPath: string, files: Record<string, string>, envVars: Record<string, string | undefined | null>): BamlRuntime
2424
reset(rootPath: string, files: Record<string, string>, envVars: Record<string, string>): void
2525
createContextManager(): RuntimeContextManager
26-
callFunction(functionName: string, args: { [name: string]: any }, ctx: RuntimeContextManager, tb?: TypeBuilder | undefined | null, cb?: ClientRegistry | undefined | null): Promise<FunctionResult>
27-
callFunctionSync(functionName: string, args: { [name: string]: any }, ctx: RuntimeContextManager, tb?: TypeBuilder | undefined | null, cb?: ClientRegistry | undefined | null): FunctionResult
28-
streamFunction(functionName: string, args: { [name: string]: any }, cb: ((err: any, param: FunctionResult) => void) | undefined, ctx: RuntimeContextManager, tb?: TypeBuilder | undefined | null, clientRegistry?: ClientRegistry | undefined | null): FunctionResultStream
29-
streamFunctionSync(functionName: string, args: { [name: string]: any }, cb: ((err: any, param: FunctionResult) => void) | undefined, ctx: RuntimeContextManager, tb?: TypeBuilder | undefined | null, clientRegistry?: ClientRegistry | undefined | null): FunctionResultStream
26+
callFunction(functionName: string, args: { [name: string]: any }, ctx: RuntimeContextManager, tb: TypeBuilder | undefined | null, cb: ClientRegistry | undefined | null, collectors: Array<Collector>): Promise<FunctionResult>
27+
callFunctionSync(functionName: string, args: { [name: string]: any }, ctx: RuntimeContextManager, tb: TypeBuilder | undefined | null, cb: ClientRegistry | undefined | null, collectors: Array<Collector>): FunctionResult
28+
streamFunction(functionName: string, args: { [name: string]: any }, cb: ((err: any, param: FunctionResult) => void) | undefined, ctx: RuntimeContextManager, tb: TypeBuilder | undefined | null, clientRegistry: ClientRegistry | undefined | null, collectors: Array<Collector>): FunctionResultStream
29+
streamFunctionSync(functionName: string, args: { [name: string]: any }, cb: ((err: any, param: FunctionResult) => void) | undefined, ctx: RuntimeContextManager, tb: TypeBuilder | undefined | null, clientRegistry: ClientRegistry | undefined | null, collectors: Array<Collector>): FunctionResultStream
3030
setLogEventCallback(func?: undefined | ((err: any, param: BamlLogEvent) => void)): void
3131
flush(): void
3232
drainStats(): TraceStats
@@ -54,6 +54,17 @@ export declare class ClientRegistry {
5454
setPrimary(primary: string): void
5555
}
5656

57+
export declare class Collector {
58+
constructor(name?: string | undefined | null)
59+
get logs(): Array<FunctionLog>
60+
get last(): FunctionLog | null
61+
id(functionLogId: string): FunctionLog | null
62+
get usage(): Usage
63+
toString(): string
64+
static __functionSpanCount(): number
65+
static __printStorage(): void
66+
}
67+
5768
export declare class EnumBuilder {
5869
value(name: string): EnumValueBuilder
5970
alias(alias?: string | undefined | null): EnumBuilder
@@ -71,6 +82,18 @@ export declare class FieldType {
7182
optional(): FieldType
7283
}
7384

85+
export declare class FunctionLog {
86+
toString(): string
87+
get id(): string
88+
get functionName(): string
89+
get logType(): string
90+
get timing(): Timing
91+
get usage(): Usage
92+
get calls(): (LLMCall | LLMStreamCall)[]
93+
get rawLlmResponse(): string | null
94+
get selectedCall(): unknown
95+
}
96+
7497
export declare class FunctionResult {
7598
isOk(): boolean
7699
parsed(allowPartials: boolean): any
@@ -81,12 +104,71 @@ export declare class FunctionResultStream {
81104
done(rctx: RuntimeContextManager): Promise<FunctionResult>
82105
}
83106

107+
export declare class HttpRequest {
108+
get bodyRaw(): string
109+
get body(): any
110+
toString(): string
111+
get url(): string
112+
get method(): string
113+
get headers(): object
114+
}
115+
export type HTTPRequest = HttpRequest
116+
117+
export declare class HttpResponse {
118+
toString(): string
119+
get status(): number
120+
get headers(): object
121+
get body(): any
122+
}
123+
export type HTTPResponse = HttpResponse
124+
125+
export declare class LlmCall {
126+
get selected(): boolean
127+
get httpRequest(): HTTPRequest | null
128+
get httpResponse(): HTTPResponse | null
129+
get usage(): Usage | null
130+
get timing(): Timing
131+
get provider(): string
132+
get clientName(): string
133+
toString(): string
134+
toString(): string
135+
}
136+
export type LLMCall = LlmCall
137+
138+
export declare class LlmStreamCall {
139+
toString(): string
140+
get httpRequest(): HTTPRequest | null
141+
get httpResponse(): HTTPResponse | null
142+
get provider(): string
143+
get clientName(): string
144+
get selected(): boolean
145+
get usage(): Usage | null
146+
get timing(): StreamTiming
147+
toString(): string
148+
}
149+
export type LLMStreamCall = LlmStreamCall
150+
84151
export declare class RuntimeContextManager {
85152
upsertTags(tags: any): void
86153
deepClone(): RuntimeContextManager
87154
contextDepth(): number
88155
}
89156

157+
export declare class StreamTiming {
158+
toString(): string
159+
get startTimeUtcMs(): number
160+
get durationMs(): number | null
161+
get timeToFirstParsedMs(): number | null
162+
get timeToFirstTokenMs(): number | null
163+
}
164+
165+
export declare class Timing {
166+
toString(): string
167+
get startTimeUtcMs(): number
168+
get durationMs(): number | null
169+
get timeToFirstParsedMs(): number | null
170+
}
171+
90172
export declare class TraceStats {
91173
get failed(): number
92174
get started(): number
@@ -117,6 +199,12 @@ export declare class TypeBuilder {
117199
toString(): string
118200
}
119201

202+
export declare class Usage {
203+
toString(): string
204+
get inputTokens(): number | null
205+
get outputTokens(): number | null
206+
}
207+
120208
export interface BamlLogEvent {
121209
metadata: LogEventMetadata
122210
prompt?: string

engine/language_client_typescript/native.js

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -368,12 +368,25 @@ module.exports.BamlSpan = nativeBinding.BamlSpan
368368
module.exports.ClassBuilder = nativeBinding.ClassBuilder
369369
module.exports.ClassPropertyBuilder = nativeBinding.ClassPropertyBuilder
370370
module.exports.ClientRegistry = nativeBinding.ClientRegistry
371+
module.exports.Collector = nativeBinding.Collector
371372
module.exports.EnumBuilder = nativeBinding.EnumBuilder
372373
module.exports.EnumValueBuilder = nativeBinding.EnumValueBuilder
373374
module.exports.FieldType = nativeBinding.FieldType
375+
module.exports.FunctionLog = nativeBinding.FunctionLog
374376
module.exports.FunctionResult = nativeBinding.FunctionResult
375377
module.exports.FunctionResultStream = nativeBinding.FunctionResultStream
378+
module.exports.HttpRequest = nativeBinding.HttpRequest
379+
module.exports.HTTPRequest = nativeBinding.HTTPRequest
380+
module.exports.HttpResponse = nativeBinding.HttpResponse
381+
module.exports.HTTPResponse = nativeBinding.HTTPResponse
382+
module.exports.LlmCall = nativeBinding.LlmCall
383+
module.exports.LLMCall = nativeBinding.LLMCall
384+
module.exports.LlmStreamCall = nativeBinding.LlmStreamCall
385+
module.exports.LLMStreamCall = nativeBinding.LLMStreamCall
376386
module.exports.RuntimeContextManager = nativeBinding.RuntimeContextManager
387+
module.exports.StreamTiming = nativeBinding.StreamTiming
388+
module.exports.Timing = nativeBinding.Timing
377389
module.exports.TraceStats = nativeBinding.TraceStats
378390
module.exports.TypeBuilder = nativeBinding.TypeBuilder
391+
module.exports.Usage = nativeBinding.Usage
379392
module.exports.invoke_runtime_cli = nativeBinding.invoke_runtime_cli

0 commit comments

Comments
 (0)