Skip to content

Commit 90ab60f

Browse files
RieranthonyRieranthony
committed
Add BYOK fallback to Cossistant credits
1 parent 651d5be commit 90ab60f

20 files changed

Lines changed: 789 additions & 325 deletions

File tree

‎apps/api/src/ai-pipeline/background-pipeline/index.test.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,21 @@ const runBackgroundKnowledgeGapReviewMock = mock((async () => ({
121121
reason: "no_candidate_gap" as const,
122122
})) as (...args: unknown[]) => Promise<any>);
123123
const createModelMock = mock((modelId: string) => modelId);
124+
const runWithOpenRouterByokFallbackMock = mock(
125+
async (params: {
126+
modelId: string;
127+
operation: (resolution: {
128+
model: string;
129+
billingSource: "cossistant";
130+
}) => Promise<unknown>;
131+
}) => ({
132+
result: await params.operation({
133+
model: createModelMock(params.modelId),
134+
billingSource: "cossistant" as const,
135+
}),
136+
billingSource: "cossistant" as const,
137+
})
138+
);
124139
const generateTextMock = mock((async () => ({
125140
output: {
126141
title: "Visitor asked for help",
@@ -181,6 +196,7 @@ mock.module("./knowledge-gap-review", () => ({
181196
mock.module("@api/lib/ai", () => ({
182197
createModel: createModelMock,
183198
generateText: generateTextMock,
199+
runWithOpenRouterByokFallback: runWithOpenRouterByokFallbackMock,
184200
Output: {
185201
object: (params: unknown) => params,
186202
},
@@ -272,6 +288,7 @@ describe("runBackgroundPipeline", () => {
272288
runGenerationRuntimeMock.mockReset();
273289
runBackgroundKnowledgeGapReviewMock.mockReset();
274290
createModelMock.mockReset();
291+
runWithOpenRouterByokFallbackMock.mockReset();
275292
generateTextMock.mockReset();
276293
resolveModelForExecutionMock.mockReset();
277294
loadCurrentConversationMock.mockReset();
@@ -371,6 +388,13 @@ describe("runBackgroundPipeline", () => {
371388
reason: "no_candidate_gap",
372389
});
373390
createModelMock.mockImplementation((modelId: string) => modelId);
391+
runWithOpenRouterByokFallbackMock.mockImplementation(async (params) => ({
392+
result: await params.operation({
393+
model: createModelMock(params.modelId),
394+
billingSource: "cossistant" as const,
395+
}),
396+
billingSource: "cossistant" as const,
397+
}));
374398
generateTextMock.mockResolvedValue({
375399
output: {
376400
title: "Visitor asked for help",

‎apps/api/src/ai-pipeline/background-pipeline/knowledge-gap-review.test.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,21 @@ const createModelForWebsiteMock = mock(async (modelId: string) => ({
1313
model: createModelMock(modelId),
1414
billingSource: "cossistant" as const,
1515
}));
16+
const runWithOpenRouterByokFallbackMock = mock(
17+
async (params: {
18+
modelId: string;
19+
operation: (resolution: {
20+
model: string;
21+
billingSource: "cossistant";
22+
}) => Promise<unknown>;
23+
}) => ({
24+
result: await params.operation({
25+
model: createModelMock(params.modelId),
26+
billingSource: "cossistant" as const,
27+
}),
28+
billingSource: "cossistant" as const,
29+
})
30+
);
1631
const generateTextMock = mock((async () => ({
1732
output: {
1833
action: "skip",
@@ -48,6 +63,7 @@ mock.module("@api/lib/ai", () => ({
4863
createModel: createModelMock,
4964
createModelForWebsite: createModelForWebsiteMock,
5065
generateText: generateTextMock,
66+
runWithOpenRouterByokFallback: runWithOpenRouterByokFallbackMock,
5167
Output: {
5268
object: (params: unknown) => params,
5369
},
@@ -158,6 +174,7 @@ describe("runBackgroundKnowledgeGapReview", () => {
158174
getConversationTimelineItemsMock.mockReset();
159175
createModelMock.mockReset();
160176
createModelForWebsiteMock.mockReset();
177+
runWithOpenRouterByokFallbackMock.mockReset();
161178
generateTextMock.mockReset();
162179
resolveModelForExecutionMock.mockReset();
163180
requestKnowledgeClarificationMock.mockReset();
@@ -171,6 +188,13 @@ describe("runBackgroundKnowledgeGapReview", () => {
171188
model: createModelMock(modelId),
172189
billingSource: "cossistant" as const,
173190
}));
191+
runWithOpenRouterByokFallbackMock.mockImplementation(async (params) => ({
192+
result: await params.operation({
193+
model: createModelMock(params.modelId),
194+
billingSource: "cossistant" as const,
195+
}),
196+
billingSource: "cossistant" as const,
197+
}));
174198
generateTextMock.mockResolvedValue({
175199
output: {
176200
action: "skip",

‎apps/api/src/ai-pipeline/background-pipeline/knowledge-gap-review.ts‎

Lines changed: 27 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import { getConversationTimelineItems } from "@api/db/queries/conversation";
22
import { getKnowledgeById } from "@api/db/queries/knowledge";
33
import { getActiveKnowledgeClarificationForConversation } from "@api/db/queries/knowledge-clarification";
4-
import { createModelForWebsite, generateText, Output } from "@api/lib/ai";
5-
import { resolveModelForExecution } from "@api/lib/ai-credits/config";
64
import {
7-
recordOpenRouterByokFailure,
8-
recordOpenRouterByokSuccess,
9-
} from "@api/lib/openrouter-byok/resolver";
5+
generateText,
6+
Output,
7+
runWithOpenRouterByokFallback,
8+
} from "@api/lib/ai";
9+
import { resolveModelForExecution } from "@api/lib/ai-credits/config";
1010
import {
1111
buildConversationClarificationContextSnapshot,
1212
buildSpecificClarificationTopicSummary,
@@ -275,17 +275,16 @@ export async function runBackgroundKnowledgeGapReview(params: {
275275
organizationId: params.input.organizationId,
276276
websiteId: params.input.websiteId,
277277
};
278-
const model = await createModelForWebsite(modelResolution.modelIdResolved, {
279-
context: openRouterContext,
280-
});
281-
let review: Awaited<ReturnType<typeof generateText>>;
282-
try {
283-
review = await generateText({
284-
model: model.model,
285-
output: Output.object({
286-
schema: BACKGROUND_KNOWLEDGE_GAP_REVIEW_OUTPUT_SCHEMA,
287-
}),
288-
system: `You decide whether an internal knowledge clarification request should be created.
278+
const { result: review } = await runWithOpenRouterByokFallback({
279+
modelId: modelResolution.modelIdResolved,
280+
options: { context: openRouterContext },
281+
operation: ({ model }) =>
282+
generateText({
283+
model,
284+
output: Output.object({
285+
schema: BACKGROUND_KNOWLEDGE_GAP_REVIEW_OUTPUT_SCHEMA,
286+
}),
287+
system: `You decide whether an internal knowledge clarification request should be created.
289288
290289
Open a clarification only when the recent knowledge-base retrieval and conversation suggest the FAQ or internal knowledge is incomplete, weak, stale, or contradicted by a teammate.
291290
@@ -294,29 +293,18 @@ Do NOT create a clarification for normal teammate handling, acknowledgements, or
294293
If you choose create:
295294
- Write a short, concrete topic summary.
296295
- Focus on the missing policy, workflow, or product detail the team should clarify.`,
297-
prompt: [
298-
`Trigger sender: ${params.intake.triggerMessage?.senderType ?? "none"}`,
299-
`Trigger visibility: ${params.intake.triggerMessage?.visibility ?? "none"}`,
300-
`Trigger text: ${triggerText || "none"}`,
301-
`Latest KB search workflow:\n${formatSearchSignals(latestWorkflowSignals)}`,
302-
`FAQ candidates to deepen:\n${formatFaqCandidates(faqCandidates)}`,
303-
`Recent transcript:\n${transcript || "- none"}`,
304-
].join("\n\n"),
305-
temperature: 0,
306-
maxOutputTokens: 220,
307-
});
308-
await recordOpenRouterByokSuccess({
309-
context: openRouterContext,
310-
billingSource: model.billingSource,
311-
});
312-
} catch (error) {
313-
await recordOpenRouterByokFailure({
314-
context: openRouterContext,
315-
billingSource: model.billingSource,
316-
error,
317-
});
318-
throw error;
319-
}
296+
prompt: [
297+
`Trigger sender: ${params.intake.triggerMessage?.senderType ?? "none"}`,
298+
`Trigger visibility: ${params.intake.triggerMessage?.visibility ?? "none"}`,
299+
`Trigger text: ${triggerText || "none"}`,
300+
`Latest KB search workflow:\n${formatSearchSignals(latestWorkflowSignals)}`,
301+
`FAQ candidates to deepen:\n${formatFaqCandidates(faqCandidates)}`,
302+
`Recent transcript:\n${transcript || "- none"}`,
303+
].join("\n\n"),
304+
temperature: 0,
305+
maxOutputTokens: 220,
306+
}),
307+
});
320308

321309
const reviewOutput = review.output;
322310
if (!(reviewOutput && reviewOutput.action !== "skip")) {

‎apps/api/src/ai-pipeline/background-pipeline/title-review.test.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,21 @@ const createModelForWebsiteMock = mock(async (modelId: string) => ({
66
model: createModelMock(modelId),
77
billingSource: "cossistant" as const,
88
}));
9+
const runWithOpenRouterByokFallbackMock = mock(
10+
async (params: {
11+
modelId: string;
12+
operation: (resolution: {
13+
model: string;
14+
billingSource: "cossistant";
15+
}) => Promise<unknown>;
16+
}) => ({
17+
result: await params.operation({
18+
model: createModelMock(params.modelId),
19+
billingSource: "cossistant" as const,
20+
}),
21+
billingSource: "cossistant" as const,
22+
})
23+
);
924
const generateTextMock = mock((async () => ({
1025
output: {
1126
title: "Visitor asked for help",
@@ -41,6 +56,7 @@ mock.module("@api/lib/ai", () => ({
4156
createModel: createModelMock,
4257
createModelForWebsite: createModelForWebsiteMock,
4358
generateText: generateTextMock,
59+
runWithOpenRouterByokFallback: runWithOpenRouterByokFallbackMock,
4460
Output: {
4561
object: (params: unknown) => params,
4662
},
@@ -176,6 +192,7 @@ describe("runBackgroundTitleReview", () => {
176192
beforeEach(() => {
177193
createModelMock.mockReset();
178194
createModelForWebsiteMock.mockReset();
195+
runWithOpenRouterByokFallbackMock.mockReset();
179196
generateTextMock.mockReset();
180197
resolveModelForExecutionMock.mockReset();
181198
logAiPipelineMock.mockReset();
@@ -189,6 +206,13 @@ describe("runBackgroundTitleReview", () => {
189206
model: createModelMock(modelId),
190207
billingSource: "cossistant" as const,
191208
}));
209+
runWithOpenRouterByokFallbackMock.mockImplementation(async (params) => ({
210+
result: await params.operation({
211+
model: createModelMock(params.modelId),
212+
billingSource: "cossistant" as const,
213+
}),
214+
billingSource: "cossistant" as const,
215+
}));
192216
resolveModelForExecutionMock.mockImplementation((modelId: string) => ({
193217
modelIdResolved: modelId,
194218
}));

‎apps/api/src/ai-pipeline/background-pipeline/title-review.ts‎

Lines changed: 24 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import type { Database } from "@api/db";
22
import type { AiAgentSelect } from "@api/db/schema/ai-agent";
33
import type { ConversationSelect } from "@api/db/schema/conversation";
4-
import { createModelForWebsite, generateText, Output } from "@api/lib/ai";
5-
import { resolveModelForExecution } from "@api/lib/ai-credits/config";
64
import {
7-
recordOpenRouterByokFailure,
8-
recordOpenRouterByokSuccess,
9-
} from "@api/lib/openrouter-byok/resolver";
5+
generateText,
6+
Output,
7+
runWithOpenRouterByokFallback,
8+
} from "@api/lib/ai";
9+
import { resolveModelForExecution } from "@api/lib/ai-credits/config";
1010
import { z } from "zod";
1111
import { logAiPipeline } from "../logger";
1212
import type {
@@ -331,28 +331,26 @@ export async function runBackgroundTitleReview(
331331
organizationId: params.organizationId,
332332
websiteId: params.websiteId,
333333
};
334-
const model = await createModelForWebsite(modelResolution.modelIdResolved, {
335-
context: openRouterContext,
336-
});
337334
try {
338-
const review = await generateText({
339-
model: model.model,
340-
output: Output.object({
341-
schema: TITLE_REVIEW_OUTPUT_SCHEMA,
342-
}),
343-
system:
344-
"You generate concise internal conversation titles for support teams. Always return a usable title. Use honest fallback titles for greetings, thanks, or generic help requests.",
345-
prompt: buildTitleReviewPrompt({
346-
conversation,
347-
websiteDefaultLanguage: params.websiteDefaultLanguage,
348-
transcript,
349-
}),
350-
temperature: 0,
351-
maxOutputTokens: 180,
352-
});
353-
await recordOpenRouterByokSuccess({
354-
context: openRouterContext,
355-
billingSource: model.billingSource,
335+
const { result: review } = await runWithOpenRouterByokFallback({
336+
modelId: modelResolution.modelIdResolved,
337+
options: { context: openRouterContext },
338+
operation: ({ model }) =>
339+
generateText({
340+
model,
341+
output: Output.object({
342+
schema: TITLE_REVIEW_OUTPUT_SCHEMA,
343+
}),
344+
system:
345+
"You generate concise internal conversation titles for support teams. Always return a usable title. Use honest fallback titles for greetings, thanks, or generic help requests.",
346+
prompt: buildTitleReviewPrompt({
347+
conversation,
348+
websiteDefaultLanguage: params.websiteDefaultLanguage,
349+
transcript,
350+
}),
351+
temperature: 0,
352+
maxOutputTokens: 180,
353+
}),
356354
});
357355

358356
const output = review.output;
@@ -380,11 +378,6 @@ export async function runBackgroundTitleReview(
380378
});
381379
}
382380
} catch (error) {
383-
await recordOpenRouterByokFailure({
384-
context: openRouterContext,
385-
billingSource: model.billingSource,
386-
error,
387-
});
388381
title = deriveFallbackTitle({
389382
conversationHistory: params.conversationHistory,
390383
triggerMessage: params.triggerMessage,

‎apps/api/src/ai-pipeline/primary-pipeline/steps/decision/smart/model-runner.ts‎

Lines changed: 16 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,8 @@
1-
import { createModelRawForWebsite, generateText, Output } from "@api/lib/ai";
21
import {
3-
recordOpenRouterByokFailure,
4-
recordOpenRouterByokSuccess,
5-
type OpenRouterBillingSource,
6-
} from "@api/lib/openrouter-byok/resolver";
2+
generateText,
3+
Output,
4+
runWithOpenRouterByokFallback,
5+
} from "@api/lib/ai";
76
import { z } from "zod";
87
import { logAiPipeline } from "../../../../logger";
98
import { observeDecision } from "./rules";
@@ -41,27 +40,23 @@ export async function runSmartDecisionModel(params: {
4140
organizationId: params.input.conversation.organizationId,
4241
websiteId: params.input.conversation.websiteId,
4342
};
44-
let billingSource: OpenRouterBillingSource | undefined;
4543
const timeout = setTimeout(() => {
4644
abortController.abort();
4745
}, modelConfig.timeoutMs);
4846

4947
try {
50-
const model = await createModelRawForWebsite(
51-
modelConfig.id,
52-
openRouterContext
53-
);
54-
billingSource = model.billingSource;
55-
const result = await generateText({
56-
model: model.model,
57-
output: Output.object({ schema: decisionOutputSchema }),
58-
prompt: params.prompt,
59-
temperature: 0,
60-
abortSignal: abortController.signal,
61-
});
62-
await recordOpenRouterByokSuccess({
63-
context: openRouterContext,
64-
billingSource,
48+
const { result } = await runWithOpenRouterByokFallback({
49+
modelId: modelConfig.id,
50+
options: { context: openRouterContext },
51+
kind: "raw",
52+
operation: ({ model }) =>
53+
generateText({
54+
model,
55+
output: Output.object({ schema: decisionOutputSchema }),
56+
prompt: params.prompt,
57+
temperature: 0,
58+
abortSignal: abortController.signal,
59+
}),
6560
});
6661

6762
if (!result.output) {
@@ -118,11 +113,6 @@ export async function runSmartDecisionModel(params: {
118113
source: "model",
119114
};
120115
} catch (error) {
121-
await recordOpenRouterByokFailure({
122-
context: openRouterContext,
123-
billingSource,
124-
error,
125-
});
126116
const isTimeout = error instanceof Error && error.name === "AbortError";
127117
lastFailure = isTimeout ? "timeout" : "error";
128118

0 commit comments

Comments
 (0)