@telemetry-dev/bedrock 0.1.0 → 0.1.2

package/src/agent-handlers.ts

@@ -13,15 +13,17 @@ import {
   PROVIDER,
   stringValue,
   type BedrockInstrumentationOptions,
+  type JsonRecord,
+  type JsonValue,
   type StreamState,
 } from "./internal.ts";
 
-export function agentInputMessages(inputText: unknown): unknown {
+export function agentInputMessages<T>(inputText: T): JsonValue {
   const text = stringValue(inputText);
   return text ? [{ role: "user", parts: [{ type: "text", content: text }] }] : undefined;
 }
 
-export function agentRequestFields(input: Record<string, unknown>): SpanFields {
+export function agentRequestFields(input: JsonRecord): SpanFields {
   return omitUndefined({
     provider: PROVIDER,
     input: agentInputMessages(input.inputText),
@@ -34,7 +36,7 @@ export function agentRequestFields(input: Record<string, unknown>): SpanFields {
   });
 }
 
-export function inlineAgentRequestFields(input: Record<string, unknown>): SpanFields {
+export function inlineAgentRequestFields(input: JsonRecord): SpanFields {
   return omitUndefined({
     input: agentInputMessages(input.inputText),
     model: stringValue(input.foundationModel),
@@ -59,28 +61,29 @@ export class AgentStreamState implements StreamState {
 
   constructor(private readonly options: BedrockInstrumentationOptions) {}
 
-  feed(event: unknown): void {
-    if (!isRecord(event)) return;
-    const streamError = modeledStreamError(event);
+  feed<T>(event: T): void {
+    const raw: unknown = event;
+    if (!isRecord(raw)) return;
+    const streamError = modeledStreamError(raw);
     if (streamError) {
       this.error = streamError.error;
       this.errorFields = streamError.fields;
       return;
     }
-    const chunk = isRecord(event.chunk) ? event.chunk : undefined;
+    const chunk = isRecord(raw.chunk) ? raw.chunk : undefined;
     const bytes = chunk?.bytes;
     if (bytes !== undefined) this.text += bytesToString(bytes) ?? "";
     const attribution = isRecord(chunk?.attribution) ? chunk.attribution : undefined;
     const citations = arrayValue<unknown>(attribution?.citations);
     if (citations) this.citationCount += citations.length;
 
-    const files = isRecord(event.files) ? arrayValue<unknown>(event.files.files) : undefined;
+    const files = isRecord(raw.files) ? arrayValue<JsonValue>(raw.files.files) : undefined;
     if (files) this.outputFiles.push(...files.map(sanitizeOutputFile));
 
-    const returnControl = isRecord(event.returnControl) ? event.returnControl : undefined;
+    const returnControl = isRecord(raw.returnControl) ? raw.returnControl : undefined;
     if (returnControl) this.returnControl = returnControl;
 
-    const traceEvent = isRecord(event.trace) ? event.trace : undefined;
+    const traceEvent = isRecord(raw.trace) ? raw.trace : undefined;
     if (traceEvent) {
       this.traceEventCount += 1;
       if (this.options.captureAgentTrace) this.traces.push(traceEvent);
@@ -119,11 +122,12 @@ export class AgentStreamState implements StreamState {
       this.errorFields,
     );
   }
-  private addUsage(value: unknown): void {
-    if (!isRecord(value)) return;
-    const inputTokens = numberValue(value.inputTokens);
-    const outputTokens = numberValue(value.outputTokens);
-    const totalTokens = numberValue(value.totalTokens);
+  private addUsage<T>(value: T): void {
+    const raw: unknown = value;
+    if (!isRecord(raw)) return;
+    const inputTokens = numberValue(raw.inputTokens);
+    const outputTokens = numberValue(raw.outputTokens);
+    const totalTokens = numberValue(raw.totalTokens);
     if (inputTokens !== undefined)
       this.usage.inputTokens = (this.usage.inputTokens ?? 0) + inputTokens;
     if (outputTokens !== undefined)
@@ -136,15 +140,16 @@ export class AgentStreamState implements StreamState {
   }
 }
 
-function sanitizeOutputFile(value: unknown): unknown {
-  if (!isRecord(value)) return value;
+function sanitizeOutputFile<T>(value: T): JsonValue {
+  const raw: unknown = value;
+  if (!isRecord(raw)) return undefined;
   return omitUndefined({
-    name: stringValue(value.name),
-    type: stringValue(value.type),
+    name: stringValue(raw.name),
+    type: stringValue(raw.type),
   });
 }
 
-export function retrieveRequestFields(input: Record<string, unknown>): SpanFields {
+export function retrieveRequestFields(input: JsonRecord): SpanFields {
   return omitUndefined({
     provider: PROVIDER,
     input: isRecord(input.retrievalQuery) ? input.retrievalQuery : undefined,
@@ -152,8 +157,9 @@ export function retrieveRequestFields(input: Record<string, unknown>): SpanField
   });
 }
 
-export function retrieveResponseFields(output: unknown): SpanFields {
-  const result = isRecord(output) ? output : {};
+export function retrieveResponseFields<T>(output: T): SpanFields {
+  const raw: unknown = output;
+  const result = isRecord(raw) ? raw : {};
   const retrievalResults = arrayValue<unknown>(result.retrievalResults);
   return mergeFields(awsMetadataFields(result.$metadata), {
     output: retrievalResults,
@@ -164,7 +170,7 @@ export function retrieveResponseFields(output: unknown): SpanFields {
   });
 }
 
-export function ragModel(input: Record<string, unknown>): string | undefined {
+export function ragModel(input: JsonRecord): string | undefined {
   const config = isRecord(input.retrieveAndGenerateConfiguration)
     ? input.retrieveAndGenerateConfiguration
     : undefined;
@@ -177,7 +183,7 @@ export function ragModel(input: Record<string, unknown>): string | undefined {
   return stringValue(kb?.modelArn) ?? stringValue(external?.modelArn);
 }
 
-export function ragRequestFields(input: Record<string, unknown>): SpanFields {
+export function ragRequestFields(input: JsonRecord): SpanFields {
   const text = isRecord(input.input) ? stringValue(input.input.text) : undefined;
   return omitUndefined({
     provider: PROVIDER,
@@ -187,8 +193,9 @@ export function ragRequestFields(input: Record<string, unknown>): SpanFields {
   });
 }
 
-export function ragResponseFields(output: unknown): SpanFields {
-  const result = isRecord(output) ? output : {};
+export function ragResponseFields<T>(output: T): SpanFields {
+  const raw: unknown = output;
+  const result = isRecord(raw) ? raw : {};
   const text = isRecord(result.output) ? stringValue(result.output.text) : undefined;
   return mergeFields(awsMetadataFields(result.$metadata), {
     output: text ? [{ role: "assistant", parts: [{ type: "text", content: text }] }] : undefined,
@@ -207,18 +214,19 @@ export class RagStreamState implements StreamState {
   private error: Error | undefined;
   private errorFields: SpanFields | undefined;
 
-  feed(event: unknown): void {
-    if (!isRecord(event)) return;
-    const streamError = modeledStreamError(event);
+  feed<T>(event: T): void {
+    const raw: unknown = event;
+    if (!isRecord(raw)) return;
+    const streamError = modeledStreamError(raw);
     if (streamError) {
       this.error = streamError.error;
       this.errorFields = streamError.fields;
       return;
     }
-    const output = isRecord(event.output) ? event.output : undefined;
+    const output = isRecord(raw.output) ? raw.output : undefined;
     this.text += stringValue(output?.text) ?? "";
-    if (event.citation) this.citationCount += 1;
-    const guardrail = isRecord(event.guardrail) ? event.guardrail : undefined;
+    if (raw.citation) this.citationCount += 1;
+    const guardrail = isRecord(raw.guardrail) ? raw.guardrail : undefined;
     this.guardrailAction = stringValue(guardrail?.action) ?? this.guardrailAction;
   }
 
@@ -239,7 +247,7 @@ export class RagStreamState implements StreamState {
   }
 }
 
-export function flowRequestFields(input: Record<string, unknown>): SpanFields {
+export function flowRequestFields(input: JsonRecord): SpanFields {
   return omitUndefined({
     provider: PROVIDER,
     input: input.inputs,
@@ -253,21 +261,22 @@ export class FlowStreamState implements StreamState {
   private error: Error | undefined;
   private errorFields: SpanFields | undefined;
 
-  feed(event: unknown): void {
-    if (!isRecord(event)) return;
-    const streamError = modeledStreamError(event);
+  feed<T>(event: T): void {
+    const raw: unknown = event;
+    if (!isRecord(raw)) return;
+    const streamError = modeledStreamError(raw);
     if (streamError) {
       this.error = streamError.error;
       this.errorFields = streamError.fields;
       return;
     }
-    const output = isRecord(event.flowOutputEvent) ? event.flowOutputEvent : undefined;
+    const output = isRecord(raw.flowOutputEvent) ? raw.flowOutputEvent : undefined;
     if (output) this.outputs.push(output.content);
-    const inputRequest = isRecord(event.flowMultiTurnInputRequestEvent)
-      ? event.flowMultiTurnInputRequestEvent
+    const inputRequest = isRecord(raw.flowMultiTurnInputRequestEvent)
+      ? raw.flowMultiTurnInputRequestEvent
       : undefined;
     if (inputRequest) this.outputs.push(inputRequest.content);
-    const completion = isRecord(event.flowCompletionEvent) ? event.flowCompletionEvent : undefined;
+    const completion = isRecord(raw.flowCompletionEvent) ? raw.flowCompletionEvent : undefined;
     this.completionReason = stringValue(completion?.completionReason) ?? this.completionReason;
   }
 
@@ -283,31 +292,33 @@ export class FlowStreamState implements StreamState {
   }
 }
 
-export function flowName(input: Record<string, unknown>): string {
+export function flowName(input: JsonRecord): string {
   return `invoke_flow ${stringValue(input.flowIdentifier) ?? "unknown"}`;
 }
 
-export function ragName(input: Record<string, unknown>): string {
+export function ragName(input: JsonRecord): string {
   return `retrieve_and_generate ${basenameArn(ragModel(input)) ?? "knowledge-base"}`;
 }
 
-function collectTrace(value: unknown, onUsage: (usage: unknown) => void): void {
-  if (!isRecord(value)) return;
-  const output = isRecord(value.modelInvocationOutput) ? value.modelInvocationOutput : undefined;
+function collectTrace<T>(value: T, onUsage: <TUsage>(usage: TUsage) => void): void {
+  const raw: unknown = value;
+  if (!isRecord(raw)) return;
+  const output = isRecord(raw.modelInvocationOutput) ? raw.modelInvocationOutput : undefined;
   if (output) {
     const metadata = isRecord(output.metadata) ? output.metadata : undefined;
     onUsage(metadata?.usage);
   }
-  for (const child of Object.values(value)) {
+  for (const child of Object.values(raw)) {
     if (isRecord(child)) collectTrace(child, onUsage);
     else if (Array.isArray(child)) child.forEach((item) => collectTrace(item, onUsage));
   }
 }
 
-function findKey(value: unknown, key: string): unknown {
-  if (!isRecord(value)) return undefined;
-  if (value[key] !== undefined) return value[key];
-  for (const child of Object.values(value)) {
+function findKey<T>(value: T, key: string): JsonValue {
+  const raw: unknown = value;
+  if (!isRecord(raw)) return undefined;
+  if (raw[key] !== undefined) return raw[key];
+  for (const child of Object.values(raw)) {
     const found = findKey(child, key);
     if (found !== undefined) return found;
   }

package/src/agents.ts

@@ -10,7 +10,7 @@ import {
 export type { BedrockInstrumentationOptions } from "./internal.ts";
 
 export interface WrappableBedrockAgentClient {
-  send: (...args: never[]) => unknown;
+  send: (...args: never[]) => any;
 }
 
 const WRAPPED = new WeakSet<object>();
@@ -33,7 +33,7 @@ export function wrapBedrockAgents<T extends WrappableBedrockAgentClient>(
 
 export function instrumentBedrockAgents(options: BedrockInstrumentationOptions = {}): void {
   if (ORIGINAL) return;
-  const proto = BedrockAgentRuntimeClient.prototype as unknown as { send: SendFn };
+  const proto = BedrockAgentRuntimeClient.prototype as { send: SendFn };
   const hadOwn = Object.prototype.hasOwnProperty.call(proto, "send");
   const original = proto.send;
   const instrumentedSend = createInstrumentedSend(original, options, AGENT_HANDLERS);

package/src/converse.ts

@@ -11,17 +11,19 @@ import {
   parseJson,
   PROVIDER,
   stringValue,
+  type JsonRecord,
+  type JsonValue,
   type StreamState,
 } from "./internal.ts";
 
-export type MessagePart = Record<string, unknown>;
+export type MessagePart = object;
 export interface NormalizedMessage {
   role: string | undefined;
   parts: MessagePart[];
   finish_reason?: string;
 }
 
-export function converseRequestFields(input: Record<string, unknown>): SpanFields {
+export function converseRequestFields(input: JsonRecord): SpanFields {
   const inference = isRecord(input.inferenceConfig) ? input.inferenceConfig : {};
   const guardrail = isRecord(input.guardrailConfig) ? input.guardrailConfig : undefined;
   return omitUndefined({
@@ -42,8 +44,8 @@ export function converseRequestFields(input: Record<string, unknown>): SpanField
   });
 }
 
-export function converseResponseFields(output: unknown): SpanFields {
-  const result = isRecord(output) ? output : {};
+export function converseResponseFields<T>(output: T): SpanFields {
+  const result: JsonRecord = isRecord(output) ? output : {};
   const usage = usageFromConverse(result.usage);
   const metrics = isRecord(result.metrics) ? result.metrics : undefined;
   const trace = isRecord(result.trace) ? result.trace : undefined;
@@ -64,7 +66,7 @@ export function converseResponseFields(output: unknown): SpanFields {
   });
 }
 
-export function usageFromConverse(value: unknown): SpanFields["usage"] | undefined {
+export function usageFromConverse<T>(value: T): SpanFields["usage"] | undefined {
   if (!isRecord(value)) return undefined;
   return omitUndefined({
     inputTokens: numberValue(value.inputTokens),
@@ -75,8 +77,8 @@ export function usageFromConverse(value: unknown): SpanFields["usage"] | undefin
   });
 }
 
-export function normalizeMessages(value: unknown): NormalizedMessage[] | undefined {
-  const messages = arrayValue<Record<string, unknown>>(value);
+export function normalizeMessages<T>(value: T): NormalizedMessage[] | undefined {
+  const messages = arrayValue<JsonRecord>(value);
   if (!messages) return undefined;
   return messages.map((message) => ({
     role: stringValue(message.role),
@@ -84,34 +86,33 @@ export function normalizeMessages(value: unknown): NormalizedMessage[] | undefin
   }));
 }
 
-export function normalizeOutputMessage(
-  value: unknown,
+export function normalizeOutputMessage<T>(
+  value: T,
   finishReason?: string,
 ): NormalizedMessage[] | undefined {
   if (!isRecord(value)) return undefined;
-  return [
-    omitUndefined({
-      role: stringValue(value.role),
-      parts: normalizeContentList(value.content),
-      finish_reason: finishReason,
-    }) as NormalizedMessage,
-  ];
+  const message = omitUndefined({
+    role: stringValue(value.role),
+    parts: normalizeContentList(value.content),
+    finish_reason: finishReason,
+  });
+  return [message];
 }
 
-export function normalizeSystem(value: unknown): MessagePart[] | undefined {
-  const blocks = arrayValue<Record<string, unknown>>(value);
+export function normalizeSystem<T>(value: T): MessagePart[] | undefined {
+  const blocks = arrayValue<JsonRecord>(value);
   if (!blocks) return undefined;
   const parts = blocks.flatMap((block) => normalizeContentBlock(block));
   return parts.length > 0 ? parts : undefined;
 }
 
-function normalizeContentList(value: unknown): MessagePart[] {
-  const blocks = arrayValue<Record<string, unknown>>(value);
+function normalizeContentList<T>(value: T): MessagePart[] {
+  const blocks = arrayValue<JsonRecord>(value);
   if (!blocks) return [];
   return blocks.flatMap((block) => normalizeContentBlock(block));
 }
 
-function normalizeContentBlock(block: Record<string, unknown>): MessagePart[] {
+function normalizeContentBlock(block: JsonRecord): MessagePart[] {
   const text = stringValue(block.text);
   if (text !== undefined)
     return [
@@ -178,11 +179,12 @@ function normalizeContentBlock(block: Record<string, unknown>): MessagePart[] {
 
   const citationsContent = isRecord(block.citationsContent) ? block.citationsContent : undefined;
   if (citationsContent) {
-    const contentBlocks = arrayValue<Record<string, unknown>>(citationsContent.content) ?? [];
-    const citations = arrayValue<unknown>(citationsContent.citations);
-    return contentBlocks.flatMap((contentBlock) =>
-      normalizeContentBlock(citations ? { ...contentBlock, citations } : contentBlock),
-    );
+    const contentBlocks = arrayValue<JsonRecord>(citationsContent.content) ?? [];
+    const citations = arrayValue<JsonValue>(citationsContent.citations);
+    return contentBlocks.flatMap((contentBlock) => {
+      if (!citations) return normalizeContentBlock(contentBlock);
+      return normalizeContentBlock({ ...contentBlock, citations });
+    });
   }
 
   const guardContent = isRecord(block.guardContent) ? block.guardContent : undefined;
@@ -199,14 +201,15 @@ function normalizeContentBlock(block: Record<string, unknown>): MessagePart[] {
   }
 }
 
-function simplifyToolResult(content: unknown): unknown {
-  const blocks = arrayValue<Record<string, unknown>>(content);
-  if (!blocks) return content;
-  if (blocks.length === 1 && typeof blocks[0]?.text === "string") return blocks[0].text;
-  return blocks.flatMap((block) => normalizeContentBlock(block));
+function simplifyToolResult<T>(content: T): JsonValue {
+  const blocks = arrayValue<JsonRecord>(content);
+  if (!blocks) return isJsonValue(content) ? content : undefined;
+  const text = stringValue(blocks[0]?.text);
+  if (blocks.length === 1 && text !== undefined) return text;
+  return blocks.flatMap((block) => normalizeContentBlock(block)) as JsonValue[];
 }
 
-function s3Uri(value: Record<string, unknown> | undefined): string | undefined {
+function s3Uri(value: JsonRecord | undefined): string | undefined {
   const bucket = stringValue(value?.bucket);
   const key = stringValue(value?.key);
   return bucket && key ? `s3://${bucket}/${key}` : undefined;
@@ -218,7 +221,7 @@ interface StreamBlock {
   toolUseId?: string;
   name?: string;
   input?: string;
-  content?: Record<string, unknown>;
+  content?: JsonRecord;
   citations?: unknown[];
 }
 
@@ -227,12 +230,12 @@ export class ConverseStreamState implements StreamState {
   private readonly blocks = new Map<number, StreamBlock>();
   private stopReason: string | undefined;
   private usage: SpanFields["usage"];
-  private metadata: Record<string, unknown> = {};
+  private metadata: JsonRecord = {};
   private responseModel: string | undefined;
   private error: Error | undefined;
   private errorFields: SpanFields | undefined;
 
-  feed(event: unknown): void {
+  feed<T>(event: T): void {
     if (!isRecord(event)) return;
     const streamError = modeledStreamError(event);
     if (streamError) {
@@ -281,14 +284,16 @@ export class ConverseStreamState implements StreamState {
       const index = numberValue(deltaEvent.contentBlockIndex) ?? 0;
       const delta = isRecord(deltaEvent.delta) ? deltaEvent.delta : {};
       const block = this.blocks.get(index) ?? { kind: "text", text: "" };
-      if (typeof delta.text === "string") {
+      const text = stringValue(delta.text);
+      if (text !== undefined) {
         block.kind = "text";
-        block.text = `${block.text ?? ""}${delta.text}`;
+        block.text = `${block.text ?? ""}${text}`;
       }
       const toolUse = isRecord(delta.toolUse) ? delta.toolUse : undefined;
-      if (toolUse && typeof toolUse.input === "string") {
+      const input = stringValue(toolUse?.input);
+      if (toolUse && input !== undefined) {
         block.kind = "tool";
-        block.input = `${block.input ?? ""}${toolUse.input}`;
+        block.input = `${block.input ?? ""}${input}`;
       }
       const reasoning = isRecord(delta.reasoningContent) ? delta.reasoningContent : undefined;
       if (reasoning) {
@@ -298,15 +303,15 @@ export class ConverseStreamState implements StreamState {
       const image = isRecord(delta.image) ? delta.image : undefined;
       if (image) {
         const content = isRecord(block.content) ? block.content : {};
-        const current = isRecord(content.image) ? content.image : {};
+        const current: JsonRecord = isRecord(content.image) ? content.image : {};
         block.kind = "content";
         block.content = { image: { ...current, ...image } };
       }
-      const toolResult = arrayValue<Record<string, unknown>>(delta.toolResult);
+      const toolResult = arrayValue<JsonRecord>(delta.toolResult);
       if (toolResult) {
         const content = isRecord(block.content) ? block.content : {};
-        const current = isRecord(content.toolResult) ? content.toolResult : {};
-        const currentContent = arrayValue<unknown>(current.content) ?? [];
+        const current: JsonRecord = isRecord(content.toolResult) ? content.toolResult : {};
+        const currentContent = arrayValue<JsonValue>(current.content) ?? [];
         block.kind = "content";
         block.content = {
           toolResult: {
@@ -370,7 +375,18 @@ export class ConverseStreamState implements StreamState {
   }
 }
 
-function blockToPart(block: StreamBlock): MessagePart {
+function isJsonValue<T>(value: T): value is T & JsonValue {
+  return (
+    value == null ||
+    String(value) === value ||
+    Number(value) === value ||
+    Boolean(value) === value ||
+    Array.isArray(value) ||
+    isRecord(value)
+  );
+}
+
+function blockToPart(block: StreamBlock) {
   if (block.kind === "tool") {
     return omitUndefined({
       type: "tool_call",

package/src/index.ts

@@ -10,7 +10,7 @@ import {
 export type { BedrockInstrumentationOptions } from "./internal.ts";
 
 export interface WrappableBedrockClient {
-  send: (...args: never[]) => unknown;
+  send: (...args: never[]) => any;
 }
 
 const WRAPPED = new WeakSet<object>();
@@ -33,7 +33,7 @@ export function wrapBedrock<T extends WrappableBedrockClient>(
 
 export function instrumentBedrock(options: BedrockInstrumentationOptions = {}): void {
   if (ORIGINAL) return;
-  const proto = BedrockRuntimeClient.prototype as unknown as { send: SendFn };
+  const proto = BedrockRuntimeClient.prototype as { send: SendFn };
   const hadOwn = Object.prototype.hasOwnProperty.call(proto, "send");
   const original = proto.send;
   const instrumentedSend = createInstrumentedSend(original, options, RUNTIME_HANDLERS);

package/src/internal-handlers.ts

@@ -1,4 +1,4 @@
-import type { OpHandler } from "./internal.ts";
+import type { JsonRecord, JsonValue, OpHandler } from "./internal.ts";
 import {
   AgentStreamState,
   FlowStreamState,
@@ -32,7 +32,7 @@ import {
   type BedrockInstrumentationOptions,
 } from "./internal.ts";
 
-export const RUNTIME_HANDLERS: Record<string, OpHandler> = {
+export const RUNTIME_HANDLERS = {
   ConverseCommand: {
     spanName: (input) => `chat ${stringValue(input.modelId) ?? "unknown"}`,
     spanType: () => "generation",
@@ -108,7 +108,7 @@ export const RUNTIME_HANDLERS: Record<string, OpHandler> = {
       }),
     }),
     onResult(result, span) {
-      const response = result as Record<string, unknown>;
+      const response = result as JsonRecord;
       endSpan(
         span,
         mergeFields(awsMetadataFields(response.$metadata), {
@@ -122,9 +122,9 @@ export const RUNTIME_HANDLERS: Record<string, OpHandler> = {
       return result;
     },
   },
-};
+} satisfies Record<string, OpHandler>;
 
-export const AGENT_HANDLERS: Record<string, OpHandler> = {
+export const AGENT_HANDLERS = {
   InvokeAgentCommand: {
     spanName: (input) => `invoke_agent ${stringValue(input.agentId) ?? "unknown"}`,
     spanType: () => "agent",
@@ -196,19 +196,19 @@ export const AGENT_HANDLERS: Record<string, OpHandler> = {
       };
     },
   },
-};
+} satisfies Record<string, OpHandler>;
 
 function streamAgentResult(key: "completion") {
   return (
-    result: unknown,
+    result: JsonValue,
     span: Parameters<OpHandler["onResult"]>[1],
     t0: number,
-    _input: Record<string, unknown>,
+    _input: JsonRecord,
     options: BedrockInstrumentationOptions,
-  ): unknown => {
-    const response = result as {
-      completion?: unknown;
-      $metadata?: unknown;
+  ) => {
+    const response = result as JsonRecord & {
+      completion?: JsonValue;
+      $metadata?: JsonValue;
       sessionId?: string;
       memoryId?: string;
     };