mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-08-08 04:36:13 +08:00
feat(server): converge legacy compatibility (#15426)
#### PR Dependency Tree * **PR #15426** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added workspace BYOK profiles with provider/model catalogs, capability validation, connection probing, credential rotation, reordering, and secure local leases. * Added Copilot route options, selectable targets, managed tiers, explicit profile/model overrides, and improved streaming with tool callbacks and abort support. * Added Copilot availability controls to prevent access when the feature is disabled. * **Changes** * Simplified Copilot configuration and removed legacy provider-specific settings. * Removed obsolete model, token-cost, transcript strategy, and provider metadata fields from public responses. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
DarkSky
@@ -20,7 +20,9 @@ chrono = { workspace = true }
|
||||
crc32fast = "1.5.0"
|
||||
doc_extractor = { workspace = true }
|
||||
file-format = { workspace = true }
|
||||
gcp_auth = "0.12.7"
|
||||
hex = { workspace = true }
|
||||
hkdf = { workspace = true }
|
||||
hmac = "0.13"
|
||||
homedir = { workspace = true }
|
||||
image = { workspace = true }
|
||||
@@ -62,7 +64,12 @@ sqlx = { workspace = true, default-features = false, features = [
|
||||
] }
|
||||
thiserror.workspace = true
|
||||
tiktoken-rs = { workspace = true }
|
||||
tokio = { workspace = true, features = ["rt-multi-thread", "sync", "time"] }
|
||||
tokio = { workspace = true, features = [
|
||||
"net",
|
||||
"rt-multi-thread",
|
||||
"sync",
|
||||
"time",
|
||||
] }
|
||||
unicode-normalization = "0.1"
|
||||
unicode_skeleton = "0.1.1"
|
||||
url = { workspace = true }
|
||||
@@ -70,6 +77,7 @@ uuid = { workspace = true, features = ["v4"] }
|
||||
v_htmlescape = { workspace = true }
|
||||
webpki-roots = "1.0"
|
||||
y-octo = { workspace = true, features = ["large_refs"] }
|
||||
zeroize = { workspace = true }
|
||||
|
||||
[target.'cfg(not(target_os = "linux"))'.dependencies]
|
||||
mimalloc = { workspace = true }
|
||||
|
||||
Vendored
+238
-148
@@ -8,6 +8,9 @@ export declare class BackendRuntime {
|
||||
acquireCoordinationLease(key: string, owner: string, ttlMs: number): Promise<CoordinationLeaseGrant | null>
|
||||
releaseCoordinationLease(key: string, owner: string, fencingToken: bigint | number): Promise<boolean>
|
||||
renewCoordinationLease(key: string, owner: string, fencingToken: bigint | number, ttlMs: number): Promise<boolean>
|
||||
executeCopilotStream(input: CopilotExecuteInput, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise<string>)): Promise<CopilotStreamHandle>
|
||||
executeCopilot(input: CopilotExecuteInput): Promise<string>
|
||||
assertCopilotRoute(input: CopilotRouteCheckInput): Promise<void>
|
||||
/**
|
||||
* Merge pending doc updates with y-octo and persist the merged snapshot.
|
||||
*
|
||||
@@ -51,21 +54,29 @@ export declare class BackendRuntime {
|
||||
getWorkspaceInviteLink(workspaceId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null>
|
||||
getWorkspaceInviteLinkById(inviteId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null>
|
||||
revokeWorkspaceInviteLink(workspaceId: string): Promise<boolean>
|
||||
createByokLocalLease(activeKey: string, leaseId: string, payload: any, ttlMs: number): Promise<RuntimeByokLocalLeaseRecord>
|
||||
getByokLocalLease(leaseId: string): Promise<RuntimeByokLocalLeaseRecord | null>
|
||||
cleanupExpiredRuntimeStates(limit: number): Promise<number>
|
||||
refreshWorkspaceAdminStatsDirty(batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRefreshResult>
|
||||
recalibrateWorkspaceAdminStats(lastSid: number, batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRecalibrationResult>
|
||||
writeWorkspaceAdminStatsDailySnapshot(owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsSnapshotResult>
|
||||
recalibrateWorkspaceAdminStatsDaily(batchLimit: number, owner: string, leaseTtlMs: number, lockRetryTimes: number, lockRetryDelayMs: number): Promise<RuntimeWorkspaceStatsDailyRecalibrationResult>
|
||||
constructor()
|
||||
constructor(privateKey?: string | undefined | null)
|
||||
start(): Promise<void>
|
||||
stop(): Promise<void>
|
||||
reloadConfig(privateKey?: string | undefined | null): Promise<void>
|
||||
health(): Promise<BackendRuntimeHealth>
|
||||
runMigrations(): Promise<void>
|
||||
listByokProfiles(workspaceId: string): Promise<Array<ByokProfileOutput>>
|
||||
createByokProfile(input: CreateByokProfileInput): Promise<ByokProfileOutput>
|
||||
replaceByokProfile(input: ReplaceByokProfileInput): Promise<ByokProfileOutput>
|
||||
rotateByokCredential(input: RotateByokCredentialInput): Promise<ByokProfileOutput>
|
||||
probeByokProfile(input: ProbeByokProfileInput): Promise<ByokProbeResultOutput>
|
||||
probeByokDraft(input: ProbeByokDraftInput): Promise<ByokProbeResultOutput>
|
||||
deleteByokProfile(workspaceId: string, profileId: string): Promise<boolean>
|
||||
reorderByokProfiles(input: ReorderByokProfilesInput): Promise<Array<ByokProfileOutput>>
|
||||
createByokLocalLease(input: CreateByokLocalLeaseInput): Promise<ByokLocalLeaseOutput>
|
||||
}
|
||||
|
||||
export declare class LlmStreamHandle {
|
||||
export declare class CopilotStreamHandle {
|
||||
abort(): void
|
||||
}
|
||||
|
||||
@@ -107,46 +118,6 @@ export declare class Tokenizer {
|
||||
count(content: string, allowedSpecial?: Array<string> | undefined | null): number
|
||||
}
|
||||
|
||||
export interface ActionEvent {
|
||||
type: ActionEventType
|
||||
actionId: string
|
||||
actionVersion: string
|
||||
stepId?: string
|
||||
status?: ActionRunStatus
|
||||
attachment?: any
|
||||
result?: any
|
||||
errorCode?: string
|
||||
errorMessage?: string
|
||||
trace?: ActionTrace
|
||||
}
|
||||
|
||||
export type ActionEventType = 'action_start'|
|
||||
'step_start'|
|
||||
'attachment'|
|
||||
'step_end'|
|
||||
'action_done'|
|
||||
'error';
|
||||
|
||||
export type ActionRunStatus = 'created'|
|
||||
'running'|
|
||||
'succeeded'|
|
||||
'failed'|
|
||||
'aborted';
|
||||
|
||||
export interface ActionRuntimeInput {
|
||||
recipeId: string
|
||||
recipeVersion?: string
|
||||
input: any
|
||||
}
|
||||
|
||||
export interface ActionTrace {
|
||||
actionId: string
|
||||
actionVersion: string
|
||||
status: ActionRunStatus
|
||||
lightweight: Array<any>
|
||||
errorCode?: string
|
||||
}
|
||||
|
||||
export declare function activateLicense(request: LicenseKeyRequest): Promise<LicenseResponse>
|
||||
|
||||
/**
|
||||
@@ -194,6 +165,15 @@ export interface BackendRuntimeHealth {
|
||||
|
||||
export declare function buildPublicRootDoc(rootDocBin: Buffer, docMetas: Array<PublicDocMetaInput>): Buffer
|
||||
|
||||
export interface BuiltInManagedTarget {
|
||||
id: string
|
||||
displayName: string
|
||||
minimumTier: BuiltInManagedTargetTier
|
||||
}
|
||||
|
||||
export type BuiltInManagedTargetTier = 'Standard'|
|
||||
'Premium';
|
||||
|
||||
export interface BuiltInPromptRenderContract {
|
||||
name: string
|
||||
renderParams: Record<string, any>
|
||||
@@ -203,20 +183,124 @@ export interface BuiltInPromptSessionContract {
|
||||
name: string
|
||||
turns: Array<PromptMessageContract>
|
||||
renderParams: Record<string, any>
|
||||
maxTokenSize: number
|
||||
}
|
||||
|
||||
export interface BuiltInPromptSpec {
|
||||
name: string
|
||||
action?: string
|
||||
model: string
|
||||
optionalModels?: Array<string>
|
||||
config?: any
|
||||
params?: Record<string, PromptParamSpec>
|
||||
builtins?: Array<PromptBuiltin>
|
||||
messages: Array<PromptSpecMessage>
|
||||
}
|
||||
|
||||
export interface BuiltInRouteOptions {
|
||||
routeId: string
|
||||
standardDefaultTargetId?: string
|
||||
premiumDefaultTargetId?: string
|
||||
choices: Array<BuiltInManagedTarget>
|
||||
}
|
||||
|
||||
export interface ByokCapabilityInput {
|
||||
input: Array<string>
|
||||
output: Array<string>
|
||||
features: Array<string>
|
||||
attachmentKinds: Array<string>
|
||||
attachmentSources: Array<string>
|
||||
}
|
||||
|
||||
export interface ByokCatalogModelOutput {
|
||||
modelId: string
|
||||
displayName: string
|
||||
recommended: boolean
|
||||
capabilities: Array<ByokCapabilityInput>
|
||||
}
|
||||
|
||||
export interface ByokCatalogOutput {
|
||||
version: string
|
||||
providers: Array<ByokCatalogProviderOutput>
|
||||
}
|
||||
|
||||
export interface ByokCatalogProviderOutput {
|
||||
provider: string
|
||||
models: Array<ByokCatalogModelOutput>
|
||||
}
|
||||
|
||||
export interface ByokEndpointInput {
|
||||
kind: string
|
||||
url?: string
|
||||
}
|
||||
|
||||
export interface ByokLocalLeaseOutput {
|
||||
leaseId: string
|
||||
expiresAtMs: number
|
||||
}
|
||||
|
||||
export interface ByokModelDeclarationInput {
|
||||
modelId: string
|
||||
enabled: boolean
|
||||
capabilities: Array<ByokCapabilityInput>
|
||||
}
|
||||
|
||||
export interface ByokModelProbeCheckOutput {
|
||||
operation: string
|
||||
status: ByokProbeStatusOutput
|
||||
}
|
||||
|
||||
export interface ByokModelProbeOutput {
|
||||
modelId: string
|
||||
checks: Array<ByokModelProbeCheckOutput>
|
||||
}
|
||||
|
||||
export interface ByokProbeCheckInput {
|
||||
modelId: string
|
||||
operation: string
|
||||
}
|
||||
|
||||
export interface ByokProbeResultOutput {
|
||||
definitionFingerprint: string
|
||||
stale: boolean
|
||||
connection: ByokProbeStatusOutput
|
||||
models: Array<ByokModelProbeOutput>
|
||||
}
|
||||
|
||||
export interface ByokProbeStatusOutput {
|
||||
kind: string
|
||||
testedAtMs?: number
|
||||
errorKind?: string
|
||||
}
|
||||
|
||||
export interface ByokProfileDefinitionInput {
|
||||
version: number
|
||||
endpoint: ByokEndpointInput
|
||||
models: Array<ByokModelDeclarationInput>
|
||||
}
|
||||
|
||||
export interface ByokProfileOrderInput {
|
||||
profileId: string
|
||||
expectedRevision: number
|
||||
}
|
||||
|
||||
export interface ByokProfileOutput {
|
||||
profileId: string
|
||||
workspaceId: string
|
||||
provider: string
|
||||
name: string
|
||||
description?: string
|
||||
definition: ByokProfileDefinitionInput
|
||||
enabled: boolean
|
||||
sortOrder: number
|
||||
revision: number
|
||||
validation?: ByokValidationOutput
|
||||
}
|
||||
|
||||
export interface ByokValidationOutput {
|
||||
definitionFingerprint: string
|
||||
credentialGeneration: number
|
||||
connection: ByokProbeStatusOutput
|
||||
models: Array<ByokModelProbeOutput>
|
||||
}
|
||||
|
||||
export interface CanonicalChatRequestContract {
|
||||
model: string
|
||||
messages: Array<PromptMessageContract>
|
||||
@@ -315,8 +399,74 @@ export interface CoordinationLeaseGrant {
|
||||
fencingToken: bigint | number
|
||||
}
|
||||
|
||||
export interface CopilotAccessProjection {
|
||||
routeAllowed: boolean
|
||||
managedTier: CopilotManagedTier
|
||||
serverByok: boolean
|
||||
localByok: boolean
|
||||
}
|
||||
|
||||
export declare function copilotActionRecipe(actionId: string, actionVersion?: string | undefined | null): string
|
||||
|
||||
export interface CopilotExecuteInput {
|
||||
slot: string
|
||||
builtInRouteId?: string
|
||||
workspaceId?: string
|
||||
userId?: string
|
||||
localLeaseId?: string
|
||||
access: CopilotAccessProjection
|
||||
managedTargetId?: string
|
||||
targetOverride?: CopilotTargetOverrideInput
|
||||
request: unknown
|
||||
}
|
||||
|
||||
export type CopilotManagedTier = 'Standard'|
|
||||
'Premium';
|
||||
|
||||
export interface CopilotRouteCheckInput {
|
||||
slot: string
|
||||
builtInRouteId?: string
|
||||
workspaceId?: string
|
||||
userId?: string
|
||||
localLeaseId?: string
|
||||
access: CopilotAccessProjection
|
||||
managedTargetId?: string
|
||||
targetOverride?: CopilotTargetOverrideInput
|
||||
}
|
||||
|
||||
export interface CopilotTargetOverrideInput {
|
||||
profileId: string
|
||||
modelId: string
|
||||
}
|
||||
|
||||
export declare function createAuthSessionRefreshToken(): AuthSessionRefreshToken
|
||||
|
||||
export interface CreateByokLocalLeaseInput {
|
||||
workspaceId: string
|
||||
userId: string
|
||||
providers: Array<CreateByokLocalLeaseProviderInput>
|
||||
}
|
||||
|
||||
export interface CreateByokLocalLeaseProviderInput {
|
||||
provider: string
|
||||
name: string
|
||||
description?: string
|
||||
credential: string
|
||||
definition: ByokProfileDefinitionInput
|
||||
enabled: boolean
|
||||
}
|
||||
|
||||
export interface CreateByokProfileInput {
|
||||
workspaceId: string
|
||||
provider: string
|
||||
name: string
|
||||
description?: string
|
||||
credential: string
|
||||
definition: ByokProfileDefinitionInput
|
||||
enabled: boolean
|
||||
actorUserId: string
|
||||
}
|
||||
|
||||
/**
|
||||
* Converts markdown content to AFFiNE-compatible y-octo document binary.
|
||||
*
|
||||
@@ -409,31 +559,11 @@ export declare function llmBuildRerankRequest(request: LlmRerankRequestContract)
|
||||
|
||||
export declare function llmCanonicalJsonSchemaHash(schema: any): string
|
||||
|
||||
export declare function llmCollectPromptMetadata(request: PromptMetadataContract): PromptMetadataResult
|
||||
|
||||
export declare function llmCompileExecutionPlan(value: any): any
|
||||
|
||||
export interface LlmCoreMessage {
|
||||
role: string
|
||||
content: Array<any>
|
||||
}
|
||||
|
||||
export declare function llmCountPromptTokens(request: PromptTokenCountContract): PromptTokenCountResult
|
||||
|
||||
export declare function llmDispatchPrepared(routesJson: string): Promise<string>
|
||||
|
||||
export declare function llmDispatchPreparedStream(routesJson: string, callback: ((err: Error | null, arg: string) => void)): LlmStreamHandle
|
||||
|
||||
export declare function llmDispatchToolLoopStream(protocol: string, backendConfigJson: string, requestJson: string, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise<string>)): LlmStreamHandle
|
||||
|
||||
export declare function llmDispatchToolLoopStreamPrepared(routesJson: string, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise<string>)): LlmStreamHandle
|
||||
|
||||
export declare function llmDispatchToolLoopStreamRouted(routesJson: string, requestJson: string, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise<string>)): LlmStreamHandle
|
||||
|
||||
export declare function llmEmbeddingDispatch(protocol: string, backendConfigJson: string, requestJson: string): Promise<string>
|
||||
|
||||
export declare function llmEmbeddingDispatchPrepared(routesJson: string): Promise<string>
|
||||
|
||||
export interface LlmEmbeddingRequestContract {
|
||||
model: string
|
||||
inputs: Array<string>
|
||||
@@ -443,9 +573,11 @@ export interface LlmEmbeddingRequestContract {
|
||||
|
||||
export declare function llmGetBuiltInPromptSpec(name: string): BuiltInPromptSpec | null
|
||||
|
||||
export declare function llmGetContractSchema(name: string): any
|
||||
export declare function llmGetBuiltInRouteOptions(name: string): BuiltInRouteOptions | null
|
||||
|
||||
export declare function llmImageDispatchPrepared(routesJson: string): Promise<string>
|
||||
export declare function llmGetByokCatalog(): ByokCatalogOutput
|
||||
|
||||
export declare function llmGetContractSchema(name: string): any
|
||||
|
||||
export interface LlmImageInputContract {
|
||||
kind: 'url' | 'data' | 'bytes'
|
||||
@@ -488,7 +620,6 @@ export interface LlmImageProviderOptionsContract {
|
||||
|
||||
export interface LlmImageRequestBuildContract {
|
||||
model: string
|
||||
protocol: 'openai_chat' | 'openai_responses' | 'openai_images' | 'anthropic' | 'gemini' | 'fal_image'
|
||||
messages: Array<PromptMessageContract>
|
||||
options?: any
|
||||
}
|
||||
@@ -511,18 +642,10 @@ export declare function llmMatchModelCapabilities(payload: CapabilityMatchReques
|
||||
|
||||
export declare function llmMatchModelRegistry(request: ModelRegistryMatchRequest): ModelRegistryMatchResponse
|
||||
|
||||
export declare function llmNormalizePreparedRoutes(value: any): any
|
||||
|
||||
export declare function llmPlanAttachmentReference(protocol: string, backendConfigJson: string, sourceJson: string): string
|
||||
|
||||
export declare function llmRenderBuiltInPrompt(request: BuiltInPromptRenderContract): PromptRenderResult
|
||||
|
||||
export declare function llmRenderBuiltInSessionPrompt(request: BuiltInPromptSessionContract): PromptSessionResult
|
||||
|
||||
export declare function llmRenderPrompt(request: PromptRenderContract): PromptRenderResult
|
||||
|
||||
export declare function llmRenderSessionPrompt(request: PromptSessionContract): PromptSessionResult
|
||||
|
||||
export interface LlmRequestContract {
|
||||
model: string
|
||||
messages: Array<LlmCoreMessage>
|
||||
@@ -537,10 +660,6 @@ export interface LlmRequestContract {
|
||||
middleware?: any
|
||||
}
|
||||
|
||||
export declare function llmRerankDispatch(protocol: string, backendConfigJson: string, requestJson: string): Promise<string>
|
||||
|
||||
export declare function llmRerankDispatchPrepared(routesJson: string): Promise<string>
|
||||
|
||||
export interface LlmRerankRequestContract {
|
||||
model: string
|
||||
query: string
|
||||
@@ -550,14 +669,6 @@ export interface LlmRerankRequestContract {
|
||||
|
||||
export declare function llmResolveModelRegistryVariant(request: ModelRegistryResolveRequest): ModelRegistryResolveResponse
|
||||
|
||||
export declare function llmResolveRequestedModelMatch(payload: RequestedModelMatchRequest): RequestedModelMatchResponse
|
||||
|
||||
export declare function llmResolveRequestIntent(protocol: string, backendConfigJson: string, intentJson: string): string
|
||||
|
||||
export declare function llmStructuredDispatch(protocol: string, backendConfigJson: string, requestJson: string): Promise<string>
|
||||
|
||||
export declare function llmStructuredDispatchPrepared(routesJson: string): Promise<string>
|
||||
|
||||
export interface LlmStructuredRequestContract {
|
||||
model: string
|
||||
messages: Array<LlmCoreMessage>
|
||||
@@ -695,6 +806,22 @@ export interface PortalResponse {
|
||||
error?: LicenseError
|
||||
}
|
||||
|
||||
export interface ProbeByokDraftInput {
|
||||
workspaceId: string
|
||||
provider: string
|
||||
credential?: string
|
||||
profileId?: string
|
||||
expectedRevision?: number
|
||||
definition: ByokProfileDefinitionInput
|
||||
checks: Array<ByokProbeCheckInput>
|
||||
}
|
||||
|
||||
export interface ProbeByokProfileInput {
|
||||
workspaceId: string
|
||||
profileId: string
|
||||
checks: Array<ByokProbeCheckInput>
|
||||
}
|
||||
|
||||
export declare function processImage(input: Buffer, maxEdge: number, keepExif: boolean): Promise<Buffer>
|
||||
|
||||
export type PromptBuiltin = 'Date'|
|
||||
@@ -705,10 +832,6 @@ export type PromptBuiltin = 'Date'|
|
||||
'HasSelected'|
|
||||
'HasCurrentDoc';
|
||||
|
||||
export interface PromptCountMessage {
|
||||
content: string
|
||||
}
|
||||
|
||||
export interface PromptMessageContract {
|
||||
role: 'system' | 'assistant' | 'user'
|
||||
content: string
|
||||
@@ -717,46 +840,16 @@ export interface PromptMessageContract {
|
||||
responseFormat?: PromptStructuredResponseContract
|
||||
}
|
||||
|
||||
export interface PromptMetadataContract {
|
||||
messages: Array<PromptMessageContract>
|
||||
}
|
||||
|
||||
export interface PromptMetadataResult {
|
||||
paramKeys: Array<string>
|
||||
templateParams: Record<string, any>
|
||||
}
|
||||
|
||||
export interface PromptParamSpec {
|
||||
default?: string
|
||||
enumValues?: Array<string>
|
||||
}
|
||||
|
||||
export interface PromptRenderContract {
|
||||
messages: Array<PromptMessageContract>
|
||||
templateParams: Record<string, any>
|
||||
renderParams: Record<string, any>
|
||||
}
|
||||
|
||||
export interface PromptRenderResult {
|
||||
messages: Array<PromptMessageContract>
|
||||
warnings: Array<string>
|
||||
}
|
||||
|
||||
export interface PromptSessionContract {
|
||||
prompt: PromptSessionPrompt
|
||||
turns: Array<PromptMessageContract>
|
||||
renderParams: Record<string, any>
|
||||
maxTokenSize: number
|
||||
}
|
||||
|
||||
export interface PromptSessionPrompt {
|
||||
action?: string
|
||||
model?: string
|
||||
promptTokens: number
|
||||
templateParams: Record<string, any>
|
||||
messages: Array<PromptMessageContract>
|
||||
}
|
||||
|
||||
export interface PromptSessionResult {
|
||||
messages: Array<PromptMessageContract>
|
||||
warnings: Array<string>
|
||||
@@ -775,15 +868,6 @@ export interface PromptStructuredResponseContract {
|
||||
strict?: boolean
|
||||
}
|
||||
|
||||
export interface PromptTokenCountContract {
|
||||
model?: string
|
||||
messages: Array<PromptCountMessage>
|
||||
}
|
||||
|
||||
export interface PromptTokenCountResult {
|
||||
tokens: number
|
||||
}
|
||||
|
||||
export interface ProviderDriverSpec {
|
||||
driverId: string
|
||||
providerType: string
|
||||
@@ -838,16 +922,22 @@ export interface RemoteMimeTypeRequest {
|
||||
timeoutMs?: number
|
||||
}
|
||||
|
||||
export interface RequestedModelMatchRequest {
|
||||
providerIds: Array<string>
|
||||
optionalModels: Array<string>
|
||||
requestedModelId?: string
|
||||
defaultModel?: string
|
||||
export interface ReorderByokProfilesInput {
|
||||
workspaceId: string
|
||||
profiles: Array<ByokProfileOrderInput>
|
||||
actorUserId: string
|
||||
}
|
||||
|
||||
export interface RequestedModelMatchResponse {
|
||||
selectedModel?: string
|
||||
matchedOptionalModel: boolean
|
||||
export interface ReplaceByokProfileInput {
|
||||
workspaceId: string
|
||||
profileId: string
|
||||
expectedRevision: number
|
||||
name: string
|
||||
description?: string
|
||||
definition: ByokProfileDefinitionInput
|
||||
credential?: string
|
||||
enabled: boolean
|
||||
actorUserId: string
|
||||
}
|
||||
|
||||
export interface RerankCandidate {
|
||||
@@ -896,7 +986,13 @@ export interface ResolveEntitlementInput {
|
||||
|
||||
export declare function resolveEntitlementV1(input: ResolveEntitlementInput): ResolvedEntitlement
|
||||
|
||||
export declare function runNativeActionRecipePreparedStream(input: ActionRuntimeInput, callback: ((err: Error | null, arg: string) => void)): LlmStreamHandle
|
||||
export interface RotateByokCredentialInput {
|
||||
workspaceId: string
|
||||
profileId: string
|
||||
expectedRevision: number
|
||||
credential: string
|
||||
actorUserId: string
|
||||
}
|
||||
|
||||
export interface RuntimeBlobCleanupExecuteResult {
|
||||
scannedCandidates: number
|
||||
@@ -943,12 +1039,6 @@ export interface RuntimeBlobMetadataBackfillResult {
|
||||
workspaceIds: Array<string>
|
||||
}
|
||||
|
||||
export interface RuntimeByokLocalLeaseRecord {
|
||||
leaseId: string
|
||||
payload: any
|
||||
expiresAtMs: number
|
||||
}
|
||||
|
||||
export interface RuntimeDocBlobRefsResult {
|
||||
scannedDocs: number
|
||||
parsedDocs: number
|
||||
|
||||
@@ -1,291 +1,98 @@
|
||||
use std::collections::HashSet;
|
||||
|
||||
use jsonschema::Draft;
|
||||
use napi::{Error, Result, Status};
|
||||
use serde::Serialize;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::{
|
||||
super::contract_schema::{transcript_input_schema, transcript_result_schema},
|
||||
ActionRecipe, ActionRecipeStep, ActionStepKind,
|
||||
};
|
||||
|
||||
fn invalid_recipe(message: impl Into<String>) -> Error {
|
||||
Error::new(Status::InvalidArg, message.into())
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ActionRecipe<'a> {
|
||||
action_id: &'a str,
|
||||
action_version: &'a str,
|
||||
slot: &'a str,
|
||||
prompt_ref: &'a str,
|
||||
response_contract: Value,
|
||||
output_projection: &'a str,
|
||||
}
|
||||
|
||||
pub fn built_in_recipes() -> Vec<ActionRecipe> {
|
||||
vec![
|
||||
action_recipe("mindmap.generate", "v1"),
|
||||
action_recipe("slides.outline", "v1"),
|
||||
action_recipe("image.filter.sketch", "v1"),
|
||||
action_recipe("image.filter.clay", "v1"),
|
||||
action_recipe("image.filter.anime", "v1"),
|
||||
action_recipe("image.filter.pixel", "v1"),
|
||||
transcript_recipe("transcript.audio.gemini", "v1"),
|
||||
]
|
||||
}
|
||||
|
||||
pub fn find_recipe(id: &str, version: Option<&str>) -> Result<ActionRecipe> {
|
||||
let catalog = load_catalog()?;
|
||||
catalog
|
||||
.into_iter()
|
||||
.find(|recipe| recipe.id == id && version.is_none_or(|version| recipe.version == version))
|
||||
.ok_or_else(|| {
|
||||
invalid_recipe(format!(
|
||||
"Action recipe not found: {}{}",
|
||||
id,
|
||||
version.map(|version| format!("@{version}")).unwrap_or_default()
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn load_catalog() -> Result<Vec<ActionRecipe>> {
|
||||
let recipes = built_in_recipes();
|
||||
validate_catalog(&recipes)?;
|
||||
Ok(recipes)
|
||||
}
|
||||
|
||||
pub fn validate_catalog(recipes: &[ActionRecipe]) -> Result<()> {
|
||||
let mut keys = HashSet::new();
|
||||
for recipe in recipes {
|
||||
validate_recipe(recipe)?;
|
||||
let key = format!("{}@{}", recipe.id, recipe.version);
|
||||
if !keys.insert(key.clone()) {
|
||||
return Err(invalid_recipe(format!("Duplicated action recipe: {key}")));
|
||||
}
|
||||
#[napi_derive::napi]
|
||||
pub fn copilot_action_recipe(action_id: String, action_version: Option<String>) -> Result<String> {
|
||||
let version = action_version.as_deref().unwrap_or("v1");
|
||||
if version != "v1" {
|
||||
return Err(Error::new(Status::InvalidArg, "Action recipe not found"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn validate_recipe(recipe: &ActionRecipe) -> Result<()> {
|
||||
if recipe.id.trim().is_empty() {
|
||||
return Err(invalid_recipe("Action recipe id is required"));
|
||||
}
|
||||
if recipe.version.trim().is_empty() {
|
||||
return Err(invalid_recipe("Action recipe version is required"));
|
||||
}
|
||||
if recipe.steps.is_empty() {
|
||||
return Err(invalid_recipe(format!(
|
||||
"Action recipe {}@{} must declare at least one step",
|
||||
recipe.id, recipe.version
|
||||
)));
|
||||
}
|
||||
compile_schema("inputSchema", &recipe.input_schema)?;
|
||||
compile_schema("outputSchema", &recipe.output_schema)?;
|
||||
|
||||
let mut step_ids = HashSet::new();
|
||||
let mut has_final = false;
|
||||
for step in &recipe.steps {
|
||||
if step.id.trim().is_empty() {
|
||||
return Err(invalid_recipe(format!(
|
||||
"Action recipe {}@{} contains a step without id",
|
||||
recipe.id, recipe.version
|
||||
)));
|
||||
}
|
||||
if !step_ids.insert(step.id.clone()) {
|
||||
return Err(invalid_recipe(format!(
|
||||
"Action recipe {}@{} contains duplicated step id {}",
|
||||
recipe.id, recipe.version, step.id
|
||||
)));
|
||||
}
|
||||
if step.kind == ActionStepKind::Final {
|
||||
has_final = true;
|
||||
}
|
||||
}
|
||||
if !has_final {
|
||||
return Err(invalid_recipe(format!(
|
||||
"Action recipe {}@{} must end with a final step",
|
||||
recipe.id, recipe.version
|
||||
)));
|
||||
}
|
||||
if recipe
|
||||
.steps
|
||||
.last()
|
||||
.is_some_and(|step| step.kind != ActionStepKind::Final)
|
||||
{
|
||||
return Err(invalid_recipe(format!(
|
||||
"Action recipe {}@{} must end with a final step",
|
||||
recipe.id, recipe.version
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn compile_schema(label: &str, schema: &Value) -> Result<()> {
|
||||
jsonschema::options()
|
||||
.with_draft(Draft::Draft7)
|
||||
.build(schema)
|
||||
.map(|_| ())
|
||||
.map_err(|error| invalid_recipe(format!("Invalid action recipe {label}: {error}")))
|
||||
}
|
||||
|
||||
fn action_recipe(id: &str, version: &str) -> ActionRecipe {
|
||||
let steps = if id.starts_with("image.filter.") {
|
||||
vec![
|
||||
ActionRecipeStep {
|
||||
id: "generate-image".to_string(),
|
||||
kind: ActionStepKind::PromptImage,
|
||||
input: Some(json!({
|
||||
"preparedRoutes": { "$state": "preparedRoutes.generate-image" },
|
||||
"outputKey": "artifact"
|
||||
})),
|
||||
state_patch: Some(json!({ "imageGenerated": true })),
|
||||
let recipe = match action_id.as_str() {
|
||||
"mindmap.generate" => structured(&action_id, "mindmap.generate", text_result_schema()),
|
||||
"slides.outline" => structured(&action_id, "slides.outline", text_result_schema()),
|
||||
"transcript.audio" => structured(
|
||||
&action_id,
|
||||
"Transcript audio structured",
|
||||
super::super::contract_schema::transcript_result_schema(),
|
||||
),
|
||||
"image.filter.sketch" | "image.filter.clay" | "image.filter.anime" | "image.filter.pixel" => ActionRecipe {
|
||||
action_id: &action_id,
|
||||
action_version: "v1",
|
||||
slot: match action_id.as_str() {
|
||||
"image.filter.sketch" => "action.image.filter.sketch",
|
||||
"image.filter.clay" => "action.image.filter.clay",
|
||||
"image.filter.anime" => "action.image.filter.anime",
|
||||
_ => "action.image.filter.pixel",
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({
|
||||
"copy": { "$state": "artifact" }
|
||||
})),
|
||||
state_patch: Some(json!({ "finalized": true })),
|
||||
},
|
||||
]
|
||||
} else if id == "slides.outline" {
|
||||
vec![
|
||||
ActionRecipeStep {
|
||||
id: "generate-structured".to_string(),
|
||||
kind: ActionStepKind::PromptStructured,
|
||||
input: Some(json!({
|
||||
"preparedRoutes": { "$state": "preparedRoutes.generate" },
|
||||
"unwrapKey": "result",
|
||||
"outputKey": "generated"
|
||||
})),
|
||||
state_patch: Some(json!({ "generatedAt": "promptStructured" })),
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "validate-json".to_string(),
|
||||
kind: ActionStepKind::ValidateJson,
|
||||
input: Some(json!({
|
||||
"value": { "$state": "generated" },
|
||||
"schema": text_action_output_schema()
|
||||
})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "project-outline".to_string(),
|
||||
kind: ActionStepKind::Transform,
|
||||
input: Some(json!({
|
||||
"slidesOutlineMarkdown": { "$state": "generated" },
|
||||
"outputKey": "outlineMarkdown"
|
||||
})),
|
||||
state_patch: Some(json!({ "projectedAt": "slidesOutlineMarkdown" })),
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({
|
||||
"copy": { "$state": "outlineMarkdown" }
|
||||
})),
|
||||
state_patch: Some(json!({ "finalized": true })),
|
||||
},
|
||||
]
|
||||
} else {
|
||||
vec![
|
||||
ActionRecipeStep {
|
||||
id: "generate-structured".to_string(),
|
||||
kind: ActionStepKind::PromptStructured,
|
||||
input: Some(json!({
|
||||
"preparedRoutes": { "$state": "preparedRoutes.generate" },
|
||||
"unwrapKey": "result",
|
||||
"outputKey": "generated"
|
||||
})),
|
||||
state_patch: Some(json!({ "generatedAt": "promptStructured" })),
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "validate-json".to_string(),
|
||||
kind: ActionStepKind::ValidateJson,
|
||||
input: Some(json!({
|
||||
"value": { "$state": "generated" },
|
||||
"schema": text_action_output_schema()
|
||||
})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({
|
||||
"copy": { "$state": "generated" }
|
||||
})),
|
||||
state_patch: Some(json!({ "finalized": true })),
|
||||
},
|
||||
]
|
||||
prompt_ref: &action_id,
|
||||
response_contract: Value::Null,
|
||||
output_projection: "first_image",
|
||||
},
|
||||
_ => return Err(Error::new(Status::InvalidArg, "Action recipe not found")),
|
||||
};
|
||||
|
||||
recipe(id, version, action_output_schema(id), steps)
|
||||
serde_json::to_string(&recipe).map_err(|error| Error::new(Status::GenericFailure, error.to_string()))
|
||||
}
|
||||
|
||||
fn transcript_recipe(id: &str, version: &str) -> ActionRecipe {
|
||||
let mut recipe = recipe(
|
||||
id,
|
||||
version,
|
||||
transcript_result_schema(),
|
||||
vec![
|
||||
ActionRecipeStep {
|
||||
id: "transcribe".to_string(),
|
||||
kind: ActionStepKind::PromptStructured,
|
||||
input: Some(json!({
|
||||
"preparedRoutes": { "$state": "preparedRoutes.transcribe" },
|
||||
"outputKey": "transcriptResult"
|
||||
})),
|
||||
state_patch: Some(json!({ "transcribedAt": "promptStructured" })),
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({
|
||||
"sourceAudio": { "$state": "sourceAudio" },
|
||||
"quality": { "$state": "quality" },
|
||||
"infos": { "$state": "infos" },
|
||||
"sliceManifest": { "$state": "sliceManifest" },
|
||||
"normalizedSegments": { "$state": "transcriptResult.normalizedSegments" },
|
||||
"normalizedTranscript": { "$state": "transcriptResult.normalizedTranscript" },
|
||||
"summaryJson": { "$state": "transcriptResult.summaryJson" },
|
||||
"providerMeta": { "$state": "transcriptResult.providerMeta" },
|
||||
"version": "transcript-result-v1",
|
||||
"strategy": id.strip_prefix("transcript.audio.").unwrap_or(id)
|
||||
})),
|
||||
state_patch: Some(json!({ "finalized": true })),
|
||||
},
|
||||
],
|
||||
);
|
||||
recipe.input_schema = transcript_input_schema();
|
||||
recipe
|
||||
}
|
||||
|
||||
fn action_output_schema(id: &str) -> Value {
|
||||
if id.starts_with("image.filter.") {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"url": { "type": "string" },
|
||||
"data_base64": { "type": "string" },
|
||||
"media_type": { "type": "string" }
|
||||
},
|
||||
"anyOf": [
|
||||
{ "required": ["url"] },
|
||||
{ "required": ["data_base64", "media_type"] }
|
||||
],
|
||||
"additionalProperties": true
|
||||
})
|
||||
} else {
|
||||
text_action_output_schema()
|
||||
}
|
||||
}
|
||||
|
||||
fn text_action_output_schema() -> Value {
|
||||
fn text_result_schema() -> Value {
|
||||
json!({
|
||||
"type": "string",
|
||||
"minLength": 1
|
||||
"type": "object",
|
||||
"properties": { "result": { "type": "string", "minLength": 1 } },
|
||||
"required": ["result"],
|
||||
"additionalProperties": false
|
||||
})
|
||||
}
|
||||
|
||||
fn recipe(id: &str, version: &str, output_schema: Value, steps: Vec<ActionRecipeStep>) -> ActionRecipe {
|
||||
fn structured<'a>(action_id: &'a str, prompt_ref: &'a str, schema: Value) -> ActionRecipe<'a> {
|
||||
ActionRecipe {
|
||||
id: id.to_string(),
|
||||
version: version.to_string(),
|
||||
input_schema: json!({}),
|
||||
output_schema,
|
||||
steps,
|
||||
action_id,
|
||||
action_version: "v1",
|
||||
slot: match action_id {
|
||||
"mindmap.generate" => "action.mindmap.generate",
|
||||
"slides.outline" => "action.slides.outline",
|
||||
"transcript.audio" => "transcript.audio",
|
||||
_ => unreachable!("structured recipe action is validated by the catalog"),
|
||||
},
|
||||
prompt_ref,
|
||||
response_contract: json!({ "schema": schema, "strict": true }),
|
||||
output_projection: if action_id == "slides.outline" {
|
||||
"slides_outline_markdown"
|
||||
} else if action_id == "transcript.audio" {
|
||||
"transcript_result"
|
||||
} else {
|
||||
"identity"
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::copilot_action_recipe;
|
||||
|
||||
#[test]
|
||||
fn recipes_only_expose_slot_prompt_contract_and_projection() {
|
||||
for (id, slot) in [
|
||||
("mindmap.generate", "action.mindmap.generate"),
|
||||
("slides.outline", "action.slides.outline"),
|
||||
("image.filter.sketch", "action.image.filter.sketch"),
|
||||
("transcript.audio", "transcript.audio"),
|
||||
] {
|
||||
let recipe = copilot_action_recipe(id.to_string(), None).unwrap();
|
||||
assert!(!recipe.contains("prepared"));
|
||||
assert!(recipe.contains(&format!("\"slot\":\"{slot}\"")));
|
||||
assert!(recipe.contains("\"promptRef\""));
|
||||
assert!(recipe.contains("\"outputProjection\""));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,152 +1,7 @@
|
||||
use napi_derive::napi;
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ActionRecipe {
|
||||
pub id: String,
|
||||
pub version: String,
|
||||
pub input_schema: Value,
|
||||
pub output_schema: Value,
|
||||
pub steps: Vec<ActionRecipeStep>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ActionRecipeStep {
|
||||
pub id: String,
|
||||
pub kind: ActionStepKind,
|
||||
#[serde(default)]
|
||||
pub input: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub state_patch: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum ActionStepKind {
|
||||
PromptStructured,
|
||||
PromptImage,
|
||||
ValidateJson,
|
||||
Transform,
|
||||
Final,
|
||||
}
|
||||
|
||||
#[napi(string_enum = "snake_case")]
|
||||
#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ActionEventType {
|
||||
ActionStart,
|
||||
StepStart,
|
||||
Attachment,
|
||||
StepEnd,
|
||||
ActionDone,
|
||||
Error,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ActionEvent {
|
||||
#[serde(rename = "type")]
|
||||
#[napi(js_name = "type")]
|
||||
pub event_type: ActionEventType,
|
||||
pub action_id: String,
|
||||
pub action_version: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub step_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub status: Option<ActionRunStatus>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub attachment: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub result: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error_code: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error_message: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub trace: Option<ActionTrace>,
|
||||
}
|
||||
|
||||
#[napi(string_enum = "snake_case")]
|
||||
#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ActionRunStatus {
|
||||
Created,
|
||||
Running,
|
||||
Succeeded,
|
||||
Failed,
|
||||
Aborted,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ActionRuntimeInput {
|
||||
pub recipe_id: String,
|
||||
#[serde(default)]
|
||||
pub recipe_version: Option<String>,
|
||||
#[serde(default)]
|
||||
pub input: Value,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ActionRuntimeOutput {
|
||||
pub result: Value,
|
||||
pub status: ActionRunStatus,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error_code: Option<String>,
|
||||
pub state: Value,
|
||||
pub steps: Vec<ActionStepRuntimeState>,
|
||||
pub trace: ActionTrace,
|
||||
pub events: Vec<ActionEvent>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ActionStepRuntimeState {
|
||||
pub id: String,
|
||||
pub input: Value,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub state_patch: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<ActionStepError>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ActionStepError {
|
||||
pub code: String,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ActionTrace {
|
||||
pub action_id: String,
|
||||
pub action_version: String,
|
||||
pub status: ActionRunStatus,
|
||||
#[serde(default)]
|
||||
pub lightweight: Vec<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error_code: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
@@ -159,8 +14,6 @@ pub struct TranscriptInputContract {
|
||||
pub infos: Option<Vec<TranscriptAudioInfo>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub slice_manifest: Option<Vec<TranscriptSliceManifestItem>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub prepared_routes: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
|
||||
@@ -232,8 +85,6 @@ pub struct TranscriptGeneratedResult {
|
||||
pub normalized_transcript: String,
|
||||
#[schemars(required)]
|
||||
pub summary_json: Option<MeetingSummary>,
|
||||
#[schemars(required)]
|
||||
pub provider_meta: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
|
||||
@@ -253,8 +104,5 @@ pub struct TranscriptResult {
|
||||
pub normalized_transcript: String,
|
||||
#[schemars(required)]
|
||||
pub summary_json: Option<MeetingSummary>,
|
||||
#[schemars(required)]
|
||||
pub provider_meta: Option<Value>,
|
||||
pub version: String,
|
||||
pub strategy: String,
|
||||
}
|
||||
|
||||
@@ -1,99 +1,5 @@
|
||||
mod catalog;
|
||||
mod contract;
|
||||
mod runtime;
|
||||
mod slides_outline;
|
||||
|
||||
use std::sync::{Arc, atomic::AtomicBool, mpsc};
|
||||
|
||||
#[cfg(test)]
|
||||
use catalog::{load_catalog, validate_catalog, validate_recipe};
|
||||
use contract::{
|
||||
ActionEvent, ActionEventType, ActionRecipe, ActionRecipeStep, ActionRunStatus, ActionRuntimeInput,
|
||||
ActionRuntimeOutput, ActionStepError, ActionStepKind, ActionStepRuntimeState, ActionTrace,
|
||||
};
|
||||
pub use catalog::copilot_action_recipe;
|
||||
pub(crate) use contract::{TranscriptGeneratedResult, TranscriptInputContract, TranscriptResult};
|
||||
use napi::{
|
||||
Result,
|
||||
threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode},
|
||||
};
|
||||
#[cfg(test)]
|
||||
use runtime::{ACTION_ABORTED_ERROR_CODE, run_action_recipe_for_test, run_action_recipe_for_test_with_control};
|
||||
use runtime::{ActionRuntimeControl, run_action_recipe_prepared_with_control};
|
||||
|
||||
use crate::llm::{LlmStreamHandle, STREAM_END_MARKER};
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn run_native_action_recipe_prepared_stream(
|
||||
input: ActionRuntimeInput,
|
||||
callback: ThreadsafeFunction<String, ()>,
|
||||
) -> Result<LlmStreamHandle> {
|
||||
let action_id = input.recipe_id.clone();
|
||||
let action_version = input.recipe_version.clone().unwrap_or_default();
|
||||
let aborted = Arc::new(AtomicBool::new(false));
|
||||
let aborted_in_worker = aborted.clone();
|
||||
let (event_sender, event_receiver) = mpsc::channel::<ActionEvent>();
|
||||
let error_sender = event_sender.clone();
|
||||
|
||||
std::thread::spawn(move || {
|
||||
if let Err(error) = run_action_recipe_prepared_with_control(
|
||||
input,
|
||||
ActionRuntimeControl {
|
||||
abort_signal: Some(aborted_in_worker.clone()),
|
||||
event_sender: Some(event_sender),
|
||||
#[cfg(test)]
|
||||
abort_after_events: None,
|
||||
#[cfg(test)]
|
||||
mock_output: None,
|
||||
},
|
||||
) {
|
||||
let _ = error_sender.send(ActionEvent {
|
||||
event_type: ActionEventType::Error,
|
||||
action_id,
|
||||
action_version,
|
||||
step_id: None,
|
||||
status: Some(ActionRunStatus::Failed),
|
||||
attachment: None,
|
||||
result: None,
|
||||
error_code: Some("action_runtime_error".to_string()),
|
||||
error_message: Some(error.reason.clone()),
|
||||
trace: None,
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
std::thread::spawn(move || {
|
||||
for event in event_receiver {
|
||||
match serde_json::to_string(&event) {
|
||||
Ok(event) => {
|
||||
let _ = callback.call(Ok(event), ThreadsafeFunctionCallMode::NonBlocking);
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = callback.call(
|
||||
Ok(
|
||||
serde_json::json!({
|
||||
"type": "error",
|
||||
"actionId": event.action_id,
|
||||
"actionVersion": event.action_version,
|
||||
"errorCode": "action_event_encode_failed",
|
||||
"errorMessage": error.to_string()
|
||||
})
|
||||
.to_string(),
|
||||
),
|
||||
ThreadsafeFunctionCallMode::NonBlocking,
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let _ = callback.call(
|
||||
Ok(STREAM_END_MARKER.to_string()),
|
||||
ThreadsafeFunctionCallMode::NonBlocking,
|
||||
);
|
||||
});
|
||||
|
||||
Ok(LlmStreamHandle { aborted })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
@@ -1,564 +0,0 @@
|
||||
use std::{
|
||||
cell::Cell,
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
mpsc::Sender,
|
||||
},
|
||||
time::Instant,
|
||||
};
|
||||
|
||||
use llm_runtime::{
|
||||
RecipeDefinition, RecipeRuntimeEvent, RecipeRuntimeOutput, RecipeRuntimeStatus, RecipeStepExecution,
|
||||
RecipeStepExecutor, StepExecutionError, execute_transform_step, execute_validate_json_step, resolve_state_ref,
|
||||
run_recipe_runtime, validate_json_schema,
|
||||
};
|
||||
use napi::{Error, Result, Status};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::{
|
||||
ActionEvent, ActionEventType, ActionRecipe, ActionRunStatus, ActionRuntimeInput, ActionRuntimeOutput,
|
||||
ActionStepError, ActionStepKind, ActionStepRuntimeState, ActionTrace, catalog::find_recipe,
|
||||
slides_outline::project_slides_outline_markdown,
|
||||
};
|
||||
use crate::llm::{
|
||||
LlmPreparedImageDispatchRoutePayload, dispatch_prepared_image_route_payloads, dispatch_prepared_structured_routes,
|
||||
};
|
||||
|
||||
pub const ACTION_ABORTED_ERROR_CODE: &str = "action_aborted";
|
||||
pub const ACTION_INVALID_STEP_ERROR_CODE: &str = "action_invalid_step";
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct ActionRuntimeControl {
|
||||
pub abort_signal: Option<Arc<AtomicBool>>,
|
||||
pub event_sender: Option<Sender<ActionEvent>>,
|
||||
#[cfg(test)]
|
||||
pub abort_after_events: Option<usize>,
|
||||
#[cfg(test)]
|
||||
pub mock_output: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ActionRuntimeState {
|
||||
pub status: ActionRunStatus,
|
||||
pub result: Value,
|
||||
pub action_state: Value,
|
||||
pub steps: Vec<ActionStepRuntimeState>,
|
||||
pub events: Vec<ActionEvent>,
|
||||
pub trace: ActionTrace,
|
||||
pub error_code: Option<String>,
|
||||
}
|
||||
|
||||
fn invalid_input(message: impl Into<String>) -> Error {
|
||||
Error::new(Status::InvalidArg, message.into())
|
||||
}
|
||||
|
||||
pub fn run_action_recipe_prepared_with_control(
|
||||
input: ActionRuntimeInput,
|
||||
control: ActionRuntimeControl,
|
||||
) -> Result<ActionRuntimeOutput> {
|
||||
let recipe = find_recipe(&input.recipe_id, input.recipe_version.as_deref())?;
|
||||
validate_value("input", &recipe.input_schema, &input.input)?;
|
||||
|
||||
run_recipe(recipe, input, control)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn run_action_recipe_for_test(
|
||||
recipe: ActionRecipe,
|
||||
input: ActionRuntimeInput,
|
||||
) -> Result<ActionRuntimeOutput> {
|
||||
validate_value("input", &recipe.input_schema, &input.input)?;
|
||||
run_recipe(recipe, input, ActionRuntimeControl::default())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn run_action_recipe_for_test_with_control(
|
||||
recipe: ActionRecipe,
|
||||
input: ActionRuntimeInput,
|
||||
control: ActionRuntimeControl,
|
||||
) -> Result<ActionRuntimeOutput> {
|
||||
validate_value("input", &recipe.input_schema, &input.input)?;
|
||||
run_recipe(recipe, input, control)
|
||||
}
|
||||
|
||||
fn run_recipe(
|
||||
recipe: ActionRecipe,
|
||||
input: ActionRuntimeInput,
|
||||
control: ActionRuntimeControl,
|
||||
) -> Result<ActionRuntimeOutput> {
|
||||
let mut runtime = Runtime::new(recipe, input, control);
|
||||
runtime.run()
|
||||
}
|
||||
|
||||
struct Runtime {
|
||||
recipe: ActionRecipe,
|
||||
state: ActionRuntimeState,
|
||||
started_at: Instant,
|
||||
control: ActionRuntimeControl,
|
||||
}
|
||||
|
||||
impl Runtime {
|
||||
fn new(recipe: ActionRecipe, input: ActionRuntimeInput, control: ActionRuntimeControl) -> Self {
|
||||
let trace = ActionTrace {
|
||||
action_id: recipe.id.clone(),
|
||||
action_version: recipe.version.clone(),
|
||||
status: ActionRunStatus::Created,
|
||||
lightweight: Vec::new(),
|
||||
error_code: None,
|
||||
};
|
||||
|
||||
Self {
|
||||
recipe,
|
||||
state: ActionRuntimeState {
|
||||
status: ActionRunStatus::Created,
|
||||
result: input.input.clone(),
|
||||
action_state: input.input,
|
||||
steps: Vec::new(),
|
||||
events: Vec::new(),
|
||||
trace,
|
||||
error_code: None,
|
||||
},
|
||||
started_at: Instant::now(),
|
||||
control,
|
||||
}
|
||||
}
|
||||
|
||||
fn run(&mut self) -> Result<ActionRuntimeOutput> {
|
||||
let recipe = self.recipe_definition();
|
||||
let action_id = self.recipe.id.clone();
|
||||
let action_version = self.recipe.version.clone();
|
||||
let output_schema = self.recipe.output_schema.clone();
|
||||
let step_patches = self
|
||||
.recipe
|
||||
.steps
|
||||
.iter()
|
||||
.map(|step| (step.id.clone(), step.state_patch.clone()))
|
||||
.collect::<std::collections::HashMap<_, _>>();
|
||||
let attachments = Arc::new(Mutex::new(Vec::new()));
|
||||
let mut executor = AffineActionStepExecutor::new(&self.control, attachments.clone());
|
||||
let mut events = Vec::new();
|
||||
let mut lightweight = Vec::new();
|
||||
let event_sender = self.control.event_sender.clone();
|
||||
let abort_signal = self.control.abort_signal.clone();
|
||||
let event_count = Cell::new(0usize);
|
||||
#[cfg(test)]
|
||||
let abort_after_events = self.control.abort_after_events;
|
||||
|
||||
let mut record = |event: ActionEvent| {
|
||||
lightweight.push(json!({
|
||||
"type": event.event_type,
|
||||
"stepId": event.step_id,
|
||||
"status": event.status
|
||||
}));
|
||||
if let Some(sender) = &event_sender {
|
||||
let _ = sender.send(event.clone());
|
||||
}
|
||||
events.push(event);
|
||||
event_count.set(events.len());
|
||||
};
|
||||
|
||||
let runtime_output = run_recipe_runtime(
|
||||
recipe,
|
||||
self.state.action_state.clone(),
|
||||
&mut executor,
|
||||
|event| {
|
||||
for action_event in map_recipe_event(&action_id, &action_version, event, &attachments) {
|
||||
record(action_event);
|
||||
}
|
||||
},
|
||||
|| {
|
||||
abort_signal
|
||||
.as_ref()
|
||||
.is_some_and(|signal| signal.load(Ordering::SeqCst))
|
||||
|| {
|
||||
#[cfg(test)]
|
||||
{
|
||||
abort_after_events.is_some_and(|max_events| event_count.get() >= max_events)
|
||||
}
|
||||
#[cfg(not(test))]
|
||||
{
|
||||
false
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
if matches!(runtime_output.status, RecipeRuntimeStatus::Succeeded) {
|
||||
validate_value("output", &output_schema, &runtime_output.result)?;
|
||||
}
|
||||
|
||||
self.state = self.action_state_from_runtime_output(runtime_output, events, lightweight, step_patches);
|
||||
self.finalize_trace();
|
||||
if let Some(event) = self
|
||||
.state
|
||||
.events
|
||||
.iter_mut()
|
||||
.rev()
|
||||
.find(|event| matches!(event.event_type, ActionEventType::ActionDone))
|
||||
{
|
||||
event.trace = Some(self.state.trace.clone());
|
||||
}
|
||||
Ok(self.output())
|
||||
}
|
||||
|
||||
fn recipe_definition(&self) -> RecipeDefinition {
|
||||
RecipeDefinition {
|
||||
id: self.recipe.id.clone(),
|
||||
version: self.recipe.version.clone(),
|
||||
steps: self
|
||||
.recipe
|
||||
.steps
|
||||
.iter()
|
||||
.map(|step| RecipeStepExecution {
|
||||
id: step.id.clone(),
|
||||
kind: action_step_kind_name(step.kind).to_string(),
|
||||
input: step.input.clone(),
|
||||
state_patch: step.state_patch.clone(),
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn action_state_from_runtime_output(
|
||||
&self,
|
||||
output: RecipeRuntimeOutput,
|
||||
events: Vec<ActionEvent>,
|
||||
lightweight: Vec<Value>,
|
||||
step_patches: std::collections::HashMap<String, Option<Value>>,
|
||||
) -> ActionRuntimeState {
|
||||
let status = recipe_status_to_action_status(&output.status);
|
||||
let error_code = output
|
||||
.trace
|
||||
.error_code
|
||||
.as_deref()
|
||||
.map(map_recipe_error_code)
|
||||
.map(ToString::to_string);
|
||||
ActionRuntimeState {
|
||||
status,
|
||||
result: output.result,
|
||||
action_state: output.state,
|
||||
steps: output
|
||||
.steps
|
||||
.into_iter()
|
||||
.map(|step| ActionStepRuntimeState {
|
||||
id: step.id.clone(),
|
||||
input: step.input.unwrap_or(Value::Null),
|
||||
output: step.output,
|
||||
state_patch: step_patches.get(&step.id).cloned().flatten(),
|
||||
error: step.error.map(ActionStepError::from),
|
||||
})
|
||||
.collect(),
|
||||
events,
|
||||
trace: ActionTrace {
|
||||
action_id: self.recipe.id.clone(),
|
||||
action_version: self.recipe.version.clone(),
|
||||
status,
|
||||
lightweight,
|
||||
error_code: error_code.clone(),
|
||||
},
|
||||
error_code,
|
||||
}
|
||||
}
|
||||
|
||||
fn output(&mut self) -> ActionRuntimeOutput {
|
||||
self.finalize_trace();
|
||||
|
||||
ActionRuntimeOutput {
|
||||
result: self.state.result.clone(),
|
||||
status: self.state.status,
|
||||
error_code: self.state.error_code.clone(),
|
||||
state: self.state.action_state.clone(),
|
||||
steps: self.state.steps.clone(),
|
||||
trace: self.state.trace.clone(),
|
||||
events: self.state.events.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn finalize_trace(&mut self) {
|
||||
self.state.trace.status = self.state.status;
|
||||
if self
|
||||
.state
|
||||
.trace
|
||||
.lightweight
|
||||
.last()
|
||||
.and_then(|event| event.get("type"))
|
||||
.is_some_and(|event_type| event_type == "action_trace")
|
||||
{
|
||||
return;
|
||||
}
|
||||
self.state.trace.lightweight.push(json!({
|
||||
"type": "action_trace",
|
||||
"actionId": self.recipe.id.clone(),
|
||||
"actionVersion": self.recipe.version.clone(),
|
||||
"status": self.state.status,
|
||||
"durationMs": self.started_at.elapsed().as_millis()
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
fn recipe_status_to_action_status(status: &RecipeRuntimeStatus) -> ActionRunStatus {
|
||||
match status {
|
||||
RecipeRuntimeStatus::Created => ActionRunStatus::Created,
|
||||
RecipeRuntimeStatus::Running => ActionRunStatus::Running,
|
||||
RecipeRuntimeStatus::Succeeded => ActionRunStatus::Succeeded,
|
||||
RecipeRuntimeStatus::Failed => ActionRunStatus::Failed,
|
||||
RecipeRuntimeStatus::Aborted => ActionRunStatus::Aborted,
|
||||
}
|
||||
}
|
||||
|
||||
fn map_recipe_error_code(code: &str) -> &str {
|
||||
match code {
|
||||
"aborted" => ACTION_ABORTED_ERROR_CODE,
|
||||
"invalid_step" | "invalid_schema" | "invalid_value" => ACTION_INVALID_STEP_ERROR_CODE,
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
fn map_recipe_event(
|
||||
action_id: &str,
|
||||
action_version: &str,
|
||||
event: &RecipeRuntimeEvent,
|
||||
attachments: &Arc<Mutex<Vec<Value>>>,
|
||||
) -> Vec<ActionEvent> {
|
||||
let status = recipe_status_to_action_status(&event.status);
|
||||
let mut events = Vec::new();
|
||||
if event.event_type == "step_end" {
|
||||
let mut pending = attachments.lock().expect("attachment queue lock");
|
||||
events.extend(pending.drain(..).map(|attachment| ActionEvent {
|
||||
event_type: ActionEventType::Attachment,
|
||||
action_id: action_id.to_string(),
|
||||
action_version: action_version.to_string(),
|
||||
step_id: None,
|
||||
status: Some(ActionRunStatus::Running),
|
||||
attachment: Some(attachment),
|
||||
result: None,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
trace: None,
|
||||
}));
|
||||
}
|
||||
|
||||
let event_type = match event.event_type.as_str() {
|
||||
"recipe_start" => ActionEventType::ActionStart,
|
||||
"step_start" => ActionEventType::StepStart,
|
||||
"step_end" => ActionEventType::StepEnd,
|
||||
"recipe_done" => ActionEventType::ActionDone,
|
||||
"error" => ActionEventType::Error,
|
||||
_ => return events,
|
||||
};
|
||||
let error = event.error.as_ref();
|
||||
events.push(ActionEvent {
|
||||
event_type,
|
||||
action_id: action_id.to_string(),
|
||||
action_version: action_version.to_string(),
|
||||
step_id: event.step_id.clone(),
|
||||
status: Some(status),
|
||||
attachment: None,
|
||||
result: event.result.clone(),
|
||||
error_code: error.map(|error| map_recipe_error_code(&error.code).to_string()),
|
||||
error_message: error.map(|error| error.message.clone()),
|
||||
trace: None,
|
||||
});
|
||||
events
|
||||
}
|
||||
|
||||
impl From<StepExecutionError> for ActionStepError {
|
||||
fn from(error: StepExecutionError) -> Self {
|
||||
let code = if error.code == "invalid_step" || error.code == "invalid_schema" || error.code == "invalid_value" {
|
||||
ACTION_INVALID_STEP_ERROR_CODE.to_string()
|
||||
} else {
|
||||
error.code
|
||||
};
|
||||
Self {
|
||||
code,
|
||||
message: error.message,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn action_step_kind_name(kind: ActionStepKind) -> &'static str {
|
||||
match kind {
|
||||
ActionStepKind::PromptStructured => "promptStructured",
|
||||
ActionStepKind::PromptImage => "promptImage",
|
||||
ActionStepKind::ValidateJson => "validateJson",
|
||||
ActionStepKind::Transform => "transform",
|
||||
ActionStepKind::Final => "final",
|
||||
}
|
||||
}
|
||||
|
||||
struct AffineActionStepExecutor<'a> {
|
||||
#[cfg(test)]
|
||||
control: &'a ActionRuntimeControl,
|
||||
#[cfg(not(test))]
|
||||
_marker: std::marker::PhantomData<&'a ()>,
|
||||
attachments: Arc<Mutex<Vec<Value>>>,
|
||||
}
|
||||
|
||||
impl<'a> AffineActionStepExecutor<'a> {
|
||||
fn new(_control: &'a ActionRuntimeControl, attachments: Arc<Mutex<Vec<Value>>>) -> Self {
|
||||
Self {
|
||||
#[cfg(test)]
|
||||
control: _control,
|
||||
#[cfg(not(test))]
|
||||
_marker: std::marker::PhantomData,
|
||||
attachments,
|
||||
}
|
||||
}
|
||||
|
||||
fn test_mock_output(&self, _step_id: &str) -> Option<&Value> {
|
||||
#[cfg(test)]
|
||||
{
|
||||
self
|
||||
.control
|
||||
.mock_output
|
||||
.as_ref()
|
||||
.and_then(|mock_output| mock_output.get(_step_id))
|
||||
.filter(|value| !value.is_null())
|
||||
}
|
||||
#[cfg(not(test))]
|
||||
{
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn prompt_structured_step(
|
||||
&self,
|
||||
step: &RecipeStepExecution,
|
||||
input: Option<Value>,
|
||||
) -> std::result::Result<Value, StepExecutionError> {
|
||||
let value = if let Some(routes) = input
|
||||
.as_ref()
|
||||
.and_then(|input| input.get("preparedRoutes"))
|
||||
.filter(|routes| !routes.is_null())
|
||||
{
|
||||
let (_provider_id, response) =
|
||||
dispatch_prepared_structured_routes(&serde_json::to_string(routes).map_err(|error| {
|
||||
StepExecutionError::new(
|
||||
"invalid_step",
|
||||
format!("Invalid promptStructured prepared routes: {error}"),
|
||||
)
|
||||
})?)
|
||||
.map_err(|error| StepExecutionError::new("invalid_step", error.reason.clone()))?;
|
||||
response.output_json.unwrap_or(Value::Null)
|
||||
} else if let Some(mock_output) = self.test_mock_output(&step.id) {
|
||||
mock_output.clone()
|
||||
} else {
|
||||
return Err(StepExecutionError::new(
|
||||
"invalid_step",
|
||||
"promptStructured requires preparedRoutes",
|
||||
));
|
||||
};
|
||||
Ok(
|
||||
input
|
||||
.as_ref()
|
||||
.and_then(|input| input.get("unwrapKey"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|key| value.get(key).cloned())
|
||||
.unwrap_or(value),
|
||||
)
|
||||
}
|
||||
|
||||
fn prompt_image_step(
|
||||
&mut self,
|
||||
step: &RecipeStepExecution,
|
||||
input: Option<Value>,
|
||||
) -> std::result::Result<Value, StepExecutionError> {
|
||||
let attachment = if let Some(routes) = input
|
||||
.as_ref()
|
||||
.and_then(|input| input.get("preparedRoutes"))
|
||||
.filter(|routes| !routes.is_null())
|
||||
{
|
||||
let payload =
|
||||
serde_json::from_value::<Vec<LlmPreparedImageDispatchRoutePayload>>(routes.clone()).map_err(|error| {
|
||||
StepExecutionError::new("invalid_step", format!("Invalid promptImage prepared routes: {error}"))
|
||||
})?;
|
||||
let (_provider_id, response) = dispatch_prepared_image_route_payloads(payload)
|
||||
.map_err(|error| StepExecutionError::new("invalid_step", error.reason.clone()))?;
|
||||
image_response_attachment(response.provider_metadata, response.images)
|
||||
.ok_or_else(|| StepExecutionError::new("invalid_step", "promptImage native dispatch produced no image"))?
|
||||
} else if let Some(mock_output) = self.test_mock_output(&step.id) {
|
||||
mock_output.clone()
|
||||
} else {
|
||||
return Err(StepExecutionError::new(
|
||||
"invalid_step",
|
||||
"promptImage requires preparedRoutes",
|
||||
));
|
||||
};
|
||||
self
|
||||
.attachments
|
||||
.lock()
|
||||
.expect("attachment queue lock")
|
||||
.push(attachment.clone());
|
||||
Ok(attachment)
|
||||
}
|
||||
|
||||
fn transform_step(&self, input: Option<Value>, state: &Value) -> std::result::Result<Value, StepExecutionError> {
|
||||
if let Some(value) = execute_transform_step(input.clone(), state)? {
|
||||
return Ok(value);
|
||||
}
|
||||
|
||||
let Some(input) = input else {
|
||||
return Ok(state.clone());
|
||||
};
|
||||
if let Some(slides_outline) = input.get("slidesOutlineMarkdown") {
|
||||
let value = resolve_state_ref(slides_outline, state);
|
||||
return project_slides_outline_markdown(&value)
|
||||
.map(Value::String)
|
||||
.map_err(|message| StepExecutionError::new("invalid_step", message));
|
||||
}
|
||||
|
||||
Ok(input)
|
||||
}
|
||||
}
|
||||
|
||||
impl RecipeStepExecutor for AffineActionStepExecutor<'_> {
|
||||
fn execute_step(
|
||||
&mut self,
|
||||
step: &RecipeStepExecution,
|
||||
input: Option<Value>,
|
||||
state: &Value,
|
||||
) -> std::result::Result<Value, StepExecutionError> {
|
||||
match step.kind.as_str() {
|
||||
"promptStructured" => self.prompt_structured_step(step, input),
|
||||
"promptImage" => self.prompt_image_step(step, input),
|
||||
"validateJson" => execute_validate_json_step(input.or_else(|| Some(state.clone()))),
|
||||
"transform" | "final" => self.transform_step(input, state),
|
||||
other => Err(StepExecutionError::new(
|
||||
"invalid_step",
|
||||
format!("Unsupported action step kind: {other}"),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn image_response_attachment(provider_metadata: Value, images: Vec<llm_adapter::core::ImageArtifact>) -> Option<Value> {
|
||||
let image = images.into_iter().next()?;
|
||||
let mut attachment = Map::new();
|
||||
if let Some(url) = image.url {
|
||||
attachment.insert("url".to_string(), Value::String(url));
|
||||
}
|
||||
if let Some(data_base64) = image.data_base64 {
|
||||
attachment.insert("data_base64".to_string(), Value::String(data_base64));
|
||||
}
|
||||
attachment.insert("media_type".to_string(), Value::String(image.media_type));
|
||||
if let Some(width) = image.width {
|
||||
attachment.insert("width".to_string(), json!(width));
|
||||
}
|
||||
if let Some(height) = image.height {
|
||||
attachment.insert("height".to_string(), json!(height));
|
||||
}
|
||||
if !image.provider_metadata.is_null() {
|
||||
attachment.insert("providerMetadata".to_string(), image.provider_metadata);
|
||||
} else if !provider_metadata.is_null() {
|
||||
attachment.insert("providerMetadata".to_string(), provider_metadata);
|
||||
}
|
||||
if !attachment.contains_key("url") && !attachment.contains_key("data_base64") {
|
||||
return None;
|
||||
}
|
||||
Some(Value::Object(attachment))
|
||||
}
|
||||
|
||||
fn validate_value(label: &str, schema: &Value, value: &Value) -> Result<()> {
|
||||
validate_json_schema(label, schema, value).map_err(|error| invalid_input(error.message))
|
||||
}
|
||||
@@ -1,240 +0,0 @@
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub(super) fn project_slides_outline_markdown(value: &Value) -> Result<String, String> {
|
||||
let text = match value {
|
||||
Value::String(text) => text.as_str(),
|
||||
Value::Object(object) => {
|
||||
if let Some(Value::String(text)) = object.get("result") {
|
||||
text
|
||||
} else if let Some(Value::String(text)) = object.get("content") {
|
||||
text
|
||||
} else if let Some(Value::String(text)) = object.get("text") {
|
||||
text
|
||||
} else {
|
||||
return Err("slidesOutlineMarkdown requires a string result".to_string());
|
||||
}
|
||||
}
|
||||
_ => return Err("slidesOutlineMarkdown requires a string result".to_string()),
|
||||
};
|
||||
|
||||
if is_markdown_list(text) {
|
||||
return Ok(text.to_string());
|
||||
}
|
||||
|
||||
let mut projected = Vec::new();
|
||||
for line in text.lines().filter(|line| !line.trim().is_empty()) {
|
||||
let item = serde_json::from_str::<Value>(line)
|
||||
.map_err(|_| "slidesOutlineMarkdown requires markdown or NDJSON object lines".to_string())?;
|
||||
if !item.is_object() {
|
||||
return Err("slidesOutlineMarkdown requires markdown or NDJSON object lines".to_string());
|
||||
}
|
||||
projected.push(render_slide_item(&item)?);
|
||||
}
|
||||
|
||||
if projected.is_empty() {
|
||||
Err("slidesOutlineMarkdown requires markdown or NDJSON object lines".to_string())
|
||||
} else {
|
||||
Ok(projected.join("\n"))
|
||||
}
|
||||
}
|
||||
|
||||
fn is_markdown_list(text: &str) -> bool {
|
||||
let mut saw_line = false;
|
||||
for line in text.lines().map(str::trim_start).filter(|line| !line.trim().is_empty()) {
|
||||
saw_line = true;
|
||||
if !(line.starts_with("- ") || line.starts_with("* ") || line.starts_with("+ ")) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
saw_line
|
||||
}
|
||||
|
||||
fn render_legacy_slide_item(item: &Value) -> Option<String> {
|
||||
let kind = item.get("type").and_then(Value::as_str)?;
|
||||
let content = item.get("content").and_then(value_to_optional_string)?;
|
||||
if content.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
match kind {
|
||||
"name" => Some(format!("- {content}")),
|
||||
"title" => Some(format!(" - {content}")),
|
||||
"content" => {
|
||||
if content.contains('\n') {
|
||||
Some(
|
||||
content
|
||||
.lines()
|
||||
.map(|line| format!(" - {line}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n"),
|
||||
)
|
||||
} else {
|
||||
Some(format!(" - {content}"))
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn render_slide_item(item: &Value) -> Result<String, String> {
|
||||
if let Some(markdown) = render_legacy_slide_item(item) {
|
||||
return Ok(markdown);
|
||||
}
|
||||
if item.get("content").and_then(Value::as_object).is_some() {
|
||||
return render_structured_slide_item(item);
|
||||
}
|
||||
if item.get("content").and_then(Value::as_str).is_some() {
|
||||
return render_labeled_string_slide_item(item);
|
||||
}
|
||||
Err("slidesOutlineMarkdown item is not a recognized slide outline object".to_string())
|
||||
}
|
||||
|
||||
fn render_labeled_string_slide_item(item: &Value) -> Result<String, String> {
|
||||
let content = item
|
||||
.get("content")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or_else(|| "slidesOutlineMarkdown labeled item requires string content".to_string())?;
|
||||
if content.trim().is_empty() {
|
||||
return Err("slidesOutlineMarkdown labeled item requires string content".to_string());
|
||||
}
|
||||
let labels = parse_labeled_segments(content);
|
||||
let title = labels
|
||||
.get("title")
|
||||
.cloned()
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| "slidesOutlineMarkdown labeled item requires Title".to_string())?;
|
||||
let keywords = labels
|
||||
.get("image keywords")
|
||||
.cloned()
|
||||
.or_else(|| labels.get("keywords").cloned())
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| "slidesOutlineMarkdown labeled item requires Image Keywords".to_string())?;
|
||||
let description = labels
|
||||
.get("description")
|
||||
.cloned()
|
||||
.or_else(|| labels.get("content").cloned())
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| "slidesOutlineMarkdown labeled item requires Description".to_string())?;
|
||||
|
||||
Ok(
|
||||
[
|
||||
format!("- {title}"),
|
||||
format!(" - {title}"),
|
||||
format!(" - {keywords}"),
|
||||
format!(" - {description}"),
|
||||
]
|
||||
.join("\n"),
|
||||
)
|
||||
}
|
||||
|
||||
fn render_structured_slide_item(item: &Value) -> Result<String, String> {
|
||||
let item_object = item
|
||||
.as_object()
|
||||
.ok_or_else(|| "slidesOutlineMarkdown structured item requires object content".to_string())?;
|
||||
let content = item
|
||||
.get("content")
|
||||
.and_then(Value::as_object)
|
||||
.ok_or_else(|| "slidesOutlineMarkdown structured item requires object content".to_string())?;
|
||||
let title = string_prop(content, &["title", "name", "page_name", "pageName"])
|
||||
.or_else(|| string_prop(item_object, &["title", "name", "page_name", "pageName", "page"]))
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| "slidesOutlineMarkdown requires slide title".to_string())?;
|
||||
let sections = content.get("sections").and_then(Value::as_array);
|
||||
let rendered_sections = if let Some(sections) = sections.filter(|sections| !sections.is_empty()) {
|
||||
sections
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, section)| render_slide_section(section, index + 1))
|
||||
.collect::<Result<Vec<_>, _>>()?
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect::<Vec<_>>()
|
||||
} else {
|
||||
render_slide_object(content)?
|
||||
};
|
||||
|
||||
Ok(
|
||||
std::iter::once(format!("- {title}"))
|
||||
.chain(rendered_sections)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n"),
|
||||
)
|
||||
}
|
||||
|
||||
fn parse_labeled_segments(text: &str) -> std::collections::HashMap<String, String> {
|
||||
text
|
||||
.split(';')
|
||||
.filter_map(|segment| {
|
||||
let (key, value) = segment.split_once(':')?;
|
||||
let key = key.trim().to_ascii_lowercase();
|
||||
let value = value.trim().to_string();
|
||||
if key.is_empty() || value.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some((key, value))
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn render_slide_section(section: &Value, index: usize) -> Result<Vec<String>, String> {
|
||||
let Some(object) = section.as_object() else {
|
||||
return Err(format!("slidesOutlineMarkdown section {index} requires object content"));
|
||||
};
|
||||
|
||||
render_slide_object(object)
|
||||
}
|
||||
|
||||
fn render_slide_object(object: &Map<String, Value>) -> Result<Vec<String>, String> {
|
||||
let title = required_string_prop(
|
||||
object,
|
||||
&["title", "name", "section", "page_name", "pageName"],
|
||||
"slide section title",
|
||||
)?;
|
||||
let keywords = string_prop(
|
||||
object,
|
||||
&["image_keywords", "imageKeywords", "keywords", "image_keywords_optional"],
|
||||
)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_else(|| title.clone());
|
||||
let content = required_string_prop(
|
||||
object,
|
||||
&["content", "description", "summary", "text"],
|
||||
"slide section content",
|
||||
)?;
|
||||
|
||||
Ok(vec![
|
||||
format!(" - {title}"),
|
||||
format!(" - {keywords}"),
|
||||
format!(" - {content}"),
|
||||
])
|
||||
}
|
||||
|
||||
fn string_prop(object: &Map<String, Value>, keys: &[&str]) -> Option<String> {
|
||||
keys
|
||||
.iter()
|
||||
.find_map(|key| object.get(*key).and_then(value_to_optional_string))
|
||||
}
|
||||
|
||||
fn required_string_prop(object: &Map<String, Value>, keys: &[&str], name: &str) -> Result<String, String> {
|
||||
string_prop(object, keys)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| format!("slidesOutlineMarkdown requires {name}"))
|
||||
}
|
||||
|
||||
fn value_to_optional_string(value: &Value) -> Option<String> {
|
||||
match value {
|
||||
Value::String(text) => Some(text.clone()),
|
||||
Value::Number(number) => Some(number.to_string()),
|
||||
Value::Array(items) => {
|
||||
let joined = items
|
||||
.iter()
|
||||
.filter_map(value_to_optional_string)
|
||||
.filter(|value| !value.is_empty())
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
Some(joined)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -1,854 +0,0 @@
|
||||
use napi::Status;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
ACTION_ABORTED_ERROR_CODE, ActionEventType, ActionRecipe, ActionRecipeStep, ActionRunStatus, ActionRuntimeControl,
|
||||
ActionRuntimeInput, ActionStepKind, load_catalog, run_action_recipe_for_test,
|
||||
run_action_recipe_for_test_with_control, run_action_recipe_prepared_with_control, validate_catalog, validate_recipe,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn validates_built_in_recipe_catalog() {
|
||||
let catalog = load_catalog().unwrap();
|
||||
let mindmap = catalog.iter().find(|recipe| recipe.id == "mindmap.generate").unwrap();
|
||||
assert!(
|
||||
mindmap
|
||||
.steps
|
||||
.iter()
|
||||
.any(|step| step.kind == ActionStepKind::PromptStructured)
|
||||
);
|
||||
assert!(
|
||||
mindmap
|
||||
.steps
|
||||
.iter()
|
||||
.any(|step| step.kind == ActionStepKind::ValidateJson)
|
||||
);
|
||||
let slides = catalog.iter().find(|recipe| recipe.id == "slides.outline").unwrap();
|
||||
assert!(
|
||||
slides
|
||||
.steps
|
||||
.iter()
|
||||
.any(|step| step.id == "project-outline" && step.kind == ActionStepKind::Transform)
|
||||
);
|
||||
assert!(catalog.iter().any(|recipe| recipe.id == "transcript.audio.gemini"));
|
||||
assert!(!catalog.iter().any(|recipe| recipe.id == "transcript.audio.local-asr"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn built_in_transcript_action_final_result_is_schema_checked() {
|
||||
let output = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "transcript.audio.gemini".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({
|
||||
"sourceAudio": { "blobId": "blob-1", "mimeType": "audio/opus" },
|
||||
"quality": null,
|
||||
"infos": [{ "url": "https://example.com/audio.opus", "mimeType": "audio/opus", "index": 0 }],
|
||||
"sliceManifest": [{
|
||||
"index": 0,
|
||||
"fileName": "audio.opus",
|
||||
"mimeType": "audio/opus",
|
||||
"startSec": 12,
|
||||
"durationSec": 30,
|
||||
"byteSize": 42
|
||||
}],
|
||||
}),
|
||||
},
|
||||
mock_control(json!({
|
||||
"transcribe": {
|
||||
"normalizedTranscript": "00:00:01 A: Hello",
|
||||
"summaryJson": {
|
||||
"title": "Sync",
|
||||
"durationMinutes": 1,
|
||||
"attendees": ["A"],
|
||||
"keyPoints": ["Hello"],
|
||||
"actionItems": [],
|
||||
"decisions": [],
|
||||
"openQuestions": [],
|
||||
"blockers": []
|
||||
},
|
||||
"providerMeta": { "provider": "gemini", "model": "gemini-3.5-flash-lite" }
|
||||
}
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(output.result["version"], json!("transcript-result-v1"));
|
||||
assert_eq!(output.result["strategy"], json!("gemini"));
|
||||
assert_eq!(output.result["normalizedSegments"], json!(null));
|
||||
assert_eq!(output.result["sourceAudio"]["blobId"], json!("blob-1"));
|
||||
assert_eq!(
|
||||
output.result["infos"][0]["url"],
|
||||
json!("https://example.com/audio.opus")
|
||||
);
|
||||
assert_eq!(output.result["sliceManifest"][0]["startSec"], json!(12));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn built_in_transcript_action_rejects_malformed_summary() {
|
||||
let error = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "transcript.audio.gemini".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
mock_control(json!({
|
||||
"transcribe": {
|
||||
"normalizedTranscript": "00:00:01 A: Hello",
|
||||
"summaryJson": { "title": "Sync" },
|
||||
"providerMeta": { "provider": "gemini", "model": "gemini-3.5-flash-lite" }
|
||||
}
|
||||
})),
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(error.reason.contains("does not match JSON schema"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn built_in_action_final_result_comes_from_prompt_output_state() {
|
||||
let output = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "mindmap.generate".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
mock_control(json!({
|
||||
"generate-structured": {
|
||||
"result": "- Root"
|
||||
}
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(output.result, json!("- Root"));
|
||||
assert_eq!(output.state["generated"], json!("- Root"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn built_in_action_unwraps_structured_text_result() {
|
||||
let output = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "mindmap.generate".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
mock_control(json!({
|
||||
"generate-structured": {
|
||||
"result": "- Root"
|
||||
}
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(output.result, json!("- Root"));
|
||||
assert_eq!(output.state["generated"], json!("- Root"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn built_in_slides_outline_projects_final_result_to_markdown() {
|
||||
let outline = [
|
||||
serde_json::to_string(&json!({
|
||||
"page": "Cover",
|
||||
"type": "cover",
|
||||
"content": {
|
||||
"title": "Apple Inc.",
|
||||
"description": "Company overview",
|
||||
"image_keywords": ["Apple logo", "Apple Park"]
|
||||
}
|
||||
}))
|
||||
.unwrap(),
|
||||
serde_json::to_string(&json!({
|
||||
"page": 2,
|
||||
"type": "content",
|
||||
"content": {
|
||||
"title": "Products",
|
||||
"sections": [{
|
||||
"title": "iPhone",
|
||||
"keywords": ["smartphone", "iOS"],
|
||||
"content": "Flagship product line"
|
||||
}]
|
||||
}
|
||||
}))
|
||||
.unwrap(),
|
||||
serde_json::to_string(&json!({
|
||||
"page": 3,
|
||||
"type": "cover",
|
||||
"content": "Page Name: Closing; Title: Outlook; Description: Future strategy; Image Keywords: roadmap, devices"
|
||||
}))
|
||||
.unwrap(),
|
||||
]
|
||||
.join("\n");
|
||||
let output = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "slides.outline".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
mock_control(json!({
|
||||
"generate-structured": {
|
||||
"result": outline
|
||||
}
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(
|
||||
output.result,
|
||||
json!(
|
||||
[
|
||||
"- Apple Inc.",
|
||||
" - Apple Inc.",
|
||||
" - Apple logo, Apple Park",
|
||||
" - Company overview",
|
||||
"- Products",
|
||||
" - iPhone",
|
||||
" - smartphone, iOS",
|
||||
" - Flagship product line",
|
||||
"- Outlook",
|
||||
" - Outlook",
|
||||
" - roadmap, devices",
|
||||
" - Future strategy"
|
||||
]
|
||||
.join("\n")
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
output
|
||||
.steps
|
||||
.iter()
|
||||
.find(|step| step.id == "project-outline")
|
||||
.and_then(|step| step.output.as_ref()),
|
||||
Some(&output.result)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slides_outline_transform_keeps_legacy_markdown_shape() {
|
||||
let outline = [
|
||||
serde_json::to_string(&json!({ "page": 1, "type": "name", "content": "Launch deck" })).unwrap(),
|
||||
serde_json::to_string(&json!({ "page": 1, "type": "title", "content": "Context" })).unwrap(),
|
||||
serde_json::to_string(&json!({ "page": 1, "type": "content", "content": "Problem\nOpportunity" })).unwrap(),
|
||||
]
|
||||
.join("\n");
|
||||
let recipe = test_recipe(vec![
|
||||
ActionRecipeStep {
|
||||
id: "project-outline".to_string(),
|
||||
kind: ActionStepKind::Transform,
|
||||
input: Some(json!({
|
||||
"slidesOutlineMarkdown": { "$state": "outline" },
|
||||
"outputKey": "outlineMarkdown"
|
||||
})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({ "copy": { "$state": "outlineMarkdown" } })),
|
||||
state_patch: None,
|
||||
},
|
||||
]);
|
||||
let output = run_action_recipe_for_test(
|
||||
recipe,
|
||||
runtime_input(json!({
|
||||
"outline": outline
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(
|
||||
output.result,
|
||||
json!(["- Launch deck", " - Context", " - Problem", " - Opportunity"].join("\n"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slides_outline_transform_rejects_unrecognized_text() {
|
||||
let recipe = test_recipe(vec![
|
||||
ActionRecipeStep {
|
||||
id: "project-outline".to_string(),
|
||||
kind: ActionStepKind::Transform,
|
||||
input: Some(json!({
|
||||
"slidesOutlineMarkdown": { "$state": "outline" },
|
||||
"outputKey": "outlineMarkdown"
|
||||
})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({ "copy": { "$state": "outlineMarkdown" } })),
|
||||
state_patch: None,
|
||||
},
|
||||
]);
|
||||
let output = run_action_recipe_for_test(
|
||||
recipe,
|
||||
runtime_input(json!({
|
||||
"outline": "not valid ndjson"
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Failed);
|
||||
assert_eq!(output.error_code, Some("action_invalid_step".to_string()));
|
||||
assert_eq!(
|
||||
output.events.last().and_then(|event| event.error_message.as_deref()),
|
||||
Some("slidesOutlineMarkdown requires markdown or NDJSON object lines")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slides_outline_transform_accepts_cover_without_image_keywords() {
|
||||
let outline = serde_json::to_string(&json!({
|
||||
"page": 1,
|
||||
"type": "cover",
|
||||
"content": {
|
||||
"title": "Launch deck",
|
||||
"description": "Overview"
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
let recipe = test_recipe(vec![
|
||||
ActionRecipeStep {
|
||||
id: "project-outline".to_string(),
|
||||
kind: ActionStepKind::Transform,
|
||||
input: Some(json!({
|
||||
"slidesOutlineMarkdown": { "$state": "outline" },
|
||||
"outputKey": "outlineMarkdown"
|
||||
})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({ "copy": { "$state": "outlineMarkdown" } })),
|
||||
state_patch: None,
|
||||
},
|
||||
]);
|
||||
let output = run_action_recipe_for_test(
|
||||
recipe,
|
||||
runtime_input(json!({
|
||||
"outline": outline
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
output.result,
|
||||
json!(
|
||||
[
|
||||
"- Launch deck",
|
||||
" - Launch deck",
|
||||
" - Launch deck",
|
||||
" - Overview"
|
||||
]
|
||||
.join("\n")
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slides_outline_transform_accepts_page_name_from_item() {
|
||||
let outline = serde_json::to_string(&json!({
|
||||
"page": 2,
|
||||
"type": "content",
|
||||
"page_name": "Workspace Benefits",
|
||||
"content": {
|
||||
"sections": [
|
||||
{
|
||||
"section": "Unified writing",
|
||||
"keywords": ["docs", "canvas"],
|
||||
"text": "AFFiNE combines documents and whiteboards."
|
||||
}
|
||||
]
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
let recipe = test_recipe(vec![
|
||||
ActionRecipeStep {
|
||||
id: "project-outline".to_string(),
|
||||
kind: ActionStepKind::Transform,
|
||||
input: Some(json!({
|
||||
"slidesOutlineMarkdown": { "$state": "outline" },
|
||||
"outputKey": "outlineMarkdown"
|
||||
})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({ "copy": { "$state": "outlineMarkdown" } })),
|
||||
state_patch: None,
|
||||
},
|
||||
]);
|
||||
let output = run_action_recipe_for_test(
|
||||
recipe,
|
||||
runtime_input(json!({
|
||||
"outline": outline
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(
|
||||
output.result,
|
||||
json!(
|
||||
[
|
||||
"- Workspace Benefits",
|
||||
" - Unified writing",
|
||||
" - docs, canvas",
|
||||
" - AFFiNE combines documents and whiteboards."
|
||||
]
|
||||
.join("\n")
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serializes_action_events_for_server_contract() {
|
||||
let output = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "mindmap.generate".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
mock_control(json!({
|
||||
"generate-structured": {
|
||||
"result": "- Root"
|
||||
}
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
let first = serde_json::to_value(output.events.first().unwrap()).unwrap();
|
||||
let last = serde_json::to_value(output.events.last().unwrap()).unwrap();
|
||||
|
||||
assert_eq!(first["type"], json!("action_start"));
|
||||
assert_eq!(last["type"], json!("action_done"));
|
||||
assert_eq!(last["status"], json!("succeeded"));
|
||||
assert_eq!(last["trace"]["status"], json!("succeeded"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn built_in_action_fails_without_routes_or_mock_output() {
|
||||
let output = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "mindmap.generate".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
ActionRuntimeControl::default(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Failed);
|
||||
assert!(
|
||||
output
|
||||
.events
|
||||
.last()
|
||||
.and_then(|event| event.error_message.as_deref())
|
||||
.unwrap_or_default()
|
||||
.contains("promptStructured requires")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn built_in_image_action_uses_prompt_image_step_output() {
|
||||
let output = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "image.filter.sketch".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
mock_control(json!({
|
||||
"generate-image": {
|
||||
"url": "https://example.com/artifact-1.png"
|
||||
}
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(output.result, json!({ "url": "https://example.com/artifact-1.png" }));
|
||||
assert_eq!(
|
||||
output.state.pointer("/artifact/url"),
|
||||
Some(&json!("https://example.com/artifact-1.png"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn built_in_image_action_accepts_inline_artifact_output() {
|
||||
let output = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "image.filter.sketch".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
mock_control(json!({
|
||||
"generate-image": {
|
||||
"data_base64": "aW1n",
|
||||
"media_type": "image/webp"
|
||||
}
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(
|
||||
output.result,
|
||||
json!({
|
||||
"data_base64": "aW1n",
|
||||
"media_type": "image/webp"
|
||||
})
|
||||
);
|
||||
assert_eq!(output.state.pointer("/artifact/data_base64"), Some(&json!("aW1n")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_recipe_without_final_step() {
|
||||
let recipe = ActionRecipe {
|
||||
id: "invalid.recipe".to_string(),
|
||||
version: "v1".to_string(),
|
||||
input_schema: json!({}),
|
||||
output_schema: json!({}),
|
||||
steps: vec![ActionRecipeStep {
|
||||
id: "start".to_string(),
|
||||
kind: ActionStepKind::ValidateJson,
|
||||
input: None,
|
||||
state_patch: None,
|
||||
}],
|
||||
};
|
||||
|
||||
let error = validate_recipe(&recipe).unwrap_err();
|
||||
assert_eq!(error.status, Status::InvalidArg);
|
||||
assert!(error.reason.contains("must end with a final step"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_duplicated_recipe_identity() {
|
||||
let recipe = ActionRecipe {
|
||||
id: "duplicated.recipe".to_string(),
|
||||
version: "v1".to_string(),
|
||||
input_schema: json!({}),
|
||||
output_schema: json!({}),
|
||||
steps: vec![ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: None,
|
||||
state_patch: None,
|
||||
}],
|
||||
};
|
||||
|
||||
let error = validate_catalog(&[recipe.clone(), recipe]).unwrap_err();
|
||||
assert_eq!(error.status, Status::InvalidArg);
|
||||
assert!(error.reason.contains("Duplicated action recipe"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_recipe_where_final_step_is_not_last() {
|
||||
let recipe = ActionRecipe {
|
||||
id: "invalid.recipe".to_string(),
|
||||
version: "v1".to_string(),
|
||||
input_schema: json!({}),
|
||||
output_schema: json!({}),
|
||||
steps: vec![
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: None,
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "after-final".to_string(),
|
||||
kind: ActionStepKind::Transform,
|
||||
input: None,
|
||||
state_patch: None,
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
let error = validate_recipe(&recipe).unwrap_err();
|
||||
assert_eq!(error.status, Status::InvalidArg);
|
||||
assert!(error.reason.contains("must end with a final step"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_json_and_prompt_projection_steps() {
|
||||
let recipe = test_recipe(vec![
|
||||
ActionRecipeStep {
|
||||
id: "prompt-structured".to_string(),
|
||||
kind: ActionStepKind::PromptStructured,
|
||||
input: Some(json!({})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "prompt-image".to_string(),
|
||||
kind: ActionStepKind::PromptImage,
|
||||
input: Some(json!({})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "validate-json".to_string(),
|
||||
kind: ActionStepKind::ValidateJson,
|
||||
input: Some(json!({
|
||||
"schema": { "type": "object", "required": ["title"] },
|
||||
"value": { "title": "Hello" }
|
||||
})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({ "copy": { "done": true } })),
|
||||
state_patch: None,
|
||||
},
|
||||
]);
|
||||
|
||||
let output = run_action_recipe_for_test_with_control(
|
||||
recipe,
|
||||
runtime_input(json!({})),
|
||||
mock_control(json!({
|
||||
"prompt-structured": { "title": "Hello" },
|
||||
"prompt-image": { "url": "https://example.com/artifact-1.png" }
|
||||
})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
output
|
||||
.events
|
||||
.iter()
|
||||
.map(|event| event.event_type)
|
||||
.filter(|event_type| matches!(event_type, ActionEventType::Attachment))
|
||||
.collect::<Vec<_>>(),
|
||||
vec![ActionEventType::Attachment]
|
||||
);
|
||||
assert_eq!(output.steps[2].output, Some(json!(true)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_prompt_steps_without_prepared_routes_or_explicit_boundary() {
|
||||
let recipe = test_recipe(vec![
|
||||
ActionRecipeStep {
|
||||
id: "prompt".to_string(),
|
||||
kind: ActionStepKind::PromptStructured,
|
||||
input: Some(json!({})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: None,
|
||||
state_patch: None,
|
||||
},
|
||||
]);
|
||||
|
||||
let output = run_action_recipe_for_test(recipe, runtime_input(json!({}))).unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Failed);
|
||||
assert_eq!(output.error_code, Some("action_invalid_step".to_string()));
|
||||
assert!(
|
||||
output
|
||||
.events
|
||||
.last()
|
||||
.and_then(|event| event.error_message.as_deref())
|
||||
.unwrap_or_default()
|
||||
.contains("requires")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_prompt_image_without_prepared_routes() {
|
||||
let recipe = test_recipe(vec![
|
||||
ActionRecipeStep {
|
||||
id: "prompt-image".to_string(),
|
||||
kind: ActionStepKind::PromptImage,
|
||||
input: Some(json!({})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: None,
|
||||
state_patch: None,
|
||||
},
|
||||
]);
|
||||
|
||||
let output = run_action_recipe_for_test(recipe, runtime_input(json!({}))).unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Failed);
|
||||
assert!(
|
||||
output
|
||||
.events
|
||||
.last()
|
||||
.and_then(|event| event.error_message.as_deref())
|
||||
.unwrap_or_default()
|
||||
.contains("preparedRoutes")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_json_distinguishes_invalid_schema_from_invalid_value() {
|
||||
let invalid_value = run_action_recipe_for_test(
|
||||
test_recipe(vec![
|
||||
ActionRecipeStep {
|
||||
id: "validate-json".to_string(),
|
||||
kind: ActionStepKind::ValidateJson,
|
||||
input: Some(json!({
|
||||
"schema": { "type": "object", "required": ["title"] },
|
||||
"value": {}
|
||||
})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({ "copy": {} })),
|
||||
state_patch: None,
|
||||
},
|
||||
]),
|
||||
runtime_input(json!({})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(invalid_value.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(invalid_value.steps[0].output, Some(json!(false)));
|
||||
|
||||
let invalid_schema = run_action_recipe_for_test(
|
||||
test_recipe(vec![
|
||||
ActionRecipeStep {
|
||||
id: "validate-json".to_string(),
|
||||
kind: ActionStepKind::ValidateJson,
|
||||
input: Some(json!({
|
||||
"schema": { "type": 1 },
|
||||
"value": {}
|
||||
})),
|
||||
state_patch: None,
|
||||
},
|
||||
ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: None,
|
||||
state_patch: None,
|
||||
},
|
||||
]),
|
||||
runtime_input(json!({})),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(invalid_schema.status, ActionRunStatus::Failed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn emits_ordered_action_events_and_final_result() {
|
||||
let output = run_action_recipe_for_test(
|
||||
test_recipe(vec![ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({ "copy": {} })),
|
||||
state_patch: Some(json!({ "finalized": true })),
|
||||
}]),
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "test.recipe".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({ "content": "hello" }),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Succeeded);
|
||||
assert_eq!(output.result, json!({}));
|
||||
assert_eq!(output.error_code, None);
|
||||
assert_eq!(output.state, json!({ "content": "hello", "finalized": true }));
|
||||
assert_eq!(output.steps.len(), 1);
|
||||
assert_eq!(output.steps[0].id, "final");
|
||||
assert_eq!(output.steps[0].output, Some(json!({})));
|
||||
assert_eq!(output.steps[0].state_patch, Some(json!({ "finalized": true })));
|
||||
assert_eq!(output.steps[0].error, None);
|
||||
assert_eq!(
|
||||
output.events.iter().map(|event| event.event_type).collect::<Vec<_>>(),
|
||||
vec![
|
||||
ActionEventType::ActionStart,
|
||||
ActionEventType::StepStart,
|
||||
ActionEventType::StepEnd,
|
||||
ActionEventType::ActionDone,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
fn runtime_input(input: serde_json::Value) -> ActionRuntimeInput {
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "test.recipe".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input,
|
||||
}
|
||||
}
|
||||
|
||||
fn mock_control(mock_output: serde_json::Value) -> ActionRuntimeControl {
|
||||
ActionRuntimeControl {
|
||||
abort_signal: None,
|
||||
event_sender: None,
|
||||
abort_after_events: None,
|
||||
mock_output: Some(mock_output),
|
||||
}
|
||||
}
|
||||
|
||||
fn test_recipe(steps: Vec<ActionRecipeStep>) -> ActionRecipe {
|
||||
ActionRecipe {
|
||||
id: "test.recipe".to_string(),
|
||||
version: "v1".to_string(),
|
||||
input_schema: json!({}),
|
||||
output_schema: json!({}),
|
||||
steps,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generates_lightweight_trace() {
|
||||
let output = run_action_recipe_for_test(
|
||||
test_recipe(vec![ActionRecipeStep {
|
||||
id: "final".to_string(),
|
||||
kind: ActionStepKind::Final,
|
||||
input: Some(json!({ "copy": {} })),
|
||||
state_patch: None,
|
||||
}]),
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "test.recipe".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.trace.status, ActionRunStatus::Succeeded);
|
||||
assert!(!output.trace.lightweight.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn abort_control_stops_runtime() {
|
||||
let output = run_action_recipe_prepared_with_control(
|
||||
ActionRuntimeInput {
|
||||
recipe_id: "image.filter.sketch".to_string(),
|
||||
recipe_version: Some("v1".to_string()),
|
||||
input: json!({}),
|
||||
},
|
||||
ActionRuntimeControl {
|
||||
abort_signal: None,
|
||||
event_sender: None,
|
||||
abort_after_events: Some(1),
|
||||
mock_output: None,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.status, ActionRunStatus::Aborted);
|
||||
assert_eq!(output.error_code, Some(ACTION_ABORTED_ERROR_CODE.to_string()));
|
||||
assert_eq!(
|
||||
output.events.last().map(|event| event.event_type),
|
||||
Some(ActionEventType::Error)
|
||||
);
|
||||
}
|
||||
@@ -2,8 +2,10 @@
|
||||
{
|
||||
"name": "Transcript audio",
|
||||
"action": "Transcript audio",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"optionalModels": ["gemini-3.5-flash-lite", "gemini-3.6-flash"],
|
||||
"managedRoute": {
|
||||
"targets": ["gemini-3.5-flash-lite"],
|
||||
"premiumTargets": ["gemini-3.6-flash"]
|
||||
},
|
||||
"config": {
|
||||
"requireContent": false,
|
||||
"requireAttachment": true,
|
||||
@@ -14,8 +16,10 @@
|
||||
{
|
||||
"name": "Transcript audio structured",
|
||||
"action": "Transcript audio structured",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"optionalModels": ["gemini-3.5-flash-lite", "gemini-3.6-flash"],
|
||||
"managedRoute": {
|
||||
"targets": ["gemini-3.5-flash-lite"],
|
||||
"premiumTargets": ["gemini-3.6-flash"]
|
||||
},
|
||||
"config": {
|
||||
"requireContent": false,
|
||||
"requireAttachment": true,
|
||||
@@ -35,7 +39,7 @@
|
||||
{
|
||||
"name": "Generate a caption",
|
||||
"action": "Generate a caption",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"config": {
|
||||
"requireContent": false,
|
||||
"requireAttachment": true
|
||||
@@ -50,7 +54,7 @@
|
||||
{
|
||||
"name": "Conversation Summary",
|
||||
"action": "Conversation Summary",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"config": {
|
||||
"requireContent": false
|
||||
},
|
||||
@@ -68,7 +72,7 @@
|
||||
{
|
||||
"name": "Summary",
|
||||
"action": "Summary",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -83,7 +87,7 @@
|
||||
{
|
||||
"name": "Summary as title",
|
||||
"action": "Summary as title",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -98,7 +102,7 @@
|
||||
{
|
||||
"name": "Summary the webpage",
|
||||
"action": "Summary the webpage",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -109,7 +113,7 @@
|
||||
{
|
||||
"name": "Explain this",
|
||||
"action": "Explain this",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"builtins": ["language"],
|
||||
"messages": [
|
||||
{
|
||||
@@ -125,7 +129,7 @@
|
||||
{
|
||||
"name": "Explain this image",
|
||||
"action": "Explain this image",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"config": {
|
||||
"requireContent": false,
|
||||
"requireAttachment": true
|
||||
@@ -144,7 +148,7 @@
|
||||
{
|
||||
"name": "Explain this code",
|
||||
"action": "Explain this code",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -159,7 +163,7 @@
|
||||
{
|
||||
"name": "Translate to",
|
||||
"action": "Translate",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"params": {
|
||||
"language": {
|
||||
"default": "English",
|
||||
@@ -192,7 +196,7 @@
|
||||
{
|
||||
"name": "Summarize the meeting structured",
|
||||
"action": "Summarize the meeting structured",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -207,7 +211,7 @@
|
||||
{
|
||||
"name": "Summarize the meeting",
|
||||
"action": "Summarize the meeting",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -222,7 +226,7 @@
|
||||
{
|
||||
"name": "Find action for summary",
|
||||
"action": "Find action for summary",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -237,7 +241,7 @@
|
||||
{
|
||||
"name": "Write an article about this",
|
||||
"action": "Write an article about this",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -252,7 +256,7 @@
|
||||
{
|
||||
"name": "Write a twitter about this",
|
||||
"action": "Write a twitter about this",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -267,7 +271,7 @@
|
||||
{
|
||||
"name": "Write a poem about this",
|
||||
"action": "Write a poem about this",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -282,7 +286,7 @@
|
||||
{
|
||||
"name": "Write a blog post about this",
|
||||
"action": "Write a blog post about this",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -297,7 +301,7 @@
|
||||
{
|
||||
"name": "Write outline",
|
||||
"action": "Write outline",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -312,7 +316,7 @@
|
||||
{
|
||||
"name": "Change tone to",
|
||||
"action": "Change tone",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"params": {
|
||||
"tone": {
|
||||
"default": "professional",
|
||||
@@ -333,7 +337,7 @@
|
||||
{
|
||||
"name": "Brainstorm ideas about this",
|
||||
"action": "Brainstorm ideas about this",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -348,7 +352,7 @@
|
||||
{
|
||||
"name": "Brainstorm mindmap",
|
||||
"action": "Brainstorm mindmap",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -363,7 +367,7 @@
|
||||
{
|
||||
"name": "Expand mind map",
|
||||
"action": "Expand mind map",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -382,7 +386,7 @@
|
||||
{
|
||||
"name": "Improve writing for it",
|
||||
"action": "Improve writing for it",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -397,7 +401,7 @@
|
||||
{
|
||||
"name": "Improve grammar for it",
|
||||
"action": "Improve grammar for it",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -412,7 +416,7 @@
|
||||
{
|
||||
"name": "Fix spelling for it",
|
||||
"action": "Fix spelling for it",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -427,7 +431,7 @@
|
||||
{
|
||||
"name": "Find action items from it",
|
||||
"action": "Find action items from it",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -442,7 +446,7 @@
|
||||
{
|
||||
"name": "Check code error",
|
||||
"action": "Check code error",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -457,7 +461,7 @@
|
||||
{
|
||||
"name": "Create a presentation",
|
||||
"action": "Create a presentation",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -472,7 +476,7 @@
|
||||
{
|
||||
"name": "Create headings",
|
||||
"action": "Create headings",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -487,7 +491,7 @@
|
||||
{
|
||||
"name": "Make it real",
|
||||
"action": "Make it real",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"managedRoute": { "targets": ["claude-sonnet-4-6"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -502,7 +506,7 @@
|
||||
{
|
||||
"name": "Make it real with text",
|
||||
"action": "Make it real with text",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"managedRoute": { "targets": ["claude-sonnet-4-6"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -517,7 +521,7 @@
|
||||
{
|
||||
"name": "Make it longer",
|
||||
"action": "Make it longer",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -532,7 +536,7 @@
|
||||
{
|
||||
"name": "Make it shorter",
|
||||
"action": "Make it shorter",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -547,7 +551,7 @@
|
||||
{
|
||||
"name": "Continue writing",
|
||||
"action": "Continue writing",
|
||||
"model": "gemini-3.5-flash-lite",
|
||||
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -562,7 +566,7 @@
|
||||
{
|
||||
"name": "Section Edit",
|
||||
"action": "Section Edit",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"managedRoute": { "targets": ["claude-sonnet-4-6"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -577,7 +581,7 @@
|
||||
{
|
||||
"name": "Generate image",
|
||||
"action": "image",
|
||||
"model": "gpt-image-1",
|
||||
"managedRoute": { "targets": ["gpt-image-1"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -588,7 +592,7 @@
|
||||
{
|
||||
"name": "Convert to Clay style",
|
||||
"action": "Convert to Clay style",
|
||||
"model": "gpt-image-1",
|
||||
"managedRoute": { "targets": ["gpt-image-1"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -599,7 +603,7 @@
|
||||
{
|
||||
"name": "Convert to Sketch style",
|
||||
"action": "Convert to Sketch style",
|
||||
"model": "gpt-image-1",
|
||||
"managedRoute": { "targets": ["gpt-image-1"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -610,7 +614,7 @@
|
||||
{
|
||||
"name": "Convert to Anime style",
|
||||
"action": "Convert to Anime style",
|
||||
"model": "gpt-image-1",
|
||||
"managedRoute": { "targets": ["gpt-image-1"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -621,7 +625,7 @@
|
||||
{
|
||||
"name": "Convert to Pixel style",
|
||||
"action": "Convert to Pixel style",
|
||||
"model": "gpt-image-1",
|
||||
"managedRoute": { "targets": ["gpt-image-1"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -632,7 +636,7 @@
|
||||
{
|
||||
"name": "Convert to sticker",
|
||||
"action": "Convert to sticker",
|
||||
"model": "gpt-image-1",
|
||||
"managedRoute": { "targets": ["gpt-image-1"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -643,7 +647,7 @@
|
||||
{
|
||||
"name": "Upscale image",
|
||||
"action": "Upscale image",
|
||||
"model": "gpt-image-1",
|
||||
"managedRoute": { "targets": ["gpt-image-1"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -654,7 +658,7 @@
|
||||
{
|
||||
"name": "Remove background",
|
||||
"action": "Remove background",
|
||||
"model": "gpt-image-1",
|
||||
"managedRoute": { "targets": ["gpt-image-1"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -665,7 +669,7 @@
|
||||
{
|
||||
"name": "debug:action:fal-teed",
|
||||
"action": "fal-teed",
|
||||
"model": "workflowutils/teed",
|
||||
"managedRoute": { "targets": ["workflowutils/teed"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -675,7 +679,7 @@
|
||||
},
|
||||
{
|
||||
"name": "Code Artifact",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"managedRoute": { "targets": ["claude-sonnet-4-6"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -689,13 +693,36 @@
|
||||
},
|
||||
{
|
||||
"name": "Chat With AFFiNE AI",
|
||||
"model": "gpt-5.6-luna",
|
||||
"optionalModels": [
|
||||
"gpt-5.6-luna",
|
||||
"gpt-5.6-terra",
|
||||
"gemini-3.6-flash",
|
||||
"claude-sonnet-4-6"
|
||||
],
|
||||
"managedRoute": {
|
||||
"targets": ["gpt-5.6-luna"],
|
||||
"premiumTargets": ["gpt-5.6-luna"],
|
||||
"selectableTargets": [
|
||||
{
|
||||
"id": "luna",
|
||||
"modelId": "gpt-5.6-luna",
|
||||
"displayName": "GPT 5.6 Luna",
|
||||
"minimumTier": "standard"
|
||||
},
|
||||
{
|
||||
"id": "terra",
|
||||
"modelId": "gpt-5.6-terra",
|
||||
"displayName": "GPT 5.6 Terra",
|
||||
"minimumTier": "premium"
|
||||
},
|
||||
{
|
||||
"id": "gemini",
|
||||
"modelId": "gemini-3.6-flash",
|
||||
"displayName": "Gemini 3.6 Flash",
|
||||
"minimumTier": "premium"
|
||||
},
|
||||
{
|
||||
"id": "claude",
|
||||
"modelId": "claude-sonnet-4-6",
|
||||
"displayName": "Claude Sonnet 4.6",
|
||||
"minimumTier": "premium"
|
||||
}
|
||||
]
|
||||
},
|
||||
"config": {
|
||||
"tools": [
|
||||
"docRead",
|
||||
@@ -708,8 +735,7 @@
|
||||
"docCompose",
|
||||
"codeArtifact",
|
||||
"blobRead"
|
||||
],
|
||||
"proModels": ["gpt-5.6-terra", "gemini-3.6-flash", "claude-sonnet-4-6"]
|
||||
]
|
||||
},
|
||||
"builtins": [
|
||||
"date",
|
||||
@@ -734,7 +760,7 @@
|
||||
{
|
||||
"name": "mindmap.generate",
|
||||
"action": "mindmap.generate",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"config": {
|
||||
"frequencyPenalty": 0.5,
|
||||
"presencePenalty": 0.5,
|
||||
@@ -759,7 +785,7 @@
|
||||
{
|
||||
"name": "slides.outline",
|
||||
"action": "slides.outline",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -778,7 +804,7 @@
|
||||
{
|
||||
"name": "image.filter.sketch",
|
||||
"action": "image.filter.sketch",
|
||||
"model": "lora/image-to-image",
|
||||
"managedRoute": { "targets": ["lora/image-to-image"] },
|
||||
"config": {
|
||||
"modelName": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"loras": [
|
||||
@@ -798,7 +824,7 @@
|
||||
{
|
||||
"name": "image.filter.clay",
|
||||
"action": "image.filter.clay",
|
||||
"model": "lora/image-to-image",
|
||||
"managedRoute": { "targets": ["lora/image-to-image"] },
|
||||
"config": {
|
||||
"modelName": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"loras": [
|
||||
@@ -818,7 +844,7 @@
|
||||
{
|
||||
"name": "image.filter.anime",
|
||||
"action": "image.filter.anime",
|
||||
"model": "lora/image-to-image",
|
||||
"managedRoute": { "targets": ["lora/image-to-image"] },
|
||||
"config": {
|
||||
"modelName": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"loras": [
|
||||
@@ -838,7 +864,7 @@
|
||||
{
|
||||
"name": "image.filter.pixel",
|
||||
"action": "image.filter.pixel",
|
||||
"model": "lora/image-to-image",
|
||||
"managedRoute": { "targets": ["lora/image-to-image"] },
|
||||
"config": {
|
||||
"modelName": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"loras": [
|
||||
@@ -858,13 +884,12 @@
|
||||
{
|
||||
"name": "workflow:presentation",
|
||||
"action": "workflow:presentation",
|
||||
"model": "slides.outline",
|
||||
"messages": []
|
||||
},
|
||||
{
|
||||
"name": "workflow:presentation:step1",
|
||||
"action": "workflow:presentation:step1",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"config": {
|
||||
"temperature": 0.7
|
||||
},
|
||||
@@ -882,7 +907,7 @@
|
||||
{
|
||||
"name": "workflow:presentation:step2",
|
||||
"action": "workflow:presentation:step2",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -901,7 +926,7 @@
|
||||
{
|
||||
"name": "workflow:presentation:step4",
|
||||
"action": "workflow:presentation:step4",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -920,13 +945,12 @@
|
||||
{
|
||||
"name": "workflow:brainstorm",
|
||||
"action": "workflow:brainstorm",
|
||||
"model": "mindmap.generate",
|
||||
"messages": []
|
||||
},
|
||||
{
|
||||
"name": "workflow:brainstorm:step1",
|
||||
"action": "workflow:brainstorm:step1",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"config": {
|
||||
"temperature": 0.7
|
||||
},
|
||||
@@ -944,7 +968,7 @@
|
||||
{
|
||||
"name": "workflow:brainstorm:step2",
|
||||
"action": "workflow:brainstorm:step2",
|
||||
"model": "gpt-5.6-luna",
|
||||
"managedRoute": { "targets": ["gpt-5.6-luna"] },
|
||||
"config": {
|
||||
"frequencyPenalty": 0.5,
|
||||
"presencePenalty": 0.5,
|
||||
@@ -969,25 +993,21 @@
|
||||
{
|
||||
"name": "workflow:image-sketch",
|
||||
"action": "workflow:image-sketch",
|
||||
"model": "image.filter.sketch",
|
||||
"messages": []
|
||||
},
|
||||
{
|
||||
"name": "workflow:image-clay",
|
||||
"action": "workflow:image-clay",
|
||||
"model": "image.filter.clay",
|
||||
"messages": []
|
||||
},
|
||||
{
|
||||
"name": "workflow:image-anime",
|
||||
"action": "workflow:image-anime",
|
||||
"model": "image.filter.anime",
|
||||
"messages": []
|
||||
},
|
||||
{
|
||||
"name": "workflow:image-pixel",
|
||||
"action": "workflow:image-pixel",
|
||||
"model": "image.filter.pixel",
|
||||
"messages": []
|
||||
}
|
||||
]
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use llm_adapter::capability::provider_default_capability_upper_bound;
|
||||
use serde::Serialize;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::{ByokCapabilityInput, contract::capability_input};
|
||||
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokCatalogModelOutput {
|
||||
pub model_id: String,
|
||||
pub display_name: String,
|
||||
pub recommended: bool,
|
||||
pub capabilities: Vec<ByokCapabilityInput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokCatalogProviderOutput {
|
||||
pub provider: String,
|
||||
pub models: Vec<ByokCatalogModelOutput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokCatalogOutput {
|
||||
pub version: String,
|
||||
pub providers: Vec<ByokCatalogProviderOutput>,
|
||||
}
|
||||
|
||||
pub fn byok_catalog() -> ByokCatalogOutput {
|
||||
let variants = llm_adapter::core::default_model_registry_variants();
|
||||
let mut providers = ["openai", "anthropic", "gemini", "fal"]
|
||||
.into_iter()
|
||||
.map(|provider| (provider, BTreeMap::new()))
|
||||
.collect::<BTreeMap<_, BTreeMap<String, ByokCatalogModelOutput>>>();
|
||||
|
||||
for variant in variants {
|
||||
let Some(provider) = provider_for_backend(&variant.backend_kind) else {
|
||||
continue;
|
||||
};
|
||||
let Some(capabilities) = provider_default_capability_upper_bound(provider, &variant.raw_model_id) else {
|
||||
continue;
|
||||
};
|
||||
providers
|
||||
.entry(provider)
|
||||
.or_default()
|
||||
.entry(variant.raw_model_id.clone())
|
||||
.or_insert_with(|| ByokCatalogModelOutput {
|
||||
model_id: variant.raw_model_id.clone(),
|
||||
display_name: variant.display_name.unwrap_or_else(|| variant.raw_model_id.clone()),
|
||||
recommended: variant
|
||||
.capabilities
|
||||
.iter()
|
||||
.any(|capability| capability.default_for_output_type == Some(true)),
|
||||
capabilities: capabilities.into_iter().map(capability_input).collect(),
|
||||
});
|
||||
}
|
||||
|
||||
let providers = providers
|
||||
.into_iter()
|
||||
.map(|(provider, models)| ByokCatalogProviderOutput {
|
||||
provider: provider.to_string(),
|
||||
models: models.into_values().collect(),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let encoded = serde_json::to_vec(&providers).expect("BYOK catalog must serialize");
|
||||
let version = Sha256::digest(encoded)
|
||||
.iter()
|
||||
.take(8)
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect();
|
||||
ByokCatalogOutput { version, providers }
|
||||
}
|
||||
|
||||
fn provider_for_backend(backend: &str) -> Option<&'static str> {
|
||||
match backend {
|
||||
"openai_chat" | "openai_responses" => Some("openai"),
|
||||
"anthropic" => Some("anthropic"),
|
||||
"gemini_api" => Some("gemini"),
|
||||
"fal" => Some("fal"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn catalog_contains_explicit_provider_default_declarations() {
|
||||
let catalog = byok_catalog();
|
||||
assert!(!catalog.version.is_empty());
|
||||
for provider in &catalog.providers {
|
||||
assert!(
|
||||
!provider.models.is_empty(),
|
||||
"{} has no catalog models",
|
||||
provider.provider
|
||||
);
|
||||
for model in &provider.models {
|
||||
assert!(!model.model_id.is_empty());
|
||||
assert!(!model.capabilities.is_empty());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,579 @@
|
||||
use std::collections::HashSet;
|
||||
|
||||
use llm_adapter::{
|
||||
capability::{
|
||||
AttachmentKind, AttachmentSource, DeclaredModelCapability, ModelFeature, ModelInput, ModelOutput,
|
||||
provider_default_capability_upper_bound, validate_capability_upper_bound, validate_declared_capability,
|
||||
},
|
||||
target::canonicalize_endpoint,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokCapabilityInput {
|
||||
pub input: Vec<String>,
|
||||
pub output: Vec<String>,
|
||||
pub features: Vec<String>,
|
||||
pub attachment_kinds: Vec<String>,
|
||||
pub attachment_sources: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokModelDeclarationInput {
|
||||
pub model_id: String,
|
||||
pub enabled: bool,
|
||||
pub capabilities: Vec<ByokCapabilityInput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokEndpointInput {
|
||||
pub kind: String,
|
||||
pub url: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokProfileDefinitionInput {
|
||||
pub version: u32,
|
||||
pub endpoint: ByokEndpointInput,
|
||||
pub models: Vec<ByokModelDeclarationInput>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct CreateByokProfileInput {
|
||||
pub workspace_id: String,
|
||||
pub provider: String,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub credential: String,
|
||||
pub definition: ByokProfileDefinitionInput,
|
||||
pub enabled: bool,
|
||||
pub actor_user_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ReplaceByokProfileInput {
|
||||
pub workspace_id: String,
|
||||
pub profile_id: String,
|
||||
pub expected_revision: i32,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub definition: ByokProfileDefinitionInput,
|
||||
pub credential: Option<String>,
|
||||
pub enabled: bool,
|
||||
pub actor_user_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct RotateByokCredentialInput {
|
||||
pub workspace_id: String,
|
||||
pub profile_id: String,
|
||||
pub expected_revision: i32,
|
||||
pub credential: String,
|
||||
pub actor_user_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokProfileOrderInput {
|
||||
pub profile_id: String,
|
||||
pub expected_revision: i32,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ReorderByokProfilesInput {
|
||||
pub workspace_id: String,
|
||||
pub profiles: Vec<ByokProfileOrderInput>,
|
||||
pub actor_user_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ProbeByokProfileInput {
|
||||
pub workspace_id: String,
|
||||
pub profile_id: String,
|
||||
pub checks: Vec<ByokProbeCheckInput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ProbeByokDraftInput {
|
||||
pub workspace_id: String,
|
||||
pub provider: String,
|
||||
pub credential: Option<String>,
|
||||
pub profile_id: Option<String>,
|
||||
pub expected_revision: Option<i32>,
|
||||
pub definition: ByokProfileDefinitionInput,
|
||||
pub checks: Vec<ByokProbeCheckInput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokProbeCheckInput {
|
||||
pub model_id: String,
|
||||
pub operation: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct CreateByokLocalLeaseProviderInput {
|
||||
pub provider: String,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub credential: String,
|
||||
pub definition: ByokProfileDefinitionInput,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct CreateByokLocalLeaseInput {
|
||||
pub workspace_id: String,
|
||||
pub user_id: String,
|
||||
pub providers: Vec<CreateByokLocalLeaseProviderInput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokLocalLeaseOutput {
|
||||
pub lease_id: String,
|
||||
pub expires_at_ms: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokProfileOutput {
|
||||
pub profile_id: String,
|
||||
pub workspace_id: String,
|
||||
pub provider: String,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub definition: ByokProfileDefinitionInput,
|
||||
pub enabled: bool,
|
||||
pub sort_order: i32,
|
||||
pub revision: i32,
|
||||
pub validation: Option<ByokValidationOutput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokProbeStatusOutput {
|
||||
pub kind: String,
|
||||
pub tested_at_ms: Option<i64>,
|
||||
pub error_kind: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokModelProbeOutput {
|
||||
pub model_id: String,
|
||||
pub checks: Vec<ByokModelProbeCheckOutput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokModelProbeCheckOutput {
|
||||
pub operation: String,
|
||||
pub status: ByokProbeStatusOutput,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokValidationOutput {
|
||||
pub definition_fingerprint: String,
|
||||
pub credential_generation: i32,
|
||||
pub connection: ByokProbeStatusOutput,
|
||||
pub models: Vec<ByokModelProbeOutput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct ByokProbeResultOutput {
|
||||
pub definition_fingerprint: String,
|
||||
pub stale: bool,
|
||||
pub connection: ByokProbeStatusOutput,
|
||||
pub models: Vec<ByokModelProbeOutput>,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Eq, Deserialize, Serialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
|
||||
pub(crate) enum ByokEndpoint {
|
||||
ProviderDefault,
|
||||
Custom { url: String },
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Eq, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub(crate) struct ByokModelDeclaration {
|
||||
pub(crate) model_id: String,
|
||||
pub(crate) enabled: bool,
|
||||
pub(crate) capabilities: Vec<DeclaredModelCapability>,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Eq, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub(crate) struct ByokProfileDefinition {
|
||||
pub(crate) version: u32,
|
||||
pub(crate) endpoint: ByokEndpoint,
|
||||
pub(crate) models: Vec<ByokModelDeclaration>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub(crate) enum ByokContractError {
|
||||
#[error("unsupported BYOK definition version")]
|
||||
Version,
|
||||
#[error("unsupported BYOK provider")]
|
||||
Provider,
|
||||
#[error("{0} is required")]
|
||||
Required(&'static str),
|
||||
#[error("duplicate {0}")]
|
||||
Duplicate(&'static str),
|
||||
#[error("invalid BYOK endpoint")]
|
||||
Endpoint,
|
||||
#[error("invalid model capability: {0}")]
|
||||
Capability(String),
|
||||
#[error("declared capability exceeds provider or model upper bound")]
|
||||
CapabilityUpperBound,
|
||||
}
|
||||
|
||||
impl ByokProfileDefinition {
|
||||
pub(crate) fn endpoint_identity(&self) -> &str {
|
||||
match &self.endpoint {
|
||||
ByokEndpoint::ProviderDefault => "default",
|
||||
ByokEndpoint::Custom { url } => url,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn validate_definition(
|
||||
provider: &str,
|
||||
input: ByokProfileDefinitionInput,
|
||||
) -> Result<ByokProfileDefinition, ByokContractError> {
|
||||
if input.version != 1 {
|
||||
return Err(ByokContractError::Version);
|
||||
}
|
||||
if !matches!(provider, "openai" | "anthropic" | "gemini" | "fal") {
|
||||
return Err(ByokContractError::Provider);
|
||||
}
|
||||
let endpoint = match (input.endpoint.kind.as_str(), input.endpoint.url) {
|
||||
("provider_default", None) => ByokEndpoint::ProviderDefault,
|
||||
("custom", Some(url)) if !url.trim().is_empty() => ByokEndpoint::Custom {
|
||||
url: canonicalize_endpoint(&url).map_err(|_| ByokContractError::Endpoint)?,
|
||||
},
|
||||
_ => return Err(ByokContractError::Endpoint),
|
||||
};
|
||||
if input.models.is_empty() {
|
||||
return Err(ByokContractError::Required("models"));
|
||||
}
|
||||
|
||||
let mut ids = HashSet::new();
|
||||
let mut models = Vec::with_capacity(input.models.len());
|
||||
for model in input.models {
|
||||
let model_id = model.model_id.trim().to_string();
|
||||
if model_id.is_empty() || model_id.len() > 512 {
|
||||
return Err(ByokContractError::Required("modelId"));
|
||||
}
|
||||
if !ids.insert(model_id.clone()) {
|
||||
return Err(ByokContractError::Duplicate("modelId"));
|
||||
}
|
||||
if model.capabilities.is_empty() {
|
||||
return Err(ByokContractError::Required("capabilities"));
|
||||
}
|
||||
let capabilities = model
|
||||
.capabilities
|
||||
.into_iter()
|
||||
.map(parse_capability)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
validate_upper_bound(provider, &endpoint, &model_id, &capabilities)?;
|
||||
models.push(ByokModelDeclaration {
|
||||
model_id,
|
||||
enabled: model.enabled,
|
||||
capabilities,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(ByokProfileDefinition {
|
||||
version: 1,
|
||||
endpoint,
|
||||
models,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_capability(input: ByokCapabilityInput) -> Result<DeclaredModelCapability, ByokContractError> {
|
||||
let capability = DeclaredModelCapability {
|
||||
input: parse_values(input.input, |value| match value {
|
||||
"text" => Some(ModelInput::Text),
|
||||
"image" => Some(ModelInput::Image),
|
||||
"audio" => Some(ModelInput::Audio),
|
||||
"file" => Some(ModelInput::File),
|
||||
_ => None,
|
||||
})?,
|
||||
output: parse_values(input.output, |value| match value {
|
||||
"text" => Some(ModelOutput::Text),
|
||||
"object" => Some(ModelOutput::Object),
|
||||
"structured" => Some(ModelOutput::Structured),
|
||||
"embedding" => Some(ModelOutput::Embedding),
|
||||
"rerank" => Some(ModelOutput::Rerank),
|
||||
"image" => Some(ModelOutput::Image),
|
||||
_ => None,
|
||||
})?,
|
||||
features: parse_values(input.features, |value| match value {
|
||||
"tool_calling" => Some(ModelFeature::ToolCalling),
|
||||
"reasoning" => Some(ModelFeature::Reasoning),
|
||||
"web_search" => Some(ModelFeature::WebSearch),
|
||||
_ => None,
|
||||
})?,
|
||||
attachment_kinds: parse_values(input.attachment_kinds, |value| match value {
|
||||
"image" => Some(AttachmentKind::Image),
|
||||
"audio" => Some(AttachmentKind::Audio),
|
||||
"file" => Some(AttachmentKind::File),
|
||||
_ => None,
|
||||
})?,
|
||||
attachment_sources: parse_values(input.attachment_sources, |value| match value {
|
||||
"url" => Some(AttachmentSource::Url),
|
||||
"data" => Some(AttachmentSource::Data),
|
||||
"bytes" => Some(AttachmentSource::Bytes),
|
||||
"file_handle" => Some(AttachmentSource::FileHandle),
|
||||
_ => None,
|
||||
})?,
|
||||
};
|
||||
validate_declared_capability(&capability).map_err(|error| ByokContractError::Capability(error.to_string()))?;
|
||||
Ok(capability)
|
||||
}
|
||||
|
||||
fn parse_values<T>(values: Vec<String>, parse: impl Fn(&str) -> Option<T>) -> Result<Vec<T>, ByokContractError> {
|
||||
values
|
||||
.into_iter()
|
||||
.map(|value| parse(&value).ok_or_else(|| ByokContractError::Capability(format!("unknown enum {value}"))))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn validate_upper_bound(
|
||||
provider: &str,
|
||||
endpoint: &ByokEndpoint,
|
||||
model_id: &str,
|
||||
capabilities: &[DeclaredModelCapability],
|
||||
) -> Result<(), ByokContractError> {
|
||||
if provider == "fal"
|
||||
&& capabilities.iter().any(|capability| {
|
||||
capability.output.iter().any(|output| *output != ModelOutput::Image)
|
||||
|| capability
|
||||
.input
|
||||
.iter()
|
||||
.any(|input| !matches!(input, ModelInput::Text | ModelInput::Image))
|
||||
})
|
||||
{
|
||||
return Err(ByokContractError::CapabilityUpperBound);
|
||||
}
|
||||
if matches!(endpoint, ByokEndpoint::Custom { .. }) {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let upper_bound =
|
||||
provider_default_capability_upper_bound(provider, model_id).ok_or(ByokContractError::CapabilityUpperBound)?;
|
||||
for capability in capabilities {
|
||||
validate_capability_upper_bound(capability, &upper_bound).map_err(|_| ByokContractError::CapabilityUpperBound)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn input_name(value: &ModelInput) -> &'static str {
|
||||
match value {
|
||||
ModelInput::Text => "text",
|
||||
ModelInput::Image => "image",
|
||||
ModelInput::Audio => "audio",
|
||||
ModelInput::File => "file",
|
||||
}
|
||||
}
|
||||
|
||||
fn output_name(value: &ModelOutput) -> &'static str {
|
||||
match value {
|
||||
ModelOutput::Text => "text",
|
||||
ModelOutput::Object => "object",
|
||||
ModelOutput::Structured => "structured",
|
||||
ModelOutput::Embedding => "embedding",
|
||||
ModelOutput::Rerank => "rerank",
|
||||
ModelOutput::Image => "image",
|
||||
}
|
||||
}
|
||||
|
||||
fn attachment_kind_name(value: &AttachmentKind) -> &'static str {
|
||||
match value {
|
||||
AttachmentKind::Image => "image",
|
||||
AttachmentKind::Audio => "audio",
|
||||
AttachmentKind::File => "file",
|
||||
}
|
||||
}
|
||||
|
||||
fn attachment_source_name(value: &AttachmentSource) -> &'static str {
|
||||
match value {
|
||||
AttachmentSource::Url => "url",
|
||||
AttachmentSource::Data => "data",
|
||||
AttachmentSource::Bytes => "bytes",
|
||||
AttachmentSource::FileHandle => "file_handle",
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ByokProfileDefinition> for ByokProfileDefinitionInput {
|
||||
fn from(definition: ByokProfileDefinition) -> Self {
|
||||
Self {
|
||||
version: definition.version,
|
||||
endpoint: match definition.endpoint {
|
||||
ByokEndpoint::ProviderDefault => ByokEndpointInput {
|
||||
kind: "provider_default".to_string(),
|
||||
url: None,
|
||||
},
|
||||
ByokEndpoint::Custom { url } => ByokEndpointInput {
|
||||
kind: "custom".to_string(),
|
||||
url: Some(url),
|
||||
},
|
||||
},
|
||||
models: definition
|
||||
.models
|
||||
.into_iter()
|
||||
.map(|model| ByokModelDeclarationInput {
|
||||
model_id: model.model_id,
|
||||
enabled: model.enabled,
|
||||
capabilities: model.capabilities.into_iter().map(capability_input).collect(),
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn capability_input(capability: DeclaredModelCapability) -> ByokCapabilityInput {
|
||||
ByokCapabilityInput {
|
||||
input: capability.input.iter().map(input_name).map(str::to_string).collect(),
|
||||
output: capability.output.iter().map(output_name).map(str::to_string).collect(),
|
||||
features: capability
|
||||
.features
|
||||
.iter()
|
||||
.map(|value| match value {
|
||||
ModelFeature::ToolCalling => "tool_calling",
|
||||
ModelFeature::Reasoning => "reasoning",
|
||||
ModelFeature::WebSearch => "web_search",
|
||||
})
|
||||
.map(str::to_string)
|
||||
.collect(),
|
||||
attachment_kinds: capability
|
||||
.attachment_kinds
|
||||
.iter()
|
||||
.map(attachment_kind_name)
|
||||
.map(str::to_string)
|
||||
.collect(),
|
||||
attachment_sources: capability
|
||||
.attachment_sources
|
||||
.iter()
|
||||
.map(attachment_source_name)
|
||||
.map(str::to_string)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn definition(model_id: &str, capabilities: Vec<ByokCapabilityInput>) -> ByokProfileDefinitionInput {
|
||||
ByokProfileDefinitionInput {
|
||||
version: 1,
|
||||
endpoint: ByokEndpointInput {
|
||||
kind: "custom".to_string(),
|
||||
url: Some("https://example.com/v1/".to_string()),
|
||||
},
|
||||
models: vec![ByokModelDeclarationInput {
|
||||
model_id: model_id.to_string(),
|
||||
enabled: true,
|
||||
capabilities,
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
fn text_capability() -> ByokCapabilityInput {
|
||||
ByokCapabilityInput {
|
||||
input: vec!["text".to_string()],
|
||||
output: vec!["text".to_string()],
|
||||
features: vec![],
|
||||
attachment_kinds: vec![],
|
||||
attachment_sources: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_opaque_model_id_and_explicit_empty_features() {
|
||||
let validated =
|
||||
validate_definition("openai", definition(" vendor/model:latest ", vec![text_capability()])).unwrap();
|
||||
assert_eq!(validated.models[0].model_id, "vendor/model:latest");
|
||||
assert!(validated.models[0].capabilities[0].features.is_empty());
|
||||
assert_eq!(validated.endpoint_identity(), "https://example.com/v1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_shapes_table() {
|
||||
let mut cases = Vec::new();
|
||||
cases.push(ByokProfileDefinitionInput {
|
||||
models: vec![],
|
||||
..definition("model", vec![text_capability()])
|
||||
});
|
||||
cases.push(definition("", vec![text_capability()]));
|
||||
cases.push(definition("model", vec![]));
|
||||
let mut empty_input = text_capability();
|
||||
empty_input.input.clear();
|
||||
cases.push(definition("model", vec![empty_input]));
|
||||
let mut duplicate = text_capability();
|
||||
duplicate.output.push("text".to_string());
|
||||
cases.push(definition("model", vec![duplicate]));
|
||||
assert!(
|
||||
cases
|
||||
.into_iter()
|
||||
.all(|case| validate_definition("openai", case).is_err())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_endpoint_tag_mismatches() {
|
||||
for endpoint in [
|
||||
ByokEndpointInput {
|
||||
kind: "provider_default".to_string(),
|
||||
url: Some("https://example.com".to_string()),
|
||||
},
|
||||
ByokEndpointInput {
|
||||
kind: "custom".to_string(),
|
||||
url: None,
|
||||
},
|
||||
ByokEndpointInput {
|
||||
kind: "custom".to_string(),
|
||||
url: Some(" ".to_string()),
|
||||
},
|
||||
] {
|
||||
let mut input = definition("model", vec![text_capability()]);
|
||||
input.endpoint = endpoint;
|
||||
assert!(validate_definition("openai", input).is_err());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
use aes_gcm::{
|
||||
Aes256Gcm, Nonce,
|
||||
aead::{Aead, KeyInit, Payload},
|
||||
};
|
||||
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use hkdf::Hkdf;
|
||||
use rand::RngCore;
|
||||
use sha2::Sha256;
|
||||
use thiserror::Error;
|
||||
use zeroize::{Zeroize, Zeroizing};
|
||||
|
||||
const PREFIX: &str = "byok:v1:";
|
||||
const INFO: &[u8] = b"AFFiNE/Copilot/BYOK/v1";
|
||||
const NONCE_LEN: usize = 12;
|
||||
const TAG_LEN: usize = 16;
|
||||
|
||||
pub(crate) struct CredentialEnvelopeKey(Zeroizing<[u8; 32]>);
|
||||
pub(crate) struct SensitiveCredential(Zeroizing<Vec<u8>>);
|
||||
|
||||
impl SensitiveCredential {
|
||||
pub(crate) fn new(value: impl Into<Vec<u8>>) -> Self {
|
||||
Self(Zeroizing::new(value.into()))
|
||||
}
|
||||
|
||||
pub(crate) fn expose(&self) -> &[u8] {
|
||||
self.0.as_slice()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub(crate) enum CredentialEnvelopeError {
|
||||
#[error("credential_unavailable")]
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
impl CredentialEnvelopeKey {
|
||||
pub(crate) fn derive(root_secret: &[u8]) -> Result<Self, CredentialEnvelopeError> {
|
||||
if root_secret.is_empty() {
|
||||
return Err(CredentialEnvelopeError::Unavailable);
|
||||
}
|
||||
let mut key = Zeroizing::new([0_u8; 32]);
|
||||
Hkdf::<Sha256>::new(None, root_secret)
|
||||
.expand(INFO, key.as_mut())
|
||||
.map_err(|_| CredentialEnvelopeError::Unavailable)?;
|
||||
Ok(Self(key))
|
||||
}
|
||||
|
||||
pub(crate) fn encrypt(
|
||||
&self,
|
||||
credential: &SensitiveCredential,
|
||||
aad: &[u8],
|
||||
) -> Result<String, CredentialEnvelopeError> {
|
||||
let cipher = Aes256Gcm::new_from_slice(self.0.as_slice()).map_err(|_| CredentialEnvelopeError::Unavailable)?;
|
||||
let mut nonce = [0_u8; NONCE_LEN];
|
||||
rand::rng().fill_bytes(&mut nonce);
|
||||
let ciphertext = cipher
|
||||
.encrypt(
|
||||
Nonce::from_slice(&nonce),
|
||||
Payload {
|
||||
msg: credential.expose(),
|
||||
aad,
|
||||
},
|
||||
)
|
||||
.map_err(|_| CredentialEnvelopeError::Unavailable)?;
|
||||
let mut body = Vec::with_capacity(NONCE_LEN + ciphertext.len());
|
||||
body.extend_from_slice(&nonce);
|
||||
body.extend_from_slice(&ciphertext);
|
||||
let encoded = URL_SAFE_NO_PAD.encode(&body);
|
||||
body.zeroize();
|
||||
Ok(format!("{PREFIX}{encoded}"))
|
||||
}
|
||||
|
||||
pub(crate) fn decrypt(&self, envelope: &str, aad: &[u8]) -> Result<SensitiveCredential, CredentialEnvelopeError> {
|
||||
let encoded = envelope
|
||||
.strip_prefix(PREFIX)
|
||||
.ok_or(CredentialEnvelopeError::Unavailable)?;
|
||||
let mut body = URL_SAFE_NO_PAD
|
||||
.decode(encoded)
|
||||
.map_err(|_| CredentialEnvelopeError::Unavailable)?;
|
||||
if body.len() < NONCE_LEN + TAG_LEN {
|
||||
body.zeroize();
|
||||
return Err(CredentialEnvelopeError::Unavailable);
|
||||
}
|
||||
let (nonce, ciphertext) = body.split_at(NONCE_LEN);
|
||||
let cipher = Aes256Gcm::new_from_slice(self.0.as_slice()).map_err(|_| CredentialEnvelopeError::Unavailable)?;
|
||||
let result = cipher
|
||||
.decrypt(Nonce::from_slice(nonce), Payload { msg: ciphertext, aad })
|
||||
.map(SensitiveCredential::new)
|
||||
.map_err(|_| CredentialEnvelopeError::Unavailable);
|
||||
body.zeroize();
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn server_aad(workspace_id: &str, profile_id: &str, provider: &str, endpoint_identity: &str) -> Vec<u8> {
|
||||
["server", workspace_id, profile_id, provider, endpoint_identity]
|
||||
.join("\0")
|
||||
.into_bytes()
|
||||
}
|
||||
|
||||
pub(crate) fn local_aad(
|
||||
workspace_id: &str,
|
||||
user_id: &str,
|
||||
lease_id: &str,
|
||||
index: usize,
|
||||
provider: &str,
|
||||
endpoint_identity: &str,
|
||||
) -> Vec<u8> {
|
||||
[
|
||||
"local".to_string(),
|
||||
workspace_id.to_string(),
|
||||
user_id.to_string(),
|
||||
lease_id.to_string(),
|
||||
index.to_string(),
|
||||
provider.to_string(),
|
||||
endpoint_identity.to_string(),
|
||||
]
|
||||
.join("\0")
|
||||
.into_bytes()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn round_trips_with_random_nonce_and_rejects_tampering() {
|
||||
let key = CredentialEnvelopeKey::derive(b"stable-root").unwrap();
|
||||
let credential = SensitiveCredential::new(b"secret".to_vec());
|
||||
let aad = server_aad("workspace", "profile", "openai", "default");
|
||||
let first = key.encrypt(&credential, &aad).unwrap();
|
||||
let second = key.encrypt(&credential, &aad).unwrap();
|
||||
assert_ne!(first, second);
|
||||
assert_eq!(key.decrypt(&first, &aad).unwrap().expose(), b"secret");
|
||||
|
||||
let mut tampered = first.into_bytes();
|
||||
let last = tampered.len() - 1;
|
||||
tampered[last] = if tampered[last] == b'A' { b'B' } else { b'A' };
|
||||
assert!(key.decrypt(std::str::from_utf8(&tampered).unwrap(), &aad).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_wrong_context_key_version_and_legacy_ciphertext() {
|
||||
let key = CredentialEnvelopeKey::derive(b"stable-root").unwrap();
|
||||
let credential = SensitiveCredential::new(b"secret".to_vec());
|
||||
let aad = server_aad("workspace", "profile", "openai", "default");
|
||||
let encrypted = key.encrypt(&credential, &aad).unwrap();
|
||||
assert!(
|
||||
key
|
||||
.decrypt(&encrypted, &server_aad("other", "profile", "openai", "default"))
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
CredentialEnvelopeKey::derive(b"other")
|
||||
.unwrap()
|
||||
.decrypt(&encrypted, &aad)
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
key
|
||||
.decrypt(&encrypted.replacen("byok:v1:", "byok:v2:", 1), &aad)
|
||||
.is_err()
|
||||
);
|
||||
assert!(key.decrypt("bGVnYWN5", &aad).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn requires_stable_root_secret() {
|
||||
assert!(CredentialEnvelopeKey::derive(b"").is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
mod catalog;
|
||||
mod contract;
|
||||
mod envelope;
|
||||
mod validation;
|
||||
|
||||
pub use catalog::{ByokCatalogModelOutput, ByokCatalogOutput, ByokCatalogProviderOutput, byok_catalog};
|
||||
pub use contract::{
|
||||
ByokCapabilityInput, ByokEndpointInput, ByokLocalLeaseOutput, ByokModelDeclarationInput, ByokModelProbeCheckOutput,
|
||||
ByokModelProbeOutput, ByokProbeCheckInput, ByokProbeResultOutput, ByokProbeStatusOutput, ByokProfileDefinitionInput,
|
||||
ByokProfileOutput, ByokValidationOutput, CreateByokLocalLeaseInput, CreateByokLocalLeaseProviderInput,
|
||||
CreateByokProfileInput, ProbeByokDraftInput, ProbeByokProfileInput, ReorderByokProfilesInput,
|
||||
ReplaceByokProfileInput, RotateByokCredentialInput,
|
||||
};
|
||||
pub(crate) use contract::{ByokEndpoint, ByokModelDeclaration, ByokProfileDefinition, validate_definition};
|
||||
pub(crate) use envelope::{CredentialEnvelopeKey, SensitiveCredential, local_aad, server_aad};
|
||||
pub(crate) use validation::{definition_fingerprint, reconcile_validation};
|
||||
@@ -0,0 +1,93 @@
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::{ByokProfileDefinition, ByokValidationOutput};
|
||||
|
||||
pub(crate) fn definition_fingerprint(definition: &ByokProfileDefinition) -> String {
|
||||
let encoded = serde_json::to_vec(definition).expect("validated BYOK definition must serialize");
|
||||
Sha256::digest(encoded)
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn reconcile_validation(
|
||||
validation: Option<ByokValidationOutput>,
|
||||
old_definition: &ByokProfileDefinition,
|
||||
definition: &ByokProfileDefinition,
|
||||
credential_generation: i32,
|
||||
credential_changed: bool,
|
||||
) -> Option<ByokValidationOutput> {
|
||||
let mut validation = validation?;
|
||||
if credential_changed || old_definition.endpoint != definition.endpoint {
|
||||
return None;
|
||||
}
|
||||
validation.models.retain(|evidence| {
|
||||
let old = old_definition
|
||||
.models
|
||||
.iter()
|
||||
.find(|model| model.model_id == evidence.model_id);
|
||||
let new = definition
|
||||
.models
|
||||
.iter()
|
||||
.find(|model| model.model_id == evidence.model_id);
|
||||
old.is_some() && old == new
|
||||
});
|
||||
validation.definition_fingerprint = definition_fingerprint(definition);
|
||||
validation.credential_generation = credential_generation;
|
||||
Some(validation)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use llm_adapter::capability::{DeclaredModelCapability, ModelInput, ModelOutput};
|
||||
|
||||
use super::*;
|
||||
use crate::llm::byok::{ByokEndpoint, ByokModelDeclaration, ByokModelProbeOutput, ByokProbeStatusOutput};
|
||||
|
||||
fn definition(models: &[&str]) -> ByokProfileDefinition {
|
||||
ByokProfileDefinition {
|
||||
version: 1,
|
||||
endpoint: ByokEndpoint::ProviderDefault,
|
||||
models: models
|
||||
.iter()
|
||||
.map(|model| ByokModelDeclaration {
|
||||
model_id: (*model).to_string(),
|
||||
enabled: true,
|
||||
capabilities: vec![DeclaredModelCapability {
|
||||
input: vec![ModelInput::Text],
|
||||
output: vec![ModelOutput::Text],
|
||||
features: vec![],
|
||||
attachment_kinds: vec![],
|
||||
attachment_sources: vec![],
|
||||
}],
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_only_unchanged_model_evidence() {
|
||||
let old = definition(&["a", "b"]);
|
||||
let new = definition(&["b", "c"]);
|
||||
let validation = ByokValidationOutput {
|
||||
definition_fingerprint: definition_fingerprint(&old),
|
||||
credential_generation: 1,
|
||||
connection: ByokProbeStatusOutput {
|
||||
kind: "verified".to_string(),
|
||||
tested_at_ms: Some(1),
|
||||
error_kind: None,
|
||||
},
|
||||
models: ["a", "b"]
|
||||
.into_iter()
|
||||
.map(|model_id| ByokModelProbeOutput {
|
||||
model_id: model_id.to_string(),
|
||||
checks: vec![],
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
let next = reconcile_validation(Some(validation), &old, &new, 1, false).unwrap();
|
||||
assert_eq!(next.models.len(), 1);
|
||||
assert_eq!(next.models[0].model_id, "b");
|
||||
assert_eq!(next.definition_fingerprint, definition_fingerprint(&new));
|
||||
}
|
||||
}
|
||||
@@ -7,8 +7,7 @@ use super::{
|
||||
action::{TranscriptGeneratedResult, TranscriptInputContract, TranscriptResult},
|
||||
core::contracts::{
|
||||
CapabilityMatchRequest, CapabilityMatchResponse, ModelConditionsContract, ModelRegistryMatchRequest,
|
||||
ModelRegistryMatchResponse, ModelRegistryResolveRequest, ModelRegistryResolveResponse, PromptRenderContract,
|
||||
PromptSessionContract, ProviderDriverSpec, RequestedModelMatchRequest, RequestedModelMatchResponse,
|
||||
ModelRegistryMatchResponse, ModelRegistryResolveRequest, ModelRegistryResolveResponse, ProviderDriverSpec,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -78,7 +77,7 @@ fn mark_definition_property_nullable(schema: &mut Value, definition: &str, prope
|
||||
|
||||
pub(crate) fn transcript_input_schema() -> Value {
|
||||
let mut schema = generated_schema_for::<TranscriptInputContract>();
|
||||
for property in ["sourceAudio", "quality", "infos", "sliceManifest", "preparedRoutes"] {
|
||||
for property in ["sourceAudio", "quality", "infos", "sliceManifest"] {
|
||||
mark_property_nullable(&mut schema, property);
|
||||
}
|
||||
mark_definition_property_nullable(&mut schema, "TranscriptAudioInfo", "index");
|
||||
@@ -88,7 +87,7 @@ pub(crate) fn transcript_input_schema() -> Value {
|
||||
|
||||
pub(crate) fn transcript_generated_result_schema() -> Value {
|
||||
let mut schema = generated_schema_for::<TranscriptGeneratedResult>();
|
||||
for property in ["normalizedSegments", "summaryJson", "providerMeta"] {
|
||||
for property in ["normalizedSegments", "summaryJson"] {
|
||||
mark_property_nullable(&mut schema, property);
|
||||
}
|
||||
mark_definition_property_nullable(&mut schema, "MeetingSummaryActionItem", "owner");
|
||||
@@ -105,7 +104,6 @@ pub(crate) fn transcript_result_schema() -> Value {
|
||||
"sliceManifest",
|
||||
"normalizedSegments",
|
||||
"summaryJson",
|
||||
"providerMeta",
|
||||
] {
|
||||
mark_property_nullable(&mut schema, property);
|
||||
}
|
||||
@@ -118,12 +116,6 @@ pub(crate) fn transcript_result_schema() -> Value {
|
||||
|
||||
fn schema_by_name(name: &str) -> Option<Value> {
|
||||
match name {
|
||||
// runtime-owned temporary native facade
|
||||
"executionPlan" => Some(generated_schema_for::<llm_runtime::SerializableExecutionPlan>()),
|
||||
// adapter-owned temporary native facade
|
||||
"preparedRoutes" => Some(generated_schema_for::<
|
||||
Vec<llm_adapter::router::SerializablePreparedRoute>,
|
||||
>()),
|
||||
// AFFiNE-native-owned N-API projection over adapter model registry/matcher
|
||||
"capabilityMatchRequest" => Some(generated_schema_for::<CapabilityMatchRequest>()),
|
||||
"capabilityMatchResponse" => Some(generated_schema_for::<CapabilityMatchResponse>()),
|
||||
@@ -133,11 +125,6 @@ fn schema_by_name(name: &str) -> Option<Value> {
|
||||
"modelRegistryResolveRequest" => Some(generated_schema_for::<ModelRegistryResolveRequest>()),
|
||||
"modelRegistryResolveResponse" => Some(generated_schema_for::<ModelRegistryResolveResponse>()),
|
||||
"providerDriverSpec" => Some(generated_schema_for::<ProviderDriverSpec>()),
|
||||
// AFFiNE-native-owned prompt facade over adapter prompt DTOs/catalog
|
||||
"promptRenderContract" => Some(generated_schema_for::<PromptRenderContract>()),
|
||||
"promptSessionContract" => Some(generated_schema_for::<PromptSessionContract>()),
|
||||
"requestedModelMatchRequest" => Some(generated_schema_for::<RequestedModelMatchRequest>()),
|
||||
"requestedModelMatchResponse" => Some(generated_schema_for::<RequestedModelMatchResponse>()),
|
||||
// runtime-owned
|
||||
"toolCallbackRequest" => Some(generated_schema_for::<llm_runtime::ToolCallbackRequest>()),
|
||||
"toolCallbackResponse" => Some(generated_schema_for::<llm_runtime::ToolCallbackResponse>()),
|
||||
@@ -176,23 +163,6 @@ pub fn llm_validate_contract(name: String, value: Value) -> Result<Value> {
|
||||
)))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_compile_execution_plan(value: Value) -> Result<Value> {
|
||||
let value = llm_validate_contract("executionPlan".to_string(), value)?;
|
||||
llm_runtime::compile_execution_plan_value(value.clone()).map_err(|error| invalid_contract(error.to_string()))?;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_normalize_prepared_routes(value: Value) -> Result<Value> {
|
||||
let value = llm_adapter::router::normalize_prepared_routes(value).map_err(|error| {
|
||||
invalid_contract(format!(
|
||||
"LLM prepared routes value does not match adapter contract: {error}"
|
||||
))
|
||||
})?;
|
||||
llm_validate_contract("preparedRoutes".to_string(), value)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
@@ -220,8 +190,7 @@ mod tests {
|
||||
"decisions": [],
|
||||
"openQuestions": [],
|
||||
"blockers": []
|
||||
},
|
||||
"providerMeta": { "provider": "gemini" }
|
||||
}
|
||||
});
|
||||
assert!(llm_validate_contract("transcriptGeneratedResult".to_string(), value).is_ok());
|
||||
}
|
||||
@@ -234,7 +203,6 @@ mod tests {
|
||||
"normalizedSegments": null,
|
||||
"normalizedTranscript": "",
|
||||
"summaryJson": null,
|
||||
"providerMeta": null,
|
||||
"extra": true
|
||||
}),
|
||||
)
|
||||
@@ -242,24 +210,6 @@ mod tests {
|
||||
assert!(error.reason.contains("does not match schema"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compiles_execution_plan_contract() {
|
||||
let value = json!({
|
||||
"routes": [{
|
||||
"providerId": "openai-main",
|
||||
"protocol": "openai_chat",
|
||||
"model": "gpt-5-mini",
|
||||
"backendConfig": { "base_url": "https://api.openai.com/v1", "auth_token": "token" }
|
||||
}],
|
||||
"request": { "kind": "text", "cond": { "modelId": "gpt-5-mini" }, "messages": [] },
|
||||
"routePolicy": { "fallbackOrder": ["openai-main"] },
|
||||
"runtimePolicy": {},
|
||||
"attachmentPolicy": { "materializeRemoteAttachments": true },
|
||||
"responsePostprocess": { "mode": "text" }
|
||||
});
|
||||
assert!(super::llm_compile_execution_plan(value).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_runtime_tool_callback_contracts() {
|
||||
assert!(
|
||||
@@ -288,100 +238,4 @@ mod tests {
|
||||
.unwrap_err();
|
||||
assert!(error.reason.contains("does not match schema"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_prompt_contracts_from_native_types() {
|
||||
assert!(
|
||||
llm_validate_contract(
|
||||
"promptRenderContract".to_string(),
|
||||
json!({
|
||||
"messages": [{ "role": "user", "content": "hello" }],
|
||||
"templateParams": {},
|
||||
"renderParams": {}
|
||||
}),
|
||||
)
|
||||
.is_ok()
|
||||
);
|
||||
assert!(
|
||||
llm_validate_contract(
|
||||
"promptSessionContract".to_string(),
|
||||
json!({
|
||||
"prompt": {
|
||||
"promptTokens": 1,
|
||||
"templateParams": {},
|
||||
"messages": [{ "role": "system", "content": "hello" }]
|
||||
},
|
||||
"turns": [],
|
||||
"renderParams": {},
|
||||
"maxTokenSize": 1000
|
||||
}),
|
||||
)
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_adapter_prepared_route_contract() {
|
||||
assert!(
|
||||
super::llm_normalize_prepared_routes(json!([
|
||||
{
|
||||
"provider_id": "openai-main",
|
||||
"protocol": "openai_chat",
|
||||
"model": "gpt-5-mini",
|
||||
"config": {
|
||||
"base_url": "https://api.openai.com/v1",
|
||||
"auth_token": "token"
|
||||
},
|
||||
"request": {
|
||||
"model": "gpt-5-mini",
|
||||
"messages": []
|
||||
}
|
||||
}
|
||||
]))
|
||||
.is_ok()
|
||||
);
|
||||
|
||||
let error = super::llm_normalize_prepared_routes(json!([
|
||||
{
|
||||
"provider_id": "openai-main",
|
||||
"protocol": "openai_chat",
|
||||
"model": "gpt-5-mini",
|
||||
"config": { "base_url": "https://api.openai.com/v1" },
|
||||
"request": {}
|
||||
}
|
||||
]))
|
||||
.unwrap_err();
|
||||
assert!(error.reason.contains("adapter contract"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn execution_plan_rejects_host_only_state() {
|
||||
let value = json!({
|
||||
"routes": [],
|
||||
"request": {
|
||||
"kind": "text",
|
||||
"cond": { "modelId": "gpt-5-mini" },
|
||||
"messages": [],
|
||||
"options": { "signal": {} }
|
||||
},
|
||||
"routePolicy": { "fallbackOrder": [] },
|
||||
"runtimePolicy": {},
|
||||
"attachmentPolicy": { "materializeRemoteAttachments": true },
|
||||
"responsePostprocess": { "mode": "text" }
|
||||
});
|
||||
let error = super::llm_compile_execution_plan(value).unwrap_err();
|
||||
assert!(error.reason.contains("request.options.signal"));
|
||||
|
||||
let value = json!({
|
||||
"routes": [],
|
||||
"request": { "kind": "text", "cond": { "modelId": "gpt-5-mini" }, "messages": [] },
|
||||
"routePolicy": { "fallbackOrder": [] },
|
||||
"runtimePolicy": {},
|
||||
"attachmentPolicy": { "materializeRemoteAttachments": true },
|
||||
"responsePostprocess": { "mode": "text" },
|
||||
"hostContext": { "signal": {} }
|
||||
});
|
||||
let error = super::llm_compile_execution_plan(value).unwrap_err();
|
||||
assert!(error.reason.contains("does not match schema"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
use napi::Result;
|
||||
|
||||
use crate::llm::core::contracts::{
|
||||
CapabilityMatchRequest, CapabilityMatchResponse, RequestedModelMatchRequest, RequestedModelMatchResponse,
|
||||
};
|
||||
use crate::llm::core::contracts::{CapabilityMatchRequest, CapabilityMatchResponse};
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_match_model_capabilities(payload: CapabilityMatchRequest) -> Result<CapabilityMatchResponse> {
|
||||
@@ -14,25 +12,7 @@ pub fn llm_match_model_capabilities(payload: CapabilityMatchRequest) -> Result<C
|
||||
.map_err(crate::llm::map_json_error)?;
|
||||
|
||||
Ok(CapabilityMatchResponse {
|
||||
model_id: llm_adapter::core::select_model_id(&models, &cond).map_err(crate::llm::host::invalid_arg)?,
|
||||
})
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_resolve_requested_model_match(payload: RequestedModelMatchRequest) -> Result<RequestedModelMatchResponse> {
|
||||
let matched_optional_model = llm_adapter::core::matches_requested_model_list(
|
||||
&payload.provider_ids,
|
||||
&payload.optional_models,
|
||||
payload.requested_model_id.as_deref(),
|
||||
);
|
||||
|
||||
Ok(RequestedModelMatchResponse {
|
||||
selected_model: if matched_optional_model {
|
||||
payload.requested_model_id
|
||||
} else {
|
||||
payload.default_model
|
||||
},
|
||||
matched_optional_model,
|
||||
model_id: llm_adapter::core::select_model_id(&models, &cond).map_err(crate::llm::invalid_arg)?,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -8,17 +8,6 @@ use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PromptRenderContract {
|
||||
pub messages: Vec<PromptMessageContract>,
|
||||
#[napi(ts_type = "Record<string, any>")]
|
||||
pub template_params: Value,
|
||||
#[napi(ts_type = "Record<string, any>")]
|
||||
pub render_params: Value,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
|
||||
pub struct PromptRenderResult {
|
||||
@@ -35,66 +24,13 @@ pub struct BuiltInPromptRenderContract {
|
||||
pub render_params: Value,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
|
||||
pub struct PromptTokenCountContract {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
pub messages: Vec<PromptCountMessage>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
|
||||
pub struct PromptTokenCountResult {
|
||||
pub tokens: u32,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
|
||||
pub struct PromptCountMessage {
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
|
||||
pub struct PromptMetadataContract {
|
||||
pub messages: Vec<PromptMessageContract>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PromptMetadataResult {
|
||||
pub param_keys: Vec<String>,
|
||||
#[napi(ts_type = "Record<string, any>")]
|
||||
pub template_params: Value,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PromptSessionContract {
|
||||
pub prompt: PromptSessionPrompt,
|
||||
pub turns: Vec<PromptMessageContract>,
|
||||
#[napi(ts_type = "Record<string, any>")]
|
||||
pub render_params: Value,
|
||||
pub max_token_size: u32,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PromptSessionPrompt {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub action: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
pub prompt_tokens: u32,
|
||||
#[napi(ts_type = "Record<string, any>")]
|
||||
pub template_params: Value,
|
||||
pub messages: Vec<PromptMessageContract>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
@@ -112,7 +48,6 @@ pub struct BuiltInPromptSessionContract {
|
||||
pub turns: Vec<PromptMessageContract>,
|
||||
#[napi(ts_type = "Record<string, any>")]
|
||||
pub render_params: Value,
|
||||
pub max_token_size: u32,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
@@ -289,29 +224,6 @@ pub struct CapabilityMatchResponse {
|
||||
pub model_id: Option<String>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, JsonSchema, Serialize, PartialEq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct RequestedModelMatchRequest {
|
||||
pub provider_ids: Vec<String>,
|
||||
pub optional_models: Vec<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub requested_model_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub default_model: Option<String>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, JsonSchema, Serialize, PartialEq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct RequestedModelMatchResponse {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub selected_model: Option<String>,
|
||||
pub matched_optional_model: bool,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, JsonSchema, Serialize, PartialEq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
@@ -636,8 +548,6 @@ pub struct LlmImageRequestContract {
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct LlmImageRequestBuildContract {
|
||||
pub model: String,
|
||||
#[napi(ts_type = "'openai_chat' | 'openai_responses' | 'openai_images' | 'anthropic' | 'gemini' | 'fal_image'")]
|
||||
pub protocol: String,
|
||||
pub messages: Vec<PromptMessageContract>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub options: Option<Value>,
|
||||
@@ -647,47 +557,7 @@ pub struct LlmImageRequestBuildContract {
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{CapabilityMatchRequest, PromptRenderContract, PromptSessionContract, ProviderDriverSpec};
|
||||
|
||||
#[test]
|
||||
fn should_roundtrip_prompt_contracts() {
|
||||
let render_value = json!({
|
||||
"messages": [{
|
||||
"role": "system",
|
||||
"content": "summarize",
|
||||
"responseFormat": {
|
||||
"type": "json_schema",
|
||||
"responseSchemaJson": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"summary": { "type": "string" }
|
||||
},
|
||||
"required": ["summary"]
|
||||
},
|
||||
"schemaHash": "abc123"
|
||||
}
|
||||
}],
|
||||
"templateParams": { "tone": "short" },
|
||||
"renderParams": { "topic": "docs" }
|
||||
});
|
||||
let session_value = json!({
|
||||
"prompt": {
|
||||
"model": "gpt-5-mini",
|
||||
"promptTokens": 12,
|
||||
"templateParams": {},
|
||||
"messages": [{ "role": "system", "content": "summarize" }]
|
||||
},
|
||||
"turns": [{ "role": "user", "content": "hello" }],
|
||||
"renderParams": { "tone": "short" },
|
||||
"maxTokenSize": 1024
|
||||
});
|
||||
|
||||
let render_contract: PromptRenderContract = serde_json::from_value(render_value.clone()).unwrap();
|
||||
let session_contract: PromptSessionContract = serde_json::from_value(session_value.clone()).unwrap();
|
||||
|
||||
assert_eq!(serde_json::to_value(render_contract).unwrap(), render_value);
|
||||
assert_eq!(serde_json::to_value(session_contract).unwrap(), session_value);
|
||||
}
|
||||
use super::{CapabilityMatchRequest, ProviderDriverSpec};
|
||||
|
||||
#[test]
|
||||
fn should_roundtrip_tool_and_runtime_contracts() {
|
||||
|
||||
@@ -21,7 +21,7 @@ pub fn llm_resolve_model_registry_variant(
|
||||
request.backend_kind.as_deref(),
|
||||
request.model_id.as_str(),
|
||||
)
|
||||
.map_err(crate::llm::host::invalid_arg)?
|
||||
.map_err(crate::llm::invalid_arg)?
|
||||
{
|
||||
Some((variant, matched_by)) => ModelRegistryResolveResponse {
|
||||
variant: Some(to_contract_variant(variant)?),
|
||||
@@ -44,7 +44,7 @@ pub fn llm_match_model_registry(request: ModelRegistryMatchRequest) -> Result<Mo
|
||||
.map_err(crate::llm::map_json_error)?;
|
||||
let response = ModelRegistryMatchResponse {
|
||||
variant: llm_adapter::core::select_model_registry_variant(&variants, request.backend_kind.as_str(), &cond)
|
||||
.map_err(crate::llm::host::invalid_arg)?
|
||||
.map_err(crate::llm::invalid_arg)?
|
||||
.map(to_contract_variant)
|
||||
.transpose()?,
|
||||
};
|
||||
|
||||
@@ -1,16 +1,12 @@
|
||||
use napi::{Error, Result, Status};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
llm::{
|
||||
core::contracts::{
|
||||
BuiltInPromptRenderContract, BuiltInPromptSessionContract, PromptMessageContract, PromptMetadataContract,
|
||||
PromptMetadataResult, PromptRenderContract, PromptRenderResult, PromptSessionContract, PromptSessionPrompt,
|
||||
PromptSessionResult, PromptTokenCountContract, PromptTokenCountResult,
|
||||
},
|
||||
prompt_catalog::{BuiltInPrompt, BuiltInPromptSpec, built_in_prompt, built_in_prompt_spec, built_in_prompt_specs},
|
||||
use crate::llm::{
|
||||
core::contracts::{
|
||||
BuiltInPromptRenderContract, BuiltInPromptSessionContract, PromptMessageContract, PromptMetadataResult,
|
||||
PromptRenderResult, PromptSessionResult,
|
||||
},
|
||||
tiktoken::{Tokenizer, from_model_name},
|
||||
prompt_catalog::{BuiltInPrompt, BuiltInPromptSpec, built_in_prompt, built_in_prompt_spec, built_in_prompt_specs},
|
||||
};
|
||||
|
||||
mod metadata;
|
||||
@@ -51,57 +47,6 @@ fn built_in_prompt_metadata(prompt: &BuiltInPrompt) -> Result<PromptMetadataResu
|
||||
.map_err(|error| invalid_arg(format!("Failed to collect built-in prompt metadata: {error}")))
|
||||
}
|
||||
|
||||
fn count_prompt_tokens(model: Option<&str>, messages: &[PromptMessageContract]) -> u32 {
|
||||
let content = messages
|
||||
.iter()
|
||||
.map(|message| message.content.as_str())
|
||||
.collect::<String>();
|
||||
prompt_tokenizer(model)
|
||||
.map(|tokenizer| tokenizer.count(content, None))
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn prompt_tokenizer(model: Option<&str>) -> Option<Tokenizer> {
|
||||
let model = model?;
|
||||
if model.starts_with("gpt") {
|
||||
return from_model_name(model.to_string());
|
||||
}
|
||||
if model.starts_with("dall") {
|
||||
return None;
|
||||
}
|
||||
|
||||
from_model_name("gpt-4".to_string())
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_render_prompt(request: PromptRenderContract) -> Result<PromptRenderResult> {
|
||||
let response = render_prompt_response(
|
||||
&request.messages,
|
||||
&value_to_map(request.template_params, "templateParams")?,
|
||||
&value_to_map(request.render_params, "renderParams")?,
|
||||
)
|
||||
.map_err(|error| invalid_arg(format!("Failed to render prompt: {error}")))?;
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_count_prompt_tokens(request: PromptTokenCountContract) -> Result<PromptTokenCountResult> {
|
||||
let content = request
|
||||
.messages
|
||||
.iter()
|
||||
.map(|message| message.content.as_str())
|
||||
.collect::<String>();
|
||||
let tokens = request
|
||||
.model
|
||||
.as_deref()
|
||||
.and_then(|model| prompt_tokenizer(Some(model)))
|
||||
.map(|tokenizer| tokenizer.count(content, None))
|
||||
.unwrap_or(0);
|
||||
|
||||
Ok(PromptTokenCountResult { tokens })
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_render_built_in_prompt(request: BuiltInPromptRenderContract) -> Result<PromptRenderResult> {
|
||||
let prompt = built_in_prompt(&request.name)
|
||||
@@ -118,46 +63,22 @@ pub fn llm_render_built_in_prompt(request: BuiltInPromptRenderContract) -> Resul
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_collect_prompt_metadata(request: PromptMetadataContract) -> Result<PromptMetadataResult> {
|
||||
let response = collect_prompt_metadata(&request.messages)
|
||||
.map_err(|error| invalid_arg(format!("Failed to collect prompt metadata: {error}")))?;
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_render_session_prompt(request: PromptSessionContract) -> Result<PromptSessionResult> {
|
||||
let template_params = value_to_map(request.prompt.template_params.clone(), "prompt.templateParams")?;
|
||||
let render_params = value_to_map(request.render_params.clone(), "renderParams")?;
|
||||
let response = render_session_prompt(&request, &template_params, &render_params)
|
||||
.map_err(|error| invalid_arg(format!("Failed to render session prompt: {error}")))?;
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_render_built_in_session_prompt(request: BuiltInPromptSessionContract) -> Result<PromptSessionResult> {
|
||||
let prompt = built_in_prompt(&request.name)
|
||||
.ok_or_else(|| invalid_arg(format!("Built-in prompt not found: {}", request.name)))?;
|
||||
let messages = built_in_prompt_messages(prompt);
|
||||
let metadata = built_in_prompt_metadata(prompt)?;
|
||||
let session_contract = PromptSessionContract {
|
||||
prompt: PromptSessionPrompt {
|
||||
action: prompt.action.clone(),
|
||||
model: Some(prompt.model.clone()),
|
||||
prompt_tokens: count_prompt_tokens(Some(prompt.model.as_str()), &messages),
|
||||
template_params: metadata.template_params,
|
||||
messages,
|
||||
},
|
||||
turns: request.turns,
|
||||
render_params: request.render_params,
|
||||
max_token_size: request.max_token_size,
|
||||
};
|
||||
let template_params = value_to_map(session_contract.prompt.template_params.clone(), "prompt.templateParams")?;
|
||||
let render_params = value_to_map(session_contract.render_params.clone(), "renderParams")?;
|
||||
let response = render_session_prompt(&session_contract, &template_params, &render_params)
|
||||
.map_err(|error| invalid_arg(format!("Failed to render built-in session prompt: {error}")))?;
|
||||
let template_params = value_to_map(metadata.template_params, "prompt.templateParams")?;
|
||||
let render_params = value_to_map(request.render_params, "renderParams")?;
|
||||
let response = render_session_prompt(
|
||||
&messages,
|
||||
prompt.action.as_deref(),
|
||||
&request.turns,
|
||||
&template_params,
|
||||
&render_params,
|
||||
)
|
||||
.map_err(|error| invalid_arg(format!("Failed to render built-in session prompt: {error}")))?;
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
@@ -171,274 +92,3 @@ pub fn llm_list_built_in_prompt_specs() -> Result<Vec<BuiltInPromptSpec>> {
|
||||
pub fn llm_get_built_in_prompt_spec(name: String) -> Result<Option<BuiltInPromptSpec>> {
|
||||
Ok(built_in_prompt_spec(&name).cloned())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use llm_adapter::core::prompt_template::{is_truthy_number, parse_template, render_tokens};
|
||||
use serde_json::json;
|
||||
|
||||
use super::{llm_collect_prompt_metadata, llm_count_prompt_tokens, llm_render_prompt, llm_render_session_prompt};
|
||||
use crate::llm::core::contracts::{
|
||||
PromptMetadataContract, PromptRenderContract, PromptSessionContract, PromptTokenCountContract,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn should_render_sections_and_current_item() {
|
||||
let tokens = parse_template("{{#links}}- {{.}}\n{{/links}}").unwrap();
|
||||
let rendered = render_tokens(
|
||||
&tokens,
|
||||
&[&json!({
|
||||
"links": ["https://affine.pro", "https://github.com/toeverything/affine"]
|
||||
})],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
rendered,
|
||||
"- https://affine.pro\n- https://github.com/toeverything/affine\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_render_prompt_with_normalized_params_and_attachments() {
|
||||
let response = llm_render_prompt(
|
||||
serde_json::from_value::<PromptRenderContract>(json!({
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "tone={{tone}}"
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "{{content}}"
|
||||
}
|
||||
],
|
||||
"templateParams": { "tone": ["formal", "casual"] },
|
||||
"renderParams": {
|
||||
"attachments": ["https://affine.pro/example.jpg"],
|
||||
"content": "hello world"
|
||||
}
|
||||
}))
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
let response = serde_json::to_value(response).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
response,
|
||||
json!({
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "tone=formal",
|
||||
"params": {
|
||||
"attachments": ["https://affine.pro/example.jpg"],
|
||||
"content": "hello world",
|
||||
"tone": "formal"
|
||||
}
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "hello world",
|
||||
"attachments": ["https://affine.pro/example.jpg"],
|
||||
"params": {
|
||||
"attachments": ["https://affine.pro/example.jpg"],
|
||||
"content": "hello world",
|
||||
"tone": "formal"
|
||||
}
|
||||
}
|
||||
],
|
||||
"warnings": ["Missing param value: tone, use default options: formal"]
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_render_host_builtins_and_js_like_variable_strings() {
|
||||
let response = llm_render_prompt(
|
||||
serde_json::from_value::<PromptRenderContract>(json!({
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "{{affine::language}}|{{tags}}|{{obj}}|{{#links}}- {{.}}\n{{/links}}"
|
||||
}
|
||||
],
|
||||
"templateParams": {},
|
||||
"renderParams": {
|
||||
"language": "French",
|
||||
"affine::language": "ignored",
|
||||
"links": ["https://affine.pro", "https://github.com/toeverything/affine"],
|
||||
"obj": { "hello": "world" },
|
||||
"tags": ["a", "b"]
|
||||
}
|
||||
}))
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
let response = serde_json::to_value(response).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
response,
|
||||
json!({
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "French|a,b|[object Object]|- https://affine.pro\n- https://github.com/toeverything/affine\n",
|
||||
"params": {
|
||||
"language": "French",
|
||||
"affine::language": "ignored",
|
||||
"links": ["https://affine.pro", "https://github.com/toeverything/affine"],
|
||||
"obj": { "hello": "world" },
|
||||
"tags": ["a", "b"]
|
||||
}
|
||||
}
|
||||
],
|
||||
"warnings": []
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_count_prompt_tokens_for_unknown_models_as_zero() {
|
||||
let response = llm_count_prompt_tokens(
|
||||
serde_json::from_value::<PromptTokenCountContract>(json!({
|
||||
"model": null,
|
||||
"messages": [{ "content": "hello" }]
|
||||
}))
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
let response = serde_json::to_value(response).unwrap();
|
||||
|
||||
assert_eq!(response, json!({ "tokens": 0 }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_count_prompt_tokens_for_non_gpt_models_with_fallback_tokenizer() {
|
||||
let response = llm_count_prompt_tokens(
|
||||
serde_json::from_value::<PromptTokenCountContract>(json!({
|
||||
"model": "claude-3-5-sonnet",
|
||||
"messages": [{ "content": "hello" }]
|
||||
}))
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(response.tokens > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_follow_js_truthiness_for_numbers() {
|
||||
assert!(!is_truthy_number(&serde_json::Number::from(0)));
|
||||
assert!(is_truthy_number(&serde_json::Number::from(1)));
|
||||
assert!(is_truthy_number(&serde_json::Number::from_f64(0.5).unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_render_session_prompt_by_merging_latest_user_content() {
|
||||
let response = llm_render_session_prompt(
|
||||
serde_json::from_value::<PromptSessionContract>(json!({
|
||||
"prompt": {
|
||||
"model": "test",
|
||||
"promptTokens": 0,
|
||||
"templateParams": {},
|
||||
"messages": [
|
||||
{ "role": "system", "content": "answer briefly" },
|
||||
{ "role": "user", "content": "{{content}}" }
|
||||
]
|
||||
},
|
||||
"turns": [
|
||||
{ "role": "user", "content": "hello", "attachments": ["https://affine.pro/hello.png"] }
|
||||
],
|
||||
"renderParams": {},
|
||||
"maxTokenSize": 1000
|
||||
}))
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
let response = serde_json::to_value(response).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
response,
|
||||
json!({
|
||||
"messages": [
|
||||
{ "role": "system", "content": "answer briefly", "params": { "content": "hello" } },
|
||||
{
|
||||
"role": "user",
|
||||
"content": "hello",
|
||||
"attachments": ["https://affine.pro/hello.png"],
|
||||
"params": { "content": "hello" }
|
||||
}
|
||||
],
|
||||
"warnings": [],
|
||||
"promptMessagePositions": [0, 1]
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_render_session_prompt_by_picking_recent_turns_under_budget() {
|
||||
let response = llm_render_session_prompt(
|
||||
serde_json::from_value::<PromptSessionContract>(json!({
|
||||
"prompt": {
|
||||
"model": "test",
|
||||
"promptTokens": 0,
|
||||
"templateParams": {},
|
||||
"messages": [
|
||||
{ "role": "system", "content": "hello {{word}}" }
|
||||
]
|
||||
},
|
||||
"turns": [
|
||||
{ "role": "user", "content": "older turn" }
|
||||
],
|
||||
"renderParams": { "word": "world" },
|
||||
"maxTokenSize": 0
|
||||
}))
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
let response = serde_json::to_value(response).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
response,
|
||||
json!({
|
||||
"messages": [
|
||||
{ "role": "system", "content": "hello world", "params": { "word": "world" } }
|
||||
],
|
||||
"warnings": [],
|
||||
"promptMessagePositions": [0]
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_collect_prompt_metadata_from_templates_and_params() {
|
||||
let response = llm_collect_prompt_metadata(
|
||||
serde_json::from_value::<PromptMetadataContract>(json!({
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "tone={{tone}}"
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "{{content}}",
|
||||
"params": { "tone": ["formal", "casual"] }
|
||||
}
|
||||
]
|
||||
}))
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
let response = serde_json::to_value(response).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
response,
|
||||
json!({
|
||||
"paramKeys": ["tone", "content"],
|
||||
"templateParams": {
|
||||
"tone": ["formal", "casual"]
|
||||
}
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -156,3 +156,68 @@ fn render_prompt_message(
|
||||
|
||||
Ok(next)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn message(role: &str, content: &str) -> PromptMessageContract {
|
||||
serde_json::from_value(json!({ "role": role, "content": content })).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_lists_and_normalizes_missing_or_invalid_params_to_defaults() {
|
||||
let messages = vec![
|
||||
message("system", "translate {{src}} to {{dest}}: {{content}}"),
|
||||
message("user", "links:\n{{#links}}- {{.}}\n{{/links}}"),
|
||||
];
|
||||
let template_params = serde_json::from_value(json!({
|
||||
"src": ["eng"],
|
||||
"dest": ["chs", "jpn"]
|
||||
}))
|
||||
.unwrap();
|
||||
let params = serde_json::from_value(json!({
|
||||
"src": "invalid",
|
||||
"content": "hello",
|
||||
"links": ["https://affine.pro", "https://github.com/toeverything/AFFiNE"]
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let rendered = render_prompt_response(&messages, &template_params, ¶ms).unwrap();
|
||||
|
||||
assert_eq!(rendered.messages[0].content, "translate eng to chs: hello");
|
||||
assert_eq!(
|
||||
rendered.messages[1].content,
|
||||
"links:\n- https://affine.pro\n- https://github.com/toeverything/AFFiNE\n"
|
||||
);
|
||||
assert_eq!(rendered.warnings.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn appends_input_attachments_only_to_user_messages() {
|
||||
let messages = vec![message("system", "system"), message("user", "{{content}}")];
|
||||
let params = serde_json::from_value(json!({
|
||||
"content": "summarize",
|
||||
"attachments": [{
|
||||
"kind": "file_handle",
|
||||
"fileHandle": "file-1",
|
||||
"mimeType": "application/pdf"
|
||||
}]
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let rendered = render_prompt_response(&messages, &Map::new(), ¶ms).unwrap();
|
||||
|
||||
assert!(rendered.messages[0].attachments.is_none());
|
||||
assert_eq!(
|
||||
rendered.messages[1].attachments,
|
||||
Some(vec![json!({
|
||||
"kind": "file_handle",
|
||||
"fileHandle": "file-1",
|
||||
"mimeType": "application/pdf"
|
||||
})])
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,22 +2,45 @@ use llm_adapter::core::prompt_template::{parse_template, template_uses_key};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::{
|
||||
super::contracts::{PromptMessageContract, PromptSessionContract, PromptSessionResult},
|
||||
super::contracts::{PromptMessageContract, PromptSessionResult},
|
||||
render::render_prompt_response,
|
||||
};
|
||||
use crate::tiktoken::{Tokenizer, from_model_name};
|
||||
|
||||
const DEFAULT_HISTORY_INPUT_BYTES: usize = 128 * 1024;
|
||||
const MESSAGE_FRAMING_BYTES: usize = 16;
|
||||
|
||||
pub(super) fn render_session_prompt(
|
||||
request: &PromptSessionContract,
|
||||
prompt_messages: &[PromptMessageContract],
|
||||
action: Option<&str>,
|
||||
turns: &[PromptMessageContract],
|
||||
template_params: &Map<String, Value>,
|
||||
params: &Map<String, Value>,
|
||||
) -> std::result::Result<PromptSessionResult, String> {
|
||||
let tokenizer = session_tokenizer(request.prompt.model.as_deref());
|
||||
let mut selected_turns = take_session_turns(request, tokenizer.as_ref())?;
|
||||
let latest_turn = selected_turns.pop();
|
||||
render_session_prompt_with_budget(
|
||||
prompt_messages,
|
||||
action,
|
||||
turns,
|
||||
template_params,
|
||||
params,
|
||||
DEFAULT_HISTORY_INPUT_BYTES,
|
||||
)
|
||||
}
|
||||
|
||||
if prompt_uses_content(&request.prompt.messages)?
|
||||
&& !selected_turns.iter().any(message_is_assistant)
|
||||
fn render_session_prompt_with_budget(
|
||||
prompt_messages: &[PromptMessageContract],
|
||||
action: Option<&str>,
|
||||
turns: &[PromptMessageContract],
|
||||
template_params: &Map<String, Value>,
|
||||
params: &Map<String, Value>,
|
||||
history_input_bytes: usize,
|
||||
) -> std::result::Result<PromptSessionResult, String> {
|
||||
let (prior_turns, latest_turn) = turns
|
||||
.split_last()
|
||||
.map(|(latest, prior)| (prior, Some(latest.clone())))
|
||||
.unwrap_or((&[], None));
|
||||
|
||||
if prompt_uses_content(prompt_messages)?
|
||||
&& !prior_turns.iter().any(message_is_assistant)
|
||||
&& let Some(last_message) = latest_turn
|
||||
.as_ref()
|
||||
.filter(|message| message_role(message) == Some("user"))
|
||||
@@ -29,16 +52,17 @@ pub(super) fn render_session_prompt(
|
||||
}
|
||||
merged_params.insert("content".to_string(), Value::String(last_message.content.clone()));
|
||||
|
||||
let rendered = render_prompt_response(&request.prompt.messages, template_params, &merged_params)?;
|
||||
let rendered = render_prompt_response(prompt_messages, template_params, &merged_params)?;
|
||||
let mut messages = rendered.messages;
|
||||
let Some(first_user_message_index) = messages
|
||||
.iter()
|
||||
.position(|message| message_role(message) == Some("user"))
|
||||
else {
|
||||
ensure_messages_fit(&messages, &[], history_input_bytes)?;
|
||||
return Ok(PromptSessionResult {
|
||||
messages,
|
||||
warnings: rendered.warnings,
|
||||
prompt_message_positions: (0..request.prompt.messages.len()).map(|index| index as u32).collect(),
|
||||
prompt_message_positions: (0..prompt_messages.len()).map(|index| index as u32).collect(),
|
||||
});
|
||||
};
|
||||
|
||||
@@ -57,9 +81,10 @@ pub(super) fn render_session_prompt(
|
||||
messages[first_user_message_index].attachments = Some(merged_attachments);
|
||||
}
|
||||
|
||||
let selected_turns = select_history_turns(&messages, prior_turns, history_input_bytes)?;
|
||||
let prior_turn_count = selected_turns.len();
|
||||
messages.splice(first_user_message_index..first_user_message_index, selected_turns);
|
||||
let prompt_message_positions = (0..request.prompt.messages.len())
|
||||
let prompt_message_positions = (0..prompt_messages.len())
|
||||
.map(|index| {
|
||||
if index < first_user_message_index {
|
||||
index as u32
|
||||
@@ -81,52 +106,63 @@ pub(super) fn render_session_prompt(
|
||||
} else {
|
||||
latest_turn.as_ref().map(message_params).unwrap_or_default()
|
||||
};
|
||||
let rendered = render_prompt_response(&request.prompt.messages, template_params, &final_params)?;
|
||||
let rendered = render_prompt_response(prompt_messages, template_params, &final_params)?;
|
||||
|
||||
let trailing_turns = selected_turns
|
||||
let latest_turns = latest_turn
|
||||
.into_iter()
|
||||
.chain(latest_turn)
|
||||
.filter(prompt_message_should_survive)
|
||||
.collect::<Vec<_>>();
|
||||
let mut messages = rendered.messages;
|
||||
messages.extend(trailing_turns);
|
||||
let selected_turns = if action.is_some() {
|
||||
ensure_messages_fit(&messages, &latest_turns, history_input_bytes)?;
|
||||
Vec::new()
|
||||
} else {
|
||||
let mut fixed = messages.clone();
|
||||
fixed.extend(latest_turns.clone());
|
||||
select_history_turns(&fixed, prior_turns, history_input_bytes)?
|
||||
};
|
||||
messages.extend(selected_turns);
|
||||
messages.extend(latest_turns);
|
||||
|
||||
Ok(PromptSessionResult {
|
||||
messages,
|
||||
warnings: rendered.warnings,
|
||||
prompt_message_positions: (0..request.prompt.messages.len()).map(|index| index as u32).collect(),
|
||||
prompt_message_positions: (0..prompt_messages.len()).map(|index| index as u32).collect(),
|
||||
})
|
||||
}
|
||||
|
||||
fn session_tokenizer(model: Option<&str>) -> Option<Tokenizer> {
|
||||
let model = model?;
|
||||
if model.starts_with("gpt") {
|
||||
return from_model_name(model.to_string());
|
||||
fn estimated_message_bytes(message: &PromptMessageContract) -> usize {
|
||||
let mut size = MESSAGE_FRAMING_BYTES
|
||||
.saturating_add(message.role.len())
|
||||
.saturating_add(message.content.len());
|
||||
for value in [
|
||||
message.attachments.as_ref().map(serde_json::to_vec),
|
||||
message.params.as_ref().map(serde_json::to_vec),
|
||||
message.response_format.as_ref().map(serde_json::to_vec),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
{
|
||||
size = size.saturating_add(value.map(|bytes| bytes.len()).unwrap_or(usize::MAX));
|
||||
}
|
||||
if model.starts_with("dall") {
|
||||
return None;
|
||||
}
|
||||
|
||||
from_model_name("gpt-4".to_string())
|
||||
size
|
||||
}
|
||||
|
||||
fn take_session_turns(
|
||||
request: &PromptSessionContract,
|
||||
tokenizer: Option<&Tokenizer>,
|
||||
fn select_history_turns(
|
||||
fixed_messages: &[PromptMessageContract],
|
||||
history: &[PromptMessageContract],
|
||||
history_input_bytes: usize,
|
||||
) -> std::result::Result<Vec<PromptMessageContract>, String> {
|
||||
if request.prompt.action.is_some() {
|
||||
return Ok(request.turns.last().cloned().into_iter().collect());
|
||||
let mut size = fixed_messages.iter().fold(0usize, |size, message| {
|
||||
size.saturating_add(estimated_message_bytes(message))
|
||||
});
|
||||
if size > history_input_bytes {
|
||||
return Err("session input exceeds history byte budget".to_string());
|
||||
}
|
||||
|
||||
let mut picked = Vec::new();
|
||||
let mut size = request.prompt.prompt_tokens;
|
||||
|
||||
for message in request.turns.iter().rev() {
|
||||
let content = message.content.as_str();
|
||||
size += tokenizer
|
||||
.map(|tokenizer| tokenizer.count(content.to_string(), None))
|
||||
.unwrap_or(0);
|
||||
if size > request.max_token_size {
|
||||
for message in history.iter().rev() {
|
||||
size = size.saturating_add(estimated_message_bytes(message));
|
||||
if size > history_input_bytes {
|
||||
break;
|
||||
}
|
||||
picked.push(message.clone());
|
||||
@@ -136,6 +172,16 @@ fn take_session_turns(
|
||||
Ok(picked)
|
||||
}
|
||||
|
||||
fn ensure_messages_fit(
|
||||
prompt_messages: &[PromptMessageContract],
|
||||
latest_turns: &[PromptMessageContract],
|
||||
history_input_bytes: usize,
|
||||
) -> std::result::Result<(), String> {
|
||||
let mut messages = prompt_messages.to_vec();
|
||||
messages.extend_from_slice(latest_turns);
|
||||
select_history_turns(&messages, &[], history_input_bytes).map(|_| ())
|
||||
}
|
||||
|
||||
fn prompt_uses_content(messages: &[PromptMessageContract]) -> std::result::Result<bool, String> {
|
||||
for message in messages {
|
||||
if template_uses_key(&parse_template(&message.content)?, "content") {
|
||||
@@ -202,3 +248,102 @@ fn attachment_has_source(attachment: &Value) -> bool {
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::{Map, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn message(role: &str, content: &str) -> PromptMessageContract {
|
||||
serde_json::from_value(json!({ "role": role, "content": content })).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn byte_estimate_is_model_independent_and_utf8_aware() {
|
||||
let ascii = estimated_message_bytes(&message("user", "abc"));
|
||||
let cjk = estimated_message_bytes(&message("user", "中文文"));
|
||||
let emoji = estimated_message_bytes(&message("user", "😀😀😀"));
|
||||
assert!(ascii < cjk);
|
||||
assert!(cjk < emoji);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_latest_turn_and_only_a_contiguous_history_suffix() {
|
||||
let prompt = vec![message("system", "system")];
|
||||
let turns = vec![
|
||||
message("user", "old"),
|
||||
message("assistant", "recent"),
|
||||
message("user", "latest"),
|
||||
];
|
||||
let fixed_bytes = estimated_message_bytes(&prompt[0]) + estimated_message_bytes(&turns[2]);
|
||||
let budget = fixed_bytes + estimated_message_bytes(&turns[1]) + 8;
|
||||
let result = render_session_prompt_with_budget(&prompt, None, &turns, &Map::new(), &Map::new(), budget).unwrap();
|
||||
assert_eq!(
|
||||
result
|
||||
.messages
|
||||
.iter()
|
||||
.map(|message| message.content.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
["system", "recent", "latest"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_oversized_rendered_prompt_or_latest_turn() {
|
||||
let prompt = vec![message("system", "{{content}}")];
|
||||
let turns = vec![message("user", "large input")];
|
||||
let error = render_session_prompt_with_budget(&prompt, None, &turns, &Map::new(), &Map::new(), 1).unwrap_err();
|
||||
assert_eq!(error, "session input exceeds history byte budget");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_latest_user_content_params_and_file_handles_into_prompt() {
|
||||
let prompt = vec![message("user", "{{content}} {{tone}}")];
|
||||
let latest = serde_json::from_value(json!({
|
||||
"role": "user",
|
||||
"content": "Summarize this file",
|
||||
"attachments": [{
|
||||
"kind": "file_handle",
|
||||
"fileHandle": "file-1",
|
||||
"mimeType": "application/pdf"
|
||||
}],
|
||||
"params": { "tone": "brief" }
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let result = render_session_prompt(&prompt, None, &[latest], &Map::new(), &Map::new()).unwrap();
|
||||
|
||||
assert_eq!(result.messages.len(), 1);
|
||||
assert_eq!(result.messages[0].content, "Summarize this file brief");
|
||||
assert_eq!(
|
||||
result.messages[0].attachments,
|
||||
Some(vec![json!({
|
||||
"kind": "file_handle",
|
||||
"fileHandle": "file-1",
|
||||
"mimeType": "application/pdf"
|
||||
})])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn action_prompt_drops_prior_history_but_keeps_latest_turn() {
|
||||
let prompt = vec![message("system", "action")];
|
||||
let turns = vec![
|
||||
message("user", "old"),
|
||||
message("assistant", "old answer"),
|
||||
message("user", "latest"),
|
||||
];
|
||||
|
||||
let result = render_session_prompt(&prompt, Some("edit"), &turns, &Map::new(), &Map::new()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
result
|
||||
.messages
|
||||
.iter()
|
||||
.map(|message| message.content.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
["action", "latest"]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use llm_adapter::core::{self as adapter_core, EmbeddingRequest, ImageInput, ImageRequest, RerankRequest};
|
||||
use llm_adapter::core::{self as adapter_core, EmbeddingRequest, ImageRequest, RerankRequest};
|
||||
use napi::Result;
|
||||
use napi_derive::napi;
|
||||
use serde::Serialize;
|
||||
@@ -8,7 +8,7 @@ use super::contracts::{
|
||||
LlmImageRequestBuildContract, LlmImageRequestContract, LlmRequestContract, LlmRerankRequestContract,
|
||||
LlmStructuredRequestContract, ModelConditionsContract, PromptMessageContract,
|
||||
};
|
||||
use crate::llm::{LlmDispatchPayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload, host::invalid_arg};
|
||||
use crate::llm::{LlmDispatchPayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload, invalid_arg};
|
||||
|
||||
mod types;
|
||||
|
||||
@@ -62,38 +62,11 @@ pub(crate) fn build_image_request(request: ImageRequest) -> Result<ImageRequest>
|
||||
}
|
||||
|
||||
pub(crate) fn build_image_request_from_messages(request: LlmImageRequestBuildContract) -> Result<ImageRequest> {
|
||||
let protocol = request.protocol.clone();
|
||||
let mut request =
|
||||
let request =
|
||||
adapter_core::build_image_request_from_prompt_messages(to_adapter(&request)?).map_err(map_builder_error)?;
|
||||
if protocol == "fal_image" {
|
||||
keep_fal_data_uri_inputs_as_urls(&mut request);
|
||||
}
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
fn keep_fal_data_uri_inputs_as_urls(request: &mut ImageRequest) {
|
||||
let ImageRequest::Edit(edit) = request else {
|
||||
return;
|
||||
};
|
||||
|
||||
for image in &mut edit.images {
|
||||
let replacement = match image {
|
||||
ImageInput::Data {
|
||||
data_base64,
|
||||
media_type,
|
||||
..
|
||||
} => Some(ImageInput::Url {
|
||||
url: format!("data:{media_type};base64,{data_base64}"),
|
||||
media_type: Some(media_type.clone()),
|
||||
}),
|
||||
_ => None,
|
||||
};
|
||||
if let Some(replacement) = replacement {
|
||||
*image = replacement;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn infer_prompt_model_conditions(messages: Vec<PromptMessageInput>) -> Result<ModelConditionsContract> {
|
||||
let messages = adapter_core::canonicalize_prompt_messages(to_adapter_prompt_messages(messages)?);
|
||||
serde_json::to_value(adapter_core::infer_model_conditions_from_prompt_messages(messages))
|
||||
@@ -471,11 +444,10 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_keep_fal_data_uri_image_inputs_as_urls() {
|
||||
fn should_canonicalize_data_uri_image_inputs() {
|
||||
let response = llm_build_image_request_from_messages(
|
||||
serde_json::from_value(json!({
|
||||
"model": "lora/image-to-image",
|
||||
"protocol": "fal_image",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": "restyle",
|
||||
@@ -493,11 +465,7 @@ mod tests {
|
||||
let response = serde_json::to_value(response).unwrap();
|
||||
assert_eq!(
|
||||
response.pointer("/images/0"),
|
||||
Some(&json!({
|
||||
"kind": "url",
|
||||
"url": "data:image/png;base64,aW1n",
|
||||
"media_type": "image/png"
|
||||
}))
|
||||
Some(&json!({ "kind": "data", "data_base64": "aW1n", "media_type": "image/png" }))
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ use super::super::contracts::{
|
||||
LlmRerankRequestContract, LlmStructuredRequestContract, RerankCandidate as ContractRerankCandidate, ToolContract,
|
||||
};
|
||||
use crate::llm::{
|
||||
LlmDispatchPayload, LlmMiddlewarePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload, host::invalid_arg,
|
||||
LlmDispatchPayload, LlmMiddlewarePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload, invalid_arg,
|
||||
map_json_error,
|
||||
};
|
||||
|
||||
|
||||
@@ -1,27 +1,5 @@
|
||||
mod dispatch;
|
||||
mod middleware;
|
||||
mod payload;
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) use dispatch::AsyncLlmDispatchPreparedTask;
|
||||
pub(crate) use dispatch::{
|
||||
dispatch_prepared_image_route_payloads, dispatch_prepared_structured_routes,
|
||||
parse_prepared_chat_routes_with_middleware, parse_prepared_chat_routes_without_middleware,
|
||||
};
|
||||
pub use dispatch::{
|
||||
llm_dispatch_prepared, llm_embedding_dispatch, llm_embedding_dispatch_prepared, llm_image_dispatch_prepared,
|
||||
llm_plan_attachment_reference, llm_rerank_dispatch, llm_rerank_dispatch_prepared, llm_resolve_request_intent,
|
||||
llm_structured_dispatch, llm_structured_dispatch_prepared,
|
||||
};
|
||||
pub(crate) use llm_adapter::middleware::StreamPipeline;
|
||||
#[cfg(test)]
|
||||
pub(crate) use middleware::resolve_request_chain;
|
||||
pub(crate) use middleware::{
|
||||
apply_request_middlewares, apply_structured_request_middlewares, backend_transport_error, map_backend_error,
|
||||
map_json_error, parse_embedding_protocol, parse_protocol, parse_rerank_protocol, parse_structured_protocol,
|
||||
resolve_stream_chain,
|
||||
};
|
||||
pub(crate) use payload::{
|
||||
LlmDispatchPayload, LlmEmbeddingDispatchPayload, LlmMiddlewarePayload, LlmPreparedImageDispatchRoutePayload,
|
||||
LlmRerankDispatchPayload, LlmRoutedBackendPayload, LlmStructuredDispatchPayload,
|
||||
LlmDispatchPayload, LlmMiddlewarePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload,
|
||||
};
|
||||
|
||||
@@ -1,15 +1,10 @@
|
||||
use llm_adapter::{
|
||||
backend::BackendConfig,
|
||||
core::{CoreRequest, EmbeddingRequest, RerankRequest, StructuredRequest},
|
||||
core::{CoreRequest, RerankRequest, StructuredRequest},
|
||||
middleware::MiddlewareConfig,
|
||||
router::SerializablePreparedRoute,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::llm::core::contracts::{
|
||||
LlmEmbeddingRequestContract, LlmImageRequestContract, LlmRequestContract, LlmRerankRequestContract,
|
||||
LlmStructuredRequestContract,
|
||||
};
|
||||
use crate::llm::core::contracts::{LlmRequestContract, LlmRerankRequestContract, LlmStructuredRequestContract};
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
@@ -41,15 +36,6 @@ pub(crate) struct LlmDispatchPayload {
|
||||
pub(crate) middleware: LlmMiddlewarePayload,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub(crate) struct LlmRoutedBackendPayload {
|
||||
pub(crate) provider_id: String,
|
||||
pub(crate) protocol: String,
|
||||
pub(crate) model: String,
|
||||
#[serde(alias = "backendConfig")]
|
||||
pub(crate) config: BackendConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
#[serde(try_from = "LlmStructuredRequestContract")]
|
||||
pub(crate) struct LlmStructuredDispatchPayload {
|
||||
@@ -59,20 +45,6 @@ pub(crate) struct LlmStructuredDispatchPayload {
|
||||
pub(crate) middleware: LlmMiddlewarePayload,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
#[serde(from = "LlmEmbeddingRequestContract")]
|
||||
pub(crate) struct LlmEmbeddingDispatchPayload {
|
||||
pub(crate) request: EmbeddingRequest,
|
||||
}
|
||||
|
||||
impl From<LlmEmbeddingRequestContract> for LlmEmbeddingDispatchPayload {
|
||||
fn from(request: LlmEmbeddingRequestContract) -> Self {
|
||||
Self {
|
||||
request: request.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
#[serde(from = "LlmRerankRequestContract")]
|
||||
pub(crate) struct LlmRerankDispatchPayload {
|
||||
@@ -87,128 +59,3 @@ impl From<LlmRerankRequestContract> for LlmRerankDispatchPayload {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) type LlmPreparedImageDispatchRoutePayload = SerializablePreparedRoute<LlmImageRequestContract>;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use llm_adapter::router::SerializablePreparedRoute;
|
||||
|
||||
use super::{
|
||||
LlmDispatchPayload, LlmPreparedImageDispatchRoutePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn prepared_chat_route_payload_deserializes_nested_request() {
|
||||
let payload = serde_json::from_value::<Vec<SerializablePreparedRoute<LlmDispatchPayload>>>(serde_json::json!([
|
||||
{
|
||||
"provider_id": "openai-primary",
|
||||
"protocol": "openai_chat",
|
||||
"model": "gpt-5-mini",
|
||||
"config": {
|
||||
"base_url": "https://api.openai.com",
|
||||
"auth_token": "test-key"
|
||||
},
|
||||
"request": {
|
||||
"model": "gpt-5-mini",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{ "type": "text", "text": "hello" }]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
]))
|
||||
.expect("prepared chat route payload should deserialize");
|
||||
|
||||
assert_eq!(payload[0].model, "gpt-5-mini");
|
||||
assert_eq!(payload[0].request.request.model, "gpt-5-mini");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepared_structured_route_payload_deserializes_nested_request() {
|
||||
let payload =
|
||||
serde_json::from_value::<Vec<SerializablePreparedRoute<LlmStructuredDispatchPayload>>>(serde_json::json!([
|
||||
{
|
||||
"provider_id": "openai-primary",
|
||||
"protocol": "openai_responses",
|
||||
"model": "gpt-5-mini",
|
||||
"config": {
|
||||
"base_url": "https://api.openai.com",
|
||||
"auth_token": "test-key"
|
||||
},
|
||||
"request": {
|
||||
"model": "gpt-5-mini",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{ "type": "text", "text": "hello" }]
|
||||
}
|
||||
],
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"summary": { "type": "string" }
|
||||
},
|
||||
"required": ["summary"]
|
||||
}
|
||||
}
|
||||
}
|
||||
]))
|
||||
.expect("prepared structured route payload should deserialize");
|
||||
|
||||
assert_eq!(payload[0].model, "gpt-5-mini");
|
||||
assert_eq!(payload[0].request.request.model, "gpt-5-mini");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepared_rerank_route_payload_deserializes_nested_request() {
|
||||
let payload =
|
||||
serde_json::from_value::<Vec<SerializablePreparedRoute<LlmRerankDispatchPayload>>>(serde_json::json!([
|
||||
{
|
||||
"provider_id": "openai-primary",
|
||||
"protocol": "openai_chat",
|
||||
"model": "gpt-5-mini",
|
||||
"config": {
|
||||
"base_url": "https://api.openai.com",
|
||||
"auth_token": "test-key"
|
||||
},
|
||||
"request": {
|
||||
"model": "gpt-5-mini",
|
||||
"query": "hello",
|
||||
"candidates": [{ "text": "world" }]
|
||||
}
|
||||
}
|
||||
]))
|
||||
.expect("prepared rerank route payload should deserialize");
|
||||
|
||||
assert_eq!(payload[0].model, "gpt-5-mini");
|
||||
assert_eq!(payload[0].request.request.model, "gpt-5-mini");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepared_image_route_payload_deserializes_nested_request() {
|
||||
let payload = serde_json::from_value::<Vec<LlmPreparedImageDispatchRoutePayload>>(serde_json::json!([
|
||||
{
|
||||
"provider_id": "openai-primary",
|
||||
"protocol": "openai_images",
|
||||
"model": "gpt-image-1",
|
||||
"config": {
|
||||
"base_url": "https://api.openai.com",
|
||||
"auth_token": "test-key",
|
||||
"request_layer": "openai_images"
|
||||
},
|
||||
"request": {
|
||||
"model": "gpt-image-1",
|
||||
"prompt": "draw",
|
||||
"operation": "generate"
|
||||
}
|
||||
}
|
||||
]))
|
||||
.expect("prepared image route payload should deserialize");
|
||||
|
||||
assert_eq!(payload[0].model, "gpt-image-1");
|
||||
assert_eq!(payload[0].request.prompt, "draw");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,19 +1,17 @@
|
||||
mod action;
|
||||
pub(crate) mod byok;
|
||||
mod contract_schema;
|
||||
mod core;
|
||||
mod ffi;
|
||||
mod host;
|
||||
mod prompt_catalog;
|
||||
pub(crate) mod route;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
pub use core::{
|
||||
capability::{llm_match_model_capabilities, llm_resolve_requested_model_match},
|
||||
pub use self::core::{
|
||||
capability::llm_match_model_capabilities,
|
||||
model_registry::{llm_match_model_registry, llm_resolve_model_registry_variant},
|
||||
prompt::{
|
||||
llm_collect_prompt_metadata, llm_count_prompt_tokens, llm_get_built_in_prompt_spec, llm_list_built_in_prompt_specs,
|
||||
llm_render_built_in_prompt, llm_render_built_in_session_prompt, llm_render_prompt, llm_render_session_prompt,
|
||||
llm_get_built_in_prompt_spec, llm_list_built_in_prompt_specs, llm_render_built_in_prompt,
|
||||
llm_render_built_in_session_prompt,
|
||||
},
|
||||
request_builder::{
|
||||
llm_build_canonical_request, llm_build_canonical_structured_request, llm_build_embedding_request,
|
||||
@@ -22,29 +20,34 @@ pub use core::{
|
||||
structured_output::{llm_canonical_json_schema_hash, llm_validate_json_schema},
|
||||
};
|
||||
|
||||
pub use action::run_native_action_recipe_prepared_stream;
|
||||
pub use contract_schema::{
|
||||
llm_compile_execution_plan, llm_get_contract_schema, llm_normalize_prepared_routes, llm_validate_contract,
|
||||
pub use action::copilot_action_recipe;
|
||||
pub use byok::{
|
||||
ByokCapabilityInput, ByokCatalogModelOutput, ByokCatalogOutput, ByokCatalogProviderOutput, ByokEndpointInput,
|
||||
ByokLocalLeaseOutput, ByokModelDeclarationInput, ByokModelProbeCheckOutput, ByokModelProbeOutput,
|
||||
ByokProbeCheckInput, ByokProbeResultOutput, ByokProbeStatusOutput, ByokProfileDefinitionInput, ByokProfileOutput,
|
||||
ByokValidationOutput, CreateByokLocalLeaseInput, CreateByokLocalLeaseProviderInput, CreateByokProfileInput,
|
||||
ProbeByokDraftInput, ProbeByokProfileInput, ReorderByokProfilesInput, ReplaceByokProfileInput,
|
||||
RotateByokCredentialInput, byok_catalog,
|
||||
};
|
||||
#[cfg(test)]
|
||||
pub(crate) use ffi::{AsyncLlmDispatchPreparedTask, resolve_request_chain};
|
||||
|
||||
#[napi_derive::napi(catch_unwind)]
|
||||
pub fn llm_get_byok_catalog() -> ByokCatalogOutput {
|
||||
byok_catalog()
|
||||
}
|
||||
pub(crate) use byok::{ByokProfileDefinition, validate_definition};
|
||||
pub use contract_schema::{llm_get_contract_schema, llm_validate_contract};
|
||||
pub(crate) use ffi::{
|
||||
LlmDispatchPayload, LlmEmbeddingDispatchPayload, LlmMiddlewarePayload, LlmPreparedImageDispatchRoutePayload,
|
||||
LlmRerankDispatchPayload, LlmRoutedBackendPayload, LlmStructuredDispatchPayload, StreamPipeline,
|
||||
apply_request_middlewares, apply_structured_request_middlewares, backend_transport_error,
|
||||
dispatch_prepared_image_route_payloads, dispatch_prepared_structured_routes, map_backend_error, map_json_error,
|
||||
parse_embedding_protocol, parse_prepared_chat_routes_with_middleware, parse_prepared_chat_routes_without_middleware,
|
||||
parse_protocol, parse_rerank_protocol, parse_structured_protocol, resolve_stream_chain,
|
||||
LlmDispatchPayload, LlmMiddlewarePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload,
|
||||
};
|
||||
pub use ffi::{
|
||||
llm_dispatch_prepared, llm_embedding_dispatch, llm_embedding_dispatch_prepared, llm_image_dispatch_prepared,
|
||||
llm_plan_attachment_reference, llm_rerank_dispatch, llm_rerank_dispatch_prepared, llm_resolve_request_intent,
|
||||
llm_structured_dispatch, llm_structured_dispatch_prepared,
|
||||
};
|
||||
pub(crate) use host::{
|
||||
LlmStreamHandle, STREAM_ABORTED_REASON, STREAM_CALLBACK_DISPATCH_FAILED_REASON, STREAM_END_MARKER, emit_error_event,
|
||||
};
|
||||
pub use host::{
|
||||
llm_dispatch_prepared_stream, llm_dispatch_tool_loop_stream, llm_dispatch_tool_loop_stream_prepared,
|
||||
llm_dispatch_tool_loop_stream_routed,
|
||||
pub use prompt_catalog::llm_get_built_in_route_options;
|
||||
pub use route::{
|
||||
CopilotAccessProjection, CopilotExecuteInput, CopilotManagedTier, CopilotRouteCheckInput, CopilotTargetOverrideInput,
|
||||
};
|
||||
|
||||
pub(crate) fn invalid_arg(message: impl Into<String>) -> napi::Error {
|
||||
napi::Error::new(napi::Status::InvalidArg, message.into())
|
||||
}
|
||||
|
||||
pub(crate) fn map_json_error(error: serde_json::Error) -> napi::Error {
|
||||
invalid_arg(error.to_string())
|
||||
}
|
||||
|
||||
@@ -46,15 +46,12 @@ pub struct PromptSpecMessage {
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BuiltInPromptSpec {
|
||||
pub name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub action: Option<String>,
|
||||
pub model: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub optional_models: Option<Vec<String>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub config: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
@@ -64,6 +61,69 @@ pub struct BuiltInPromptSpec {
|
||||
pub messages: Vec<PromptSpecMessage>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct PromptCatalogSpec {
|
||||
name: String,
|
||||
#[serde(default)]
|
||||
action: Option<String>,
|
||||
#[serde(default)]
|
||||
managed_route: Option<BuiltInManagedRouteSpec>,
|
||||
#[serde(default)]
|
||||
config: Option<Value>,
|
||||
#[serde(default)]
|
||||
params: Option<BTreeMap<String, PromptParamSpec>>,
|
||||
#[serde(default)]
|
||||
builtins: Option<Vec<PromptBuiltin>>,
|
||||
messages: Vec<PromptSpecMessage>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct BuiltInManagedRouteSpec {
|
||||
targets: Vec<String>,
|
||||
#[serde(default)]
|
||||
premium_targets: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
selectable_targets: Vec<BuiltInManagedTargetSpec>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct BuiltInManagedTargetSpec {
|
||||
id: String,
|
||||
model_id: String,
|
||||
display_name: String,
|
||||
minimum_tier: BuiltInManagedTargetTier,
|
||||
}
|
||||
|
||||
#[napi(string_enum)]
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum BuiltInManagedTargetTier {
|
||||
Standard,
|
||||
Premium,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BuiltInManagedTarget {
|
||||
pub id: String,
|
||||
pub display_name: String,
|
||||
pub minimum_tier: BuiltInManagedTargetTier,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BuiltInRouteOptions {
|
||||
pub route_id: String,
|
||||
pub standard_default_target_id: Option<String>,
|
||||
pub premium_default_target_id: Option<String>,
|
||||
pub choices: Vec<BuiltInManagedTarget>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct BuiltInPromptMessage {
|
||||
@@ -73,20 +133,29 @@ pub(crate) struct BuiltInPromptMessage {
|
||||
pub(crate) params: Option<Map<String, Value>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct BuiltInPrompt {
|
||||
pub(crate) name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) action: Option<String>,
|
||||
pub(crate) model: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) optional_models: Option<Vec<String>>,
|
||||
pub(crate) managed_targets: Vec<String>,
|
||||
pub(crate) managed_premium_targets: Option<Vec<String>>,
|
||||
#[serde(skip)]
|
||||
pub(crate) managed_selectable_targets: Vec<BuiltInManagedTargetDefinition>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) config: Option<Value>,
|
||||
pub(crate) messages: Vec<BuiltInPromptMessage>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct BuiltInManagedTargetDefinition {
|
||||
pub(crate) id: String,
|
||||
pub(crate) model_id: String,
|
||||
pub(crate) display_name: String,
|
||||
pub(crate) minimum_tier: BuiltInManagedTargetTier,
|
||||
}
|
||||
|
||||
struct PromptCatalog {
|
||||
specs: Vec<BuiltInPromptSpec>,
|
||||
prompts: Vec<BuiltInPrompt>,
|
||||
@@ -112,16 +181,89 @@ pub(crate) fn built_in_prompt(name: &str) -> Option<&'static BuiltInPrompt> {
|
||||
.and_then(|index| BUILTIN_PROMPT_CATALOG.prompts.get(*index))
|
||||
}
|
||||
|
||||
pub(crate) fn built_in_managed_targets(name: &str, premium: bool) -> Option<&'static [String]> {
|
||||
built_in_prompt(name).and_then(|prompt| {
|
||||
let targets = if premium {
|
||||
prompt
|
||||
.managed_premium_targets
|
||||
.as_deref()
|
||||
.unwrap_or(prompt.managed_targets.as_slice())
|
||||
} else {
|
||||
prompt.managed_targets.as_slice()
|
||||
};
|
||||
(!targets.is_empty()).then_some(targets)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn built_in_managed_target(
|
||||
name: &str,
|
||||
target_id: &str,
|
||||
premium: bool,
|
||||
) -> Option<&'static BuiltInManagedTargetDefinition> {
|
||||
built_in_prompt(name)?
|
||||
.managed_selectable_targets
|
||||
.iter()
|
||||
.find(|target| target.id == target_id && (premium || target.minimum_tier == BuiltInManagedTargetTier::Standard))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn llm_get_built_in_route_options(name: String) -> Option<BuiltInRouteOptions> {
|
||||
let prompt = built_in_prompt(&name)?;
|
||||
if prompt.managed_selectable_targets.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let target_id_for_model = |model: Option<&String>| {
|
||||
model.and_then(|model| {
|
||||
prompt
|
||||
.managed_selectable_targets
|
||||
.iter()
|
||||
.find(|target| &target.model_id == model)
|
||||
.map(|target| target.id.clone())
|
||||
})
|
||||
};
|
||||
Some(BuiltInRouteOptions {
|
||||
route_id: prompt.name.clone(),
|
||||
standard_default_target_id: target_id_for_model(prompt.managed_targets.first()),
|
||||
premium_default_target_id: target_id_for_model(
|
||||
prompt
|
||||
.managed_premium_targets
|
||||
.as_ref()
|
||||
.and_then(|targets| targets.first())
|
||||
.or_else(|| prompt.managed_targets.first()),
|
||||
),
|
||||
choices: prompt
|
||||
.managed_selectable_targets
|
||||
.iter()
|
||||
.map(|target| BuiltInManagedTarget {
|
||||
id: target.id.clone(),
|
||||
display_name: target.display_name.clone(),
|
||||
minimum_tier: target.minimum_tier,
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
}
|
||||
|
||||
impl PromptCatalog {
|
||||
fn load() -> Result<Self, String> {
|
||||
let partials: BTreeMap<String, String> =
|
||||
serde_json::from_str(PROMPT_PARTIALS_SOURCE).map_err(|error| format!("invalid prompt partials JSON: {error}"))?;
|
||||
let specs: Vec<BuiltInPromptSpec> =
|
||||
let catalog_specs: Vec<PromptCatalogSpec> =
|
||||
serde_json::from_str(PROMPT_SPECS_SOURCE).map_err(|error| format!("invalid prompt spec JSON: {error}"))?;
|
||||
let prompts = specs
|
||||
let prompts = catalog_specs
|
||||
.iter()
|
||||
.map(|spec| compile_prompt_spec(spec, &partials))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let specs = catalog_specs
|
||||
.into_iter()
|
||||
.map(|spec| BuiltInPromptSpec {
|
||||
name: spec.name,
|
||||
action: spec.action,
|
||||
config: spec.config.filter(|value| !value.is_null()),
|
||||
params: spec.params,
|
||||
builtins: spec.builtins,
|
||||
messages: spec.messages,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(Self {
|
||||
specs_by_name: specs
|
||||
@@ -140,7 +282,17 @@ impl PromptCatalog {
|
||||
}
|
||||
}
|
||||
|
||||
fn compile_prompt_spec(spec: &BuiltInPromptSpec, partials: &BTreeMap<String, String>) -> Result<BuiltInPrompt, String> {
|
||||
fn compile_prompt_spec(spec: &PromptCatalogSpec, partials: &BTreeMap<String, String>) -> Result<BuiltInPrompt, String> {
|
||||
if spec
|
||||
.managed_route
|
||||
.as_ref()
|
||||
.is_some_and(|route| !valid_managed_route(route))
|
||||
{
|
||||
return Err(format!("Prompt \"{}\" has an invalid managed route", spec.name));
|
||||
}
|
||||
if !spec.messages.is_empty() && spec.managed_route.is_none() {
|
||||
return Err(format!("Executable prompt \"{}\" requires a managed route", spec.name));
|
||||
}
|
||||
let resolved_templates = spec
|
||||
.messages
|
||||
.iter()
|
||||
@@ -186,13 +338,71 @@ fn compile_prompt_spec(spec: &BuiltInPromptSpec, partials: &BTreeMap<String, Str
|
||||
Ok(BuiltInPrompt {
|
||||
name: spec.name.clone(),
|
||||
action: spec.action.clone(),
|
||||
model: spec.model.clone(),
|
||||
optional_models: spec.optional_models.clone(),
|
||||
managed_targets: spec
|
||||
.managed_route
|
||||
.as_ref()
|
||||
.map(|route| route.targets.clone())
|
||||
.unwrap_or_default(),
|
||||
managed_premium_targets: spec
|
||||
.managed_route
|
||||
.as_ref()
|
||||
.and_then(|route| route.premium_targets.clone()),
|
||||
managed_selectable_targets: spec
|
||||
.managed_route
|
||||
.as_ref()
|
||||
.map(|route| {
|
||||
route
|
||||
.selectable_targets
|
||||
.iter()
|
||||
.map(|target| BuiltInManagedTargetDefinition {
|
||||
id: target.id.clone(),
|
||||
model_id: target.model_id.clone(),
|
||||
display_name: target.display_name.clone(),
|
||||
minimum_tier: target.minimum_tier,
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default(),
|
||||
config: spec.config.clone().filter(|value| !value.is_null()),
|
||||
messages,
|
||||
})
|
||||
}
|
||||
|
||||
fn valid_managed_route(route: &BuiltInManagedRouteSpec) -> bool {
|
||||
if route.targets.is_empty()
|
||||
|| route.targets.iter().any(|target| target.trim().is_empty())
|
||||
|| route
|
||||
.premium_targets
|
||||
.as_ref()
|
||||
.is_some_and(|targets| targets.is_empty() || targets.iter().any(|target| target.trim().is_empty()))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let ids = route
|
||||
.selectable_targets
|
||||
.iter()
|
||||
.map(|target| target.id.as_str())
|
||||
.collect::<BTreeSet<_>>();
|
||||
let models = route
|
||||
.selectable_targets
|
||||
.iter()
|
||||
.map(|target| target.model_id.as_str())
|
||||
.collect::<BTreeSet<_>>();
|
||||
if route.selectable_targets.iter().any(|target| {
|
||||
target.id.trim().is_empty() || target.model_id.trim().is_empty() || target.display_name.trim().is_empty()
|
||||
}) || ids.len() != route.selectable_targets.len()
|
||||
|| models.len() != route.selectable_targets.len()
|
||||
{
|
||||
return false;
|
||||
}
|
||||
route.selectable_targets.is_empty()
|
||||
|| route
|
||||
.targets
|
||||
.iter()
|
||||
.chain(route.premium_targets.iter().flatten())
|
||||
.all(|model| models.contains(model.as_str()))
|
||||
}
|
||||
|
||||
fn normalize_prompt_param(spec: &PromptParamSpec) -> Value {
|
||||
match spec.enum_values.as_ref() {
|
||||
Some(values) if !values.is_empty() => {
|
||||
@@ -250,7 +460,7 @@ fn resolve_prompt_template(template: &str, partials: &BTreeMap<String, String>)
|
||||
Err("Prompt partial expansion exceeded maximum depth".to_string())
|
||||
}
|
||||
|
||||
fn validate_builtins(spec: &BuiltInPromptSpec, templates: &[String]) -> Result<(), String> {
|
||||
fn validate_builtins(spec: &PromptCatalogSpec, templates: &[String]) -> Result<(), String> {
|
||||
let declared = spec
|
||||
.builtins
|
||||
.clone()
|
||||
@@ -355,26 +565,35 @@ mod tests {
|
||||
);
|
||||
|
||||
let chat = built_in_prompt("Chat With AFFiNE AI").expect("chat prompt");
|
||||
assert_eq!(chat.model, "gpt-5.6-luna");
|
||||
assert_eq!(chat.managed_targets, ["gpt-5.6-luna"]);
|
||||
assert_eq!(
|
||||
chat
|
||||
.optional_models
|
||||
.as_ref()
|
||||
.map(|models| models.iter().map(String::as_str).collect::<Vec<_>>()),
|
||||
Some(vec![
|
||||
"gpt-5.6-luna",
|
||||
"gpt-5.6-terra",
|
||||
"gemini-3.6-flash",
|
||||
"claude-sonnet-4-6"
|
||||
])
|
||||
.managed_premium_targets
|
||||
.as_deref()
|
||||
.map(|targets| targets.iter().map(String::as_str).collect::<Vec<_>>()),
|
||||
Some(vec!["gpt-5.6-luna"])
|
||||
);
|
||||
let options = llm_get_built_in_route_options(chat.name.clone()).expect("chat route options");
|
||||
assert_eq!(options.standard_default_target_id.as_deref(), Some("luna"));
|
||||
assert_eq!(options.premium_default_target_id.as_deref(), Some("luna"));
|
||||
assert_eq!(options.choices.len(), 4);
|
||||
assert_eq!(
|
||||
built_in_managed_target(&chat.name, "terra", false).map(|target| target.model_id.as_str()),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
chat.config.as_ref().and_then(|config| config.get("proModels")),
|
||||
Some(&serde_json::json!([
|
||||
"gpt-5.6-terra",
|
||||
"gemini-3.6-flash",
|
||||
"claude-sonnet-4-6"
|
||||
]))
|
||||
built_in_managed_target(&chat.name, "terra", true).map(|target| target.model_id.as_str()),
|
||||
Some("gpt-5.6-terra")
|
||||
);
|
||||
|
||||
let transcript = built_in_prompt("Transcript audio structured").expect("transcript prompt");
|
||||
assert_eq!(transcript.managed_targets, ["gemini-3.5-flash-lite"]);
|
||||
assert_eq!(
|
||||
transcript
|
||||
.managed_premium_targets
|
||||
.as_deref()
|
||||
.map(|targets| targets.iter().map(String::as_str).collect::<Vec<_>>()),
|
||||
Some(vec!["gemini-3.6-flash"])
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,288 @@
|
||||
use llm_adapter::capability::{
|
||||
AttachmentKind, AttachmentSource, ModelFeature, ModelInput, ModelOutput, ModelRequirements,
|
||||
};
|
||||
|
||||
use crate::llm::{
|
||||
prompt_catalog::{built_in_managed_target, built_in_managed_targets},
|
||||
route::CopilotManagedTier,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum RouteOperation {
|
||||
Chat,
|
||||
Structured,
|
||||
Embedding,
|
||||
Rerank,
|
||||
Image,
|
||||
Transcription,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub(crate) struct CatalogSlot {
|
||||
pub(crate) id: &'static str,
|
||||
pub(crate) operation: RouteOperation,
|
||||
pub(crate) requirements: ModelRequirements,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum QuotaPolicy {
|
||||
Metered,
|
||||
Internal,
|
||||
System,
|
||||
}
|
||||
|
||||
pub(crate) fn quota_policy(slot: &CatalogSlot, built_in_route_id: Option<&str>) -> QuotaPolicy {
|
||||
if matches!(slot.id, "index.embedding" | "search.rerank") {
|
||||
QuotaPolicy::System
|
||||
} else if built_in_route_id == Some("Summary as title") {
|
||||
QuotaPolicy::Internal
|
||||
} else {
|
||||
QuotaPolicy::Metered
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn managed_targets(
|
||||
slot: &CatalogSlot,
|
||||
built_in_route_id: Option<&str>,
|
||||
managed_tier: CopilotManagedTier,
|
||||
) -> Option<Vec<String>> {
|
||||
if let Some(route_id) = built_in_route_id {
|
||||
return built_in_managed_targets(route_id, managed_tier == CopilotManagedTier::Premium).map(<[String]>::to_vec);
|
||||
}
|
||||
match slot.id {
|
||||
"index.embedding" => Some(vec!["gemini-embedding-001".to_string()]),
|
||||
"search.rerank" => Some(vec!["gpt-4o-mini".to_string()]),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn managed_selected_target(
|
||||
built_in_route_id: Option<&str>,
|
||||
target_id: &str,
|
||||
managed_tier: CopilotManagedTier,
|
||||
) -> Option<String> {
|
||||
built_in_managed_target(
|
||||
built_in_route_id?,
|
||||
target_id,
|
||||
managed_tier == CopilotManagedTier::Premium,
|
||||
)
|
||||
.map(|target| target.model_id.clone())
|
||||
}
|
||||
|
||||
pub(crate) fn slot(id: &str) -> Option<CatalogSlot> {
|
||||
let (canonical, operation, input, output, features, attachment_kinds, attachment_sources) = match id {
|
||||
"chat.default" | "prompt.text" => (
|
||||
if id == "chat.default" {
|
||||
"chat.default"
|
||||
} else {
|
||||
"prompt.text"
|
||||
},
|
||||
RouteOperation::Chat,
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Text],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"chat.structured" | "prompt.structured" => (
|
||||
if id == "chat.structured" {
|
||||
"chat.structured"
|
||||
} else {
|
||||
"prompt.structured"
|
||||
},
|
||||
RouteOperation::Structured,
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Structured],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"action.mindmap.generate" => (
|
||||
"action.mindmap.generate",
|
||||
RouteOperation::Structured,
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Structured],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"action.slides.outline" => (
|
||||
"action.slides.outline",
|
||||
RouteOperation::Structured,
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Structured],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"image.generate"
|
||||
| "action.image.filter.sketch"
|
||||
| "action.image.filter.clay"
|
||||
| "action.image.filter.anime"
|
||||
| "action.image.filter.pixel" => (
|
||||
match id {
|
||||
"image.generate" => "image.generate",
|
||||
"action.image.filter.sketch" => "action.image.filter.sketch",
|
||||
"action.image.filter.clay" => "action.image.filter.clay",
|
||||
"action.image.filter.anime" => "action.image.filter.anime",
|
||||
_ => "action.image.filter.pixel",
|
||||
},
|
||||
RouteOperation::Image,
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Image],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"index.embedding" => (
|
||||
"index.embedding",
|
||||
RouteOperation::Embedding,
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Embedding],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"search.rerank" => (
|
||||
"search.rerank",
|
||||
RouteOperation::Rerank,
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Rerank],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"transcript.audio" => (
|
||||
"transcript.audio",
|
||||
RouteOperation::Transcription,
|
||||
vec![ModelInput::Audio],
|
||||
vec![ModelOutput::Structured],
|
||||
vec![],
|
||||
vec![AttachmentKind::Audio],
|
||||
vec![
|
||||
AttachmentSource::Url,
|
||||
AttachmentSource::Data,
|
||||
AttachmentSource::Bytes,
|
||||
AttachmentSource::FileHandle,
|
||||
],
|
||||
),
|
||||
_ => return None,
|
||||
};
|
||||
Some(CatalogSlot {
|
||||
id: canonical,
|
||||
operation,
|
||||
requirements: ModelRequirements {
|
||||
input,
|
||||
output,
|
||||
features,
|
||||
attachment_kinds,
|
||||
attachment_sources,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn with_request_requirements(
|
||||
mut slot: CatalogSlot,
|
||||
needs_tools: bool,
|
||||
attachment_kinds: Vec<AttachmentKind>,
|
||||
attachment_sources: Vec<AttachmentSource>,
|
||||
) -> CatalogSlot {
|
||||
if needs_tools {
|
||||
slot.requirements.features.push(ModelFeature::ToolCalling);
|
||||
}
|
||||
for kind in attachment_kinds {
|
||||
let input = match kind {
|
||||
AttachmentKind::Image => ModelInput::Image,
|
||||
AttachmentKind::Audio => ModelInput::Audio,
|
||||
AttachmentKind::File => ModelInput::File,
|
||||
};
|
||||
if !slot.requirements.input.contains(&input) {
|
||||
slot.requirements.input.push(input);
|
||||
}
|
||||
if !slot.requirements.attachment_kinds.contains(&kind) {
|
||||
slot.requirements.attachment_kinds.push(kind);
|
||||
}
|
||||
}
|
||||
for source in attachment_sources {
|
||||
if !slot.requirements.attachment_sources.contains(&source) {
|
||||
slot.requirements.attachment_sources.push(source);
|
||||
}
|
||||
}
|
||||
slot
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn inventory_slots_have_one_operation_and_explicit_requirements() {
|
||||
let slots = [
|
||||
"chat.default",
|
||||
"chat.structured",
|
||||
"prompt.text",
|
||||
"prompt.structured",
|
||||
"action.mindmap.generate",
|
||||
"action.slides.outline",
|
||||
"action.image.filter.sketch",
|
||||
"action.image.filter.clay",
|
||||
"action.image.filter.anime",
|
||||
"action.image.filter.pixel",
|
||||
"image.generate",
|
||||
"index.embedding",
|
||||
"search.rerank",
|
||||
"transcript.audio",
|
||||
];
|
||||
for id in slots {
|
||||
let slot = slot(id).expect("inventory slot must exist");
|
||||
assert!(!slot.requirements.input.is_empty());
|
||||
assert!(!slot.requirements.output.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn managed_routes_are_built_in_by_prompt_or_system_slot() {
|
||||
let prompt_slot = slot("prompt.text").unwrap();
|
||||
assert_eq!(
|
||||
managed_targets(&prompt_slot, Some("Summary as title"), CopilotManagedTier::Standard).unwrap(),
|
||||
["gpt-5.6-luna"]
|
||||
);
|
||||
assert!(
|
||||
managed_targets(
|
||||
&prompt_slot,
|
||||
Some("workflow:presentation"),
|
||||
CopilotManagedTier::Standard
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(
|
||||
managed_targets(
|
||||
&prompt_slot,
|
||||
Some("Transcript audio structured"),
|
||||
CopilotManagedTier::Premium
|
||||
)
|
||||
.unwrap(),
|
||||
["gemini-3.6-flash"]
|
||||
);
|
||||
assert_eq!(
|
||||
managed_targets(&slot("index.embedding").unwrap(), None, CopilotManagedTier::Standard).unwrap(),
|
||||
["gemini-embedding-001"]
|
||||
);
|
||||
assert_eq!(
|
||||
managed_targets(&slot("search.rerank").unwrap(), None, CopilotManagedTier::Standard).unwrap(),
|
||||
["gpt-4o-mini"]
|
||||
);
|
||||
assert!(matches!(
|
||||
quota_policy(&prompt_slot, Some("Summary as title")),
|
||||
QuotaPolicy::Internal
|
||||
));
|
||||
assert!(matches!(
|
||||
quota_policy(&prompt_slot, Some("Chat With AFFiNE AI")),
|
||||
QuotaPolicy::Metered
|
||||
));
|
||||
assert!(matches!(
|
||||
quota_policy(&slot("index.embedding").unwrap(), None),
|
||||
QuotaPolicy::System
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
#[napi_derive::napi(string_enum)]
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CopilotManagedTier {
|
||||
Standard,
|
||||
Premium,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct CopilotAccessProjection {
|
||||
pub route_allowed: bool,
|
||||
pub managed_tier: CopilotManagedTier,
|
||||
pub server_byok: bool,
|
||||
pub local_byok: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct CopilotTargetOverrideInput {
|
||||
pub profile_id: String,
|
||||
pub model_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct CopilotRouteCheckInput {
|
||||
pub slot: String,
|
||||
pub built_in_route_id: Option<String>,
|
||||
pub workspace_id: Option<String>,
|
||||
pub user_id: Option<String>,
|
||||
pub local_lease_id: Option<String>,
|
||||
pub access: CopilotAccessProjection,
|
||||
pub managed_target_id: Option<String>,
|
||||
pub target_override: Option<CopilotTargetOverrideInput>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct CopilotExecuteInput {
|
||||
pub slot: String,
|
||||
pub built_in_route_id: Option<String>,
|
||||
pub workspace_id: Option<String>,
|
||||
pub user_id: Option<String>,
|
||||
pub local_lease_id: Option<String>,
|
||||
pub access: CopilotAccessProjection,
|
||||
pub managed_target_id: Option<String>,
|
||||
pub target_override: Option<CopilotTargetOverrideInput>,
|
||||
#[napi(ts_type = "unknown")]
|
||||
pub request: serde_json::Value,
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
mod catalog;
|
||||
mod contract;
|
||||
mod policy;
|
||||
|
||||
pub(crate) use catalog::{
|
||||
CatalogSlot, QuotaPolicy, RouteOperation, managed_selected_target, managed_targets, quota_policy, slot,
|
||||
with_request_requirements,
|
||||
};
|
||||
pub use contract::{
|
||||
CopilotAccessProjection, CopilotExecuteInput, CopilotManagedTier, CopilotRouteCheckInput, CopilotTargetOverrideInput,
|
||||
};
|
||||
pub(crate) use policy::{
|
||||
AuthorizedProfileRef, AuthorizedTargetRef, CredentialRef, Deployment, ProfileSource, RouteDecision,
|
||||
RouteDecisionReason, RoutePolicyInput, TargetOverride, decide,
|
||||
};
|
||||
@@ -0,0 +1,311 @@
|
||||
use llm_adapter::capability::declared_model_matches;
|
||||
|
||||
use super::CatalogSlot;
|
||||
use crate::llm::byok::ByokProfileDefinition;
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum Deployment {
|
||||
Cloud,
|
||||
SelfHosted,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum ProfileSource {
|
||||
Server,
|
||||
Local,
|
||||
Managed,
|
||||
}
|
||||
|
||||
pub(crate) struct AuthorizedProfileRef {
|
||||
pub(crate) profile_id: String,
|
||||
pub(crate) source: ProfileSource,
|
||||
pub(crate) provider: String,
|
||||
pub(crate) definition: ByokProfileDefinition,
|
||||
pub(crate) sort_order: i32,
|
||||
pub(crate) credential_ref: CredentialRef,
|
||||
}
|
||||
|
||||
pub(crate) enum CredentialRef {
|
||||
Envelope { encrypted: String, aad: Vec<u8> },
|
||||
Managed { profile_id: String },
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub(crate) struct TargetOverride {
|
||||
pub(crate) profile_id: String,
|
||||
pub(crate) model_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum RouteDecisionReason {
|
||||
ByokDisabled,
|
||||
AccessUnavailable,
|
||||
ExplicitTargetUnavailable,
|
||||
NoCompatibleTarget,
|
||||
ManagedPresetUnavailable,
|
||||
}
|
||||
|
||||
pub(crate) enum RouteDecision {
|
||||
Ready(Vec<AuthorizedTargetRef>),
|
||||
Denied(RouteDecisionReason),
|
||||
NoRoute(RouteDecisionReason),
|
||||
}
|
||||
|
||||
pub(crate) struct AuthorizedTargetRef {
|
||||
pub(crate) profile_index: usize,
|
||||
pub(crate) model_index: usize,
|
||||
}
|
||||
|
||||
pub(crate) struct RoutePolicyInput<'a> {
|
||||
pub(crate) slot: &'a CatalogSlot,
|
||||
pub(crate) deployment: Deployment,
|
||||
pub(crate) byok_enabled: bool,
|
||||
pub(crate) access_available: bool,
|
||||
pub(crate) profiles: &'a [AuthorizedProfileRef],
|
||||
pub(crate) target_override: Option<&'a TargetOverride>,
|
||||
pub(crate) target_override_managed: bool,
|
||||
}
|
||||
|
||||
pub(crate) fn decide(input: RoutePolicyInput<'_>) -> RouteDecision {
|
||||
if input.deployment == Deployment::SelfHosted && !input.byok_enabled {
|
||||
return RouteDecision::NoRoute(RouteDecisionReason::ByokDisabled);
|
||||
}
|
||||
|
||||
if input.target_override_managed {
|
||||
if input.deployment != Deployment::Cloud {
|
||||
return RouteDecision::Denied(RouteDecisionReason::ExplicitTargetUnavailable);
|
||||
}
|
||||
if !input.access_available {
|
||||
return RouteDecision::Denied(RouteDecisionReason::AccessUnavailable);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(target) = input.target_override {
|
||||
let mut selected = compatible_targets(&input, input.target_override_managed);
|
||||
selected.retain(|candidate| {
|
||||
let profile = &input.profiles[candidate.profile_index];
|
||||
let model = &profile.definition.models[candidate.model_index];
|
||||
profile.profile_id == target.profile_id && model.model_id == target.model_id
|
||||
});
|
||||
return if selected.is_empty() {
|
||||
RouteDecision::Denied(RouteDecisionReason::ExplicitTargetUnavailable)
|
||||
} else {
|
||||
RouteDecision::Ready(selected)
|
||||
};
|
||||
}
|
||||
let byok = compatible_targets(&input, false);
|
||||
if !byok.is_empty() {
|
||||
return RouteDecision::Ready(byok);
|
||||
}
|
||||
if !input.access_available {
|
||||
return RouteDecision::Denied(RouteDecisionReason::AccessUnavailable);
|
||||
}
|
||||
if input.deployment == Deployment::Cloud {
|
||||
let managed = compatible_targets(&input, true);
|
||||
if managed.is_empty() {
|
||||
RouteDecision::NoRoute(RouteDecisionReason::ManagedPresetUnavailable)
|
||||
} else {
|
||||
RouteDecision::Ready(managed)
|
||||
}
|
||||
} else {
|
||||
RouteDecision::NoRoute(RouteDecisionReason::NoCompatibleTarget)
|
||||
}
|
||||
}
|
||||
|
||||
fn compatible_targets(input: &RoutePolicyInput<'_>, managed: bool) -> Vec<AuthorizedTargetRef> {
|
||||
let mut profiles = input
|
||||
.profiles
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, profile)| (profile.source == ProfileSource::Managed) == managed)
|
||||
.collect::<Vec<_>>();
|
||||
profiles.sort_by_key(|(_, profile)| profile.sort_order);
|
||||
profiles
|
||||
.into_iter()
|
||||
.flat_map(|(profile_index, profile)| {
|
||||
profile
|
||||
.definition
|
||||
.models
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, model)| model.enabled && declared_model_matches(&model.capabilities, &input.slot.requirements))
|
||||
.map(move |(model_index, _)| AuthorizedTargetRef {
|
||||
profile_index,
|
||||
model_index,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use llm_adapter::capability::{DeclaredModelCapability, ModelInput, ModelOutput};
|
||||
|
||||
use super::*;
|
||||
use crate::llm::{
|
||||
byok::{ByokEndpoint, ByokModelDeclaration},
|
||||
route::catalog,
|
||||
};
|
||||
|
||||
fn profile(id: &str, source: ProfileSource, model: &str, output: ModelOutput) -> AuthorizedProfileRef {
|
||||
AuthorizedProfileRef {
|
||||
profile_id: id.to_string(),
|
||||
source,
|
||||
provider: "openai".to_string(),
|
||||
definition: ByokProfileDefinition {
|
||||
version: 1,
|
||||
endpoint: ByokEndpoint::Custom {
|
||||
url: "https://example.test/v1".to_string(),
|
||||
},
|
||||
models: vec![ByokModelDeclaration {
|
||||
model_id: model.to_string(),
|
||||
enabled: true,
|
||||
capabilities: vec![DeclaredModelCapability {
|
||||
input: vec![ModelInput::Text],
|
||||
output: vec![output],
|
||||
features: vec![],
|
||||
attachment_kinds: vec![],
|
||||
attachment_sources: vec![],
|
||||
}],
|
||||
}],
|
||||
},
|
||||
sort_order: 0,
|
||||
credential_ref: CredentialRef::Managed {
|
||||
profile_id: id.to_string(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deployment_matrix_and_unsupported_only_fallback() {
|
||||
let slot = catalog::slot("chat.default").unwrap();
|
||||
let cases = [
|
||||
(
|
||||
Deployment::Cloud,
|
||||
true,
|
||||
vec![profile("managed", ProfileSource::Managed, "A", ModelOutput::Text)],
|
||||
true,
|
||||
),
|
||||
(
|
||||
Deployment::SelfHosted,
|
||||
false,
|
||||
vec![profile("managed", ProfileSource::Managed, "A", ModelOutput::Text)],
|
||||
false,
|
||||
),
|
||||
(Deployment::SelfHosted, true, vec![], false),
|
||||
(
|
||||
Deployment::Cloud,
|
||||
true,
|
||||
vec![
|
||||
profile("byok", ProfileSource::Server, "B", ModelOutput::Image),
|
||||
profile("managed", ProfileSource::Managed, "A", ModelOutput::Text),
|
||||
],
|
||||
true,
|
||||
),
|
||||
];
|
||||
for (deployment, byok_enabled, profiles, ready) in cases {
|
||||
let decision = decide(RoutePolicyInput {
|
||||
slot: &slot,
|
||||
deployment,
|
||||
byok_enabled,
|
||||
access_available: true,
|
||||
profiles: &profiles,
|
||||
target_override: None,
|
||||
target_override_managed: false,
|
||||
});
|
||||
assert_eq!(matches!(decision, RouteDecision::Ready(_)), ready);
|
||||
}
|
||||
|
||||
let profiles = vec![
|
||||
profile("byok", ProfileSource::Server, "B", ModelOutput::Text),
|
||||
profile("managed", ProfileSource::Managed, "A", ModelOutput::Text),
|
||||
];
|
||||
assert!(matches!(
|
||||
decide(RoutePolicyInput {
|
||||
slot: &slot,
|
||||
deployment: Deployment::Cloud,
|
||||
byok_enabled: true,
|
||||
access_available: false,
|
||||
profiles: &profiles,
|
||||
target_override: None,
|
||||
target_override_managed: false,
|
||||
}),
|
||||
RouteDecision::Ready(_)
|
||||
));
|
||||
assert!(matches!(
|
||||
decide(RoutePolicyInput {
|
||||
slot: &slot,
|
||||
deployment: Deployment::Cloud,
|
||||
byok_enabled: true,
|
||||
access_available: false,
|
||||
profiles: &profiles[1..],
|
||||
target_override: None,
|
||||
target_override_managed: false,
|
||||
}),
|
||||
RouteDecision::Denied(RouteDecisionReason::AccessUnavailable)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn override_is_complete_and_custom_model_is_not_replaced_by_managed_model() {
|
||||
let slot = catalog::slot("chat.default").unwrap();
|
||||
let profiles = vec![
|
||||
profile("byok", ProfileSource::Server, "vendor/model:B", ModelOutput::Text),
|
||||
profile("managed", ProfileSource::Managed, "model:A", ModelOutput::Text),
|
||||
];
|
||||
let target = TargetOverride {
|
||||
profile_id: "byok".to_string(),
|
||||
model_id: "vendor/model:B".to_string(),
|
||||
};
|
||||
let RouteDecision::Ready(candidates) = decide(RoutePolicyInput {
|
||||
slot: &slot,
|
||||
deployment: Deployment::Cloud,
|
||||
byok_enabled: true,
|
||||
access_available: true,
|
||||
profiles: &profiles,
|
||||
target_override: Some(&target),
|
||||
target_override_managed: false,
|
||||
}) else {
|
||||
panic!("override should resolve");
|
||||
};
|
||||
assert_eq!(
|
||||
profiles[candidates[0].profile_index].definition.models[candidates[0].model_index].model_id,
|
||||
"vendor/model:B"
|
||||
);
|
||||
|
||||
let managed_target = TargetOverride {
|
||||
profile_id: "managed".to_string(),
|
||||
model_id: "model:A".to_string(),
|
||||
};
|
||||
let RouteDecision::Ready(candidates) = decide(RoutePolicyInput {
|
||||
slot: &slot,
|
||||
deployment: Deployment::Cloud,
|
||||
byok_enabled: true,
|
||||
access_available: true,
|
||||
profiles: &profiles,
|
||||
target_override: Some(&managed_target),
|
||||
target_override_managed: true,
|
||||
}) else {
|
||||
panic!("managed selection should resolve");
|
||||
};
|
||||
assert!(matches!(
|
||||
profiles[candidates[0].profile_index].source,
|
||||
ProfileSource::Managed
|
||||
));
|
||||
|
||||
let mut disabled = profile("disabled", ProfileSource::Server, "model:C", ModelOutput::Text);
|
||||
disabled.definition.models[0].enabled = false;
|
||||
assert!(matches!(
|
||||
decide(RoutePolicyInput {
|
||||
slot: &slot,
|
||||
deployment: Deployment::SelfHosted,
|
||||
byok_enabled: true,
|
||||
access_available: true,
|
||||
profiles: &[disabled],
|
||||
target_override: None,
|
||||
target_override_managed: false,
|
||||
}),
|
||||
RouteDecision::NoRoute(RouteDecisionReason::NoCompatibleTarget)
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
use std::{
|
||||
net::{IpAddr, Ipv4Addr},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use super::{RuntimeError, RuntimeResult};
|
||||
use crate::{
|
||||
llm::byok::{ByokEndpoint, ByokProfileDefinition},
|
||||
runtime::config::CopilotByokRuntimeConfig,
|
||||
};
|
||||
|
||||
const DNS_RESOLUTION_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
pub(super) async fn admit_endpoint(
|
||||
definition: &ByokProfileDefinition,
|
||||
policy: &CopilotByokRuntimeConfig,
|
||||
) -> RuntimeResult<()> {
|
||||
let ByokEndpoint::Custom { url } = &definition.endpoint else {
|
||||
return Ok(());
|
||||
};
|
||||
if !policy.allow_custom_endpoint {
|
||||
return Err(RuntimeError::invalid_input("custom BYOK endpoints are disabled"));
|
||||
}
|
||||
if policy.allow_private_endpoint {
|
||||
return Ok(());
|
||||
}
|
||||
let parsed = url::Url::parse(url).map_err(|_| RuntimeError::invalid_input("invalid BYOK endpoint"))?;
|
||||
let host = parsed
|
||||
.host_str()
|
||||
.ok_or_else(|| RuntimeError::invalid_input("invalid BYOK endpoint"))?;
|
||||
if host.eq_ignore_ascii_case("localhost") {
|
||||
return Err(RuntimeError::invalid_input("private BYOK endpoints are disabled"));
|
||||
}
|
||||
let port = parsed.port_or_known_default().unwrap_or(443);
|
||||
let addresses = tokio::time::timeout(DNS_RESOLUTION_TIMEOUT, tokio::net::lookup_host((host, port)))
|
||||
.await
|
||||
.map_err(|_| RuntimeError::invalid_input("BYOK endpoint DNS resolution timed out"))?
|
||||
.map_err(|_| RuntimeError::invalid_input("BYOK endpoint DNS resolution failed"))?;
|
||||
let mut resolved = false;
|
||||
for address in addresses {
|
||||
resolved = true;
|
||||
if is_private_address(address.ip()) {
|
||||
return Err(RuntimeError::invalid_input("private BYOK endpoints are disabled"));
|
||||
}
|
||||
}
|
||||
if !resolved {
|
||||
return Err(RuntimeError::invalid_input("BYOK endpoint DNS resolution failed"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_private_address(address: IpAddr) -> bool {
|
||||
match address {
|
||||
IpAddr::V4(address) => {
|
||||
address.is_private()
|
||||
|| address.is_loopback()
|
||||
|| address.is_link_local()
|
||||
|| address.is_broadcast()
|
||||
|| address.is_documentation()
|
||||
|| address.is_unspecified()
|
||||
|| address.octets()[0] == 0
|
||||
|| Ipv4Addr::new(100, 64, 0, 0) <= address && address <= Ipv4Addr::new(100, 127, 255, 255)
|
||||
}
|
||||
IpAddr::V6(address) => {
|
||||
address.is_loopback()
|
||||
|| address.is_unspecified()
|
||||
|| address.is_unique_local()
|
||||
|| address.is_unicast_link_local()
|
||||
|| address
|
||||
.to_ipv4_mapped()
|
||||
.is_some_and(|address| is_private_address(IpAddr::V4(address)))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,198 @@
|
||||
use hmac::{Hmac, KeyInit, Mac};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::Sha256;
|
||||
use sqlx::{PgPool, Row};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::{RuntimeError, RuntimeResult, admit_endpoint, envelope_key, require_text, token_hash};
|
||||
use crate::{
|
||||
llm::{
|
||||
ByokLocalLeaseOutput, ByokProfileDefinition, CreateByokLocalLeaseInput,
|
||||
byok::{SensitiveCredential, local_aad},
|
||||
validate_definition,
|
||||
},
|
||||
runtime::config::CopilotByokRuntimeConfig,
|
||||
};
|
||||
|
||||
const LOCAL_LEASE_PURPOSE: &str = "copilot_byok_local_lease";
|
||||
const LOCAL_LEASE_ACTIVE_PURPOSE: &str = "copilot_byok_local_lease:active";
|
||||
const LOCAL_LEASE_TTL_MS: i64 = 10 * 60 * 1000;
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub(crate) struct LocalLeasePayload {
|
||||
pub(crate) version: u32,
|
||||
pub(crate) workspace_id: String,
|
||||
pub(crate) user_id: String,
|
||||
pub(crate) providers: Vec<LocalLeaseProvider>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub(crate) struct LocalLeaseProvider {
|
||||
pub(crate) provider: String,
|
||||
pub(crate) name: String,
|
||||
pub(crate) description: Option<String>,
|
||||
pub(crate) encrypted_credential: String,
|
||||
pub(crate) definition: ByokProfileDefinition,
|
||||
pub(crate) enabled: bool,
|
||||
}
|
||||
|
||||
pub(in super::super) async fn create(
|
||||
pool: &PgPool,
|
||||
root_secret: &[u8],
|
||||
policy: &CopilotByokRuntimeConfig,
|
||||
input: CreateByokLocalLeaseInput,
|
||||
) -> RuntimeResult<ByokLocalLeaseOutput> {
|
||||
require_text(&input.workspace_id, "workspaceId")?;
|
||||
require_text(&input.user_id, "userId")?;
|
||||
if input.providers.is_empty() {
|
||||
return Err(RuntimeError::invalid_input("providers is required"));
|
||||
}
|
||||
let lease_id = Uuid::new_v4().to_string();
|
||||
let key = envelope_key(root_secret)?;
|
||||
let mut fingerprint =
|
||||
Hmac::<Sha256>::new_from_slice(root_secret).map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
|
||||
fingerprint.update(input.workspace_id.as_bytes());
|
||||
fingerprint.update(&[0]);
|
||||
fingerprint.update(input.user_id.as_bytes());
|
||||
|
||||
let mut providers = Vec::with_capacity(input.providers.len());
|
||||
for (index, provider) in input.providers.into_iter().enumerate() {
|
||||
require_text(&provider.name, "name")?;
|
||||
require_text(&provider.credential, "credential")?;
|
||||
let definition = validate_definition(&provider.provider, provider.definition)
|
||||
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
|
||||
admit_endpoint(&definition, policy).await?;
|
||||
fingerprint.update(&[0]);
|
||||
fingerprint.update(provider.provider.as_bytes());
|
||||
fingerprint.update(&[0]);
|
||||
fingerprint.update(provider.credential.as_bytes());
|
||||
fingerprint.update(&[0]);
|
||||
fingerprint.update(
|
||||
&serde_json::to_vec(&definition)
|
||||
.map_err(|error| RuntimeError::json("serialize BYOK local definition failed", error))?,
|
||||
);
|
||||
let encrypted_credential = key
|
||||
.encrypt(
|
||||
&SensitiveCredential::new(provider.credential.into_bytes()),
|
||||
&local_aad(
|
||||
&input.workspace_id,
|
||||
&input.user_id,
|
||||
&lease_id,
|
||||
index,
|
||||
&provider.provider,
|
||||
definition.endpoint_identity(),
|
||||
),
|
||||
)
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
|
||||
providers.push(LocalLeaseProvider {
|
||||
provider: provider.provider,
|
||||
name: provider.name,
|
||||
description: provider.description,
|
||||
encrypted_credential,
|
||||
definition,
|
||||
enabled: provider.enabled,
|
||||
});
|
||||
}
|
||||
|
||||
let active_key = hex::encode(fingerprint.finalize().into_bytes());
|
||||
let payload = serde_json::to_value(LocalLeasePayload {
|
||||
version: 1,
|
||||
workspace_id: input.workspace_id,
|
||||
user_id: input.user_id,
|
||||
providers,
|
||||
})
|
||||
.map_err(|error| RuntimeError::json("serialize BYOK local lease failed", error))?;
|
||||
let mut tx = pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("create BYOK local lease transaction failed", error))?;
|
||||
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
|
||||
.bind(&active_key)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("lock BYOK local lease failed", error))?;
|
||||
if let Some(row) = sqlx::query(
|
||||
r#"
|
||||
SELECT payload->>'leaseId' AS lease_id
|
||||
FROM runtime_states
|
||||
WHERE purpose = $1 AND token_hash = $2
|
||||
AND consumed_at IS NULL AND expires_at > clock_timestamp()
|
||||
FOR UPDATE
|
||||
"#,
|
||||
)
|
||||
.bind(LOCAL_LEASE_ACTIVE_PURPOSE)
|
||||
.bind(token_hash(&active_key))
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("read active BYOK local lease failed", error))?
|
||||
{
|
||||
let existing_lease_id: String = row.get("lease_id");
|
||||
if let Some(expires_at_ms) = sqlx::query_scalar::<_, i64>(
|
||||
r#"
|
||||
SELECT (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT
|
||||
FROM runtime_states
|
||||
WHERE purpose = $1 AND token_hash = $2
|
||||
AND consumed_at IS NULL AND expires_at > clock_timestamp()
|
||||
FOR UPDATE
|
||||
"#,
|
||||
)
|
||||
.bind(LOCAL_LEASE_PURPOSE)
|
||||
.bind(token_hash(&existing_lease_id))
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("read active BYOK local lease payload failed", error))?
|
||||
{
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("reuse BYOK local lease commit failed", error))?;
|
||||
return Ok(ByokLocalLeaseOutput {
|
||||
lease_id: existing_lease_id,
|
||||
expires_at_ms,
|
||||
});
|
||||
}
|
||||
}
|
||||
sqlx::query("DELETE FROM runtime_states WHERE purpose = $1 AND token_hash = $2")
|
||||
.bind(LOCAL_LEASE_ACTIVE_PURPOSE)
|
||||
.bind(token_hash(&active_key))
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("delete stale BYOK local lease active record failed", error))?;
|
||||
let expires_at_ms = sqlx::query_scalar::<_, i64>(
|
||||
r#"
|
||||
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
|
||||
VALUES ($1, $2, $3, $4, clock_timestamp() + ($5 * INTERVAL '1 millisecond'))
|
||||
RETURNING (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT
|
||||
"#,
|
||||
)
|
||||
.bind(LOCAL_LEASE_PURPOSE)
|
||||
.bind(token_hash(&lease_id))
|
||||
.bind(&active_key)
|
||||
.bind(payload)
|
||||
.bind(LOCAL_LEASE_TTL_MS as f64)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("create BYOK local lease failed", error))?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
|
||||
VALUES ($1, $2, $3, jsonb_build_object('leaseId', $4::text), clock_timestamp() + ($5 * INTERVAL '1 millisecond'))
|
||||
"#,
|
||||
)
|
||||
.bind(LOCAL_LEASE_ACTIVE_PURPOSE)
|
||||
.bind(token_hash(&active_key))
|
||||
.bind(&active_key)
|
||||
.bind(&lease_id)
|
||||
.bind(LOCAL_LEASE_TTL_MS as f64)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("create BYOK local lease active record failed", error))?;
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("create BYOK local lease commit failed", error))?;
|
||||
Ok(ByokLocalLeaseOutput {
|
||||
lease_id,
|
||||
expires_at_ms,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
mod admission;
|
||||
mod local;
|
||||
mod probe;
|
||||
mod profile;
|
||||
|
||||
use admission::admit_endpoint;
|
||||
pub(super) use local::{LocalLeasePayload, create as create_local_lease};
|
||||
pub(super) use profile::{create, delete, list, probe_draft, probe_profile, reorder, replace, rotate};
|
||||
use profile::{envelope_key, require_text};
|
||||
|
||||
use super::{RuntimeError, RuntimeResult, backend_provider, byok_endpoint, executable_protocol, token_hash};
|
||||
@@ -0,0 +1,465 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use llm_adapter::{
|
||||
backend::{BackendError, DefaultHttpClient},
|
||||
capability::{
|
||||
AttachmentKind, AttachmentSource, ModelFeature, ModelInput, ModelOutput, ModelRequirements, declared_model_matches,
|
||||
},
|
||||
core::{
|
||||
CoreContent, CoreMessage, CoreRequest, CoreRole, CoreToolDefinition, EmbeddingRequest, ImageOptions,
|
||||
ImageProviderOptions, ImageRequest, RerankCandidate, RerankRequest, StructuredRequest,
|
||||
},
|
||||
router::{ExecutablePreparedRoute, ExecutableRequest, dispatch_prepared_route},
|
||||
target::{BackendCredential, BackendOperation, BackendTargetInput, EgressPolicy, compile_backend_target},
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
use super::{RuntimeError, RuntimeResult, backend_provider, byok_endpoint, executable_protocol};
|
||||
use crate::{
|
||||
llm::{
|
||||
ByokModelProbeCheckOutput, ByokModelProbeOutput, ByokProbeCheckInput, ByokProbeResultOutput, ByokProbeStatusOutput,
|
||||
byok::{ByokEndpoint, ByokProfileDefinition, SensitiveCredential, definition_fingerprint},
|
||||
},
|
||||
runtime::config::CopilotByokRuntimeConfig,
|
||||
};
|
||||
|
||||
pub(super) async fn execute_probe(
|
||||
provider: &str,
|
||||
definition: &ByokProfileDefinition,
|
||||
credential: SensitiveCredential,
|
||||
policy: &CopilotByokRuntimeConfig,
|
||||
checks: Vec<ByokProbeCheckInput>,
|
||||
) -> RuntimeResult<ByokProbeResultOutput> {
|
||||
let tested_at_ms = chrono::Utc::now().timestamp_millis();
|
||||
let connection_error = connection_probe(provider, &definition.endpoint, &credential, policy).await;
|
||||
let connection = status(tested_at_ms, connection_error.as_deref());
|
||||
let mut requested = HashSet::new();
|
||||
for check in checks {
|
||||
if !requested.insert((check.model_id.clone(), check.operation.clone())) {
|
||||
return Err(RuntimeError::invalid_input("duplicate BYOK probe check"));
|
||||
}
|
||||
if !matches!(
|
||||
check.operation.as_str(),
|
||||
"chat" | "structured" | "tools" | "vision" | "embedding" | "rerank" | "image" | "transcript"
|
||||
) {
|
||||
return Err(RuntimeError::invalid_input("unknown BYOK probe operation"));
|
||||
}
|
||||
}
|
||||
|
||||
let mut models = Vec::new();
|
||||
for model in &definition.models {
|
||||
let model_checks = requested
|
||||
.iter()
|
||||
.filter(|(model_id, _)| model_id == &model.model_id)
|
||||
.map(|(_, operation)| operation.clone())
|
||||
.collect::<Vec<_>>();
|
||||
if model_checks.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let mut outputs = Vec::with_capacity(model_checks.len());
|
||||
for operation in model_checks {
|
||||
let probe_status = if connection_error.is_some() {
|
||||
not_tested()
|
||||
} else if !model.enabled {
|
||||
failed(tested_at_ms, "model_disabled")
|
||||
} else if !declared_model_matches(&model.capabilities, &requirements(&operation)) {
|
||||
failed(tested_at_ms, "capability_not_declared")
|
||||
} else if matches!(operation.as_str(), "vision" | "transcript") {
|
||||
not_tested()
|
||||
} else {
|
||||
let provider = provider.to_string();
|
||||
let endpoint = definition.endpoint.clone();
|
||||
let model_id = model.model_id.clone();
|
||||
let credential = String::from_utf8(credential.expose().to_vec())
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
|
||||
let operation_for_task = operation.clone();
|
||||
let allow_private = policy.allow_custom_endpoint && policy.allow_private_endpoint;
|
||||
tokio::task::spawn_blocking(move || {
|
||||
dispatch_check(
|
||||
&provider,
|
||||
&endpoint,
|
||||
&model_id,
|
||||
credential,
|
||||
&operation_for_task,
|
||||
allow_private,
|
||||
)
|
||||
})
|
||||
.await
|
||||
.map_err(|error| RuntimeError::invalid_state(format!("BYOK model probe task failed: {error}")))?
|
||||
};
|
||||
outputs.push(ByokModelProbeCheckOutput {
|
||||
operation,
|
||||
status: probe_status,
|
||||
});
|
||||
}
|
||||
models.push(ByokModelProbeOutput {
|
||||
model_id: model.model_id.clone(),
|
||||
checks: outputs,
|
||||
});
|
||||
}
|
||||
if requested
|
||||
.iter()
|
||||
.any(|(model_id, _)| !definition.models.iter().any(|model| &model.model_id == model_id))
|
||||
{
|
||||
return Err(RuntimeError::invalid_input("BYOK probe model not found"));
|
||||
}
|
||||
|
||||
Ok(ByokProbeResultOutput {
|
||||
definition_fingerprint: definition_fingerprint(definition),
|
||||
stale: false,
|
||||
connection,
|
||||
models,
|
||||
})
|
||||
}
|
||||
|
||||
async fn connection_probe(
|
||||
provider: &str,
|
||||
endpoint: &ByokEndpoint,
|
||||
credential: &SensitiveCredential,
|
||||
policy: &CopilotByokRuntimeConfig,
|
||||
) -> Option<String> {
|
||||
let credential = match std::str::from_utf8(credential.expose()) {
|
||||
Ok(value) => value.to_string(),
|
||||
Err(_) => return Some("credential_unavailable".to_string()),
|
||||
};
|
||||
let (url, headers) = probe_request(provider, endpoint, credential);
|
||||
let allow_private = policy.allow_custom_endpoint && policy.allow_private_endpoint;
|
||||
let result = tokio::task::spawn_blocking(move || {
|
||||
safefetch::safe_fetch(&safefetch::SafeFetchRequest {
|
||||
url,
|
||||
method: Some(safefetch::SafeFetchMethod::Get),
|
||||
headers: Some(headers.clone()),
|
||||
body: None,
|
||||
timeout_ms: Some(10_000),
|
||||
max_redirects: Some(0),
|
||||
max_bytes: Some(1024 * 1024),
|
||||
allowed_headers: Some(headers.keys().cloned().collect()),
|
||||
allowed_hosts: None,
|
||||
allow_http: Some(allow_private),
|
||||
allow_private_target_origin: Some(allow_private),
|
||||
ech_config_list: None,
|
||||
})
|
||||
})
|
||||
.await;
|
||||
match result {
|
||||
Ok(Ok(response))
|
||||
if (200..300).contains(&response.status) && valid_connection_response(provider, &response.body) =>
|
||||
{
|
||||
None
|
||||
}
|
||||
Ok(Ok(response)) => Some(http_error_kind(response.status).to_string()),
|
||||
_ => Some("transport".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_check(
|
||||
provider: &str,
|
||||
endpoint: &ByokEndpoint,
|
||||
model_id: &str,
|
||||
credential: String,
|
||||
operation: &str,
|
||||
allow_private: bool,
|
||||
) -> ByokProbeStatusOutput {
|
||||
let checked_at = chrono::Utc::now().timestamp_millis();
|
||||
let operation_kind = match operation {
|
||||
"chat" | "tools" => BackendOperation::Chat,
|
||||
"structured" => BackendOperation::Structured,
|
||||
"embedding" => BackendOperation::Embedding,
|
||||
"rerank" => BackendOperation::Rerank,
|
||||
"image" => BackendOperation::Image,
|
||||
_ => return not_tested(),
|
||||
};
|
||||
let target = compile_backend_target(BackendTargetInput {
|
||||
provider: match backend_provider(provider) {
|
||||
Ok(provider) => provider,
|
||||
Err(_) => return failed(checked_at, "unsupported_provider"),
|
||||
},
|
||||
operation: operation_kind,
|
||||
endpoint: byok_endpoint(provider, endpoint),
|
||||
model: model_id.to_string(),
|
||||
credential: BackendCredential::new(credential),
|
||||
timeout_ms: Some(15_000),
|
||||
egress_policy: if allow_private {
|
||||
EgressPolicy::AllowPrivate
|
||||
} else {
|
||||
EgressPolicy::PublicOnly
|
||||
},
|
||||
});
|
||||
let target = match target {
|
||||
Ok(target) => target,
|
||||
Err(_) => return failed(checked_at, "unsupported_operation"),
|
||||
};
|
||||
let route = ExecutablePreparedRoute::new(
|
||||
executable_protocol(target.protocol),
|
||||
target.model,
|
||||
target.config,
|
||||
probe_request_for_operation(operation),
|
||||
);
|
||||
let route = match route {
|
||||
Ok(route) => route,
|
||||
Err(_) => return failed(checked_at, "invalid_probe_request"),
|
||||
};
|
||||
match dispatch_prepared_route(&DefaultHttpClient::default(), &route) {
|
||||
Ok(_) => verified(checked_at),
|
||||
Err(error) => failed(checked_at, backend_error_kind(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
fn probe_request_for_operation(operation: &str) -> ExecutableRequest {
|
||||
let message = CoreMessage {
|
||||
role: CoreRole::User,
|
||||
content: vec![CoreContent::Text {
|
||||
text: "Reply with OK.".to_string(),
|
||||
}],
|
||||
};
|
||||
match operation {
|
||||
"chat" | "tools" => ExecutableRequest::Chat(CoreRequest {
|
||||
model: String::new(),
|
||||
messages: vec![message],
|
||||
stream: false,
|
||||
max_tokens: Some(8),
|
||||
temperature: Some(0.0),
|
||||
tools: if operation == "tools" {
|
||||
vec![CoreToolDefinition {
|
||||
name: "byok_probe".to_string(),
|
||||
description: Some("Probe tool compatibility".to_string()),
|
||||
parameters: json!({ "type": "object", "properties": {} }),
|
||||
}]
|
||||
} else {
|
||||
vec![]
|
||||
},
|
||||
tool_choice: None,
|
||||
include: None,
|
||||
reasoning: None,
|
||||
response_schema: None,
|
||||
}),
|
||||
"structured" => ExecutableRequest::Structured(StructuredRequest {
|
||||
model: String::new(),
|
||||
messages: vec![message],
|
||||
schema: json!({
|
||||
"type": "object",
|
||||
"properties": { "ok": { "type": "boolean" } },
|
||||
"required": ["ok"],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
max_tokens: Some(16),
|
||||
temperature: Some(0.0),
|
||||
reasoning: None,
|
||||
strict: Some(true),
|
||||
response_mime_type: Some("application/json".to_string()),
|
||||
}),
|
||||
"embedding" => ExecutableRequest::Embedding(EmbeddingRequest {
|
||||
model: String::new(),
|
||||
inputs: vec!["BYOK probe".to_string()],
|
||||
dimensions: None,
|
||||
task_type: None,
|
||||
}),
|
||||
"rerank" => ExecutableRequest::Rerank(RerankRequest {
|
||||
model: String::new(),
|
||||
query: "probe".to_string(),
|
||||
candidates: vec![RerankCandidate {
|
||||
id: None,
|
||||
text: "probe".to_string(),
|
||||
}],
|
||||
top_n: Some(1),
|
||||
}),
|
||||
"image" => ExecutableRequest::Image(Box::new(ImageRequest::generate(
|
||||
String::new(),
|
||||
"A single black pixel".to_string(),
|
||||
ImageOptions::default(),
|
||||
ImageProviderOptions::default(),
|
||||
))),
|
||||
_ => unreachable!("validated probe operation"),
|
||||
}
|
||||
}
|
||||
|
||||
fn requirements(operation: &str) -> ModelRequirements {
|
||||
let (input, output, features, attachment_kinds, attachment_sources) = match operation {
|
||||
"chat" => (vec![ModelInput::Text], vec![ModelOutput::Text], vec![], vec![], vec![]),
|
||||
"structured" => (
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Structured],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"tools" => (
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Text],
|
||||
vec![ModelFeature::ToolCalling],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"vision" => (
|
||||
vec![ModelInput::Text, ModelInput::Image],
|
||||
vec![ModelOutput::Text],
|
||||
vec![],
|
||||
vec![AttachmentKind::Image],
|
||||
vec![AttachmentSource::Data],
|
||||
),
|
||||
"embedding" => (
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Embedding],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"rerank" => (
|
||||
vec![ModelInput::Text],
|
||||
vec![ModelOutput::Rerank],
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
),
|
||||
"image" => (vec![ModelInput::Text], vec![ModelOutput::Image], vec![], vec![], vec![]),
|
||||
"transcript" => (
|
||||
vec![ModelInput::Audio],
|
||||
vec![ModelOutput::Structured],
|
||||
vec![],
|
||||
vec![AttachmentKind::Audio],
|
||||
vec![AttachmentSource::Data],
|
||||
),
|
||||
_ => unreachable!("validated probe operation"),
|
||||
};
|
||||
ModelRequirements {
|
||||
input,
|
||||
output,
|
||||
features,
|
||||
attachment_kinds,
|
||||
attachment_sources,
|
||||
}
|
||||
}
|
||||
|
||||
fn probe_request(provider: &str, endpoint: &ByokEndpoint, credential: String) -> (String, HashMap<String, String>) {
|
||||
let base = match endpoint {
|
||||
ByokEndpoint::Custom { url } => url.as_str(),
|
||||
ByokEndpoint::ProviderDefault => match provider {
|
||||
"openai" => "https://api.openai.com/v1",
|
||||
"anthropic" => "https://api.anthropic.com/v1",
|
||||
"gemini" => "https://generativelanguage.googleapis.com/v1beta",
|
||||
"fal" => "https://api.fal.ai/v1",
|
||||
_ => unreachable!("validated provider"),
|
||||
},
|
||||
};
|
||||
let mut headers = HashMap::new();
|
||||
match provider {
|
||||
"openai" => {
|
||||
headers.insert("authorization".to_string(), format!("Bearer {credential}"));
|
||||
}
|
||||
"anthropic" => {
|
||||
headers.insert("x-api-key".to_string(), credential);
|
||||
headers.insert("anthropic-version".to_string(), "2023-06-01".to_string());
|
||||
}
|
||||
"gemini" => {
|
||||
headers.insert("x-goog-api-key".to_string(), credential);
|
||||
}
|
||||
"fal" => {
|
||||
headers.insert("authorization".to_string(), format!("Key {credential}"));
|
||||
}
|
||||
_ => unreachable!("validated provider"),
|
||||
}
|
||||
let suffix = if provider == "fal" { "models?limit=10" } else { "models" };
|
||||
(format!("{}/{suffix}", base.trim_end_matches('/')), headers)
|
||||
}
|
||||
|
||||
fn valid_connection_response(provider: &str, body: &[u8]) -> bool {
|
||||
let Ok(body) = serde_json::from_slice::<serde_json::Value>(body) else {
|
||||
return false;
|
||||
};
|
||||
match provider {
|
||||
"openai" | "anthropic" => body.get("data").is_some_and(serde_json::Value::is_array),
|
||||
"gemini" => body.get("models").is_some_and(serde_json::Value::is_array),
|
||||
"fal" => body.get("error").is_none(),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn status(tested_at_ms: i64, error: Option<&str>) -> ByokProbeStatusOutput {
|
||||
match error {
|
||||
Some(error) => failed(tested_at_ms, error),
|
||||
None => verified(tested_at_ms),
|
||||
}
|
||||
}
|
||||
|
||||
fn verified(tested_at_ms: i64) -> ByokProbeStatusOutput {
|
||||
ByokProbeStatusOutput {
|
||||
kind: "verified".to_string(),
|
||||
tested_at_ms: Some(tested_at_ms),
|
||||
error_kind: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn failed(tested_at_ms: i64, error: &str) -> ByokProbeStatusOutput {
|
||||
ByokProbeStatusOutput {
|
||||
kind: "failed".to_string(),
|
||||
tested_at_ms: Some(tested_at_ms),
|
||||
error_kind: Some(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn not_tested() -> ByokProbeStatusOutput {
|
||||
ByokProbeStatusOutput {
|
||||
kind: "not_tested".to_string(),
|
||||
tested_at_ms: None,
|
||||
error_kind: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn http_error_kind(status: u16) -> &'static str {
|
||||
match status {
|
||||
401 => "authentication",
|
||||
403 => "permission",
|
||||
404 => "not_found",
|
||||
429 => "rate_limited",
|
||||
500..=599 => "unavailable",
|
||||
_ => "rejected",
|
||||
}
|
||||
}
|
||||
|
||||
fn backend_error_kind(error: &BackendError) -> &'static str {
|
||||
match error {
|
||||
BackendError::UpstreamStatus { status, .. } => http_error_kind(*status),
|
||||
BackendError::Transport { .. } => "transport",
|
||||
BackendError::Timeout { .. } => "timeout",
|
||||
BackendError::InvalidConfig { .. } | BackendError::InvalidRequest { .. } => "unsupported_operation",
|
||||
BackendError::InvalidResponse { .. }
|
||||
| BackendError::InvalidStructuredOutput { .. }
|
||||
| BackendError::Json(_)
|
||||
| BackendError::Stream(_) => "invalid_response",
|
||||
BackendError::NoBackendAvailable => "unavailable",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use llm_adapter::target::BackendEndpoint;
|
||||
|
||||
#[test]
|
||||
fn connection_probe_errors_are_low_information() {
|
||||
let (url, headers) = probe_request("openai", &ByokEndpoint::ProviderDefault, "secret".to_string());
|
||||
assert_eq!(url, "https://api.openai.com/v1/models");
|
||||
assert_eq!(headers.get("authorization").map(String::as_str), Some("Bearer secret"));
|
||||
assert_eq!(http_error_kind(401), "authentication");
|
||||
assert_eq!(http_error_kind(403), "permission");
|
||||
assert_eq!(http_error_kind(429), "rate_limited");
|
||||
assert_eq!(http_error_kind(503), "unavailable");
|
||||
|
||||
let custom = ByokEndpoint::Custom {
|
||||
url: "http://127.0.0.1:1234/v1".to_string(),
|
||||
};
|
||||
assert_eq!(
|
||||
probe_request("openai", &custom, "secret".to_string()).0,
|
||||
"http://127.0.0.1:1234/v1/models"
|
||||
);
|
||||
assert_eq!(
|
||||
byok_endpoint("openai", &custom),
|
||||
BackendEndpoint::Custom("http://127.0.0.1:1234".to_string())
|
||||
);
|
||||
assert!(valid_connection_response("openai", br#"{"data":[]}"#));
|
||||
assert!(!valid_connection_response(
|
||||
"openai",
|
||||
br#"{"error":"Unexpected endpoint"}"#
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,577 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use sqlx::{FromRow, PgPool};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::{RuntimeError, RuntimeResult, admit_endpoint};
|
||||
use crate::{
|
||||
llm::{
|
||||
ByokProfileDefinition, ByokProfileOutput, ByokValidationOutput, CreateByokProfileInput, ProbeByokDraftInput,
|
||||
ProbeByokProfileInput, ReorderByokProfilesInput, ReplaceByokProfileInput, RotateByokCredentialInput,
|
||||
byok::{CredentialEnvelopeKey, SensitiveCredential, reconcile_validation, server_aad},
|
||||
validate_definition,
|
||||
},
|
||||
runtime::config::CopilotByokRuntimeConfig,
|
||||
};
|
||||
|
||||
#[derive(FromRow)]
|
||||
struct ProfileRow {
|
||||
id: String,
|
||||
workspace_id: String,
|
||||
provider: String,
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
encrypted_api_key: String,
|
||||
definition: serde_json::Value,
|
||||
sort_order: i32,
|
||||
enabled: bool,
|
||||
revision: i32,
|
||||
credential_generation: i32,
|
||||
validation: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(FromRow)]
|
||||
struct ProfileAdmissionRow {
|
||||
provider: String,
|
||||
revision: i32,
|
||||
}
|
||||
|
||||
pub(in super::super) async fn list(pool: &PgPool, workspace_id: &str) -> RuntimeResult<Vec<ByokProfileOutput>> {
|
||||
let rows = sqlx::query_as::<_, ProfileRow>(
|
||||
r#"
|
||||
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
|
||||
definition, sort_order, enabled, revision, credential_generation, validation
|
||||
FROM ai_workspace_byok_configs
|
||||
WHERE workspace_id = $1
|
||||
ORDER BY sort_order ASC, created_at ASC
|
||||
"#,
|
||||
)
|
||||
.bind(workspace_id)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("list BYOK profiles failed", error))?;
|
||||
rows.into_iter().map(profile_output).collect()
|
||||
}
|
||||
|
||||
pub(in super::super) async fn create(
|
||||
pool: &PgPool,
|
||||
root_secret: &[u8],
|
||||
policy: &CopilotByokRuntimeConfig,
|
||||
input: CreateByokProfileInput,
|
||||
) -> RuntimeResult<ByokProfileOutput> {
|
||||
require_text(&input.workspace_id, "workspaceId")?;
|
||||
require_text(&input.name, "name")?;
|
||||
require_text(&input.credential, "credential")?;
|
||||
require_text(&input.actor_user_id, "actorUserId")?;
|
||||
let definition = validate_definition(&input.provider, input.definition)
|
||||
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
|
||||
admit_endpoint(&definition, policy).await?;
|
||||
let key = envelope_key(root_secret)?;
|
||||
let profile_id = Uuid::new_v4().to_string();
|
||||
let aad = server_aad(
|
||||
&input.workspace_id,
|
||||
&profile_id,
|
||||
&input.provider,
|
||||
definition.endpoint_identity(),
|
||||
);
|
||||
let encrypted = key
|
||||
.encrypt(&SensitiveCredential::new(input.credential.into_bytes()), &aad)
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
|
||||
let definition_json =
|
||||
serde_json::to_value(&definition).map_err(|error| RuntimeError::json("serialize BYOK definition failed", error))?;
|
||||
|
||||
let mut tx = pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("create BYOK profile transaction failed", error))?;
|
||||
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
|
||||
.bind(&input.workspace_id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("lock BYOK profile order failed", error))?;
|
||||
let sort_order = sqlx::query_scalar::<_, i32>(
|
||||
"SELECT COALESCE(MAX(sort_order) + 1, 0)::int FROM ai_workspace_byok_configs WHERE workspace_id = $1",
|
||||
)
|
||||
.bind(&input.workspace_id)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("resolve BYOK profile order failed", error))?;
|
||||
let row = sqlx::query_as::<_, ProfileRow>(
|
||||
r#"
|
||||
INSERT INTO ai_workspace_byok_configs (
|
||||
id, workspace_id, provider, name, description, encrypted_api_key,
|
||||
definition, sort_order, enabled, created_by, updated_by, created_at, updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $10, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
|
||||
RETURNING id, workspace_id, provider, name, description, encrypted_api_key,
|
||||
definition, sort_order, enabled, revision, credential_generation, validation
|
||||
"#,
|
||||
)
|
||||
.bind(&profile_id)
|
||||
.bind(&input.workspace_id)
|
||||
.bind(&input.provider)
|
||||
.bind(&input.name)
|
||||
.bind(&input.description)
|
||||
.bind(encrypted)
|
||||
.bind(definition_json)
|
||||
.bind(sort_order)
|
||||
.bind(input.enabled)
|
||||
.bind(&input.actor_user_id)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("create BYOK profile failed", error))?;
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("create BYOK profile commit failed", error))?;
|
||||
profile_output(row)
|
||||
}
|
||||
|
||||
pub(in super::super) async fn replace(
|
||||
pool: &PgPool,
|
||||
root_secret: &[u8],
|
||||
policy: &CopilotByokRuntimeConfig,
|
||||
input: ReplaceByokProfileInput,
|
||||
) -> RuntimeResult<ByokProfileOutput> {
|
||||
require_text(&input.workspace_id, "workspaceId")?;
|
||||
require_text(&input.name, "name")?;
|
||||
require_text(&input.actor_user_id, "actorUserId")?;
|
||||
if let Some(credential) = input.credential.as_deref() {
|
||||
require_text(credential, "credential")?;
|
||||
}
|
||||
if input.expected_revision < 1 {
|
||||
return Err(RuntimeError::invalid_input("expectedRevision is required"));
|
||||
}
|
||||
let admission = select_profile_for_admission(pool, &input.workspace_id, &input.profile_id).await?;
|
||||
if admission.revision != input.expected_revision {
|
||||
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
|
||||
}
|
||||
let definition = validate_definition(&admission.provider, input.definition)
|
||||
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
|
||||
admit_endpoint(&definition, policy).await?;
|
||||
|
||||
let mut tx = pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("replace BYOK profile transaction failed", error))?;
|
||||
let old = select_profile_for_update(&mut tx, &input.workspace_id, &input.profile_id).await?;
|
||||
if old.revision != input.expected_revision {
|
||||
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
|
||||
}
|
||||
let old_definition = parse_definition(old.definition.clone())?;
|
||||
let key = envelope_key(root_secret)?;
|
||||
let credential_changed = input.credential.is_some();
|
||||
let credential = if let Some(credential) = input.credential {
|
||||
SensitiveCredential::new(credential.into_bytes())
|
||||
} else {
|
||||
key
|
||||
.decrypt(
|
||||
&old.encrypted_api_key,
|
||||
&server_aad(
|
||||
&old.workspace_id,
|
||||
&old.id,
|
||||
&old.provider,
|
||||
old_definition.endpoint_identity(),
|
||||
),
|
||||
)
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?
|
||||
};
|
||||
let encrypted = key
|
||||
.encrypt(
|
||||
&credential,
|
||||
&server_aad(
|
||||
&old.workspace_id,
|
||||
&old.id,
|
||||
&old.provider,
|
||||
definition.endpoint_identity(),
|
||||
),
|
||||
)
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
|
||||
let definition_json =
|
||||
serde_json::to_value(&definition).map_err(|error| RuntimeError::json("serialize BYOK definition failed", error))?;
|
||||
let credential_generation = old.credential_generation + i32::from(credential_changed);
|
||||
let validation = reconcile_validation(
|
||||
parse_validation(old.validation)?,
|
||||
&old_definition,
|
||||
&definition,
|
||||
credential_generation,
|
||||
credential_changed,
|
||||
)
|
||||
.map(serde_json::to_value)
|
||||
.transpose()
|
||||
.map_err(|error| RuntimeError::json("serialize BYOK validation failed", error))?;
|
||||
let row = sqlx::query_as::<_, ProfileRow>(
|
||||
r#"
|
||||
UPDATE ai_workspace_byok_configs
|
||||
SET name = $3, description = $4, definition = $5, encrypted_api_key = $6,
|
||||
enabled = $7, credential_generation = $8, validation = $9,
|
||||
revision = revision + 1, updated_by = $10, updated_at = CURRENT_TIMESTAMP
|
||||
WHERE workspace_id = $1 AND id = $2 AND revision = $11
|
||||
RETURNING id, workspace_id, provider, name, description, encrypted_api_key,
|
||||
definition, sort_order, enabled, revision, credential_generation, validation
|
||||
"#,
|
||||
)
|
||||
.bind(&input.workspace_id)
|
||||
.bind(&input.profile_id)
|
||||
.bind(&input.name)
|
||||
.bind(&input.description)
|
||||
.bind(definition_json)
|
||||
.bind(encrypted)
|
||||
.bind(input.enabled)
|
||||
.bind(credential_generation)
|
||||
.bind(validation)
|
||||
.bind(&input.actor_user_id)
|
||||
.bind(input.expected_revision)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("replace BYOK profile failed", error))?
|
||||
.ok_or_else(|| RuntimeError::invalid_input("byok_revision_conflict"))?;
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("replace BYOK profile commit failed", error))?;
|
||||
profile_output(row)
|
||||
}
|
||||
|
||||
async fn select_profile_for_admission(
|
||||
pool: &PgPool,
|
||||
workspace_id: &str,
|
||||
profile_id: &str,
|
||||
) -> RuntimeResult<ProfileAdmissionRow> {
|
||||
sqlx::query_as::<_, ProfileAdmissionRow>(
|
||||
"SELECT provider, revision FROM ai_workspace_byok_configs WHERE workspace_id = $1 AND id = $2",
|
||||
)
|
||||
.bind(workspace_id)
|
||||
.bind(profile_id)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("load BYOK profile admission data failed", error))?
|
||||
.ok_or_else(|| RuntimeError::invalid_input("BYOK profile not found"))
|
||||
}
|
||||
|
||||
pub(in super::super) async fn rotate(
|
||||
pool: &PgPool,
|
||||
root_secret: &[u8],
|
||||
input: RotateByokCredentialInput,
|
||||
) -> RuntimeResult<ByokProfileOutput> {
|
||||
require_text(&input.credential, "credential")?;
|
||||
require_text(&input.actor_user_id, "actorUserId")?;
|
||||
if input.expected_revision < 1 {
|
||||
return Err(RuntimeError::invalid_input("expectedRevision is required"));
|
||||
}
|
||||
let mut tx = pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("rotate BYOK credential transaction failed", error))?;
|
||||
let profile = select_profile_for_update(&mut tx, &input.workspace_id, &input.profile_id).await?;
|
||||
if profile.revision != input.expected_revision {
|
||||
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
|
||||
}
|
||||
let definition = parse_definition(profile.definition.clone())?;
|
||||
let encrypted = envelope_key(root_secret)?
|
||||
.encrypt(
|
||||
&SensitiveCredential::new(input.credential.into_bytes()),
|
||||
&server_aad(
|
||||
&profile.workspace_id,
|
||||
&profile.id,
|
||||
&profile.provider,
|
||||
definition.endpoint_identity(),
|
||||
),
|
||||
)
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
|
||||
let row = sqlx::query_as::<_, ProfileRow>(
|
||||
r#"
|
||||
UPDATE ai_workspace_byok_configs
|
||||
SET encrypted_api_key = $3, credential_generation = credential_generation + 1,
|
||||
validation = NULL, revision = revision + 1, updated_by = $4,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
WHERE workspace_id = $1 AND id = $2 AND revision = $5
|
||||
RETURNING id, workspace_id, provider, name, description, encrypted_api_key,
|
||||
definition, sort_order, enabled, revision, credential_generation, validation
|
||||
"#,
|
||||
)
|
||||
.bind(&input.workspace_id)
|
||||
.bind(&input.profile_id)
|
||||
.bind(encrypted)
|
||||
.bind(&input.actor_user_id)
|
||||
.bind(input.expected_revision)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("rotate BYOK credential failed", error))?
|
||||
.ok_or_else(|| RuntimeError::invalid_input("byok_revision_conflict"))?;
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("rotate BYOK credential commit failed", error))?;
|
||||
profile_output(row)
|
||||
}
|
||||
|
||||
pub(in super::super) async fn delete(pool: &PgPool, workspace_id: &str, profile_id: &str) -> RuntimeResult<bool> {
|
||||
let affected = sqlx::query("DELETE FROM ai_workspace_byok_configs WHERE workspace_id = $1 AND id = $2")
|
||||
.bind(workspace_id)
|
||||
.bind(profile_id)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("delete BYOK profile failed", error))?
|
||||
.rows_affected();
|
||||
Ok(affected == 1)
|
||||
}
|
||||
|
||||
pub(in super::super) async fn reorder(
|
||||
pool: &PgPool,
|
||||
input: ReorderByokProfilesInput,
|
||||
) -> RuntimeResult<Vec<ByokProfileOutput>> {
|
||||
require_text(&input.workspace_id, "workspaceId")?;
|
||||
require_text(&input.actor_user_id, "actorUserId")?;
|
||||
let unique = input
|
||||
.profiles
|
||||
.iter()
|
||||
.map(|profile| profile.profile_id.as_str())
|
||||
.collect::<HashSet<_>>();
|
||||
if unique.len() != input.profiles.len() {
|
||||
return Err(RuntimeError::invalid_input("duplicate BYOK profile id"));
|
||||
}
|
||||
let mut tx = pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("reorder BYOK profiles transaction failed", error))?;
|
||||
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
|
||||
.bind(&input.workspace_id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("lock BYOK profile order failed", error))?;
|
||||
let current = sqlx::query_as::<_, (String, i32)>(
|
||||
"SELECT id, revision FROM ai_workspace_byok_configs WHERE workspace_id = $1 ORDER BY sort_order, created_at",
|
||||
)
|
||||
.bind(&input.workspace_id)
|
||||
.fetch_all(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("read BYOK profile order failed", error))?;
|
||||
let current_set = current
|
||||
.iter()
|
||||
.map(|(profile_id, _)| profile_id.as_str())
|
||||
.collect::<HashSet<_>>();
|
||||
if current_set != unique {
|
||||
return Err(RuntimeError::invalid_input(
|
||||
"BYOK profile order must contain every server profile",
|
||||
));
|
||||
}
|
||||
let revisions = current.into_iter().collect::<HashMap<_, _>>();
|
||||
if input
|
||||
.profiles
|
||||
.iter()
|
||||
.any(|profile| revisions.get(&profile.profile_id).copied() != Some(profile.expected_revision))
|
||||
{
|
||||
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
|
||||
}
|
||||
for (sort_order, profile) in input.profiles.iter().enumerate() {
|
||||
let updated = sqlx::query(
|
||||
"UPDATE ai_workspace_byok_configs SET sort_order = $3, revision = revision + 1, updated_by = $4, updated_at = \
|
||||
CURRENT_TIMESTAMP WHERE workspace_id = $1 AND id = $2 AND revision = $5",
|
||||
)
|
||||
.bind(&input.workspace_id)
|
||||
.bind(&profile.profile_id)
|
||||
.bind(sort_order as i32)
|
||||
.bind(&input.actor_user_id)
|
||||
.bind(profile.expected_revision)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("update BYOK profile order failed", error))?
|
||||
.rows_affected();
|
||||
if updated != 1 {
|
||||
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
|
||||
}
|
||||
}
|
||||
let rows = sqlx::query_as::<_, ProfileRow>(
|
||||
r#"
|
||||
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
|
||||
definition, sort_order, enabled, revision, credential_generation, validation
|
||||
FROM ai_workspace_byok_configs
|
||||
WHERE workspace_id = $1
|
||||
ORDER BY sort_order, created_at
|
||||
"#,
|
||||
)
|
||||
.bind(&input.workspace_id)
|
||||
.fetch_all(&mut *tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("read reordered BYOK profiles failed", error))?;
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("reorder BYOK profiles commit failed", error))?;
|
||||
rows.into_iter().map(profile_output).collect()
|
||||
}
|
||||
|
||||
pub(in super::super) async fn probe_profile(
|
||||
pool: &PgPool,
|
||||
root_secret: &[u8],
|
||||
policy: &CopilotByokRuntimeConfig,
|
||||
input: ProbeByokProfileInput,
|
||||
) -> RuntimeResult<crate::llm::ByokProbeResultOutput> {
|
||||
let profile = sqlx::query_as::<_, ProfileRow>(
|
||||
r#"
|
||||
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
|
||||
definition, sort_order, enabled, revision, credential_generation, validation
|
||||
FROM ai_workspace_byok_configs
|
||||
WHERE workspace_id = $1 AND id = $2
|
||||
"#,
|
||||
)
|
||||
.bind(&input.workspace_id)
|
||||
.bind(&input.profile_id)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("read BYOK profile for probe failed", error))?
|
||||
.ok_or_else(|| RuntimeError::invalid_input("BYOK profile not found"))?;
|
||||
let definition = parse_definition(profile.definition.clone())?;
|
||||
admit_endpoint(&definition, policy).await?;
|
||||
let credential = envelope_key(root_secret)?
|
||||
.decrypt(
|
||||
&profile.encrypted_api_key,
|
||||
&server_aad(
|
||||
&profile.workspace_id,
|
||||
&profile.id,
|
||||
&profile.provider,
|
||||
definition.endpoint_identity(),
|
||||
),
|
||||
)
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
|
||||
let mut result =
|
||||
super::probe::execute_probe(&profile.provider, &definition, credential, policy, input.checks).await?;
|
||||
let validation = ByokValidationOutput {
|
||||
definition_fingerprint: result.definition_fingerprint.clone(),
|
||||
credential_generation: profile.credential_generation,
|
||||
connection: result.connection.clone(),
|
||||
models: result.models.clone(),
|
||||
};
|
||||
let validation =
|
||||
serde_json::to_value(validation).map_err(|error| RuntimeError::json("serialize BYOK validation failed", error))?;
|
||||
let updated = sqlx::query(
|
||||
"UPDATE ai_workspace_byok_configs SET validation = $3 WHERE workspace_id = $1 AND id = $2 AND revision = $4 AND \
|
||||
credential_generation = $5",
|
||||
)
|
||||
.bind(&input.workspace_id)
|
||||
.bind(&input.profile_id)
|
||||
.bind(validation)
|
||||
.bind(profile.revision)
|
||||
.bind(profile.credential_generation)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("record BYOK profile probe failed", error))?
|
||||
.rows_affected();
|
||||
result.stale = updated != 1;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub(in super::super) async fn probe_draft(
|
||||
pool: &PgPool,
|
||||
root_secret: &[u8],
|
||||
policy: &CopilotByokRuntimeConfig,
|
||||
input: ProbeByokDraftInput,
|
||||
) -> RuntimeResult<crate::llm::ByokProbeResultOutput> {
|
||||
let definition = validate_definition(&input.provider, input.definition)
|
||||
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
|
||||
admit_endpoint(&definition, policy).await?;
|
||||
let credential = match (input.credential, input.profile_id, input.expected_revision) {
|
||||
(Some(credential), None, None) => {
|
||||
require_text(&credential, "credential")?;
|
||||
SensitiveCredential::new(credential.into_bytes())
|
||||
}
|
||||
(None, Some(profile_id), Some(expected_revision)) => {
|
||||
let profile = sqlx::query_as::<_, ProfileRow>(
|
||||
r#"
|
||||
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
|
||||
definition, sort_order, enabled, revision, credential_generation, validation
|
||||
FROM ai_workspace_byok_configs
|
||||
WHERE workspace_id = $1 AND id = $2
|
||||
"#,
|
||||
)
|
||||
.bind(&input.workspace_id)
|
||||
.bind(&profile_id)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("read BYOK profile for draft probe failed", error))?
|
||||
.ok_or_else(|| RuntimeError::invalid_input("BYOK profile not found"))?;
|
||||
if profile.provider != input.provider || profile.revision != expected_revision {
|
||||
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
|
||||
}
|
||||
let stored_definition = parse_definition(profile.definition)?;
|
||||
envelope_key(root_secret)?
|
||||
.decrypt(
|
||||
&profile.encrypted_api_key,
|
||||
&server_aad(
|
||||
&profile.workspace_id,
|
||||
&profile.id,
|
||||
&profile.provider,
|
||||
stored_definition.endpoint_identity(),
|
||||
),
|
||||
)
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?
|
||||
}
|
||||
_ => {
|
||||
return Err(RuntimeError::invalid_input(
|
||||
"draft probe requires either credential or stored profile revision",
|
||||
));
|
||||
}
|
||||
};
|
||||
super::probe::execute_probe(&input.provider, &definition, credential, policy, input.checks).await
|
||||
}
|
||||
|
||||
async fn select_profile_for_update(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
workspace_id: &str,
|
||||
profile_id: &str,
|
||||
) -> RuntimeResult<ProfileRow> {
|
||||
sqlx::query_as::<_, ProfileRow>(
|
||||
r#"
|
||||
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
|
||||
definition, sort_order, enabled, revision, credential_generation, validation
|
||||
FROM ai_workspace_byok_configs
|
||||
WHERE workspace_id = $1 AND id = $2
|
||||
FOR UPDATE
|
||||
"#,
|
||||
)
|
||||
.bind(workspace_id)
|
||||
.bind(profile_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("read BYOK profile failed", error))?
|
||||
.ok_or_else(|| RuntimeError::invalid_input("BYOK profile not found"))
|
||||
}
|
||||
|
||||
pub(super) fn envelope_key(root_secret: &[u8]) -> RuntimeResult<CredentialEnvelopeKey> {
|
||||
CredentialEnvelopeKey::derive(root_secret)
|
||||
.map_err(|_| RuntimeError::invalid_state("stable crypto.privateKey is required for persistent BYOK"))
|
||||
}
|
||||
|
||||
fn parse_definition(value: serde_json::Value) -> RuntimeResult<ByokProfileDefinition> {
|
||||
serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid stored BYOK definition", error))
|
||||
}
|
||||
|
||||
fn profile_output(row: ProfileRow) -> RuntimeResult<ByokProfileOutput> {
|
||||
let definition = parse_definition(row.definition)?;
|
||||
let validation = parse_validation(row.validation)?;
|
||||
Ok(ByokProfileOutput {
|
||||
profile_id: row.id,
|
||||
workspace_id: row.workspace_id,
|
||||
provider: row.provider,
|
||||
name: row.name,
|
||||
description: row.description,
|
||||
definition: definition.into(),
|
||||
enabled: row.enabled,
|
||||
sort_order: row.sort_order,
|
||||
revision: row.revision,
|
||||
validation,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_validation(value: Option<serde_json::Value>) -> RuntimeResult<Option<ByokValidationOutput>> {
|
||||
value
|
||||
.map(|value| {
|
||||
serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid stored BYOK validation", error))
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
pub(super) fn require_text(value: &str, field: &'static str) -> RuntimeResult<()> {
|
||||
if value.trim().is_empty() {
|
||||
Err(RuntimeError::invalid_input(format!("{field} is required")))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,3 @@
|
||||
pub(super) const BYOK_LOCAL_LEASE_ACTIVE_PURPOSE: &str = "copilot_byok_local_lease:active";
|
||||
pub(super) const BYOK_LOCAL_LEASE_PURPOSE: &str = "copilot_byok_local_lease";
|
||||
pub(super) const MAGIC_LINK_OTP_PURPOSE: &str = "magic_link_otp";
|
||||
pub(super) const MAX_MAGIC_LINK_OTP_ATTEMPTS: i32 = 10;
|
||||
pub(super) const WORKSPACE_INVITE_LINK_ID_PURPOSE: &str = "workspace_invite_link:id";
|
||||
|
||||
@@ -0,0 +1,266 @@
|
||||
use llm_adapter::capability::provider_default_capability_upper_bound;
|
||||
use sqlx::{FromRow, PgPool, Row};
|
||||
|
||||
use super::super::{LocalLeasePayload, RuntimeError, RuntimeResult, token_hash};
|
||||
use crate::{
|
||||
llm::{
|
||||
CopilotAccessProjection,
|
||||
byok::{ByokEndpoint, ByokModelDeclaration, ByokProfileDefinition, local_aad, server_aad},
|
||||
route::{self, AuthorizedProfileRef, CatalogSlot, CredentialRef, ProfileSource},
|
||||
},
|
||||
runtime::{CopilotManagedProfileConfig, CopilotRuntimeConfig},
|
||||
};
|
||||
|
||||
#[derive(FromRow)]
|
||||
struct ServerProfileRow {
|
||||
id: String,
|
||||
workspace_id: String,
|
||||
provider: String,
|
||||
encrypted_api_key: String,
|
||||
definition: serde_json::Value,
|
||||
sort_order: i32,
|
||||
}
|
||||
|
||||
pub(super) struct ProfileLoadInput<'a> {
|
||||
pub(super) slot: &'a CatalogSlot,
|
||||
pub(super) built_in_route_id: Option<&'a str>,
|
||||
pub(super) workspace_id: Option<&'a str>,
|
||||
pub(super) user_id: Option<&'a str>,
|
||||
pub(super) local_lease_id: Option<&'a str>,
|
||||
pub(super) access: &'a CopilotAccessProjection,
|
||||
pub(super) managed_target_id: Option<&'a str>,
|
||||
}
|
||||
|
||||
pub(super) async fn load_profiles(
|
||||
pool: &PgPool,
|
||||
config: &CopilotRuntimeConfig,
|
||||
input: ProfileLoadInput<'_>,
|
||||
) -> RuntimeResult<Vec<AuthorizedProfileRef>> {
|
||||
let mut profiles = Vec::new();
|
||||
if let Some(workspace_id) = input.workspace_id
|
||||
&& input.access.server_byok
|
||||
{
|
||||
profiles.extend(load_server_profiles(pool, workspace_id).await?);
|
||||
}
|
||||
if let (Some(workspace_id), Some(user_id), Some(lease_id)) = (input.workspace_id, input.user_id, input.local_lease_id)
|
||||
&& input.access.local_byok
|
||||
{
|
||||
profiles.extend(load_local_profiles(pool, workspace_id, user_id, lease_id).await?);
|
||||
}
|
||||
profiles.extend(load_managed_profiles(
|
||||
config,
|
||||
input.slot,
|
||||
input.built_in_route_id,
|
||||
input.access.managed_tier,
|
||||
input.managed_target_id,
|
||||
)?);
|
||||
Ok(profiles)
|
||||
}
|
||||
|
||||
async fn load_server_profiles(pool: &PgPool, workspace_id: &str) -> RuntimeResult<Vec<AuthorizedProfileRef>> {
|
||||
let rows = sqlx::query_as::<_, ServerProfileRow>(
|
||||
r#"
|
||||
SELECT id, workspace_id, provider, encrypted_api_key, definition, sort_order
|
||||
FROM ai_workspace_byok_configs
|
||||
WHERE workspace_id = $1 AND enabled = TRUE
|
||||
ORDER BY sort_order ASC, created_at ASC
|
||||
"#,
|
||||
)
|
||||
.bind(workspace_id)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("load authorized BYOK profiles failed", error))?;
|
||||
rows
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
let definition: ByokProfileDefinition = serde_json::from_value(row.definition)
|
||||
.map_err(|error| RuntimeError::json("invalid stored BYOK definition", error))?;
|
||||
let aad = server_aad(
|
||||
&row.workspace_id,
|
||||
&row.id,
|
||||
&row.provider,
|
||||
definition.endpoint_identity(),
|
||||
);
|
||||
Ok(AuthorizedProfileRef {
|
||||
profile_id: row.id,
|
||||
source: ProfileSource::Server,
|
||||
provider: row.provider,
|
||||
definition,
|
||||
sort_order: row.sort_order,
|
||||
credential_ref: CredentialRef::Envelope {
|
||||
encrypted: row.encrypted_api_key,
|
||||
aad,
|
||||
},
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn load_local_profiles(
|
||||
pool: &PgPool,
|
||||
workspace_id: &str,
|
||||
user_id: &str,
|
||||
lease_id: &str,
|
||||
) -> RuntimeResult<Vec<AuthorizedProfileRef>> {
|
||||
let payload = sqlx::query(
|
||||
r#"
|
||||
SELECT payload
|
||||
FROM runtime_states
|
||||
WHERE purpose = 'copilot_byok_local_lease' AND token_hash = $1
|
||||
AND consumed_at IS NULL AND expires_at > clock_timestamp()
|
||||
"#,
|
||||
)
|
||||
.bind(token_hash(lease_id))
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(|error| RuntimeError::database("load BYOK local lease failed", error))?
|
||||
.map(|row| row.get::<serde_json::Value, _>("payload"));
|
||||
let Some(payload) = payload else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let payload: LocalLeasePayload =
|
||||
serde_json::from_value(payload).map_err(|error| RuntimeError::json("invalid BYOK local lease", error))?;
|
||||
if payload.version != 1 || payload.workspace_id != workspace_id || payload.user_id != user_id {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Ok(
|
||||
payload
|
||||
.providers
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.filter(|(_, provider)| provider.enabled)
|
||||
.map(|(index, provider)| {
|
||||
let aad = local_aad(
|
||||
workspace_id,
|
||||
user_id,
|
||||
lease_id,
|
||||
index,
|
||||
&provider.provider,
|
||||
provider.definition.endpoint_identity(),
|
||||
);
|
||||
AuthorizedProfileRef {
|
||||
profile_id: format!("{lease_id}:{index}"),
|
||||
source: ProfileSource::Local,
|
||||
provider: provider.provider,
|
||||
definition: provider.definition,
|
||||
sort_order: index as i32,
|
||||
credential_ref: CredentialRef::Envelope {
|
||||
encrypted: provider.encrypted_credential,
|
||||
aad,
|
||||
},
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
fn load_managed_profiles(
|
||||
config: &CopilotRuntimeConfig,
|
||||
slot: &CatalogSlot,
|
||||
built_in_route_id: Option<&str>,
|
||||
managed_tier: route::CopilotManagedTier,
|
||||
managed_target_id: Option<&str>,
|
||||
) -> RuntimeResult<Vec<AuthorizedProfileRef>> {
|
||||
let targets = if let Some(target_id) = managed_target_id {
|
||||
vec![
|
||||
route::managed_selected_target(built_in_route_id, target_id, managed_tier)
|
||||
.ok_or_else(|| RuntimeError::invalid_input("managed_target_unavailable"))?,
|
||||
]
|
||||
} else if let Some(targets) = route::managed_targets(slot, built_in_route_id, managed_tier) {
|
||||
targets
|
||||
} else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
targets
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, model_id)| {
|
||||
let matches = config
|
||||
.providers
|
||||
.profiles
|
||||
.iter()
|
||||
.filter(|profile| profile.enabled && profile.models.iter().any(|model| model == model_id))
|
||||
.collect::<Vec<_>>();
|
||||
let [profile] = matches.as_slice() else {
|
||||
return Err(RuntimeError::invalid_state(if matches.is_empty() {
|
||||
"built-in managed route model is unavailable"
|
||||
} else {
|
||||
"built-in managed route model matches multiple profiles"
|
||||
}));
|
||||
};
|
||||
let capabilities = provider_default_capability_upper_bound(&profile.provider, model_id)
|
||||
.ok_or_else(|| RuntimeError::invalid_state("built-in managed route model is incompatible with its profile"))?;
|
||||
Ok(AuthorizedProfileRef {
|
||||
profile_id: profile.id.clone(),
|
||||
source: ProfileSource::Managed,
|
||||
provider: profile.provider.clone(),
|
||||
definition: ByokProfileDefinition {
|
||||
version: 1,
|
||||
endpoint: managed_endpoint(profile)?,
|
||||
models: vec![ByokModelDeclaration {
|
||||
model_id: model_id.clone(),
|
||||
enabled: true,
|
||||
capabilities,
|
||||
}],
|
||||
},
|
||||
sort_order: index as i32,
|
||||
credential_ref: CredentialRef::Managed {
|
||||
profile_id: profile.id.clone(),
|
||||
},
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn managed_endpoint(profile: &CopilotManagedProfileConfig) -> RuntimeResult<ByokEndpoint> {
|
||||
if let Some(base_url) = profile.config.get("baseURL").and_then(serde_json::Value::as_str) {
|
||||
return Ok(ByokEndpoint::Custom {
|
||||
url: llm_adapter::target::canonicalize_endpoint(base_url)
|
||||
.map_err(|error| RuntimeError::invalid_state(error.to_string()))?,
|
||||
});
|
||||
}
|
||||
let endpoint = match profile.provider.as_str() {
|
||||
"geminiVertex" | "anthropicVertex" => {
|
||||
let location = required_config_text(profile, "location")?;
|
||||
let project = required_config_text(profile, "project")?;
|
||||
let publisher = if profile.provider == "geminiVertex" {
|
||||
"google"
|
||||
} else {
|
||||
"anthropic"
|
||||
};
|
||||
format!(
|
||||
"https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/publishers/{publisher}"
|
||||
)
|
||||
}
|
||||
"cloudflareWorkersAi" => format!(
|
||||
"https://api.cloudflare.com/client/v4/accounts/{}/ai",
|
||||
required_config_text(profile, "accountId")?
|
||||
),
|
||||
_ => return Ok(ByokEndpoint::ProviderDefault),
|
||||
};
|
||||
Ok(ByokEndpoint::Custom { url: endpoint })
|
||||
}
|
||||
|
||||
pub(super) fn managed_profile<'a>(
|
||||
config: &'a CopilotRuntimeConfig,
|
||||
profile_id: &str,
|
||||
) -> RuntimeResult<&'a CopilotManagedProfileConfig> {
|
||||
config
|
||||
.providers
|
||||
.profiles
|
||||
.iter()
|
||||
.find(|profile| profile.id == profile_id && profile.enabled)
|
||||
.ok_or_else(|| RuntimeError::invalid_state("managed copilot credential unavailable"))
|
||||
}
|
||||
|
||||
pub(super) fn required_config_text<'a>(
|
||||
profile: &'a CopilotManagedProfileConfig,
|
||||
field: &'static str,
|
||||
) -> RuntimeResult<&'a str> {
|
||||
profile
|
||||
.config
|
||||
.get(field)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.ok_or_else(|| RuntimeError::invalid_state(format!("managed copilot profile requires {field}")))
|
||||
}
|
||||
@@ -0,0 +1,407 @@
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
sync::Arc,
|
||||
};
|
||||
|
||||
use gcp_auth::{CustomServiceAccount, TokenProvider};
|
||||
use llm_adapter::{
|
||||
backend::{BackendError, DefaultHttpClient},
|
||||
capability::{AttachmentKind, AttachmentSource},
|
||||
core::{CoreContent, ImageInput, ImageRequest},
|
||||
router::{ExecutablePreparedRoute, ExecutableProtocol, ExecutableRequest, ExecutableResponse},
|
||||
target::{
|
||||
BackendCredential, BackendEndpoint, BackendOperation, BackendProtocol, BackendProvider, BackendTargetInput,
|
||||
EgressPolicy, compile_backend_target,
|
||||
},
|
||||
};
|
||||
use llm_runtime::{CompiledPlan, CompiledRoute, RuntimeRouteEvent, RuntimeUsage, dispatch_compiled_plan};
|
||||
use serde::Serialize;
|
||||
use uuid::Uuid;
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
use super::{COPILOT_REQUEST_TIMEOUT, RuntimeError, RuntimeResult, context};
|
||||
use crate::{
|
||||
llm::{
|
||||
byok::{ByokEndpoint, CredentialEnvelopeKey},
|
||||
route::{
|
||||
AuthorizedProfileRef, AuthorizedTargetRef, CatalogSlot, CredentialRef, RouteOperation, with_request_requirements,
|
||||
},
|
||||
},
|
||||
runtime::{BackendRuntimeConfig, CopilotManagedProfileConfig},
|
||||
};
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "snake_case", tag = "type")]
|
||||
pub(super) enum ProductEvent {
|
||||
RouteSelected { route: RouteIdentity },
|
||||
RouteFailed { route: RouteIdentity, error_kind: String },
|
||||
Usage { route: RouteIdentity, usage: ProductUsage },
|
||||
}
|
||||
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(super) struct RouteIdentity {
|
||||
profile_id: String,
|
||||
source: &'static str,
|
||||
provider: String,
|
||||
model: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(untagged)]
|
||||
pub(super) enum ProductUsage {
|
||||
Tokens(llm_adapter::core::CoreUsage),
|
||||
Image(llm_adapter::core::ImageUsage),
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(super) struct CopilotExecutionResult {
|
||||
events: Vec<ProductEvent>,
|
||||
result: serde_json::Value,
|
||||
}
|
||||
|
||||
pub(super) struct CompiledExecution {
|
||||
pub(super) plan: CompiledPlan,
|
||||
identities: HashMap<String, RouteIdentity>,
|
||||
}
|
||||
|
||||
impl CompiledExecution {
|
||||
pub(super) fn project(&self, event: RuntimeRouteEvent) -> RuntimeResult<ProductEvent> {
|
||||
project_event(event, &self.identities)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn request_and_slot(
|
||||
slot: CatalogSlot,
|
||||
request: serde_json::Value,
|
||||
) -> RuntimeResult<(CatalogSlot, ExecutableRequest)> {
|
||||
let executable = match slot.operation {
|
||||
RouteOperation::Chat => ExecutableRequest::Chat(parse_request(request)?),
|
||||
RouteOperation::Structured | RouteOperation::Transcription => {
|
||||
ExecutableRequest::Structured(parse_request(request)?)
|
||||
}
|
||||
RouteOperation::Embedding => ExecutableRequest::Embedding(parse_request(request)?),
|
||||
RouteOperation::Rerank => ExecutableRequest::Rerank(parse_request(request)?),
|
||||
RouteOperation::Image => ExecutableRequest::Image(Box::new(parse_request(request)?)),
|
||||
};
|
||||
let (needs_tools, attachment_kinds, attachment_sources) = request_requirements(&executable);
|
||||
Ok((
|
||||
with_request_requirements(slot, needs_tools, attachment_kinds, attachment_sources),
|
||||
executable,
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) fn execute(
|
||||
config: Arc<BackendRuntimeConfig>,
|
||||
slot: CatalogSlot,
|
||||
request: ExecutableRequest,
|
||||
profiles: Vec<AuthorizedProfileRef>,
|
||||
candidates: Vec<AuthorizedTargetRef>,
|
||||
managed_credentials: HashMap<String, Zeroizing<String>>,
|
||||
) -> RuntimeResult<CopilotExecutionResult> {
|
||||
let execution = compile_execution(&config, slot, request, &profiles, &candidates, &managed_credentials)?;
|
||||
let mut runtime_events = Vec::new();
|
||||
let response = dispatch_compiled_plan(&DefaultHttpClient::default(), &execution.plan, |event| {
|
||||
runtime_events.push(event)
|
||||
})
|
||||
.map_err(|error| RuntimeError::invalid_state(error.to_string()))?;
|
||||
let events = runtime_events
|
||||
.into_iter()
|
||||
.map(|event| execution.project(event))
|
||||
.collect::<RuntimeResult<Vec<_>>>()?;
|
||||
Ok(CopilotExecutionResult {
|
||||
events,
|
||||
result: response_value(response)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn compile_execution(
|
||||
config: &BackendRuntimeConfig,
|
||||
slot: CatalogSlot,
|
||||
request: ExecutableRequest,
|
||||
profiles: &[AuthorizedProfileRef],
|
||||
candidates: &[AuthorizedTargetRef],
|
||||
managed_credentials: &HashMap<String, Zeroizing<String>>,
|
||||
) -> RuntimeResult<CompiledExecution> {
|
||||
let key = CredentialEnvelopeKey::derive(config.private_key.as_bytes())
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
|
||||
let mut identities = HashMap::new();
|
||||
let mut routes = Vec::with_capacity(candidates.len());
|
||||
for candidate in candidates {
|
||||
let profile = profiles
|
||||
.get(candidate.profile_index)
|
||||
.ok_or_else(|| RuntimeError::invalid_state("invalid authorized route profile"))?;
|
||||
let model = profile
|
||||
.definition
|
||||
.models
|
||||
.get(candidate.model_index)
|
||||
.ok_or_else(|| RuntimeError::invalid_state("invalid authorized route model"))?;
|
||||
let credential = resolve_credential(&key, profile, managed_credentials)?;
|
||||
let target = compile_backend_target(BackendTargetInput {
|
||||
provider: provider(&profile.provider)?,
|
||||
operation: operation(slot.operation),
|
||||
endpoint: endpoint(&profile.provider, &profile.definition.endpoint),
|
||||
model: model.model_id.clone(),
|
||||
credential: BackendCredential::new(credential),
|
||||
timeout_ms: Some(COPILOT_REQUEST_TIMEOUT.as_millis() as u64),
|
||||
egress_policy: if profile.source != crate::llm::route::ProfileSource::Managed
|
||||
&& config.copilot.byok.allow_private_endpoint
|
||||
{
|
||||
EgressPolicy::AllowPrivate
|
||||
} else {
|
||||
EgressPolicy::PublicOnly
|
||||
},
|
||||
})
|
||||
.map_err(|error| RuntimeError::invalid_state(error.to_string()))?;
|
||||
let route_id = Uuid::new_v4().to_string();
|
||||
identities.insert(
|
||||
route_id.clone(),
|
||||
RouteIdentity {
|
||||
profile_id: profile.profile_id.clone(),
|
||||
source: match profile.source {
|
||||
crate::llm::route::ProfileSource::Server => "server",
|
||||
crate::llm::route::ProfileSource::Local => "local",
|
||||
crate::llm::route::ProfileSource::Managed => "affine_cloud",
|
||||
},
|
||||
provider: profile.provider.clone(),
|
||||
model: target.model.clone(),
|
||||
},
|
||||
);
|
||||
let protocol = protocol(target.protocol);
|
||||
let route = ExecutablePreparedRoute::new(protocol, target.model, target.config, request.clone())
|
||||
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
|
||||
routes.push(CompiledRoute::new(route_id, route));
|
||||
}
|
||||
Ok(CompiledExecution {
|
||||
plan: CompiledPlan::new(routes).map_err(|error| RuntimeError::invalid_state(error.to_string()))?,
|
||||
identities,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_request<T: serde::de::DeserializeOwned>(value: serde_json::Value) -> RuntimeResult<T> {
|
||||
serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid copilot execution request", error))
|
||||
}
|
||||
|
||||
fn request_requirements(request: &ExecutableRequest) -> (bool, Vec<AttachmentKind>, Vec<AttachmentSource>) {
|
||||
let mut kinds = HashSet::new();
|
||||
let mut sources = HashSet::new();
|
||||
let needs_tools = match request {
|
||||
ExecutableRequest::Chat(request) => {
|
||||
collect_message_attachments(&request.messages, &mut kinds, &mut sources);
|
||||
!request.tools.is_empty()
|
||||
}
|
||||
ExecutableRequest::Structured(request) => {
|
||||
collect_message_attachments(&request.messages, &mut kinds, &mut sources);
|
||||
false
|
||||
}
|
||||
ExecutableRequest::Image(request) => {
|
||||
if let ImageRequest::Edit(request) = request.as_ref() {
|
||||
kinds.insert(AttachmentKind::Image);
|
||||
for image in &request.images {
|
||||
sources.insert(match image {
|
||||
ImageInput::Url { .. } => AttachmentSource::Url,
|
||||
ImageInput::Data { .. } => AttachmentSource::Data,
|
||||
ImageInput::Bytes { .. } => AttachmentSource::Bytes,
|
||||
});
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
ExecutableRequest::Embedding(_) | ExecutableRequest::Rerank(_) => false,
|
||||
};
|
||||
(needs_tools, kinds.into_iter().collect(), sources.into_iter().collect())
|
||||
}
|
||||
|
||||
fn collect_message_attachments(
|
||||
messages: &[llm_adapter::core::CoreMessage],
|
||||
kinds: &mut HashSet<AttachmentKind>,
|
||||
sources: &mut HashSet<AttachmentSource>,
|
||||
) {
|
||||
for content in messages.iter().flat_map(|message| &message.content) {
|
||||
let source = match content {
|
||||
CoreContent::Image { source } => {
|
||||
kinds.insert(AttachmentKind::Image);
|
||||
source
|
||||
}
|
||||
CoreContent::Audio { source } => {
|
||||
kinds.insert(AttachmentKind::Audio);
|
||||
source
|
||||
}
|
||||
CoreContent::File { source } => {
|
||||
kinds.insert(AttachmentKind::File);
|
||||
source
|
||||
}
|
||||
_ => continue,
|
||||
};
|
||||
let source_kind = if source.get("url").is_some() {
|
||||
AttachmentSource::Url
|
||||
} else if source.get("data").is_some() || source.get("data_base64").is_some() {
|
||||
AttachmentSource::Data
|
||||
} else if source.get("bytes").is_some() {
|
||||
AttachmentSource::Bytes
|
||||
} else {
|
||||
AttachmentSource::FileHandle
|
||||
};
|
||||
sources.insert(source_kind);
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_credential(
|
||||
key: &CredentialEnvelopeKey,
|
||||
profile: &AuthorizedProfileRef,
|
||||
managed_credentials: &HashMap<String, Zeroizing<String>>,
|
||||
) -> RuntimeResult<String> {
|
||||
match &profile.credential_ref {
|
||||
CredentialRef::Envelope { encrypted, aad } => key
|
||||
.decrypt(encrypted, aad)
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))
|
||||
.and_then(|credential| {
|
||||
String::from_utf8(credential.expose().to_vec())
|
||||
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))
|
||||
}),
|
||||
CredentialRef::Managed { profile_id } => managed_credentials
|
||||
.get(profile_id)
|
||||
.map(|credential| credential.as_str().to_string())
|
||||
.ok_or_else(|| RuntimeError::invalid_state("managed copilot credential unavailable")),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn managed_credential(
|
||||
profile: &CopilotManagedProfileConfig,
|
||||
token_provider: Option<Arc<dyn TokenProvider>>,
|
||||
) -> RuntimeResult<String> {
|
||||
if matches!(profile.provider.as_str(), "geminiVertex" | "anthropicVertex") {
|
||||
return token_provider
|
||||
.ok_or_else(|| RuntimeError::invalid_state("managed Vertex credential unavailable"))?
|
||||
.token(&["https://www.googleapis.com/auth/cloud-platform"])
|
||||
.await
|
||||
.map(|token| token.as_str().to_string())
|
||||
.map_err(|_| RuntimeError::invalid_state("managed Vertex credential unavailable"));
|
||||
}
|
||||
let field = if profile.provider == "cloudflareWorkersAi" {
|
||||
"apiToken"
|
||||
} else {
|
||||
"apiKey"
|
||||
};
|
||||
Ok(context::required_config_text(profile, field)?.to_string())
|
||||
}
|
||||
|
||||
pub(super) async fn create_vertex_token_provider(
|
||||
profile: &CopilotManagedProfileConfig,
|
||||
) -> RuntimeResult<Arc<dyn TokenProvider>> {
|
||||
if let Some(credentials) = profile.config.pointer("/googleAuthOptions/credentials") {
|
||||
let project = context::required_config_text(profile, "project")?.to_string();
|
||||
let mut credentials = credentials.clone();
|
||||
let object = credentials
|
||||
.as_object_mut()
|
||||
.ok_or_else(|| RuntimeError::invalid_state("managed Vertex credentials must be an object"))?;
|
||||
object
|
||||
.entry("type")
|
||||
.or_insert_with(|| serde_json::Value::String("service_account".to_string()));
|
||||
object
|
||||
.entry("project_id")
|
||||
.or_insert_with(|| serde_json::Value::String(project));
|
||||
object
|
||||
.entry("token_uri")
|
||||
.or_insert_with(|| serde_json::Value::String("https://oauth2.googleapis.com/token".to_string()));
|
||||
let json = serde_json::to_string(&credentials)
|
||||
.map_err(|error| RuntimeError::json("serialize managed Vertex credentials failed", error))?;
|
||||
return CustomServiceAccount::from_json(&json)
|
||||
.map(|provider| Arc::new(provider) as Arc<dyn TokenProvider>)
|
||||
.map_err(|_| RuntimeError::invalid_state("managed Vertex credential unavailable"));
|
||||
}
|
||||
gcp_auth::provider()
|
||||
.await
|
||||
.map_err(|_| RuntimeError::invalid_state("managed Vertex credential unavailable"))
|
||||
}
|
||||
|
||||
pub(in crate::runtime::backend_runtime) fn provider(value: &str) -> RuntimeResult<BackendProvider> {
|
||||
match value {
|
||||
"openai" => Ok(BackendProvider::OpenAi),
|
||||
"anthropic" => Ok(BackendProvider::Anthropic),
|
||||
"anthropicVertex" => Ok(BackendProvider::AnthropicVertex),
|
||||
"gemini" => Ok(BackendProvider::Gemini),
|
||||
"geminiVertex" => Ok(BackendProvider::GeminiVertex),
|
||||
"cloudflareWorkersAi" => Ok(BackendProvider::CloudflareWorkersAi),
|
||||
"fal" => Ok(BackendProvider::Fal),
|
||||
_ => Err(RuntimeError::invalid_state("unsupported copilot provider")),
|
||||
}
|
||||
}
|
||||
|
||||
fn operation(value: RouteOperation) -> BackendOperation {
|
||||
match value {
|
||||
RouteOperation::Chat => BackendOperation::Chat,
|
||||
RouteOperation::Structured | RouteOperation::Transcription => BackendOperation::Structured,
|
||||
RouteOperation::Embedding => BackendOperation::Embedding,
|
||||
RouteOperation::Rerank => BackendOperation::Rerank,
|
||||
RouteOperation::Image => BackendOperation::Image,
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::runtime::backend_runtime) fn endpoint(provider: &str, value: &ByokEndpoint) -> BackendEndpoint {
|
||||
match (provider, value) {
|
||||
("anthropic", ByokEndpoint::ProviderDefault) => BackendEndpoint::Custom("https://api.anthropic.com".to_string()),
|
||||
("openai" | "anthropic", ByokEndpoint::Custom { url }) => {
|
||||
BackendEndpoint::Custom(url.strip_suffix("/v1").unwrap_or(url).to_string())
|
||||
}
|
||||
(_, ByokEndpoint::ProviderDefault) => BackendEndpoint::ProviderDefault,
|
||||
(_, ByokEndpoint::Custom { url }) => BackendEndpoint::Custom(url.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::runtime::backend_runtime) fn protocol(value: BackendProtocol) -> ExecutableProtocol {
|
||||
match value {
|
||||
BackendProtocol::Chat(value) => ExecutableProtocol::Chat(value),
|
||||
BackendProtocol::Structured(value) => ExecutableProtocol::Structured(value),
|
||||
BackendProtocol::Embedding(value) => ExecutableProtocol::Embedding(value),
|
||||
BackendProtocol::Rerank(value) => ExecutableProtocol::Rerank(value),
|
||||
BackendProtocol::Image(value) => ExecutableProtocol::Image(value),
|
||||
}
|
||||
}
|
||||
|
||||
fn project_event(event: RuntimeRouteEvent, identities: &HashMap<String, RouteIdentity>) -> RuntimeResult<ProductEvent> {
|
||||
match event {
|
||||
RuntimeRouteEvent::Selected { route_id } => Ok(ProductEvent::RouteSelected {
|
||||
route: identity(identities, &route_id)?,
|
||||
}),
|
||||
RuntimeRouteEvent::Failed { route_id, error_kind } => Ok(ProductEvent::RouteFailed {
|
||||
route: identity(identities, &route_id)?,
|
||||
error_kind,
|
||||
}),
|
||||
RuntimeRouteEvent::Usage { route_id, usage } => Ok(ProductEvent::Usage {
|
||||
route: identity(identities, &route_id)?,
|
||||
usage: match usage {
|
||||
RuntimeUsage::Tokens(usage) => ProductUsage::Tokens(usage),
|
||||
RuntimeUsage::Image(usage) => ProductUsage::Image(usage),
|
||||
},
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn identity(identities: &HashMap<String, RouteIdentity>, route_id: &str) -> RuntimeResult<RouteIdentity> {
|
||||
identities
|
||||
.get(route_id)
|
||||
.cloned()
|
||||
.ok_or_else(|| RuntimeError::invalid_state("runtime emitted an unknown route id"))
|
||||
}
|
||||
|
||||
fn response_value(response: ExecutableResponse) -> RuntimeResult<serde_json::Value> {
|
||||
match response {
|
||||
ExecutableResponse::Chat(response) => serialize_response(response),
|
||||
ExecutableResponse::Structured(response) => serialize_response(response),
|
||||
ExecutableResponse::Embedding(response) => serialize_response(response),
|
||||
ExecutableResponse::Rerank(response) => serialize_response(response),
|
||||
ExecutableResponse::Image(response) => serialize_response(response),
|
||||
}
|
||||
}
|
||||
|
||||
fn serialize_response(value: impl Serialize) -> RuntimeResult<serde_json::Value> {
|
||||
serde_json::to_value(value).map_err(|error| RuntimeError::json("serialize copilot response failed", error))
|
||||
}
|
||||
|
||||
impl From<BackendError> for RuntimeError {
|
||||
fn from(error: BackendError) -> Self {
|
||||
RuntimeError::invalid_state(error.to_string())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
mod context;
|
||||
mod dispatch;
|
||||
mod stream;
|
||||
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Arc, RwLock},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use gcp_auth::TokenProvider;
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::sync::OnceCell;
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
pub(in crate::runtime::backend_runtime) use dispatch::{
|
||||
endpoint as byok_endpoint, protocol as executable_protocol, provider as backend_provider,
|
||||
};
|
||||
|
||||
use super::{BackendRuntime, RuntimeError, RuntimeResult, to_napi_error};
|
||||
use crate::{
|
||||
llm::{
|
||||
CopilotExecuteInput, CopilotRouteCheckInput,
|
||||
route::{self, AuthorizedProfileRef, AuthorizedTargetRef, CredentialRef},
|
||||
},
|
||||
runtime::{BackendRuntimeConfig, CopilotManagedProfileConfig},
|
||||
};
|
||||
|
||||
pub(super) type ManagedTokenProviderCache = RwLock<HashMap<String, Arc<OnceCell<Arc<dyn TokenProvider>>>>>;
|
||||
pub(super) const COPILOT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30 * 60);
|
||||
|
||||
struct AuthorizedCopilotRoute {
|
||||
config: std::sync::Arc<BackendRuntimeConfig>,
|
||||
slot: route::CatalogSlot,
|
||||
profiles: Vec<route::AuthorizedProfileRef>,
|
||||
candidates: Vec<route::AuthorizedTargetRef>,
|
||||
}
|
||||
|
||||
#[napi_derive::napi]
|
||||
impl BackendRuntime {
|
||||
#[napi]
|
||||
pub async fn execute_copilot(&self, input: CopilotExecuteInput) -> napi::Result<String> {
|
||||
self.execute_copilot_inner(input).await.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn assert_copilot_route(&self, input: CopilotRouteCheckInput) -> napi::Result<()> {
|
||||
let slot = route::slot(&input.slot).ok_or_else(|| RuntimeError::invalid_input("unknown copilot route slot"));
|
||||
self
|
||||
.authorize_copilot_route(input, slot.map_err(to_napi_error)?)
|
||||
.await
|
||||
.map(|_| ())
|
||||
.map_err(to_napi_error)
|
||||
}
|
||||
}
|
||||
|
||||
impl BackendRuntime {
|
||||
async fn execute_copilot_inner(&self, input: CopilotExecuteInput) -> RuntimeResult<String> {
|
||||
let (config, slot, request, profiles, candidates, managed_credentials) = self.prepare_copilot(input).await?;
|
||||
let output = tokio::task::spawn_blocking(move || {
|
||||
dispatch::execute(config, slot, request, profiles, candidates, managed_credentials)
|
||||
})
|
||||
.await
|
||||
.map_err(|error| RuntimeError::invalid_state(format!("copilot execution task failed: {error}")))??;
|
||||
serde_json::to_string(&output).map_err(|error| RuntimeError::json("serialize copilot execution failed", error))
|
||||
}
|
||||
|
||||
pub(super) async fn prepare_copilot(
|
||||
&self,
|
||||
input: CopilotExecuteInput,
|
||||
) -> RuntimeResult<stream::PreparedCopilotExecution> {
|
||||
let CopilotExecuteInput {
|
||||
slot,
|
||||
built_in_route_id,
|
||||
workspace_id,
|
||||
user_id,
|
||||
local_lease_id,
|
||||
access,
|
||||
managed_target_id,
|
||||
target_override,
|
||||
request,
|
||||
} = input;
|
||||
let base_slot = route::slot(&slot).ok_or_else(|| RuntimeError::invalid_input("unknown copilot route slot"))?;
|
||||
let (slot, request) = dispatch::request_and_slot(base_slot, request)?;
|
||||
let authorized = self
|
||||
.authorize_copilot_route(
|
||||
CopilotRouteCheckInput {
|
||||
slot: slot.id.to_string(),
|
||||
built_in_route_id,
|
||||
workspace_id,
|
||||
user_id,
|
||||
local_lease_id,
|
||||
access,
|
||||
managed_target_id,
|
||||
target_override,
|
||||
},
|
||||
slot,
|
||||
)
|
||||
.await?;
|
||||
let managed_credentials = self
|
||||
.resolve_managed_credentials(&authorized.config, &authorized.profiles, &authorized.candidates)
|
||||
.await?;
|
||||
Ok((
|
||||
authorized.config,
|
||||
authorized.slot,
|
||||
request,
|
||||
authorized.profiles,
|
||||
authorized.candidates,
|
||||
managed_credentials,
|
||||
))
|
||||
}
|
||||
|
||||
async fn resolve_managed_credentials(
|
||||
&self,
|
||||
config: &BackendRuntimeConfig,
|
||||
profiles: &[AuthorizedProfileRef],
|
||||
candidates: &[AuthorizedTargetRef],
|
||||
) -> RuntimeResult<HashMap<String, Zeroizing<String>>> {
|
||||
let mut credentials = HashMap::new();
|
||||
for candidate in candidates {
|
||||
let profile = profiles
|
||||
.get(candidate.profile_index)
|
||||
.ok_or_else(|| RuntimeError::invalid_state("invalid authorized route profile"))?;
|
||||
let CredentialRef::Managed { profile_id } = &profile.credential_ref else {
|
||||
continue;
|
||||
};
|
||||
if credentials.contains_key(profile_id) {
|
||||
continue;
|
||||
}
|
||||
let managed = context::managed_profile(&config.copilot, profile_id)?;
|
||||
let token_provider = if matches!(managed.provider.as_str(), "geminiVertex" | "anthropicVertex") {
|
||||
Some(self.managed_token_provider(managed).await?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
credentials.insert(
|
||||
profile_id.clone(),
|
||||
Zeroizing::new(dispatch::managed_credential(managed, token_provider).await?),
|
||||
);
|
||||
}
|
||||
Ok(credentials)
|
||||
}
|
||||
|
||||
async fn managed_token_provider(
|
||||
&self,
|
||||
profile: &CopilotManagedProfileConfig,
|
||||
) -> RuntimeResult<Arc<dyn TokenProvider>> {
|
||||
let config = serde_json::to_vec(&profile.config)
|
||||
.map_err(|error| RuntimeError::json("serialize managed Vertex profile failed", error))?;
|
||||
let cache_key = format!(
|
||||
"{}:{}:{}",
|
||||
profile.id,
|
||||
profile.provider,
|
||||
hex::encode(Sha256::digest(config))
|
||||
);
|
||||
let cell = {
|
||||
let mut providers = self
|
||||
.managed_token_providers
|
||||
.write()
|
||||
.map_err(|_| RuntimeError::invalid_state("managed token provider cache lock poisoned"))?;
|
||||
Arc::clone(providers.entry(cache_key).or_insert_with(|| Arc::new(OnceCell::new())))
|
||||
};
|
||||
cell
|
||||
.get_or_try_init(|| dispatch::create_vertex_token_provider(profile))
|
||||
.await
|
||||
.map(Arc::clone)
|
||||
}
|
||||
|
||||
async fn authorize_copilot_route(
|
||||
&self,
|
||||
input: CopilotRouteCheckInput,
|
||||
slot: route::CatalogSlot,
|
||||
) -> RuntimeResult<AuthorizedCopilotRoute> {
|
||||
let config = self.config()?;
|
||||
if !config.copilot.enabled {
|
||||
return Err(RuntimeError::invalid_state("copilot_disabled"));
|
||||
}
|
||||
let deployment = if std::env::var("DEPLOYMENT_TYPE").as_deref() == Ok("selfhosted") {
|
||||
route::Deployment::SelfHosted
|
||||
} else {
|
||||
route::Deployment::Cloud
|
||||
};
|
||||
let profiles = context::load_profiles(
|
||||
&self.pool().await?,
|
||||
&config.copilot,
|
||||
context::ProfileLoadInput {
|
||||
slot: &slot,
|
||||
built_in_route_id: input.built_in_route_id.as_deref(),
|
||||
workspace_id: input.workspace_id.as_deref(),
|
||||
user_id: input.user_id.as_deref(),
|
||||
local_lease_id: input.local_lease_id.as_deref(),
|
||||
access: &input.access,
|
||||
managed_target_id: input.managed_target_id.as_deref(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
if input.managed_target_id.is_some() && input.target_override.is_some() {
|
||||
return Err(RuntimeError::invalid_input("multiple_target_selections"));
|
||||
}
|
||||
let target_override = if let Some(target_id) = input.managed_target_id.as_deref() {
|
||||
let model_id =
|
||||
route::managed_selected_target(input.built_in_route_id.as_deref(), target_id, input.access.managed_tier)
|
||||
.ok_or_else(|| RuntimeError::invalid_input("managed_target_unavailable"))?;
|
||||
let profile = profiles
|
||||
.iter()
|
||||
.find(|profile| {
|
||||
profile.source == route::ProfileSource::Managed
|
||||
&& profile.definition.models.iter().any(|model| model.model_id == model_id)
|
||||
})
|
||||
.ok_or_else(|| RuntimeError::invalid_state("managed_target_unavailable"))?;
|
||||
Some(route::TargetOverride {
|
||||
profile_id: profile.profile_id.clone(),
|
||||
model_id,
|
||||
})
|
||||
} else {
|
||||
input.target_override.map(|target| route::TargetOverride {
|
||||
profile_id: target.profile_id,
|
||||
model_id: target.model_id,
|
||||
})
|
||||
};
|
||||
let candidates = match route::decide(route::RoutePolicyInput {
|
||||
slot: &slot,
|
||||
deployment,
|
||||
byok_enabled: config.copilot.byok.enabled,
|
||||
access_available: input.access.route_allowed
|
||||
|| route::quota_policy(&slot, input.built_in_route_id.as_deref()) != route::QuotaPolicy::Metered,
|
||||
profiles: &profiles,
|
||||
target_override: target_override.as_ref(),
|
||||
target_override_managed: input.managed_target_id.is_some(),
|
||||
}) {
|
||||
route::RouteDecision::Ready(candidates) => candidates,
|
||||
route::RouteDecision::Denied(reason) => {
|
||||
return Err(RuntimeError::invalid_input(reason_name(reason)));
|
||||
}
|
||||
route::RouteDecision::NoRoute(reason) => {
|
||||
return Err(RuntimeError::invalid_state(reason_name(reason)));
|
||||
}
|
||||
};
|
||||
Ok(AuthorizedCopilotRoute {
|
||||
config,
|
||||
slot,
|
||||
profiles,
|
||||
candidates,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn reason_name(reason: route::RouteDecisionReason) -> &'static str {
|
||||
match reason {
|
||||
route::RouteDecisionReason::ByokDisabled => "byok_disabled",
|
||||
route::RouteDecisionReason::AccessUnavailable => "access_unavailable",
|
||||
route::RouteDecisionReason::ExplicitTargetUnavailable => "explicit_target_unavailable",
|
||||
route::RouteDecisionReason::NoCompatibleTarget => "no_compatible_target",
|
||||
route::RouteDecisionReason::ManagedPresetUnavailable => "managed_preset_unavailable",
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,261 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
mpsc,
|
||||
},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use llm_adapter::{
|
||||
backend::{BackendError, DefaultHttpClient},
|
||||
core::CoreMessage,
|
||||
router::ExecutableRequest,
|
||||
};
|
||||
use llm_runtime::{
|
||||
AccumulatedToolCall, RuntimeRouteEvent, ToolCallbackRequest, ToolCallbackResponse, ToolExecutionResult,
|
||||
ToolLoopEvent, dispatch_compiled_round, run_tool_loop,
|
||||
};
|
||||
use napi::{
|
||||
JsValue, Result, Status,
|
||||
bindgen_prelude::{CallbackContext, PromiseRaw, Unknown},
|
||||
threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode},
|
||||
};
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
use super::{BackendRuntime, COPILOT_REQUEST_TIMEOUT, RuntimeError, dispatch, to_napi_error};
|
||||
use crate::{
|
||||
llm::{
|
||||
CopilotExecuteInput,
|
||||
route::{AuthorizedProfileRef, AuthorizedTargetRef, CatalogSlot},
|
||||
},
|
||||
runtime::BackendRuntimeConfig,
|
||||
};
|
||||
|
||||
pub(super) type PreparedCopilotExecution = (
|
||||
Arc<BackendRuntimeConfig>,
|
||||
CatalogSlot,
|
||||
ExecutableRequest,
|
||||
Vec<AuthorizedProfileRef>,
|
||||
Vec<AuthorizedTargetRef>,
|
||||
HashMap<String, Zeroizing<String>>,
|
||||
);
|
||||
|
||||
const STREAM_END: &str = "__AFFINE_COPILOT_STREAM_END__";
|
||||
const TOOL_CALLBACK_TIMEOUT: Duration = Duration::from_secs(5 * 60);
|
||||
const TOOL_CALLBACK_POLL_INTERVAL: Duration = Duration::from_millis(100);
|
||||
|
||||
#[napi_derive::napi]
|
||||
pub struct CopilotStreamHandle {
|
||||
aborted: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
#[napi_derive::napi]
|
||||
impl CopilotStreamHandle {
|
||||
#[napi]
|
||||
pub fn abort(&self) {
|
||||
self.aborted.store(true, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
#[napi_derive::napi]
|
||||
impl BackendRuntime {
|
||||
#[napi]
|
||||
pub async fn execute_copilot_stream(
|
||||
&self,
|
||||
input: CopilotExecuteInput,
|
||||
max_steps: u32,
|
||||
callback: ThreadsafeFunction<String, ()>,
|
||||
tool_callback: ThreadsafeFunction<String, PromiseRaw<'static, String>>,
|
||||
) -> Result<CopilotStreamHandle> {
|
||||
let (config, slot, request, profiles, candidates, managed_credentials) =
|
||||
self.prepare_copilot(input).await.map_err(to_napi_error)?;
|
||||
let messages = match &request {
|
||||
ExecutableRequest::Chat(request) => request.messages.clone(),
|
||||
_ => {
|
||||
return Err(to_napi_error(RuntimeError::invalid_input(
|
||||
"copilot stream requires a chat slot",
|
||||
)));
|
||||
}
|
||||
};
|
||||
let mut execution =
|
||||
dispatch::compile_execution(&config, slot, request, &profiles, &candidates, &managed_credentials)
|
||||
.map_err(to_napi_error)?;
|
||||
let aborted = Arc::new(AtomicBool::new(false));
|
||||
let worker_aborted = aborted.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let deadline = Instant::now() + COPILOT_REQUEST_TIMEOUT;
|
||||
let result = run_stream(
|
||||
&mut execution,
|
||||
messages,
|
||||
max_steps.max(1) as usize,
|
||||
&callback,
|
||||
&tool_callback,
|
||||
&worker_aborted,
|
||||
deadline,
|
||||
);
|
||||
if let Err(message) = result
|
||||
&& !worker_aborted.load(Ordering::Relaxed)
|
||||
{
|
||||
let _ = emit_json(
|
||||
&callback,
|
||||
&serde_json::json!({
|
||||
"type": "error",
|
||||
"errorKind": "dispatch",
|
||||
"message": message,
|
||||
}),
|
||||
);
|
||||
}
|
||||
let _ = callback.call(Ok(STREAM_END.to_string()), ThreadsafeFunctionCallMode::Blocking);
|
||||
});
|
||||
Ok(CopilotStreamHandle { aborted })
|
||||
}
|
||||
}
|
||||
|
||||
fn run_stream(
|
||||
execution: &mut dispatch::CompiledExecution,
|
||||
mut messages: Vec<CoreMessage>,
|
||||
max_steps: usize,
|
||||
callback: &ThreadsafeFunction<String, ()>,
|
||||
tool_callback: &ThreadsafeFunction<String, PromiseRaw<'static, String>>,
|
||||
aborted: &AtomicBool,
|
||||
deadline: Instant,
|
||||
) -> std::result::Result<(), String> {
|
||||
let result = run_tool_loop(
|
||||
&mut messages,
|
||||
max_steps,
|
||||
|messages| {
|
||||
let mut route_events = Vec::new();
|
||||
let result = dispatch_compiled_round(
|
||||
&DefaultHttpClient::default(),
|
||||
&mut execution.plan,
|
||||
messages,
|
||||
|| aborted.load(Ordering::Relaxed) || Instant::now() >= deadline,
|
||||
|event| emit_json(callback, event).map_err(transport_error),
|
||||
|event: RuntimeRouteEvent| route_events.push(event),
|
||||
);
|
||||
for event in route_events {
|
||||
let event = execution.project(event).map_err(|error| error.to_string())?;
|
||||
emit_json(callback, &event)?;
|
||||
}
|
||||
result.map_err(|error| error.to_string())
|
||||
},
|
||||
|call: &AccumulatedToolCall| execute_tool(tool_callback, call, aborted, deadline),
|
||||
|event: &ToolLoopEvent| emit_json(callback, event),
|
||||
|| "tool loop reached max steps".to_string(),
|
||||
);
|
||||
if !aborted.load(Ordering::Relaxed) && Instant::now() >= deadline {
|
||||
Err("copilot stream deadline exceeded".to_string())
|
||||
} else {
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
fn emit_json(
|
||||
callback: &ThreadsafeFunction<String, ()>,
|
||||
value: &impl serde::Serialize,
|
||||
) -> std::result::Result<(), String> {
|
||||
let value = serde_json::to_string(value).map_err(|error| error.to_string())?;
|
||||
let status = callback.call(Ok(value), ThreadsafeFunctionCallMode::Blocking);
|
||||
if status == Status::Ok {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!("copilot stream callback failed: {status}"))
|
||||
}
|
||||
}
|
||||
|
||||
fn transport_error(message: String) -> BackendError {
|
||||
BackendError::Transport { message }
|
||||
}
|
||||
|
||||
fn execute_tool(
|
||||
callback: &ThreadsafeFunction<String, PromiseRaw<'static, String>>,
|
||||
call: &AccumulatedToolCall,
|
||||
aborted: &AtomicBool,
|
||||
stream_deadline: Instant,
|
||||
) -> std::result::Result<ToolExecutionResult, String> {
|
||||
let request = serde_json::to_string(&ToolCallbackRequest {
|
||||
call_id: call.id.clone(),
|
||||
name: call.name.clone(),
|
||||
args: call.args.clone(),
|
||||
raw_arguments_text: call.raw_arguments_text.clone(),
|
||||
argument_parse_error: call.argument_parse_error.clone(),
|
||||
})
|
||||
.map_err(|error| error.to_string())?;
|
||||
let (sender, receiver) = mpsc::sync_channel(1);
|
||||
let sender = Arc::new(Mutex::new(Some(sender)));
|
||||
let callback_sender = sender.clone();
|
||||
let status = callback.call_with_return_value(
|
||||
Ok(request),
|
||||
ThreadsafeFunctionCallMode::NonBlocking,
|
||||
move |promise, _env| {
|
||||
match promise {
|
||||
Ok(promise) => {
|
||||
let success_sender = callback_sender.clone();
|
||||
let failure_sender = callback_sender.clone();
|
||||
match promise.then(move |ctx| {
|
||||
send_tool_result(
|
||||
&success_sender,
|
||||
serde_json::from_str::<ToolCallbackResponse>(&ctx.value).map_err(|error| error.to_string()),
|
||||
);
|
||||
Ok(())
|
||||
}) {
|
||||
Ok(promise) => {
|
||||
if let Err(error) = promise.catch(move |ctx: CallbackContext<Unknown>| {
|
||||
let message = ctx.value.coerce_to_string()?.into_utf8()?.as_str()?.to_string();
|
||||
send_tool_result(&failure_sender, Err(message));
|
||||
Ok(())
|
||||
}) {
|
||||
send_tool_result(&callback_sender, Err(error.to_string()));
|
||||
}
|
||||
}
|
||||
Err(error) => send_tool_result(&callback_sender, Err(error.to_string())),
|
||||
}
|
||||
}
|
||||
Err(error) => send_tool_result(&callback_sender, Err(error.to_string())),
|
||||
}
|
||||
Ok(())
|
||||
},
|
||||
);
|
||||
if status != Status::Ok {
|
||||
return Err(format!("copilot tool callback failed: {status}"));
|
||||
}
|
||||
let tool_deadline = std::cmp::min(stream_deadline, Instant::now() + TOOL_CALLBACK_TIMEOUT);
|
||||
let response = loop {
|
||||
if aborted.load(Ordering::Relaxed) {
|
||||
return Err("copilot stream aborted".to_string());
|
||||
}
|
||||
let now = Instant::now();
|
||||
if now >= tool_deadline {
|
||||
return Err("copilot tool callback deadline exceeded".to_string());
|
||||
}
|
||||
match receiver.recv_timeout(std::cmp::min(TOOL_CALLBACK_POLL_INTERVAL, tool_deadline - now)) {
|
||||
Ok(response) => break response?,
|
||||
Err(mpsc::RecvTimeoutError::Timeout) => continue,
|
||||
Err(mpsc::RecvTimeoutError::Disconnected) => {
|
||||
return Err("copilot tool callback closed before completion".to_string());
|
||||
}
|
||||
}
|
||||
};
|
||||
if !response.args.is_object() {
|
||||
return Err("copilot tool callback args must be an object".to_string());
|
||||
}
|
||||
Ok(ToolExecutionResult {
|
||||
call_id: response.call_id,
|
||||
name: response.name,
|
||||
arguments: response.args,
|
||||
arguments_text: response.raw_arguments_text,
|
||||
arguments_error: response.argument_parse_error,
|
||||
output: response.output,
|
||||
is_error: response.is_error,
|
||||
})
|
||||
}
|
||||
|
||||
type ToolResultSender = Arc<Mutex<Option<mpsc::SyncSender<std::result::Result<ToolCallbackResponse, String>>>>>;
|
||||
|
||||
fn send_tool_result(sender: &ToolResultSender, result: std::result::Result<ToolCallbackResponse, String>) {
|
||||
if let Some(sender) = sender.lock().expect("tool callback sender poisoned").take() {
|
||||
let _ = sender.send(result);
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,7 @@
|
||||
mod byok;
|
||||
mod constants;
|
||||
mod coordination_lease;
|
||||
mod copilot;
|
||||
mod doc_compactor;
|
||||
mod doc_storage;
|
||||
mod gate;
|
||||
@@ -9,8 +11,13 @@ mod runtime_state;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
mod workspace_stats;
|
||||
use std::{sync::RwLock, time::Duration};
|
||||
use std::{
|
||||
sync::{Arc, RwLock},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use byok::LocalLeasePayload;
|
||||
use copilot::{backend_provider, byok_endpoint, executable_protocol};
|
||||
use napi::Result;
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::{PgPool, Row, postgres::PgPoolOptions};
|
||||
@@ -22,6 +29,11 @@ pub(super) use super::{
|
||||
BackendRuntimeConfig, InviteQuotaConfig, RuntimeError, RuntimeResult, migrations::migrate_runtime_tables, napi_error,
|
||||
to_napi_error,
|
||||
};
|
||||
use crate::llm::{
|
||||
ByokLocalLeaseOutput, ByokProbeResultOutput, ByokProfileOutput, CreateByokLocalLeaseInput, CreateByokProfileInput,
|
||||
ProbeByokDraftInput, ProbeByokProfileInput, ReorderByokProfilesInput, ReplaceByokProfileInput,
|
||||
RotateByokCredentialInput,
|
||||
};
|
||||
|
||||
pub(super) fn token_hash(token: &str) -> String {
|
||||
hex::encode(Sha256::digest(token.as_bytes()))
|
||||
@@ -29,17 +41,20 @@ pub(super) fn token_hash(token: &str) -> String {
|
||||
|
||||
#[napi_derive::napi]
|
||||
pub struct BackendRuntime {
|
||||
config: RwLock<BackendRuntimeConfig>,
|
||||
config: RwLock<Arc<BackendRuntimeConfig>>,
|
||||
pool: Mutex<Option<PgPool>>,
|
||||
managed_token_providers: copilot::ManagedTokenProviderCache,
|
||||
}
|
||||
|
||||
#[napi_derive::napi]
|
||||
impl BackendRuntime {
|
||||
#[napi(constructor)]
|
||||
pub fn new() -> Result<Self> {
|
||||
pub fn new(private_key: Option<String>) -> Result<Self> {
|
||||
let config = BackendRuntimeConfig::from_config_files(private_key).map_err(to_napi_error)?;
|
||||
Ok(Self {
|
||||
config: RwLock::new(BackendRuntimeConfig::from_config_files().map_err(to_napi_error)?),
|
||||
config: RwLock::new(Arc::new(config)),
|
||||
pool: Mutex::new(None),
|
||||
managed_token_providers: Default::default(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -54,11 +69,12 @@ impl BackendRuntime {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let database_url = self.config()?.database_url;
|
||||
let config = self.config()?;
|
||||
let database_url = &config.database_url;
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(5)
|
||||
.acquire_timeout(Duration::from_secs(5))
|
||||
.connect(&database_url)
|
||||
.connect(database_url)
|
||||
.await
|
||||
.map_err(|err| RuntimeError::database("BackendRuntime failed to connect postgres", err))?;
|
||||
|
||||
@@ -83,6 +99,18 @@ impl BackendRuntime {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn reload_config(&self, private_key: Option<String>) -> Result<()> {
|
||||
let pool = self.pool().await.map_err(to_napi_error)?;
|
||||
let active_private_key = self.config().map_err(to_napi_error)?.private_key.to_string();
|
||||
let config = BackendRuntimeConfig::from_config_files(private_key.or(Some(active_private_key)))
|
||||
.map_err(to_napi_error)?
|
||||
.with_db_overrides(&pool)
|
||||
.await
|
||||
.map_err(to_napi_error)?;
|
||||
self.update_config(config).map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn health(&self) -> Result<BackendRuntimeHealth> {
|
||||
let pool = self.pool.lock().await.as_ref().cloned();
|
||||
@@ -107,6 +135,98 @@ impl BackendRuntime {
|
||||
migrate_runtime_tables(&pool).await.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn list_byok_profiles(&self, workspace_id: String) -> Result<Vec<ByokProfileOutput>> {
|
||||
byok::list(&self.pool().await?, &workspace_id)
|
||||
.await
|
||||
.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn create_byok_profile(&self, input: CreateByokProfileInput) -> Result<ByokProfileOutput> {
|
||||
let config = self.config()?;
|
||||
byok::create(
|
||||
&self.pool().await?,
|
||||
config.private_key.as_bytes(),
|
||||
&config.copilot.byok,
|
||||
input,
|
||||
)
|
||||
.await
|
||||
.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn replace_byok_profile(&self, input: ReplaceByokProfileInput) -> Result<ByokProfileOutput> {
|
||||
let config = self.config()?;
|
||||
byok::replace(
|
||||
&self.pool().await?,
|
||||
config.private_key.as_bytes(),
|
||||
&config.copilot.byok,
|
||||
input,
|
||||
)
|
||||
.await
|
||||
.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn rotate_byok_credential(&self, input: RotateByokCredentialInput) -> Result<ByokProfileOutput> {
|
||||
let config = self.config()?;
|
||||
byok::rotate(&self.pool().await?, config.private_key.as_bytes(), input)
|
||||
.await
|
||||
.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn probe_byok_profile(&self, input: ProbeByokProfileInput) -> Result<ByokProbeResultOutput> {
|
||||
let config = self.config()?;
|
||||
byok::probe_profile(
|
||||
&self.pool().await?,
|
||||
config.private_key.as_bytes(),
|
||||
&config.copilot.byok,
|
||||
input,
|
||||
)
|
||||
.await
|
||||
.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn probe_byok_draft(&self, input: ProbeByokDraftInput) -> Result<ByokProbeResultOutput> {
|
||||
let config = self.config()?;
|
||||
byok::probe_draft(
|
||||
&self.pool().await?,
|
||||
config.private_key.as_bytes(),
|
||||
&config.copilot.byok,
|
||||
input,
|
||||
)
|
||||
.await
|
||||
.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn delete_byok_profile(&self, workspace_id: String, profile_id: String) -> Result<bool> {
|
||||
byok::delete(&self.pool().await?, &workspace_id, &profile_id)
|
||||
.await
|
||||
.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn reorder_byok_profiles(&self, input: ReorderByokProfilesInput) -> Result<Vec<ByokProfileOutput>> {
|
||||
byok::reorder(&self.pool().await?, input).await.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn create_byok_local_lease(&self, input: CreateByokLocalLeaseInput) -> Result<ByokLocalLeaseOutput> {
|
||||
let config = self.config()?;
|
||||
byok::create_local_lease(
|
||||
&self.pool().await?,
|
||||
config.private_key.as_bytes(),
|
||||
&config.copilot.byok,
|
||||
input,
|
||||
)
|
||||
.await
|
||||
.map_err(to_napi_error)
|
||||
}
|
||||
|
||||
pub(crate) async fn pool(&self) -> RuntimeResult<PgPool> {
|
||||
self
|
||||
.pool
|
||||
@@ -117,19 +237,24 @@ impl BackendRuntime {
|
||||
.ok_or_else(|| RuntimeError::invalid_state("BackendRuntime must be started before using postgres operations"))
|
||||
}
|
||||
|
||||
pub(crate) fn config(&self) -> RuntimeResult<BackendRuntimeConfig> {
|
||||
pub(crate) fn config(&self) -> RuntimeResult<Arc<BackendRuntimeConfig>> {
|
||||
self
|
||||
.config
|
||||
.read()
|
||||
.map(|config| config.clone())
|
||||
.map(|config| Arc::clone(&config))
|
||||
.map_err(|_| RuntimeError::invalid_state("BackendRuntime config lock poisoned"))
|
||||
}
|
||||
|
||||
fn update_config(&self, config: BackendRuntimeConfig) -> RuntimeResult<()> {
|
||||
self
|
||||
.managed_token_providers
|
||||
.write()
|
||||
.map_err(|_| RuntimeError::invalid_state("managed token provider cache lock poisoned"))?
|
||||
.clear();
|
||||
*self
|
||||
.config
|
||||
.write()
|
||||
.map_err(|_| RuntimeError::invalid_state("BackendRuntime config lock poisoned"))? = config;
|
||||
.map_err(|_| RuntimeError::invalid_state("BackendRuntime config lock poisoned"))? = Arc::new(config);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -75,8 +75,9 @@ impl BackendRuntime {
|
||||
&self,
|
||||
input: RuntimeMailDeliveryQuotaInput,
|
||||
) -> Result<RuntimeMailDeliveryQuotaDecision> {
|
||||
let config = self.config()?.invite_quota;
|
||||
let Some(class) = mail_class(&input.mail_name, &config) else {
|
||||
let runtime_config = self.config()?;
|
||||
let config = &runtime_config.invite_quota;
|
||||
let Some(class) = mail_class(&input.mail_name, config) else {
|
||||
return Ok(RuntimeMailDeliveryQuotaDecision {
|
||||
allowed: false,
|
||||
reservation_id: None,
|
||||
@@ -91,7 +92,7 @@ impl BackendRuntime {
|
||||
});
|
||||
};
|
||||
let pool = self.pool().await?;
|
||||
let scopes = build_mail_scopes(&input, class, &config);
|
||||
let scopes = build_mail_scopes(&input, class, config);
|
||||
match reserve_scopes(&pool, "mail_delivery", input.request_id.as_deref(), scopes).await? {
|
||||
Ok(reservation) => Ok(RuntimeMailDeliveryQuotaDecision {
|
||||
allowed: true,
|
||||
|
||||
+6
-5
@@ -293,14 +293,15 @@ impl BackendRuntime {
|
||||
if input.target_count <= 0 {
|
||||
return Err(napi_error("target_count must be positive"));
|
||||
}
|
||||
let config = self.config()?.invite_quota;
|
||||
let runtime_config = self.config()?;
|
||||
let config = &runtime_config.invite_quota;
|
||||
let pool = self.pool().await?;
|
||||
let now: DateTime<Utc> = sqlx::query_scalar("SELECT clock_timestamp()")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.map_err(|err| RuntimeError::database("failed to read database clock", err))?;
|
||||
let actor = load_actor(&pool, &input.actor_user_id).await?;
|
||||
let actor_subject = subject_hash(&actor.email, &config);
|
||||
let actor_subject = subject_hash(&actor.email, config);
|
||||
if let Some(status) = active_subject_status(&pool, &actor_subject).await?
|
||||
&& matches!(status.as_str(), "banned" | "quarantined")
|
||||
{
|
||||
@@ -387,14 +388,14 @@ impl BackendRuntime {
|
||||
action_required: None,
|
||||
});
|
||||
}
|
||||
if let Some(abuse_decision) = high_confidence_invite_abuse(&input, &actor, &config) {
|
||||
if let Some(abuse_decision) = high_confidence_invite_abuse(&input, &actor, config) {
|
||||
let reason = abuse_decision.reason;
|
||||
let scope_key = match abuse_decision.subject_kind {
|
||||
"workspace" => format!("invite:workspace_subject:{}", abuse_decision.subject_key),
|
||||
"source_prefix_domain" => format!("invite:source_cohort_subject:{}", abuse_decision.subject_key),
|
||||
_ => format!("invite:actor_subject:{}", abuse_decision.subject_key),
|
||||
};
|
||||
let action_required = record_invite_abuse_action(&pool, &input, &actor, abuse_decision, &config).await?;
|
||||
let action_required = record_invite_abuse_action(&pool, &input, &actor, abuse_decision, config).await?;
|
||||
return Ok(RuntimeWorkspaceInviteQuotaDecision {
|
||||
allowed: false,
|
||||
reservation_id: None,
|
||||
@@ -411,7 +412,7 @@ impl BackendRuntime {
|
||||
|
||||
let workspace = load_workspace(&pool, &input.workspace_id).await?;
|
||||
let activity = load_invite_activity(&pool, &input.actor_user_id, &input.workspace_id).await?;
|
||||
let scopes = build_invite_scopes(&input, &actor, &workspace, "a, &activity, &config, now)?;
|
||||
let scopes = build_invite_scopes(&input, &actor, &workspace, "a, &activity, config, now)?;
|
||||
match reserve_scopes(&pool, "workspace_invite", input.request_id.as_deref(), scopes).await? {
|
||||
Ok(reservation) => Ok(RuntimeWorkspaceInviteQuotaDecision {
|
||||
allowed: true,
|
||||
|
||||
-128
@@ -1,128 +0,0 @@
|
||||
use super::{
|
||||
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, BYOK_LOCAL_LEASE_PURPOSE, Result, RuntimeByokLocalLeaseRecord, RuntimeError,
|
||||
dto::{RuntimeStateInsertPayload, RuntimeStatePayloadRow, RuntimeStateRows},
|
||||
};
|
||||
|
||||
pub(super) async fn get(rows: &RuntimeStateRows, lease_id: String) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
|
||||
get_lease_by_id(rows, &lease_id).await
|
||||
}
|
||||
|
||||
pub(super) async fn create(
|
||||
rows: &RuntimeStateRows,
|
||||
active_key: String,
|
||||
lease_id: String,
|
||||
payload: serde_json::Value,
|
||||
ttl_ms: i64,
|
||||
) -> Result<RuntimeByokLocalLeaseRecord> {
|
||||
if ttl_ms <= 0 {
|
||||
return Err(RuntimeError::invalid_input("BYOK local lease ttl must be positive"));
|
||||
}
|
||||
|
||||
let mut tx = rows.begin("RuntimeState BYOK local lease").await?;
|
||||
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
|
||||
.bind(&active_key)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| RuntimeError::database("RuntimeState BYOK local lease active lock failed", err))?;
|
||||
|
||||
if let Some(active) = rows
|
||||
.active_payload_with_expires_for_update_in_tx(
|
||||
&mut tx,
|
||||
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE,
|
||||
&active_key,
|
||||
"RuntimeState BYOK local lease active get",
|
||||
)
|
||||
.await?
|
||||
{
|
||||
let existing_lease = match active.payload.get("leaseId").and_then(serde_json::Value::as_str) {
|
||||
Some(existing_lease_id) => get_lease_by_id_in_tx(rows, &mut tx, existing_lease_id).await?,
|
||||
None => None,
|
||||
};
|
||||
if let Some(lease) = existing_lease {
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| RuntimeError::database("RuntimeState BYOK local lease transaction commit failed", err))?;
|
||||
return Ok(lease);
|
||||
}
|
||||
|
||||
rows
|
||||
.delete_by_key_in_tx(
|
||||
&mut tx,
|
||||
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE,
|
||||
&active_key,
|
||||
"RuntimeState BYOK local lease stale active delete",
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
let expires_at_ms = rows
|
||||
.insert_payload_returning_expires_in_tx(
|
||||
&mut tx,
|
||||
RuntimeStateInsertPayload {
|
||||
purpose: BYOK_LOCAL_LEASE_PURPOSE,
|
||||
token: &lease_id,
|
||||
lookup_key: &active_key,
|
||||
payload: &payload,
|
||||
ttl_ms,
|
||||
context: "RuntimeState BYOK local lease create",
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let active_payload = serde_json::json!({ "leaseId": lease_id });
|
||||
rows
|
||||
.insert_payload_returning_expires_in_tx(
|
||||
&mut tx,
|
||||
RuntimeStateInsertPayload {
|
||||
purpose: BYOK_LOCAL_LEASE_ACTIVE_PURPOSE,
|
||||
token: &active_key,
|
||||
lookup_key: &active_key,
|
||||
payload: &active_payload,
|
||||
ttl_ms,
|
||||
context: "RuntimeState BYOK local lease active create",
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| RuntimeError::database("RuntimeState BYOK local lease transaction commit failed", err))?;
|
||||
|
||||
Ok(RuntimeByokLocalLeaseRecord {
|
||||
lease_id,
|
||||
payload,
|
||||
expires_at_ms,
|
||||
})
|
||||
}
|
||||
|
||||
async fn get_lease_by_id(rows: &RuntimeStateRows, lease_id: &str) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
|
||||
rows
|
||||
.active_payload_with_expires(BYOK_LOCAL_LEASE_PURPOSE, lease_id, "RuntimeState BYOK local lease get")
|
||||
.await?
|
||||
.map(|row| record_from_row(lease_id, row))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
async fn get_lease_by_id_in_tx(
|
||||
rows: &RuntimeStateRows,
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
lease_id: &str,
|
||||
) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
|
||||
rows
|
||||
.active_payload_with_expires_for_update_in_tx(
|
||||
tx,
|
||||
BYOK_LOCAL_LEASE_PURPOSE,
|
||||
lease_id,
|
||||
"RuntimeState BYOK local lease get",
|
||||
)
|
||||
.await?
|
||||
.map(|row| record_from_row(lease_id, row))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn record_from_row(lease_id: &str, row: RuntimeStatePayloadRow) -> Result<RuntimeByokLocalLeaseRecord> {
|
||||
Ok(RuntimeByokLocalLeaseRecord {
|
||||
lease_id: lease_id.to_string(),
|
||||
payload: row.payload,
|
||||
expires_at_ms: row.expires_at_ms,
|
||||
})
|
||||
}
|
||||
@@ -1,18 +1,14 @@
|
||||
use super::{BackendRuntime, RuntimeError, RuntimeResult, napi_error};
|
||||
pub(super) use super::{
|
||||
constants::{
|
||||
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, BYOK_LOCAL_LEASE_PURPOSE, MAGIC_LINK_OTP_PURPOSE, MAX_MAGIC_LINK_OTP_ATTEMPTS,
|
||||
WORKSPACE_INVITE_LINK_ID_PURPOSE, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
|
||||
MAGIC_LINK_OTP_PURPOSE, MAX_MAGIC_LINK_OTP_ATTEMPTS, WORKSPACE_INVITE_LINK_ID_PURPOSE,
|
||||
WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
|
||||
},
|
||||
token_hash,
|
||||
types::{
|
||||
RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord,
|
||||
RuntimeWorkspaceInviteLinkRecord,
|
||||
},
|
||||
types::{RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord, RuntimeWorkspaceInviteLinkRecord},
|
||||
};
|
||||
|
||||
mod auth_challenge;
|
||||
mod byok_local_lease;
|
||||
mod dto;
|
||||
mod invite_link;
|
||||
mod magic_link_otp;
|
||||
@@ -197,28 +193,6 @@ impl BackendRuntime {
|
||||
.map_err(napi::Error::from)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn create_byok_local_lease(
|
||||
&self,
|
||||
active_key: String,
|
||||
lease_id: String,
|
||||
payload: serde_json::Value,
|
||||
ttl_ms: i64,
|
||||
) -> napi::Result<RuntimeByokLocalLeaseRecord> {
|
||||
RuntimeStateStore::new(self.pool().await?)
|
||||
.create_byok_local_lease(active_key, lease_id, payload, ttl_ms)
|
||||
.await
|
||||
.map_err(napi::Error::from)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn get_byok_local_lease(&self, lease_id: String) -> napi::Result<Option<RuntimeByokLocalLeaseRecord>> {
|
||||
RuntimeStateStore::new(self.pool().await?)
|
||||
.get_byok_local_lease(lease_id)
|
||||
.await
|
||||
.map_err(napi::Error::from)
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn cleanup_expired_runtime_states(&self, limit: i64) -> napi::Result<i64> {
|
||||
if limit <= 0 {
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
use sqlx::PgPool;
|
||||
|
||||
use super::{
|
||||
Result, RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord,
|
||||
RuntimeWorkspaceInviteLinkRecord, auth_challenge, byok_local_lease, dto::RuntimeStateRows, invite_link,
|
||||
magic_link_otp, verification_token,
|
||||
Result, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord, RuntimeWorkspaceInviteLinkRecord,
|
||||
auth_challenge, dto::RuntimeStateRows, invite_link, magic_link_otp, verification_token,
|
||||
};
|
||||
|
||||
pub(super) struct RuntimeStateStore {
|
||||
@@ -121,18 +120,4 @@ impl RuntimeStateStore {
|
||||
pub(super) async fn revoke_workspace_invite_link(&self, workspace_id: String) -> Result<bool> {
|
||||
invite_link::revoke(&self.rows, workspace_id).await
|
||||
}
|
||||
|
||||
pub(super) async fn create_byok_local_lease(
|
||||
&self,
|
||||
active_key: String,
|
||||
lease_id: String,
|
||||
payload: serde_json::Value,
|
||||
ttl_ms: i64,
|
||||
) -> Result<RuntimeByokLocalLeaseRecord> {
|
||||
byok_local_lease::create(&self.rows, active_key, lease_id, payload, ttl_ms).await
|
||||
}
|
||||
|
||||
pub(super) async fn get_byok_local_lease(&self, lease_id: String) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
|
||||
byok_local_lease::get(&self.rows, lease_id).await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -97,11 +97,14 @@ async fn runtime_from_database_url() -> AnyResult<Option<BackendRuntime>> {
|
||||
.context("cleanup invite abuse subjects for backend runtime tests")?;
|
||||
|
||||
Ok(Some(BackendRuntime {
|
||||
config: std::sync::RwLock::new(BackendRuntimeConfig {
|
||||
config: std::sync::RwLock::new(std::sync::Arc::new(BackendRuntimeConfig {
|
||||
database_url,
|
||||
invite_quota: Default::default(),
|
||||
}),
|
||||
private_key: std::sync::Arc::new(zeroize::Zeroizing::new("test-private-key".to_string())),
|
||||
copilot: Default::default(),
|
||||
})),
|
||||
pool: Mutex::new(Some(pool)),
|
||||
managed_token_providers: Default::default(),
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -248,6 +251,7 @@ async fn runtime_gate_sql_semantics_are_atomic_and_ttl_bound() {
|
||||
let runtime = BackendRuntime {
|
||||
config: std::sync::RwLock::new(runtime.config().unwrap()),
|
||||
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
|
||||
managed_token_providers: Default::default(),
|
||||
};
|
||||
tasks.push(tokio::spawn(async move {
|
||||
runtime
|
||||
@@ -579,6 +583,7 @@ async fn coordination_lease_sql_semantics_are_fenced_and_ttl_bound() {
|
||||
let runtime = BackendRuntime {
|
||||
config: std::sync::RwLock::new(runtime.config().unwrap()),
|
||||
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
|
||||
managed_token_providers: Default::default(),
|
||||
};
|
||||
tasks.push(tokio::spawn(async move {
|
||||
runtime
|
||||
@@ -783,6 +788,7 @@ async fn verification_token_sql_state_machine_handles_keep_verify_and_cleanup()
|
||||
let runtime = BackendRuntime {
|
||||
config: std::sync::RwLock::new(runtime.config().unwrap()),
|
||||
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
|
||||
managed_token_providers: Default::default(),
|
||||
};
|
||||
let token = concurrent_token.clone();
|
||||
tasks.push(tokio::spawn(async move {
|
||||
|
||||
@@ -2,18 +2,71 @@ use std::{
|
||||
collections::BTreeMap,
|
||||
env, fs,
|
||||
path::{Path, PathBuf},
|
||||
sync::Arc,
|
||||
};
|
||||
|
||||
use llm_adapter::capability::provider_default_capability_upper_bound;
|
||||
use serde::Deserialize;
|
||||
use serde_json::Map;
|
||||
use sqlx::{PgPool, Row};
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
use super::{RuntimeError, RuntimeResult};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct BackendRuntimeConfig {
|
||||
pub(crate) database_url: String,
|
||||
pub(crate) invite_quota: InviteQuotaConfig,
|
||||
pub(crate) private_key: Arc<Zeroizing<String>>,
|
||||
pub(crate) copilot: CopilotRuntimeConfig,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", default)]
|
||||
pub(crate) struct CopilotRuntimeConfig {
|
||||
pub(crate) enabled: bool,
|
||||
pub(crate) byok: CopilotByokRuntimeConfig,
|
||||
pub(crate) providers: CopilotProvidersRuntimeConfig,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", default)]
|
||||
pub(crate) struct CopilotByokRuntimeConfig {
|
||||
pub(crate) enabled: bool,
|
||||
pub(crate) allow_custom_endpoint: bool,
|
||||
pub(crate) allow_private_endpoint: bool,
|
||||
}
|
||||
|
||||
impl Default for CopilotByokRuntimeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true,
|
||||
allow_custom_endpoint: false,
|
||||
allow_private_endpoint: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", default)]
|
||||
pub(crate) struct CopilotProvidersRuntimeConfig {
|
||||
pub(crate) profiles: Vec<CopilotManagedProfileConfig>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct CopilotManagedProfileConfig {
|
||||
pub(crate) id: String,
|
||||
#[serde(rename = "type")]
|
||||
pub(crate) provider: String,
|
||||
#[serde(default = "enabled_by_default")]
|
||||
pub(crate) enabled: bool,
|
||||
#[serde(default)]
|
||||
pub(crate) models: Vec<String>,
|
||||
pub(crate) config: serde_json::Value,
|
||||
}
|
||||
|
||||
fn enabled_by_default() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -45,35 +98,114 @@ impl Default for InviteQuotaConfig {
|
||||
}
|
||||
|
||||
impl BackendRuntimeConfig {
|
||||
pub(crate) fn from_config_files() -> RuntimeResult<Self> {
|
||||
pub(crate) fn from_config_files(private_key: Option<String>) -> RuntimeResult<Self> {
|
||||
let app_config = app_config_from_config_files()?;
|
||||
let database_url = database_url_from_env()
|
||||
.or(app_config.database_url())
|
||||
.unwrap_or_else(|| "postgresql://localhost:5432/affine".to_string());
|
||||
Ok(Self {
|
||||
Self {
|
||||
database_url,
|
||||
invite_quota: app_config.invite_quota_config(),
|
||||
})
|
||||
private_key: Arc::new(Zeroizing::new(
|
||||
private_key
|
||||
.filter(|key| !key.trim().is_empty())
|
||||
.or_else(private_key_from_env)
|
||||
.or_else(|| app_config.crypto.as_ref().and_then(|crypto| crypto.private_key.clone()))
|
||||
.unwrap_or_default(),
|
||||
)),
|
||||
copilot: app_config.copilot.unwrap_or_default(),
|
||||
}
|
||||
.validated()
|
||||
}
|
||||
|
||||
pub(crate) async fn with_db_overrides(&self, pool: &PgPool) -> RuntimeResult<Self> {
|
||||
let mut app_config = app_config_from_config_files()?;
|
||||
app_config.apply_file_config(load_app_config_overrides_from_db(pool).await?);
|
||||
Ok(Self {
|
||||
let app_config_value = app_config_value_from_config_files()?;
|
||||
let db_overrides = load_app_config_overrides_from_db(pool).await?;
|
||||
self.apply_db_overrides(app_config_value, db_overrides)
|
||||
}
|
||||
|
||||
fn apply_db_overrides(
|
||||
&self,
|
||||
mut app_config_value: serde_json::Value,
|
||||
db_overrides: serde_json::Value,
|
||||
) -> RuntimeResult<Self> {
|
||||
let db_private_key = db_overrides
|
||||
.pointer("/crypto/privateKey")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::to_string)
|
||||
.and_then(non_empty_string);
|
||||
merge_config_value(&mut app_config_value, db_overrides);
|
||||
let app_config = deserialize_app_config(app_config_value)?;
|
||||
Self {
|
||||
// The DB override is loaded after this connection already exists, so it
|
||||
// must not rewrite the active datasource URL.
|
||||
database_url: self.database_url.clone(),
|
||||
invite_quota: app_config.invite_quota_config(),
|
||||
})
|
||||
private_key: db_private_key
|
||||
.map(|key| Arc::new(Zeroizing::new(key)))
|
||||
.unwrap_or_else(|| Arc::clone(&self.private_key)),
|
||||
copilot: app_config.copilot.unwrap_or_else(|| self.copilot.clone()),
|
||||
}
|
||||
.validated()
|
||||
}
|
||||
|
||||
fn validated(self) -> RuntimeResult<Self> {
|
||||
if self.copilot.enabled && self.copilot.byok.enabled && self.private_key.is_empty() {
|
||||
return Err(RuntimeError::invalid_state(
|
||||
"stable crypto.privateKey is required when persistent BYOK is enabled",
|
||||
));
|
||||
}
|
||||
validate_copilot_config(&self.copilot)?;
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
struct AppConfigFile {
|
||||
db: Option<DbConfigFile>,
|
||||
fn validate_copilot_config(config: &CopilotRuntimeConfig) -> RuntimeResult<()> {
|
||||
let mut profile_ids = std::collections::HashSet::new();
|
||||
for profile in &config.providers.profiles {
|
||||
if profile.id.trim().is_empty() || !profile_ids.insert(profile.id.as_str()) {
|
||||
return Err(RuntimeError::invalid_state(
|
||||
"managed copilot profile ids must be non-empty and unique",
|
||||
));
|
||||
}
|
||||
if profile.provider.trim().is_empty() {
|
||||
return Err(RuntimeError::invalid_state(
|
||||
"managed copilot profile provider is required",
|
||||
));
|
||||
}
|
||||
if profile.models.is_empty() {
|
||||
return Err(RuntimeError::invalid_state(
|
||||
"managed copilot profile models must be non-empty",
|
||||
));
|
||||
}
|
||||
let mut models = std::collections::HashSet::new();
|
||||
for model in &profile.models {
|
||||
if model.trim().is_empty() || !models.insert(model.as_str()) {
|
||||
return Err(RuntimeError::invalid_state(
|
||||
"managed copilot profile models must be non-empty and unique",
|
||||
));
|
||||
}
|
||||
provider_default_capability_upper_bound(&profile.provider, model)
|
||||
.ok_or_else(|| RuntimeError::invalid_state("managed copilot profile model is unsupported"))?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[derive(Default, Deserialize)]
|
||||
struct AppConfigFile {
|
||||
db: Option<DbConfigFile>,
|
||||
crypto: Option<CryptoConfigFile>,
|
||||
copilot: Option<CopilotRuntimeConfig>,
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct CryptoConfigFile {
|
||||
private_key: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct DbConfigFile {
|
||||
datasource_url: Option<String>,
|
||||
@@ -97,31 +229,67 @@ fn database_url_from_env() -> Option<String> {
|
||||
env::var("DATABASE_URL").ok().and_then(non_empty_string)
|
||||
}
|
||||
|
||||
fn private_key_from_env() -> Option<String> {
|
||||
env::var("AFFINE_PRIVATE_KEY").ok().and_then(non_empty_string)
|
||||
}
|
||||
|
||||
fn non_empty_string(value: String) -> Option<String> {
|
||||
if value.trim().is_empty() { None } else { Some(value) }
|
||||
}
|
||||
|
||||
fn app_config_from_config_files() -> RuntimeResult<AppConfigFile> {
|
||||
let mut merged = AppConfigFile::default();
|
||||
deserialize_app_config(app_config_value_from_config_files()?)
|
||||
}
|
||||
|
||||
fn app_config_value_from_config_files() -> RuntimeResult<serde_json::Value> {
|
||||
let mut merged = serde_json::Value::Object(Map::new());
|
||||
for path in config_json_paths() {
|
||||
if !path.exists() {
|
||||
continue;
|
||||
}
|
||||
let raw = fs::read_to_string(&path).map_err(|err| RuntimeError::io("failed to read config file", err))?;
|
||||
let config: AppConfigFile =
|
||||
serde_json::from_str(&raw).map_err(|err| RuntimeError::json("failed to parse config file", err))?;
|
||||
merged.apply_file_config(config);
|
||||
let value = serde_json::from_str(&raw).map_err(|err| RuntimeError::json("failed to parse config file", err))?;
|
||||
merge_config_value(&mut merged, expand_module_config_paths(value));
|
||||
}
|
||||
|
||||
Ok(merged)
|
||||
}
|
||||
|
||||
impl AppConfigFile {
|
||||
fn apply_file_config(&mut self, config: AppConfigFile) {
|
||||
if config.db.is_some() {
|
||||
self.db = config.db;
|
||||
fn expand_module_config_paths(mut value: serde_json::Value) -> serde_json::Value {
|
||||
if let Some(root) = value.as_object_mut() {
|
||||
for module in root.values_mut().filter_map(serde_json::Value::as_object_mut) {
|
||||
let entries = std::mem::take(module);
|
||||
for (path, value) in entries {
|
||||
insert_flat_override(module, &path, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
value
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn app_config_from_module_json(value: serde_json::Value) -> RuntimeResult<AppConfigFile> {
|
||||
deserialize_app_config(expand_module_config_paths(value))
|
||||
}
|
||||
|
||||
fn deserialize_app_config(value: serde_json::Value) -> RuntimeResult<AppConfigFile> {
|
||||
serde_json::from_value(value).map_err(|err| RuntimeError::json("failed to parse config file", err))
|
||||
}
|
||||
|
||||
fn merge_config_value(base: &mut serde_json::Value, overrides: serde_json::Value) {
|
||||
match (base, overrides) {
|
||||
(serde_json::Value::Object(base), serde_json::Value::Object(overrides)) => {
|
||||
for (key, value) in overrides {
|
||||
if let Some(existing) = base.get_mut(&key) {
|
||||
merge_config_value(existing, value);
|
||||
} else {
|
||||
base.insert(key, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
(base, overrides) => *base = overrides,
|
||||
}
|
||||
}
|
||||
|
||||
fn default_mail_class_mapping() -> BTreeMap<String, String> {
|
||||
@@ -161,40 +329,60 @@ fn default_mail_class_mapping() -> BTreeMap<String, String> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn load_app_config_overrides_from_db(pool: &PgPool) -> RuntimeResult<AppConfigFile> {
|
||||
async fn load_app_config_overrides_from_db(pool: &PgPool) -> RuntimeResult<serde_json::Value> {
|
||||
let rows = match sqlx::query("SELECT id, value FROM app_configs").fetch_all(pool).await {
|
||||
Ok(rows) => rows,
|
||||
Err(sqlx::Error::Database(err)) if err.code().as_deref() == Some("42P01") => return Ok(AppConfigFile::default()),
|
||||
Err(sqlx::Error::Database(err)) if err.code().as_deref() == Some("42P01") => {
|
||||
return Ok(serde_json::Value::Object(Map::new()));
|
||||
}
|
||||
Err(err) => return Err(RuntimeError::database("failed to load app config overrides", err)),
|
||||
};
|
||||
|
||||
app_config_from_flat_overrides(rows.into_iter().map(|row| {
|
||||
Ok(app_config_value_from_flat_overrides(rows.into_iter().map(|row| {
|
||||
let id: String = row.get("id");
|
||||
let value: serde_json::Value = row.get("value");
|
||||
(id, value)
|
||||
}))
|
||||
})))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn app_config_from_flat_overrides<I, S>(rows: I) -> RuntimeResult<AppConfigFile>
|
||||
where
|
||||
I: IntoIterator<Item = (S, serde_json::Value)>,
|
||||
S: AsRef<str>,
|
||||
{
|
||||
deserialize_app_config(app_config_value_from_flat_overrides(rows))
|
||||
}
|
||||
|
||||
fn app_config_value_from_flat_overrides<I, S>(rows: I) -> serde_json::Value
|
||||
where
|
||||
I: IntoIterator<Item = (S, serde_json::Value)>,
|
||||
S: AsRef<str>,
|
||||
{
|
||||
let mut root = Map::new();
|
||||
for (path, value) in rows {
|
||||
let Some((module, key)) = path.as_ref().split_once('.') else {
|
||||
continue;
|
||||
};
|
||||
root
|
||||
.entry(module.to_string())
|
||||
.or_insert_with(|| serde_json::Value::Object(Map::new()));
|
||||
if let Some(serde_json::Value::Object(module_object)) = root.get_mut(module) {
|
||||
module_object.insert(key.to_string(), value);
|
||||
}
|
||||
insert_flat_override(&mut root, path.as_ref(), value);
|
||||
}
|
||||
|
||||
serde_json::from_value(serde_json::Value::Object(root))
|
||||
.map_err(|err| RuntimeError::json("invalid app config overrides", err))
|
||||
serde_json::Value::Object(root)
|
||||
}
|
||||
|
||||
fn insert_flat_override(root: &mut Map<String, serde_json::Value>, path: &str, value: serde_json::Value) {
|
||||
let mut parts = path.split('.').peekable();
|
||||
let mut current = root;
|
||||
while let Some(part) = parts.next() {
|
||||
if parts.peek().is_none() {
|
||||
current.insert(part.to_string(), value);
|
||||
return;
|
||||
}
|
||||
let entry = current
|
||||
.entry(part.to_string())
|
||||
.or_insert_with(|| serde_json::Value::Object(Map::new()));
|
||||
if !entry.is_object() {
|
||||
*entry = serde_json::Value::Object(Map::new());
|
||||
}
|
||||
current = entry.as_object_mut().expect("override node must be an object");
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn config_json_paths() -> Vec<PathBuf> {
|
||||
@@ -273,6 +461,79 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expands_module_config_paths_from_json_files() {
|
||||
let app_config = app_config_from_module_json(serde_json::json!({
|
||||
"copilot": {
|
||||
"enabled": true,
|
||||
"byok.enabled": false,
|
||||
"providers.profiles": [{
|
||||
"id": "managed-openai",
|
||||
"type": "openai",
|
||||
"models": ["gpt-5.6-luna"],
|
||||
"config": { "apiKey": "test" }
|
||||
}]
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
let copilot = app_config.copilot.unwrap();
|
||||
|
||||
assert!(copilot.enabled);
|
||||
assert!(!copilot.byok.enabled);
|
||||
assert_eq!(copilot.providers.profiles.len(), 1);
|
||||
assert_eq!(copilot.providers.profiles[0].id, "managed-openai");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_database_config_preserves_file_config_siblings() {
|
||||
let mut file_config = expand_module_config_paths(serde_json::json!({
|
||||
"copilot": {
|
||||
"enabled": true,
|
||||
"byok": { "enabled": true, "allowCustomEndpoint": true },
|
||||
"providers": {
|
||||
"profiles": [{
|
||||
"id": "managed-openai",
|
||||
"type": "openai",
|
||||
"models": ["gpt-5.6-luna"],
|
||||
"config": { "apiKey": "test" }
|
||||
}]
|
||||
}
|
||||
}
|
||||
}));
|
||||
let database_config = app_config_value_from_flat_overrides([("copilot.byok.enabled", serde_json::json!(false))]);
|
||||
|
||||
merge_config_value(&mut file_config, database_config);
|
||||
let copilot = deserialize_app_config(file_config).unwrap().copilot.unwrap();
|
||||
|
||||
assert!(copilot.enabled);
|
||||
assert!(!copilot.byok.enabled);
|
||||
assert!(copilot.byok.allow_custom_endpoint);
|
||||
assert_eq!(copilot.providers.profiles.len(), 1);
|
||||
assert_eq!(copilot.providers.profiles[0].id, "managed-openai");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn database_config_only_replaces_an_active_private_key_explicitly() {
|
||||
let active = BackendRuntimeConfig {
|
||||
database_url: "postgresql://active".to_string(),
|
||||
invite_quota: InviteQuotaConfig::default(),
|
||||
private_key: Arc::new(Zeroizing::new("active-private-key".to_string())),
|
||||
copilot: CopilotRuntimeConfig::default(),
|
||||
};
|
||||
let empty = serde_json::Value::Object(Map::new());
|
||||
|
||||
let unchanged = active.apply_db_overrides(empty.clone(), empty.clone()).unwrap();
|
||||
assert_eq!(unchanged.private_key.as_str(), "active-private-key");
|
||||
|
||||
let overridden = active
|
||||
.apply_db_overrides(
|
||||
empty,
|
||||
app_config_value_from_flat_overrides([("crypto.privateKey", serde_json::json!("database-private-key"))]),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(overridden.private_key.as_str(), "database-private-key");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invite_quota_policy_is_internal_not_app_configurable() {
|
||||
let app_config = app_config_from_flat_overrides([
|
||||
|
||||
@@ -6,5 +6,5 @@ pub(crate) mod error;
|
||||
pub(crate) mod migrations;
|
||||
pub(crate) mod types;
|
||||
|
||||
pub(crate) use config::{BackendRuntimeConfig, InviteQuotaConfig};
|
||||
pub(crate) use config::{BackendRuntimeConfig, CopilotManagedProfileConfig, CopilotRuntimeConfig, InviteQuotaConfig};
|
||||
pub(crate) use error::{RuntimeError, RuntimeResult, napi_error, to_napi_error};
|
||||
|
||||
@@ -138,13 +138,6 @@ pub struct RuntimeWorkspaceInviteLinkRecord {
|
||||
pub expires_at_ms: i64,
|
||||
}
|
||||
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct RuntimeByokLocalLeaseRecord {
|
||||
pub lease_id: String,
|
||||
pub payload: serde_json::Value,
|
||||
pub expires_at_ms: i64,
|
||||
}
|
||||
|
||||
#[napi_derive::napi(object)]
|
||||
pub struct RuntimeDocHistoryInput {
|
||||
pub workspace_id: String,
|
||||
|
||||
Reference in New Issue
Block a user