mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
Add support for generic sampling parameters (#7568)
This commit is contained in:
parent
fed6009cc9
commit
25a2c8dcfe
12 changed files with 223 additions and 2 deletions
|
|
@ -59,6 +59,7 @@ export type ProxyAssistantMessageEvent =
|
|||
type ProxySerializableStreamOptions = Pick<
|
||||
SimpleStreamOptions,
|
||||
| "temperature"
|
||||
| "samplingParams"
|
||||
| "maxTokens"
|
||||
| "reasoning"
|
||||
| "cacheRetention"
|
||||
|
|
@ -101,6 +102,7 @@ export interface ProxyStreamOptions extends ProxySerializableStreamOptions {
|
|||
function buildProxyRequestOptions(options: ProxyStreamOptions): ProxySerializableStreamOptions {
|
||||
return {
|
||||
temperature: options.temperature,
|
||||
samplingParams: options.samplingParams,
|
||||
maxTokens: options.maxTokens,
|
||||
reasoning: options.reasoning,
|
||||
cacheRetention: options.cacheRetention,
|
||||
|
|
|
|||
|
|
@ -321,5 +321,10 @@ function buildParams(
|
|||
}
|
||||
}
|
||||
|
||||
// Last so custom keys override the named request fields.
|
||||
if (options?.samplingParams) {
|
||||
Object.assign(params, options.samplingParams);
|
||||
}
|
||||
|
||||
return params;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -855,6 +855,11 @@ function buildParams(
|
|||
}
|
||||
}
|
||||
|
||||
// Last so custom keys override the named request fields.
|
||||
if (options?.samplingParams) {
|
||||
Object.assign(params, options.samplingParams);
|
||||
}
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -327,6 +327,11 @@ function buildParams(
|
|||
if (model.provider === "xai") params.include = ["reasoning.encrypted_content"];
|
||||
}
|
||||
|
||||
// Last so custom keys override the named request fields.
|
||||
if (options?.samplingParams) {
|
||||
Object.assign(params, options.samplingParams);
|
||||
}
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -24,8 +24,13 @@ export function buildBaseOptions(
|
|||
options?: SimpleStreamOptions,
|
||||
apiKey?: string,
|
||||
): StreamOptions {
|
||||
const samplingParams =
|
||||
model.samplingParams || options?.samplingParams
|
||||
? { ...model.samplingParams, ...options?.samplingParams }
|
||||
: undefined;
|
||||
return {
|
||||
temperature: options?.temperature,
|
||||
samplingParams,
|
||||
maxTokens: clampMaxTokensToContext(model, context, options?.maxTokens ?? model.maxTokens),
|
||||
signal: options?.signal,
|
||||
apiKey: apiKey || options?.apiKey,
|
||||
|
|
|
|||
|
|
@ -116,6 +116,14 @@ export interface ProviderResponse {
|
|||
|
||||
export interface StreamOptions {
|
||||
temperature?: number;
|
||||
/**
|
||||
* Arbitrary sampling parameters merged into the request body as-is, after the named request
|
||||
* fields, so keys here override them. Lets custom OpenAI-compatible servers (llama.cpp, vLLM,
|
||||
* SGLang, ...) receive parameters pi does not model, e.g. `top_p`, `top_k`, `min_p`,
|
||||
* `repetition_penalty`. Merged over `Model.samplingParams` per key. Only applied by
|
||||
* OpenAI-compatible adapters (completions, responses, Azure responses); other APIs ignore it.
|
||||
*/
|
||||
samplingParams?: Record<string, unknown>;
|
||||
maxTokens?: number;
|
||||
signal?: AbortSignal;
|
||||
apiKey?: string;
|
||||
|
|
@ -778,6 +786,8 @@ export interface Model<TApi extends Api> {
|
|||
cost: ModelCost;
|
||||
contextWindow: number;
|
||||
maxTokens: number;
|
||||
/** Default sampling parameters for this model. See {@link StreamOptions.samplingParams}; per-request keys override these. */
|
||||
samplingParams?: Record<string, unknown>;
|
||||
headers?: Record<string, string>;
|
||||
/** Compatibility overrides for OpenAI-compatible APIs. If not set, auto-detected from baseUrl. */
|
||||
compat?: TApi extends "openai-completions"
|
||||
|
|
|
|||
128
packages/ai/test/sampling-options.test.ts
Normal file
128
packages/ai/test/sampling-options.test.ts
Normal file
|
|
@ -0,0 +1,128 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import { streamSimple } from "../src/compat.ts";
|
||||
import type { Api, Context, Model, SimpleStreamOptions } from "../src/types.ts";
|
||||
|
||||
interface SamplingPayload {
|
||||
temperature?: number;
|
||||
top_p?: number;
|
||||
top_k?: number;
|
||||
min_p?: number;
|
||||
}
|
||||
|
||||
class PayloadCaptured extends Error {
|
||||
constructor() {
|
||||
super("payload captured");
|
||||
this.name = "PayloadCaptured";
|
||||
}
|
||||
}
|
||||
|
||||
function makeContext(): Context {
|
||||
return {
|
||||
messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
|
||||
};
|
||||
}
|
||||
|
||||
function makeCompletionsModel(overrides?: Partial<Model<"openai-completions">>): Model<"openai-completions"> {
|
||||
return {
|
||||
id: "custom-model",
|
||||
name: "Custom Model",
|
||||
api: "openai-completions",
|
||||
provider: "custom-provider",
|
||||
baseUrl: "http://127.0.0.1:9/v1",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 16384,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function makeAnthropicModel(): Model<"anthropic-messages"> {
|
||||
return {
|
||||
id: "vendor--claude",
|
||||
name: "Vendor Proxy Claude",
|
||||
api: "anthropic-messages",
|
||||
provider: "vendor-proxy",
|
||||
baseUrl: "http://127.0.0.1:9",
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 200000,
|
||||
maxTokens: 32000,
|
||||
};
|
||||
}
|
||||
|
||||
async function capturePayload(model: Model<Api>, options?: SimpleStreamOptions): Promise<SamplingPayload> {
|
||||
let capturedPayload: SamplingPayload | undefined;
|
||||
|
||||
const s = streamSimple(model, makeContext(), {
|
||||
...options,
|
||||
apiKey: "fake-key",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload as SamplingPayload;
|
||||
throw new PayloadCaptured();
|
||||
},
|
||||
});
|
||||
|
||||
await s.result();
|
||||
|
||||
if (!capturedPayload) {
|
||||
throw new Error("Expected payload to be captured before request failure");
|
||||
}
|
||||
|
||||
return capturedPayload;
|
||||
}
|
||||
|
||||
describe("sampling params", () => {
|
||||
it("merges stream-option sampling params into the request body", async () => {
|
||||
const payload = await capturePayload(makeCompletionsModel(), {
|
||||
samplingParams: { top_p: 0.95, top_k: 0, min_p: 0 },
|
||||
});
|
||||
|
||||
expect(payload.top_p).toBe(0.95);
|
||||
expect(payload.top_k).toBe(0);
|
||||
expect(payload.min_p).toBe(0);
|
||||
});
|
||||
|
||||
it("omits sampling params when neither options nor model set them", async () => {
|
||||
const payload = await capturePayload(makeCompletionsModel());
|
||||
|
||||
expect(payload.temperature).toBeUndefined();
|
||||
expect(payload.top_p).toBeUndefined();
|
||||
});
|
||||
|
||||
it("applies model-level sampling params", async () => {
|
||||
const payload = await capturePayload(makeCompletionsModel({ samplingParams: { temperature: 1, top_p: 0.95 } }));
|
||||
|
||||
expect(payload.temperature).toBe(1);
|
||||
expect(payload.top_p).toBe(0.95);
|
||||
});
|
||||
|
||||
it("merges stream-option keys over model-level keys", async () => {
|
||||
const payload = await capturePayload(makeCompletionsModel({ samplingParams: { top_p: 0.95, min_p: 0.05 } }), {
|
||||
samplingParams: { top_p: 0.5 },
|
||||
});
|
||||
|
||||
expect(payload.top_p).toBe(0.5);
|
||||
expect(payload.min_p).toBe(0.05);
|
||||
});
|
||||
|
||||
it("overrides named request fields", async () => {
|
||||
const payload = await capturePayload(makeCompletionsModel(), {
|
||||
temperature: 0,
|
||||
samplingParams: { temperature: 1 },
|
||||
});
|
||||
|
||||
expect(payload.temperature).toBe(1);
|
||||
});
|
||||
|
||||
it("is ignored by non-OpenAI-compatible APIs", async () => {
|
||||
const payload = await capturePayload(makeAnthropicModel(), {
|
||||
samplingParams: { top_p: 0.9, top_k: 40 },
|
||||
});
|
||||
|
||||
expect(payload.top_p).toBeUndefined();
|
||||
expect(payload.top_k).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
|
@ -206,6 +206,7 @@ If your command is slow, expensive, rate-limited, or should keep using a previou
|
|||
| `input` | No | `["text"]` | Input types: `["text"]` or `["text", "image"]` |
|
||||
| `contextWindow` | No | `128000` | Context window size in tokens |
|
||||
| `maxTokens` | No | `16384` | Maximum output tokens |
|
||||
| `samplingParams` | No | omitted | Sampling parameters merged verbatim into every request body (see below) |
|
||||
| `cost` | No | all zeros | Per-million-token rates with optional request-wide input pricing tiers |
|
||||
| `compat` | No | provider `compat` | Provider compatibility overrides. Merged with provider-level `compat` when both are set. |
|
||||
|
||||
|
|
@ -235,6 +236,24 @@ Current behavior:
|
|||
- `/model`, `--list-models`, and the interactive footer display entries by model `id`.
|
||||
- The configured `name` is used for model matching and secondary model detail text. It does not replace the footer/status-bar model id.
|
||||
|
||||
### Sampling Parameters
|
||||
|
||||
`samplingParams` is a free-form object merged verbatim into every request body for the model, after the fields pi sets itself, so its keys win. Use it to send sampling parameters pi does not model — including server-specific ones like llama.cpp's `min_p` or vLLM's `top_k`:
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "deepseek-v4-flash",
|
||||
"samplingParams": {
|
||||
"temperature": 1.0,
|
||||
"top_p": 0.95,
|
||||
"top_k": 0,
|
||||
"min_p": 0.0
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Only OpenAI-compatible APIs apply it (`openai-completions`, `openai-responses`, `azure-openai-responses`); other APIs ignore it. Keys override pi's named request fields (for example a `temperature` key here beats the request-level temperature), so prefer it as the single source of sampling truth for a model. In `modelOverrides`, `samplingParams` merges per key with the base model's value.
|
||||
|
||||
### Thinking Level Map
|
||||
|
||||
Use `thinkingLevelMap` on a model to describe model-specific thinking controls. Keys are pi thinking levels: `off`, `minimal`, `low`, `medium`, `high`, `xhigh`, `max`. Maps may contain holes; for example, a model can expose `high` and `max` without exposing `xhigh`.
|
||||
|
|
@ -338,7 +357,7 @@ Use `modelOverrides` to customize built-in models and matching extension-registe
|
|||
}
|
||||
```
|
||||
|
||||
`modelOverrides` supports these fields per model: `name`, `reasoning`, `thinkingLevelMap`, `input`, `cost` (partial), `contextWindow`, `maxTokens`, `headers`, `compat`.
|
||||
`modelOverrides` supports these fields per model: `name`, `reasoning`, `thinkingLevelMap`, `input`, `cost` (partial), `contextWindow`, `maxTokens`, `samplingParams` (merged per key), `headers`, `compat`.
|
||||
|
||||
Direct OpenAI GPT-5.6 Sol, Terra, and Luna default to a `272000` context window so requests remain within OpenAI's short-context pricing tier. To opt into OpenAI's 1.05M context window, increase it for each model you use:
|
||||
|
||||
|
|
|
|||
|
|
@ -164,6 +164,7 @@ const ModelDefinitionSchema = Type.Object({
|
|||
cost: Type.Optional(ModelCostSchema),
|
||||
contextWindow: Type.Optional(Type.Number()),
|
||||
maxTokens: Type.Optional(Type.Number()),
|
||||
samplingParams: Type.Optional(Type.Record(Type.String(), Type.Unknown())),
|
||||
headers: Type.Optional(Type.Record(Type.String(), Type.String())),
|
||||
compat: Type.Optional(ProviderCompatSchema),
|
||||
});
|
||||
|
|
@ -184,6 +185,7 @@ const ModelOverrideSchema = Type.Object({
|
|||
),
|
||||
contextWindow: Type.Optional(Type.Number()),
|
||||
maxTokens: Type.Optional(Type.Number()),
|
||||
samplingParams: Type.Optional(Type.Record(Type.String(), Type.Unknown())),
|
||||
headers: Type.Optional(Type.Record(Type.String(), Type.String())),
|
||||
compat: Type.Optional(ProviderCompatSchema),
|
||||
});
|
||||
|
|
|
|||
|
|
@ -61,6 +61,7 @@ export interface ProviderConfigInput {
|
|||
cost: Model<Api>["cost"];
|
||||
contextWindow: number;
|
||||
maxTokens: number;
|
||||
samplingParams?: Record<string, unknown>;
|
||||
headers?: Record<string, string>;
|
||||
compat?: Model<Api>["compat"];
|
||||
}>;
|
||||
|
|
@ -117,6 +118,9 @@ function applyModelOverride(model: Model<Api>, override: ModelsJsonModelOverride
|
|||
: model.cost,
|
||||
contextWindow: override.contextWindow ?? model.contextWindow,
|
||||
maxTokens: override.maxTokens ?? model.maxTokens,
|
||||
samplingParams: override.samplingParams
|
||||
? { ...model.samplingParams, ...override.samplingParams }
|
||||
: model.samplingParams,
|
||||
compat: mergeCompat(model.compat, override.compat),
|
||||
};
|
||||
}
|
||||
|
|
@ -153,6 +157,7 @@ function modelFromJson(
|
|||
cost: definition.cost ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: definition.contextWindow ?? 128000,
|
||||
maxTokens: definition.maxTokens ?? 16384,
|
||||
samplingParams: definition.samplingParams,
|
||||
headers: undefined,
|
||||
compat: mergeCompat(providerConfig.compat, definition.compat),
|
||||
};
|
||||
|
|
|
|||
|
|
@ -713,6 +713,39 @@ describe("ModelRegistry", () => {
|
|||
expect(opus?.name).not.toBe("Custom Sonnet Name");
|
||||
});
|
||||
|
||||
test("custom model and model override carry sampling params", async () => {
|
||||
writeRawModelsJson({
|
||||
openrouter: {
|
||||
baseUrl: "https://my-proxy.example.com/v1",
|
||||
api: "openai-completions",
|
||||
models: [
|
||||
{
|
||||
id: "custom/sampling-model",
|
||||
samplingParams: { temperature: 1, top_p: 0.95, top_k: 0 },
|
||||
},
|
||||
],
|
||||
modelOverrides: {
|
||||
"anthropic/claude-sonnet-4": {
|
||||
samplingParams: { top_p: 0.9 },
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const registry = await createModelRegistry(authStorage, modelsJsonPath);
|
||||
const models = getModelsForProvider(registry, "openrouter");
|
||||
|
||||
const custom = models.find((m) => m.id === "custom/sampling-model");
|
||||
expect(custom?.samplingParams).toEqual({ temperature: 1, top_p: 0.95, top_k: 0 });
|
||||
|
||||
const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4");
|
||||
expect(sonnet?.samplingParams).toEqual({ top_p: 0.9 });
|
||||
|
||||
// Models without sampling config keep it unset.
|
||||
const opus = models.find((m) => m.id === "anthropic/claude-opus-4");
|
||||
expect(opus?.samplingParams).toBeUndefined();
|
||||
});
|
||||
|
||||
test("model override with compat.openRouterRouting", async () => {
|
||||
writeRawModelsJson({
|
||||
openrouter: {
|
||||
|
|
|
|||
|
|
@ -32,7 +32,8 @@ type _ProtocolModelInputsFitAi = Assert<ProtocolModelInput extends AiModelInput
|
|||
/**
|
||||
* Enumerate mapped and intentionally omitted pi-ai fields so additions fail compilation here.
|
||||
* Provider replay metadata, diagnostics, cache-write retention splits, model transport settings,
|
||||
* pricing tiers, and deferred-tool availability remain intentionally server-side.
|
||||
* model sampling defaults, pricing tiers, and deferred-tool availability remain intentionally
|
||||
* server-side.
|
||||
*/
|
||||
type _AiTextContentFieldsAccountedFor = Assert<ExactKeys<AiTextContent, "type" | "text" | "textSignature">>;
|
||||
type _AiThinkingContentFieldsAccountedFor = Assert<
|
||||
|
|
@ -68,6 +69,7 @@ type _AiModelFieldsAccountedFor = Assert<
|
|||
| "cost"
|
||||
| "contextWindow"
|
||||
| "maxTokens"
|
||||
| "samplingParams"
|
||||
| "headers"
|
||||
| "compat"
|
||||
>
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue