llmz 0.1.4 → 0.2.0

package/dist/{execute-ZENBZBWL.js → execute-HIRXXR6X.js}

@@ -1,6 +1,8 @@
 import {
-  runAsyncFunction
-} from "./chunk-EFEF2AL7.js";
+  hasTopLevelReturn,
+  runAsyncFunction,
+  warmupVM
+} from "./chunk-DH35WO5G.js";
 import {
   Context,
   ErrorExecutionResult,
@@ -9,8 +11,11 @@ import {
   SuccessExecutionResult,
   _CustomModelClient,
   parseExit
-} from "./chunk-TIZWOIDF.js";
-import "./chunk-2O3UTPVH.js";
+} from "./chunk-33AB6SVX.js";
+import {
+  StreamingMessageParser,
+  toParsedAssistantResponse
+} from "./chunk-E5NHK3FS.js";
 import "./chunk-IRCB4FRM.js";
 import "./chunk-TOGYDN6U.js";
 import {
@@ -27,11 +32,13 @@ import {
   cleanStackTrace
 } from "./chunk-YEAWWJSJ.js";
 import "./chunk-XVVYNOTX.js";
-import "./chunk-3T5IU7OU.js";
-import "./chunk-XTEHT2TK.js";
+import "./chunk-KWQTYXAT.js";
+import {
+  createJsxComponent
+} from "./chunk-XTEHT2TK.js";
 import {
   truncateWrappedContent
-} from "./chunk-L74UZONX.js";
+} from "./chunk-HM5H556L.js";
 import {
   getErrorMessage,
   init,
@@ -39,9 +46,7 @@ import {
 } from "./chunk-YX27XM67.js";
 import {
   clamp_default,
-  isEqual_default,
-  isPlainObject_default,
-  omit_default
+  isEqual_default
 } from "./chunk-34GJUS6U.js";
 
 // src/runtime/execute.ts
@@ -83,26 +88,78 @@ var RESPONSE_LENGTH_BUFFER = {
   MAX_TOKENS: 16e3,
   PERCENTAGE: 0.1
 };
+var STREAM_INACTIVITY_TIMEOUT = 6e4;
 var getModelOutputLimit = (inputLength) => clamp_default(
   RESPONSE_LENGTH_BUFFER.PERCENTAGE * inputLength,
   RESPONSE_LENGTH_BUFFER.MIN_TOKENS,
   RESPONSE_LENGTH_BUFFER.MAX_TOKENS
 );
-var generateCode = async ({ iteration, ctx, cognitive, controller, metadata }) => {
-  var _a, _b, _c;
+var LeadingFenceFilter = class {
+  _buffer = "";
+  _done = false;
+  push(text) {
+    if (this._done) {
+      return text;
+    }
+    this._buffer += text;
+    const newline = this._buffer.indexOf("\n");
+    if (newline === -1) {
+      return "";
+    }
+    this._done = true;
+    const firstLine = this._buffer.slice(0, newline);
+    const rest = this._buffer.slice(newline + 1);
+    this._buffer = "";
+    return firstLine.trim().startsWith("```") ? rest : `${firstLine}
+${rest}`;
+  }
+  flush() {
+    if (this._done) {
+      return "";
+    }
+    this._done = true;
+    const buffered = this._buffer;
+    this._buffer = "";
+    return buffered.trim().startsWith("```") ? "" : buffered;
+  }
+};
+var generateCode = async ({
+  iteration,
+  ctx,
+  cognitive,
+  controller,
+  metadata,
+  onSend,
+  onSendDelta,
+  onRunStart,
+  onRunComplete
+}) => {
+  var _a, _b, _c, _d;
   const startedAt = Date.now();
   const traces = iteration.traces;
   const modelRef = Array.isArray(iteration.model) ? iteration.model[0] : iteration.model;
   const model = await cognitive.getModelDetails(modelRef).catch((thrown) => {
     throw new CognitiveError(`Failed to fetch model details for model "${modelRef}": ${getErrorMessage(thrown)}`);
   });
-  const modelLimit = Math.max(model.input.maxTokens, 8e3);
+  let modelLimit = Math.max(model.input.maxTokens, 8e3);
+  if (ctx.maxTokens) {
+    modelLimit = Math.min(ctx.maxTokens, modelLimit);
+  }
   const responseLengthBuffer = getModelOutputLimit(modelLimit);
-  const messages = truncateWrappedContent({
-    messages: iteration.messages,
-    tokenLimit: modelLimit - responseLengthBuffer,
-    throwOnFailure: true
-  }).filter((x) => typeof x.content !== "string" || x.content.trim().length > 0);
+  if (iteration.tokens) {
+    iteration.tokens.limit = modelLimit;
+  }
+  let messages;
+  try {
+    messages = truncateWrappedContent({
+      messages: iteration.messages,
+      tokenLimit: modelLimit - responseLengthBuffer,
+      throwOnFailure: true
+    }).filter((x) => typeof x.content !== "string" || x.content.trim().length > 0);
+  } catch (thrown) {
+    const cap = ctx.maxTokens ? ` (context window capped at ${modelLimit} tokens by options.maxTokens \u2014 consider raising or removing it)` : ` (model context window: ${modelLimit} tokens)`;
+    throw new CognitiveError(`The prompt does not fit in the context window${cap}: ${getErrorMessage(thrown)}`);
+  }
   iteration.messages = messages;
   traces.push({
     type: "llm_call_started",
@@ -110,7 +167,7 @@ var generateCode = async ({ iteration, ctx, cognitive, controller, metadata }) =
     ended_at: startedAt,
     model: model.ref
   });
-  const output = await cognitive.generateContent({
+  const input = {
     signal: controller.signal,
     systemPrompt: (_a = messages.find((x) => x.role === "system")) == null ? void 0 : _a.content,
     model: iteration.model,
@@ -120,15 +177,123 @@ var generateCode = async ({ iteration, ctx, cognitive, controller, metadata }) =
     messages: messages.filter((x) => x.role !== "system"),
     stopSequences: ctx.version.getStopTokens(),
     meta: metadata ? { metadata } : void 0
-  }).catch((thrown) => {
-    throw new CognitiveError(`LLM generation failed: ${getErrorMessage(thrown)}`);
-  });
-  const out = typeof ((_c = (_b = output.output.choices) == null ? void 0 : _b[0]) == null ? void 0 : _c.content) === "string" ? output.output.choices[0].content : null;
-  if (!out) {
-    throw new CognitiveError("LLM did not return any text output");
+  };
+  let output;
+  let raw;
+  let assistantResponse;
+  let timeToFirstToken;
+  let timeToLastToken;
+  const liveItems = /* @__PURE__ */ new Map();
+  const liveContent = /* @__PURE__ */ new Map();
+  let codeGenerationTraced = false;
+  let runCompleted = false;
+  const dispatchSends = async (events) => {
+    for (const event of events) {
+      if (event.type === "item-start") {
+        liveItems.set(event.item.id, event.item);
+        if (event.item.kind === "run" && !codeGenerationTraced) {
+          codeGenerationTraced = true;
+          traces.push({ type: "code_generation_started", started_at: Date.now() });
+          onRunStart == null ? void 0 : onRunStart();
+        }
+      } else if (event.type === "body-delta" && onSendDelta) {
+        const item = liveItems.get(event.itemId);
+        if ((item == null ? void 0 : item.kind) !== "send") {
+          continue;
+        }
+        const content = (liveContent.get(item.id) ?? "") + event.delta;
+        liveContent.set(item.id, content);
+        try {
+          await onSendDelta({
+            id: `${iteration.id}:${item.id}`,
+            component: item.name,
+            props: item.props,
+            delta: event.delta,
+            content
+          });
+        } catch (err) {
+        }
+      } else if (event.type === "item-complete") {
+        if (event.item.kind === "send" && onSend) {
+          await onSend({ name: event.item.name, props: event.item.props, body: event.item.body });
+        } else if (event.item.kind === "run" && event.item.status === "complete" && !runCompleted) {
+          runCompleted = true;
+          onRunComplete == null ? void 0 : onRunComplete((event.item.body ?? "").trim());
+        }
+      }
+    }
+  };
+  if (typeof cognitive.generateContentStream === "function") {
+    const parser = new StreamingMessageParser();
+    const fence = new LeadingFenceFilter();
+    const streamController = createJoinedAbortController([controller.signal]);
+    const requestedAt = Date.now();
+    const stream = cognitive.generateContentStream({
+      ...input,
+      signal: streamController.signal,
+      // Passed through to the cognitive v2 request: fall back to the next
+      // model/provider when the first token takes too long
+      ...ctx.maxTimeToFirstToken ? { options: { maxTimeToFirstToken: ctx.maxTimeToFirstToken } } : {}
+    });
+    const inactivityTimeout = Math.max(STREAM_INACTIVITY_TIMEOUT, ctx.maxTimeToFirstToken ?? 0);
+    const nextChunk = async () => {
+      let timer;
+      const stalled = new Promise((_, reject) => {
+        timer = setTimeout(() => {
+          streamController.abort("LLM stream stalled");
+          reject(new Error(`LLM stream stalled: no data received for ${inactivityTimeout}ms`));
+        }, inactivityTimeout);
+      });
+      try {
+        return await Promise.race([stream.next(), stalled]);
+      } finally {
+        clearTimeout(timer);
+      }
+    };
+    raw = "";
+    while (true) {
+      let chunk;
+      try {
+        chunk = await nextChunk();
+      } catch (thrown) {
+        throw new CognitiveError(`LLM generation failed: ${getErrorMessage(thrown)}`);
+      }
+      if (chunk.done) {
+        output = chunk.value;
+        break;
+      }
+      const delta = (_b = chunk.value) == null ? void 0 : _b.output;
+      if (!delta) {
+        continue;
+      }
+      timeToLastToken = Date.now() - requestedAt;
+      timeToFirstToken ??= timeToLastToken;
+      raw += delta;
+      await dispatchSends(parser.push(fence.push(delta)));
+    }
+    const remaining = fence.flush();
+    if (remaining) {
+      await dispatchSends(parser.push(remaining));
+    }
+    await dispatchSends(parser.finish());
+    assistantResponse = toParsedAssistantResponse(parser.items, raw);
+  } else {
+    output = await cognitive.generateContent(input).catch((thrown) => {
+      throw new CognitiveError(`LLM generation failed: ${getErrorMessage(thrown)}`);
+    });
+    const out = typeof ((_d = (_c = output.output.choices) == null ? void 0 : _c[0]) == null ? void 0 : _d.content) === "string" ? output.output.choices[0].content : null;
+    if (!out) {
+      throw new CognitiveError("LLM did not return any text output");
+    }
+    raw = out;
+    assistantResponse = ctx.version.parseAssistantResponse(out);
+    for (const send of assistantResponse.sends) {
+      await (onSend == null ? void 0 : onSend(send));
+    }
   }
-  const assistantResponse = ctx.version.parseAssistantResponse(out);
-  iteration.code = assistantResponse.code.trim();
+  iteration.code = assistantResponse.code;
+  iteration.sends = assistantResponse.sends;
+  iteration.next = assistantResponse.next;
   iteration.llm = {
     cached: output.meta.cached || false,
     ended_at: Date.now(),
@@ -138,14 +303,21 @@ var generateCode = async ({ iteration, ctx, cognitive, controller, metadata }) =
     spend: output.meta.cost.input + output.meta.cost.output,
     output: assistantResponse.raw,
     model: `${output.meta.model.integration}:${output.meta.model.model}`,
+    time_to_first_token: timeToFirstToken,
+    time_to_last_token: timeToLastToken,
     usage: output.output.usage
   };
+  if (iteration.tokens) {
+    iteration.tokens.input = output.meta.tokens.input;
+    iteration.tokens.output = output.meta.tokens.output;
+    iteration.tokens.total = output.meta.tokens.input + output.meta.tokens.output;
+  }
   traces.push({
     type: "llm_call_success",
     started_at: startedAt,
     ended_at: iteration.llm.ended_at,
     model: model.ref,
-    code: iteration.code
+    code: iteration.code ?? ""
   });
 };
 
@@ -222,35 +394,50 @@ var interpretVMResult = async ({
     });
     return;
   }
-  let returnValue = result.success && result.return_value ? result.return_value : null;
-  const returnAction = returnValue == null ? void 0 : returnValue.action;
-  if (returnAction === "think") {
-    const variables = omit_default(returnValue ?? {}, "action");
-    if (isPlainObject_default(variables) && Object.keys(variables).length > 0) {
-      iteration.end({
-        type: "thinking_requested",
-        thinking_requested: {
-          variables,
-          reason: "Thinking requested"
-        }
-      });
-      return;
+  const returnValue = result.return_value;
+  const codeRequestsInspection = hasTopLevelReturn(iteration.code ?? "");
+  if (iteration.next && !codeRequestsInspection) {
+    await applyNextExit({ iteration, controller, onExit });
+    return;
+  }
+  const localVariables = Object.fromEntries(
+    Object.entries(result.variables ?? {}).filter(([, value]) => value !== "[[non-primitive]]")
+  );
+  iteration.end({
+    type: "thinking_requested",
+    thinking_requested: {
+      reason: iteration.next ? `Code execution completed. The \u25A0next=${iteration.next.name} exit was NOT applied because your code returns a value \u2014 inspect it below, then end your next response with \u25A0next.` : "Code execution completed",
+      variables: returnValue !== void 0 ? returnValue : Object.keys(localVariables).length ? localVariables : iteration.variables
     }
-    iteration.end({
-      type: "thinking_requested",
-      thinking_requested: {
-        reason: "Thinking requested",
-        variables: iteration.variables
+  });
+};
+var applyNextExit = async ({
+  iteration,
+  controller,
+  onExit
+}) => {
+  const next = iteration.next;
+  if (!next) {
+    throw new Error("applyNextExit called without a \u25A0next block on the iteration");
+  }
+  const attempt = (props) => parseExit(Object.keys(props).length ? { action: next.name, value: props } : { action: next.name }, iteration.exits);
+  let parsedExit = attempt(next.props);
+  const keys = Object.keys(next.props);
+  if (!parsedExit.success && keys.length === 1 && ["props", "value", "data"].includes(keys[0])) {
+    const inner = next.props[keys[0]];
+    if (inner && typeof inner === "object" && !Array.isArray(inner)) {
+      const unwrapped = attempt(inner);
+      if (unwrapped.success) {
+        parsedExit = unwrapped;
       }
-    });
-    return;
+    }
   }
-  const parsedExit = parseExit(returnValue, iteration.exits);
+  const returnValue = { action: next.name, value: parsedExit.success ? parsedExit.value : next.props };
   if (!parsedExit.success) {
     iteration.end({
       type: "exit_error",
       exit_error: {
-        exit: (returnValue == null ? void 0 : returnValue.action) ?? "n/a",
+        exit: next.name,
         message: parsedExit.error,
         return_value: returnValue
       }
@@ -258,12 +445,14 @@ var interpretVMResult = async ({
     return;
   }
   const returnExit = parsedExit.exit;
-  returnValue = { action: returnExit.name, value: parsedExit.value };
   try {
-    await (onExit == null ? void 0 : onExit({
-      exit: returnExit,
-      result: returnValue == null ? void 0 : returnValue.value
-    }));
+    await (onExit == null ? void 0 : onExit(
+      {
+        exit: returnExit,
+        result: parsedExit.value
+      },
+      controller
+    ));
   } catch (err) {
     iteration.end({
       type: "exit_error",
@@ -279,10 +468,9 @@ var interpretVMResult = async ({
     type: "exit_success",
     exit_success: {
       exit_name: returnExit.name,
-      return_value: returnValue == null ? void 0 : returnValue.value
+      return_value: parsedExit.value
     }
   });
-  return;
 };
 
 // src/runtime/utils.ts
@@ -535,7 +723,7 @@ var executeContext = async (props) => {
   return result;
 };
 var executeContextInternal = async (props) => {
-  var _a, _b;
+  var _a, _b, _c, _d;
   const controller = createJoinedAbortController([props.signal]);
   const { onIterationStart, onIterationEnd, onTrace, onExit, onBeforeExecution, onAfterTool, onBeforeTool } = props;
   const client = props.client ?? new Client();
@@ -547,6 +735,8 @@ var executeContextInternal = async (props) => {
     tools: props.tools,
     loop: (_a = props.options) == null ? void 0 : _a.loop,
     timeout: (_b = props.options) == null ? void 0 : _b.timeout,
+    maxTokens: (_c = props.options) == null ? void 0 : _c.maxTokens,
+    maxTimeToFirstToken: (_d = props.options) == null ? void 0 : _d.maxTimeToFirstToken,
     exits: props.exits,
     snapshot: props.snapshot,
     model: props.model,
@@ -561,7 +751,7 @@ var executeContextInternal = async (props) => {
       const iteration = await ctx.nextIteration();
       const unsubscribeTrace = iteration.traces.onPush((traces) => {
         for (const trace of traces) {
-          onTrace == null ? void 0 : onTrace({ trace, iteration: ctx.iterations.length });
+          onTrace == null ? void 0 : onTrace({ trace, iteration: ctx.iterations.length, controller });
         }
       });
       try {
@@ -773,7 +963,58 @@ var executeIteration = async ({
   onAfterTool,
   metadata
 }) => {
-  await generateCode({ iteration, ctx, cognitive, controller, metadata });
+  var _a, _b;
+  const runCode = (code) => {
+    const vmContext = buildVMContext({ ctx, iteration, controller, onBeforeTool, onAfterTool });
+    return runAsyncFunction(vmContext, code, iteration.traces, controller.signal, ctx.timeout).catch((err) => {
+      return {
+        success: false,
+        error: err instanceof Error ? err : new Error(getErrorMessage(err)),
+        lines_executed: [],
+        traces: [],
+        variables: {}
+      };
+    });
+  };
+  const canExecuteEarly = typeof onBeforeExecution !== "function";
+  let earlyExecution;
+  await generateCode({
+    iteration,
+    ctx,
+    cognitive,
+    controller,
+    metadata,
+    // Pre-warm the VM while the model is still writing the ■run block
+    onRunStart: () => warmupVM(),
+    onRunComplete: canExecuteEarly ? (code) => {
+      if (!code.length || controller.signal.aborted || earlyExecution) {
+        return;
+      }
+      iteration.code = code;
+      earlyExecution = { code, started_at: Date.now(), promise: runCode(code) };
+    } : void 0,
+    onSend: async (send) => {
+      if (!ctx.chat) {
+        return;
+      }
+      const sendStartedAt = Date.now();
+      const component = createJsxComponent({
+        type: send.name,
+        props: send.props,
+        children: send.body ? [send.body] : []
+      });
+      try {
+        await ctx.chat.handler(component);
+      } catch (err) {
+        throw new Error(`Error while sending message (\u25A0send=${send.name}): ${getErrorMessage(err)}`);
+      }
+      iteration.traces.push({ type: "yield", value: component, started_at: sendStartedAt, ended_at: Date.now() });
+    },
+    onSendDelta: ((_a = ctx.chat) == null ? void 0 : _a.onMessageDelta) ? (delta) => ctx.chat.onMessageDelta(delta) : void 0
+  });
+  if (earlyExecution) {
+    iteration.code = earlyExecution.code;
+  }
   if (typeof onBeforeExecution === "function") {
     try {
       const hookRes = await onBeforeExecution(iteration, controller);
@@ -804,7 +1045,6 @@ var executeIteration = async ({
     }
   }
   const traces = iteration.traces;
-  const vmContext = buildVMContext({ ctx, iteration, controller, onBeforeTool, onAfterTool });
   if (controller.signal.aborted) {
     traces.push({
       type: "abort_signal",
@@ -819,22 +1059,26 @@ var executeIteration = async ({
     });
     return;
   }
-  const startedAt = Date.now();
-  const result = await runAsyncFunction(
-    vmContext,
-    iteration.code ?? "",
-    traces,
-    controller.signal,
-    ctx.timeout
-  ).catch((err) => {
-    return {
-      success: false,
-      error: err instanceof Error ? err : new Error(getErrorMessage(err)),
-      lines_executed: [],
-      traces: [],
-      variables: {}
-    };
-  });
+  if (!iteration.code) {
+    if (iteration.next) {
+      await applyNextExit({ iteration, controller, onExit });
+      return;
+    }
+    if (ctx.chat && ((_b = iteration.sends) == null ? void 0 : _b.length)) {
+      iteration.next = { name: "listen", props: {} };
+      await applyNextExit({ iteration, controller, onExit });
+      return;
+    }
+    iteration.end({
+      type: "invalid_code_error",
+      invalid_code_error: {
+        message: "The response did not include a \u25A0run block or a \u25A0next exit. Reply using \u25A0 blocks and end your response with \u25A0run or \u25A0next=<exit>."
+      }
+    });
+    return;
+  }
+  const startedAt = (earlyExecution == null ? void 0 : earlyExecution.started_at) ?? Date.now();
+  const result = earlyExecution ? await earlyExecution.promise : await runCode(iteration.code ?? "");
   await interpretVMResult({
     iteration,
     result,