mirror of
https://github.com/whyour/qinglong.git
synced 2026-09-21 09:58:46 +08:00
feat(ql3): establish 3.0 incubation baseline
This commit is contained in:
whyour
@@ -0,0 +1,656 @@
|
||||
import {
|
||||
transitionStepRunMutation,
|
||||
type StepRunStatus,
|
||||
} from '@qinglong/runtime-core/step-run';
|
||||
|
||||
import type {
|
||||
ModelInvocationAuditDisposition,
|
||||
ModelInvocationAuditRecord,
|
||||
ModelInvocationAuditSink,
|
||||
ModelUsage,
|
||||
} from '../model-gateway/model';
|
||||
import {
|
||||
MAX_MODEL_INVOCATION_RECOVERY_PAGE_SIZE,
|
||||
ModelInvocationConflictError,
|
||||
ModelInvocationRepositoryUnavailableError,
|
||||
createModelInvocationCompletionCommand,
|
||||
createModelInvocationMutationIdentity,
|
||||
createModelInvocationStartCommand,
|
||||
normalizeModelInvocationCompletionRecord,
|
||||
normalizeModelInvocationStartRecord,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationRepository,
|
||||
type ModelInvocationStartRecord,
|
||||
} from './modelInvocation';
|
||||
import {
|
||||
isQuotaAwareModelInvocationRepository,
|
||||
normalizeModelInvocationQuotaAdmission,
|
||||
type ModelInvocationQuotaAdmission,
|
||||
} from '../usage/usageQuota';
|
||||
import {
|
||||
isPricingAwareModelInvocationRepository,
|
||||
normalizeModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceQuote,
|
||||
} from '../pricing/pricing';
|
||||
import {
|
||||
isPluginPackagePromptOutputCompletionRepository,
|
||||
type PluginPackagePromptOutputCompletionRepository,
|
||||
} from '../prompt-output/pluginPackagePromptOutputCompletion';
|
||||
import {
|
||||
PluginPackagePromptOutputArtifactConflictError,
|
||||
normalizePluginPackagePromptOutputArtifact,
|
||||
pluginPackagePromptOutputArtifactReference,
|
||||
type PluginPackagePromptOutputArtifact,
|
||||
type PluginPackagePromptOutputArtifactReference,
|
||||
} from '../prompt-output/pluginPackagePromptOutputArtifact';
|
||||
|
||||
const MAX_COORDINATOR_ATTEMPTS = 3;
|
||||
|
||||
interface CompletionTransition {
|
||||
readonly to: StepRunStatus;
|
||||
readonly outputRef?: string;
|
||||
readonly resultCode?: string;
|
||||
readonly errorSummary?: string;
|
||||
}
|
||||
|
||||
export interface ModelInvocationRecoverySummary {
|
||||
readonly observedAtMs: number;
|
||||
readonly scanned: number;
|
||||
readonly recovered: number;
|
||||
readonly alreadyCompleted: number;
|
||||
readonly failed: number;
|
||||
readonly hasMore: boolean;
|
||||
}
|
||||
|
||||
function sameUsage(
|
||||
left: Readonly<ModelUsage> | null,
|
||||
right: Readonly<ModelUsage> | null,
|
||||
): boolean {
|
||||
return JSON.stringify(left) === JSON.stringify(right);
|
||||
}
|
||||
|
||||
function assertStartMatchesAudit(
|
||||
startValue: ModelInvocationStartRecord,
|
||||
audit: Readonly<ModelInvocationAuditRecord>,
|
||||
): Readonly<ModelInvocationStartRecord> {
|
||||
const start = normalizeModelInvocationStartRecord(startValue);
|
||||
if (
|
||||
start.invocationId !== audit.requestId ||
|
||||
start.projectId !== audit.projectId ||
|
||||
start.runId !== audit.runId ||
|
||||
start.stepRunId !== audit.stepRunId ||
|
||||
start.traceId !== audit.traceId ||
|
||||
start.provider !== audit.provider ||
|
||||
start.model !== audit.model ||
|
||||
start.policyRevision !== audit.policyRevision ||
|
||||
start.requestDigest !== audit.requestDigest ||
|
||||
start.inputBytes !== audit.inputBytes ||
|
||||
start.maxOutputTokens !== audit.maxOutputTokens ||
|
||||
start.deadlineAtMs !== audit.deadlineAtMs
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return start;
|
||||
}
|
||||
|
||||
function completionTransition(
|
||||
audit: Readonly<ModelInvocationAuditRecord>,
|
||||
successOutputRef?: string,
|
||||
): Readonly<CompletionTransition> {
|
||||
if (audit.phase === 'completed') {
|
||||
return Object.freeze({
|
||||
to: 'succeeded',
|
||||
outputRef: successOutputRef ?? `model-invocation:${audit.requestId}`,
|
||||
});
|
||||
}
|
||||
if (audit.errorCode === 'MODEL_INVOCATION_DEADLINE_EXCEEDED') {
|
||||
return Object.freeze({
|
||||
to: 'timed_out',
|
||||
resultCode: 'model_deadline_exceeded',
|
||||
errorSummary: 'Model invocation deadline exceeded',
|
||||
});
|
||||
}
|
||||
if (
|
||||
audit.errorCode === 'MODEL_INVOCATION_ABORTED' ||
|
||||
audit.errorCode === 'MODEL_STREAM_CANCELLED' ||
|
||||
audit.errorCode === 'MODEL_INVOCATION_OUTCOME_UNKNOWN'
|
||||
) {
|
||||
return Object.freeze({
|
||||
to: 'lost',
|
||||
resultCode: 'model_outcome_unknown',
|
||||
errorSummary: 'Model invocation outcome is unknown',
|
||||
});
|
||||
}
|
||||
return Object.freeze({
|
||||
to: 'failed',
|
||||
resultCode: 'model_provider_failed',
|
||||
errorSummary: 'Model invocation failed',
|
||||
});
|
||||
}
|
||||
|
||||
function expectedOutcome(
|
||||
transition: Readonly<CompletionTransition>,
|
||||
): ModelInvocationCompletionRecord['outcome'] {
|
||||
if (transition.to === 'succeeded') return 'succeeded';
|
||||
if (transition.to === 'timed_out') return 'timed_out';
|
||||
if (transition.to === 'lost') return 'outcome_unknown';
|
||||
return 'failed';
|
||||
}
|
||||
|
||||
function assertCompletionMatchesAudit(
|
||||
completionValue: ModelInvocationCompletionRecord,
|
||||
start: Readonly<ModelInvocationStartRecord>,
|
||||
audit: Readonly<ModelInvocationAuditRecord>,
|
||||
): Readonly<ModelInvocationCompletionRecord> {
|
||||
const completion = normalizeModelInvocationCompletionRecord(completionValue);
|
||||
const transition = completionTransition(audit);
|
||||
if (
|
||||
completion.invocationId !== start.invocationId ||
|
||||
completion.projectId !== start.projectId ||
|
||||
completion.runId !== start.runId ||
|
||||
completion.stepRunId !== start.stepRunId ||
|
||||
completion.traceId !== start.traceId ||
|
||||
completion.startDigest !== start.startDigest ||
|
||||
completion.outcome !== expectedOutcome(transition) ||
|
||||
completion.outputBytes !== audit.outputBytes ||
|
||||
!sameUsage(completion.usage, audit.usage) ||
|
||||
completion.errorCode !== audit.errorCode
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return completion;
|
||||
}
|
||||
|
||||
function identity(record: Readonly<ModelInvocationAuditRecord>): Readonly<{
|
||||
projectId: string;
|
||||
runId: string;
|
||||
stepRunId: string;
|
||||
}> {
|
||||
return Object.freeze({
|
||||
projectId: record.projectId,
|
||||
runId: record.runId,
|
||||
stepRunId: record.stepRunId,
|
||||
});
|
||||
}
|
||||
|
||||
export class DurableModelInvocationCoordinator
|
||||
implements ModelInvocationAuditSink
|
||||
{
|
||||
constructor(private readonly repository: ModelInvocationRepository) {
|
||||
if (
|
||||
!repository ||
|
||||
typeof repository.findStart !== 'function' ||
|
||||
typeof repository.findCompletion !== 'function' ||
|
||||
typeof repository.readAuthority !== 'function' ||
|
||||
typeof repository.admit !== 'function' ||
|
||||
typeof repository.complete !== 'function'
|
||||
) {
|
||||
throw new ModelInvocationRepositoryUnavailableError();
|
||||
}
|
||||
}
|
||||
|
||||
async record(
|
||||
record: Readonly<ModelInvocationAuditRecord>,
|
||||
): Promise<Readonly<ModelInvocationAuditDisposition>> {
|
||||
return record.phase === 'admitted'
|
||||
? this.#admit(record)
|
||||
: this.#complete(record);
|
||||
}
|
||||
|
||||
async recordWithQuota(
|
||||
record: Readonly<ModelInvocationAuditRecord>,
|
||||
admissionValue: Readonly<ModelInvocationQuotaAdmission>,
|
||||
): Promise<Readonly<ModelInvocationAuditDisposition>> {
|
||||
const admission = normalizeModelInvocationQuotaAdmission(admissionValue);
|
||||
if (
|
||||
record.phase !== 'admitted' ||
|
||||
record.requestId !== admission.invocationId ||
|
||||
record.projectId !== admission.projectId ||
|
||||
record.policyRevision !== admission.modelPolicyRevision ||
|
||||
!isQuotaAwareModelInvocationRepository(this.repository)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return this.#admit(record, admission);
|
||||
}
|
||||
|
||||
async recordWithPricing(
|
||||
record: Readonly<ModelInvocationAuditRecord>,
|
||||
quoteValue: Readonly<ModelInvocationPriceQuote>,
|
||||
admissionValue?: Readonly<ModelInvocationQuotaAdmission>,
|
||||
): Promise<Readonly<ModelInvocationAuditDisposition>> {
|
||||
const quote = normalizeModelInvocationPriceQuote(quoteValue);
|
||||
const admission =
|
||||
admissionValue === undefined
|
||||
? undefined
|
||||
: normalizeModelInvocationQuotaAdmission(admissionValue);
|
||||
if (
|
||||
record.phase !== 'admitted' ||
|
||||
record.requestId !== quote.invocationId ||
|
||||
record.projectId !== quote.projectId ||
|
||||
record.policyRevision !== quote.modelPolicyRevision ||
|
||||
record.provider !== quote.provider ||
|
||||
record.model !== quote.model ||
|
||||
(admission !== undefined &&
|
||||
(admission.invocationId !== quote.invocationId ||
|
||||
admission.projectId !== quote.projectId ||
|
||||
admission.modelPolicyRevision !== quote.modelPolicyRevision)) ||
|
||||
!isPricingAwareModelInvocationRepository(this.repository)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return this.#admit(record, admission, quote);
|
||||
}
|
||||
|
||||
async recordWithPromptOutputArtifact(
|
||||
record: Readonly<ModelInvocationAuditRecord>,
|
||||
artifactValue: PluginPackagePromptOutputArtifact,
|
||||
): Promise<
|
||||
Readonly<{
|
||||
status: 'created' | 'existing';
|
||||
reference: Readonly<PluginPackagePromptOutputArtifactReference>;
|
||||
}>
|
||||
> {
|
||||
const artifact = normalizePluginPackagePromptOutputArtifact(artifactValue);
|
||||
if (
|
||||
record.phase !== 'completed' ||
|
||||
record.requestId !== artifact.invocationId ||
|
||||
record.projectId !== artifact.projectId ||
|
||||
record.runId !== artifact.runId ||
|
||||
record.stepRunId !== artifact.stepRunId ||
|
||||
record.provider !== artifact.provider ||
|
||||
record.model !== artifact.model ||
|
||||
record.outputBytes !== artifact.outputBytes ||
|
||||
!isPluginPackagePromptOutputCompletionRepository(this.repository)
|
||||
) {
|
||||
throw new PluginPackagePromptOutputArtifactConflictError();
|
||||
}
|
||||
const result = await this.#complete(record, artifact);
|
||||
if (!result.reference) {
|
||||
throw new PluginPackagePromptOutputArtifactConflictError();
|
||||
}
|
||||
return Object.freeze({
|
||||
status: result.status,
|
||||
reference: result.reference,
|
||||
});
|
||||
}
|
||||
|
||||
async #admit(
|
||||
audit: Readonly<ModelInvocationAuditRecord>,
|
||||
admission?: Readonly<ModelInvocationQuotaAdmission>,
|
||||
quote?: Readonly<ModelInvocationPriceQuote>,
|
||||
): Promise<Readonly<ModelInvocationAuditDisposition>> {
|
||||
const existing = await this.repository.findStart(audit.requestId);
|
||||
if (existing) {
|
||||
assertStartMatchesAudit(existing, audit);
|
||||
if (admission) {
|
||||
if (!isQuotaAwareModelInvocationRepository(this.repository)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const reservation = await this.repository.findQuotaReservation(
|
||||
audit.requestId,
|
||||
);
|
||||
if (
|
||||
!reservation ||
|
||||
reservation.admissionDigest !== admission.admissionDigest
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
}
|
||||
if (quote) {
|
||||
if (!isPricingAwareModelInvocationRepository(this.repository)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const storedQuote = await this.repository.findPriceQuote(
|
||||
audit.requestId,
|
||||
);
|
||||
if (!storedQuote || storedQuote.quoteDigest !== quote.quoteDigest) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
}
|
||||
return Object.freeze({ status: 'existing' });
|
||||
}
|
||||
for (let attempt = 0; attempt < MAX_COORDINATOR_ATTEMPTS; attempt += 1) {
|
||||
const authority = await this.repository.readAuthority(identity(audit));
|
||||
if (
|
||||
!authority ||
|
||||
authority.stepRun.status !== 'ready' ||
|
||||
authority.stepRun.kind !== 'model'
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const mutationIdentity = createModelInvocationMutationIdentity(
|
||||
audit.requestId,
|
||||
'start',
|
||||
);
|
||||
const command = createModelInvocationStartCommand(
|
||||
audit,
|
||||
transitionStepRunMutation(
|
||||
authority.stepRun,
|
||||
{
|
||||
expectedVersion: authority.stepRun.version,
|
||||
expectedDigest: authority.stepRun.stepRunDigest,
|
||||
mutationId: mutationIdentity.mutationId,
|
||||
to: 'running',
|
||||
atMs: audit.occurredAtMs,
|
||||
},
|
||||
{
|
||||
expectedRunVersion: authority.runVersion,
|
||||
expectedRunEventSequence: authority.runEventSequence,
|
||||
eventId: mutationIdentity.eventId,
|
||||
dedupeKey: mutationIdentity.dedupeKey,
|
||||
actor: { type: 'executor', id: 'model-gateway' },
|
||||
},
|
||||
),
|
||||
);
|
||||
try {
|
||||
const result = quote
|
||||
? await (() => {
|
||||
if (!isPricingAwareModelInvocationRepository(this.repository)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return this.repository.admitWithPricing(
|
||||
command,
|
||||
quote,
|
||||
admission,
|
||||
);
|
||||
})()
|
||||
: admission
|
||||
? await (() => {
|
||||
if (!isQuotaAwareModelInvocationRepository(this.repository)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return this.repository.admitWithQuota(command, admission);
|
||||
})()
|
||||
: await this.repository.admit(command);
|
||||
return Object.freeze({ status: result.status });
|
||||
} catch (error) {
|
||||
const stored = await this.#startAfterFailure(
|
||||
audit,
|
||||
error,
|
||||
admission,
|
||||
quote,
|
||||
);
|
||||
if (stored) return Object.freeze({ status: 'existing' });
|
||||
if (
|
||||
!(error instanceof ModelInvocationConflictError) ||
|
||||
attempt + 1 >= MAX_COORDINATOR_ATTEMPTS
|
||||
) {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
|
||||
async #complete(
|
||||
audit: Readonly<ModelInvocationAuditRecord>,
|
||||
artifactValue?: Readonly<PluginPackagePromptOutputArtifact>,
|
||||
): Promise<
|
||||
Readonly<
|
||||
ModelInvocationAuditDisposition & {
|
||||
reference?: Readonly<PluginPackagePromptOutputArtifactReference>;
|
||||
}
|
||||
>
|
||||
> {
|
||||
const artifact = artifactValue
|
||||
? normalizePluginPackagePromptOutputArtifact(artifactValue)
|
||||
: undefined;
|
||||
const artifactRepository = artifact
|
||||
? (this.repository as ModelInvocationRepository &
|
||||
PluginPackagePromptOutputCompletionRepository)
|
||||
: undefined;
|
||||
const startValue = await this.repository.findStart(audit.requestId);
|
||||
if (!startValue) throw new ModelInvocationConflictError();
|
||||
const start = assertStartMatchesAudit(startValue, audit);
|
||||
const existing = await this.repository.findCompletion(audit.requestId);
|
||||
if (existing) {
|
||||
assertCompletionMatchesAudit(existing, start, audit);
|
||||
if (artifact && artifactRepository) {
|
||||
const stored = await artifactRepository.findPromptOutputArtifact(
|
||||
artifact.artifactId,
|
||||
);
|
||||
if (!stored || JSON.stringify(stored) !== JSON.stringify(artifact)) {
|
||||
throw new PluginPackagePromptOutputArtifactConflictError();
|
||||
}
|
||||
return Object.freeze({
|
||||
status: 'existing' as const,
|
||||
reference: pluginPackagePromptOutputArtifactReference(stored),
|
||||
});
|
||||
}
|
||||
return Object.freeze({ status: 'existing' as const });
|
||||
}
|
||||
for (let attempt = 0; attempt < MAX_COORDINATOR_ATTEMPTS; attempt += 1) {
|
||||
const authority = await this.repository.readAuthority(identity(audit));
|
||||
if (
|
||||
!authority ||
|
||||
authority.stepRun.status !== 'running' ||
|
||||
authority.stepRun.version !== start.startedStepRunVersion ||
|
||||
authority.stepRun.stepRunDigest !== start.startedStepRunDigest
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const transition = completionTransition(audit, artifact?.artifactId);
|
||||
const mutationIdentity = createModelInvocationMutationIdentity(
|
||||
audit.requestId,
|
||||
'completion',
|
||||
);
|
||||
const command = createModelInvocationCompletionCommand(
|
||||
start,
|
||||
audit,
|
||||
transitionStepRunMutation(
|
||||
authority.stepRun,
|
||||
{
|
||||
expectedVersion: authority.stepRun.version,
|
||||
expectedDigest: authority.stepRun.stepRunDigest,
|
||||
mutationId: mutationIdentity.mutationId,
|
||||
to: transition.to,
|
||||
atMs: audit.occurredAtMs,
|
||||
...(transition.outputRef === undefined
|
||||
? {}
|
||||
: { outputRef: transition.outputRef }),
|
||||
...(transition.resultCode === undefined
|
||||
? {}
|
||||
: { resultCode: transition.resultCode }),
|
||||
...(transition.errorSummary === undefined
|
||||
? {}
|
||||
: { errorSummary: transition.errorSummary }),
|
||||
},
|
||||
{
|
||||
expectedRunVersion: authority.runVersion,
|
||||
expectedRunEventSequence: authority.runEventSequence,
|
||||
eventId: mutationIdentity.eventId,
|
||||
dedupeKey: mutationIdentity.dedupeKey,
|
||||
actor: { type: 'executor', id: 'model-gateway' },
|
||||
},
|
||||
),
|
||||
artifact?.artifactId,
|
||||
);
|
||||
try {
|
||||
const pricingAware = isPricingAwareModelInvocationRepository(
|
||||
this.repository,
|
||||
);
|
||||
const quote = pricingAware
|
||||
? await this.repository.findPriceQuote(audit.requestId)
|
||||
: null;
|
||||
const quotaAware = isQuotaAwareModelInvocationRepository(
|
||||
this.repository,
|
||||
);
|
||||
const reservation = quotaAware
|
||||
? await this.repository.findQuotaReservation(audit.requestId)
|
||||
: null;
|
||||
const result =
|
||||
artifactRepository && artifact
|
||||
? await artifactRepository.completeWithPromptOutputArtifact(
|
||||
command,
|
||||
artifact,
|
||||
)
|
||||
: pricingAware && quote
|
||||
? await this.repository.completeWithPricing(command)
|
||||
: quotaAware && reservation
|
||||
? await this.repository.completeWithQuota(command)
|
||||
: await this.repository.complete(command);
|
||||
return Object.freeze({
|
||||
status: result.status,
|
||||
...(artifact && artifactRepository
|
||||
? {
|
||||
reference: pluginPackagePromptOutputArtifactReference(artifact),
|
||||
}
|
||||
: {}),
|
||||
});
|
||||
} catch (error) {
|
||||
const stored = await this.#completionAfterFailure(
|
||||
start,
|
||||
audit,
|
||||
error,
|
||||
artifact,
|
||||
artifactRepository,
|
||||
);
|
||||
if (stored) {
|
||||
return Object.freeze({
|
||||
status: 'existing' as const,
|
||||
...(artifact
|
||||
? {
|
||||
reference:
|
||||
pluginPackagePromptOutputArtifactReference(artifact),
|
||||
}
|
||||
: {}),
|
||||
});
|
||||
}
|
||||
if (
|
||||
!(error instanceof ModelInvocationConflictError) ||
|
||||
attempt + 1 >= MAX_COORDINATOR_ATTEMPTS
|
||||
) {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
|
||||
async #startAfterFailure(
|
||||
audit: Readonly<ModelInvocationAuditRecord>,
|
||||
original: unknown,
|
||||
admission?: Readonly<ModelInvocationQuotaAdmission>,
|
||||
quote?: Readonly<ModelInvocationPriceQuote>,
|
||||
): Promise<Readonly<ModelInvocationStartRecord> | null> {
|
||||
try {
|
||||
const stored = await this.repository.findStart(audit.requestId);
|
||||
if (!stored) return null;
|
||||
const start = assertStartMatchesAudit(stored, audit);
|
||||
if (admission) {
|
||||
if (!isQuotaAwareModelInvocationRepository(this.repository)) {
|
||||
throw original;
|
||||
}
|
||||
const reservation = await this.repository.findQuotaReservation(
|
||||
audit.requestId,
|
||||
);
|
||||
if (
|
||||
!reservation ||
|
||||
reservation.admissionDigest !== admission.admissionDigest
|
||||
) {
|
||||
throw original;
|
||||
}
|
||||
}
|
||||
if (quote) {
|
||||
if (!isPricingAwareModelInvocationRepository(this.repository)) {
|
||||
throw original;
|
||||
}
|
||||
const storedQuote = await this.repository.findPriceQuote(
|
||||
audit.requestId,
|
||||
);
|
||||
if (!storedQuote || storedQuote.quoteDigest !== quote.quoteDigest) {
|
||||
throw original;
|
||||
}
|
||||
}
|
||||
return start;
|
||||
} catch {
|
||||
throw original;
|
||||
}
|
||||
}
|
||||
|
||||
async #completionAfterFailure(
|
||||
start: Readonly<ModelInvocationStartRecord>,
|
||||
audit: Readonly<ModelInvocationAuditRecord>,
|
||||
original: unknown,
|
||||
artifact?: Readonly<PluginPackagePromptOutputArtifact>,
|
||||
artifactRepository?: PluginPackagePromptOutputCompletionRepository,
|
||||
): Promise<Readonly<ModelInvocationCompletionRecord> | null> {
|
||||
try {
|
||||
const stored = await this.repository.findCompletion(audit.requestId);
|
||||
if (!stored) return null;
|
||||
const completion = assertCompletionMatchesAudit(stored, start, audit);
|
||||
if (artifact) {
|
||||
if (!artifactRepository) throw original;
|
||||
const storedArtifact =
|
||||
await artifactRepository.findPromptOutputArtifact(
|
||||
artifact.artifactId,
|
||||
);
|
||||
if (
|
||||
!storedArtifact ||
|
||||
JSON.stringify(storedArtifact) !== JSON.stringify(artifact)
|
||||
) {
|
||||
throw original;
|
||||
}
|
||||
}
|
||||
return completion;
|
||||
} catch {
|
||||
throw original;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export class DurableModelInvocationRecovery {
|
||||
constructor(
|
||||
private readonly repository: ModelInvocationRepository,
|
||||
private readonly coordinator = new DurableModelInvocationCoordinator(
|
||||
repository,
|
||||
),
|
||||
) {}
|
||||
|
||||
async recover(
|
||||
limit = MAX_MODEL_INVOCATION_RECOVERY_PAGE_SIZE,
|
||||
): Promise<Readonly<ModelInvocationRecoverySummary>> {
|
||||
const page = await this.repository.listIncomplete(limit);
|
||||
let recovered = 0;
|
||||
let alreadyCompleted = 0;
|
||||
let failed = 0;
|
||||
for (const start of page.candidates) {
|
||||
try {
|
||||
const result = await this.coordinator.record(
|
||||
Object.freeze({
|
||||
phase: 'failed',
|
||||
projectId: start.projectId,
|
||||
runId: start.runId,
|
||||
stepRunId: start.stepRunId,
|
||||
traceId: start.traceId,
|
||||
requestId: start.invocationId,
|
||||
provider: start.provider,
|
||||
model: start.model,
|
||||
policyRevision: start.policyRevision,
|
||||
requestDigest: start.requestDigest,
|
||||
deadlineAtMs: start.deadlineAtMs,
|
||||
inputBytes: start.inputBytes,
|
||||
maxOutputTokens: start.maxOutputTokens,
|
||||
outputBytes: 0,
|
||||
usage: null,
|
||||
errorCode: 'MODEL_INVOCATION_OUTCOME_UNKNOWN',
|
||||
occurredAtMs: page.observedAtMs,
|
||||
}),
|
||||
);
|
||||
if (result.status === 'created') recovered += 1;
|
||||
else alreadyCompleted += 1;
|
||||
} catch {
|
||||
failed += 1;
|
||||
}
|
||||
}
|
||||
return Object.freeze({
|
||||
observedAtMs: page.observedAtMs,
|
||||
scanned: page.candidates.length,
|
||||
recovered,
|
||||
alreadyCompleted,
|
||||
failed,
|
||||
hasMore: page.hasMore,
|
||||
});
|
||||
}
|
||||
}
|
||||
+305
@@ -0,0 +1,305 @@
|
||||
import type { DatabaseSync } from 'node:sqlite';
|
||||
|
||||
import {
|
||||
normalizeModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceQuote,
|
||||
} from '../../pricing/pricing';
|
||||
import {
|
||||
ModelInvocationProjectQuotaExceededError,
|
||||
createModelInvocationQuotaReservation,
|
||||
normalizeModelInvocationQuotaAdmission,
|
||||
type ModelInvocationQuotaAdmission,
|
||||
type ModelInvocationQuotaReservation,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
ModelInvocationConflictError,
|
||||
normalizeModelInvocationStartCommand,
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationStartCommand,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
|
||||
import type { LocalModelInvocationOperationAuthority, Row } from './authority';
|
||||
import {
|
||||
assertLocalFeatureActive,
|
||||
enqueueLocalModelInvocation,
|
||||
integer,
|
||||
unavailable,
|
||||
} from './authority';
|
||||
import { parsePriceQuote, parseQuotaReservation, parseStart } from './codec';
|
||||
import {
|
||||
applyMutation,
|
||||
assertCurrent,
|
||||
insertPriceQuote,
|
||||
insertQuotaReservation,
|
||||
insertStart,
|
||||
quotaWindowUsage,
|
||||
} from './mutations';
|
||||
import { priceQuoteRows, quotaReservationRows, startRows } from './queries';
|
||||
|
||||
export function admitOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
): Promise<Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>> {
|
||||
const command = normalizeModelInvocationStartCommand(commandValue);
|
||||
const start = command.start;
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
let began = false;
|
||||
try {
|
||||
client.exec('BEGIN IMMEDIATE');
|
||||
began = true;
|
||||
const existing = startRows(
|
||||
client,
|
||||
`start.invocation_id = ? OR
|
||||
start.mutation_id = ? OR start.run_event_id = ?`,
|
||||
[start.invocationId, start.stepRunMutationId, start.runEventId],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseStart(existing[0]);
|
||||
if (JSON.stringify(stored) !== JSON.stringify(start)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
assertLocalFeatureActive(client);
|
||||
assertCurrent(client, command.stepRunMutation, start.projectId);
|
||||
applyMutation(client, command.stepRunMutation);
|
||||
insertStart(client, start);
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'created' as const, record: start });
|
||||
} catch (error) {
|
||||
if (began && client.isTransaction) {
|
||||
try {
|
||||
client.exec('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the original transaction failure.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
export function admitWithQuotaOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
admissionValue: ModelInvocationQuotaAdmission,
|
||||
): Promise<Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>> {
|
||||
const command = normalizeModelInvocationStartCommand(commandValue);
|
||||
const admission = normalizeModelInvocationQuotaAdmission(admissionValue);
|
||||
const start = command.start;
|
||||
if (
|
||||
admission.invocationId !== start.invocationId ||
|
||||
admission.projectId !== start.projectId ||
|
||||
admission.modelPolicyRevision !== start.policyRevision
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
let began = false;
|
||||
try {
|
||||
client.exec('BEGIN IMMEDIATE');
|
||||
began = true;
|
||||
const existing = startRows(
|
||||
client,
|
||||
`start.invocation_id = ? OR
|
||||
start.mutation_id = ? OR start.run_event_id = ?`,
|
||||
[start.invocationId, start.stepRunMutationId, start.runEventId],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseStart(existing[0]);
|
||||
const reservations = quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = ?',
|
||||
[start.invocationId],
|
||||
);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(start) ||
|
||||
reservations.length !== 1 ||
|
||||
parseQuotaReservation(reservations[0]!).admissionDigest !==
|
||||
admission.admissionDigest
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
assertLocalFeatureActive(client);
|
||||
const observed = client
|
||||
.prepare(
|
||||
`SELECT CAST(unixepoch('subsec') * 1000 AS INTEGER)
|
||||
AS "observedAtMs"`,
|
||||
)
|
||||
.get() as Row | undefined;
|
||||
if (!observed) throw unavailable();
|
||||
const reservation = createModelInvocationQuotaReservation(
|
||||
admission,
|
||||
integer(observed, 'observedAtMs'),
|
||||
);
|
||||
const usage = quotaWindowUsage(
|
||||
client,
|
||||
start.projectId,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowMs,
|
||||
);
|
||||
if (
|
||||
usage.invocationCount + 1 > reservation.maxInvocations ||
|
||||
usage.effectiveTokens + reservation.reservedTokens >
|
||||
reservation.maxTokens ||
|
||||
(reservation.maxCostMicros !== null &&
|
||||
(usage.unknownCostInvocations !== 0 ||
|
||||
usage.effectiveCostMicros + reservation.reservedCostMicros! >
|
||||
reservation.maxCostMicros))
|
||||
) {
|
||||
throw new ModelInvocationProjectQuotaExceededError();
|
||||
}
|
||||
assertCurrent(client, command.stepRunMutation, start.projectId);
|
||||
applyMutation(client, command.stepRunMutation);
|
||||
insertStart(client, start);
|
||||
insertQuotaReservation(client, reservation);
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'created' as const, record: start });
|
||||
} catch (error) {
|
||||
if (began && client.isTransaction) {
|
||||
try {
|
||||
client.exec('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the original transaction failure.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
export function admitWithPricingOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
quoteValue: ModelInvocationPriceQuote,
|
||||
admissionValue?: ModelInvocationQuotaAdmission,
|
||||
): Promise<Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>> {
|
||||
const command = normalizeModelInvocationStartCommand(commandValue);
|
||||
const quote = normalizeModelInvocationPriceQuote(quoteValue);
|
||||
const admission =
|
||||
admissionValue === undefined
|
||||
? undefined
|
||||
: normalizeModelInvocationQuotaAdmission(admissionValue);
|
||||
const start = command.start;
|
||||
if (
|
||||
quote.invocationId !== start.invocationId ||
|
||||
quote.projectId !== start.projectId ||
|
||||
quote.modelPolicyRevision !== start.policyRevision ||
|
||||
quote.provider !== start.provider ||
|
||||
quote.model !== start.model ||
|
||||
quote.maxOutputTokens !== start.maxOutputTokens ||
|
||||
(admission !== undefined &&
|
||||
(admission.invocationId !== start.invocationId ||
|
||||
admission.projectId !== start.projectId ||
|
||||
admission.modelPolicyRevision !== start.policyRevision ||
|
||||
(admission.maxCostMicros !== null &&
|
||||
admission.reservedCostMicros !== quote.reservedCostMicros)))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
let began = false;
|
||||
try {
|
||||
client.exec('BEGIN IMMEDIATE');
|
||||
began = true;
|
||||
const existing = startRows(
|
||||
client,
|
||||
`start.invocation_id = ? OR
|
||||
start.mutation_id = ? OR start.run_event_id = ?`,
|
||||
[start.invocationId, start.stepRunMutationId, start.runEventId],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseStart(existing[0]);
|
||||
const quotes = priceQuoteRows(client, 'quote.invocation_id = ?', [
|
||||
start.invocationId,
|
||||
]);
|
||||
const reservations = quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = ?',
|
||||
[start.invocationId],
|
||||
);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(start) ||
|
||||
quotes.length !== 1 ||
|
||||
JSON.stringify(parsePriceQuote(quotes[0]!)) !==
|
||||
JSON.stringify(quote) ||
|
||||
reservations.length !== (admission ? 1 : 0) ||
|
||||
(admission !== undefined &&
|
||||
parseQuotaReservation(reservations[0]!).admissionDigest !==
|
||||
admission.admissionDigest)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
assertLocalFeatureActive(client);
|
||||
let reservation: Readonly<ModelInvocationQuotaReservation> | undefined;
|
||||
if (admission) {
|
||||
const observed = client
|
||||
.prepare(
|
||||
`SELECT CAST(unixepoch('subsec') * 1000 AS INTEGER)
|
||||
AS "observedAtMs"`,
|
||||
)
|
||||
.get() as Row | undefined;
|
||||
if (!observed) throw unavailable();
|
||||
reservation = createModelInvocationQuotaReservation(
|
||||
admission,
|
||||
integer(observed, 'observedAtMs'),
|
||||
);
|
||||
const usage = quotaWindowUsage(
|
||||
client,
|
||||
start.projectId,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowMs,
|
||||
);
|
||||
if (
|
||||
usage.invocationCount + 1 > reservation.maxInvocations ||
|
||||
usage.effectiveTokens + reservation.reservedTokens >
|
||||
reservation.maxTokens ||
|
||||
(reservation.maxCostMicros !== null &&
|
||||
(usage.unknownCostInvocations !== 0 ||
|
||||
usage.effectiveCostMicros + reservation.reservedCostMicros! >
|
||||
reservation.maxCostMicros))
|
||||
) {
|
||||
throw new ModelInvocationProjectQuotaExceededError();
|
||||
}
|
||||
}
|
||||
assertCurrent(client, command.stepRunMutation, start.projectId);
|
||||
applyMutation(client, command.stepRunMutation);
|
||||
insertStart(client, start);
|
||||
insertPriceQuote(client, quote);
|
||||
if (reservation) {
|
||||
insertQuotaReservation(client, reservation);
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'created' as const, record: start });
|
||||
} catch (error) {
|
||||
if (began && client.isTransaction) {
|
||||
try {
|
||||
client.exec('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the original transaction failure.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
import type { DatabaseSync } from 'node:sqlite';
|
||||
|
||||
import { assertLocalModelInvocationFeatureActive } from '../../feature-activation/localModelInvocationFeatureActivation';
|
||||
import {
|
||||
PluginPackagePromptOutputArtifactConflictError,
|
||||
PluginPackagePromptOutputArtifactUnavailableError,
|
||||
} from '../../prompt-output/pluginPackagePromptOutputArtifact';
|
||||
import { ModelInvocationUsageSummaryLimitExceededError } from '../../usage/usageLedger';
|
||||
import { ModelInvocationProjectQuotaExceededError } from '../../usage/usageQuota';
|
||||
import {
|
||||
MAX_MODEL_INVOCATION_RECOVERY_PAGE_SIZE,
|
||||
ModelInvocationConflictError,
|
||||
ModelInvocationRepositoryUnavailableError,
|
||||
} from '../modelInvocation';
|
||||
|
||||
export type Row = Record<string, unknown>;
|
||||
|
||||
export interface LocalModelInvocationOperationAuthority {
|
||||
readonly client: DatabaseSync;
|
||||
enqueue<T>(
|
||||
work: () => Promise<T>,
|
||||
rejection: (reason: 'closed' | 'busy') => Error,
|
||||
): Promise<T>;
|
||||
}
|
||||
|
||||
export const IDENTIFIER_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}$/;
|
||||
export const TERMINAL_RUN_STATUSES = new Set([
|
||||
'succeeded',
|
||||
'failed',
|
||||
'cancelled',
|
||||
'timed_out',
|
||||
]);
|
||||
|
||||
export function unavailable(
|
||||
cause?: unknown,
|
||||
): ModelInvocationRepositoryUnavailableError {
|
||||
return new ModelInvocationRepositoryUnavailableError({
|
||||
cause: cause instanceof Error ? cause : undefined,
|
||||
});
|
||||
}
|
||||
|
||||
export function assertLocalFeatureActive(client: DatabaseSync): void {
|
||||
try {
|
||||
assertLocalModelInvocationFeatureActive(client);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
}
|
||||
|
||||
export function identifier(value: unknown): string {
|
||||
if (typeof value !== 'string' || !IDENTIFIER_PATTERN.test(value)) {
|
||||
throw unavailable();
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
export function text(row: Row, key: string): string {
|
||||
const value = row[key];
|
||||
if (typeof value !== 'string') throw unavailable();
|
||||
return value;
|
||||
}
|
||||
|
||||
export function integer(row: Row, key: string): number {
|
||||
const value = row[key];
|
||||
if (!Number.isSafeInteger(value) || (value as number) < 0) {
|
||||
throw unavailable();
|
||||
}
|
||||
return value as number;
|
||||
}
|
||||
|
||||
export function nullableInteger(row: Row, key: string): number | null {
|
||||
return row[key] === null ? null : integer(row, key);
|
||||
}
|
||||
|
||||
export function recoveryLimit(value: number): number {
|
||||
if (
|
||||
!Number.isSafeInteger(value) ||
|
||||
value < 1 ||
|
||||
value > MAX_MODEL_INVOCATION_RECOVERY_PAGE_SIZE
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
export function sqliteCode(error: unknown): string {
|
||||
if (!error || typeof error !== 'object') return '';
|
||||
const value = (error as { code?: unknown }).code;
|
||||
return typeof value === 'string' ? value : '';
|
||||
}
|
||||
|
||||
export function mapStorageError(error: unknown): Error {
|
||||
if (
|
||||
error instanceof ModelInvocationConflictError ||
|
||||
error instanceof ModelInvocationRepositoryUnavailableError ||
|
||||
error instanceof ModelInvocationUsageSummaryLimitExceededError ||
|
||||
error instanceof ModelInvocationProjectQuotaExceededError ||
|
||||
error instanceof PluginPackagePromptOutputArtifactConflictError ||
|
||||
error instanceof PluginPackagePromptOutputArtifactUnavailableError
|
||||
) {
|
||||
return error;
|
||||
}
|
||||
const code = sqliteCode(error);
|
||||
if (code.startsWith('ERR_SQLITE_CONSTRAINT')) {
|
||||
return new ModelInvocationConflictError();
|
||||
}
|
||||
return unavailable(error);
|
||||
}
|
||||
|
||||
export class PrivateLocalAuthority
|
||||
implements LocalModelInvocationOperationAuthority
|
||||
{
|
||||
readonly client: DatabaseSync;
|
||||
#tail: Promise<void> = Promise.resolve();
|
||||
#pending = 0;
|
||||
|
||||
constructor(client: DatabaseSync) {
|
||||
this.client = client;
|
||||
}
|
||||
|
||||
enqueue<T>(
|
||||
work: () => Promise<T>,
|
||||
rejection: (reason: 'closed' | 'busy') => Error,
|
||||
): Promise<T> {
|
||||
if (this.#pending >= 64) return Promise.reject(rejection('busy'));
|
||||
this.#pending += 1;
|
||||
const result = this.#tail.then(work, work);
|
||||
this.#tail = result.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
);
|
||||
return result.finally(() => {
|
||||
this.#pending -= 1;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
export function isAuthority(
|
||||
value: LocalModelInvocationOperationAuthority | DatabaseSync,
|
||||
): value is LocalModelInvocationOperationAuthority {
|
||||
return (
|
||||
!!value &&
|
||||
typeof value === 'object' &&
|
||||
'client' in value &&
|
||||
'enqueue' in value &&
|
||||
typeof value.enqueue === 'function'
|
||||
);
|
||||
}
|
||||
|
||||
export function enqueueLocalModelInvocation<T>(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
work: () => T,
|
||||
): Promise<T> {
|
||||
return authority.enqueue(async () => {
|
||||
try {
|
||||
return work();
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}, unavailable);
|
||||
}
|
||||
@@ -0,0 +1,342 @@
|
||||
import {
|
||||
normalizeStepRunRecord,
|
||||
type StepRunRecord,
|
||||
} from '@qinglong/runtime-core/step-run';
|
||||
import {
|
||||
normalizeModelInvocationPriceQuote,
|
||||
normalizeModelInvocationPriceSettlement,
|
||||
type ModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceSettlement,
|
||||
} from '../../pricing/pricing';
|
||||
import {
|
||||
normalizeModelInvocationUsageLedgerRecord,
|
||||
type ModelInvocationUsageLedgerRecord,
|
||||
} from '../../usage/usageLedger';
|
||||
import {
|
||||
normalizeModelInvocationQuotaReservation,
|
||||
normalizeModelInvocationQuotaSettlement,
|
||||
type ModelInvocationQuotaReservation,
|
||||
type ModelInvocationQuotaSettlement,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
normalizeModelInvocationCompletionRecord,
|
||||
normalizeModelInvocationStartRecord,
|
||||
type ModelInvocationAuthoritySnapshot,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
import {
|
||||
normalizeModelInvocationResolutionRecord,
|
||||
type ModelInvocationResolutionRecord,
|
||||
} from '../modelInvocationResolution';
|
||||
|
||||
import type { Row } from './authority';
|
||||
import { integer, nullableInteger, text, unavailable } from './authority';
|
||||
|
||||
export function parseStart(row: Row): Readonly<ModelInvocationStartRecord> {
|
||||
let start: Readonly<ModelInvocationStartRecord>;
|
||||
try {
|
||||
start = normalizeModelInvocationStartRecord(
|
||||
JSON.parse(text(row, 'recordJson')) as ModelInvocationStartRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
start.invocationId !== text(row, 'invocationId') ||
|
||||
start.projectId !== text(row, 'projectId') ||
|
||||
start.runId !== text(row, 'runId') ||
|
||||
start.stepRunId !== text(row, 'stepRunId') ||
|
||||
start.traceId !== text(row, 'traceId') ||
|
||||
start.provider !== text(row, 'provider') ||
|
||||
start.model !== text(row, 'model') ||
|
||||
start.policyRevision !== text(row, 'policyRevision') ||
|
||||
start.requestDigest !== text(row, 'requestDigest') ||
|
||||
start.inputBytes !== integer(row, 'inputBytes') ||
|
||||
start.maxOutputTokens !== integer(row, 'maxOutputTokens') ||
|
||||
start.deadlineAtMs !== integer(row, 'deadlineAtMs') ||
|
||||
start.admittedAtMs !== integer(row, 'admittedAtMs') ||
|
||||
start.stepRunMutationId !== text(row, 'mutationId') ||
|
||||
start.stepRunMutationDigest !== text(row, 'mutationDigest') ||
|
||||
start.startedStepRunDigest !== text(row, 'stepRunDigest') ||
|
||||
start.startedStepRunVersion !== integer(row, 'stepRunVersion') ||
|
||||
start.runEventId !== text(row, 'runEventId') ||
|
||||
start.startDigest !== text(row, 'startDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return start;
|
||||
}
|
||||
|
||||
export function parseCompletion(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationCompletionRecord> {
|
||||
let completion: Readonly<ModelInvocationCompletionRecord>;
|
||||
try {
|
||||
completion = normalizeModelInvocationCompletionRecord(
|
||||
JSON.parse(text(row, 'recordJson')) as ModelInvocationCompletionRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
completion.invocationId !== text(row, 'invocationId') ||
|
||||
completion.projectId !== text(row, 'projectId') ||
|
||||
completion.runId !== text(row, 'runId') ||
|
||||
completion.stepRunId !== text(row, 'stepRunId') ||
|
||||
completion.traceId !== text(row, 'traceId') ||
|
||||
completion.startDigest !== text(row, 'startDigest') ||
|
||||
completion.outcome !== text(row, 'outcome') ||
|
||||
completion.outputBytes !== integer(row, 'outputBytes') ||
|
||||
completion.errorCode !== row.errorCode ||
|
||||
completion.completedAtMs !== integer(row, 'completedAtMs') ||
|
||||
completion.stepRunMutationId !== text(row, 'mutationId') ||
|
||||
completion.stepRunMutationDigest !== text(row, 'mutationDigest') ||
|
||||
completion.completedStepRunDigest !== text(row, 'stepRunDigest') ||
|
||||
completion.completedStepRunVersion !== integer(row, 'stepRunVersion') ||
|
||||
completion.runEventId !== text(row, 'runEventId') ||
|
||||
completion.completionDigest !== text(row, 'completionDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return completion;
|
||||
}
|
||||
|
||||
export function parseUsage(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationUsageLedgerRecord> {
|
||||
let usage: Readonly<ModelInvocationUsageLedgerRecord>;
|
||||
try {
|
||||
usage = normalizeModelInvocationUsageLedgerRecord(
|
||||
JSON.parse(text(row, 'recordJson')) as ModelInvocationUsageLedgerRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
usage.invocationId !== text(row, 'invocationId') ||
|
||||
usage.projectId !== text(row, 'projectId') ||
|
||||
usage.runId !== text(row, 'runId') ||
|
||||
usage.stepRunId !== text(row, 'stepRunId') ||
|
||||
usage.traceId !== text(row, 'traceId') ||
|
||||
usage.provider !== text(row, 'provider') ||
|
||||
usage.model !== text(row, 'model') ||
|
||||
usage.policyRevision !== text(row, 'policyRevision') ||
|
||||
usage.completionDigest !== text(row, 'completionDigest') ||
|
||||
usage.outcome !== text(row, 'outcome') ||
|
||||
usage.settledAtMs !== integer(row, 'settledAtMs') ||
|
||||
usage.inputBytes !== integer(row, 'inputBytes') ||
|
||||
usage.outputBytes !== integer(row, 'outputBytes') ||
|
||||
usage.inputTokens !== integer(row, 'inputTokens') ||
|
||||
usage.outputTokens !== integer(row, 'outputTokens') ||
|
||||
usage.totalTokens !== integer(row, 'totalTokens') ||
|
||||
usage.costMicros !== nullableInteger(row, 'costMicros') ||
|
||||
usage.ledgerDigest !== text(row, 'ledgerDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return usage;
|
||||
}
|
||||
|
||||
export function parseQuotaReservation(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationQuotaReservation> {
|
||||
try {
|
||||
return normalizeModelInvocationQuotaReservation(
|
||||
JSON.parse(text(row, 'recordJson')) as ModelInvocationQuotaReservation,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
}
|
||||
|
||||
export function parseQuotaSettlement(
|
||||
row: Row,
|
||||
reservation: Readonly<ModelInvocationQuotaReservation>,
|
||||
completion: Readonly<ModelInvocationCompletionRecord>,
|
||||
): Readonly<ModelInvocationQuotaSettlement> {
|
||||
try {
|
||||
return normalizeModelInvocationQuotaSettlement(
|
||||
JSON.parse(text(row, 'recordJson')) as ModelInvocationQuotaSettlement,
|
||||
reservation,
|
||||
completion,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
}
|
||||
|
||||
export function parsePriceQuote(row: Row): Readonly<ModelInvocationPriceQuote> {
|
||||
try {
|
||||
return normalizeModelInvocationPriceQuote(
|
||||
JSON.parse(text(row, 'recordJson')) as ModelInvocationPriceQuote,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
}
|
||||
|
||||
export function parsePriceSettlement(
|
||||
row: Row,
|
||||
quote: Readonly<ModelInvocationPriceQuote>,
|
||||
completion: Readonly<ModelInvocationCompletionRecord>,
|
||||
): Readonly<ModelInvocationPriceSettlement> {
|
||||
try {
|
||||
return normalizeModelInvocationPriceSettlement(
|
||||
JSON.parse(text(row, 'recordJson')) as ModelInvocationPriceSettlement,
|
||||
quote,
|
||||
completion,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
}
|
||||
|
||||
export function parseResolution(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationResolutionRecord> {
|
||||
let resolution: Readonly<ModelInvocationResolutionRecord>;
|
||||
try {
|
||||
resolution = normalizeModelInvocationResolutionRecord(
|
||||
JSON.parse(text(row, 'recordJson')) as ModelInvocationResolutionRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
resolution.resolutionId !== text(row, 'resolutionId') ||
|
||||
resolution.invocationId !== text(row, 'invocationId') ||
|
||||
resolution.projectId !== text(row, 'projectId') ||
|
||||
resolution.runId !== text(row, 'runId') ||
|
||||
resolution.stepRunId !== text(row, 'stepRunId') ||
|
||||
resolution.traceId !== text(row, 'traceId') ||
|
||||
resolution.completionDigest !== text(row, 'completionDigest') ||
|
||||
resolution.decision !== text(row, 'decision') ||
|
||||
resolution.resolvedByUserId !== text(row, 'resolvedByUserId') ||
|
||||
resolution.resolvedAtMs !== integer(row, 'resolvedAtMs') ||
|
||||
resolution.stepRunMutationId !== text(row, 'mutationId') ||
|
||||
resolution.stepRunMutationDigest !== text(row, 'mutationDigest') ||
|
||||
resolution.resolvedStepRunDigest !== text(row, 'stepRunDigest') ||
|
||||
resolution.resolvedStepRunVersion !== integer(row, 'stepRunVersion') ||
|
||||
resolution.runEventId !== text(row, 'runEventId') ||
|
||||
resolution.resolutionDigest !== text(row, 'resolutionDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return resolution;
|
||||
}
|
||||
|
||||
export function parseAuthority(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationAuthoritySnapshot> {
|
||||
let stepRun: Readonly<StepRunRecord>;
|
||||
try {
|
||||
stepRun = normalizeStepRunRecord(
|
||||
JSON.parse(text(row, 'stepRunJson')) as StepRunRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
stepRun.id !== text(row, 'stepRunId') ||
|
||||
stepRun.runId !== text(row, 'runId') ||
|
||||
stepRun.kind !== 'model' ||
|
||||
stepRun.status !== text(row, 'stepStatus') ||
|
||||
stepRun.version !== integer(row, 'stepVersion') ||
|
||||
stepRun.stepRunDigest !== text(row, 'stepDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return Object.freeze({
|
||||
projectId: text(row, 'projectId'),
|
||||
runId: text(row, 'runId'),
|
||||
runVersion: integer(row, 'runVersion'),
|
||||
runEventSequence: integer(row, 'runEventSequence'),
|
||||
stepRun,
|
||||
});
|
||||
}
|
||||
|
||||
export const START_SELECT = `
|
||||
start.invocation_id AS "invocationId",
|
||||
start.project_id AS "projectId",
|
||||
start.run_id AS "runId",
|
||||
start.step_run_id AS "stepRunId",
|
||||
start.trace_id AS "traceId",
|
||||
start.provider AS "provider",
|
||||
start.model AS "model",
|
||||
start.policy_revision AS "policyRevision",
|
||||
start.request_digest AS "requestDigest",
|
||||
start.input_bytes AS "inputBytes",
|
||||
start.max_output_tokens AS "maxOutputTokens",
|
||||
start.deadline_at_ms AS "deadlineAtMs",
|
||||
start.admitted_at_ms AS "admittedAtMs",
|
||||
start.mutation_id AS "mutationId",
|
||||
start.mutation_digest AS "mutationDigest",
|
||||
start.run_event_id AS "runEventId",
|
||||
start.start_digest AS "startDigest",
|
||||
start.record_json AS "recordJson",
|
||||
mutation.step_run_digest AS "stepRunDigest",
|
||||
json_extract(mutation.step_run_json, '$.version') AS "stepRunVersion"
|
||||
`;
|
||||
|
||||
export const COMPLETION_SELECT = `
|
||||
completion.invocation_id AS "invocationId",
|
||||
completion.project_id AS "projectId",
|
||||
completion.run_id AS "runId",
|
||||
completion.step_run_id AS "stepRunId",
|
||||
completion.trace_id AS "traceId",
|
||||
completion.start_digest AS "startDigest",
|
||||
completion.outcome AS "outcome",
|
||||
completion.output_bytes AS "outputBytes",
|
||||
completion.error_code AS "errorCode",
|
||||
completion.completed_at_ms AS "completedAtMs",
|
||||
completion.mutation_id AS "mutationId",
|
||||
completion.mutation_digest AS "mutationDigest",
|
||||
completion.run_event_id AS "runEventId",
|
||||
completion.completion_digest AS "completionDigest",
|
||||
completion.record_json AS "recordJson",
|
||||
mutation.step_run_digest AS "stepRunDigest",
|
||||
json_extract(mutation.step_run_json, '$.version') AS "stepRunVersion"
|
||||
`;
|
||||
|
||||
export const USAGE_SELECT = `
|
||||
usage.invocation_id AS "invocationId",
|
||||
usage.project_id AS "projectId",
|
||||
usage.run_id AS "runId",
|
||||
usage.step_run_id AS "stepRunId",
|
||||
usage.trace_id AS "traceId",
|
||||
usage.provider AS "provider",
|
||||
usage.model AS "model",
|
||||
usage.policy_revision AS "policyRevision",
|
||||
usage.completion_digest AS "completionDigest",
|
||||
usage.outcome AS "outcome",
|
||||
usage.settled_at_ms AS "settledAtMs",
|
||||
usage.input_bytes AS "inputBytes",
|
||||
usage.output_bytes AS "outputBytes",
|
||||
usage.input_tokens AS "inputTokens",
|
||||
usage.output_tokens AS "outputTokens",
|
||||
usage.total_tokens AS "totalTokens",
|
||||
usage.cost_micros AS "costMicros",
|
||||
usage.ledger_digest AS "ledgerDigest",
|
||||
usage.record_json AS "recordJson"
|
||||
`;
|
||||
|
||||
export const RESOLUTION_SELECT = `
|
||||
resolution.resolution_id AS "resolutionId",
|
||||
resolution.invocation_id AS "invocationId",
|
||||
resolution.project_id AS "projectId",
|
||||
resolution.run_id AS "runId",
|
||||
resolution.step_run_id AS "stepRunId",
|
||||
resolution.trace_id AS "traceId",
|
||||
resolution.completion_digest AS "completionDigest",
|
||||
resolution.decision AS "decision",
|
||||
resolution.resolved_by_user_id AS "resolvedByUserId",
|
||||
resolution.resolved_at_ms AS "resolvedAtMs",
|
||||
resolution.mutation_id AS "mutationId",
|
||||
resolution.mutation_digest AS "mutationDigest",
|
||||
resolution.run_event_id AS "runEventId",
|
||||
resolution.resolution_digest AS "resolutionDigest",
|
||||
resolution.record_json AS "recordJson",
|
||||
mutation.step_run_digest AS "stepRunDigest",
|
||||
json_extract(mutation.step_run_json, '$.version') AS "stepRunVersion"
|
||||
`;
|
||||
+538
@@ -0,0 +1,538 @@
|
||||
import type { DatabaseSync } from 'node:sqlite';
|
||||
|
||||
import { createModelInvocationPriceSettlement } from '../../pricing/pricing';
|
||||
import { type PluginPackagePromptOutputArtifact } from '../../prompt-output/pluginPackagePromptOutputArtifact';
|
||||
import {
|
||||
assertPluginPackagePromptOutputCompletionBinding,
|
||||
type CommitPluginPackagePromptOutputResult,
|
||||
} from '../../prompt-output/pluginPackagePromptOutputCompletion';
|
||||
import {
|
||||
putLocalPluginPackagePromptOutputArtifactInTransaction,
|
||||
readLocalPluginPackagePromptOutputArtifactInTransaction,
|
||||
} from '../../prompt-output/storage/localPluginPackagePromptOutputArtifactRepository';
|
||||
import { createModelInvocationUsageLedgerRecord } from '../../usage/usageLedger';
|
||||
import { createModelInvocationQuotaSettlement } from '../../usage/usageQuota';
|
||||
import {
|
||||
ModelInvocationConflictError,
|
||||
normalizeModelInvocationCompletionCommand,
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationCompletionCommand,
|
||||
type ModelInvocationCompletionRecord,
|
||||
} from '../modelInvocation';
|
||||
|
||||
import type { LocalModelInvocationOperationAuthority } from './authority';
|
||||
import { enqueueLocalModelInvocation } from './authority';
|
||||
import {
|
||||
parseCompletion,
|
||||
parsePriceQuote,
|
||||
parsePriceSettlement,
|
||||
parseQuotaReservation,
|
||||
parseQuotaSettlement,
|
||||
parseStart,
|
||||
parseUsage,
|
||||
} from './codec';
|
||||
import {
|
||||
applyMutation,
|
||||
assertCurrent,
|
||||
insertCompletion,
|
||||
insertPriceSettlement,
|
||||
insertQuotaSettlement,
|
||||
insertUsage,
|
||||
} from './mutations';
|
||||
import {
|
||||
completionRows,
|
||||
priceQuoteRows,
|
||||
priceSettlementRows,
|
||||
quotaReservationRows,
|
||||
quotaSettlementRows,
|
||||
startRows,
|
||||
usageRows,
|
||||
} from './queries';
|
||||
|
||||
export function completeOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
const command = normalizeModelInvocationCompletionCommand(commandValue);
|
||||
const completion = command.completion;
|
||||
const expectedUsage = createModelInvocationUsageLedgerRecord(
|
||||
command.start,
|
||||
completion,
|
||||
);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
let began = false;
|
||||
try {
|
||||
client.exec('BEGIN IMMEDIATE');
|
||||
began = true;
|
||||
const existing = completionRows(
|
||||
client,
|
||||
`completion.invocation_id = ? OR
|
||||
completion.mutation_id = ? OR completion.run_event_id = ?`,
|
||||
[
|
||||
completion.invocationId,
|
||||
completion.stepRunMutationId,
|
||||
completion.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseCompletion(existing[0]);
|
||||
if (JSON.stringify(stored) !== JSON.stringify(completion)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const usage = usageRows(client, 'usage.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
usage.length !== (expectedUsage ? 1 : 0) ||
|
||||
(expectedUsage &&
|
||||
JSON.stringify(parseUsage(usage[0]!)) !==
|
||||
JSON.stringify(expectedUsage))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
const starts = startRows(client, 'start.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
starts.length !== 1 ||
|
||||
JSON.stringify(parseStart(starts[0]!)) !== JSON.stringify(command.start)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
assertCurrent(client, command.stepRunMutation, completion.projectId);
|
||||
applyMutation(client, command.stepRunMutation);
|
||||
insertCompletion(client, completion);
|
||||
if (expectedUsage) insertUsage(client, expectedUsage);
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: completion,
|
||||
});
|
||||
} catch (error) {
|
||||
if (began && client.isTransaction) {
|
||||
try {
|
||||
client.exec('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the original transaction failure.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
export function completeWithQuotaOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
const command = normalizeModelInvocationCompletionCommand(commandValue);
|
||||
const completion = command.completion;
|
||||
const expectedUsage = createModelInvocationUsageLedgerRecord(
|
||||
command.start,
|
||||
completion,
|
||||
);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
let began = false;
|
||||
try {
|
||||
client.exec('BEGIN IMMEDIATE');
|
||||
began = true;
|
||||
const reservationRows = quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = ?',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (reservationRows.length !== 1) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const reservation = parseQuotaReservation(reservationRows[0]!);
|
||||
const expectedSettlement = createModelInvocationQuotaSettlement(
|
||||
reservation,
|
||||
completion,
|
||||
);
|
||||
const existing = completionRows(
|
||||
client,
|
||||
`completion.invocation_id = ? OR
|
||||
completion.mutation_id = ? OR completion.run_event_id = ?`,
|
||||
[
|
||||
completion.invocationId,
|
||||
completion.stepRunMutationId,
|
||||
completion.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseCompletion(existing[0]);
|
||||
const usage = usageRows(client, 'usage.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
const settlements = quotaSettlementRows(
|
||||
client,
|
||||
'settlement.invocation_id = ?',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(completion) ||
|
||||
usage.length !== (expectedUsage ? 1 : 0) ||
|
||||
(expectedUsage &&
|
||||
JSON.stringify(parseUsage(usage[0]!)) !==
|
||||
JSON.stringify(expectedUsage)) ||
|
||||
settlements.length !== 1 ||
|
||||
JSON.stringify(
|
||||
parseQuotaSettlement(settlements[0]!, reservation, completion),
|
||||
) !== JSON.stringify(expectedSettlement)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
const starts = startRows(client, 'start.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
starts.length !== 1 ||
|
||||
JSON.stringify(parseStart(starts[0]!)) !== JSON.stringify(command.start)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
assertCurrent(client, command.stepRunMutation, completion.projectId);
|
||||
applyMutation(client, command.stepRunMutation);
|
||||
insertCompletion(client, completion);
|
||||
if (expectedUsage) insertUsage(client, expectedUsage);
|
||||
insertQuotaSettlement(client, expectedSettlement);
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: completion,
|
||||
});
|
||||
} catch (error) {
|
||||
if (began && client.isTransaction) {
|
||||
try {
|
||||
client.exec('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the original transaction failure.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
export function completeWithPricingOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
const command = normalizeModelInvocationCompletionCommand(commandValue);
|
||||
const completion = command.completion;
|
||||
const expectedUsage = createModelInvocationUsageLedgerRecord(
|
||||
command.start,
|
||||
completion,
|
||||
);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
let began = false;
|
||||
try {
|
||||
client.exec('BEGIN IMMEDIATE');
|
||||
began = true;
|
||||
const quoteRows = priceQuoteRows(client, 'quote.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (quoteRows.length !== 1) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const quote = parsePriceQuote(quoteRows[0]!);
|
||||
if (
|
||||
quote.invocationId !== command.start.invocationId ||
|
||||
quote.projectId !== command.start.projectId ||
|
||||
quote.modelPolicyRevision !== command.start.policyRevision ||
|
||||
quote.provider !== command.start.provider ||
|
||||
quote.model !== command.start.model
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const expectedPriceSettlement = createModelInvocationPriceSettlement(
|
||||
quote,
|
||||
completion,
|
||||
);
|
||||
const reservationRows = quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = ?',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (reservationRows.length > 1) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const reservation = reservationRows[0]
|
||||
? parseQuotaReservation(reservationRows[0])
|
||||
: null;
|
||||
const expectedQuotaSettlement = reservation
|
||||
? createModelInvocationQuotaSettlement(reservation, completion)
|
||||
: null;
|
||||
const existing = completionRows(
|
||||
client,
|
||||
`completion.invocation_id = ? OR
|
||||
completion.mutation_id = ? OR completion.run_event_id = ?`,
|
||||
[
|
||||
completion.invocationId,
|
||||
completion.stepRunMutationId,
|
||||
completion.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseCompletion(existing[0]);
|
||||
const usage = usageRows(client, 'usage.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
const priceSettlements = priceSettlementRows(
|
||||
client,
|
||||
'settlement.invocation_id = ?',
|
||||
[completion.invocationId],
|
||||
);
|
||||
const quotaSettlements = quotaSettlementRows(
|
||||
client,
|
||||
'settlement.invocation_id = ?',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(completion) ||
|
||||
usage.length !== (expectedUsage ? 1 : 0) ||
|
||||
(expectedUsage &&
|
||||
JSON.stringify(parseUsage(usage[0]!)) !==
|
||||
JSON.stringify(expectedUsage)) ||
|
||||
priceSettlements.length !== (expectedPriceSettlement ? 1 : 0) ||
|
||||
(expectedPriceSettlement &&
|
||||
JSON.stringify(
|
||||
parsePriceSettlement(priceSettlements[0]!, quote, completion),
|
||||
) !== JSON.stringify(expectedPriceSettlement)) ||
|
||||
quotaSettlements.length !== (expectedQuotaSettlement ? 1 : 0) ||
|
||||
(expectedQuotaSettlement &&
|
||||
JSON.stringify(
|
||||
parseQuotaSettlement(
|
||||
quotaSettlements[0]!,
|
||||
reservation!,
|
||||
completion,
|
||||
),
|
||||
) !== JSON.stringify(expectedQuotaSettlement))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
const starts = startRows(client, 'start.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
starts.length !== 1 ||
|
||||
JSON.stringify(parseStart(starts[0]!)) !== JSON.stringify(command.start)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
assertCurrent(client, command.stepRunMutation, completion.projectId);
|
||||
applyMutation(client, command.stepRunMutation);
|
||||
insertCompletion(client, completion);
|
||||
if (expectedUsage) insertUsage(client, expectedUsage);
|
||||
if (expectedPriceSettlement) {
|
||||
insertPriceSettlement(client, expectedPriceSettlement);
|
||||
}
|
||||
if (expectedQuotaSettlement) {
|
||||
insertQuotaSettlement(client, expectedQuotaSettlement);
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: completion,
|
||||
});
|
||||
} catch (error) {
|
||||
if (began && client.isTransaction) {
|
||||
try {
|
||||
client.exec('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the original transaction failure.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
export function completeWithPromptOutputArtifactOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
artifactValue: PluginPackagePromptOutputArtifact,
|
||||
): Promise<Readonly<CommitPluginPackagePromptOutputResult>> {
|
||||
const command = normalizeModelInvocationCompletionCommand(commandValue);
|
||||
const completion = command.completion;
|
||||
const binding = assertPluginPackagePromptOutputCompletionBinding(
|
||||
command,
|
||||
artifactValue,
|
||||
);
|
||||
const expectedUsage = createModelInvocationUsageLedgerRecord(
|
||||
command.start,
|
||||
completion,
|
||||
);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
let began = false;
|
||||
try {
|
||||
client.exec('BEGIN IMMEDIATE');
|
||||
began = true;
|
||||
const quoteRows = priceQuoteRows(client, 'quote.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (quoteRows.length > 1) throw new ModelInvocationConflictError();
|
||||
const quote = quoteRows[0] ? parsePriceQuote(quoteRows[0]) : null;
|
||||
if (
|
||||
quote &&
|
||||
(quote.invocationId !== command.start.invocationId ||
|
||||
quote.projectId !== command.start.projectId ||
|
||||
quote.modelPolicyRevision !== command.start.policyRevision ||
|
||||
quote.provider !== command.start.provider ||
|
||||
quote.model !== command.start.model)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const expectedPriceSettlement = quote
|
||||
? createModelInvocationPriceSettlement(quote, completion)
|
||||
: null;
|
||||
const reservationRows = quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = ?',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (reservationRows.length > 1) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const reservation = reservationRows[0]
|
||||
? parseQuotaReservation(reservationRows[0])
|
||||
: null;
|
||||
const expectedQuotaSettlement = reservation
|
||||
? createModelInvocationQuotaSettlement(reservation, completion)
|
||||
: null;
|
||||
const existing = completionRows(
|
||||
client,
|
||||
`completion.invocation_id = ? OR
|
||||
completion.mutation_id = ? OR completion.run_event_id = ?`,
|
||||
[
|
||||
completion.invocationId,
|
||||
completion.stepRunMutationId,
|
||||
completion.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseCompletion(existing[0]);
|
||||
const storedArtifact =
|
||||
readLocalPluginPackagePromptOutputArtifactInTransaction(
|
||||
client,
|
||||
binding.artifact.artifactId,
|
||||
);
|
||||
const usage = usageRows(client, 'usage.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
const priceSettlements = priceSettlementRows(
|
||||
client,
|
||||
'settlement.invocation_id = ?',
|
||||
[completion.invocationId],
|
||||
);
|
||||
const quotaSettlements = quotaSettlementRows(
|
||||
client,
|
||||
'settlement.invocation_id = ?',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(completion) ||
|
||||
!storedArtifact ||
|
||||
JSON.stringify(storedArtifact) !== JSON.stringify(binding.artifact) ||
|
||||
usage.length !== (expectedUsage ? 1 : 0) ||
|
||||
(expectedUsage &&
|
||||
JSON.stringify(parseUsage(usage[0]!)) !==
|
||||
JSON.stringify(expectedUsage)) ||
|
||||
priceSettlements.length !== (expectedPriceSettlement ? 1 : 0) ||
|
||||
(expectedPriceSettlement &&
|
||||
JSON.stringify(
|
||||
parsePriceSettlement(priceSettlements[0]!, quote!, completion),
|
||||
) !== JSON.stringify(expectedPriceSettlement)) ||
|
||||
quotaSettlements.length !== (expectedQuotaSettlement ? 1 : 0) ||
|
||||
(expectedQuotaSettlement &&
|
||||
JSON.stringify(
|
||||
parseQuotaSettlement(
|
||||
quotaSettlements[0]!,
|
||||
reservation!,
|
||||
completion,
|
||||
),
|
||||
) !== JSON.stringify(expectedQuotaSettlement))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({
|
||||
status: 'existing' as const,
|
||||
record: stored,
|
||||
artifact: storedArtifact,
|
||||
reference: binding.reference,
|
||||
});
|
||||
}
|
||||
const starts = startRows(client, 'start.invocation_id = ?', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
starts.length !== 1 ||
|
||||
JSON.stringify(parseStart(starts[0]!)) !== JSON.stringify(command.start)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
assertCurrent(client, command.stepRunMutation, completion.projectId);
|
||||
const artifact = putLocalPluginPackagePromptOutputArtifactInTransaction(
|
||||
client,
|
||||
binding.artifact,
|
||||
).artifact;
|
||||
applyMutation(client, command.stepRunMutation);
|
||||
insertCompletion(client, completion);
|
||||
if (expectedUsage) insertUsage(client, expectedUsage);
|
||||
if (expectedPriceSettlement) {
|
||||
insertPriceSettlement(client, expectedPriceSettlement);
|
||||
}
|
||||
if (expectedQuotaSettlement) {
|
||||
insertQuotaSettlement(client, expectedQuotaSettlement);
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: completion,
|
||||
artifact,
|
||||
reference: binding.reference,
|
||||
});
|
||||
} catch (error) {
|
||||
if (began && client.isTransaction) {
|
||||
try {
|
||||
client.exec('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the original transaction failure.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,479 @@
|
||||
import type { DatabaseSync } from 'node:sqlite';
|
||||
|
||||
import { type StepRunMutation } from '@qinglong/runtime-core/step-run';
|
||||
import {
|
||||
type ModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceSettlement,
|
||||
} from '../../pricing/pricing';
|
||||
import { type ModelInvocationUsageLedgerRecord } from '../../usage/usageLedger';
|
||||
import {
|
||||
type ModelInvocationQuotaReservation,
|
||||
type ModelInvocationQuotaSettlement,
|
||||
type ModelInvocationQuotaWindowUsage,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
ModelInvocationConflictError,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
import { type ModelInvocationResolutionRecord } from '../modelInvocationResolution';
|
||||
|
||||
import type { Row } from './authority';
|
||||
import { TERMINAL_RUN_STATUSES, integer, text, unavailable } from './authority';
|
||||
|
||||
export function updateStepRun(
|
||||
client: DatabaseSync,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): void {
|
||||
const step = mutation.stepRun;
|
||||
const result = client
|
||||
.prepare(
|
||||
`UPDATE "StepRuns"
|
||||
SET status = ?, version = ?, attempt_count = ?, output_ref = ?,
|
||||
approval_request_id = ?, ready_at_ms = ?, started_at_ms = ?,
|
||||
finished_at_ms = ?, result_code = ?, error_summary = ?,
|
||||
updated_at_ms = ?, last_mutation_id = ?, step_run_digest = ?,
|
||||
step_run_json = ?
|
||||
WHERE id = ? AND run_id = ? AND version = ?
|
||||
AND step_run_digest = ? AND status = ?`,
|
||||
)
|
||||
.run(
|
||||
step.status,
|
||||
step.version,
|
||||
step.attemptCount,
|
||||
step.outputRef,
|
||||
step.approvalRequestId,
|
||||
step.readyAtMs,
|
||||
step.startedAtMs,
|
||||
step.finishedAtMs,
|
||||
step.resultCode,
|
||||
step.errorSummary,
|
||||
step.updatedAtMs,
|
||||
step.lastMutationId,
|
||||
step.stepRunDigest,
|
||||
JSON.stringify(step),
|
||||
step.id,
|
||||
step.runId,
|
||||
mutation.expectedStepRunVersion,
|
||||
mutation.expectedStepRunDigest,
|
||||
mutation.previousStatus,
|
||||
);
|
||||
if (result.changes !== 1) throw new ModelInvocationConflictError();
|
||||
}
|
||||
|
||||
export function updateRun(
|
||||
client: DatabaseSync,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): void {
|
||||
const result = client
|
||||
.prepare(
|
||||
`UPDATE "Runs"
|
||||
SET version = version + 1, event_sequence = event_sequence + 1
|
||||
WHERE id = ? AND version = ? AND event_sequence = ?`,
|
||||
)
|
||||
.run(
|
||||
mutation.runId,
|
||||
mutation.expectedRunVersion,
|
||||
mutation.expectedRunEventSequence,
|
||||
);
|
||||
if (result.changes !== 1) throw new ModelInvocationConflictError();
|
||||
}
|
||||
|
||||
export function insertRunEvent(
|
||||
client: DatabaseSync,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): void {
|
||||
const event = mutation.event;
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "RunEvents" (
|
||||
id, run_id, sequence, type, dedupe_key, actor_type, actor_id,
|
||||
attempt_id, step_run_id, payload, created_at_ms
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?)`,
|
||||
)
|
||||
.run(
|
||||
event.id,
|
||||
event.runId,
|
||||
event.sequence,
|
||||
event.type,
|
||||
event.dedupeKey!,
|
||||
event.actorType,
|
||||
event.actorId ?? null,
|
||||
mutation.stepRun.id,
|
||||
JSON.stringify(event.payload),
|
||||
event.createdAtMs,
|
||||
);
|
||||
}
|
||||
|
||||
export function insertMutation(
|
||||
client: DatabaseSync,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): void {
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "StepRunMutations" (
|
||||
mutation_id, mutation_digest, run_id, step_run_id,
|
||||
step_run_digest, event_id, event_sequence, run_version,
|
||||
step_run_json, committed_at_ms
|
||||
) VALUES (
|
||||
?, ?, ?, ?, ?, ?, ?, ?, ?,
|
||||
CAST(unixepoch('subsec') * 1000 AS INTEGER)
|
||||
)`,
|
||||
)
|
||||
.run(
|
||||
mutation.mutationId,
|
||||
mutation.mutationDigest,
|
||||
mutation.runId,
|
||||
mutation.stepRun.id,
|
||||
mutation.stepRun.stepRunDigest,
|
||||
mutation.event.id,
|
||||
mutation.event.sequence,
|
||||
mutation.expectedRunVersion + 1,
|
||||
JSON.stringify(mutation.stepRun),
|
||||
);
|
||||
}
|
||||
|
||||
export function applyMutation(
|
||||
client: DatabaseSync,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): void {
|
||||
updateStepRun(client, mutation);
|
||||
updateRun(client, mutation);
|
||||
insertRunEvent(client, mutation);
|
||||
insertMutation(client, mutation);
|
||||
}
|
||||
|
||||
export function insertStart(
|
||||
client: DatabaseSync,
|
||||
start: Readonly<ModelInvocationStartRecord>,
|
||||
): void {
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "ModelInvocationStarts" (
|
||||
invocation_id, project_id, run_id, step_run_id, trace_id,
|
||||
provider, model, policy_revision, request_digest, input_bytes,
|
||||
max_output_tokens, deadline_at_ms, admitted_at_ms, mutation_id,
|
||||
mutation_digest, run_event_id, start_digest, record_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
)
|
||||
.run(
|
||||
start.invocationId,
|
||||
start.projectId,
|
||||
start.runId,
|
||||
start.stepRunId,
|
||||
start.traceId,
|
||||
start.provider,
|
||||
start.model,
|
||||
start.policyRevision,
|
||||
start.requestDigest,
|
||||
start.inputBytes,
|
||||
start.maxOutputTokens,
|
||||
start.deadlineAtMs,
|
||||
start.admittedAtMs,
|
||||
start.stepRunMutationId,
|
||||
start.stepRunMutationDigest,
|
||||
start.runEventId,
|
||||
start.startDigest,
|
||||
JSON.stringify(start),
|
||||
);
|
||||
}
|
||||
|
||||
export function insertCompletion(
|
||||
client: DatabaseSync,
|
||||
completion: Readonly<ModelInvocationCompletionRecord>,
|
||||
): void {
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "ModelInvocationCompletions" (
|
||||
invocation_id, project_id, run_id, step_run_id, trace_id,
|
||||
start_digest, outcome, output_bytes, error_code, completed_at_ms,
|
||||
mutation_id, mutation_digest, run_event_id, completion_digest,
|
||||
record_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
)
|
||||
.run(
|
||||
completion.invocationId,
|
||||
completion.projectId,
|
||||
completion.runId,
|
||||
completion.stepRunId,
|
||||
completion.traceId,
|
||||
completion.startDigest,
|
||||
completion.outcome,
|
||||
completion.outputBytes,
|
||||
completion.errorCode,
|
||||
completion.completedAtMs,
|
||||
completion.stepRunMutationId,
|
||||
completion.stepRunMutationDigest,
|
||||
completion.runEventId,
|
||||
completion.completionDigest,
|
||||
JSON.stringify(completion),
|
||||
);
|
||||
}
|
||||
|
||||
export function insertUsage(
|
||||
client: DatabaseSync,
|
||||
usage: Readonly<ModelInvocationUsageLedgerRecord>,
|
||||
): void {
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "ModelInvocationUsageLedger" (
|
||||
invocation_id, project_id, run_id, step_run_id, trace_id,
|
||||
provider, model, policy_revision, completion_digest, outcome,
|
||||
settled_at_ms, input_bytes, output_bytes, input_tokens,
|
||||
output_tokens, total_tokens, cost_micros, ledger_digest, record_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
)
|
||||
.run(
|
||||
usage.invocationId,
|
||||
usage.projectId,
|
||||
usage.runId,
|
||||
usage.stepRunId,
|
||||
usage.traceId,
|
||||
usage.provider,
|
||||
usage.model,
|
||||
usage.policyRevision,
|
||||
usage.completionDigest,
|
||||
usage.outcome,
|
||||
usage.settledAtMs,
|
||||
usage.inputBytes,
|
||||
usage.outputBytes,
|
||||
usage.inputTokens,
|
||||
usage.outputTokens,
|
||||
usage.totalTokens,
|
||||
usage.costMicros,
|
||||
usage.ledgerDigest,
|
||||
JSON.stringify(usage),
|
||||
);
|
||||
}
|
||||
|
||||
export function insertQuotaReservation(
|
||||
client: DatabaseSync,
|
||||
reservation: Readonly<ModelInvocationQuotaReservation>,
|
||||
): void {
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "ModelInvocationQuotaReservations" (
|
||||
invocation_id, project_id, model_policy_revision,
|
||||
quota_policy_revision, window_ms, window_start_ms, window_end_ms,
|
||||
max_invocations, max_tokens, max_cost_micros, reserved_tokens,
|
||||
reserved_cost_micros, reserved_at_ms, admission_digest,
|
||||
reservation_digest, record_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
)
|
||||
.run(
|
||||
reservation.invocationId,
|
||||
reservation.projectId,
|
||||
reservation.modelPolicyRevision,
|
||||
reservation.quotaPolicyRevision,
|
||||
reservation.windowMs,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowEndMs,
|
||||
reservation.maxInvocations,
|
||||
reservation.maxTokens,
|
||||
reservation.maxCostMicros,
|
||||
reservation.reservedTokens,
|
||||
reservation.reservedCostMicros,
|
||||
reservation.reservedAtMs,
|
||||
reservation.admissionDigest,
|
||||
reservation.reservationDigest,
|
||||
JSON.stringify(reservation),
|
||||
);
|
||||
}
|
||||
|
||||
export function insertQuotaSettlement(
|
||||
client: DatabaseSync,
|
||||
settlement: Readonly<ModelInvocationQuotaSettlement>,
|
||||
): void {
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "ModelInvocationQuotaSettlements" (
|
||||
invocation_id, project_id, reservation_digest, completion_digest,
|
||||
effective_tokens, effective_cost_micros,
|
||||
retained_token_reservation, retained_cost_reservation,
|
||||
settled_at_ms, settlement_digest, record_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
)
|
||||
.run(
|
||||
settlement.invocationId,
|
||||
settlement.projectId,
|
||||
settlement.reservationDigest,
|
||||
settlement.completionDigest,
|
||||
settlement.effectiveTokens,
|
||||
settlement.effectiveCostMicros,
|
||||
settlement.retainedTokenReservation ? 1 : 0,
|
||||
settlement.retainedCostReservation ? 1 : 0,
|
||||
settlement.settledAtMs,
|
||||
settlement.settlementDigest,
|
||||
JSON.stringify(settlement),
|
||||
);
|
||||
}
|
||||
|
||||
export function insertPriceQuote(
|
||||
client: DatabaseSync,
|
||||
quote: Readonly<ModelInvocationPriceQuote>,
|
||||
): void {
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "ModelInvocationPriceQuotes" (
|
||||
invocation_id, project_id, model_policy_revision, provider, model,
|
||||
price_revision, currency, input_micros_per_million_tokens,
|
||||
output_micros_per_million_tokens, max_total_tokens, max_output_tokens,
|
||||
reserved_cost_micros, catalog_digest, quote_digest, record_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
)
|
||||
.run(
|
||||
quote.invocationId,
|
||||
quote.projectId,
|
||||
quote.modelPolicyRevision,
|
||||
quote.provider,
|
||||
quote.model,
|
||||
quote.priceRevision,
|
||||
quote.currency,
|
||||
quote.inputMicrosPerMillionTokens,
|
||||
quote.outputMicrosPerMillionTokens,
|
||||
quote.maxTotalTokens,
|
||||
quote.maxOutputTokens,
|
||||
quote.reservedCostMicros,
|
||||
quote.catalogDigest,
|
||||
quote.quoteDigest,
|
||||
JSON.stringify(quote),
|
||||
);
|
||||
}
|
||||
|
||||
export function insertPriceSettlement(
|
||||
client: DatabaseSync,
|
||||
settlement: Readonly<ModelInvocationPriceSettlement>,
|
||||
): void {
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "ModelInvocationPriceSettlements" (
|
||||
invocation_id, project_id, quote_digest, completion_digest, currency,
|
||||
input_tokens, output_tokens, cost_micros, settled_at_ms,
|
||||
settlement_digest, record_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
)
|
||||
.run(
|
||||
settlement.invocationId,
|
||||
settlement.projectId,
|
||||
settlement.quoteDigest,
|
||||
settlement.completionDigest,
|
||||
settlement.currency,
|
||||
settlement.inputTokens,
|
||||
settlement.outputTokens,
|
||||
settlement.costMicros,
|
||||
settlement.settledAtMs,
|
||||
settlement.settlementDigest,
|
||||
JSON.stringify(settlement),
|
||||
);
|
||||
}
|
||||
|
||||
export function quotaWindowUsage(
|
||||
client: DatabaseSync,
|
||||
projectId: string,
|
||||
windowStartMs: number,
|
||||
windowMs: number,
|
||||
): Readonly<ModelInvocationQuotaWindowUsage> {
|
||||
const row = client
|
||||
.prepare(
|
||||
`SELECT
|
||||
COUNT(*) AS "invocationCount",
|
||||
COALESCE(SUM(COALESCE(
|
||||
settlement.effective_tokens, reservation.reserved_tokens
|
||||
)), 0) AS "effectiveTokens",
|
||||
COALESCE(SUM(COALESCE(
|
||||
settlement.effective_cost_micros,
|
||||
reservation.reserved_cost_micros,
|
||||
0
|
||||
)), 0) AS "effectiveCostMicros",
|
||||
COALESCE(SUM(CASE WHEN
|
||||
settlement.effective_cost_micros IS NULL
|
||||
AND reservation.reserved_cost_micros IS NULL
|
||||
THEN 1 ELSE 0 END), 0) AS "unknownCostInvocations"
|
||||
FROM "ModelInvocationQuotaReservations" AS reservation
|
||||
LEFT JOIN "ModelInvocationQuotaSettlements" AS settlement
|
||||
ON settlement.invocation_id = reservation.invocation_id
|
||||
WHERE reservation.project_id = ?
|
||||
AND reservation.window_start_ms = ?
|
||||
AND reservation.window_ms = ?`,
|
||||
)
|
||||
.get(projectId, windowStartMs, windowMs) as Row | undefined;
|
||||
if (!row) throw unavailable();
|
||||
return Object.freeze({
|
||||
projectId,
|
||||
windowStartMs,
|
||||
windowEndMs: windowStartMs + windowMs,
|
||||
invocationCount: integer(row, 'invocationCount'),
|
||||
effectiveTokens: integer(row, 'effectiveTokens'),
|
||||
effectiveCostMicros: integer(row, 'effectiveCostMicros'),
|
||||
unknownCostInvocations: integer(row, 'unknownCostInvocations'),
|
||||
});
|
||||
}
|
||||
|
||||
export function insertResolution(
|
||||
client: DatabaseSync,
|
||||
resolution: Readonly<ModelInvocationResolutionRecord>,
|
||||
): void {
|
||||
client
|
||||
.prepare(
|
||||
`INSERT INTO "ModelInvocationResolutions" (
|
||||
resolution_id, invocation_id, project_id, run_id, step_run_id,
|
||||
trace_id, completion_digest, decision, resolved_by_user_id,
|
||||
resolved_at_ms, mutation_id, mutation_digest, run_event_id,
|
||||
resolution_digest, record_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
)
|
||||
.run(
|
||||
resolution.resolutionId,
|
||||
resolution.invocationId,
|
||||
resolution.projectId,
|
||||
resolution.runId,
|
||||
resolution.stepRunId,
|
||||
resolution.traceId,
|
||||
resolution.completionDigest,
|
||||
resolution.decision,
|
||||
resolution.resolvedByUserId,
|
||||
resolution.resolvedAtMs,
|
||||
resolution.stepRunMutationId,
|
||||
resolution.stepRunMutationDigest,
|
||||
resolution.runEventId,
|
||||
resolution.resolutionDigest,
|
||||
JSON.stringify(resolution),
|
||||
);
|
||||
}
|
||||
|
||||
export function assertCurrent(
|
||||
client: DatabaseSync,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
projectId: string,
|
||||
): void {
|
||||
const rows = client
|
||||
.prepare(
|
||||
`SELECT
|
||||
step.kind AS "stepKind", step.status AS "stepStatus",
|
||||
step.version AS "stepVersion",
|
||||
step.step_run_digest AS "stepDigest",
|
||||
run.project_id AS "projectId", run.status AS "runStatus",
|
||||
run.version AS "runVersion",
|
||||
run.event_sequence AS "runEventSequence"
|
||||
FROM "StepRuns" AS step
|
||||
JOIN "Runs" AS run ON run.id = step.run_id
|
||||
WHERE step.id = ? AND step.run_id = ?
|
||||
LIMIT 2`,
|
||||
)
|
||||
.all(mutation.stepRun.id, mutation.runId) as Row[];
|
||||
const row = rows[0];
|
||||
if (
|
||||
rows.length !== 1 ||
|
||||
!row ||
|
||||
text(row, 'stepKind') !== 'model' ||
|
||||
text(row, 'stepStatus') !== mutation.previousStatus ||
|
||||
integer(row, 'stepVersion') !== mutation.expectedStepRunVersion ||
|
||||
text(row, 'stepDigest') !== mutation.expectedStepRunDigest ||
|
||||
text(row, 'projectId') !== projectId ||
|
||||
integer(row, 'runVersion') !== mutation.expectedRunVersion ||
|
||||
integer(row, 'runEventSequence') !== mutation.expectedRunEventSequence ||
|
||||
TERMINAL_RUN_STATUSES.has(text(row, 'runStatus'))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
import type { DatabaseSync } from 'node:sqlite';
|
||||
|
||||
import type { Row } from './authority';
|
||||
import {
|
||||
COMPLETION_SELECT,
|
||||
RESOLUTION_SELECT,
|
||||
START_SELECT,
|
||||
USAGE_SELECT,
|
||||
} from './codec';
|
||||
|
||||
export function startRows(
|
||||
client: DatabaseSync,
|
||||
where: string,
|
||||
values: readonly (string | number)[],
|
||||
): readonly Row[] {
|
||||
return client
|
||||
.prepare(
|
||||
`SELECT ${START_SELECT}
|
||||
FROM "ModelInvocationStarts" AS start
|
||||
JOIN "StepRunMutations" AS mutation
|
||||
ON mutation.mutation_id = start.mutation_id
|
||||
JOIN "RunEvents" AS event ON event.id = start.run_event_id
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
)
|
||||
.all(...values) as Row[];
|
||||
}
|
||||
|
||||
export function completionRows(
|
||||
client: DatabaseSync,
|
||||
where: string,
|
||||
values: readonly (string | number)[],
|
||||
): readonly Row[] {
|
||||
return client
|
||||
.prepare(
|
||||
`SELECT ${COMPLETION_SELECT}
|
||||
FROM "ModelInvocationCompletions" AS completion
|
||||
JOIN "StepRunMutations" AS mutation
|
||||
ON mutation.mutation_id = completion.mutation_id
|
||||
JOIN "RunEvents" AS event ON event.id = completion.run_event_id
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
)
|
||||
.all(...values) as Row[];
|
||||
}
|
||||
|
||||
export function usageRows(
|
||||
client: DatabaseSync,
|
||||
where: string,
|
||||
values: readonly (string | number)[],
|
||||
limit = 2,
|
||||
): readonly Row[] {
|
||||
return client
|
||||
.prepare(
|
||||
`SELECT ${USAGE_SELECT}
|
||||
FROM "ModelInvocationUsageLedger" AS usage
|
||||
WHERE ${where}
|
||||
ORDER BY usage.settled_at_ms, usage.invocation_id
|
||||
LIMIT ?`,
|
||||
)
|
||||
.all(...values, limit) as Row[];
|
||||
}
|
||||
|
||||
export function quotaReservationRows(
|
||||
client: DatabaseSync,
|
||||
where: string,
|
||||
values: readonly (string | number)[],
|
||||
limit = 2,
|
||||
): readonly Row[] {
|
||||
return client
|
||||
.prepare(
|
||||
`SELECT reservation.record_json AS "recordJson"
|
||||
FROM "ModelInvocationQuotaReservations" AS reservation
|
||||
WHERE ${where}
|
||||
ORDER BY reservation.window_start_ms, reservation.invocation_id
|
||||
LIMIT ?`,
|
||||
)
|
||||
.all(...values, limit) as Row[];
|
||||
}
|
||||
|
||||
export function quotaSettlementRows(
|
||||
client: DatabaseSync,
|
||||
where: string,
|
||||
values: readonly (string | number)[],
|
||||
): readonly Row[] {
|
||||
return client
|
||||
.prepare(
|
||||
`SELECT settlement.record_json AS "recordJson"
|
||||
FROM "ModelInvocationQuotaSettlements" AS settlement
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
)
|
||||
.all(...values) as Row[];
|
||||
}
|
||||
|
||||
export function priceQuoteRows(
|
||||
client: DatabaseSync,
|
||||
where: string,
|
||||
values: readonly (string | number)[],
|
||||
): readonly Row[] {
|
||||
return client
|
||||
.prepare(
|
||||
`SELECT quote.record_json AS "recordJson"
|
||||
FROM "ModelInvocationPriceQuotes" AS quote
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
)
|
||||
.all(...values) as Row[];
|
||||
}
|
||||
|
||||
export function priceSettlementRows(
|
||||
client: DatabaseSync,
|
||||
where: string,
|
||||
values: readonly (string | number)[],
|
||||
): readonly Row[] {
|
||||
return client
|
||||
.prepare(
|
||||
`SELECT settlement.record_json AS "recordJson"
|
||||
FROM "ModelInvocationPriceSettlements" AS settlement
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
)
|
||||
.all(...values) as Row[];
|
||||
}
|
||||
|
||||
export function resolutionRows(
|
||||
client: DatabaseSync,
|
||||
where: string,
|
||||
values: readonly (string | number)[],
|
||||
): readonly Row[] {
|
||||
return client
|
||||
.prepare(
|
||||
`SELECT ${RESOLUTION_SELECT}
|
||||
FROM "ModelInvocationResolutions" AS resolution
|
||||
JOIN "StepRunMutations" AS mutation
|
||||
ON mutation.mutation_id = resolution.mutation_id
|
||||
JOIN "RunEvents" AS event ON event.id = resolution.run_event_id
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
)
|
||||
.all(...values) as Row[];
|
||||
}
|
||||
+356
@@ -0,0 +1,356 @@
|
||||
import type { DatabaseSync } from 'node:sqlite';
|
||||
|
||||
import {
|
||||
type ModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceSettlement,
|
||||
} from '../../pricing/pricing';
|
||||
import { type PluginPackagePromptOutputArtifact } from '../../prompt-output/pluginPackagePromptOutputArtifact';
|
||||
import type { PluginPackagePromptOutputArtifactTombstone } from '../../prompt-output/pluginPackagePromptOutputRetention';
|
||||
import { readLocalPluginPackagePromptOutputArtifactInTransaction } from '../../prompt-output/storage/localPluginPackagePromptOutputArtifactRepository';
|
||||
import { readLocalPluginPackagePromptOutputArtifactTombstoneInTransaction } from '../../prompt-output/storage/localPluginPackagePromptOutputRetentionRepository';
|
||||
import {
|
||||
MAX_MODEL_INVOCATION_USAGE_SUMMARY_ROWS,
|
||||
ModelInvocationUsageSummaryLimitExceededError,
|
||||
normalizeModelInvocationUsageLedgerQuery,
|
||||
normalizeModelInvocationUsageLedgerSummaryQuery,
|
||||
type ModelInvocationUsageLedgerPage,
|
||||
type ModelInvocationUsageLedgerQuery,
|
||||
type ModelInvocationUsageLedgerRecord,
|
||||
type ModelInvocationUsageLedgerSummary,
|
||||
type ModelInvocationUsageLedgerSummaryQuery,
|
||||
} from '../../usage/usageLedger';
|
||||
import {
|
||||
type ModelInvocationQuotaReservation,
|
||||
type ModelInvocationQuotaSettlement,
|
||||
type ModelInvocationQuotaWindowUsage,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
|
||||
import type { LocalModelInvocationOperationAuthority, Row } from './authority';
|
||||
import {
|
||||
enqueueLocalModelInvocation,
|
||||
identifier,
|
||||
integer,
|
||||
unavailable,
|
||||
} from './authority';
|
||||
import {
|
||||
parseCompletion,
|
||||
parsePriceQuote,
|
||||
parsePriceSettlement,
|
||||
parseQuotaReservation,
|
||||
parseQuotaSettlement,
|
||||
parseStart,
|
||||
parseUsage,
|
||||
} from './codec';
|
||||
import { quotaWindowUsage } from './mutations';
|
||||
import {
|
||||
completionRows,
|
||||
priceQuoteRows,
|
||||
priceSettlementRows,
|
||||
quotaReservationRows,
|
||||
quotaSettlementRows,
|
||||
startRows,
|
||||
usageRows,
|
||||
} from './queries';
|
||||
|
||||
export function findStartOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationStartRecord> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const rows = startRows(client, 'start.invocation_id = ?', [invocationId]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseStart(rows[0]) : null;
|
||||
});
|
||||
}
|
||||
|
||||
export function findCompletionOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationCompletionRecord> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const rows = completionRows(client, 'completion.invocation_id = ?', [
|
||||
invocationId,
|
||||
]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseCompletion(rows[0]) : null;
|
||||
});
|
||||
}
|
||||
|
||||
export function findPromptOutputArtifactOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
artifactIdValue: string,
|
||||
): Promise<Readonly<PluginPackagePromptOutputArtifact> | null> {
|
||||
const artifactId = identifier(artifactIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () =>
|
||||
readLocalPluginPackagePromptOutputArtifactInTransaction(client, artifactId),
|
||||
);
|
||||
}
|
||||
|
||||
export function findPromptOutputArtifactTombstoneOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
artifactIdValue: string,
|
||||
): Promise<Readonly<PluginPackagePromptOutputArtifactTombstone> | null> {
|
||||
const artifactId = identifier(artifactIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () =>
|
||||
readLocalPluginPackagePromptOutputArtifactTombstoneInTransaction(
|
||||
client,
|
||||
artifactId,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
export function findUsageOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerRecord> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const rows = usageRows(client, 'usage.invocation_id = ?', [invocationId]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseUsage(rows[0]) : null;
|
||||
});
|
||||
}
|
||||
|
||||
export function findPriceQuoteOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationPriceQuote> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const rows = priceQuoteRows(client, 'quote.invocation_id = ?', [
|
||||
invocationId,
|
||||
]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parsePriceQuote(rows[0]) : null;
|
||||
});
|
||||
}
|
||||
|
||||
export function findPriceSettlementOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationPriceSettlement> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const quotes = priceQuoteRows(client, 'quote.invocation_id = ?', [
|
||||
invocationId,
|
||||
]);
|
||||
const completions = completionRows(client, 'completion.invocation_id = ?', [
|
||||
invocationId,
|
||||
]);
|
||||
const settlements = priceSettlementRows(
|
||||
client,
|
||||
'settlement.invocation_id = ?',
|
||||
[invocationId],
|
||||
);
|
||||
if (quotes.length > 1 || completions.length > 1 || settlements.length > 1) {
|
||||
throw unavailable();
|
||||
}
|
||||
if (!settlements[0]) return null;
|
||||
if (!quotes[0] || !completions[0]) throw unavailable();
|
||||
return parsePriceSettlement(
|
||||
settlements[0],
|
||||
parsePriceQuote(quotes[0]),
|
||||
parseCompletion(completions[0]),
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
export function findQuotaReservationOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationQuotaReservation> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const rows = quotaReservationRows(client, 'reservation.invocation_id = ?', [
|
||||
invocationId,
|
||||
]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseQuotaReservation(rows[0]) : null;
|
||||
});
|
||||
}
|
||||
|
||||
export function findQuotaSettlementOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationQuotaSettlement> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const reservations = quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = ?',
|
||||
[invocationId],
|
||||
);
|
||||
const completions = completionRows(client, 'completion.invocation_id = ?', [
|
||||
invocationId,
|
||||
]);
|
||||
const settlements = quotaSettlementRows(
|
||||
client,
|
||||
'settlement.invocation_id = ?',
|
||||
[invocationId],
|
||||
);
|
||||
if (
|
||||
reservations.length > 1 ||
|
||||
completions.length > 1 ||
|
||||
settlements.length > 1
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
if (!settlements[0]) return null;
|
||||
if (!reservations[0] || !completions[0]) throw unavailable();
|
||||
return parseQuotaSettlement(
|
||||
settlements[0],
|
||||
parseQuotaReservation(reservations[0]),
|
||||
parseCompletion(completions[0]),
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
export function readQuotaWindowUsageOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
projectIdValue: string,
|
||||
atMsValue?: number,
|
||||
): Promise<Readonly<ModelInvocationQuotaWindowUsage> | null> {
|
||||
const projectId = identifier(projectIdValue);
|
||||
if (
|
||||
atMsValue !== undefined &&
|
||||
(!Number.isSafeInteger(atMsValue) || atMsValue < 0)
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const atMs =
|
||||
atMsValue ??
|
||||
integer(
|
||||
client
|
||||
.prepare(
|
||||
`SELECT CAST(unixepoch('subsec') * 1000 AS INTEGER) AS "atMs"`,
|
||||
)
|
||||
.get() as Row,
|
||||
'atMs',
|
||||
);
|
||||
const row = client
|
||||
.prepare(
|
||||
`SELECT reservation.record_json AS "recordJson"
|
||||
FROM "ModelInvocationQuotaReservations" AS reservation
|
||||
WHERE reservation.project_id = ?
|
||||
AND reservation.window_start_ms <= ?
|
||||
AND reservation.window_end_ms > ?
|
||||
ORDER BY reservation.reserved_at_ms DESC,
|
||||
reservation.invocation_id DESC
|
||||
LIMIT 1`,
|
||||
)
|
||||
.get(projectId, atMs, atMs) as Row | undefined;
|
||||
if (!row) return null;
|
||||
const reservation = parseQuotaReservation(row);
|
||||
return quotaWindowUsage(
|
||||
client,
|
||||
projectId,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowMs,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
export function listProjectUsageOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
queryValue: ModelInvocationUsageLedgerQuery,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerPage>> {
|
||||
const query = normalizeModelInvocationUsageLedgerQuery(queryValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const cursor = query.after;
|
||||
const rows = usageRows(
|
||||
client,
|
||||
`usage.project_id = ?
|
||||
AND usage.settled_at_ms >= ? AND usage.settled_at_ms < ?
|
||||
${
|
||||
cursor
|
||||
? `AND (
|
||||
usage.settled_at_ms > ? OR
|
||||
(usage.settled_at_ms = ? AND usage.invocation_id > ?)
|
||||
)`
|
||||
: ''
|
||||
}`,
|
||||
[
|
||||
query.projectId,
|
||||
query.fromMsInclusive,
|
||||
query.toMsExclusive,
|
||||
...(cursor
|
||||
? [cursor.settledAtMs, cursor.settledAtMs, cursor.invocationId]
|
||||
: []),
|
||||
],
|
||||
query.limit + 1,
|
||||
);
|
||||
return Object.freeze({
|
||||
records: Object.freeze(
|
||||
rows.slice(0, query.limit).map((row) => parseUsage(row)),
|
||||
),
|
||||
hasMore: rows.length > query.limit,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
export function summarizeProjectUsageOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
queryValue: ModelInvocationUsageLedgerSummaryQuery,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerSummary>> {
|
||||
const query = normalizeModelInvocationUsageLedgerSummaryQuery(queryValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const row = client
|
||||
.prepare(
|
||||
`SELECT
|
||||
COUNT(*) AS "invocationCount",
|
||||
COALESCE(SUM(input_tokens), 0) AS "inputTokens",
|
||||
COALESCE(SUM(output_tokens), 0) AS "outputTokens",
|
||||
COALESCE(SUM(total_tokens), 0) AS "totalTokens",
|
||||
COALESCE(SUM(cost_micros), 0) AS "knownCostMicros",
|
||||
COALESCE(SUM(CASE WHEN cost_micros IS NULL THEN 1 ELSE 0 END), 0)
|
||||
AS "unknownCostInvocations"
|
||||
FROM (
|
||||
SELECT input_tokens, output_tokens, total_tokens, cost_micros
|
||||
FROM "ModelInvocationUsageLedger"
|
||||
WHERE project_id = ?
|
||||
AND settled_at_ms >= ? AND settled_at_ms < ?
|
||||
ORDER BY settled_at_ms, invocation_id
|
||||
LIMIT ?
|
||||
) AS bounded_usage`,
|
||||
)
|
||||
.get(
|
||||
query.projectId,
|
||||
query.fromMsInclusive,
|
||||
query.toMsExclusive,
|
||||
MAX_MODEL_INVOCATION_USAGE_SUMMARY_ROWS + 1,
|
||||
) as Row | undefined;
|
||||
if (!row) throw unavailable();
|
||||
if (
|
||||
integer(row, 'invocationCount') > MAX_MODEL_INVOCATION_USAGE_SUMMARY_ROWS
|
||||
) {
|
||||
throw new ModelInvocationUsageSummaryLimitExceededError();
|
||||
}
|
||||
return Object.freeze({
|
||||
invocationCount: integer(row, 'invocationCount'),
|
||||
inputTokens: integer(row, 'inputTokens'),
|
||||
outputTokens: integer(row, 'outputTokens'),
|
||||
totalTokens: integer(row, 'totalTokens'),
|
||||
knownCostMicros: integer(row, 'knownCostMicros'),
|
||||
unknownCostInvocations: integer(row, 'unknownCostInvocations'),
|
||||
});
|
||||
});
|
||||
}
|
||||
+198
@@ -0,0 +1,198 @@
|
||||
import type { DatabaseSync } from 'node:sqlite';
|
||||
|
||||
import {
|
||||
ModelInvocationConflictError,
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationAuthoritySnapshot,
|
||||
type ModelInvocationRecoveryPage,
|
||||
} from '../modelInvocation';
|
||||
import {
|
||||
normalizeModelInvocationResolutionCommand,
|
||||
type ModelInvocationResolutionCommand,
|
||||
type ModelInvocationResolutionRecord,
|
||||
} from '../modelInvocationResolution';
|
||||
|
||||
import type { LocalModelInvocationOperationAuthority, Row } from './authority';
|
||||
import {
|
||||
TERMINAL_RUN_STATUSES,
|
||||
enqueueLocalModelInvocation,
|
||||
identifier,
|
||||
integer,
|
||||
recoveryLimit,
|
||||
text,
|
||||
unavailable,
|
||||
} from './authority';
|
||||
import {
|
||||
START_SELECT,
|
||||
parseAuthority,
|
||||
parseCompletion,
|
||||
parseResolution,
|
||||
parseStart,
|
||||
} from './codec';
|
||||
import { applyMutation, assertCurrent, insertResolution } from './mutations';
|
||||
import { completionRows, resolutionRows } from './queries';
|
||||
|
||||
export function findResolutionOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationResolutionRecord> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const rows = resolutionRows(client, 'resolution.invocation_id = ?', [
|
||||
invocationId,
|
||||
]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseResolution(rows[0]) : null;
|
||||
});
|
||||
}
|
||||
|
||||
export function readAuthorityOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
identity: Readonly<{
|
||||
projectId: string;
|
||||
runId: string;
|
||||
stepRunId: string;
|
||||
}>,
|
||||
): Promise<Readonly<ModelInvocationAuthoritySnapshot> | null> {
|
||||
const projectId = identifier(identity?.projectId);
|
||||
const runId = identifier(identity?.runId);
|
||||
const stepRunId = identifier(identity?.stepRunId);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const rows = client
|
||||
.prepare(
|
||||
`SELECT
|
||||
run.project_id AS "projectId", run.id AS "runId",
|
||||
run.status AS "runStatus", run.version AS "runVersion",
|
||||
run.event_sequence AS "runEventSequence",
|
||||
step.id AS "stepRunId", step.status AS "stepStatus",
|
||||
step.version AS "stepVersion",
|
||||
step.step_run_digest AS "stepDigest",
|
||||
step.step_run_json AS "stepRunJson"
|
||||
FROM "Runs" AS run
|
||||
JOIN "StepRuns" AS step ON step.run_id = run.id
|
||||
WHERE run.project_id = ? AND run.id = ? AND step.id = ?
|
||||
LIMIT 2`,
|
||||
)
|
||||
.all(projectId, runId, stepRunId) as Row[];
|
||||
if (rows.length > 1) throw unavailable();
|
||||
const row = rows[0];
|
||||
if (!row || TERMINAL_RUN_STATUSES.has(text(row, 'runStatus'))) {
|
||||
return null;
|
||||
}
|
||||
return parseAuthority(row);
|
||||
});
|
||||
}
|
||||
|
||||
export function listIncompleteOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
limitValue: number,
|
||||
): Promise<Readonly<ModelInvocationRecoveryPage>> {
|
||||
const limit = recoveryLimit(limitValue);
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
const observed = client
|
||||
.prepare(
|
||||
`SELECT CAST(unixepoch('subsec') * 1000 AS INTEGER)
|
||||
AS "observedAtMs"`,
|
||||
)
|
||||
.get() as Row | undefined;
|
||||
if (!observed) throw unavailable();
|
||||
const observedAtMs = integer(observed, 'observedAtMs');
|
||||
const rows = client
|
||||
.prepare(
|
||||
`SELECT ${START_SELECT}
|
||||
FROM "ModelInvocationStarts" AS start
|
||||
JOIN "StepRunMutations" AS mutation
|
||||
ON mutation.mutation_id = start.mutation_id
|
||||
JOIN "RunEvents" AS event ON event.id = start.run_event_id
|
||||
LEFT JOIN "ModelInvocationCompletions" AS completion
|
||||
ON completion.invocation_id = start.invocation_id
|
||||
WHERE completion.invocation_id IS NULL
|
||||
AND start.deadline_at_ms <= ?
|
||||
ORDER BY start.deadline_at_ms, start.invocation_id
|
||||
LIMIT ?`,
|
||||
)
|
||||
.all(observedAtMs, limit + 1) as Row[];
|
||||
const hasMore = rows.length > limit;
|
||||
return Object.freeze({
|
||||
observedAtMs,
|
||||
candidates: Object.freeze(
|
||||
rows.slice(0, limit).map((row) => parseStart(row)),
|
||||
),
|
||||
hasMore,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
export function resolveOperation(
|
||||
authority: LocalModelInvocationOperationAuthority,
|
||||
client: DatabaseSync,
|
||||
commandValue: ModelInvocationResolutionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationResolutionRecord>>
|
||||
> {
|
||||
const command = normalizeModelInvocationResolutionCommand(commandValue);
|
||||
const resolution = command.resolution;
|
||||
return enqueueLocalModelInvocation(authority, () => {
|
||||
let began = false;
|
||||
try {
|
||||
client.exec('BEGIN IMMEDIATE');
|
||||
began = true;
|
||||
const existing = resolutionRows(
|
||||
client,
|
||||
`resolution.invocation_id = ? OR
|
||||
resolution.resolution_id = ? OR
|
||||
resolution.mutation_id = ? OR resolution.run_event_id = ?`,
|
||||
[
|
||||
resolution.invocationId,
|
||||
resolution.resolutionId,
|
||||
resolution.stepRunMutationId,
|
||||
resolution.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseResolution(existing[0]);
|
||||
if (JSON.stringify(stored) !== JSON.stringify(resolution)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
const completions = completionRows(
|
||||
client,
|
||||
'completion.invocation_id = ?',
|
||||
[resolution.invocationId],
|
||||
);
|
||||
if (
|
||||
completions.length !== 1 ||
|
||||
JSON.stringify(parseCompletion(completions[0]!)) !==
|
||||
JSON.stringify(command.completion) ||
|
||||
command.completion.outcome !== 'outcome_unknown'
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
assertCurrent(client, command.stepRunMutation, resolution.projectId);
|
||||
applyMutation(client, command.stepRunMutation);
|
||||
insertResolution(client, resolution);
|
||||
client.exec('COMMIT');
|
||||
began = false;
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: resolution,
|
||||
});
|
||||
} catch (error) {
|
||||
if (began && client.isTransaction) {
|
||||
try {
|
||||
client.exec('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the original transaction failure.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,328 @@
|
||||
import type { DatabaseSync } from 'node:sqlite';
|
||||
|
||||
import {
|
||||
type ModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceSettlement,
|
||||
type PricingAwareModelInvocationRepository,
|
||||
} from '../../pricing/pricing';
|
||||
import { type PluginPackagePromptOutputArtifact } from '../../prompt-output/pluginPackagePromptOutputArtifact';
|
||||
import {
|
||||
type CommitPluginPackagePromptOutputResult,
|
||||
type PluginPackagePromptOutputCompletionRepository,
|
||||
} from '../../prompt-output/pluginPackagePromptOutputCompletion';
|
||||
import type { PluginPackagePromptOutputArtifactTombstone } from '../../prompt-output/pluginPackagePromptOutputRetention';
|
||||
import {
|
||||
type ModelInvocationUsageLedgerPage,
|
||||
type ModelInvocationUsageLedgerQuery,
|
||||
type ModelInvocationUsageLedgerRecord,
|
||||
type ModelInvocationUsageLedgerRepository,
|
||||
type ModelInvocationUsageLedgerSummary,
|
||||
type ModelInvocationUsageLedgerSummaryQuery,
|
||||
} from '../../usage/usageLedger';
|
||||
import {
|
||||
type ModelInvocationQuotaAdmission,
|
||||
type ModelInvocationQuotaReservation,
|
||||
type ModelInvocationQuotaSettlement,
|
||||
type ModelInvocationQuotaWindowUsage,
|
||||
type QuotaAwareModelInvocationRepository,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationAuthoritySnapshot,
|
||||
type ModelInvocationCompletionCommand,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationRecoveryPage,
|
||||
type ModelInvocationRepository,
|
||||
type ModelInvocationStartCommand,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
import {
|
||||
type ModelInvocationResolutionCommand,
|
||||
type ModelInvocationResolutionRecord,
|
||||
type ModelInvocationResolutionRepository,
|
||||
} from '../modelInvocationResolution';
|
||||
|
||||
import {
|
||||
admitOperation,
|
||||
admitWithPricingOperation,
|
||||
admitWithQuotaOperation,
|
||||
} from './admissionOperations';
|
||||
import {
|
||||
PrivateLocalAuthority,
|
||||
isAuthority,
|
||||
type LocalModelInvocationOperationAuthority,
|
||||
} from './authority';
|
||||
import {
|
||||
completeOperation,
|
||||
completeWithPricingOperation,
|
||||
completeWithPromptOutputArtifactOperation,
|
||||
completeWithQuotaOperation,
|
||||
} from './completionOperations';
|
||||
import {
|
||||
findCompletionOperation,
|
||||
findPriceQuoteOperation,
|
||||
findPriceSettlementOperation,
|
||||
findPromptOutputArtifactOperation,
|
||||
findPromptOutputArtifactTombstoneOperation,
|
||||
findQuotaReservationOperation,
|
||||
findQuotaSettlementOperation,
|
||||
findStartOperation,
|
||||
findUsageOperation,
|
||||
listProjectUsageOperation,
|
||||
readQuotaWindowUsageOperation,
|
||||
summarizeProjectUsageOperation,
|
||||
} from './readOperations';
|
||||
import {
|
||||
findResolutionOperation,
|
||||
listIncompleteOperation,
|
||||
readAuthorityOperation,
|
||||
resolveOperation,
|
||||
} from './recoveryResolutionOperations';
|
||||
|
||||
export class LocalModelInvocationRepository
|
||||
implements
|
||||
ModelInvocationRepository,
|
||||
ModelInvocationResolutionRepository,
|
||||
ModelInvocationUsageLedgerRepository,
|
||||
QuotaAwareModelInvocationRepository,
|
||||
PricingAwareModelInvocationRepository,
|
||||
PluginPackagePromptOutputCompletionRepository
|
||||
{
|
||||
readonly #authority: LocalModelInvocationOperationAuthority;
|
||||
readonly #client: DatabaseSync;
|
||||
|
||||
constructor(
|
||||
authority: LocalModelInvocationOperationAuthority | DatabaseSync,
|
||||
) {
|
||||
this.#authority = isAuthority(authority)
|
||||
? authority
|
||||
: new PrivateLocalAuthority(authority);
|
||||
this.#client = this.#authority.client;
|
||||
}
|
||||
|
||||
findStart(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationStartRecord> | null> {
|
||||
return findStartOperation(this.#authority, this.#client, invocationIdValue);
|
||||
}
|
||||
|
||||
findCompletion(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationCompletionRecord> | null> {
|
||||
return findCompletionOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
invocationIdValue,
|
||||
);
|
||||
}
|
||||
|
||||
findPromptOutputArtifact(
|
||||
artifactIdValue: string,
|
||||
): Promise<Readonly<PluginPackagePromptOutputArtifact> | null> {
|
||||
return findPromptOutputArtifactOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
artifactIdValue,
|
||||
);
|
||||
}
|
||||
|
||||
findPromptOutputArtifactTombstone(
|
||||
artifactIdValue: string,
|
||||
): Promise<Readonly<PluginPackagePromptOutputArtifactTombstone> | null> {
|
||||
return findPromptOutputArtifactTombstoneOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
artifactIdValue,
|
||||
);
|
||||
}
|
||||
|
||||
findUsage(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerRecord> | null> {
|
||||
return findUsageOperation(this.#authority, this.#client, invocationIdValue);
|
||||
}
|
||||
|
||||
findPriceQuote(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationPriceQuote> | null> {
|
||||
return findPriceQuoteOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
invocationIdValue,
|
||||
);
|
||||
}
|
||||
|
||||
findPriceSettlement(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationPriceSettlement> | null> {
|
||||
return findPriceSettlementOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
invocationIdValue,
|
||||
);
|
||||
}
|
||||
|
||||
findQuotaReservation(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationQuotaReservation> | null> {
|
||||
return findQuotaReservationOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
invocationIdValue,
|
||||
);
|
||||
}
|
||||
|
||||
findQuotaSettlement(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationQuotaSettlement> | null> {
|
||||
return findQuotaSettlementOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
invocationIdValue,
|
||||
);
|
||||
}
|
||||
|
||||
readQuotaWindowUsage(
|
||||
projectIdValue: string,
|
||||
atMsValue?: number,
|
||||
): Promise<Readonly<ModelInvocationQuotaWindowUsage> | null> {
|
||||
return readQuotaWindowUsageOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
projectIdValue,
|
||||
atMsValue,
|
||||
);
|
||||
}
|
||||
|
||||
listProjectUsage(
|
||||
queryValue: ModelInvocationUsageLedgerQuery,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerPage>> {
|
||||
return listProjectUsageOperation(this.#authority, this.#client, queryValue);
|
||||
}
|
||||
|
||||
summarizeProjectUsage(
|
||||
queryValue: ModelInvocationUsageLedgerSummaryQuery,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerSummary>> {
|
||||
return summarizeProjectUsageOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
queryValue,
|
||||
);
|
||||
}
|
||||
|
||||
findResolution(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationResolutionRecord> | null> {
|
||||
return findResolutionOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
invocationIdValue,
|
||||
);
|
||||
}
|
||||
|
||||
readAuthority(
|
||||
identity: Readonly<{
|
||||
projectId: string;
|
||||
runId: string;
|
||||
stepRunId: string;
|
||||
}>,
|
||||
): Promise<Readonly<ModelInvocationAuthoritySnapshot> | null> {
|
||||
return readAuthorityOperation(this.#authority, this.#client, identity);
|
||||
}
|
||||
|
||||
listIncomplete(
|
||||
limitValue: number,
|
||||
): Promise<Readonly<ModelInvocationRecoveryPage>> {
|
||||
return listIncompleteOperation(this.#authority, this.#client, limitValue);
|
||||
}
|
||||
|
||||
admit(
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>
|
||||
> {
|
||||
return admitOperation(this.#authority, this.#client, commandValue);
|
||||
}
|
||||
|
||||
admitWithQuota(
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
admissionValue: ModelInvocationQuotaAdmission,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>
|
||||
> {
|
||||
return admitWithQuotaOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
commandValue,
|
||||
admissionValue,
|
||||
);
|
||||
}
|
||||
|
||||
admitWithPricing(
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
quoteValue: ModelInvocationPriceQuote,
|
||||
admissionValue?: ModelInvocationQuotaAdmission,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>
|
||||
> {
|
||||
return admitWithPricingOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
commandValue,
|
||||
quoteValue,
|
||||
admissionValue,
|
||||
);
|
||||
}
|
||||
|
||||
complete(
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
return completeOperation(this.#authority, this.#client, commandValue);
|
||||
}
|
||||
|
||||
completeWithQuota(
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
return completeWithQuotaOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
commandValue,
|
||||
);
|
||||
}
|
||||
|
||||
completeWithPricing(
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
return completeWithPricingOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
commandValue,
|
||||
);
|
||||
}
|
||||
|
||||
completeWithPromptOutputArtifact(
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
artifactValue: PluginPackagePromptOutputArtifact,
|
||||
): Promise<Readonly<CommitPluginPackagePromptOutputResult>> {
|
||||
return completeWithPromptOutputArtifactOperation(
|
||||
this.#authority,
|
||||
this.#client,
|
||||
commandValue,
|
||||
artifactValue,
|
||||
);
|
||||
}
|
||||
|
||||
resolve(
|
||||
commandValue: ModelInvocationResolutionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationResolutionRecord>>
|
||||
> {
|
||||
return resolveOperation(this.#authority, this.#client, commandValue);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
export { LocalModelInvocationRepository } from './local-model-invocation-repository/repository';
|
||||
export type { LocalModelInvocationOperationAuthority } from './local-model-invocation-repository/authority';
|
||||
@@ -0,0 +1,151 @@
|
||||
import type { StepRunStatus } from '@qinglong/runtime-core/step-run';
|
||||
|
||||
import type { ModelInvocationAuditRecord } from '../../model-gateway/model';
|
||||
import { normalizeModelUsage } from '../../model-gateway/validation';
|
||||
import type { ModelInvocationOutcome } from './contracts';
|
||||
import {
|
||||
ERROR_CODE_PATTERN,
|
||||
dataRecord,
|
||||
exactKeys,
|
||||
identifier,
|
||||
integer,
|
||||
invalid,
|
||||
requestDigest,
|
||||
} from './common';
|
||||
|
||||
export function normalizeAdmissionAudit(
|
||||
value: ModelInvocationAuditRecord,
|
||||
): Readonly<ModelInvocationAuditRecord> {
|
||||
const candidate = normalizeAuditCommon(value);
|
||||
if (
|
||||
candidate.phase !== 'admitted' ||
|
||||
candidate.outputBytes !== 0 ||
|
||||
candidate.usage !== null ||
|
||||
candidate.errorCode !== null
|
||||
) {
|
||||
invalid('admission audit facts are invalid');
|
||||
}
|
||||
return candidate;
|
||||
}
|
||||
|
||||
export function normalizeAuditCommon(
|
||||
value: ModelInvocationAuditRecord,
|
||||
): Readonly<ModelInvocationAuditRecord> {
|
||||
const candidate = dataRecord(value, 'audit record');
|
||||
exactKeys(
|
||||
candidate,
|
||||
[
|
||||
'deadlineAtMs',
|
||||
'errorCode',
|
||||
'inputBytes',
|
||||
'maxOutputTokens',
|
||||
'model',
|
||||
'occurredAtMs',
|
||||
'outputBytes',
|
||||
'phase',
|
||||
'policyRevision',
|
||||
'projectId',
|
||||
'provider',
|
||||
'requestDigest',
|
||||
'requestId',
|
||||
'runId',
|
||||
'stepRunId',
|
||||
'traceId',
|
||||
'usage',
|
||||
],
|
||||
'audit record',
|
||||
);
|
||||
if (!['admitted', 'completed', 'failed'].includes(value.phase)) {
|
||||
invalid('audit phase is invalid');
|
||||
}
|
||||
const usage = value.usage === null ? null : normalizeModelUsage(value.usage);
|
||||
const errorCode =
|
||||
value.errorCode === null
|
||||
? null
|
||||
: typeof value.errorCode === 'string' &&
|
||||
ERROR_CODE_PATTERN.test(value.errorCode)
|
||||
? value.errorCode
|
||||
: invalid('audit error code is invalid');
|
||||
return Object.freeze({
|
||||
phase: value.phase,
|
||||
projectId: identifier(value.projectId, 'audit project id'),
|
||||
runId: identifier(value.runId, 'audit Run id'),
|
||||
stepRunId: identifier(value.stepRunId, 'audit StepRun id'),
|
||||
traceId: identifier(value.traceId, 'audit trace id'),
|
||||
requestId: identifier(value.requestId, 'audit request id'),
|
||||
provider: identifier(value.provider, 'audit provider'),
|
||||
model: identifier(value.model, 'audit model'),
|
||||
policyRevision: identifier(value.policyRevision, 'audit policy revision'),
|
||||
requestDigest: requestDigest(value.requestDigest),
|
||||
deadlineAtMs: integer(
|
||||
value.deadlineAtMs,
|
||||
0,
|
||||
Number.MAX_SAFE_INTEGER,
|
||||
'audit deadline',
|
||||
),
|
||||
inputBytes: integer(value.inputBytes, 1, 256 * 1024, 'audit input bytes'),
|
||||
maxOutputTokens: integer(
|
||||
value.maxOutputTokens,
|
||||
1,
|
||||
32_768,
|
||||
'audit max output tokens',
|
||||
),
|
||||
outputBytes: integer(
|
||||
value.outputBytes,
|
||||
0,
|
||||
1024 * 1024,
|
||||
'audit output bytes',
|
||||
),
|
||||
usage,
|
||||
errorCode,
|
||||
occurredAtMs: integer(
|
||||
value.occurredAtMs,
|
||||
0,
|
||||
Number.MAX_SAFE_INTEGER,
|
||||
'audit time',
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
export function outcomeFor(audit: Readonly<ModelInvocationAuditRecord>): Readonly<{
|
||||
outcome: ModelInvocationOutcome;
|
||||
stepStatus: StepRunStatus;
|
||||
resultCode?: string;
|
||||
errorSummary?: string;
|
||||
}> {
|
||||
if (audit.phase === 'completed') {
|
||||
if (!audit.usage || audit.errorCode !== null) {
|
||||
invalid('successful completion facts are invalid');
|
||||
}
|
||||
return Object.freeze({ outcome: 'succeeded', stepStatus: 'succeeded' });
|
||||
}
|
||||
if (audit.phase !== 'failed' || audit.errorCode === null) {
|
||||
invalid('failed completion facts are invalid');
|
||||
}
|
||||
if (audit.errorCode === 'MODEL_INVOCATION_DEADLINE_EXCEEDED') {
|
||||
return Object.freeze({
|
||||
outcome: 'timed_out',
|
||||
stepStatus: 'timed_out',
|
||||
resultCode: 'model_deadline_exceeded',
|
||||
errorSummary: 'Model invocation deadline exceeded',
|
||||
});
|
||||
}
|
||||
if (
|
||||
audit.errorCode === 'MODEL_INVOCATION_ABORTED' ||
|
||||
audit.errorCode === 'MODEL_STREAM_CANCELLED' ||
|
||||
audit.errorCode === 'MODEL_INVOCATION_OUTCOME_UNKNOWN'
|
||||
) {
|
||||
return Object.freeze({
|
||||
outcome: 'outcome_unknown',
|
||||
stepStatus: 'lost',
|
||||
resultCode: 'model_outcome_unknown',
|
||||
errorSummary: 'Model invocation outcome is unknown',
|
||||
});
|
||||
}
|
||||
return Object.freeze({
|
||||
outcome: 'failed',
|
||||
stepStatus: 'failed',
|
||||
resultCode: 'model_provider_failed',
|
||||
errorSummary: 'Model invocation failed',
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
import { Buffer } from 'node:buffer';
|
||||
import { createHash } from 'node:crypto';
|
||||
|
||||
import {
|
||||
MAX_MODEL_INVOCATION_RECORD_JSON_BYTES,
|
||||
MODEL_INVOCATION_MUTATION_PHASES,
|
||||
InvalidModelInvocationError,
|
||||
type ModelInvocationMutationPhase,
|
||||
} from './contracts';
|
||||
|
||||
const IDENTITY_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}$/;
|
||||
export const ERROR_CODE_PATTERN = /^[A-Z][A-Z0-9_]{0,63}$/;
|
||||
const SHA256_PATTERN = /^sha256:[0-9a-f]{64}$/;
|
||||
const DIGEST_PATTERN = /^[0-9a-f]{64}$/;
|
||||
export const START_DIGEST_DOMAIN = Buffer.from(
|
||||
'qinglong/model-invocation-start-digest@v1\0',
|
||||
'utf8',
|
||||
);
|
||||
export const START_COMMAND_DIGEST_DOMAIN = Buffer.from(
|
||||
'qinglong/model-invocation-start-command-digest@v1\0',
|
||||
'utf8',
|
||||
);
|
||||
export const COMPLETION_DIGEST_DOMAIN = Buffer.from(
|
||||
'qinglong/model-invocation-completion-digest@v1\0',
|
||||
'utf8',
|
||||
);
|
||||
export const COMPLETION_COMMAND_DIGEST_DOMAIN = Buffer.from(
|
||||
'qinglong/model-invocation-completion-command-digest@v1\0',
|
||||
'utf8',
|
||||
);
|
||||
const MUTATION_IDENTITY_DOMAIN = Buffer.from(
|
||||
'qinglong/model-invocation-mutation-identity@v1\0',
|
||||
'utf8',
|
||||
);
|
||||
|
||||
export function invalid(message: string): never {
|
||||
throw new InvalidModelInvocationError(message);
|
||||
}
|
||||
|
||||
export function dataRecord(value: unknown, label: string): Record<string, unknown> {
|
||||
if (
|
||||
!value ||
|
||||
typeof value !== 'object' ||
|
||||
Array.isArray(value) ||
|
||||
Object.getPrototypeOf(value) !== Object.prototype
|
||||
) {
|
||||
return invalid(`${label} must be a plain object`);
|
||||
}
|
||||
return value as Record<string, unknown>;
|
||||
}
|
||||
|
||||
export function exactKeys(
|
||||
value: object,
|
||||
expected: readonly string[],
|
||||
label: string,
|
||||
): void {
|
||||
const actual = Object.keys(value).sort();
|
||||
const canonical = [...expected].sort();
|
||||
if (
|
||||
actual.length !== canonical.length ||
|
||||
actual.some((key, index) => key !== canonical[index])
|
||||
) {
|
||||
invalid(`${label} shape is invalid`);
|
||||
}
|
||||
}
|
||||
|
||||
export function identifier(value: unknown, label: string): string {
|
||||
if (typeof value !== 'string' || !IDENTITY_PATTERN.test(value)) {
|
||||
return invalid(`${label} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
export function digest(value: unknown, label: string): string {
|
||||
if (typeof value !== 'string' || !DIGEST_PATTERN.test(value)) {
|
||||
return invalid(`${label} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
export function requestDigest(value: unknown): string {
|
||||
if (typeof value !== 'string' || !SHA256_PATTERN.test(value)) {
|
||||
return invalid('request digest is invalid');
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
export function createModelInvocationMutationIdentity(
|
||||
invocationIdValue: string,
|
||||
phaseValue: ModelInvocationMutationPhase,
|
||||
): Readonly<{
|
||||
mutationId: string;
|
||||
eventId: string;
|
||||
dedupeKey: string;
|
||||
}> {
|
||||
const invocationId = identifier(invocationIdValue, 'invocation id');
|
||||
if (!MODEL_INVOCATION_MUTATION_PHASES.includes(phaseValue)) {
|
||||
invalid('mutation phase is invalid');
|
||||
}
|
||||
const identityDigest = createHash('sha256')
|
||||
.update(MUTATION_IDENTITY_DOMAIN)
|
||||
.update(phaseValue, 'utf8')
|
||||
.update('\0', 'utf8')
|
||||
.update(invocationId, 'utf8')
|
||||
.digest('hex');
|
||||
const eventHex =
|
||||
identityDigest.slice(0, 12) +
|
||||
'8' +
|
||||
identityDigest.slice(13, 16) +
|
||||
'8' +
|
||||
identityDigest.slice(17, 32);
|
||||
return Object.freeze({
|
||||
mutationId: `ql3mi.${phaseValue}.mutation.${identityDigest}`,
|
||||
eventId: `${eventHex.slice(0, 8)}-${eventHex.slice(8, 12)}-${eventHex.slice(
|
||||
12,
|
||||
16,
|
||||
)}-${eventHex.slice(16, 20)}-${eventHex.slice(20, 32)}`,
|
||||
dedupeKey: `ql3mi.${phaseValue}.dedupe.${identityDigest}`,
|
||||
});
|
||||
}
|
||||
|
||||
export function integer(
|
||||
value: unknown,
|
||||
minimum: number,
|
||||
maximum: number,
|
||||
label: string,
|
||||
): number {
|
||||
if (
|
||||
!Number.isSafeInteger(value) ||
|
||||
(value as number) < minimum ||
|
||||
(value as number) > maximum
|
||||
) {
|
||||
return invalid(`${label} is invalid`);
|
||||
}
|
||||
return value as number;
|
||||
}
|
||||
|
||||
export function hash(domain: Uint8Array, value: unknown): string {
|
||||
return createHash('sha256')
|
||||
.update(domain)
|
||||
.update(JSON.stringify(value))
|
||||
.digest('hex');
|
||||
}
|
||||
|
||||
export function assertJsonBudget(value: unknown, label: string): void {
|
||||
if (
|
||||
Buffer.byteLength(JSON.stringify(value), 'utf8') >
|
||||
MAX_MODEL_INVOCATION_RECORD_JSON_BYTES
|
||||
) {
|
||||
invalid(`${label} exceeds its JSON budget`);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,269 @@
|
||||
import { normalizeStepRunMutation, type StepRunMutation } from '@qinglong/runtime-core/step-run';
|
||||
|
||||
import type { ModelInvocationAuditRecord } from '../../model-gateway/model';
|
||||
import { normalizeModelUsage } from '../../model-gateway/validation';
|
||||
import {
|
||||
MODEL_INVOCATION_COMPLETION_COMMAND_SCHEMA,
|
||||
MODEL_INVOCATION_COMPLETION_SCHEMA,
|
||||
MODEL_INVOCATION_OUTCOMES,
|
||||
type ModelInvocationCompletionCommand,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationStartRecord,
|
||||
} from './contracts';
|
||||
import {
|
||||
COMPLETION_COMMAND_DIGEST_DOMAIN,
|
||||
COMPLETION_DIGEST_DOMAIN,
|
||||
ERROR_CODE_PATTERN,
|
||||
assertJsonBudget,
|
||||
createModelInvocationMutationIdentity,
|
||||
dataRecord,
|
||||
digest,
|
||||
exactKeys,
|
||||
hash,
|
||||
identifier,
|
||||
integer,
|
||||
invalid,
|
||||
} from './common';
|
||||
import { normalizeAuditCommon, outcomeFor } from './audit';
|
||||
import { normalizeModelInvocationStartRecord } from './startProtocol';
|
||||
|
||||
function completionWithoutDigest(
|
||||
value: Readonly<ModelInvocationCompletionRecord>,
|
||||
): Omit<ModelInvocationCompletionRecord, 'completionDigest'> {
|
||||
const { completionDigest: _completionDigest, ...unsigned } = value;
|
||||
return unsigned;
|
||||
}
|
||||
|
||||
export function normalizeModelInvocationCompletionRecord(
|
||||
value: ModelInvocationCompletionRecord,
|
||||
): Readonly<ModelInvocationCompletionRecord> {
|
||||
const candidate = dataRecord(value, 'completion record');
|
||||
exactKeys(
|
||||
candidate,
|
||||
[
|
||||
'completedAtMs',
|
||||
'completedStepRunDigest',
|
||||
'completedStepRunVersion',
|
||||
'completionDigest',
|
||||
'errorCode',
|
||||
'invocationId',
|
||||
'outcome',
|
||||
'outputBytes',
|
||||
'projectId',
|
||||
'runEventId',
|
||||
'runId',
|
||||
'schema',
|
||||
'startDigest',
|
||||
'stepRunId',
|
||||
'stepRunMutationDigest',
|
||||
'stepRunMutationId',
|
||||
'traceId',
|
||||
'usage',
|
||||
],
|
||||
'completion record',
|
||||
);
|
||||
if (
|
||||
value.schema !== MODEL_INVOCATION_COMPLETION_SCHEMA ||
|
||||
!MODEL_INVOCATION_OUTCOMES.includes(value.outcome)
|
||||
) {
|
||||
invalid('completion schema or outcome is invalid');
|
||||
}
|
||||
const usage = value.usage === null ? null : normalizeModelUsage(value.usage);
|
||||
const errorCode =
|
||||
value.errorCode === null
|
||||
? null
|
||||
: typeof value.errorCode === 'string' &&
|
||||
ERROR_CODE_PATTERN.test(value.errorCode)
|
||||
? value.errorCode
|
||||
: invalid('completion error code is invalid');
|
||||
const normalized = Object.freeze({
|
||||
schema: MODEL_INVOCATION_COMPLETION_SCHEMA,
|
||||
invocationId: identifier(value.invocationId, 'invocation id'),
|
||||
projectId: identifier(value.projectId, 'project id'),
|
||||
runId: identifier(value.runId, 'Run id'),
|
||||
stepRunId: identifier(value.stepRunId, 'StepRun id'),
|
||||
traceId: identifier(value.traceId, 'trace id'),
|
||||
startDigest: digest(value.startDigest, 'start digest'),
|
||||
outcome: value.outcome,
|
||||
outputBytes: integer(value.outputBytes, 0, 1024 * 1024, 'output bytes'),
|
||||
usage,
|
||||
errorCode,
|
||||
completedStepRunVersion: integer(
|
||||
value.completedStepRunVersion,
|
||||
3,
|
||||
2_147_483_647,
|
||||
'completed StepRun version',
|
||||
),
|
||||
stepRunMutationId: identifier(
|
||||
value.stepRunMutationId,
|
||||
'StepRun mutation id',
|
||||
),
|
||||
stepRunMutationDigest: digest(
|
||||
value.stepRunMutationDigest,
|
||||
'StepRun mutation digest',
|
||||
),
|
||||
completedStepRunDigest: digest(
|
||||
value.completedStepRunDigest,
|
||||
'completed StepRun digest',
|
||||
),
|
||||
runEventId: identifier(value.runEventId, 'RunEvent id'),
|
||||
completedAtMs: integer(
|
||||
value.completedAtMs,
|
||||
0,
|
||||
Number.MAX_SAFE_INTEGER,
|
||||
'completed time',
|
||||
),
|
||||
completionDigest: digest(value.completionDigest, 'completion digest'),
|
||||
});
|
||||
const identity = createModelInvocationMutationIdentity(
|
||||
normalized.invocationId,
|
||||
'completion',
|
||||
);
|
||||
if (
|
||||
(normalized.outcome === 'succeeded') !==
|
||||
(normalized.usage !== null && normalized.errorCode === null) ||
|
||||
(normalized.outcome !== 'succeeded' && normalized.errorCode === null) ||
|
||||
normalized.stepRunMutationId !== identity.mutationId ||
|
||||
normalized.runEventId !== identity.eventId ||
|
||||
hash(COMPLETION_DIGEST_DOMAIN, completionWithoutDigest(normalized)) !==
|
||||
normalized.completionDigest
|
||||
) {
|
||||
invalid('completion facts or digest are invalid');
|
||||
}
|
||||
assertJsonBudget(normalized, 'completion record');
|
||||
return normalized;
|
||||
}
|
||||
|
||||
export function createModelInvocationCompletionCommand(
|
||||
startValue: ModelInvocationStartRecord,
|
||||
auditValue: ModelInvocationAuditRecord,
|
||||
mutationValue: StepRunMutation,
|
||||
successOutputRefValue?: string,
|
||||
): Readonly<ModelInvocationCompletionCommand> {
|
||||
const start = normalizeModelInvocationStartRecord(startValue);
|
||||
const audit = normalizeAuditCommon(auditValue);
|
||||
const mutation = normalizeStepRunMutation(mutationValue);
|
||||
const outcome = outcomeFor(audit);
|
||||
const identity = createModelInvocationMutationIdentity(
|
||||
audit.requestId,
|
||||
'completion',
|
||||
);
|
||||
const successOutputRef =
|
||||
successOutputRefValue ?? `model-invocation:${start.invocationId}`;
|
||||
if (
|
||||
audit.requestId !== start.invocationId ||
|
||||
audit.projectId !== start.projectId ||
|
||||
audit.runId !== start.runId ||
|
||||
audit.stepRunId !== start.stepRunId ||
|
||||
audit.traceId !== start.traceId ||
|
||||
audit.provider !== start.provider ||
|
||||
audit.model !== start.model ||
|
||||
audit.policyRevision !== start.policyRevision ||
|
||||
audit.requestDigest !== start.requestDigest ||
|
||||
audit.deadlineAtMs !== start.deadlineAtMs ||
|
||||
audit.inputBytes !== start.inputBytes ||
|
||||
audit.maxOutputTokens !== start.maxOutputTokens ||
|
||||
audit.occurredAtMs < start.admittedAtMs ||
|
||||
mutation.previousStatus !== 'running' ||
|
||||
mutation.mutationId !== identity.mutationId ||
|
||||
mutation.expectedStepRunVersion !== start.startedStepRunVersion ||
|
||||
mutation.expectedStepRunDigest !== start.startedStepRunDigest ||
|
||||
mutation.stepRun.runId !== start.runId ||
|
||||
mutation.stepRun.id !== start.stepRunId ||
|
||||
mutation.stepRun.kind !== 'model' ||
|
||||
mutation.stepRun.status !== outcome.stepStatus ||
|
||||
mutation.stepRun.updatedAtMs !== audit.occurredAtMs ||
|
||||
mutation.event.id !== identity.eventId ||
|
||||
mutation.event.dedupeKey !== identity.dedupeKey ||
|
||||
mutation.event.type !== `step.${outcome.stepStatus}` ||
|
||||
(outcome.stepStatus === 'succeeded'
|
||||
? mutation.stepRun.outputRef !== successOutputRef
|
||||
: mutation.stepRun.resultCode !== outcome.resultCode ||
|
||||
mutation.stepRun.errorSummary !== outcome.errorSummary)
|
||||
) {
|
||||
invalid('completion identity or StepRun mutation is not exact');
|
||||
}
|
||||
const unsigned = Object.freeze({
|
||||
schema: MODEL_INVOCATION_COMPLETION_SCHEMA,
|
||||
invocationId: start.invocationId,
|
||||
projectId: start.projectId,
|
||||
runId: start.runId,
|
||||
stepRunId: start.stepRunId,
|
||||
traceId: start.traceId,
|
||||
startDigest: start.startDigest,
|
||||
outcome: outcome.outcome,
|
||||
outputBytes: audit.outputBytes,
|
||||
usage: audit.usage,
|
||||
errorCode: audit.errorCode,
|
||||
completedStepRunVersion: mutation.stepRun.version,
|
||||
stepRunMutationId: mutation.mutationId,
|
||||
stepRunMutationDigest: mutation.mutationDigest,
|
||||
completedStepRunDigest: mutation.stepRun.stepRunDigest,
|
||||
runEventId: mutation.event.id,
|
||||
completedAtMs: audit.occurredAtMs,
|
||||
});
|
||||
const completion = normalizeModelInvocationCompletionRecord({
|
||||
...unsigned,
|
||||
completionDigest: hash(COMPLETION_DIGEST_DOMAIN, unsigned),
|
||||
});
|
||||
const commandUnsigned = Object.freeze({
|
||||
schema: MODEL_INVOCATION_COMPLETION_COMMAND_SCHEMA,
|
||||
start,
|
||||
completion,
|
||||
stepRunMutation: mutation,
|
||||
});
|
||||
return Object.freeze({
|
||||
...commandUnsigned,
|
||||
commandDigest: hash(COMPLETION_COMMAND_DIGEST_DOMAIN, commandUnsigned),
|
||||
});
|
||||
}
|
||||
|
||||
export function normalizeModelInvocationCompletionCommand(
|
||||
value: ModelInvocationCompletionCommand,
|
||||
): Readonly<ModelInvocationCompletionCommand> {
|
||||
const candidate = dataRecord(value, 'completion command');
|
||||
exactKeys(
|
||||
candidate,
|
||||
['commandDigest', 'completion', 'schema', 'start', 'stepRunMutation'],
|
||||
'completion command',
|
||||
);
|
||||
if (value.schema !== MODEL_INVOCATION_COMPLETION_COMMAND_SCHEMA) {
|
||||
invalid('completion command schema is invalid');
|
||||
}
|
||||
const start = normalizeModelInvocationStartRecord(value.start);
|
||||
const completion = normalizeModelInvocationCompletionRecord(value.completion);
|
||||
const mutation = normalizeStepRunMutation(value.stepRunMutation);
|
||||
const canonical = createModelInvocationCompletionCommand(
|
||||
start,
|
||||
{
|
||||
phase: completion.outcome === 'succeeded' ? 'completed' : 'failed',
|
||||
projectId: completion.projectId,
|
||||
runId: completion.runId,
|
||||
stepRunId: completion.stepRunId,
|
||||
traceId: completion.traceId,
|
||||
requestId: completion.invocationId,
|
||||
provider: start.provider,
|
||||
model: start.model,
|
||||
policyRevision: start.policyRevision,
|
||||
requestDigest: start.requestDigest,
|
||||
deadlineAtMs: start.deadlineAtMs,
|
||||
inputBytes: start.inputBytes,
|
||||
maxOutputTokens: start.maxOutputTokens,
|
||||
outputBytes: completion.outputBytes,
|
||||
usage: completion.usage,
|
||||
errorCode: completion.errorCode,
|
||||
occurredAtMs: completion.completedAtMs,
|
||||
},
|
||||
mutation,
|
||||
completion.outcome === 'succeeded'
|
||||
? mutation.stepRun.outputRef ?? undefined
|
||||
: undefined,
|
||||
);
|
||||
if (
|
||||
canonical.completion.completionDigest !== completion.completionDigest ||
|
||||
digest(value.commandDigest, 'command digest') !== canonical.commandDigest
|
||||
) {
|
||||
invalid('completion command is not canonical');
|
||||
}
|
||||
return canonical;
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
import type { StepRunRecord, StepRunMutation } from '@qinglong/runtime-core/step-run';
|
||||
import type { ModelUsage } from '../../model-gateway/model';
|
||||
|
||||
export const MODEL_INVOCATION_START_SCHEMA =
|
||||
'qinglong/model-invocation-start@v1' as const;
|
||||
export const MODEL_INVOCATION_START_COMMAND_SCHEMA =
|
||||
'qinglong/model-invocation-start-command@v1' as const;
|
||||
export const MODEL_INVOCATION_COMPLETION_SCHEMA =
|
||||
'qinglong/model-invocation-completion@v1' as const;
|
||||
export const MODEL_INVOCATION_COMPLETION_COMMAND_SCHEMA =
|
||||
'qinglong/model-invocation-completion-command@v1' as const;
|
||||
|
||||
export const MAX_MODEL_INVOCATION_RECORD_JSON_BYTES = 24 * 1024;
|
||||
export const MAX_MODEL_INVOCATION_RECOVERY_PAGE_SIZE = 128;
|
||||
export const MODEL_INVOCATION_OUTCOMES = [
|
||||
'succeeded',
|
||||
'failed',
|
||||
'timed_out',
|
||||
'outcome_unknown',
|
||||
] as const;
|
||||
export const MODEL_INVOCATION_MUTATION_PHASES = [
|
||||
'start',
|
||||
'completion',
|
||||
'resolution',
|
||||
] as const;
|
||||
|
||||
export type ModelInvocationOutcome = (typeof MODEL_INVOCATION_OUTCOMES)[number];
|
||||
export type ModelInvocationMutationPhase =
|
||||
(typeof MODEL_INVOCATION_MUTATION_PHASES)[number];
|
||||
|
||||
export interface ModelInvocationStartRecord {
|
||||
readonly schema: typeof MODEL_INVOCATION_START_SCHEMA;
|
||||
readonly invocationId: string;
|
||||
readonly projectId: string;
|
||||
readonly runId: string;
|
||||
readonly stepRunId: string;
|
||||
readonly traceId: string;
|
||||
readonly provider: string;
|
||||
readonly model: string;
|
||||
readonly policyRevision: string;
|
||||
readonly requestDigest: string;
|
||||
readonly inputBytes: number;
|
||||
readonly maxOutputTokens: number;
|
||||
readonly deadlineAtMs: number;
|
||||
readonly startedStepRunVersion: number;
|
||||
readonly stepRunMutationId: string;
|
||||
readonly stepRunMutationDigest: string;
|
||||
readonly startedStepRunDigest: string;
|
||||
readonly runEventId: string;
|
||||
readonly admittedAtMs: number;
|
||||
readonly startDigest: string;
|
||||
}
|
||||
|
||||
export interface ModelInvocationStartCommand {
|
||||
readonly schema: typeof MODEL_INVOCATION_START_COMMAND_SCHEMA;
|
||||
readonly start: Readonly<ModelInvocationStartRecord>;
|
||||
readonly stepRunMutation: Readonly<StepRunMutation>;
|
||||
readonly commandDigest: string;
|
||||
}
|
||||
|
||||
export interface ModelInvocationCompletionRecord {
|
||||
readonly schema: typeof MODEL_INVOCATION_COMPLETION_SCHEMA;
|
||||
readonly invocationId: string;
|
||||
readonly projectId: string;
|
||||
readonly runId: string;
|
||||
readonly stepRunId: string;
|
||||
readonly traceId: string;
|
||||
readonly startDigest: string;
|
||||
readonly outcome: ModelInvocationOutcome;
|
||||
readonly outputBytes: number;
|
||||
readonly usage: Readonly<ModelUsage> | null;
|
||||
readonly errorCode: string | null;
|
||||
readonly completedStepRunVersion: number;
|
||||
readonly stepRunMutationId: string;
|
||||
readonly stepRunMutationDigest: string;
|
||||
readonly completedStepRunDigest: string;
|
||||
readonly runEventId: string;
|
||||
readonly completedAtMs: number;
|
||||
readonly completionDigest: string;
|
||||
}
|
||||
|
||||
export interface ModelInvocationCompletionCommand {
|
||||
readonly schema: typeof MODEL_INVOCATION_COMPLETION_COMMAND_SCHEMA;
|
||||
readonly start: Readonly<ModelInvocationStartRecord>;
|
||||
readonly completion: Readonly<ModelInvocationCompletionRecord>;
|
||||
readonly stepRunMutation: Readonly<StepRunMutation>;
|
||||
readonly commandDigest: string;
|
||||
}
|
||||
|
||||
export interface CommitModelInvocationResult<T> {
|
||||
readonly status: 'created' | 'existing';
|
||||
readonly record: Readonly<T>;
|
||||
}
|
||||
|
||||
export interface ModelInvocationAuthoritySnapshot {
|
||||
readonly projectId: string;
|
||||
readonly runId: string;
|
||||
readonly runVersion: number;
|
||||
readonly runEventSequence: number;
|
||||
readonly stepRun: Readonly<StepRunRecord>;
|
||||
}
|
||||
|
||||
export interface ModelInvocationRecoveryPage {
|
||||
readonly observedAtMs: number;
|
||||
readonly candidates: readonly Readonly<ModelInvocationStartRecord>[];
|
||||
readonly hasMore: boolean;
|
||||
}
|
||||
|
||||
export interface ModelInvocationRepository {
|
||||
findStart(
|
||||
invocationId: string,
|
||||
): Promise<Readonly<ModelInvocationStartRecord> | null>;
|
||||
findCompletion(
|
||||
invocationId: string,
|
||||
): Promise<Readonly<ModelInvocationCompletionRecord> | null>;
|
||||
readAuthority(
|
||||
identity: Readonly<{
|
||||
projectId: string;
|
||||
runId: string;
|
||||
stepRunId: string;
|
||||
}>,
|
||||
): Promise<Readonly<ModelInvocationAuthoritySnapshot> | null>;
|
||||
listIncomplete(limit: number): Promise<Readonly<ModelInvocationRecoveryPage>>;
|
||||
admit(
|
||||
command: ModelInvocationStartCommand,
|
||||
): Promise<Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>>;
|
||||
complete(
|
||||
command: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
>;
|
||||
}
|
||||
|
||||
export class InvalidModelInvocationError extends TypeError {
|
||||
readonly code = 'MODEL_INVOCATION_INVALID';
|
||||
|
||||
constructor(message: string) {
|
||||
super(`Model invocation is invalid: ${message}`);
|
||||
this.name = 'InvalidModelInvocationError';
|
||||
}
|
||||
}
|
||||
|
||||
export class ModelInvocationConflictError extends Error {
|
||||
readonly code = 'MODEL_INVOCATION_CONFLICT';
|
||||
|
||||
constructor() {
|
||||
super('Model invocation conflicts with durable state');
|
||||
this.name = 'ModelInvocationConflictError';
|
||||
}
|
||||
}
|
||||
|
||||
export class ModelInvocationRepositoryUnavailableError extends Error {
|
||||
readonly code = 'MODEL_INVOCATION_REPOSITORY_UNAVAILABLE';
|
||||
|
||||
constructor(options?: ErrorOptions) {
|
||||
super('Model invocation repository is unavailable', options);
|
||||
this.name = 'ModelInvocationRepositoryUnavailableError';
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
import { normalizeStepRunMutation, type StepRunMutation } from '@qinglong/runtime-core/step-run';
|
||||
|
||||
import type { ModelInvocationAuditRecord } from '../../model-gateway/model';
|
||||
import {
|
||||
MODEL_INVOCATION_START_COMMAND_SCHEMA,
|
||||
MODEL_INVOCATION_START_SCHEMA,
|
||||
type ModelInvocationStartCommand,
|
||||
type ModelInvocationStartRecord,
|
||||
} from './contracts';
|
||||
import {
|
||||
START_COMMAND_DIGEST_DOMAIN,
|
||||
START_DIGEST_DOMAIN,
|
||||
assertJsonBudget,
|
||||
createModelInvocationMutationIdentity,
|
||||
dataRecord,
|
||||
digest,
|
||||
exactKeys,
|
||||
hash,
|
||||
identifier,
|
||||
integer,
|
||||
invalid,
|
||||
requestDigest,
|
||||
} from './common';
|
||||
import { normalizeAdmissionAudit } from './audit';
|
||||
|
||||
function startWithoutDigest(
|
||||
value: Readonly<ModelInvocationStartRecord>,
|
||||
): Omit<ModelInvocationStartRecord, 'startDigest'> {
|
||||
const { startDigest: _startDigest, ...unsigned } = value;
|
||||
return unsigned;
|
||||
}
|
||||
|
||||
export function normalizeModelInvocationStartRecord(
|
||||
value: ModelInvocationStartRecord,
|
||||
): Readonly<ModelInvocationStartRecord> {
|
||||
const candidate = dataRecord(value, 'start record');
|
||||
exactKeys(
|
||||
candidate,
|
||||
[
|
||||
'admittedAtMs',
|
||||
'deadlineAtMs',
|
||||
'inputBytes',
|
||||
'invocationId',
|
||||
'maxOutputTokens',
|
||||
'model',
|
||||
'policyRevision',
|
||||
'projectId',
|
||||
'provider',
|
||||
'requestDigest',
|
||||
'runEventId',
|
||||
'runId',
|
||||
'schema',
|
||||
'startDigest',
|
||||
'startedStepRunDigest',
|
||||
'startedStepRunVersion',
|
||||
'stepRunId',
|
||||
'stepRunMutationDigest',
|
||||
'stepRunMutationId',
|
||||
'traceId',
|
||||
],
|
||||
'start record',
|
||||
);
|
||||
if (value.schema !== MODEL_INVOCATION_START_SCHEMA) {
|
||||
invalid('start schema is invalid');
|
||||
}
|
||||
const normalized = Object.freeze({
|
||||
schema: MODEL_INVOCATION_START_SCHEMA,
|
||||
invocationId: identifier(value.invocationId, 'invocation id'),
|
||||
projectId: identifier(value.projectId, 'project id'),
|
||||
runId: identifier(value.runId, 'Run id'),
|
||||
stepRunId: identifier(value.stepRunId, 'StepRun id'),
|
||||
traceId: identifier(value.traceId, 'trace id'),
|
||||
provider: identifier(value.provider, 'provider'),
|
||||
model: identifier(value.model, 'model'),
|
||||
policyRevision: identifier(value.policyRevision, 'policy revision'),
|
||||
requestDigest: requestDigest(value.requestDigest),
|
||||
inputBytes: integer(value.inputBytes, 1, 256 * 1024, 'input bytes'),
|
||||
maxOutputTokens: integer(
|
||||
value.maxOutputTokens,
|
||||
1,
|
||||
32_768,
|
||||
'max output tokens',
|
||||
),
|
||||
deadlineAtMs: integer(
|
||||
value.deadlineAtMs,
|
||||
0,
|
||||
Number.MAX_SAFE_INTEGER,
|
||||
'deadline',
|
||||
),
|
||||
startedStepRunVersion: integer(
|
||||
value.startedStepRunVersion,
|
||||
2,
|
||||
2_147_483_647,
|
||||
'started StepRun version',
|
||||
),
|
||||
stepRunMutationId: identifier(
|
||||
value.stepRunMutationId,
|
||||
'StepRun mutation id',
|
||||
),
|
||||
stepRunMutationDigest: digest(
|
||||
value.stepRunMutationDigest,
|
||||
'StepRun mutation digest',
|
||||
),
|
||||
startedStepRunDigest: digest(
|
||||
value.startedStepRunDigest,
|
||||
'started StepRun digest',
|
||||
),
|
||||
runEventId: identifier(value.runEventId, 'RunEvent id'),
|
||||
admittedAtMs: integer(
|
||||
value.admittedAtMs,
|
||||
0,
|
||||
Number.MAX_SAFE_INTEGER,
|
||||
'admitted time',
|
||||
),
|
||||
startDigest: digest(value.startDigest, 'start digest'),
|
||||
});
|
||||
const identity = createModelInvocationMutationIdentity(
|
||||
normalized.invocationId,
|
||||
'start',
|
||||
);
|
||||
if (
|
||||
normalized.deadlineAtMs <= normalized.admittedAtMs ||
|
||||
normalized.stepRunMutationId !== identity.mutationId ||
|
||||
normalized.runEventId !== identity.eventId ||
|
||||
hash(START_DIGEST_DOMAIN, startWithoutDigest(normalized)) !==
|
||||
normalized.startDigest
|
||||
) {
|
||||
invalid('start time or digest is invalid');
|
||||
}
|
||||
assertJsonBudget(normalized, 'start record');
|
||||
return normalized;
|
||||
}
|
||||
|
||||
function assertStartMutation(
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
audit: Readonly<ModelInvocationAuditRecord>,
|
||||
): void {
|
||||
const identity = createModelInvocationMutationIdentity(
|
||||
audit.requestId,
|
||||
'start',
|
||||
);
|
||||
if (
|
||||
mutation.previousStatus !== 'ready' ||
|
||||
mutation.mutationId !== identity.mutationId ||
|
||||
mutation.stepRun.kind !== 'model' ||
|
||||
mutation.stepRun.status !== 'running' ||
|
||||
mutation.stepRun.runId !== audit.runId ||
|
||||
mutation.stepRun.id !== audit.stepRunId ||
|
||||
mutation.stepRun.updatedAtMs !== audit.occurredAtMs ||
|
||||
mutation.event.id !== identity.eventId ||
|
||||
mutation.event.dedupeKey !== identity.dedupeKey ||
|
||||
mutation.event.type !== 'step.running'
|
||||
) {
|
||||
invalid('start StepRun mutation is not exact');
|
||||
}
|
||||
}
|
||||
|
||||
export function createModelInvocationStartCommand(
|
||||
auditValue: ModelInvocationAuditRecord,
|
||||
mutationValue: StepRunMutation,
|
||||
): Readonly<ModelInvocationStartCommand> {
|
||||
const audit = normalizeAdmissionAudit(auditValue);
|
||||
const mutation = normalizeStepRunMutation(mutationValue);
|
||||
assertStartMutation(mutation, audit);
|
||||
const unsigned = Object.freeze({
|
||||
schema: MODEL_INVOCATION_START_SCHEMA,
|
||||
invocationId: audit.requestId,
|
||||
projectId: audit.projectId,
|
||||
runId: audit.runId,
|
||||
stepRunId: audit.stepRunId,
|
||||
traceId: audit.traceId,
|
||||
provider: audit.provider,
|
||||
model: audit.model,
|
||||
policyRevision: audit.policyRevision,
|
||||
requestDigest: audit.requestDigest,
|
||||
inputBytes: audit.inputBytes,
|
||||
maxOutputTokens: audit.maxOutputTokens,
|
||||
deadlineAtMs: audit.deadlineAtMs,
|
||||
startedStepRunVersion: mutation.stepRun.version,
|
||||
stepRunMutationId: mutation.mutationId,
|
||||
stepRunMutationDigest: mutation.mutationDigest,
|
||||
startedStepRunDigest: mutation.stepRun.stepRunDigest,
|
||||
runEventId: mutation.event.id,
|
||||
admittedAtMs: audit.occurredAtMs,
|
||||
});
|
||||
const start = normalizeModelInvocationStartRecord({
|
||||
...unsigned,
|
||||
startDigest: hash(START_DIGEST_DOMAIN, unsigned),
|
||||
});
|
||||
const commandUnsigned = Object.freeze({
|
||||
schema: MODEL_INVOCATION_START_COMMAND_SCHEMA,
|
||||
start,
|
||||
stepRunMutation: mutation,
|
||||
});
|
||||
return Object.freeze({
|
||||
...commandUnsigned,
|
||||
commandDigest: hash(START_COMMAND_DIGEST_DOMAIN, commandUnsigned),
|
||||
});
|
||||
}
|
||||
|
||||
export function normalizeModelInvocationStartCommand(
|
||||
value: ModelInvocationStartCommand,
|
||||
): Readonly<ModelInvocationStartCommand> {
|
||||
const candidate = dataRecord(value, 'start command');
|
||||
exactKeys(
|
||||
candidate,
|
||||
['commandDigest', 'schema', 'start', 'stepRunMutation'],
|
||||
'start command',
|
||||
);
|
||||
if (value.schema !== MODEL_INVOCATION_START_COMMAND_SCHEMA) {
|
||||
invalid('start command schema is invalid');
|
||||
}
|
||||
const start = normalizeModelInvocationStartRecord(value.start);
|
||||
const mutation = normalizeStepRunMutation(value.stepRunMutation);
|
||||
const canonical = createModelInvocationStartCommand(
|
||||
{
|
||||
phase: 'admitted',
|
||||
projectId: start.projectId,
|
||||
runId: start.runId,
|
||||
stepRunId: start.stepRunId,
|
||||
traceId: start.traceId,
|
||||
requestId: start.invocationId,
|
||||
provider: start.provider,
|
||||
model: start.model,
|
||||
policyRevision: start.policyRevision,
|
||||
requestDigest: start.requestDigest,
|
||||
deadlineAtMs: start.deadlineAtMs,
|
||||
inputBytes: start.inputBytes,
|
||||
maxOutputTokens: start.maxOutputTokens,
|
||||
outputBytes: 0,
|
||||
usage: null,
|
||||
errorCode: null,
|
||||
occurredAtMs: start.admittedAtMs,
|
||||
},
|
||||
mutation,
|
||||
);
|
||||
if (
|
||||
canonical.start.startDigest !== start.startDigest ||
|
||||
digest(value.commandDigest, 'command digest') !== canonical.commandDigest
|
||||
) {
|
||||
invalid('start command is not canonical');
|
||||
}
|
||||
return canonical;
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
export {
|
||||
MAX_MODEL_INVOCATION_RECORD_JSON_BYTES,
|
||||
MAX_MODEL_INVOCATION_RECOVERY_PAGE_SIZE,
|
||||
MODEL_INVOCATION_COMPLETION_COMMAND_SCHEMA,
|
||||
MODEL_INVOCATION_COMPLETION_SCHEMA,
|
||||
MODEL_INVOCATION_MUTATION_PHASES,
|
||||
MODEL_INVOCATION_OUTCOMES,
|
||||
MODEL_INVOCATION_START_COMMAND_SCHEMA,
|
||||
MODEL_INVOCATION_START_SCHEMA,
|
||||
InvalidModelInvocationError,
|
||||
ModelInvocationConflictError,
|
||||
ModelInvocationRepositoryUnavailableError,
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationAuthoritySnapshot,
|
||||
type ModelInvocationCompletionCommand,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationMutationPhase,
|
||||
type ModelInvocationOutcome,
|
||||
type ModelInvocationRecoveryPage,
|
||||
type ModelInvocationRepository,
|
||||
type ModelInvocationStartCommand,
|
||||
type ModelInvocationStartRecord,
|
||||
} from './model-invocation/contracts';
|
||||
export { createModelInvocationMutationIdentity } from './model-invocation/common';
|
||||
export {
|
||||
createModelInvocationStartCommand,
|
||||
normalizeModelInvocationStartCommand,
|
||||
normalizeModelInvocationStartRecord,
|
||||
} from './model-invocation/startProtocol';
|
||||
export {
|
||||
createModelInvocationCompletionCommand,
|
||||
normalizeModelInvocationCompletionCommand,
|
||||
normalizeModelInvocationCompletionRecord,
|
||||
} from './model-invocation/completionProtocol';
|
||||
@@ -0,0 +1,584 @@
|
||||
import { Buffer } from 'node:buffer';
|
||||
import { createHash } from 'node:crypto';
|
||||
|
||||
// Model Invocation owns durable ambiguity resolution alongside its transaction contract.
|
||||
|
||||
import {
|
||||
normalizeStepRunMutation,
|
||||
transitionStepRunMutation,
|
||||
type StepRunMutation,
|
||||
type StepRunStatus,
|
||||
} from '@qinglong/runtime-core/step-run';
|
||||
|
||||
import {
|
||||
MAX_MODEL_INVOCATION_RECORD_JSON_BYTES,
|
||||
ModelInvocationConflictError,
|
||||
ModelInvocationRepositoryUnavailableError,
|
||||
createModelInvocationMutationIdentity,
|
||||
normalizeModelInvocationCompletionRecord,
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationRepository,
|
||||
} from './modelInvocation';
|
||||
|
||||
export const MODEL_INVOCATION_RESOLUTION_SCHEMA =
|
||||
'qinglong/model-invocation-resolution@v1' as const;
|
||||
export const MODEL_INVOCATION_RESOLUTION_COMMAND_SCHEMA =
|
||||
'qinglong/model-invocation-resolution-command@v1' as const;
|
||||
export const MODEL_INVOCATION_RESOLUTION_DECISIONS = [
|
||||
'retry',
|
||||
'fail',
|
||||
'cancel',
|
||||
] as const;
|
||||
|
||||
export type ModelInvocationResolutionDecision =
|
||||
(typeof MODEL_INVOCATION_RESOLUTION_DECISIONS)[number];
|
||||
|
||||
export interface ModelInvocationResolutionRecord {
|
||||
readonly schema: typeof MODEL_INVOCATION_RESOLUTION_SCHEMA;
|
||||
readonly resolutionId: string;
|
||||
readonly invocationId: string;
|
||||
readonly projectId: string;
|
||||
readonly runId: string;
|
||||
readonly stepRunId: string;
|
||||
readonly traceId: string;
|
||||
readonly completionDigest: string;
|
||||
readonly decision: ModelInvocationResolutionDecision;
|
||||
readonly resolvedByUserId: string;
|
||||
readonly resolvedStepRunVersion: number;
|
||||
readonly stepRunMutationId: string;
|
||||
readonly stepRunMutationDigest: string;
|
||||
readonly resolvedStepRunDigest: string;
|
||||
readonly runEventId: string;
|
||||
readonly resolvedAtMs: number;
|
||||
readonly resolutionDigest: string;
|
||||
}
|
||||
|
||||
export interface ModelInvocationResolutionCommand {
|
||||
readonly schema: typeof MODEL_INVOCATION_RESOLUTION_COMMAND_SCHEMA;
|
||||
readonly completion: Readonly<ModelInvocationCompletionRecord>;
|
||||
readonly resolution: Readonly<ModelInvocationResolutionRecord>;
|
||||
readonly stepRunMutation: Readonly<StepRunMutation>;
|
||||
readonly commandDigest: string;
|
||||
}
|
||||
|
||||
export interface ModelInvocationResolutionRepository
|
||||
extends ModelInvocationRepository {
|
||||
findResolution(
|
||||
invocationId: string,
|
||||
): Promise<Readonly<ModelInvocationResolutionRecord> | null>;
|
||||
resolve(
|
||||
command: ModelInvocationResolutionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationResolutionRecord>>
|
||||
>;
|
||||
}
|
||||
|
||||
export class InvalidModelInvocationResolutionError extends TypeError {
|
||||
readonly code = 'MODEL_INVOCATION_RESOLUTION_INVALID';
|
||||
|
||||
constructor(message: string) {
|
||||
super(`Model invocation resolution is invalid: ${message}`);
|
||||
this.name = 'InvalidModelInvocationResolutionError';
|
||||
}
|
||||
}
|
||||
|
||||
const IDENTITY_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}$/;
|
||||
const DIGEST_PATTERN = /^[0-9a-f]{64}$/;
|
||||
const RESOLUTION_DIGEST_DOMAIN = Buffer.from(
|
||||
'qinglong/model-invocation-resolution-digest@v1\0',
|
||||
'utf8',
|
||||
);
|
||||
const RESOLUTION_COMMAND_DIGEST_DOMAIN = Buffer.from(
|
||||
'qinglong/model-invocation-resolution-command-digest@v1\0',
|
||||
'utf8',
|
||||
);
|
||||
|
||||
function invalid(message: string): never {
|
||||
throw new InvalidModelInvocationResolutionError(message);
|
||||
}
|
||||
|
||||
function dataRecord(value: unknown, label: string): Record<string, unknown> {
|
||||
if (
|
||||
!value ||
|
||||
typeof value !== 'object' ||
|
||||
Array.isArray(value) ||
|
||||
Object.getPrototypeOf(value) !== Object.prototype
|
||||
) {
|
||||
return invalid(`${label} must be a plain object`);
|
||||
}
|
||||
return value as Record<string, unknown>;
|
||||
}
|
||||
|
||||
function exactKeys(
|
||||
value: object,
|
||||
expected: readonly string[],
|
||||
label: string,
|
||||
): void {
|
||||
const actual = Object.keys(value).sort();
|
||||
const canonical = [...expected].sort();
|
||||
if (
|
||||
actual.length !== canonical.length ||
|
||||
actual.some((key, index) => key !== canonical[index])
|
||||
) {
|
||||
invalid(`${label} shape is invalid`);
|
||||
}
|
||||
}
|
||||
|
||||
function identifier(value: unknown, label: string): string {
|
||||
if (typeof value !== 'string' || !IDENTITY_PATTERN.test(value)) {
|
||||
return invalid(`${label} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function digest(value: unknown, label: string): string {
|
||||
if (typeof value !== 'string' || !DIGEST_PATTERN.test(value)) {
|
||||
return invalid(`${label} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function integer(
|
||||
value: unknown,
|
||||
minimum: number,
|
||||
maximum: number,
|
||||
label: string,
|
||||
): number {
|
||||
if (
|
||||
!Number.isSafeInteger(value) ||
|
||||
(value as number) < minimum ||
|
||||
(value as number) > maximum
|
||||
) {
|
||||
return invalid(`${label} is invalid`);
|
||||
}
|
||||
return value as number;
|
||||
}
|
||||
|
||||
function hash(domain: Uint8Array, value: unknown): string {
|
||||
return createHash('sha256')
|
||||
.update(domain)
|
||||
.update(JSON.stringify(value))
|
||||
.digest('hex');
|
||||
}
|
||||
|
||||
function withoutDigest(
|
||||
value: Readonly<ModelInvocationResolutionRecord>,
|
||||
): Omit<ModelInvocationResolutionRecord, 'resolutionDigest'> {
|
||||
const { resolutionDigest: _resolutionDigest, ...unsigned } = value;
|
||||
return unsigned;
|
||||
}
|
||||
|
||||
export function resolutionTransition(
|
||||
decision: ModelInvocationResolutionDecision,
|
||||
): Readonly<{
|
||||
to: StepRunStatus;
|
||||
resultCode?: string;
|
||||
errorSummary?: string;
|
||||
}> {
|
||||
if (decision === 'retry') return Object.freeze({ to: 'ready' });
|
||||
if (decision === 'fail') {
|
||||
return Object.freeze({
|
||||
to: 'failed',
|
||||
resultCode: 'model_outcome_rejected',
|
||||
errorSummary: 'Unknown model outcome rejected by operator',
|
||||
});
|
||||
}
|
||||
if (decision === 'cancel') {
|
||||
return Object.freeze({
|
||||
to: 'cancelled',
|
||||
resultCode: 'model_outcome_cancelled',
|
||||
});
|
||||
}
|
||||
return invalid('decision is invalid');
|
||||
}
|
||||
|
||||
export function normalizeModelInvocationResolutionRecord(
|
||||
value: ModelInvocationResolutionRecord,
|
||||
): Readonly<ModelInvocationResolutionRecord> {
|
||||
const candidate = dataRecord(value, 'resolution record');
|
||||
exactKeys(
|
||||
candidate,
|
||||
[
|
||||
'completionDigest',
|
||||
'decision',
|
||||
'invocationId',
|
||||
'projectId',
|
||||
'resolutionDigest',
|
||||
'resolutionId',
|
||||
'resolvedAtMs',
|
||||
'resolvedByUserId',
|
||||
'resolvedStepRunDigest',
|
||||
'resolvedStepRunVersion',
|
||||
'runEventId',
|
||||
'runId',
|
||||
'schema',
|
||||
'stepRunId',
|
||||
'stepRunMutationDigest',
|
||||
'stepRunMutationId',
|
||||
'traceId',
|
||||
],
|
||||
'resolution record',
|
||||
);
|
||||
if (
|
||||
value.schema !== MODEL_INVOCATION_RESOLUTION_SCHEMA ||
|
||||
!MODEL_INVOCATION_RESOLUTION_DECISIONS.includes(value.decision)
|
||||
) {
|
||||
invalid('resolution schema or decision is invalid');
|
||||
}
|
||||
const normalized = Object.freeze({
|
||||
schema: MODEL_INVOCATION_RESOLUTION_SCHEMA,
|
||||
resolutionId: identifier(value.resolutionId, 'resolution id'),
|
||||
invocationId: identifier(value.invocationId, 'invocation id'),
|
||||
projectId: identifier(value.projectId, 'project id'),
|
||||
runId: identifier(value.runId, 'Run id'),
|
||||
stepRunId: identifier(value.stepRunId, 'StepRun id'),
|
||||
traceId: identifier(value.traceId, 'trace id'),
|
||||
completionDigest: digest(value.completionDigest, 'completion digest'),
|
||||
decision: value.decision,
|
||||
resolvedByUserId: identifier(value.resolvedByUserId, 'resolving user id'),
|
||||
resolvedStepRunVersion: integer(
|
||||
value.resolvedStepRunVersion,
|
||||
4,
|
||||
2_147_483_647,
|
||||
'resolved StepRun version',
|
||||
),
|
||||
stepRunMutationId: identifier(
|
||||
value.stepRunMutationId,
|
||||
'StepRun mutation id',
|
||||
),
|
||||
stepRunMutationDigest: digest(
|
||||
value.stepRunMutationDigest,
|
||||
'StepRun mutation digest',
|
||||
),
|
||||
resolvedStepRunDigest: digest(
|
||||
value.resolvedStepRunDigest,
|
||||
'resolved StepRun digest',
|
||||
),
|
||||
runEventId: identifier(value.runEventId, 'RunEvent id'),
|
||||
resolvedAtMs: integer(
|
||||
value.resolvedAtMs,
|
||||
0,
|
||||
Number.MAX_SAFE_INTEGER,
|
||||
'resolved time',
|
||||
),
|
||||
resolutionDigest: digest(value.resolutionDigest, 'resolution digest'),
|
||||
});
|
||||
const identity = createModelInvocationMutationIdentity(
|
||||
normalized.invocationId,
|
||||
'resolution',
|
||||
);
|
||||
if (
|
||||
normalized.resolutionId !== identity.dedupeKey ||
|
||||
normalized.stepRunMutationId !== identity.mutationId ||
|
||||
normalized.runEventId !== identity.eventId ||
|
||||
hash(RESOLUTION_DIGEST_DOMAIN, withoutDigest(normalized)) !==
|
||||
normalized.resolutionDigest
|
||||
) {
|
||||
invalid('resolution identity or digest is invalid');
|
||||
}
|
||||
if (
|
||||
Buffer.byteLength(JSON.stringify(normalized), 'utf8') >
|
||||
MAX_MODEL_INVOCATION_RECORD_JSON_BYTES
|
||||
) {
|
||||
invalid('resolution record exceeds its JSON budget');
|
||||
}
|
||||
return normalized;
|
||||
}
|
||||
|
||||
export function createModelInvocationResolutionCommand(
|
||||
completionValue: ModelInvocationCompletionRecord,
|
||||
decision: ModelInvocationResolutionDecision,
|
||||
resolvedByUserIdValue: string,
|
||||
mutationValue: StepRunMutation,
|
||||
): Readonly<ModelInvocationResolutionCommand> {
|
||||
const completion = normalizeModelInvocationCompletionRecord(completionValue);
|
||||
const transition = resolutionTransition(decision);
|
||||
const resolvedByUserId = identifier(
|
||||
resolvedByUserIdValue,
|
||||
'resolving user id',
|
||||
);
|
||||
const mutation = normalizeStepRunMutation(mutationValue);
|
||||
const identity = createModelInvocationMutationIdentity(
|
||||
completion.invocationId,
|
||||
'resolution',
|
||||
);
|
||||
if (
|
||||
completion.outcome !== 'outcome_unknown' ||
|
||||
mutation.previousStatus !== 'lost' ||
|
||||
mutation.expectedStepRunVersion !== completion.completedStepRunVersion ||
|
||||
mutation.expectedStepRunDigest !== completion.completedStepRunDigest ||
|
||||
mutation.mutationId !== identity.mutationId ||
|
||||
mutation.runId !== completion.runId ||
|
||||
mutation.stepRun.id !== completion.stepRunId ||
|
||||
mutation.stepRun.runId !== completion.runId ||
|
||||
mutation.stepRun.kind !== 'model' ||
|
||||
mutation.stepRun.status !== transition.to ||
|
||||
mutation.event.id !== identity.eventId ||
|
||||
mutation.event.dedupeKey !== identity.dedupeKey ||
|
||||
mutation.event.type !== `step.${transition.to}` ||
|
||||
mutation.event.actorType !== 'user' ||
|
||||
mutation.event.actorId !== resolvedByUserId ||
|
||||
mutation.stepRun.updatedAtMs < completion.completedAtMs ||
|
||||
(transition.to === 'ready'
|
||||
? mutation.stepRun.outputRef !== null ||
|
||||
mutation.stepRun.resultCode !== null ||
|
||||
mutation.stepRun.errorSummary !== null
|
||||
: mutation.stepRun.resultCode !== transition.resultCode ||
|
||||
mutation.stepRun.errorSummary !==
|
||||
(transition.errorSummary === undefined
|
||||
? null
|
||||
: transition.errorSummary))
|
||||
) {
|
||||
invalid('resolution completion or StepRun mutation is not exact');
|
||||
}
|
||||
const unsigned = Object.freeze({
|
||||
schema: MODEL_INVOCATION_RESOLUTION_SCHEMA,
|
||||
resolutionId: identity.dedupeKey,
|
||||
invocationId: completion.invocationId,
|
||||
projectId: completion.projectId,
|
||||
runId: completion.runId,
|
||||
stepRunId: completion.stepRunId,
|
||||
traceId: completion.traceId,
|
||||
completionDigest: completion.completionDigest,
|
||||
decision,
|
||||
resolvedByUserId,
|
||||
resolvedStepRunVersion: mutation.stepRun.version,
|
||||
stepRunMutationId: mutation.mutationId,
|
||||
stepRunMutationDigest: mutation.mutationDigest,
|
||||
resolvedStepRunDigest: mutation.stepRun.stepRunDigest,
|
||||
runEventId: mutation.event.id,
|
||||
resolvedAtMs: mutation.stepRun.updatedAtMs,
|
||||
});
|
||||
const resolution = normalizeModelInvocationResolutionRecord({
|
||||
...unsigned,
|
||||
resolutionDigest: hash(RESOLUTION_DIGEST_DOMAIN, unsigned),
|
||||
});
|
||||
const commandUnsigned = Object.freeze({
|
||||
schema: MODEL_INVOCATION_RESOLUTION_COMMAND_SCHEMA,
|
||||
completion,
|
||||
resolution,
|
||||
stepRunMutation: mutation,
|
||||
});
|
||||
return Object.freeze({
|
||||
...commandUnsigned,
|
||||
commandDigest: hash(RESOLUTION_COMMAND_DIGEST_DOMAIN, commandUnsigned),
|
||||
});
|
||||
}
|
||||
|
||||
export function normalizeModelInvocationResolutionCommand(
|
||||
value: ModelInvocationResolutionCommand,
|
||||
): Readonly<ModelInvocationResolutionCommand> {
|
||||
const candidate = dataRecord(value, 'resolution command');
|
||||
exactKeys(
|
||||
candidate,
|
||||
['commandDigest', 'completion', 'resolution', 'schema', 'stepRunMutation'],
|
||||
'resolution command',
|
||||
);
|
||||
if (value.schema !== MODEL_INVOCATION_RESOLUTION_COMMAND_SCHEMA) {
|
||||
invalid('resolution command schema is invalid');
|
||||
}
|
||||
const completion = normalizeModelInvocationCompletionRecord(value.completion);
|
||||
const resolution = normalizeModelInvocationResolutionRecord(value.resolution);
|
||||
const mutation = normalizeStepRunMutation(value.stepRunMutation);
|
||||
const canonical = createModelInvocationResolutionCommand(
|
||||
completion,
|
||||
resolution.decision,
|
||||
resolution.resolvedByUserId,
|
||||
mutation,
|
||||
);
|
||||
if (
|
||||
canonical.resolution.resolutionDigest !== resolution.resolutionDigest ||
|
||||
digest(value.commandDigest, 'command digest') !== canonical.commandDigest
|
||||
) {
|
||||
invalid('resolution command is not canonical');
|
||||
}
|
||||
return canonical;
|
||||
}
|
||||
|
||||
export function assertResolutionMatchesDecision(
|
||||
resolutionValue: ModelInvocationResolutionRecord,
|
||||
completionValue: ModelInvocationCompletionRecord,
|
||||
decision: ModelInvocationResolutionDecision,
|
||||
resolvedByUserId: string,
|
||||
): Readonly<ModelInvocationResolutionRecord> {
|
||||
const resolution = normalizeModelInvocationResolutionRecord(resolutionValue);
|
||||
const completion = normalizeModelInvocationCompletionRecord(completionValue);
|
||||
if (
|
||||
resolution.invocationId !== completion.invocationId ||
|
||||
resolution.projectId !== completion.projectId ||
|
||||
resolution.runId !== completion.runId ||
|
||||
resolution.stepRunId !== completion.stepRunId ||
|
||||
resolution.traceId !== completion.traceId ||
|
||||
resolution.completionDigest !== completion.completionDigest ||
|
||||
resolution.decision !== decision ||
|
||||
resolution.resolvedByUserId !== resolvedByUserId
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return resolution;
|
||||
}
|
||||
|
||||
export interface ResolveModelInvocationOptions {
|
||||
readonly invocationId: string;
|
||||
readonly decision: ModelInvocationResolutionDecision;
|
||||
readonly resolvedByUserId: string;
|
||||
readonly resolvedAtMs: number;
|
||||
}
|
||||
|
||||
const MAX_RESOLUTION_COORDINATOR_ATTEMPTS = 3;
|
||||
|
||||
export class DurableModelInvocationResolutionCoordinator {
|
||||
constructor(
|
||||
private readonly repository: ModelInvocationResolutionRepository,
|
||||
) {
|
||||
if (
|
||||
!repository ||
|
||||
typeof repository.findCompletion !== 'function' ||
|
||||
typeof repository.findResolution !== 'function' ||
|
||||
typeof repository.readAuthority !== 'function' ||
|
||||
typeof repository.resolve !== 'function'
|
||||
) {
|
||||
throw new ModelInvocationRepositoryUnavailableError();
|
||||
}
|
||||
}
|
||||
|
||||
async resolve(
|
||||
optionsValue: ResolveModelInvocationOptions,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationResolutionRecord>>
|
||||
> {
|
||||
const options = dataRecord(optionsValue, 'resolution options');
|
||||
exactKeys(
|
||||
options,
|
||||
['decision', 'invocationId', 'resolvedAtMs', 'resolvedByUserId'],
|
||||
'resolution options',
|
||||
);
|
||||
const invocationId = identifier(optionsValue.invocationId, 'invocation id');
|
||||
const decision = optionsValue.decision;
|
||||
resolutionTransition(decision);
|
||||
const resolvedByUserId = identifier(
|
||||
optionsValue.resolvedByUserId,
|
||||
'resolving user id',
|
||||
);
|
||||
const resolvedAtMs = integer(
|
||||
optionsValue.resolvedAtMs,
|
||||
0,
|
||||
Number.MAX_SAFE_INTEGER,
|
||||
'resolved time',
|
||||
);
|
||||
const completion = await this.repository.findCompletion(invocationId);
|
||||
if (!completion || completion.outcome !== 'outcome_unknown') {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const existing = await this.repository.findResolution(invocationId);
|
||||
if (existing) {
|
||||
return Object.freeze({
|
||||
status: 'existing',
|
||||
record: assertResolutionMatchesDecision(
|
||||
existing,
|
||||
completion,
|
||||
decision,
|
||||
resolvedByUserId,
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
for (
|
||||
let attempt = 0;
|
||||
attempt < MAX_RESOLUTION_COORDINATOR_ATTEMPTS;
|
||||
attempt += 1
|
||||
) {
|
||||
const authority = await this.repository.readAuthority({
|
||||
projectId: completion.projectId,
|
||||
runId: completion.runId,
|
||||
stepRunId: completion.stepRunId,
|
||||
});
|
||||
if (
|
||||
!authority ||
|
||||
authority.stepRun.kind !== 'model' ||
|
||||
authority.stepRun.status !== 'lost' ||
|
||||
authority.stepRun.version !== completion.completedStepRunVersion ||
|
||||
authority.stepRun.stepRunDigest !== completion.completedStepRunDigest
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const transition = resolutionTransition(decision);
|
||||
const identity = createModelInvocationMutationIdentity(
|
||||
invocationId,
|
||||
'resolution',
|
||||
);
|
||||
const command = createModelInvocationResolutionCommand(
|
||||
completion,
|
||||
decision,
|
||||
resolvedByUserId,
|
||||
transitionStepRunMutation(
|
||||
authority.stepRun,
|
||||
{
|
||||
expectedVersion: authority.stepRun.version,
|
||||
expectedDigest: authority.stepRun.stepRunDigest,
|
||||
mutationId: identity.mutationId,
|
||||
to: transition.to,
|
||||
atMs: resolvedAtMs,
|
||||
...(transition.resultCode === undefined
|
||||
? {}
|
||||
: { resultCode: transition.resultCode }),
|
||||
...(transition.errorSummary === undefined
|
||||
? {}
|
||||
: { errorSummary: transition.errorSummary }),
|
||||
},
|
||||
{
|
||||
expectedRunVersion: authority.runVersion,
|
||||
expectedRunEventSequence: authority.runEventSequence,
|
||||
eventId: identity.eventId,
|
||||
dedupeKey: identity.dedupeKey,
|
||||
actor: { type: 'user', id: resolvedByUserId },
|
||||
},
|
||||
),
|
||||
);
|
||||
try {
|
||||
return await this.repository.resolve(command);
|
||||
} catch (error) {
|
||||
const stored = await this.#resolutionAfterFailure(
|
||||
completion,
|
||||
decision,
|
||||
resolvedByUserId,
|
||||
error,
|
||||
);
|
||||
if (stored) {
|
||||
return Object.freeze({ status: 'existing', record: stored });
|
||||
}
|
||||
if (
|
||||
!(error instanceof ModelInvocationConflictError) ||
|
||||
attempt + 1 >= MAX_RESOLUTION_COORDINATOR_ATTEMPTS
|
||||
) {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
|
||||
async #resolutionAfterFailure(
|
||||
completion: Readonly<ModelInvocationCompletionRecord>,
|
||||
decision: ModelInvocationResolutionDecision,
|
||||
resolvedByUserId: string,
|
||||
original: unknown,
|
||||
): Promise<Readonly<ModelInvocationResolutionRecord> | null> {
|
||||
try {
|
||||
const stored = await this.repository.findResolution(
|
||||
completion.invocationId,
|
||||
);
|
||||
return stored
|
||||
? assertResolutionMatchesDecision(
|
||||
stored,
|
||||
completion,
|
||||
decision,
|
||||
resolvedByUserId,
|
||||
)
|
||||
: null;
|
||||
} catch {
|
||||
throw original;
|
||||
}
|
||||
}
|
||||
}
|
||||
+258
@@ -0,0 +1,258 @@
|
||||
import type { PostgresPool } from '@qinglong/runtime-core';
|
||||
|
||||
import {
|
||||
normalizeModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceQuote,
|
||||
} from '../../pricing/pricing';
|
||||
import {
|
||||
ModelInvocationProjectQuotaExceededError,
|
||||
createModelInvocationQuotaReservation,
|
||||
normalizeModelInvocationQuotaAdmission,
|
||||
type ModelInvocationQuotaAdmission,
|
||||
type ModelInvocationQuotaReservation,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
ModelInvocationConflictError,
|
||||
normalizeModelInvocationStartCommand,
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationStartCommand,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
|
||||
import type { Row } from './authority';
|
||||
import { integer, unavailable } from './authority';
|
||||
import { parsePriceQuote, parseQuotaReservation, parseStart } from './codec';
|
||||
import {
|
||||
applyMutation,
|
||||
assertCurrent,
|
||||
insertPriceQuote,
|
||||
insertQuotaReservation,
|
||||
insertStart,
|
||||
quotaWindowUsage,
|
||||
} from './mutations';
|
||||
import { priceQuoteRows, quotaReservationRows, startRows } from './queries';
|
||||
import { runPostgresModelInvocationTransaction } from './transaction';
|
||||
|
||||
export async function admitOperation(
|
||||
pool: PostgresPool,
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
): Promise<Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>> {
|
||||
const command = normalizeModelInvocationStartCommand(commandValue);
|
||||
const start = command.start;
|
||||
return runPostgresModelInvocationTransaction(pool, async (client) => {
|
||||
const existing = await startRows(
|
||||
client,
|
||||
`start.invocation_id = $1 OR
|
||||
start.mutation_id = $2 OR start.run_event_id = $3`,
|
||||
[start.invocationId, start.stepRunMutationId, start.runEventId],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseStart(existing[0]);
|
||||
if (JSON.stringify(stored) !== JSON.stringify(start)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
await assertCurrent(client, command.stepRunMutation, start.projectId);
|
||||
await applyMutation(client, command.stepRunMutation);
|
||||
await insertStart(client, start);
|
||||
return Object.freeze({ status: 'created' as const, record: start });
|
||||
});
|
||||
}
|
||||
|
||||
export async function admitWithQuotaOperation(
|
||||
pool: PostgresPool,
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
admissionValue: ModelInvocationQuotaAdmission,
|
||||
): Promise<Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>> {
|
||||
const command = normalizeModelInvocationStartCommand(commandValue);
|
||||
const admission = normalizeModelInvocationQuotaAdmission(admissionValue);
|
||||
const start = command.start;
|
||||
if (
|
||||
admission.invocationId !== start.invocationId ||
|
||||
admission.projectId !== start.projectId ||
|
||||
admission.modelPolicyRevision !== start.policyRevision
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return runPostgresModelInvocationTransaction(pool, async (client) => {
|
||||
const existing = await startRows(
|
||||
client,
|
||||
`start.invocation_id = $1 OR
|
||||
start.mutation_id = $2 OR start.run_event_id = $3`,
|
||||
[start.invocationId, start.stepRunMutationId, start.runEventId],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseStart(existing[0]);
|
||||
const reservations = await quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = $1',
|
||||
[start.invocationId],
|
||||
);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(start) ||
|
||||
reservations.length !== 1 ||
|
||||
parseQuotaReservation(reservations[0]!).admissionDigest !==
|
||||
admission.admissionDigest
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
const observation = await client.query<Row>(
|
||||
`SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint
|
||||
AS "observedAtMs"`,
|
||||
);
|
||||
const observedRow = observation.rows[0];
|
||||
if (observation.rows.length !== 1 || !observedRow) throw unavailable();
|
||||
const reservation = createModelInvocationQuotaReservation(
|
||||
admission,
|
||||
integer(observedRow, 'observedAtMs'),
|
||||
);
|
||||
await client.query(
|
||||
`SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`,
|
||||
[
|
||||
JSON.stringify([
|
||||
start.projectId,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowMs,
|
||||
]),
|
||||
],
|
||||
);
|
||||
const usage = await quotaWindowUsage(
|
||||
client,
|
||||
start.projectId,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowMs,
|
||||
);
|
||||
if (
|
||||
usage.invocationCount + 1 > reservation.maxInvocations ||
|
||||
usage.effectiveTokens + reservation.reservedTokens >
|
||||
reservation.maxTokens ||
|
||||
(reservation.maxCostMicros !== null &&
|
||||
(usage.unknownCostInvocations !== 0 ||
|
||||
usage.effectiveCostMicros + reservation.reservedCostMicros! >
|
||||
reservation.maxCostMicros))
|
||||
) {
|
||||
throw new ModelInvocationProjectQuotaExceededError();
|
||||
}
|
||||
await assertCurrent(client, command.stepRunMutation, start.projectId);
|
||||
await applyMutation(client, command.stepRunMutation);
|
||||
await insertStart(client, start);
|
||||
await insertQuotaReservation(client, reservation);
|
||||
return Object.freeze({ status: 'created' as const, record: start });
|
||||
});
|
||||
}
|
||||
|
||||
export async function admitWithPricingOperation(
|
||||
pool: PostgresPool,
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
quoteValue: ModelInvocationPriceQuote,
|
||||
admissionValue?: ModelInvocationQuotaAdmission,
|
||||
): Promise<Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>> {
|
||||
const command = normalizeModelInvocationStartCommand(commandValue);
|
||||
const quote = normalizeModelInvocationPriceQuote(quoteValue);
|
||||
const admission =
|
||||
admissionValue === undefined
|
||||
? undefined
|
||||
: normalizeModelInvocationQuotaAdmission(admissionValue);
|
||||
const start = command.start;
|
||||
if (
|
||||
quote.invocationId !== start.invocationId ||
|
||||
quote.projectId !== start.projectId ||
|
||||
quote.modelPolicyRevision !== start.policyRevision ||
|
||||
quote.provider !== start.provider ||
|
||||
quote.model !== start.model ||
|
||||
quote.maxOutputTokens !== start.maxOutputTokens ||
|
||||
(admission !== undefined &&
|
||||
(admission.invocationId !== start.invocationId ||
|
||||
admission.projectId !== start.projectId ||
|
||||
admission.modelPolicyRevision !== start.policyRevision ||
|
||||
(admission.maxCostMicros !== null &&
|
||||
admission.reservedCostMicros !== quote.reservedCostMicros)))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return runPostgresModelInvocationTransaction(pool, async (client) => {
|
||||
const existing = await startRows(
|
||||
client,
|
||||
`start.invocation_id = $1 OR
|
||||
start.mutation_id = $2 OR start.run_event_id = $3`,
|
||||
[start.invocationId, start.stepRunMutationId, start.runEventId],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const [quotes, reservations] = await Promise.all([
|
||||
priceQuoteRows(client, 'quote.invocation_id = $1', [
|
||||
start.invocationId,
|
||||
]),
|
||||
quotaReservationRows(client, 'reservation.invocation_id = $1', [
|
||||
start.invocationId,
|
||||
]),
|
||||
]);
|
||||
const stored = parseStart(existing[0]);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(start) ||
|
||||
quotes.length !== 1 ||
|
||||
JSON.stringify(parsePriceQuote(quotes[0]!)) !== JSON.stringify(quote) ||
|
||||
reservations.length !== (admission ? 1 : 0) ||
|
||||
(admission !== undefined &&
|
||||
parseQuotaReservation(reservations[0]!).admissionDigest !==
|
||||
admission.admissionDigest)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
let reservation: Readonly<ModelInvocationQuotaReservation> | undefined;
|
||||
if (admission) {
|
||||
const observation = await client.query<Row>(
|
||||
`SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint
|
||||
AS "observedAtMs"`,
|
||||
);
|
||||
const observedRow = observation.rows[0];
|
||||
if (observation.rows.length !== 1 || !observedRow) {
|
||||
throw unavailable();
|
||||
}
|
||||
reservation = createModelInvocationQuotaReservation(
|
||||
admission,
|
||||
integer(observedRow, 'observedAtMs'),
|
||||
);
|
||||
await client.query(
|
||||
`SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`,
|
||||
[
|
||||
JSON.stringify([
|
||||
start.projectId,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowMs,
|
||||
]),
|
||||
],
|
||||
);
|
||||
const usage = await quotaWindowUsage(
|
||||
client,
|
||||
start.projectId,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowMs,
|
||||
);
|
||||
if (
|
||||
usage.invocationCount + 1 > reservation.maxInvocations ||
|
||||
usage.effectiveTokens + reservation.reservedTokens >
|
||||
reservation.maxTokens ||
|
||||
(reservation.maxCostMicros !== null &&
|
||||
(usage.unknownCostInvocations !== 0 ||
|
||||
usage.effectiveCostMicros + reservation.reservedCostMicros! >
|
||||
reservation.maxCostMicros))
|
||||
) {
|
||||
throw new ModelInvocationProjectQuotaExceededError();
|
||||
}
|
||||
}
|
||||
await assertCurrent(client, command.stepRunMutation, start.projectId);
|
||||
await applyMutation(client, command.stepRunMutation);
|
||||
await insertStart(client, start);
|
||||
await insertPriceQuote(client, quote);
|
||||
if (reservation) await insertQuotaReservation(client, reservation);
|
||||
return Object.freeze({ status: 'created' as const, record: start });
|
||||
});
|
||||
}
|
||||
+118
@@ -0,0 +1,118 @@
|
||||
import type { PostgresQueryable } from '@qinglong/runtime-core';
|
||||
|
||||
import {
|
||||
PluginPackagePromptOutputArtifactConflictError,
|
||||
PluginPackagePromptOutputArtifactUnavailableError,
|
||||
} from '../../prompt-output/pluginPackagePromptOutputArtifact';
|
||||
import { ModelInvocationUsageSummaryLimitExceededError } from '../../usage/usageLedger';
|
||||
import { ModelInvocationProjectQuotaExceededError } from '../../usage/usageQuota';
|
||||
import {
|
||||
MAX_MODEL_INVOCATION_RECOVERY_PAGE_SIZE,
|
||||
ModelInvocationConflictError,
|
||||
ModelInvocationRepositoryUnavailableError,
|
||||
} from '../modelInvocation';
|
||||
|
||||
export type Row = Record<string, unknown>;
|
||||
export type Queryable = Pick<PostgresQueryable, 'query'>;
|
||||
|
||||
export const IDENTIFIER_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}$/;
|
||||
export const TERMINAL_RUN_STATUSES = new Set([
|
||||
'succeeded',
|
||||
'failed',
|
||||
'cancelled',
|
||||
'timed_out',
|
||||
]);
|
||||
export const RETRYABLE_SQL_STATES = new Set(['40001', '40P01']);
|
||||
export const MAX_TRANSACTION_ATTEMPTS = 3;
|
||||
|
||||
export function unavailable(
|
||||
cause?: unknown,
|
||||
): ModelInvocationRepositoryUnavailableError {
|
||||
return new ModelInvocationRepositoryUnavailableError({
|
||||
cause: cause instanceof Error ? cause : undefined,
|
||||
});
|
||||
}
|
||||
|
||||
export function identifier(value: unknown): string {
|
||||
if (typeof value !== 'string' || !IDENTIFIER_PATTERN.test(value)) {
|
||||
throw unavailable();
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
export function text(row: Row, key: string): string {
|
||||
const value = row[key];
|
||||
if (typeof value !== 'string') throw unavailable();
|
||||
return value;
|
||||
}
|
||||
|
||||
export function integer(row: Row, key: string): number {
|
||||
const value = row[key];
|
||||
if (typeof value === 'number' && Number.isSafeInteger(value) && value >= 0) {
|
||||
return value;
|
||||
}
|
||||
if (typeof value === 'string' && /^(0|[1-9]\d*)$/.test(value)) {
|
||||
const parsed = Number(value);
|
||||
if (Number.isSafeInteger(parsed)) return parsed;
|
||||
}
|
||||
throw unavailable();
|
||||
}
|
||||
|
||||
export function nullableInteger(row: Row, key: string): number | null {
|
||||
return row[key] === null ? null : integer(row, key);
|
||||
}
|
||||
|
||||
export function recoveryLimit(value: number): number {
|
||||
if (
|
||||
!Number.isSafeInteger(value) ||
|
||||
value < 1 ||
|
||||
value > MAX_MODEL_INVOCATION_RECOVERY_PAGE_SIZE
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
export function jsonObject(row: Row, key: string): Record<string, unknown> {
|
||||
const value = row[key];
|
||||
if (
|
||||
value &&
|
||||
typeof value === 'object' &&
|
||||
!Array.isArray(value) &&
|
||||
Object.getPrototypeOf(value) === Object.prototype
|
||||
) {
|
||||
return value as Record<string, unknown>;
|
||||
}
|
||||
if (typeof value === 'string') {
|
||||
try {
|
||||
return JSON.parse(value) as Record<string, unknown>;
|
||||
} catch {
|
||||
throw unavailable();
|
||||
}
|
||||
}
|
||||
throw unavailable();
|
||||
}
|
||||
|
||||
export function sqlState(error: unknown): string {
|
||||
if (!error || typeof error !== 'object') return '';
|
||||
const value = (error as { code?: unknown }).code;
|
||||
return typeof value === 'string' ? value : '';
|
||||
}
|
||||
|
||||
export function mapStorageError(error: unknown): Error {
|
||||
if (
|
||||
error instanceof ModelInvocationConflictError ||
|
||||
error instanceof ModelInvocationRepositoryUnavailableError ||
|
||||
error instanceof ModelInvocationUsageSummaryLimitExceededError ||
|
||||
error instanceof ModelInvocationProjectQuotaExceededError ||
|
||||
error instanceof PluginPackagePromptOutputArtifactConflictError ||
|
||||
error instanceof PluginPackagePromptOutputArtifactUnavailableError
|
||||
) {
|
||||
return error;
|
||||
}
|
||||
const state = sqlState(error);
|
||||
if (state === '23503' || state === '23505' || state === '23514') {
|
||||
return new ModelInvocationConflictError();
|
||||
}
|
||||
return unavailable(error);
|
||||
}
|
||||
@@ -0,0 +1,367 @@
|
||||
import {
|
||||
normalizeStepRunRecord,
|
||||
type StepRunRecord,
|
||||
} from '@qinglong/runtime-core/step-run';
|
||||
|
||||
import {
|
||||
normalizeModelInvocationPriceQuote,
|
||||
normalizeModelInvocationPriceSettlement,
|
||||
type ModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceSettlement,
|
||||
} from '../../pricing/pricing';
|
||||
import {
|
||||
normalizeModelInvocationUsageLedgerRecord,
|
||||
type ModelInvocationUsageLedgerRecord,
|
||||
} from '../../usage/usageLedger';
|
||||
import {
|
||||
normalizeModelInvocationQuotaReservation,
|
||||
normalizeModelInvocationQuotaSettlement,
|
||||
type ModelInvocationQuotaReservation,
|
||||
type ModelInvocationQuotaSettlement,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
normalizeModelInvocationCompletionRecord,
|
||||
normalizeModelInvocationStartRecord,
|
||||
type ModelInvocationAuthoritySnapshot,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
import {
|
||||
normalizeModelInvocationResolutionRecord,
|
||||
type ModelInvocationResolutionRecord,
|
||||
} from '../modelInvocationResolution';
|
||||
|
||||
import type { Row } from './authority';
|
||||
import {
|
||||
integer,
|
||||
jsonObject,
|
||||
nullableInteger,
|
||||
text,
|
||||
unavailable,
|
||||
} from './authority';
|
||||
|
||||
export function parseStart(row: Row): Readonly<ModelInvocationStartRecord> {
|
||||
let start: Readonly<ModelInvocationStartRecord>;
|
||||
try {
|
||||
start = normalizeModelInvocationStartRecord(
|
||||
jsonObject(row, 'recordJson') as unknown as ModelInvocationStartRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
start.invocationId !== text(row, 'invocationId') ||
|
||||
start.projectId !== text(row, 'projectId') ||
|
||||
start.runId !== text(row, 'runId') ||
|
||||
start.stepRunId !== text(row, 'stepRunId') ||
|
||||
start.traceId !== text(row, 'traceId') ||
|
||||
start.provider !== text(row, 'provider') ||
|
||||
start.model !== text(row, 'model') ||
|
||||
start.policyRevision !== text(row, 'policyRevision') ||
|
||||
start.requestDigest !== text(row, 'requestDigest') ||
|
||||
start.inputBytes !== integer(row, 'inputBytes') ||
|
||||
start.maxOutputTokens !== integer(row, 'maxOutputTokens') ||
|
||||
start.deadlineAtMs !== integer(row, 'deadlineAtMs') ||
|
||||
start.admittedAtMs !== integer(row, 'admittedAtMs') ||
|
||||
start.stepRunMutationId !== text(row, 'mutationId') ||
|
||||
start.stepRunMutationDigest !== text(row, 'mutationDigest') ||
|
||||
start.startedStepRunDigest !== text(row, 'stepRunDigest') ||
|
||||
start.startedStepRunVersion !== integer(row, 'stepRunVersion') ||
|
||||
start.runEventId !== text(row, 'runEventId') ||
|
||||
start.startDigest !== text(row, 'startDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return start;
|
||||
}
|
||||
|
||||
export function parseCompletion(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationCompletionRecord> {
|
||||
let completion: Readonly<ModelInvocationCompletionRecord>;
|
||||
try {
|
||||
completion = normalizeModelInvocationCompletionRecord(
|
||||
jsonObject(
|
||||
row,
|
||||
'recordJson',
|
||||
) as unknown as ModelInvocationCompletionRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
completion.invocationId !== text(row, 'invocationId') ||
|
||||
completion.projectId !== text(row, 'projectId') ||
|
||||
completion.runId !== text(row, 'runId') ||
|
||||
completion.stepRunId !== text(row, 'stepRunId') ||
|
||||
completion.traceId !== text(row, 'traceId') ||
|
||||
completion.startDigest !== text(row, 'startDigest') ||
|
||||
completion.outcome !== text(row, 'outcome') ||
|
||||
completion.outputBytes !== integer(row, 'outputBytes') ||
|
||||
completion.errorCode !== row.errorCode ||
|
||||
completion.completedAtMs !== integer(row, 'completedAtMs') ||
|
||||
completion.stepRunMutationId !== text(row, 'mutationId') ||
|
||||
completion.stepRunMutationDigest !== text(row, 'mutationDigest') ||
|
||||
completion.completedStepRunDigest !== text(row, 'stepRunDigest') ||
|
||||
completion.completedStepRunVersion !== integer(row, 'stepRunVersion') ||
|
||||
completion.runEventId !== text(row, 'runEventId') ||
|
||||
completion.completionDigest !== text(row, 'completionDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return completion;
|
||||
}
|
||||
|
||||
export function parseUsage(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationUsageLedgerRecord> {
|
||||
let usage: Readonly<ModelInvocationUsageLedgerRecord>;
|
||||
try {
|
||||
usage = normalizeModelInvocationUsageLedgerRecord(
|
||||
jsonObject(
|
||||
row,
|
||||
'recordJson',
|
||||
) as unknown as ModelInvocationUsageLedgerRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
usage.invocationId !== text(row, 'invocationId') ||
|
||||
usage.projectId !== text(row, 'projectId') ||
|
||||
usage.runId !== text(row, 'runId') ||
|
||||
usage.stepRunId !== text(row, 'stepRunId') ||
|
||||
usage.traceId !== text(row, 'traceId') ||
|
||||
usage.provider !== text(row, 'provider') ||
|
||||
usage.model !== text(row, 'model') ||
|
||||
usage.policyRevision !== text(row, 'policyRevision') ||
|
||||
usage.completionDigest !== text(row, 'completionDigest') ||
|
||||
usage.outcome !== text(row, 'outcome') ||
|
||||
usage.settledAtMs !== integer(row, 'settledAtMs') ||
|
||||
usage.inputBytes !== integer(row, 'inputBytes') ||
|
||||
usage.outputBytes !== integer(row, 'outputBytes') ||
|
||||
usage.inputTokens !== integer(row, 'inputTokens') ||
|
||||
usage.outputTokens !== integer(row, 'outputTokens') ||
|
||||
usage.totalTokens !== integer(row, 'totalTokens') ||
|
||||
usage.costMicros !== nullableInteger(row, 'costMicros') ||
|
||||
usage.ledgerDigest !== text(row, 'ledgerDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return usage;
|
||||
}
|
||||
|
||||
export function parseQuotaReservation(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationQuotaReservation> {
|
||||
try {
|
||||
return normalizeModelInvocationQuotaReservation(
|
||||
jsonObject(
|
||||
row,
|
||||
'recordJson',
|
||||
) as unknown as ModelInvocationQuotaReservation,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
}
|
||||
|
||||
export function parseQuotaSettlement(
|
||||
row: Row,
|
||||
reservation: Readonly<ModelInvocationQuotaReservation>,
|
||||
completion: Readonly<ModelInvocationCompletionRecord>,
|
||||
): Readonly<ModelInvocationQuotaSettlement> {
|
||||
try {
|
||||
return normalizeModelInvocationQuotaSettlement(
|
||||
jsonObject(
|
||||
row,
|
||||
'recordJson',
|
||||
) as unknown as ModelInvocationQuotaSettlement,
|
||||
reservation,
|
||||
completion,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
}
|
||||
|
||||
export function parsePriceQuote(row: Row): Readonly<ModelInvocationPriceQuote> {
|
||||
try {
|
||||
return normalizeModelInvocationPriceQuote(
|
||||
jsonObject(row, 'recordJson') as unknown as ModelInvocationPriceQuote,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
}
|
||||
|
||||
export function parsePriceSettlement(
|
||||
row: Row,
|
||||
quote: Readonly<ModelInvocationPriceQuote>,
|
||||
completion: Readonly<ModelInvocationCompletionRecord>,
|
||||
): Readonly<ModelInvocationPriceSettlement> {
|
||||
try {
|
||||
return normalizeModelInvocationPriceSettlement(
|
||||
jsonObject(
|
||||
row,
|
||||
'recordJson',
|
||||
) as unknown as ModelInvocationPriceSettlement,
|
||||
quote,
|
||||
completion,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
}
|
||||
|
||||
export function parseResolution(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationResolutionRecord> {
|
||||
let resolution: Readonly<ModelInvocationResolutionRecord>;
|
||||
try {
|
||||
resolution = normalizeModelInvocationResolutionRecord(
|
||||
jsonObject(
|
||||
row,
|
||||
'recordJson',
|
||||
) as unknown as ModelInvocationResolutionRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
resolution.resolutionId !== text(row, 'resolutionId') ||
|
||||
resolution.invocationId !== text(row, 'invocationId') ||
|
||||
resolution.projectId !== text(row, 'projectId') ||
|
||||
resolution.runId !== text(row, 'runId') ||
|
||||
resolution.stepRunId !== text(row, 'stepRunId') ||
|
||||
resolution.traceId !== text(row, 'traceId') ||
|
||||
resolution.completionDigest !== text(row, 'completionDigest') ||
|
||||
resolution.decision !== text(row, 'decision') ||
|
||||
resolution.resolvedByUserId !== text(row, 'resolvedByUserId') ||
|
||||
resolution.resolvedAtMs !== integer(row, 'resolvedAtMs') ||
|
||||
resolution.stepRunMutationId !== text(row, 'mutationId') ||
|
||||
resolution.stepRunMutationDigest !== text(row, 'mutationDigest') ||
|
||||
resolution.resolvedStepRunDigest !== text(row, 'stepRunDigest') ||
|
||||
resolution.resolvedStepRunVersion !== integer(row, 'stepRunVersion') ||
|
||||
resolution.runEventId !== text(row, 'runEventId') ||
|
||||
resolution.resolutionDigest !== text(row, 'resolutionDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return resolution;
|
||||
}
|
||||
|
||||
export function parseAuthority(
|
||||
row: Row,
|
||||
): Readonly<ModelInvocationAuthoritySnapshot> {
|
||||
let stepRun: Readonly<StepRunRecord>;
|
||||
try {
|
||||
stepRun = normalizeStepRunRecord(
|
||||
jsonObject(row, 'stepRunJson') as unknown as StepRunRecord,
|
||||
);
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
if (
|
||||
stepRun.id !== text(row, 'stepRunId') ||
|
||||
stepRun.runId !== text(row, 'runId') ||
|
||||
stepRun.kind !== 'model' ||
|
||||
stepRun.status !== text(row, 'stepStatus') ||
|
||||
stepRun.version !== integer(row, 'stepVersion') ||
|
||||
stepRun.stepRunDigest !== text(row, 'stepDigest')
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
return Object.freeze({
|
||||
projectId: text(row, 'projectId'),
|
||||
runId: text(row, 'runId'),
|
||||
runVersion: integer(row, 'runVersion'),
|
||||
runEventSequence: integer(row, 'runEventSequence'),
|
||||
stepRun,
|
||||
});
|
||||
}
|
||||
|
||||
export const START_SELECT = `
|
||||
start.invocation_id AS "invocationId",
|
||||
start.project_id AS "projectId",
|
||||
start.run_id AS "runId",
|
||||
start.step_run_id AS "stepRunId",
|
||||
start.trace_id AS "traceId",
|
||||
start.provider,
|
||||
start.model,
|
||||
start.policy_revision AS "policyRevision",
|
||||
start.request_digest AS "requestDigest",
|
||||
start.input_bytes AS "inputBytes",
|
||||
start.max_output_tokens AS "maxOutputTokens",
|
||||
start.deadline_at_ms AS "deadlineAtMs",
|
||||
start.admitted_at_ms AS "admittedAtMs",
|
||||
start.mutation_id AS "mutationId",
|
||||
start.mutation_digest AS "mutationDigest",
|
||||
start.run_event_id AS "runEventId",
|
||||
start.start_digest AS "startDigest",
|
||||
start.record_json AS "recordJson",
|
||||
mutation.step_run_digest AS "stepRunDigest",
|
||||
mutation.step_run_json->>'version' AS "stepRunVersion"
|
||||
`;
|
||||
|
||||
export const COMPLETION_SELECT = `
|
||||
completion.invocation_id AS "invocationId",
|
||||
completion.project_id AS "projectId",
|
||||
completion.run_id AS "runId",
|
||||
completion.step_run_id AS "stepRunId",
|
||||
completion.trace_id AS "traceId",
|
||||
completion.start_digest AS "startDigest",
|
||||
completion.outcome,
|
||||
completion.output_bytes AS "outputBytes",
|
||||
completion.error_code AS "errorCode",
|
||||
completion.completed_at_ms AS "completedAtMs",
|
||||
completion.mutation_id AS "mutationId",
|
||||
completion.mutation_digest AS "mutationDigest",
|
||||
completion.run_event_id AS "runEventId",
|
||||
completion.completion_digest AS "completionDigest",
|
||||
completion.record_json AS "recordJson",
|
||||
mutation.step_run_digest AS "stepRunDigest",
|
||||
mutation.step_run_json->>'version' AS "stepRunVersion"
|
||||
`;
|
||||
|
||||
export const USAGE_SELECT = `
|
||||
usage.invocation_id AS "invocationId",
|
||||
usage.project_id AS "projectId",
|
||||
usage.run_id AS "runId",
|
||||
usage.step_run_id AS "stepRunId",
|
||||
usage.trace_id AS "traceId",
|
||||
usage.provider,
|
||||
usage.model,
|
||||
usage.policy_revision AS "policyRevision",
|
||||
usage.completion_digest AS "completionDigest",
|
||||
usage.outcome,
|
||||
usage.settled_at_ms AS "settledAtMs",
|
||||
usage.input_bytes AS "inputBytes",
|
||||
usage.output_bytes AS "outputBytes",
|
||||
usage.input_tokens AS "inputTokens",
|
||||
usage.output_tokens AS "outputTokens",
|
||||
usage.total_tokens AS "totalTokens",
|
||||
usage.cost_micros AS "costMicros",
|
||||
usage.ledger_digest AS "ledgerDigest",
|
||||
usage.record_json AS "recordJson"
|
||||
`;
|
||||
|
||||
export const RESOLUTION_SELECT = `
|
||||
resolution.resolution_id AS "resolutionId",
|
||||
resolution.invocation_id AS "invocationId",
|
||||
resolution.project_id AS "projectId",
|
||||
resolution.run_id AS "runId",
|
||||
resolution.step_run_id AS "stepRunId",
|
||||
resolution.trace_id AS "traceId",
|
||||
resolution.completion_digest AS "completionDigest",
|
||||
resolution.decision,
|
||||
resolution.resolved_by_user_id AS "resolvedByUserId",
|
||||
resolution.resolved_at_ms AS "resolvedAtMs",
|
||||
resolution.mutation_id AS "mutationId",
|
||||
resolution.mutation_digest AS "mutationDigest",
|
||||
resolution.run_event_id AS "runEventId",
|
||||
resolution.resolution_digest AS "resolutionDigest",
|
||||
resolution.record_json AS "recordJson",
|
||||
mutation.step_run_digest AS "stepRunDigest",
|
||||
mutation.step_run_json->>'version' AS "stepRunVersion"
|
||||
`;
|
||||
+459
@@ -0,0 +1,459 @@
|
||||
import type { PostgresPool } from '@qinglong/runtime-core';
|
||||
|
||||
import { createModelInvocationPriceSettlement } from '../../pricing/pricing';
|
||||
import { type PluginPackagePromptOutputArtifact } from '../../prompt-output/pluginPackagePromptOutputArtifact';
|
||||
import {
|
||||
assertPluginPackagePromptOutputCompletionBinding,
|
||||
type CommitPluginPackagePromptOutputResult,
|
||||
} from '../../prompt-output/pluginPackagePromptOutputCompletion';
|
||||
import {
|
||||
putPostgresPluginPackagePromptOutputArtifactInTransaction,
|
||||
readPostgresPluginPackagePromptOutputArtifactInTransaction,
|
||||
} from '../../prompt-output/storage/postgresPluginPackagePromptOutputArtifactRepository';
|
||||
import { createModelInvocationUsageLedgerRecord } from '../../usage/usageLedger';
|
||||
import { createModelInvocationQuotaSettlement } from '../../usage/usageQuota';
|
||||
import {
|
||||
ModelInvocationConflictError,
|
||||
normalizeModelInvocationCompletionCommand,
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationCompletionCommand,
|
||||
type ModelInvocationCompletionRecord,
|
||||
} from '../modelInvocation';
|
||||
|
||||
import {
|
||||
parseCompletion,
|
||||
parsePriceQuote,
|
||||
parsePriceSettlement,
|
||||
parseQuotaReservation,
|
||||
parseQuotaSettlement,
|
||||
parseStart,
|
||||
parseUsage,
|
||||
} from './codec';
|
||||
import {
|
||||
applyMutation,
|
||||
assertCurrent,
|
||||
insertCompletion,
|
||||
insertPriceSettlement,
|
||||
insertQuotaSettlement,
|
||||
insertUsage,
|
||||
} from './mutations';
|
||||
import {
|
||||
completionRows,
|
||||
priceQuoteRows,
|
||||
priceSettlementRows,
|
||||
quotaReservationRows,
|
||||
quotaSettlementRows,
|
||||
startRows,
|
||||
usageRows,
|
||||
} from './queries';
|
||||
import { runPostgresModelInvocationTransaction } from './transaction';
|
||||
|
||||
export async function completeOperation(
|
||||
pool: PostgresPool,
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
const command = normalizeModelInvocationCompletionCommand(commandValue);
|
||||
const completion = command.completion;
|
||||
const expectedUsage = createModelInvocationUsageLedgerRecord(
|
||||
command.start,
|
||||
completion,
|
||||
);
|
||||
return runPostgresModelInvocationTransaction(pool, async (client) => {
|
||||
const existing = await completionRows(
|
||||
client,
|
||||
`completion.invocation_id = $1 OR
|
||||
completion.mutation_id = $2 OR completion.run_event_id = $3`,
|
||||
[
|
||||
completion.invocationId,
|
||||
completion.stepRunMutationId,
|
||||
completion.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseCompletion(existing[0]);
|
||||
if (JSON.stringify(stored) !== JSON.stringify(completion)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const usage = await usageRows(client, 'usage.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
usage.length !== (expectedUsage ? 1 : 0) ||
|
||||
(expectedUsage &&
|
||||
JSON.stringify(parseUsage(usage[0]!)) !==
|
||||
JSON.stringify(expectedUsage))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
const starts = await startRows(client, 'start.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
starts.length !== 1 ||
|
||||
JSON.stringify(parseStart(starts[0]!)) !== JSON.stringify(command.start)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
await assertCurrent(client, command.stepRunMutation, completion.projectId);
|
||||
await applyMutation(client, command.stepRunMutation);
|
||||
await insertCompletion(client, completion);
|
||||
if (expectedUsage) await insertUsage(client, expectedUsage);
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: completion,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
export async function completeWithQuotaOperation(
|
||||
pool: PostgresPool,
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
const command = normalizeModelInvocationCompletionCommand(commandValue);
|
||||
const completion = command.completion;
|
||||
const expectedUsage = createModelInvocationUsageLedgerRecord(
|
||||
command.start,
|
||||
completion,
|
||||
);
|
||||
return runPostgresModelInvocationTransaction(pool, async (client) => {
|
||||
const reservationRows = await quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = $1',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (reservationRows.length !== 1) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const reservation = parseQuotaReservation(reservationRows[0]!);
|
||||
const expectedSettlement = createModelInvocationQuotaSettlement(
|
||||
reservation,
|
||||
completion,
|
||||
);
|
||||
const existing = await completionRows(
|
||||
client,
|
||||
`completion.invocation_id = $1 OR
|
||||
completion.mutation_id = $2 OR completion.run_event_id = $3`,
|
||||
[
|
||||
completion.invocationId,
|
||||
completion.stepRunMutationId,
|
||||
completion.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseCompletion(existing[0]);
|
||||
const usage = await usageRows(client, 'usage.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
const settlements = await quotaSettlementRows(
|
||||
client,
|
||||
'settlement.invocation_id = $1',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(completion) ||
|
||||
usage.length !== (expectedUsage ? 1 : 0) ||
|
||||
(expectedUsage &&
|
||||
JSON.stringify(parseUsage(usage[0]!)) !==
|
||||
JSON.stringify(expectedUsage)) ||
|
||||
settlements.length !== 1 ||
|
||||
JSON.stringify(
|
||||
parseQuotaSettlement(settlements[0]!, reservation, completion),
|
||||
) !== JSON.stringify(expectedSettlement)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
const starts = await startRows(client, 'start.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
starts.length !== 1 ||
|
||||
JSON.stringify(parseStart(starts[0]!)) !== JSON.stringify(command.start)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
await assertCurrent(client, command.stepRunMutation, completion.projectId);
|
||||
await applyMutation(client, command.stepRunMutation);
|
||||
await insertCompletion(client, completion);
|
||||
if (expectedUsage) await insertUsage(client, expectedUsage);
|
||||
await insertQuotaSettlement(client, expectedSettlement);
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: completion,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
export async function completeWithPricingOperation(
|
||||
pool: PostgresPool,
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
const command = normalizeModelInvocationCompletionCommand(commandValue);
|
||||
const completion = command.completion;
|
||||
const expectedUsage = createModelInvocationUsageLedgerRecord(
|
||||
command.start,
|
||||
completion,
|
||||
);
|
||||
return runPostgresModelInvocationTransaction(pool, async (client) => {
|
||||
const quoteRows = await priceQuoteRows(client, 'quote.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (quoteRows.length !== 1) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const quote = parsePriceQuote(quoteRows[0]!);
|
||||
if (
|
||||
quote.invocationId !== command.start.invocationId ||
|
||||
quote.projectId !== command.start.projectId ||
|
||||
quote.modelPolicyRevision !== command.start.policyRevision ||
|
||||
quote.provider !== command.start.provider ||
|
||||
quote.model !== command.start.model
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const expectedPriceSettlement = createModelInvocationPriceSettlement(
|
||||
quote,
|
||||
completion,
|
||||
);
|
||||
const reservationRows = await quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = $1',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (reservationRows.length > 1) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const reservation = reservationRows[0]
|
||||
? parseQuotaReservation(reservationRows[0])
|
||||
: null;
|
||||
const expectedQuotaSettlement = reservation
|
||||
? createModelInvocationQuotaSettlement(reservation, completion)
|
||||
: null;
|
||||
const existing = await completionRows(
|
||||
client,
|
||||
`completion.invocation_id = $1 OR
|
||||
completion.mutation_id = $2 OR completion.run_event_id = $3`,
|
||||
[
|
||||
completion.invocationId,
|
||||
completion.stepRunMutationId,
|
||||
completion.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const [usage, priceSettlements, quotaSettlements] = await Promise.all([
|
||||
usageRows(client, 'usage.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]),
|
||||
priceSettlementRows(client, 'settlement.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]),
|
||||
quotaSettlementRows(client, 'settlement.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]),
|
||||
]);
|
||||
const stored = parseCompletion(existing[0]);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(completion) ||
|
||||
usage.length !== (expectedUsage ? 1 : 0) ||
|
||||
(expectedUsage &&
|
||||
JSON.stringify(parseUsage(usage[0]!)) !==
|
||||
JSON.stringify(expectedUsage)) ||
|
||||
priceSettlements.length !== (expectedPriceSettlement ? 1 : 0) ||
|
||||
(expectedPriceSettlement &&
|
||||
JSON.stringify(
|
||||
parsePriceSettlement(priceSettlements[0]!, quote, completion),
|
||||
) !== JSON.stringify(expectedPriceSettlement)) ||
|
||||
quotaSettlements.length !== (expectedQuotaSettlement ? 1 : 0) ||
|
||||
(expectedQuotaSettlement &&
|
||||
JSON.stringify(
|
||||
parseQuotaSettlement(
|
||||
quotaSettlements[0]!,
|
||||
reservation!,
|
||||
completion,
|
||||
),
|
||||
) !== JSON.stringify(expectedQuotaSettlement))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
const starts = await startRows(client, 'start.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
starts.length !== 1 ||
|
||||
JSON.stringify(parseStart(starts[0]!)) !== JSON.stringify(command.start)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
await assertCurrent(client, command.stepRunMutation, completion.projectId);
|
||||
await applyMutation(client, command.stepRunMutation);
|
||||
await insertCompletion(client, completion);
|
||||
if (expectedUsage) await insertUsage(client, expectedUsage);
|
||||
if (expectedPriceSettlement) {
|
||||
await insertPriceSettlement(client, expectedPriceSettlement);
|
||||
}
|
||||
if (expectedQuotaSettlement) {
|
||||
await insertQuotaSettlement(client, expectedQuotaSettlement);
|
||||
}
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: completion,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
export async function completeWithPromptOutputArtifactOperation(
|
||||
pool: PostgresPool,
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
artifactValue: PluginPackagePromptOutputArtifact,
|
||||
): Promise<Readonly<CommitPluginPackagePromptOutputResult>> {
|
||||
const command = normalizeModelInvocationCompletionCommand(commandValue);
|
||||
const completion = command.completion;
|
||||
const binding = assertPluginPackagePromptOutputCompletionBinding(
|
||||
command,
|
||||
artifactValue,
|
||||
);
|
||||
const expectedUsage = createModelInvocationUsageLedgerRecord(
|
||||
command.start,
|
||||
completion,
|
||||
);
|
||||
return runPostgresModelInvocationTransaction(pool, async (client) => {
|
||||
const quoteRows = await priceQuoteRows(client, 'quote.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (quoteRows.length > 1) throw new ModelInvocationConflictError();
|
||||
const quote = quoteRows[0] ? parsePriceQuote(quoteRows[0]) : null;
|
||||
if (
|
||||
quote &&
|
||||
(quote.invocationId !== command.start.invocationId ||
|
||||
quote.projectId !== command.start.projectId ||
|
||||
quote.modelPolicyRevision !== command.start.policyRevision ||
|
||||
quote.provider !== command.start.provider ||
|
||||
quote.model !== command.start.model)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const expectedPriceSettlement = quote
|
||||
? createModelInvocationPriceSettlement(quote, completion)
|
||||
: null;
|
||||
const reservationRows = await quotaReservationRows(
|
||||
client,
|
||||
'reservation.invocation_id = $1',
|
||||
[completion.invocationId],
|
||||
);
|
||||
if (reservationRows.length > 1) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
const reservation = reservationRows[0]
|
||||
? parseQuotaReservation(reservationRows[0])
|
||||
: null;
|
||||
const expectedQuotaSettlement = reservation
|
||||
? createModelInvocationQuotaSettlement(reservation, completion)
|
||||
: null;
|
||||
const existing = await completionRows(
|
||||
client,
|
||||
`completion.invocation_id = $1 OR
|
||||
completion.mutation_id = $2 OR completion.run_event_id = $3`,
|
||||
[
|
||||
completion.invocationId,
|
||||
completion.stepRunMutationId,
|
||||
completion.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const [storedArtifact, usage, priceSettlements, quotaSettlements] =
|
||||
await Promise.all([
|
||||
readPostgresPluginPackagePromptOutputArtifactInTransaction(
|
||||
client,
|
||||
binding.artifact.artifactId,
|
||||
),
|
||||
usageRows(client, 'usage.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]),
|
||||
priceSettlementRows(client, 'settlement.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]),
|
||||
quotaSettlementRows(client, 'settlement.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]),
|
||||
]);
|
||||
const stored = parseCompletion(existing[0]);
|
||||
if (
|
||||
JSON.stringify(stored) !== JSON.stringify(completion) ||
|
||||
!storedArtifact ||
|
||||
JSON.stringify(storedArtifact) !== JSON.stringify(binding.artifact) ||
|
||||
usage.length !== (expectedUsage ? 1 : 0) ||
|
||||
(expectedUsage &&
|
||||
JSON.stringify(parseUsage(usage[0]!)) !==
|
||||
JSON.stringify(expectedUsage)) ||
|
||||
priceSettlements.length !== (expectedPriceSettlement ? 1 : 0) ||
|
||||
(expectedPriceSettlement &&
|
||||
JSON.stringify(
|
||||
parsePriceSettlement(priceSettlements[0]!, quote!, completion),
|
||||
) !== JSON.stringify(expectedPriceSettlement)) ||
|
||||
quotaSettlements.length !== (expectedQuotaSettlement ? 1 : 0) ||
|
||||
(expectedQuotaSettlement &&
|
||||
JSON.stringify(
|
||||
parseQuotaSettlement(
|
||||
quotaSettlements[0]!,
|
||||
reservation!,
|
||||
completion,
|
||||
),
|
||||
) !== JSON.stringify(expectedQuotaSettlement))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return Object.freeze({
|
||||
status: 'existing' as const,
|
||||
record: stored,
|
||||
artifact: storedArtifact,
|
||||
reference: binding.reference,
|
||||
});
|
||||
}
|
||||
const starts = await startRows(client, 'start.invocation_id = $1', [
|
||||
completion.invocationId,
|
||||
]);
|
||||
if (
|
||||
starts.length !== 1 ||
|
||||
JSON.stringify(parseStart(starts[0]!)) !== JSON.stringify(command.start)
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
await assertCurrent(client, command.stepRunMutation, completion.projectId);
|
||||
const artifact = (
|
||||
await putPostgresPluginPackagePromptOutputArtifactInTransaction(
|
||||
client,
|
||||
binding.artifact,
|
||||
)
|
||||
).artifact;
|
||||
await applyMutation(client, command.stepRunMutation);
|
||||
await insertCompletion(client, completion);
|
||||
if (expectedUsage) await insertUsage(client, expectedUsage);
|
||||
if (expectedPriceSettlement) {
|
||||
await insertPriceSettlement(client, expectedPriceSettlement);
|
||||
}
|
||||
if (expectedQuotaSettlement) {
|
||||
await insertQuotaSettlement(client, expectedQuotaSettlement);
|
||||
}
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: completion,
|
||||
artifact,
|
||||
reference: binding.reference,
|
||||
});
|
||||
});
|
||||
}
|
||||
+490
@@ -0,0 +1,490 @@
|
||||
import type { PostgresClient } from '@qinglong/runtime-core';
|
||||
import { type StepRunMutation } from '@qinglong/runtime-core/step-run';
|
||||
|
||||
import { POSTGRES_MODEL_INVOCATION_SCHEMA } from '../../migration/modelInvocationMigration';
|
||||
import {
|
||||
type ModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceSettlement,
|
||||
} from '../../pricing/pricing';
|
||||
import { type ModelInvocationUsageLedgerRecord } from '../../usage/usageLedger';
|
||||
import {
|
||||
type ModelInvocationQuotaReservation,
|
||||
type ModelInvocationQuotaSettlement,
|
||||
type ModelInvocationQuotaWindowUsage,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
ModelInvocationConflictError,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
import { type ModelInvocationResolutionRecord } from '../modelInvocationResolution';
|
||||
|
||||
import type { Queryable, Row } from './authority';
|
||||
import { TERMINAL_RUN_STATUSES, integer, text, unavailable } from './authority';
|
||||
|
||||
export async function updateStepRun(
|
||||
client: PostgresClient,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): Promise<void> {
|
||||
const step = mutation.stepRun;
|
||||
const result = await client.query(
|
||||
`UPDATE "ql3"."step_runs"
|
||||
SET status = $1, version = $2, attempt_count = $3, output_ref = $4,
|
||||
approval_request_id = $5, ready_at_ms = $6, started_at_ms = $7,
|
||||
finished_at_ms = $8, result_code = $9, error_summary = $10,
|
||||
updated_at_ms = $11, last_mutation_id = $12,
|
||||
step_run_digest = $13, step_run_json = $14::jsonb
|
||||
WHERE id = $15 AND run_id = $16 AND version = $17
|
||||
AND step_run_digest = $18 AND status = $19`,
|
||||
[
|
||||
step.status,
|
||||
step.version,
|
||||
step.attemptCount,
|
||||
step.outputRef,
|
||||
step.approvalRequestId,
|
||||
step.readyAtMs,
|
||||
step.startedAtMs,
|
||||
step.finishedAtMs,
|
||||
step.resultCode,
|
||||
step.errorSummary,
|
||||
step.updatedAtMs,
|
||||
step.lastMutationId,
|
||||
step.stepRunDigest,
|
||||
JSON.stringify(step),
|
||||
step.id,
|
||||
step.runId,
|
||||
mutation.expectedStepRunVersion,
|
||||
mutation.expectedStepRunDigest,
|
||||
mutation.previousStatus,
|
||||
],
|
||||
);
|
||||
if (result.rowCount !== 1) throw new ModelInvocationConflictError();
|
||||
}
|
||||
|
||||
export async function updateRun(
|
||||
client: PostgresClient,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): Promise<void> {
|
||||
const result = await client.query(
|
||||
`UPDATE "ql3"."runs"
|
||||
SET version = version + 1, event_sequence = event_sequence + 1
|
||||
WHERE id = $1 AND version = $2 AND event_sequence = $3`,
|
||||
[
|
||||
mutation.runId,
|
||||
mutation.expectedRunVersion,
|
||||
mutation.expectedRunEventSequence,
|
||||
],
|
||||
);
|
||||
if (result.rowCount !== 1) throw new ModelInvocationConflictError();
|
||||
}
|
||||
|
||||
export async function insertRunEvent(
|
||||
client: PostgresClient,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): Promise<void> {
|
||||
const event = mutation.event;
|
||||
await client.query(
|
||||
`INSERT INTO "ql3"."run_events" (
|
||||
id, run_id, sequence, type, dedupe_key, actor_type, actor_id,
|
||||
attempt_id, step_run_id, payload, created_at_ms
|
||||
) VALUES ($1, $2, $3, $4, $5, $6, $7, NULL, $8, $9::jsonb, $10)`,
|
||||
[
|
||||
event.id,
|
||||
event.runId,
|
||||
event.sequence,
|
||||
event.type,
|
||||
event.dedupeKey,
|
||||
event.actorType,
|
||||
event.actorId ?? null,
|
||||
mutation.stepRun.id,
|
||||
JSON.stringify(event.payload),
|
||||
event.createdAtMs,
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function insertMutation(
|
||||
client: PostgresClient,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): Promise<void> {
|
||||
await client.query(
|
||||
`INSERT INTO "ql3"."step_run_mutations" (
|
||||
mutation_id, mutation_digest, run_id, step_run_id,
|
||||
step_run_digest, event_id, event_sequence, run_version,
|
||||
step_run_json, committed_at_ms
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9::jsonb,
|
||||
floor(extract(epoch FROM transaction_timestamp()) * 1000)::bigint
|
||||
)`,
|
||||
[
|
||||
mutation.mutationId,
|
||||
mutation.mutationDigest,
|
||||
mutation.runId,
|
||||
mutation.stepRun.id,
|
||||
mutation.stepRun.stepRunDigest,
|
||||
mutation.event.id,
|
||||
mutation.event.sequence,
|
||||
mutation.expectedRunVersion + 1,
|
||||
JSON.stringify(mutation.stepRun),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function applyMutation(
|
||||
client: PostgresClient,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
): Promise<void> {
|
||||
await updateStepRun(client, mutation);
|
||||
await updateRun(client, mutation);
|
||||
await insertRunEvent(client, mutation);
|
||||
await insertMutation(client, mutation);
|
||||
}
|
||||
|
||||
export async function insertStart(
|
||||
client: PostgresClient,
|
||||
start: Readonly<ModelInvocationStartRecord>,
|
||||
): Promise<void> {
|
||||
await client.query(
|
||||
`INSERT INTO "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_starts" (
|
||||
invocation_id, project_id, run_id, step_run_id, trace_id,
|
||||
provider, model, policy_revision, request_digest, input_bytes,
|
||||
max_output_tokens, deadline_at_ms, admitted_at_ms, mutation_id,
|
||||
mutation_digest, run_event_id, start_digest, record_json
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
|
||||
$14, $15, $16, $17, $18::jsonb
|
||||
)`,
|
||||
[
|
||||
start.invocationId,
|
||||
start.projectId,
|
||||
start.runId,
|
||||
start.stepRunId,
|
||||
start.traceId,
|
||||
start.provider,
|
||||
start.model,
|
||||
start.policyRevision,
|
||||
start.requestDigest,
|
||||
start.inputBytes,
|
||||
start.maxOutputTokens,
|
||||
start.deadlineAtMs,
|
||||
start.admittedAtMs,
|
||||
start.stepRunMutationId,
|
||||
start.stepRunMutationDigest,
|
||||
start.runEventId,
|
||||
start.startDigest,
|
||||
JSON.stringify(start),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function insertCompletion(
|
||||
client: PostgresClient,
|
||||
completion: Readonly<ModelInvocationCompletionRecord>,
|
||||
): Promise<void> {
|
||||
await client.query(
|
||||
`INSERT INTO "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_completions" (
|
||||
invocation_id, project_id, run_id, step_run_id, trace_id,
|
||||
start_digest, outcome, output_bytes, error_code, completed_at_ms,
|
||||
mutation_id, mutation_digest, run_event_id, completion_digest,
|
||||
record_json
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
|
||||
$14, $15::jsonb
|
||||
)`,
|
||||
[
|
||||
completion.invocationId,
|
||||
completion.projectId,
|
||||
completion.runId,
|
||||
completion.stepRunId,
|
||||
completion.traceId,
|
||||
completion.startDigest,
|
||||
completion.outcome,
|
||||
completion.outputBytes,
|
||||
completion.errorCode,
|
||||
completion.completedAtMs,
|
||||
completion.stepRunMutationId,
|
||||
completion.stepRunMutationDigest,
|
||||
completion.runEventId,
|
||||
completion.completionDigest,
|
||||
JSON.stringify(completion),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function insertUsage(
|
||||
client: PostgresClient,
|
||||
usage: Readonly<ModelInvocationUsageLedgerRecord>,
|
||||
): Promise<void> {
|
||||
await client.query(
|
||||
`INSERT INTO "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_usage_ledger" (
|
||||
invocation_id, project_id, run_id, step_run_id, trace_id,
|
||||
provider, model, policy_revision, completion_digest, outcome,
|
||||
settled_at_ms, input_bytes, output_bytes, input_tokens,
|
||||
output_tokens, total_tokens, cost_micros, ledger_digest, record_json
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
|
||||
$14, $15, $16, $17, $18, $19::jsonb
|
||||
)`,
|
||||
[
|
||||
usage.invocationId,
|
||||
usage.projectId,
|
||||
usage.runId,
|
||||
usage.stepRunId,
|
||||
usage.traceId,
|
||||
usage.provider,
|
||||
usage.model,
|
||||
usage.policyRevision,
|
||||
usage.completionDigest,
|
||||
usage.outcome,
|
||||
usage.settledAtMs,
|
||||
usage.inputBytes,
|
||||
usage.outputBytes,
|
||||
usage.inputTokens,
|
||||
usage.outputTokens,
|
||||
usage.totalTokens,
|
||||
usage.costMicros,
|
||||
usage.ledgerDigest,
|
||||
JSON.stringify(usage),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function insertQuotaReservation(
|
||||
client: PostgresClient,
|
||||
reservation: Readonly<ModelInvocationQuotaReservation>,
|
||||
): Promise<void> {
|
||||
await client.query(
|
||||
`INSERT INTO "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_quota_reservations" (
|
||||
invocation_id, project_id, model_policy_revision,
|
||||
quota_policy_revision, window_ms, window_start_ms, window_end_ms,
|
||||
max_invocations, max_tokens, max_cost_micros, reserved_tokens,
|
||||
reserved_cost_micros, reserved_at_ms, admission_digest,
|
||||
reservation_digest, record_json
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
|
||||
$14, $15, $16::jsonb
|
||||
)`,
|
||||
[
|
||||
reservation.invocationId,
|
||||
reservation.projectId,
|
||||
reservation.modelPolicyRevision,
|
||||
reservation.quotaPolicyRevision,
|
||||
reservation.windowMs,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowEndMs,
|
||||
reservation.maxInvocations,
|
||||
reservation.maxTokens,
|
||||
reservation.maxCostMicros,
|
||||
reservation.reservedTokens,
|
||||
reservation.reservedCostMicros,
|
||||
reservation.reservedAtMs,
|
||||
reservation.admissionDigest,
|
||||
reservation.reservationDigest,
|
||||
JSON.stringify(reservation),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function insertQuotaSettlement(
|
||||
client: PostgresClient,
|
||||
settlement: Readonly<ModelInvocationQuotaSettlement>,
|
||||
): Promise<void> {
|
||||
await client.query(
|
||||
`INSERT INTO "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_quota_settlements" (
|
||||
invocation_id, project_id, reservation_digest, completion_digest,
|
||||
effective_tokens, effective_cost_micros,
|
||||
retained_token_reservation, retained_cost_reservation,
|
||||
settled_at_ms, settlement_digest, record_json
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::jsonb
|
||||
)`,
|
||||
[
|
||||
settlement.invocationId,
|
||||
settlement.projectId,
|
||||
settlement.reservationDigest,
|
||||
settlement.completionDigest,
|
||||
settlement.effectiveTokens,
|
||||
settlement.effectiveCostMicros,
|
||||
settlement.retainedTokenReservation,
|
||||
settlement.retainedCostReservation,
|
||||
settlement.settledAtMs,
|
||||
settlement.settlementDigest,
|
||||
JSON.stringify(settlement),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function insertPriceQuote(
|
||||
client: PostgresClient,
|
||||
quote: Readonly<ModelInvocationPriceQuote>,
|
||||
): Promise<void> {
|
||||
await client.query(
|
||||
`INSERT INTO "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_price_quotes" (
|
||||
invocation_id, project_id, model_policy_revision, provider, model,
|
||||
price_revision, currency, input_micros_per_million_tokens,
|
||||
output_micros_per_million_tokens, max_total_tokens, max_output_tokens,
|
||||
reserved_cost_micros, catalog_digest, quote_digest, record_json
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14,
|
||||
$15::jsonb
|
||||
)`,
|
||||
[
|
||||
quote.invocationId,
|
||||
quote.projectId,
|
||||
quote.modelPolicyRevision,
|
||||
quote.provider,
|
||||
quote.model,
|
||||
quote.priceRevision,
|
||||
quote.currency,
|
||||
quote.inputMicrosPerMillionTokens,
|
||||
quote.outputMicrosPerMillionTokens,
|
||||
quote.maxTotalTokens,
|
||||
quote.maxOutputTokens,
|
||||
quote.reservedCostMicros,
|
||||
quote.catalogDigest,
|
||||
quote.quoteDigest,
|
||||
JSON.stringify(quote),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function insertPriceSettlement(
|
||||
client: PostgresClient,
|
||||
settlement: Readonly<ModelInvocationPriceSettlement>,
|
||||
): Promise<void> {
|
||||
await client.query(
|
||||
`INSERT INTO "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_price_settlements" (
|
||||
invocation_id, project_id, quote_digest, completion_digest, currency,
|
||||
input_tokens, output_tokens, cost_micros, settled_at_ms,
|
||||
settlement_digest, record_json
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::jsonb
|
||||
)`,
|
||||
[
|
||||
settlement.invocationId,
|
||||
settlement.projectId,
|
||||
settlement.quoteDigest,
|
||||
settlement.completionDigest,
|
||||
settlement.currency,
|
||||
settlement.inputTokens,
|
||||
settlement.outputTokens,
|
||||
settlement.costMicros,
|
||||
settlement.settledAtMs,
|
||||
settlement.settlementDigest,
|
||||
JSON.stringify(settlement),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function quotaWindowUsage(
|
||||
queryable: Queryable,
|
||||
projectId: string,
|
||||
windowStartMs: number,
|
||||
windowMs: number,
|
||||
): Promise<Readonly<ModelInvocationQuotaWindowUsage>> {
|
||||
const result = await queryable.query<Row>(
|
||||
`SELECT
|
||||
COUNT(*)::text AS "invocationCount",
|
||||
COALESCE(SUM(COALESCE(
|
||||
settlement.effective_tokens, reservation.reserved_tokens
|
||||
)), 0)::text AS "effectiveTokens",
|
||||
COALESCE(SUM(COALESCE(
|
||||
settlement.effective_cost_micros,
|
||||
reservation.reserved_cost_micros,
|
||||
0
|
||||
)), 0)::text AS "effectiveCostMicros",
|
||||
COUNT(*) FILTER (
|
||||
WHERE settlement.effective_cost_micros IS NULL
|
||||
AND reservation.reserved_cost_micros IS NULL
|
||||
)::text AS "unknownCostInvocations"
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_quota_reservations"
|
||||
AS reservation
|
||||
LEFT JOIN "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_quota_settlements"
|
||||
AS settlement ON settlement.invocation_id = reservation.invocation_id
|
||||
WHERE reservation.project_id = $1
|
||||
AND reservation.window_start_ms = $2
|
||||
AND reservation.window_ms = $3`,
|
||||
[projectId, windowStartMs, windowMs],
|
||||
);
|
||||
const row = result.rows[0];
|
||||
if (result.rows.length !== 1 || !row) throw unavailable();
|
||||
return Object.freeze({
|
||||
projectId,
|
||||
windowStartMs,
|
||||
windowEndMs: windowStartMs + windowMs,
|
||||
invocationCount: integer(row, 'invocationCount'),
|
||||
effectiveTokens: integer(row, 'effectiveTokens'),
|
||||
effectiveCostMicros: integer(row, 'effectiveCostMicros'),
|
||||
unknownCostInvocations: integer(row, 'unknownCostInvocations'),
|
||||
});
|
||||
}
|
||||
|
||||
export async function insertResolution(
|
||||
client: PostgresClient,
|
||||
resolution: Readonly<ModelInvocationResolutionRecord>,
|
||||
): Promise<void> {
|
||||
await client.query(
|
||||
`INSERT INTO "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_resolutions" (
|
||||
resolution_id, invocation_id, project_id, run_id, step_run_id,
|
||||
trace_id, completion_digest, decision, resolved_by_user_id,
|
||||
resolved_at_ms, mutation_id, mutation_digest, run_event_id,
|
||||
resolution_digest, record_json
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
|
||||
$14, $15::jsonb
|
||||
)`,
|
||||
[
|
||||
resolution.resolutionId,
|
||||
resolution.invocationId,
|
||||
resolution.projectId,
|
||||
resolution.runId,
|
||||
resolution.stepRunId,
|
||||
resolution.traceId,
|
||||
resolution.completionDigest,
|
||||
resolution.decision,
|
||||
resolution.resolvedByUserId,
|
||||
resolution.resolvedAtMs,
|
||||
resolution.stepRunMutationId,
|
||||
resolution.stepRunMutationDigest,
|
||||
resolution.runEventId,
|
||||
resolution.resolutionDigest,
|
||||
JSON.stringify(resolution),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
export async function assertCurrent(
|
||||
client: PostgresClient,
|
||||
mutation: Readonly<StepRunMutation>,
|
||||
projectId: string,
|
||||
): Promise<void> {
|
||||
const result = await client.query<Row>(
|
||||
`SELECT
|
||||
step.kind AS "stepKind", step.status AS "stepStatus",
|
||||
step.version AS "stepVersion", step.step_run_digest AS "stepDigest",
|
||||
run.project_id AS "projectId", run.status AS "runStatus",
|
||||
run.version AS "runVersion",
|
||||
run.event_sequence AS "runEventSequence"
|
||||
FROM "ql3"."step_runs" AS step
|
||||
JOIN "ql3"."runs" AS run ON run.id = step.run_id
|
||||
WHERE step.id = $1 AND step.run_id = $2
|
||||
LIMIT 2
|
||||
FOR UPDATE OF step, run`,
|
||||
[mutation.stepRun.id, mutation.runId],
|
||||
);
|
||||
const row = result.rows[0];
|
||||
if (
|
||||
result.rows.length !== 1 ||
|
||||
!row ||
|
||||
text(row, 'stepKind') !== 'model' ||
|
||||
text(row, 'stepStatus') !== mutation.previousStatus ||
|
||||
integer(row, 'stepVersion') !== mutation.expectedStepRunVersion ||
|
||||
text(row, 'stepDigest') !== mutation.expectedStepRunDigest ||
|
||||
text(row, 'projectId') !== projectId ||
|
||||
integer(row, 'runVersion') !== mutation.expectedRunVersion ||
|
||||
integer(row, 'runEventSequence') !== mutation.expectedRunEventSequence ||
|
||||
TERMINAL_RUN_STATUSES.has(text(row, 'runStatus'))
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
import { POSTGRES_MODEL_INVOCATION_SCHEMA } from '../../migration/modelInvocationMigration';
|
||||
|
||||
import type { Queryable, Row } from './authority';
|
||||
import {
|
||||
COMPLETION_SELECT,
|
||||
RESOLUTION_SELECT,
|
||||
START_SELECT,
|
||||
USAGE_SELECT,
|
||||
} from './codec';
|
||||
|
||||
export async function startRows(
|
||||
queryable: Queryable,
|
||||
where: string,
|
||||
values: readonly unknown[],
|
||||
): Promise<readonly Row[]> {
|
||||
const result = await queryable.query<Row>(
|
||||
`SELECT ${START_SELECT}
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_starts" AS start
|
||||
JOIN "ql3"."step_run_mutations" AS mutation
|
||||
ON mutation.mutation_id = start.mutation_id
|
||||
JOIN "ql3"."run_events" AS event ON event.id = start.run_event_id
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
values,
|
||||
);
|
||||
return result.rows;
|
||||
}
|
||||
|
||||
export async function completionRows(
|
||||
queryable: Queryable,
|
||||
where: string,
|
||||
values: readonly unknown[],
|
||||
): Promise<readonly Row[]> {
|
||||
const result = await queryable.query<Row>(
|
||||
`SELECT ${COMPLETION_SELECT}
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_completions" AS completion
|
||||
JOIN "ql3"."step_run_mutations" AS mutation
|
||||
ON mutation.mutation_id = completion.mutation_id
|
||||
JOIN "ql3"."run_events" AS event ON event.id = completion.run_event_id
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
values,
|
||||
);
|
||||
return result.rows;
|
||||
}
|
||||
|
||||
export async function usageRows(
|
||||
queryable: Queryable,
|
||||
where: string,
|
||||
values: readonly unknown[],
|
||||
limit = 2,
|
||||
): Promise<readonly Row[]> {
|
||||
const result = await queryable.query<Row>(
|
||||
`SELECT ${USAGE_SELECT}
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_usage_ledger" AS usage
|
||||
WHERE ${where}
|
||||
ORDER BY usage.settled_at_ms, usage.invocation_id
|
||||
LIMIT $${values.length + 1}`,
|
||||
[...values, limit],
|
||||
);
|
||||
return result.rows;
|
||||
}
|
||||
|
||||
export async function quotaReservationRows(
|
||||
queryable: Queryable,
|
||||
where: string,
|
||||
values: readonly unknown[],
|
||||
limit = 2,
|
||||
): Promise<readonly Row[]> {
|
||||
const result = await queryable.query<Row>(
|
||||
`SELECT reservation.record_json AS "recordJson"
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_quota_reservations"
|
||||
AS reservation
|
||||
WHERE ${where}
|
||||
ORDER BY reservation.window_start_ms, reservation.invocation_id
|
||||
LIMIT $${values.length + 1}`,
|
||||
[...values, limit],
|
||||
);
|
||||
return result.rows;
|
||||
}
|
||||
|
||||
export async function quotaSettlementRows(
|
||||
queryable: Queryable,
|
||||
where: string,
|
||||
values: readonly unknown[],
|
||||
): Promise<readonly Row[]> {
|
||||
const result = await queryable.query<Row>(
|
||||
`SELECT settlement.record_json AS "recordJson"
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_quota_settlements"
|
||||
AS settlement
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
values,
|
||||
);
|
||||
return result.rows;
|
||||
}
|
||||
|
||||
export async function priceQuoteRows(
|
||||
queryable: Queryable,
|
||||
where: string,
|
||||
values: readonly unknown[],
|
||||
): Promise<readonly Row[]> {
|
||||
const result = await queryable.query<Row>(
|
||||
`SELECT quote.record_json AS "recordJson"
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_price_quotes"
|
||||
AS quote
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
values,
|
||||
);
|
||||
return result.rows;
|
||||
}
|
||||
|
||||
export async function priceSettlementRows(
|
||||
queryable: Queryable,
|
||||
where: string,
|
||||
values: readonly unknown[],
|
||||
): Promise<readonly Row[]> {
|
||||
const result = await queryable.query<Row>(
|
||||
`SELECT settlement.record_json AS "recordJson"
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_price_settlements"
|
||||
AS settlement
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
values,
|
||||
);
|
||||
return result.rows;
|
||||
}
|
||||
|
||||
export async function resolutionRows(
|
||||
queryable: Queryable,
|
||||
where: string,
|
||||
values: readonly unknown[],
|
||||
): Promise<readonly Row[]> {
|
||||
const result = await queryable.query<Row>(
|
||||
`SELECT ${RESOLUTION_SELECT}
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_resolutions" AS resolution
|
||||
JOIN "ql3"."step_run_mutations" AS mutation
|
||||
ON mutation.mutation_id = resolution.mutation_id
|
||||
JOIN "ql3"."run_events" AS event ON event.id = resolution.run_event_id
|
||||
WHERE ${where}
|
||||
LIMIT 2`,
|
||||
values,
|
||||
);
|
||||
return result.rows;
|
||||
}
|
||||
+366
@@ -0,0 +1,366 @@
|
||||
import type { PostgresPool } from '@qinglong/runtime-core';
|
||||
|
||||
import { POSTGRES_MODEL_INVOCATION_SCHEMA } from '../../migration/modelInvocationMigration';
|
||||
import {
|
||||
type ModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceSettlement,
|
||||
} from '../../pricing/pricing';
|
||||
import { type PluginPackagePromptOutputArtifact } from '../../prompt-output/pluginPackagePromptOutputArtifact';
|
||||
import type { PluginPackagePromptOutputArtifactTombstone } from '../../prompt-output/pluginPackagePromptOutputRetention';
|
||||
import { readPostgresPluginPackagePromptOutputArtifactInTransaction } from '../../prompt-output/storage/postgresPluginPackagePromptOutputArtifactRepository';
|
||||
import { readPostgresPluginPackagePromptOutputArtifactTombstoneInTransaction } from '../../prompt-output/storage/postgresPluginPackagePromptOutputRetentionRepository';
|
||||
import {
|
||||
MAX_MODEL_INVOCATION_USAGE_SUMMARY_ROWS,
|
||||
ModelInvocationUsageSummaryLimitExceededError,
|
||||
normalizeModelInvocationUsageLedgerQuery,
|
||||
normalizeModelInvocationUsageLedgerSummaryQuery,
|
||||
type ModelInvocationUsageLedgerPage,
|
||||
type ModelInvocationUsageLedgerQuery,
|
||||
type ModelInvocationUsageLedgerRecord,
|
||||
type ModelInvocationUsageLedgerSummary,
|
||||
type ModelInvocationUsageLedgerSummaryQuery,
|
||||
} from '../../usage/usageLedger';
|
||||
import {
|
||||
type ModelInvocationQuotaReservation,
|
||||
type ModelInvocationQuotaSettlement,
|
||||
type ModelInvocationQuotaWindowUsage,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
|
||||
import type { Row } from './authority';
|
||||
import { identifier, integer, mapStorageError, unavailable } from './authority';
|
||||
import {
|
||||
parseCompletion,
|
||||
parsePriceQuote,
|
||||
parsePriceSettlement,
|
||||
parseQuotaReservation,
|
||||
parseQuotaSettlement,
|
||||
parseStart,
|
||||
parseUsage,
|
||||
} from './codec';
|
||||
import { quotaWindowUsage } from './mutations';
|
||||
import {
|
||||
completionRows,
|
||||
priceQuoteRows,
|
||||
priceSettlementRows,
|
||||
quotaReservationRows,
|
||||
quotaSettlementRows,
|
||||
startRows,
|
||||
usageRows,
|
||||
} from './queries';
|
||||
|
||||
export async function findStartOperation(
|
||||
pool: PostgresPool,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationStartRecord> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
try {
|
||||
const rows = await startRows(pool, 'start.invocation_id = $1', [
|
||||
invocationId,
|
||||
]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseStart(rows[0]) : null;
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function findCompletionOperation(
|
||||
pool: PostgresPool,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationCompletionRecord> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
try {
|
||||
const rows = await completionRows(pool, 'completion.invocation_id = $1', [
|
||||
invocationId,
|
||||
]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseCompletion(rows[0]) : null;
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function findPromptOutputArtifactOperation(
|
||||
pool: PostgresPool,
|
||||
artifactIdValue: string,
|
||||
): Promise<Readonly<PluginPackagePromptOutputArtifact> | null> {
|
||||
const artifactId = identifier(artifactIdValue);
|
||||
try {
|
||||
return await readPostgresPluginPackagePromptOutputArtifactInTransaction(
|
||||
pool,
|
||||
artifactId,
|
||||
);
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function findPromptOutputArtifactTombstoneOperation(
|
||||
pool: PostgresPool,
|
||||
artifactIdValue: string,
|
||||
): Promise<Readonly<PluginPackagePromptOutputArtifactTombstone> | null> {
|
||||
const artifactId = identifier(artifactIdValue);
|
||||
try {
|
||||
return await readPostgresPluginPackagePromptOutputArtifactTombstoneInTransaction(
|
||||
pool,
|
||||
artifactId,
|
||||
);
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function findUsageOperation(
|
||||
pool: PostgresPool,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerRecord> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
try {
|
||||
const rows = await usageRows(pool, 'usage.invocation_id = $1', [
|
||||
invocationId,
|
||||
]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseUsage(rows[0]) : null;
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function findPriceQuoteOperation(
|
||||
pool: PostgresPool,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationPriceQuote> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
try {
|
||||
const rows = await priceQuoteRows(pool, 'quote.invocation_id = $1', [
|
||||
invocationId,
|
||||
]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parsePriceQuote(rows[0]) : null;
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function findPriceSettlementOperation(
|
||||
pool: PostgresPool,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationPriceSettlement> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
try {
|
||||
const [quotes, completions, settlements] = await Promise.all([
|
||||
priceQuoteRows(pool, 'quote.invocation_id = $1', [invocationId]),
|
||||
completionRows(pool, 'completion.invocation_id = $1', [invocationId]),
|
||||
priceSettlementRows(pool, 'settlement.invocation_id = $1', [
|
||||
invocationId,
|
||||
]),
|
||||
]);
|
||||
if (quotes.length > 1 || completions.length > 1 || settlements.length > 1) {
|
||||
throw unavailable();
|
||||
}
|
||||
if (!settlements[0]) return null;
|
||||
if (!quotes[0] || !completions[0]) throw unavailable();
|
||||
return parsePriceSettlement(
|
||||
settlements[0],
|
||||
parsePriceQuote(quotes[0]),
|
||||
parseCompletion(completions[0]),
|
||||
);
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function findQuotaReservationOperation(
|
||||
pool: PostgresPool,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationQuotaReservation> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
try {
|
||||
const rows = await quotaReservationRows(
|
||||
pool,
|
||||
'reservation.invocation_id = $1',
|
||||
[invocationId],
|
||||
);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseQuotaReservation(rows[0]) : null;
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function findQuotaSettlementOperation(
|
||||
pool: PostgresPool,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationQuotaSettlement> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
try {
|
||||
const [reservations, completions, settlements] = await Promise.all([
|
||||
quotaReservationRows(pool, 'reservation.invocation_id = $1', [
|
||||
invocationId,
|
||||
]),
|
||||
completionRows(pool, 'completion.invocation_id = $1', [invocationId]),
|
||||
quotaSettlementRows(pool, 'settlement.invocation_id = $1', [
|
||||
invocationId,
|
||||
]),
|
||||
]);
|
||||
if (
|
||||
reservations.length > 1 ||
|
||||
completions.length > 1 ||
|
||||
settlements.length > 1
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
if (!settlements[0]) return null;
|
||||
if (!reservations[0] || !completions[0]) throw unavailable();
|
||||
return parseQuotaSettlement(
|
||||
settlements[0],
|
||||
parseQuotaReservation(reservations[0]),
|
||||
parseCompletion(completions[0]),
|
||||
);
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function readQuotaWindowUsageOperation(
|
||||
pool: PostgresPool,
|
||||
projectIdValue: string,
|
||||
atMsValue?: number,
|
||||
): Promise<Readonly<ModelInvocationQuotaWindowUsage> | null> {
|
||||
const projectId = identifier(projectIdValue);
|
||||
if (
|
||||
atMsValue !== undefined &&
|
||||
(!Number.isSafeInteger(atMsValue) || atMsValue < 0)
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
try {
|
||||
let atMs = atMsValue;
|
||||
if (atMs === undefined) {
|
||||
const observation = await pool.query<Row>(
|
||||
`SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint
|
||||
AS "atMs"`,
|
||||
);
|
||||
const row = observation.rows[0];
|
||||
if (observation.rows.length !== 1 || !row) throw unavailable();
|
||||
atMs = integer(row, 'atMs');
|
||||
}
|
||||
const result = await pool.query<Row>(
|
||||
`SELECT reservation.record_json AS "recordJson"
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_quota_reservations"
|
||||
AS reservation
|
||||
WHERE reservation.project_id = $1
|
||||
AND reservation.window_start_ms <= $2
|
||||
AND reservation.window_end_ms > $3
|
||||
ORDER BY reservation.reserved_at_ms DESC,
|
||||
reservation.invocation_id DESC
|
||||
LIMIT 1`,
|
||||
[projectId, atMs, atMs],
|
||||
);
|
||||
const row = result.rows[0];
|
||||
if (!row) return null;
|
||||
const reservation = parseQuotaReservation(row);
|
||||
return quotaWindowUsage(
|
||||
pool,
|
||||
projectId,
|
||||
reservation.windowStartMs,
|
||||
reservation.windowMs,
|
||||
);
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function listProjectUsageOperation(
|
||||
pool: PostgresPool,
|
||||
queryValue: ModelInvocationUsageLedgerQuery,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerPage>> {
|
||||
const query = normalizeModelInvocationUsageLedgerQuery(queryValue);
|
||||
const cursor = query.after;
|
||||
try {
|
||||
const rows = await usageRows(
|
||||
pool,
|
||||
`usage.project_id = $1
|
||||
AND usage.settled_at_ms >= $2 AND usage.settled_at_ms < $3
|
||||
${
|
||||
cursor
|
||||
? `AND (
|
||||
usage.settled_at_ms > $4 OR
|
||||
(usage.settled_at_ms = $5 AND usage.invocation_id > $6)
|
||||
)`
|
||||
: ''
|
||||
}`,
|
||||
[
|
||||
query.projectId,
|
||||
query.fromMsInclusive,
|
||||
query.toMsExclusive,
|
||||
...(cursor
|
||||
? [cursor.settledAtMs, cursor.settledAtMs, cursor.invocationId]
|
||||
: []),
|
||||
],
|
||||
query.limit + 1,
|
||||
);
|
||||
return Object.freeze({
|
||||
records: Object.freeze(
|
||||
rows.slice(0, query.limit).map((row) => parseUsage(row)),
|
||||
),
|
||||
hasMore: rows.length > query.limit,
|
||||
});
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function summarizeProjectUsageOperation(
|
||||
pool: PostgresPool,
|
||||
queryValue: ModelInvocationUsageLedgerSummaryQuery,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerSummary>> {
|
||||
const query = normalizeModelInvocationUsageLedgerSummaryQuery(queryValue);
|
||||
try {
|
||||
const result = await pool.query<Row>(
|
||||
`SELECT
|
||||
COUNT(*)::text AS "invocationCount",
|
||||
COALESCE(SUM(input_tokens), 0)::text AS "inputTokens",
|
||||
COALESCE(SUM(output_tokens), 0)::text AS "outputTokens",
|
||||
COALESCE(SUM(total_tokens), 0)::text AS "totalTokens",
|
||||
COALESCE(SUM(cost_micros), 0)::text AS "knownCostMicros",
|
||||
COUNT(*) FILTER (WHERE cost_micros IS NULL)::text
|
||||
AS "unknownCostInvocations"
|
||||
FROM (
|
||||
SELECT input_tokens, output_tokens, total_tokens, cost_micros
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_usage_ledger"
|
||||
WHERE project_id = $1
|
||||
AND settled_at_ms >= $2 AND settled_at_ms < $3
|
||||
ORDER BY settled_at_ms, invocation_id
|
||||
LIMIT $4
|
||||
) AS bounded_usage`,
|
||||
[
|
||||
query.projectId,
|
||||
query.fromMsInclusive,
|
||||
query.toMsExclusive,
|
||||
MAX_MODEL_INVOCATION_USAGE_SUMMARY_ROWS + 1,
|
||||
],
|
||||
);
|
||||
const row = result.rows[0];
|
||||
if (result.rows.length !== 1 || !row) throw unavailable();
|
||||
if (
|
||||
integer(row, 'invocationCount') > MAX_MODEL_INVOCATION_USAGE_SUMMARY_ROWS
|
||||
) {
|
||||
throw new ModelInvocationUsageSummaryLimitExceededError();
|
||||
}
|
||||
return Object.freeze({
|
||||
invocationCount: integer(row, 'invocationCount'),
|
||||
inputTokens: integer(row, 'inputTokens'),
|
||||
outputTokens: integer(row, 'outputTokens'),
|
||||
totalTokens: integer(row, 'totalTokens'),
|
||||
knownCostMicros: integer(row, 'knownCostMicros'),
|
||||
unknownCostInvocations: integer(row, 'unknownCostInvocations'),
|
||||
});
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
+183
@@ -0,0 +1,183 @@
|
||||
import type { PostgresPool } from '@qinglong/runtime-core';
|
||||
|
||||
import { POSTGRES_MODEL_INVOCATION_SCHEMA } from '../../migration/modelInvocationMigration';
|
||||
import {
|
||||
ModelInvocationConflictError,
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationAuthoritySnapshot,
|
||||
type ModelInvocationRecoveryPage,
|
||||
} from '../modelInvocation';
|
||||
import {
|
||||
normalizeModelInvocationResolutionCommand,
|
||||
type ModelInvocationResolutionCommand,
|
||||
type ModelInvocationResolutionRecord,
|
||||
} from '../modelInvocationResolution';
|
||||
|
||||
import type { Row } from './authority';
|
||||
import {
|
||||
TERMINAL_RUN_STATUSES,
|
||||
identifier,
|
||||
integer,
|
||||
mapStorageError,
|
||||
recoveryLimit,
|
||||
text,
|
||||
unavailable,
|
||||
} from './authority';
|
||||
import {
|
||||
START_SELECT,
|
||||
parseAuthority,
|
||||
parseCompletion,
|
||||
parseResolution,
|
||||
parseStart,
|
||||
} from './codec';
|
||||
import { applyMutation, assertCurrent, insertResolution } from './mutations';
|
||||
import { completionRows, resolutionRows } from './queries';
|
||||
import { runPostgresModelInvocationTransaction } from './transaction';
|
||||
|
||||
export async function findResolutionOperation(
|
||||
pool: PostgresPool,
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationResolutionRecord> | null> {
|
||||
const invocationId = identifier(invocationIdValue);
|
||||
try {
|
||||
const rows = await resolutionRows(pool, 'resolution.invocation_id = $1', [
|
||||
invocationId,
|
||||
]);
|
||||
if (rows.length > 1) throw unavailable();
|
||||
return rows[0] ? parseResolution(rows[0]) : null;
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function readAuthorityOperation(
|
||||
pool: PostgresPool,
|
||||
identity: Readonly<{
|
||||
projectId: string;
|
||||
runId: string;
|
||||
stepRunId: string;
|
||||
}>,
|
||||
): Promise<Readonly<ModelInvocationAuthoritySnapshot> | null> {
|
||||
const projectId = identifier(identity?.projectId);
|
||||
const runId = identifier(identity?.runId);
|
||||
const stepRunId = identifier(identity?.stepRunId);
|
||||
try {
|
||||
const result = await pool.query<Row>(
|
||||
`SELECT
|
||||
run.project_id AS "projectId", run.id AS "runId",
|
||||
run.status AS "runStatus", run.version AS "runVersion",
|
||||
run.event_sequence AS "runEventSequence",
|
||||
step.id AS "stepRunId", step.status AS "stepStatus",
|
||||
step.version AS "stepVersion",
|
||||
step.step_run_digest AS "stepDigest",
|
||||
step.step_run_json AS "stepRunJson"
|
||||
FROM "ql3"."runs" AS run
|
||||
JOIN "ql3"."step_runs" AS step ON step.run_id = run.id
|
||||
WHERE run.project_id = $1 AND run.id = $2 AND step.id = $3
|
||||
LIMIT 2`,
|
||||
[projectId, runId, stepRunId],
|
||||
);
|
||||
if (result.rows.length > 1) throw unavailable();
|
||||
const row = result.rows[0];
|
||||
if (!row || TERMINAL_RUN_STATUSES.has(text(row, 'runStatus'))) {
|
||||
return null;
|
||||
}
|
||||
return parseAuthority(row);
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function listIncompleteOperation(
|
||||
pool: PostgresPool,
|
||||
limitValue: number,
|
||||
): Promise<Readonly<ModelInvocationRecoveryPage>> {
|
||||
const limit = recoveryLimit(limitValue);
|
||||
try {
|
||||
const observation = await pool.query<Row>(
|
||||
`SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint
|
||||
AS "observedAtMs"`,
|
||||
);
|
||||
const observedRow = observation.rows[0];
|
||||
if (observation.rows.length !== 1 || !observedRow) {
|
||||
throw unavailable();
|
||||
}
|
||||
const observedAtMs = integer(observedRow, 'observedAtMs');
|
||||
const result = await pool.query<Row>(
|
||||
`SELECT ${START_SELECT}
|
||||
FROM "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_starts" AS start
|
||||
JOIN "ql3"."step_run_mutations" AS mutation
|
||||
ON mutation.mutation_id = start.mutation_id
|
||||
JOIN "ql3"."run_events" AS event ON event.id = start.run_event_id
|
||||
LEFT JOIN "${POSTGRES_MODEL_INVOCATION_SCHEMA}"."model_invocation_completions" AS completion
|
||||
ON completion.invocation_id = start.invocation_id
|
||||
WHERE completion.invocation_id IS NULL
|
||||
AND start.deadline_at_ms <= $1
|
||||
ORDER BY start.deadline_at_ms, start.invocation_id
|
||||
LIMIT $2`,
|
||||
[observedAtMs, limit + 1],
|
||||
);
|
||||
const hasMore = result.rows.length > limit;
|
||||
return Object.freeze({
|
||||
observedAtMs,
|
||||
candidates: Object.freeze(
|
||||
result.rows.slice(0, limit).map((row) => parseStart(row)),
|
||||
),
|
||||
hasMore,
|
||||
});
|
||||
} catch (error) {
|
||||
throw mapStorageError(error);
|
||||
}
|
||||
}
|
||||
|
||||
export async function resolveOperation(
|
||||
pool: PostgresPool,
|
||||
commandValue: ModelInvocationResolutionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationResolutionRecord>>
|
||||
> {
|
||||
const command = normalizeModelInvocationResolutionCommand(commandValue);
|
||||
const resolution = command.resolution;
|
||||
return runPostgresModelInvocationTransaction(pool, async (client) => {
|
||||
const existing = await resolutionRows(
|
||||
client,
|
||||
`resolution.invocation_id = $1 OR
|
||||
resolution.resolution_id = $2 OR
|
||||
resolution.mutation_id = $3 OR resolution.run_event_id = $4`,
|
||||
[
|
||||
resolution.invocationId,
|
||||
resolution.resolutionId,
|
||||
resolution.stepRunMutationId,
|
||||
resolution.runEventId,
|
||||
],
|
||||
);
|
||||
if (existing.length > 1) throw new ModelInvocationConflictError();
|
||||
if (existing[0]) {
|
||||
const stored = parseResolution(existing[0]);
|
||||
if (JSON.stringify(stored) !== JSON.stringify(resolution)) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
return Object.freeze({ status: 'existing' as const, record: stored });
|
||||
}
|
||||
const completions = await completionRows(
|
||||
client,
|
||||
'completion.invocation_id = $1',
|
||||
[resolution.invocationId],
|
||||
);
|
||||
if (
|
||||
completions.length !== 1 ||
|
||||
JSON.stringify(parseCompletion(completions[0]!)) !==
|
||||
JSON.stringify(command.completion) ||
|
||||
command.completion.outcome !== 'outcome_unknown'
|
||||
) {
|
||||
throw new ModelInvocationConflictError();
|
||||
}
|
||||
await assertCurrent(client, command.stepRunMutation, resolution.projectId);
|
||||
await applyMutation(client, command.stepRunMutation);
|
||||
await insertResolution(client, resolution);
|
||||
return Object.freeze({
|
||||
status: 'created' as const,
|
||||
record: resolution,
|
||||
});
|
||||
});
|
||||
}
|
||||
+269
@@ -0,0 +1,269 @@
|
||||
import type { PostgresPool } from '@qinglong/runtime-core';
|
||||
|
||||
import {
|
||||
type ModelInvocationPriceQuote,
|
||||
type ModelInvocationPriceSettlement,
|
||||
type PricingAwareModelInvocationRepository,
|
||||
} from '../../pricing/pricing';
|
||||
import { type PluginPackagePromptOutputArtifact } from '../../prompt-output/pluginPackagePromptOutputArtifact';
|
||||
import {
|
||||
type CommitPluginPackagePromptOutputResult,
|
||||
type PluginPackagePromptOutputCompletionRepository,
|
||||
} from '../../prompt-output/pluginPackagePromptOutputCompletion';
|
||||
import type { PluginPackagePromptOutputArtifactTombstone } from '../../prompt-output/pluginPackagePromptOutputRetention';
|
||||
import {
|
||||
type ModelInvocationUsageLedgerPage,
|
||||
type ModelInvocationUsageLedgerQuery,
|
||||
type ModelInvocationUsageLedgerRecord,
|
||||
type ModelInvocationUsageLedgerRepository,
|
||||
type ModelInvocationUsageLedgerSummary,
|
||||
type ModelInvocationUsageLedgerSummaryQuery,
|
||||
} from '../../usage/usageLedger';
|
||||
import {
|
||||
type ModelInvocationQuotaAdmission,
|
||||
type ModelInvocationQuotaReservation,
|
||||
type ModelInvocationQuotaSettlement,
|
||||
type ModelInvocationQuotaWindowUsage,
|
||||
type QuotaAwareModelInvocationRepository,
|
||||
} from '../../usage/usageQuota';
|
||||
import {
|
||||
type CommitModelInvocationResult,
|
||||
type ModelInvocationAuthoritySnapshot,
|
||||
type ModelInvocationCompletionCommand,
|
||||
type ModelInvocationCompletionRecord,
|
||||
type ModelInvocationRecoveryPage,
|
||||
type ModelInvocationRepository,
|
||||
type ModelInvocationStartCommand,
|
||||
type ModelInvocationStartRecord,
|
||||
} from '../modelInvocation';
|
||||
import {
|
||||
type ModelInvocationResolutionCommand,
|
||||
type ModelInvocationResolutionRecord,
|
||||
type ModelInvocationResolutionRepository,
|
||||
} from '../modelInvocationResolution';
|
||||
|
||||
import {
|
||||
admitOperation,
|
||||
admitWithPricingOperation,
|
||||
admitWithQuotaOperation,
|
||||
} from './admissionOperations';
|
||||
import { unavailable } from './authority';
|
||||
import {
|
||||
completeOperation,
|
||||
completeWithPricingOperation,
|
||||
completeWithPromptOutputArtifactOperation,
|
||||
completeWithQuotaOperation,
|
||||
} from './completionOperations';
|
||||
import {
|
||||
findCompletionOperation,
|
||||
findPriceQuoteOperation,
|
||||
findPriceSettlementOperation,
|
||||
findPromptOutputArtifactOperation,
|
||||
findPromptOutputArtifactTombstoneOperation,
|
||||
findQuotaReservationOperation,
|
||||
findQuotaSettlementOperation,
|
||||
findStartOperation,
|
||||
findUsageOperation,
|
||||
listProjectUsageOperation,
|
||||
readQuotaWindowUsageOperation,
|
||||
summarizeProjectUsageOperation,
|
||||
} from './readOperations';
|
||||
import {
|
||||
findResolutionOperation,
|
||||
listIncompleteOperation,
|
||||
readAuthorityOperation,
|
||||
resolveOperation,
|
||||
} from './recoveryResolutionOperations';
|
||||
|
||||
export class PostgresModelInvocationRepository
|
||||
implements
|
||||
ModelInvocationRepository,
|
||||
ModelInvocationResolutionRepository,
|
||||
ModelInvocationUsageLedgerRepository,
|
||||
QuotaAwareModelInvocationRepository,
|
||||
PricingAwareModelInvocationRepository,
|
||||
PluginPackagePromptOutputCompletionRepository
|
||||
{
|
||||
constructor(private readonly pool: PostgresPool) {
|
||||
if (
|
||||
!pool ||
|
||||
typeof pool.query !== 'function' ||
|
||||
typeof pool.connect !== 'function'
|
||||
) {
|
||||
throw unavailable();
|
||||
}
|
||||
}
|
||||
|
||||
async findStart(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationStartRecord> | null> {
|
||||
return findStartOperation(this.pool, invocationIdValue);
|
||||
}
|
||||
|
||||
async findCompletion(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationCompletionRecord> | null> {
|
||||
return findCompletionOperation(this.pool, invocationIdValue);
|
||||
}
|
||||
|
||||
async findPromptOutputArtifact(
|
||||
artifactIdValue: string,
|
||||
): Promise<Readonly<PluginPackagePromptOutputArtifact> | null> {
|
||||
return findPromptOutputArtifactOperation(this.pool, artifactIdValue);
|
||||
}
|
||||
|
||||
async findPromptOutputArtifactTombstone(
|
||||
artifactIdValue: string,
|
||||
): Promise<Readonly<PluginPackagePromptOutputArtifactTombstone> | null> {
|
||||
return findPromptOutputArtifactTombstoneOperation(
|
||||
this.pool,
|
||||
artifactIdValue,
|
||||
);
|
||||
}
|
||||
|
||||
async findUsage(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerRecord> | null> {
|
||||
return findUsageOperation(this.pool, invocationIdValue);
|
||||
}
|
||||
|
||||
async findPriceQuote(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationPriceQuote> | null> {
|
||||
return findPriceQuoteOperation(this.pool, invocationIdValue);
|
||||
}
|
||||
|
||||
async findPriceSettlement(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationPriceSettlement> | null> {
|
||||
return findPriceSettlementOperation(this.pool, invocationIdValue);
|
||||
}
|
||||
|
||||
async findQuotaReservation(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationQuotaReservation> | null> {
|
||||
return findQuotaReservationOperation(this.pool, invocationIdValue);
|
||||
}
|
||||
|
||||
async findQuotaSettlement(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationQuotaSettlement> | null> {
|
||||
return findQuotaSettlementOperation(this.pool, invocationIdValue);
|
||||
}
|
||||
|
||||
async readQuotaWindowUsage(
|
||||
projectIdValue: string,
|
||||
atMsValue?: number,
|
||||
): Promise<Readonly<ModelInvocationQuotaWindowUsage> | null> {
|
||||
return readQuotaWindowUsageOperation(this.pool, projectIdValue, atMsValue);
|
||||
}
|
||||
|
||||
async listProjectUsage(
|
||||
queryValue: ModelInvocationUsageLedgerQuery,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerPage>> {
|
||||
return listProjectUsageOperation(this.pool, queryValue);
|
||||
}
|
||||
|
||||
async summarizeProjectUsage(
|
||||
queryValue: ModelInvocationUsageLedgerSummaryQuery,
|
||||
): Promise<Readonly<ModelInvocationUsageLedgerSummary>> {
|
||||
return summarizeProjectUsageOperation(this.pool, queryValue);
|
||||
}
|
||||
|
||||
async findResolution(
|
||||
invocationIdValue: string,
|
||||
): Promise<Readonly<ModelInvocationResolutionRecord> | null> {
|
||||
return findResolutionOperation(this.pool, invocationIdValue);
|
||||
}
|
||||
|
||||
async readAuthority(
|
||||
identity: Readonly<{
|
||||
projectId: string;
|
||||
runId: string;
|
||||
stepRunId: string;
|
||||
}>,
|
||||
): Promise<Readonly<ModelInvocationAuthoritySnapshot> | null> {
|
||||
return readAuthorityOperation(this.pool, identity);
|
||||
}
|
||||
|
||||
async listIncomplete(
|
||||
limitValue: number,
|
||||
): Promise<Readonly<ModelInvocationRecoveryPage>> {
|
||||
return listIncompleteOperation(this.pool, limitValue);
|
||||
}
|
||||
|
||||
async admit(
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>
|
||||
> {
|
||||
return admitOperation(this.pool, commandValue);
|
||||
}
|
||||
|
||||
async admitWithQuota(
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
admissionValue: ModelInvocationQuotaAdmission,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>
|
||||
> {
|
||||
return admitWithQuotaOperation(this.pool, commandValue, admissionValue);
|
||||
}
|
||||
|
||||
async admitWithPricing(
|
||||
commandValue: ModelInvocationStartCommand,
|
||||
quoteValue: ModelInvocationPriceQuote,
|
||||
admissionValue?: ModelInvocationQuotaAdmission,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationStartRecord>>
|
||||
> {
|
||||
return admitWithPricingOperation(
|
||||
this.pool,
|
||||
commandValue,
|
||||
quoteValue,
|
||||
admissionValue,
|
||||
);
|
||||
}
|
||||
|
||||
async complete(
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
return completeOperation(this.pool, commandValue);
|
||||
}
|
||||
|
||||
async completeWithQuota(
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
return completeWithQuotaOperation(this.pool, commandValue);
|
||||
}
|
||||
|
||||
async completeWithPricing(
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationCompletionRecord>>
|
||||
> {
|
||||
return completeWithPricingOperation(this.pool, commandValue);
|
||||
}
|
||||
|
||||
async completeWithPromptOutputArtifact(
|
||||
commandValue: ModelInvocationCompletionCommand,
|
||||
artifactValue: PluginPackagePromptOutputArtifact,
|
||||
): Promise<Readonly<CommitPluginPackagePromptOutputResult>> {
|
||||
return completeWithPromptOutputArtifactOperation(
|
||||
this.pool,
|
||||
commandValue,
|
||||
artifactValue,
|
||||
);
|
||||
}
|
||||
|
||||
async resolve(
|
||||
commandValue: ModelInvocationResolutionCommand,
|
||||
): Promise<
|
||||
Readonly<CommitModelInvocationResult<ModelInvocationResolutionRecord>>
|
||||
> {
|
||||
return resolveOperation(this.pool, commandValue);
|
||||
}
|
||||
}
|
||||
+64
@@ -0,0 +1,64 @@
|
||||
import type { PostgresClient, PostgresPool } from '@qinglong/runtime-core';
|
||||
|
||||
import {
|
||||
MAX_TRANSACTION_ATTEMPTS,
|
||||
RETRYABLE_SQL_STATES,
|
||||
mapStorageError,
|
||||
sqlState,
|
||||
unavailable,
|
||||
} from './authority';
|
||||
|
||||
export async function begin(client: PostgresClient): Promise<void> {
|
||||
await client.query('BEGIN ISOLATION LEVEL SERIALIZABLE');
|
||||
await client.query(`SELECT set_config('statement_timeout', $1, true)`, [
|
||||
'5s',
|
||||
]);
|
||||
await client.query(`SELECT set_config('lock_timeout', $1, true)`, ['2s']);
|
||||
await client.query(
|
||||
`SELECT set_config('idle_in_transaction_session_timeout', $1, true)`,
|
||||
['5s'],
|
||||
);
|
||||
}
|
||||
|
||||
export async function rollback(client: PostgresClient): Promise<void> {
|
||||
try {
|
||||
await client.query('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the original transaction failure.
|
||||
}
|
||||
}
|
||||
|
||||
export async function runPostgresModelInvocationTransaction<T>(
|
||||
pool: PostgresPool,
|
||||
work: (client: PostgresClient) => Promise<T>,
|
||||
): Promise<T> {
|
||||
for (let attempt = 0; attempt < MAX_TRANSACTION_ATTEMPTS; attempt += 1) {
|
||||
let client: PostgresClient;
|
||||
try {
|
||||
client = await pool.connect();
|
||||
} catch (error) {
|
||||
throw unavailable(error);
|
||||
}
|
||||
let began = false;
|
||||
try {
|
||||
await begin(client);
|
||||
began = true;
|
||||
const result = await work(client);
|
||||
await client.query('COMMIT');
|
||||
began = false;
|
||||
return result;
|
||||
} catch (error) {
|
||||
if (began) await rollback(client);
|
||||
if (
|
||||
RETRYABLE_SQL_STATES.has(sqlState(error)) &&
|
||||
attempt + 1 < MAX_TRANSACTION_ATTEMPTS
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
throw mapStorageError(error);
|
||||
} finally {
|
||||
client.release();
|
||||
}
|
||||
}
|
||||
throw unavailable();
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
export { PostgresModelInvocationRepository } from './postgres-model-invocation-repository/repository';
|
||||
Reference in New Issue
Block a user