feat(agents): forward generation params (temperature/top_p/stop/...) to the serving adapter

packages/appkit/src/agents/databricks.ts

@@ -27,6 +27,44 @@ function isRecord(value: unknown): value is Record<string, unknown> {
   return typeof value === "object" && value !== null;
 }
 
+/**
+ * Optional generation parameters forwarded to the OpenAI-compatible serving
+ * request body. Names match the serving API wire keys. Only keys that are set
+ * are sent — undefined values are omitted so the endpoint applies its own
+ * defaults. Ranges are not validated here; the serving endpoint validates.
+ */
+export interface GenerationParams {
+  /** Sampling temperature. */
+  temperature?: number;
+  /** Nucleus sampling probability mass (`top_p`). */
+  top_p?: number;
+  /** Stop sequence(s) that end generation. */
+  stop?: string | string[];
+  /** Penalize tokens by frequency. */
+  frequency_penalty?: number;
+  /** Penalize tokens by prior presence. */
+  presence_penalty?: number;
+}
+
+const GENERATION_PARAM_KEYS = [
+  "temperature",
+  "top_p",
+  "stop",
+  "frequency_penalty",
+  "presence_penalty",
+] as const satisfies readonly (keyof GenerationParams)[];
+
+/** Copy only the set generation params onto the request body. */
+function applyGenerationParams(
+  body: Record<string, unknown>,
+  params: GenerationParams,
+): void {
+  for (const key of GENERATION_PARAM_KEYS) {
+    const value = params[key];
+    if (value !== undefined) body[key] = value;
+  }
+}
+
 function extractLlamaToolJsonSlice(text: string): string | undefined {
   const start = text.indexOf("[{");
   if (start < 0) return undefined;
@@ -84,6 +122,8 @@ interface RawFetchAdapterOptions {
   authenticate: () => Promise<Record<string, string>>;
   maxSteps?: number;
   maxTokens?: number;
+  /** Optional generation params forwarded to the serving request body. */
+  generationParams?: GenerationParams;
   /** Max length of one SSE line (including an incomplete tail in the buffer). */
   maxSseLineChars?: number;
   /** Max total length of assistant `delta.content` across the stream. */
@@ -102,6 +142,7 @@ interface StreamBodyAdapterOptions {
   streamBody: StreamBody;
   maxSteps?: number;
   maxTokens?: number;
+  generationParams?: GenerationParams;
   maxSseLineChars?: number;
   maxStreamTextChars?: number;
   maxToolArgumentsChars?: number;
@@ -135,6 +176,7 @@ interface ServingEndpointOptions {
   endpointName: string;
   maxSteps?: number;
   maxTokens?: number;
+  generationParams?: GenerationParams;
   maxSseLineChars?: number;
   maxStreamTextChars?: number;
   maxToolArgumentsChars?: number;
@@ -143,6 +185,7 @@ interface ServingEndpointOptions {
 interface ModelServingOptions {
   maxSteps?: number;
   maxTokens?: number;
+  generationParams?: GenerationParams;
   workspaceClient?: WorkspaceClientLike;
   maxSseLineChars?: number;
   maxStreamTextChars?: number;
@@ -238,13 +281,15 @@ export class DatabricksAdapter implements AgentAdapter {
   private streamBody: StreamBody;
   private maxSteps: number;
   private maxTokens: number;
+  private generationParams: GenerationParams;
   private maxSseLineChars: number;
   private maxStreamTextChars: number;
   private maxToolArgumentsChars: number;
 
   constructor(options: DatabricksAdapterOptions) {
     this.maxSteps = options.maxSteps ?? 10;
     this.maxTokens = options.maxTokens ?? 4096;
+    this.generationParams = options.generationParams ?? {};
     this.maxSseLineChars =
       options.maxSseLineChars ?? DEFAULT_MAX_SSE_LINE_CHARS;
     this.maxStreamTextChars =
@@ -298,6 +343,7 @@ export class DatabricksAdapter implements AgentAdapter {
       endpointName,
       maxSteps,
       maxTokens,
+      generationParams,
       maxSseLineChars,
       maxStreamTextChars,
       maxToolArgumentsChars,
@@ -315,6 +361,7 @@ export class DatabricksAdapter implements AgentAdapter {
         ),
       maxSteps,
       maxTokens,
+      generationParams,
       maxSseLineChars,
       maxStreamTextChars,
       maxToolArgumentsChars,
@@ -370,6 +417,7 @@ export class DatabricksAdapter implements AgentAdapter {
       endpointName: resolvedEndpoint,
       maxSteps: options?.maxSteps,
       maxTokens: options?.maxTokens,
+      generationParams: options?.generationParams,
       maxSseLineChars: options?.maxSseLineChars,
       maxStreamTextChars: options?.maxStreamTextChars,
       maxToolArgumentsChars: options?.maxToolArgumentsChars,
@@ -495,6 +543,8 @@ export class DatabricksAdapter implements AgentAdapter {
       max_tokens: this.maxTokens,
     };
 
+    applyGenerationParams(body, this.generationParams);
+
     if (tools.length > 0) {
       body.tools = tools;
     }

packages/appkit/src/agents/tests/databricks.test.ts

@@ -1,6 +1,10 @@
 import type { AgentEvent, AgentToolDefinition, Message } from "shared";
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
-import { DatabricksAdapter, parseTextToolCalls } from "../databricks";
+import {
+  DatabricksAdapter,
+  type GenerationParams,
+  parseTextToolCalls,
+} from "../databricks";
 
 const mockAuthenticate = vi
   .fn()
@@ -93,6 +97,7 @@ function createAdapter(overrides?: {
   authenticate?: () => Promise<Record<string, string>>;
   maxSteps?: number;
   maxTokens?: number;
+  generationParams?: GenerationParams;
   maxSseLineChars?: number;
   maxStreamTextChars?: number;
   maxToolArgumentsChars?: number;
@@ -566,6 +571,56 @@ describe("DatabricksAdapter", () => {
     });
   });
 
+  test("forwards set generation params to the request body", async () => {
+    globalThis.fetch = mockFetch([textDelta("Hi"), sseChunk("[DONE]")]);
+
+    const adapter = createAdapter({
+      generationParams: {
+        temperature: 0.2,
+        top_p: 0.9,
+        stop: ["END"],
+        frequency_penalty: 0.5,
+        presence_penalty: 0.1,
+      },
+    });
+
+    for await (const _ of adapter.run(
+      { messages: createTestMessages(), tools: [], threadId: "t1" },
+      { executeTool: vi.fn() },
+    )) {
+      // drain
+    }
+
+    const [, init] = (globalThis.fetch as any).mock.calls[0];
+    const body = JSON.parse(init.body);
+    expect(body.temperature).toBe(0.2);
+    expect(body.top_p).toBe(0.9);
+    expect(body.stop).toEqual(["END"]);
+    expect(body.frequency_penalty).toBe(0.5);
+    expect(body.presence_penalty).toBe(0.1);
+  });
+
+  test("omits generation param keys that are not set", async () => {
+    globalThis.fetch = mockFetch([textDelta("Hi"), sseChunk("[DONE]")]);
+
+    const adapter = createAdapter({ generationParams: { temperature: 0.7 } });
+
+    for await (const _ of adapter.run(
+      { messages: createTestMessages(), tools: [], threadId: "t1" },
+      { executeTool: vi.fn() },
+    )) {
+      // drain
+    }
+
+    const [, init] = (globalThis.fetch as any).mock.calls[0];
+    const body = JSON.parse(init.body);
+    expect(body.temperature).toBe(0.7);
+    expect(body).not.toHaveProperty("top_p");
+    expect(body).not.toHaveProperty("stop");
+    expect(body).not.toHaveProperty("frequency_penalty");
+    expect(body).not.toHaveProperty("presence_penalty");
+  });
+
   test("forwards tool thread fields from input messages to the request body", async () => {
     globalThis.fetch = mockFetch([textDelta("Done"), sseChunk("[DONE]")]);
 

packages/appkit/src/beta.ts

@@ -18,7 +18,11 @@ export type {
   ToolAnnotations,
   ToolProvider,
 } from "shared";
-export { DatabricksAdapter, parseTextToolCalls } from "./agents/databricks";
+export {
+  DatabricksAdapter,
+  type GenerationParams,
+  parseTextToolCalls,
+} from "./agents/databricks";
 
 // Agent runtime
 export { createAgent } from "./core/agent/create-agent";