-
Notifications
You must be signed in to change notification settings - Fork 1.9k
fix(compaction): enable custom provider remote compaction #3106
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 1 commit
b7aefe0
d370b37
aa463f9
d1066ae
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -15,7 +15,7 @@ | |
| import { ProviderHttpError } from "@oh-my-pi/pi-ai/errors"; | ||
| import { parseTextSignature } from "@oh-my-pi/pi-ai/providers/openai-shared"; | ||
| import { transformMessages } from "@oh-my-pi/pi-ai/providers/transform-messages"; | ||
| import type { AssistantMessage, FetchImpl, Message, Model } from "@oh-my-pi/pi-ai/types"; | ||
| import type { Api, AssistantMessage, FetchImpl, Message, Model } from "@oh-my-pi/pi-ai/types"; | ||
| import { | ||
| getOpenAIResponsesHistoryItems, | ||
| getOpenAIResponsesHistoryPayload, | ||
|
|
@@ -86,12 +86,22 @@ export interface RemoteCompactionResponse { | |
| // OpenAI provider gating + endpoint resolution | ||
| // ============================================================================ | ||
|
|
||
| function isOpenAiRemoteCompactionApi(api: Api | undefined): boolean { | ||
| return api === "openai-responses" || api === "azure-openai-responses" || api === "openai-codex-responses"; | ||
| } | ||
|
|
||
| export function shouldUseOpenAiRemoteCompaction(model: Model): boolean { | ||
| return model.provider === "openai" || model.provider === "openai-codex"; | ||
| if (model.remoteCompaction?.enabled === false) return false; | ||
| if (model.provider === "openai" || model.provider === "openai-codex") return true; | ||
| if (model.remoteCompaction?.enabled !== true) return false; | ||
| return isOpenAiRemoteCompactionApi(model.remoteCompaction.api ?? model.api); | ||
| } | ||
|
|
||
| function resolveOpenAiCompactEndpoint(model: Model): string { | ||
| if (model.provider === "openai-codex") { | ||
| const configuredEndpoint = model.remoteCompaction?.endpoint; | ||
| if (configuredEndpoint && configuredEndpoint.length > 0) return configuredEndpoint; | ||
| const compactionApi = model.remoteCompaction?.api ?? model.api; | ||
| if (model.provider === "openai-codex" || compactionApi === "openai-codex-responses") { | ||
| return resolveOpenAiCodexCompactEndpoint(model.baseUrl); | ||
| } | ||
|
|
||
|
|
@@ -444,7 +454,6 @@ export function buildOpenAiNativeHistory( | |
| // ============================================================================ | ||
| // Endpoint requests | ||
| // ============================================================================ | ||
|
|
||
| export async function requestOpenAiRemoteCompaction( | ||
| model: Model, | ||
| apiKey: string, | ||
|
|
@@ -454,8 +463,9 @@ export async function requestOpenAiRemoteCompaction( | |
| opts?: { fetch?: FetchImpl; timeoutMs?: number }, | ||
| ): Promise<OpenAiRemoteCompactionResponse> { | ||
| const endpoint = resolveOpenAiCompactEndpoint(model); | ||
| const requestModel = model.remoteCompaction?.model ?? model.requestModelId ?? model.id; | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When Useful? React with 👍 / 👎. |
||
| const request: OpenAiRemoteCompactionRequest = { | ||
| model: model.id, | ||
| model: requestModel, | ||
| input: trimOpenAiCompactInput(compactInput, model.contextWindow ?? Number.POSITIVE_INFINITY, instructions), | ||
| instructions, | ||
| }; | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,7 +1,15 @@ | ||
| import { execSync } from "node:child_process"; | ||
| import * as path from "node:path"; | ||
| import { registerCustomApi, unregisterCustomApis } from "@oh-my-pi/pi-ai/api-registry"; | ||
| import type { Api, Context, Model, ModelSpec, SimpleStreamOptions, ThinkingConfig } from "@oh-my-pi/pi-ai/types"; | ||
| import type { | ||
| Api, | ||
| Context, | ||
| Model, | ||
| ModelSpec, | ||
| RemoteCompactionConfig, | ||
| SimpleStreamOptions, | ||
| ThinkingConfig, | ||
| } from "@oh-my-pi/pi-ai/types"; | ||
| import type { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream"; | ||
| import { buildModel } from "@oh-my-pi/pi-catalog/build"; | ||
| import { isVertexExpressOpenAIUrl } from "@oh-my-pi/pi-catalog/hosts"; | ||
|
|
@@ -95,6 +103,7 @@ interface ProviderOverride { | |
| apiKey?: string; | ||
| authHeader?: boolean; | ||
| compat?: ModelSpec<Api>["compat"]; | ||
| remoteCompaction?: RemoteCompactionConfig<Api>; | ||
| transport?: Model<Api>["transport"]; | ||
| } | ||
|
|
||
|
|
@@ -127,7 +136,7 @@ interface ProviderOverride { | |
| export function mergeDiscoveredModel<TApi extends Api>( | ||
| model: Model<TApi>, | ||
| existing: Model<Api> | undefined, | ||
| providerOverride?: Pick<ProviderOverride, "baseUrl" | "headers" | "transport">, | ||
| providerOverride?: Pick<ProviderOverride, "baseUrl" | "headers" | "remoteCompaction" | "transport">, | ||
| ): Model<TApi> { | ||
| if (existing) { | ||
| const supportsTools = model.supportsTools ?? existing.supportsTools; | ||
|
|
@@ -136,6 +145,10 @@ export function mergeDiscoveredModel<TApi extends Api>( | |
| baseUrl: providerOverride?.baseUrl ?? model.baseUrl ?? existing.baseUrl, | ||
| headers: existing.headers ? { ...existing.headers, ...model.headers } : model.headers, | ||
| transport: providerOverride?.transport ?? existing.transport ?? model.transport, | ||
| remoteCompaction: mergeRemoteCompactionConfig( | ||
| existing.remoteCompaction ?? model.remoteCompaction, | ||
| providerOverride?.remoteCompaction, | ||
| ), | ||
| ...(supportsTools !== undefined ? { supportsTools } : {}), | ||
| compat: model.compatConfig, | ||
| } as ModelSpec<TApi>); | ||
|
|
@@ -146,6 +159,7 @@ export function mergeDiscoveredModel<TApi extends Api>( | |
| baseUrl: providerOverride.baseUrl ?? model.baseUrl, | ||
| headers: providerOverride.headers ? { ...model.headers, ...providerOverride.headers } : model.headers, | ||
| ...(providerOverride.transport !== undefined ? { transport: providerOverride.transport } : {}), | ||
| remoteCompaction: mergeRemoteCompactionConfig(model.remoteCompaction, providerOverride.remoteCompaction), | ||
| compat: model.compatConfig, | ||
| } as ModelSpec<TApi>); | ||
| } | ||
|
|
@@ -353,6 +367,15 @@ function mergeCompat<TBase extends object, TOverride extends object>( | |
| return merged as TBase & TOverride; | ||
| } | ||
|
|
||
| function mergeRemoteCompactionConfig( | ||
| baseConfig: RemoteCompactionConfig<Api> | undefined, | ||
| overrideConfig: RemoteCompactionConfig<Api> | undefined, | ||
| ): RemoteCompactionConfig<Api> | undefined { | ||
| if (!baseConfig) return overrideConfig; | ||
| if (!overrideConfig) return baseConfig; | ||
| return { ...baseConfig, ...overrideConfig }; | ||
| } | ||
|
|
||
| /** | ||
| * Project a built model back to spec shape for the model-manager/cache | ||
| * boundary: sparse compat comes from `compatConfig`, never from the resolved | ||
|
|
@@ -380,6 +403,8 @@ interface ModelPatch { | |
| headers?: Record<string, string>; | ||
| compat?: ModelSpec<Api>["compat"]; | ||
| contextPromotionTarget?: string; | ||
| compactionModel?: string; | ||
| remoteCompaction?: RemoteCompactionConfig<Api>; | ||
| premiumMultiplier?: number; | ||
| } | ||
|
|
||
|
|
@@ -403,6 +428,10 @@ function applyModelPatch(base: Model<Api>, patch: ModelPatch, transport: ModelTr | |
| if (patch.maxTokens !== undefined) result.maxTokens = patch.maxTokens; | ||
| if (patch.omitMaxOutputTokens !== undefined) result.omitMaxOutputTokens = patch.omitMaxOutputTokens; | ||
| if (patch.contextPromotionTarget !== undefined) result.contextPromotionTarget = patch.contextPromotionTarget; | ||
| if (patch.compactionModel !== undefined) result.compactionModel = patch.compactionModel; | ||
| if (patch.remoteCompaction !== undefined) { | ||
| result.remoteCompaction = mergeRemoteCompactionConfig(base.remoteCompaction, patch.remoteCompaction); | ||
| } | ||
| if (patch.premiumMultiplier !== undefined) result.premiumMultiplier = patch.premiumMultiplier; | ||
| if (patch.cost) { | ||
| result.cost = { | ||
|
|
@@ -475,8 +504,6 @@ function mergeAuthHeader( | |
| /** | ||
| * Decide whether a custom-yaml model should force OAuth-style request shaping. | ||
| * - Explicit `auth: oauth` → force on. | ||
| * - Explicit `auth: apiKey` / `auth: none` → leave unset (auto-detect by key prefix). | ||
| * - No `auth` specified and `api: anthropic-messages` → default on. Custom Anthropic | ||
| * endpoints are typically Claude-Code-style proxies (e.g. CLIProxyAPI) that expect | ||
| * the cloaked request shape regardless of how the proxy itself is authenticated. | ||
| * - Otherwise → unset. | ||
|
|
@@ -497,6 +524,7 @@ function buildCustomModelOverlay( | |
| authHeader: boolean | undefined, | ||
| providerCompat: ModelSpec<Api>["compat"] | undefined, | ||
| providerAuth: ProviderAuthMode | undefined, | ||
| providerRemoteCompaction: RemoteCompactionConfig<Api> | undefined, | ||
| modelDef: CustomModelDefinitionLike, | ||
| ): CustomModelOverlay | undefined { | ||
| const api = modelDef.api ?? providerApi; | ||
|
|
@@ -518,6 +546,8 @@ function buildCustomModelOverlay( | |
| headers: mergeCustomModelHeaders(providerHeaders, modelDef.headers, authHeader, providerApiKey), | ||
| compat: mergeCompat(providerCompat, modelDef.compat), | ||
| contextPromotionTarget: modelDef.contextPromotionTarget, | ||
| compactionModel: modelDef.compactionModel, | ||
| remoteCompaction: mergeRemoteCompactionConfig(providerRemoteCompaction, modelDef.remoteCompaction), | ||
| premiumMultiplier: modelDef.premiumMultiplier, | ||
| isOAuth: resolveCustomModelIsOAuth(api, providerAuth), | ||
| }; | ||
|
|
@@ -558,6 +588,8 @@ function finalizeCustomModel(model: CustomModelOverlay, options: CustomModelBuil | |
| omitMaxOutputTokens: resolvedModel.omitMaxOutputTokens ?? reference?.omitMaxOutputTokens, | ||
| compat: mergeCompat(reference?.compatConfig, resolvedModel.compat), | ||
| contextPromotionTarget: resolvedModel.contextPromotionTarget, | ||
| compactionModel: resolvedModel.compactionModel, | ||
| remoteCompaction: resolvedModel.remoteCompaction, | ||
| premiumMultiplier: resolvedModel.premiumMultiplier, | ||
| isOAuth: resolvedModel.isOAuth, | ||
| } as ModelSpec<Api>); | ||
|
|
@@ -1127,7 +1159,6 @@ export class ModelRegistry { | |
| const discoverableProviders: DiscoveryProviderConfig[] = []; | ||
| const providerEntries = Object.entries(value.providers ?? {}); | ||
| const configuredProviders = new Set(Object.keys(value.providers ?? {})); | ||
|
|
||
| for (const [providerName, providerConfig] of providerEntries) { | ||
| const resolvedProviderHeaders = resolveConfigHeaders(providerConfig.headers); | ||
| // Always set overrides when baseUrl/headers/apiKey/authHeader/compat/disableStrictTools/transport are present | ||
|
|
@@ -1138,6 +1169,7 @@ export class ModelRegistry { | |
| providerConfig.authHeader !== undefined || | ||
| providerConfig.compat || | ||
| providerConfig.disableStrictTools || | ||
| providerConfig.remoteCompaction || | ||
| providerConfig.transport | ||
| ) { | ||
| const disableStrictCompat = providerConfig.disableStrictTools ? { disableStrictTools: true } : undefined; | ||
|
|
@@ -1147,6 +1179,7 @@ export class ModelRegistry { | |
| apiKey: providerConfig.apiKey, | ||
| authHeader: providerConfig.authHeader, | ||
| compat: mergeCompat(providerConfig.compat, disableStrictCompat), | ||
| remoteCompaction: providerConfig.remoteCompaction, | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Adding this provider-level override makes discovery-backed providers look configurable, but warm-start cache loading in Useful? React with 👍 / 👎. |
||
| transport: providerConfig.transport, | ||
| }); | ||
| } | ||
|
|
@@ -1538,12 +1571,18 @@ export class ModelRegistry { | |
| authHeader: override.authHeader ?? baseOverride?.authHeader, | ||
| headers: override.headers ? { ...(baseOverride?.headers ?? {}), ...override.headers } : baseOverride?.headers, | ||
| compat: override.compat ? mergeCompat(baseOverride?.compat, override.compat) : baseOverride?.compat, | ||
| remoteCompaction: mergeRemoteCompactionConfig(baseOverride?.remoteCompaction, override.remoteCompaction), | ||
| transport: override.transport ?? baseOverride?.transport, | ||
| }; | ||
| } | ||
| #applyProviderTransportOverride<T extends { baseUrl?: string; headers?: Record<string, string> }>( | ||
| #applyProviderTransportOverride< | ||
| T extends { baseUrl?: string; headers?: Record<string, string>; remoteCompaction?: RemoteCompactionConfig<Api> }, | ||
| >( | ||
| entry: T, | ||
| override: Pick<ProviderOverride, "baseUrl" | "headers" | "authHeader" | "apiKey" | "transport">, | ||
| override: Pick< | ||
| ProviderOverride, | ||
| "baseUrl" | "headers" | "authHeader" | "apiKey" | "remoteCompaction" | "transport" | ||
| >, | ||
| ): T { | ||
| const headers = mergeAuthHeader( | ||
| override.headers ? { ...entry.headers, ...override.headers } : entry.headers, | ||
|
|
@@ -1557,6 +1596,7 @@ export class ModelRegistry { | |
| // Preserve the model's existing transport when the override omits one; | ||
| // providers without a `transport` field keep the default per-API dispatch. | ||
| ...(override.transport !== undefined ? { transport: override.transport } : {}), | ||
| remoteCompaction: mergeRemoteCompactionConfig(entry.remoteCompaction, override.remoteCompaction), | ||
| }; | ||
| } | ||
| #applyRuntimeProviderOverrides(models: Model<Api>[]): Model<Api>[] { | ||
|
|
@@ -1641,7 +1681,6 @@ export class ModelRegistry { | |
|
|
||
| #parseModels(config: ModelsConfig): CustomModelOverlay[] { | ||
| const models: CustomModelOverlay[] = []; | ||
|
|
||
| for (const [providerName, providerConfig] of Object.entries(config.providers ?? {})) { | ||
| const modelDefs = providerConfig.models ?? []; | ||
| if (modelDefs.length === 0) continue; // Override-only, no custom models | ||
|
|
@@ -1662,6 +1701,7 @@ export class ModelRegistry { | |
| providerConfig.authHeader, | ||
| providerCompat, | ||
| (providerConfig.auth as ProviderAuthMode | undefined) ?? undefined, | ||
| providerConfig.remoteCompaction, | ||
| modelDef as CustomModelDefinitionLike, | ||
| ); | ||
| if (!model) continue; | ||
|
|
@@ -2069,6 +2109,7 @@ export class ModelRegistry { | |
| config.authHeader, | ||
| config.compat, | ||
| undefined, | ||
| config.remoteCompaction, | ||
| modelDef as CustomModelDefinitionLike, | ||
| ); | ||
| if (!overlay) { | ||
|
|
@@ -2136,6 +2177,7 @@ export class ModelRegistry { | |
| providerAuthHeader, | ||
| providerCompat, | ||
| undefined, | ||
| config.remoteCompaction, | ||
| modelDef as CustomModelDefinitionLike, | ||
| ); | ||
| if (overlay) results.push(finalizeCustomModel(overlay, { useDefaults: true })); | ||
|
|
@@ -2221,6 +2263,7 @@ export interface ProviderConfigInput { | |
| streamSimple?: (model: Model<Api>, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream; | ||
| headers?: Record<string, string>; | ||
| compat?: ModelSpec<Api>["compat"]; | ||
| remoteCompaction?: RemoteCompactionConfig<Api>; | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Adding Useful? React with 👍 / 👎. |
||
| authHeader?: boolean; | ||
| /** Streaming transport override — see {@link Model.transport}. */ | ||
| transport?: Model<Api>["transport"]; | ||
|
|
@@ -2255,6 +2298,8 @@ export interface ProviderConfigInput { | |
| headers?: Record<string, string>; | ||
| compat?: ModelSpec<Api>["compat"]; | ||
| contextPromotionTarget?: string; | ||
| compactionModel?: string; | ||
| remoteCompaction?: RemoteCompactionConfig<Api>; | ||
| premiumMultiplier?: number; | ||
| }>; | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When a custom/provider model opts in with
remoteCompaction.api: "azure-openai-responses", this check sends it throughrequestOpenAiRemoteCompaction, but that helper still builds an OpenAI-style URL/header set (/v1/responses/compactandAuthorization: Bearer). I checked the repo's Azure Responses request builder inpackages/ai/src/providers/azure-openai-responses.ts, which uses anapi-keyheader and?api-version=...on${baseUrl}/responses; Azure compaction requests using the normal Azure config will therefore fail and fall back to local summarization instead of using the configured native path.Useful? React with 👍 / 👎.