@neutrome/lilsdk-ts 0.3.1 → 0.3.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@neutrome/lilsdk-ts",
-  "version": "0.3.1",
+  "version": "0.3.2",
   "type": "module",
   "exports": {
     ".": "./src/index.ts",

package/src/managed/index.ts

@@ -5,11 +5,14 @@ export {
   createTwoPassExecutor,
   INTERNAL_DRAFT_CALL_ID,
   INTERNAL_DRAFT_TOOL_NAME,
-  invokeModel,
-  streamModel,
+  invokeExecutor,
+  streamExecutor,
 } from "./twoPassExecutor.ts";
 
 export type {
   InternalDraft,
+  ModelExecutor,
   TwoPassExecutorOptions,
+  TwoPassSettings,
+  TwoPassSettingsResolver,
 } from "./twoPassExecutor.ts";

package/src/managed/two-pass-request.ts

@@ -0,0 +1,164 @@
+import {
+  appendInstructions,
+  contentText,
+  insertBefore,
+  messageText,
+  messages,
+  removeRange,
+  type Program,
+} from "@neutrome/lil-engine";
+import { Opcode } from "@neutrome/lil-engine";
+import {
+  appendToolInteraction,
+  makeMessage,
+  prependSystemPrompt,
+} from "../synthetic/index.ts";
+import type { ExecutorContext } from "../types.ts";
+
+export const INTERNAL_DRAFT_TOOL_NAME = "knowledge";
+export const INTERNAL_DRAFT_CALL_ID = "knowledge_0";
+
+export type InternalDraft =
+  | string
+  | readonly string[]
+  | readonly { text: string }[];
+
+export type TwoPassSettings = {
+  reasoningLevel?: string;
+  systemPrompt?: string;
+};
+
+export type TwoPassSettingsResolver = (
+  request: Program,
+  ctx: ExecutorContext,
+) => TwoPassSettings | null | Promise<TwoPassSettings | null>;
+
+export type TwoPassRequestOptions = {
+  maxTotalContextLength?: number;
+  buildFinalRequest?: (request: Program, draft: string) => Program;
+};
+
+export function appendInternalDraft(
+  request: Program,
+  draft: InternalDraft,
+  options: { callId?: string } = {},
+): Program {
+  const callId = options.callId ?? INTERNAL_DRAFT_CALL_ID;
+  const drafts = normalizeDrafts(draft);
+  if (drafts.length === 0) return request;
+  return drafts.reduce(
+    (program, text, index) =>
+      appendToolInteraction(program, {
+        callId: drafts.length === 1 ? callId : `${callId}_${index}`,
+        name: INTERNAL_DRAFT_TOOL_NAME,
+        args: {},
+        result: text,
+      }),
+    request,
+  );
+}
+
+export function buildDraftRequest(
+  request: Program,
+  settings: TwoPassSettings,
+): Program {
+  const withPrompt = settings.systemPrompt
+    ? replaceSystemPrompt(request, settings.systemPrompt)
+    : request;
+  return applySettings(withPrompt, settings);
+}
+
+export function buildFinalRequest(
+  options: TwoPassRequestOptions,
+  request: Program,
+  settings: TwoPassSettings,
+  draft: string,
+): Program {
+  const limitedDraft = limitDraft(options, request, settings, draft);
+  const withDraft = options.buildFinalRequest
+    ? options.buildFinalRequest(request, limitedDraft)
+    : appendInternalDraft(request, limitedDraft);
+  const withPrompt = settings.systemPrompt
+    ? prependSystemPrompt(withDraft, settings.systemPrompt)
+    : withDraft;
+  return applySettings(withPrompt, settings);
+}
+
+export function validateMaxContextLength(value: number | undefined): void {
+  if (value !== undefined && (!Number.isSafeInteger(value) || value < 0)) {
+    throw new Error("maxTotalContextLength must be a non-negative integer");
+  }
+}
+
+function replaceSystemPrompt(request: Program, systemPrompt: string): Program {
+  const systemSpans = messages(request).filter(
+    (span) => span.role === "system",
+  );
+  const previousSystemPrompt = systemSpans
+    .map((span) => messageText(request, span).trim())
+    .filter(Boolean)
+    .join("\n\n");
+  let draft = request;
+  for (const span of systemSpans.slice().reverse()) {
+    draft = removeRange(draft, span.start, span.end);
+  }
+  draft = prependSystemPrompt(draft, systemPrompt);
+  if (!previousSystemPrompt) return draft;
+  const firstNonSystem = messages(draft).find((span) => span.role !== "system");
+  return insertBefore(
+    draft,
+    firstNonSystem?.start ?? draft.code.length,
+    makeMessage("user", previousSystemPrompt).code,
+  );
+}
+
+function applySettings(request: Program, settings: TwoPassSettings): Program {
+  if (!settings.reasoningLevel) return request;
+  return appendInstructions(
+    {
+      ...request,
+      code: request.code.filter(
+        (instruction) => instruction.opcode !== Opcode.SET_REASON_EFFORT,
+      ),
+    },
+    [
+      {
+        opcode: Opcode.SET_REASON_EFFORT,
+        value: { kind: "string", value: settings.reasoningLevel },
+      },
+    ],
+  );
+}
+
+function limitDraft(
+  options: TwoPassRequestOptions,
+  request: Program,
+  settings: TwoPassSettings,
+  draft: string,
+): string {
+  if (options.maxTotalContextLength === undefined) return draft;
+  const base = buildFinalRequestWithoutDraft(options, request, settings);
+  const available = options.maxTotalContextLength - contentText(base).length;
+  return available > 0 ? draft.slice(0, available) : "";
+}
+
+function buildFinalRequestWithoutDraft(
+  options: TwoPassRequestOptions,
+  request: Program,
+  settings: TwoPassSettings,
+): Program {
+  const withDraft = options.buildFinalRequest
+    ? options.buildFinalRequest(request, "")
+    : request;
+  const withPrompt = settings.systemPrompt
+    ? prependSystemPrompt(withDraft, settings.systemPrompt)
+    : withDraft;
+  return applySettings(withPrompt, settings);
+}
+
+function normalizeDrafts(draft: InternalDraft): string[] {
+  const drafts = Array.isArray(draft)
+    ? draft.map((item) => (typeof item === "string" ? item : item.text))
+    : [draft];
+  return drafts.map((text) => text.trim()).filter((text) => text.length > 0);
+}

package/src/managed/twoPassExecutor.ts

@@ -4,79 +4,69 @@ import {
   deltaText,
   finishReason,
   hasToolDelta,
-  insertBefore,
   lastMessageRole,
-  messageText,
-  messages,
-  removeRange,
   setModel,
   type Program,
 } from "@neutrome/lil-engine";
 import { streamReasoningDelta } from "../output.ts";
-import {
-  appendToolInteraction,
-  makeMessage,
-  prependSystemPrompt,
-} from "../synthetic/index.ts";
 import type { Executor, ExecutorContext, InvokeOptions } from "../types.ts";
+import {
+  appendInternalDraft,
+  buildDraftRequest,
+  buildFinalRequest,
+  INTERNAL_DRAFT_CALL_ID,
+  INTERNAL_DRAFT_TOOL_NAME,
+  validateMaxContextLength,
+  type InternalDraft,
+  type TwoPassRequestOptions,
+  type TwoPassSettings,
+  type TwoPassSettingsResolver,
+} from "./two-pass-request.ts";
+
+export type {
+  InternalDraft,
+  TwoPassSettings,
+  TwoPassSettingsResolver,
+} from "./two-pass-request.ts";
+export {
+  appendInternalDraft,
+  INTERNAL_DRAFT_CALL_ID,
+  INTERNAL_DRAFT_TOOL_NAME,
+};
 
-export const INTERNAL_DRAFT_TOOL_NAME = "knowledge";
-export const INTERNAL_DRAFT_CALL_ID = "knowledge_0";
-
-export type InternalDraft =
-  | string
-  | readonly string[]
-  | readonly { text: string }[];
+export type ModelExecutor = string | Executor;
 
-export type TwoPassExecutorOptions = {
-  draftModel: string;
-  finalModel: string;
-  draftSystemPrompt?: string;
-  finalSystemPrompt?: string;
+export type TwoPassExecutorOptions = TwoPassRequestOptions & {
+  draftModel: ModelExecutor;
+  finalModel: ModelExecutor;
+  resolveDraftSettings?: TwoPassSettingsResolver;
+  resolveFinalSettings?: TwoPassSettingsResolver;
   reasoningIntro?: string;
   reasoningSeparator?: string;
-  buildFinalRequest?: (request: Program, draft: string) => Program;
 };
 
-export function invokeModel(
+export function invokeExecutor(
   ctx: ExecutorContext,
   request: Program,
-  model: string,
+  executor: ModelExecutor,
   options?: InvokeOptions,
 ): Promise<Program> {
-  return ctx.invoke(setModel(request, model), options);
+  return typeof executor === "string"
+    ? ctx.invoke(setModel(request, executor), options)
+    : executor.execute(request, ctx);
 }
 
-export async function* streamModel(
+export async function* streamExecutor(
   ctx: ExecutorContext,
   request: Program,
-  model: string,
+  executor: ModelExecutor,
   options?: InvokeOptions,
 ): AsyncIterable<Program> {
-  yield* ctx.invokeStream(setModel(request, model), options);
-}
-
-export function appendInternalDraft(
-  request: Program,
-  draft: InternalDraft,
-  options: { callId?: string } = {},
-): Program {
-  const callId = options.callId ?? INTERNAL_DRAFT_CALL_ID;
-  const drafts = normalizeDrafts(draft);
-  if (drafts.length === 0) {
-    return request;
+  if (typeof executor === "string") {
+    yield* ctx.invokeStream(setModel(request, executor), options);
+    return;
   }
-
-  return drafts.reduce(
-    (program, text, index) =>
-      appendToolInteraction(program, {
-        callId: drafts.length === 1 ? callId : `${callId}_${index}`,
-        name: INTERNAL_DRAFT_TOOL_NAME,
-        args: {},
-        result: text,
-      }),
-    request,
-  );
+  yield* executor.stream(request, ctx);
 }
 
 export function createTwoPassExecutor(
@@ -84,65 +74,67 @@ export function createTwoPassExecutor(
 ): Executor {
   const reasoningIntro = options.reasoningIntro ?? "Let me think...\n\n";
   const reasoningSeparator = options.reasoningSeparator ?? "\n\n";
+  validateMaxContextLength(options.maxTotalContextLength);
 
   return {
     async execute(request, ctx) {
-      const draftRequest = buildDraftRequest(options, request);
-      if (lastMessageRole(request) === "tool") {
-        return invokeModel(ctx, draftRequest, options.draftModel);
+      const settings = await resolveSettings(options, request, ctx);
+      if (!settings.draft) {
+        return invokeFinal(options, request, ctx, settings.final!);
+      }
+      const draftRequest = buildDraftRequest(request, settings.draft);
+      if (!settings.final || lastMessageRole(request) === "tool") {
+        return invokeExecutor(ctx, draftRequest, options.draftModel);
       }
 
-      const draftResponse = await invokeModel(
+      const draftResponse = await invokeExecutor(
         ctx,
         draftRequest,
         options.draftModel,
       );
-      if (callData(draftResponse).length > 0) {
-        return draftResponse;
-      }
+      if (callData(draftResponse).length > 0) return draftResponse;
 
-      return invokeModel(
+      return invokeFinal(
+        options,
+        request,
         ctx,
-        buildFinalRequest(options, request, contentText(draftResponse)),
-        options.finalModel,
+        settings.final,
+        contentText(draftResponse),
       );
     },
 
     async *stream(request, ctx) {
-      const draftRequest = buildDraftRequest(options, request);
-      if (lastMessageRole(request) === "tool") {
-        yield* streamModel(ctx, draftRequest, options.draftModel);
+      const settings = await resolveSettings(options, request, ctx);
+      if (!settings.draft) {
+        yield* streamFinal(options, request, ctx, settings.final!);
+        return;
+      }
+      const draftRequest = buildDraftRequest(request, settings.draft);
+      if (!settings.final || lastMessageRole(request) === "tool") {
+        yield* streamExecutor(ctx, draftRequest, options.draftModel);
         return;
       }
 
       let transcript = "";
       let emittedReasoning = false;
       let toolMode = false;
-      const draftStream = streamModel(ctx, draftRequest, options.draftModel);
-
-      for await (const chunk of draftStream) {
-        if (ctx.signal.aborted) {
-          return;
-        }
-
+      for await (const chunk of streamExecutor(
+        ctx,
+        draftRequest,
+        options.draftModel,
+      )) {
+        if (ctx.signal.aborted) return;
         if (toolMode) {
           yield chunk;
           continue;
         }
-
-        const toolRequested =
-          hasToolDelta(chunk) || finishReason(chunk) === "tool_calls";
-        if (toolRequested) {
+        if (hasToolDelta(chunk) || finishReason(chunk) === "tool_calls") {
           toolMode = true;
           yield chunk;
           continue;
         }
-
         const text = deltaText(chunk);
-        if (!text) {
-          continue;
-        }
-
+        if (!text) continue;
         transcript += text;
         if (!emittedReasoning) {
           emittedReasoning = true;
@@ -151,76 +143,58 @@ export function createTwoPassExecutor(
         yield streamReasoningDelta(text);
       }
 
-      if (toolMode || ctx.signal.aborted) {
-        return;
-      }
-
+      if (toolMode || ctx.signal.aborted) return;
       if (emittedReasoning && reasoningSeparator) {
         yield streamReasoningDelta(reasoningSeparator);
       }
-
-      yield* streamModel(
-        ctx,
-        buildFinalRequest(options, request, transcript),
-        options.finalModel,
-      );
+      yield* streamFinal(options, request, ctx, settings.final, transcript);
     },
   };
 }
 
-function buildDraftRequest(
+async function resolveSettings(
   options: TwoPassExecutorOptions,
   request: Program,
-): Program {
-  if (!options.draftSystemPrompt) {
-    return request;
-  }
-
-  const systemSpans = messages(request).filter(
-    (span) => span.role === "system",
-  );
-  const previousSystemPrompt = systemSpans
-    .map((span) => messageText(request, span).trim())
-    .filter((text) => text.length > 0)
-    .join("\n\n");
-
-  let draft = request;
-  for (const span of systemSpans.slice().reverse()) {
-    draft = removeRange(draft, span.start, span.end);
-  }
-
-  draft = prependSystemPrompt(draft, options.draftSystemPrompt);
-
-  if (!previousSystemPrompt) {
-    return draft;
+  ctx: ExecutorContext,
+): Promise<{ draft: TwoPassSettings | null; final: TwoPassSettings | null }> {
+  const [draft, final] = await Promise.all([
+    options.resolveDraftSettings
+      ? options.resolveDraftSettings(request, ctx)
+      : {},
+    options.resolveFinalSettings
+      ? options.resolveFinalSettings(request, ctx)
+      : {},
+  ]);
+  if (!draft && !final) {
+    throw new Error("Two-pass executor requires at least one enabled pass");
   }
-
-  const firstNonSystem = messages(draft).find((span) => span.role !== "system");
-  return insertBefore(
-    draft,
-    firstNonSystem?.start ?? draft.code.length,
-    makeMessage("user", previousSystemPrompt).code,
-  );
+  return { draft, final };
 }
 
-function buildFinalRequest(
+function invokeFinal(
   options: TwoPassExecutorOptions,
   request: Program,
-  draft: string,
-): Program {
-  if (options.buildFinalRequest) {
-    return options.buildFinalRequest(request, draft);
-  }
-
-  const withDraft = appendInternalDraft(request, draft);
-  return options.finalSystemPrompt
-    ? prependSystemPrompt(withDraft, options.finalSystemPrompt)
-    : withDraft;
+  ctx: ExecutorContext,
+  settings: TwoPassSettings,
+  draft = "",
+): Promise<Program> {
+  return invokeExecutor(
+    ctx,
+    buildFinalRequest(options, request, settings, draft),
+    options.finalModel,
+  );
 }
 
-function normalizeDrafts(draft: InternalDraft): string[] {
-  const drafts = Array.isArray(draft)
-    ? draft.map((item) => (typeof item === "string" ? item : item.text))
-    : [draft];
-  return drafts.map((text) => text.trim()).filter((text) => text.length > 0);
+async function* streamFinal(
+  options: TwoPassExecutorOptions,
+  request: Program,
+  ctx: ExecutorContext,
+  settings: TwoPassSettings,
+  draft = "",
+): AsyncIterable<Program> {
+  yield* streamExecutor(
+    ctx,
+    buildFinalRequest(options, request, settings, draft),
+    options.finalModel,
+  );
 }

package/test/lilsdk-ts.test.ts

@@ -27,9 +27,9 @@ import { fallback, retry, runGoal } from "../src/loops/index.ts";
 import {
   appendInternalDraft,
   createTwoPassExecutor,
-  invokeModel,
+  invokeExecutor,
   INTERNAL_DRAFT_TOOL_NAME,
-  streamModel,
+  streamExecutor,
 } from "../src/managed/twoPassExecutor.ts";
 import type { Executor, ExecutorContext, OutputSink } from "../src/types.ts";
 
@@ -182,7 +182,7 @@ describe("@neutrome/lilsdk-ts", () => {
     });
   });
 
-  it("builds two-pass executors from model and prompt config", async () => {
+  it("builds two-pass executors from model and settings resolvers", async () => {
     const request = parseChatCompletionsRequest(
       encoder.encode(
         JSON.stringify({
@@ -214,7 +214,10 @@ describe("@neutrome/lilsdk-ts", () => {
     const executor = createTwoPassExecutor({
       draftModel: "smart-model",
       finalModel: "base-model",
-      finalSystemPrompt: "system prompt",
+      resolveFinalSettings: () => ({
+        reasoningLevel: "high",
+        systemPrompt: "system prompt",
+      }),
     });
 
     const result = await executor.execute(request, ctx);
@@ -232,6 +235,7 @@ describe("@neutrome/lilsdk-ts", () => {
       role: "system",
       content: "system prompt",
     });
+    expect(finalRequest.reasoning_effort).toBe("high");
     expect(finalRequest.messages[2].tool_calls[0].function).toEqual({
       name: INTERNAL_DRAFT_TOOL_NAME,
       arguments: "{}",
@@ -272,7 +276,7 @@ describe("@neutrome/lilsdk-ts", () => {
     const executor = createTwoPassExecutor({
       draftModel: "smart-model",
       finalModel: "base-model",
-      draftSystemPrompt: "draft system",
+      resolveDraftSettings: () => ({ systemPrompt: "draft system" }),
     });
 
     await executor.execute(request, ctx);
@@ -287,7 +291,7 @@ describe("@neutrome/lilsdk-ts", () => {
     ]);
   });
 
-  it("invokes and streams another model through executor context helpers", async () => {
+  it("invokes and streams string or custom executors", async () => {
     const calls: Array<{ model: string; streaming: boolean }> = [];
     const ctx = buildExecutorContext({
       async invoke(program) {
@@ -319,9 +323,9 @@ describe("@neutrome/lilsdk-ts", () => {
         }),
       ),
     );
-    const first = await invokeModel(ctx, request, "smart-model");
+    const first = await invokeExecutor(ctx, request, "smart-model");
     const chunks: Program[] = [];
-    for await (const chunk of streamModel(ctx, request, "base-model")) {
+    for await (const chunk of streamExecutor(ctx, request, "base-model")) {
       chunks.push(chunk);
     }
 
@@ -334,6 +338,81 @@ describe("@neutrome/lilsdk-ts", () => {
       decoder.decode(emitChatCompletionsStreamChunk(chunks[0]!)),
     ).toContain("hello");
     expect(getModel(request)).toBe("virtual-model");
+
+    const customExecutor = executorFromExecute(async () =>
+      appendAssistantMessage({ code: [], buffers: [] }, "custom answer"),
+    );
+    const custom = await invokeExecutor(ctx, request, customExecutor);
+    expect(contentText(custom)).toBe("custom answer");
+
+    const twoPass = createTwoPassExecutor({
+      draftModel: customExecutor,
+      finalModel: customExecutor,
+    });
+    expect(contentText(await twoPass.execute(request, ctx))).toBe(
+      "custom answer",
+    );
+  });
+
+  it("can skip the draft pass and limits the final context length", async () => {
+    const request = parseChatCompletionsRequest(
+      encoder.encode(
+        JSON.stringify({
+          model: "virtual-model",
+          messages: [{ role: "user", content: "request" }],
+        }),
+      ),
+    );
+    const seen: Program[] = [];
+    const ctx = buildExecutorContext({
+      async invoke(program) {
+        seen.push(program);
+        return appendAssistantMessage(
+          { code: [], buffers: [] },
+          seen.length === 1 ? "a very long draft" : "final",
+        );
+      },
+      async *invokeStream() {
+        throw new Error("streaming path is not used in this test");
+      },
+    });
+
+    const executor = createTwoPassExecutor({
+      draftModel: "smart-model",
+      finalModel: "base-model",
+      maxTotalContextLength: 10,
+      resolveDraftSettings: () => ({ systemPrompt: "draft" }),
+    });
+    await executor.execute(request, ctx);
+
+    const finalRequest = JSON.parse(
+      decoder.decode(emitChatCompletionsRequest(seen[1]!)),
+    );
+    expect(finalRequest.messages.at(-1)?.content).toBe("a v");
+
+    let receivedRequest: Program | undefined;
+    let receivedContext: ExecutorContext | undefined;
+    const directFinal = createTwoPassExecutor({
+      draftModel: "unused",
+      finalModel: "base-model",
+      resolveDraftSettings: (incomingRequest, incomingContext) => {
+        receivedRequest = incomingRequest;
+        receivedContext = incomingContext;
+        return null;
+      },
+    });
+    await directFinal.execute(request, ctx);
+    expect(getModel(seen.at(-1)!)).toBe("base-model");
+    expect(receivedRequest).toBe(request);
+    expect(receivedContext).toBe(ctx);
+
+    const draftOnly = createTwoPassExecutor({
+      draftModel: "smart-model",
+      finalModel: "unused",
+      resolveFinalSettings: () => null,
+    });
+    await draftOnly.execute(request, ctx);
+    expect(getModel(seen.at(-1)!)).toBe("smart-model");
   });
 
   it("writes reasoning helpers to a sink", async () => {