diff --git a/apps/webapp/app/bullmq/start-workers.ts b/apps/webapp/app/bullmq/start-workers.ts index 3c643e6d5..10bda11f7 100644 --- a/apps/webapp/app/bullmq/start-workers.ts +++ b/apps/webapp/app/bullmq/start-workers.ts @@ -138,26 +138,30 @@ export async function initWorkers(): Promise { // Log worker startup logger.log("\nšŸš€ Starting BullMQ workers..."); logger.log("ā”€".repeat(80)); - logger.log(`āœ“ Ingest worker: ${ingestWorker.name} (concurrency: 1)`); logger.log( - `āœ“ Document ingest worker: ${preprocessWorker.name} (concurrency: 3)`, + `āœ“ Ingest worker: ${ingestWorker.name} (concurrency: ${ingestWorker.opts.concurrency ?? 1})`, ); logger.log( - `āœ“ Conversation title worker: ${conversationTitleWorker.name} (concurrency: 10)`, + `āœ“ Document ingest worker: ${preprocessWorker.name} (concurrency: ${preprocessWorker.opts.concurrency ?? 1})`, ); logger.log( - `āœ“ Session compaction worker: ${sessionCompactionWorker.name} (concurrency: 3)`, + `āœ“ Conversation title worker: ${conversationTitleWorker.name} (concurrency: ${conversationTitleWorker.opts.concurrency ?? 1})`, ); logger.log( - `āœ“ Label assignment worker: ${labelAssignmentWorker.name} (concurrency: 5)`, + `āœ“ Session compaction worker: ${sessionCompactionWorker.name} (concurrency: ${sessionCompactionWorker.opts.concurrency ?? 1})`, ); logger.log( - `āœ“ Title generation worker: ${titleGenerationWorker.name} (concurrency: 10)`, + `āœ“ Label assignment worker: ${labelAssignmentWorker.name} (concurrency: ${labelAssignmentWorker.opts.concurrency ?? 1})`, ); logger.log( - `āœ“ Integration run worker: ${integrationRunWorker.name} (concurrency: 3)`, + `āœ“ Title generation worker: ${titleGenerationWorker.name} (concurrency: ${titleGenerationWorker.opts.concurrency ?? 1})`, + ); + logger.log( + `āœ“ Integration run worker: ${integrationRunWorker.name} (concurrency: ${integrationRunWorker.opts.concurrency ?? 1})`, + ); + logger.log( + `āœ“ Scratchpad scan worker: ${scratchpadScanWorker.name} (concurrency: ${scratchpadScanWorker.opts.concurrency ?? 1})`, ); - logger.log(`āœ“ Scratchpad scan worker: ${scratchpadScanWorker.name} (concurrency: 5)`); logger.log(`āœ“ Reminder scheduler: reminder-queue + followup-queue`); logger.log(`āœ“ Scheduled task scheduler: scheduled-task-queue`); logger.log("ā”€".repeat(80)); diff --git a/apps/webapp/app/bullmq/workers/index.ts b/apps/webapp/app/bullmq/workers/index.ts index 2fc16051d..50804b0a6 100644 --- a/apps/webapp/app/bullmq/workers/index.ts +++ b/apps/webapp/app/bullmq/workers/index.ts @@ -37,6 +37,7 @@ import { enqueueGraphResolution, } from "~/lib/queue-adapter.server"; import { logger } from "~/services/logger.service"; +import { getBurstSafeBullmqConcurrency } from "~/services/llm-provider.server"; import { type PersonaGenerationPayload, processPersonaGeneration, @@ -106,7 +107,8 @@ export const preprocessWorker = new Worker( // Callback to enqueue individual chunk ingestion jobs enqueueIngestEpisode, // Callback to enqueue session compaction for conversations - enqueueSessionCompaction, + (compactionPayload, delayMs) => + enqueueSessionCompaction(compactionPayload, delayMs), ); if (!result?.success) { throw new Error(result?.error || "Episode preprocessing failed"); @@ -147,7 +149,10 @@ export const ingestWorker = new Worker( }, { connection: getRedisConnection(), - concurrency: env.BULLMQ_CONCURRENCY_INGEST, // Global limit for ingestion jobs + concurrency: getBurstSafeBullmqConcurrency( + "BULLMQ_CONCURRENCY_INGEST", + env.BULLMQ_CONCURRENCY_INGEST, + ), // Keep proxy/self-hosted chat from being dogpiled by background ingestion unless explicitly overridden }, ); @@ -162,7 +167,10 @@ export const conversationTitleWorker = new Worker( }, { connection: getRedisConnection(), - concurrency: env.BULLMQ_CONCURRENCY_CONVERSATION_TITLE, // Process title creations in parallel + concurrency: getBurstSafeBullmqConcurrency( + "BULLMQ_CONCURRENCY_CONVERSATION_TITLE", + env.BULLMQ_CONCURRENCY_CONVERSATION_TITLE, + ), // Strict providers benefit from serial title generation unless explicitly overridden }, ); @@ -177,7 +185,10 @@ export const sessionCompactionWorker = new Worker( }, { connection: getRedisConnection(), - concurrency: env.BULLMQ_CONCURRENCY_SESSION_COMPACTION, // Process compactions in parallel + concurrency: getBurstSafeBullmqConcurrency( + "BULLMQ_CONCURRENCY_SESSION_COMPACTION", + env.BULLMQ_CONCURRENCY_SESSION_COMPACTION, + ), // Compaction is background polish; keep it from competing with chat unless explicitly overridden }, ); @@ -193,7 +204,10 @@ export const labelAssignmentWorker = new Worker( }, { connection: getRedisConnection(), - concurrency: env.BULLMQ_CONCURRENCY_LABEL_ASSIGNMENT, // Process label assignments in parallel + concurrency: getBurstSafeBullmqConcurrency( + "BULLMQ_CONCURRENCY_LABEL_ASSIGNMENT", + env.BULLMQ_CONCURRENCY_LABEL_ASSIGNMENT, + ), // Labels are non-critical background work for proxy mode unless explicitly overridden }, ); @@ -209,7 +223,10 @@ export const titleGenerationWorker = new Worker( }, { connection: getRedisConnection(), - concurrency: env.BULLMQ_CONCURRENCY_TITLE_GENERATION, // Process title generations in parallel + concurrency: getBurstSafeBullmqConcurrency( + "BULLMQ_CONCURRENCY_TITLE_GENERATION", + env.BULLMQ_CONCURRENCY_TITLE_GENERATION, + ), // Keep follow-up title generation from stampeding strict providers unless explicitly overridden }, ); @@ -225,7 +242,10 @@ export const personaGenerationWorker = new Worker( }, { connection: getRedisConnection(), - concurrency: env.BULLMQ_CONCURRENCY_PERSONA_GENERATION, // Persona is CPU-intensive + concurrency: getBurstSafeBullmqConcurrency( + "BULLMQ_CONCURRENCY_PERSONA_GENERATION", + env.BULLMQ_CONCURRENCY_PERSONA_GENERATION, + ), // Serialize expensive background LLM work for burst-sensitive setups unless explicitly overridden }, ); @@ -241,7 +261,10 @@ export const graphResolutionWorker = new Worker( }, { connection: getRedisConnection(), - concurrency: env.BULLMQ_CONCURRENCY_GRAPH_RESOLUTION, // Graph resolution concurrency + concurrency: getBurstSafeBullmqConcurrency( + "BULLMQ_CONCURRENCY_GRAPH_RESOLUTION", + env.BULLMQ_CONCURRENCY_GRAPH_RESOLUTION, + ), // Proxy mode needs stricter background pacing unless explicitly overridden }, ); diff --git a/apps/webapp/app/components/chat-panel/global-chat-panel.client.tsx b/apps/webapp/app/components/chat-panel/global-chat-panel.client.tsx index ccc36052e..c52e7095c 100644 --- a/apps/webapp/app/components/chat-panel/global-chat-panel.client.tsx +++ b/apps/webapp/app/components/chat-panel/global-chat-panel.client.tsx @@ -25,6 +25,7 @@ interface GlobalChatPanelProps { onClose: () => void; models: LLMModel[]; integrationAccountMap: Record; + enableBurstSafeReplyRecovery?: boolean; } // Minimal chat input ā€” creates a conversation and hands off to ConversationView @@ -186,12 +187,14 @@ export function GlobalChatPanel({ onClose, models, integrationAccountMap, + enableBurstSafeReplyRecovery = false, }: GlobalChatPanelProps) { const { pinnedConversationId } = useChatPanel()!; const navigate = useNavigate(); const [activeConversation, setActiveConversation] = useState<{ conversationId: string; + status?: string; history: Array<{ id: string; userType: string; @@ -231,6 +234,7 @@ export function GlobalChatPanel({ ) => { setActiveConversation({ conversationId, + status: "pending", history: [ { id: historyId, @@ -256,6 +260,7 @@ export function GlobalChatPanel({ const conv = historyFetcher.data.conversation; setActiveConversation({ conversationId: pendingHistoryId.current, + status: conv.status, history: conv.ConversationHistory ?? [], }); pendingHistoryId.current = null; @@ -314,7 +319,9 @@ export function GlobalChatPanel({ conversationId={activeConversation.conversationId} history={activeConversation.history} autoRegenerate + enableBurstSafeFirstReplyRecovery={enableBurstSafeReplyRecovery} integrationAccountMap={integrationAccountMap} + conversationStatus={activeConversation.status} models={models} /> diff --git a/apps/webapp/app/components/conversation/conversation-view.client.tsx b/apps/webapp/app/components/conversation/conversation-view.client.tsx index 18e603fc1..0fe52acdf 100644 --- a/apps/webapp/app/components/conversation/conversation-view.client.tsx +++ b/apps/webapp/app/components/conversation/conversation-view.client.tsx @@ -1,5 +1,5 @@ import { useCallback, useEffect, useRef, useState } from "react"; -import { useFetcher } from "@remix-run/react"; +import { useFetcher, useRevalidator } from "@remix-run/react"; import { useLocalCommonState } from "~/hooks/use-local-state"; import { useChat, type UIMessage } from "@ai-sdk/react"; import { @@ -36,6 +36,8 @@ interface ConversationViewProps { integrationFrontendMap?: Record; /** When true, auto-triggers regenerate if history has only 1 message */ autoRegenerate?: boolean; + /** When true, add burst-safe latest-reply retries for strict proxy/self-hosted providers */ + enableBurstSafeFirstReplyRecovery?: boolean; /** DB conversation status ā€” input is disabled when "running" */ conversationStatus?: string; models?: LLMModel[]; @@ -48,13 +50,88 @@ export function ConversationView({ integrationAccountMap = {}, integrationFrontendMap = {}, autoRegenerate = false, + enableBurstSafeFirstReplyRecovery = false, conversationStatus, models: modelsProp = [], }: ConversationViewProps) { + const [runtimeKey, setRuntimeKey] = useState(0); + const historySignature = history + .map((entry) => `${entry.id}:${entry.userType}:${entry.createdAt ?? ""}`) + .join("|"); + const previousHistorySignatureRef = useRef(historySignature); + // Burst/provider recovery shim: in strict proxy/self-hosted modes a late stream + // failure can leave a persisted trailing user turn with no assistant reply. We + // allow a one-time runtime remount per persisted turn so useChat can recover + // without relying on a full page refresh. + const forcedRehydrateKeyRef = useRef(null); + + useEffect(() => { + if (previousHistorySignatureRef.current !== historySignature) { + previousHistorySignatureRef.current = historySignature; + setRuntimeKey((current) => current + 1); + } + }, [historySignature]); + + return ( + { + if (forcedRehydrateKeyRef.current === recoveryKey) { + return; + } + forcedRehydrateKeyRef.current = recoveryKey; + setRuntimeKey((current) => current + 1); + }} + /> + ); +} + +interface ConversationViewRuntimeProps extends ConversationViewProps { + onForceRehydrate?: (recoveryKey: string) => void; +} + +function ConversationViewRuntime({ + conversationId, + history, + className, + integrationAccountMap = {}, + integrationFrontendMap = {}, + autoRegenerate = false, + enableBurstSafeFirstReplyRecovery = false, + conversationStatus, + models: modelsProp = [], + onForceRehydrate, +}: ConversationViewRuntimeProps) { const readFetcher = useFetcher(); + const revalidator = useRevalidator(); const scrollContainerRef = useRef(null); const composerRef = useRef(null); const messageRefs = useRef<(HTMLDivElement | null)[]>([]); + const autoRegeneratedConversationRef = useRef(null); + const autoRetryTimeoutRef = useRef(null); + const stalledTurnWatchdogRef = useRef(null); + const autoRetryAttemptRef = useRef(0); + const regenerateRef = useRef<(() => void) | null>(null); + const clearErrorRef = useRef<(() => void) | null>(null); + const pendingUserTurnRef = useRef(false); + const [autoRetryStatus, setAutoRetryStatus] = useState<{ + nextDelayMs: number | null; + attempt: number; + exhausted: boolean; + }>({ + nextDelayMs: null, + attempt: 0, + exhausted: false, + }); // initialize to history.length so mount doesn't trigger the scroll effect const prevMessageCountRef = useRef(history.length); // spacer height = scroll container clientHeight so any message can scroll to top @@ -74,6 +151,57 @@ export function ConversationView({ const handleModelChange = (modelId: string) => { setSelectedModelId(modelId); }; + + const initialHistoryHasPendingUserTurn = + history.length > 0 && + history[history.length - 1]?.userType !== UserTypeEnum.Agent; + pendingUserTurnRef.current = initialHistoryHasPendingUserTurn; + + const scheduleAutoRetry = useCallback(() => { + if ( + !enableBurstSafeFirstReplyRecovery || + !pendingUserTurnRef.current + ) { + return; + } + + if (autoRetryTimeoutRef.current) { + return; + } + + const retryDelaysMs = [5000, 15000, 30000, 60000]; + const delayMs = retryDelaysMs[autoRetryAttemptRef.current]; + if (delayMs == null) { + setAutoRetryStatus({ + nextDelayMs: null, + attempt: autoRetryAttemptRef.current, + exhausted: true, + }); + return; + } + + autoRetryAttemptRef.current += 1; + setAutoRetryStatus({ + nextDelayMs: delayMs, + attempt: autoRetryAttemptRef.current, + exhausted: false, + }); + + if (autoRetryTimeoutRef.current) { + window.clearTimeout(autoRetryTimeoutRef.current); + } + + autoRetryTimeoutRef.current = window.setTimeout(() => { + autoRetryTimeoutRef.current = null; + setAutoRetryStatus({ + nextDelayMs: null, + attempt: autoRetryAttemptRef.current, + exhausted: false, + }); + clearErrorRef.current?.(); + regenerateRef.current?.(); + }, delayMs); + }, [enableBurstSafeFirstReplyRecovery]); // toolCallId ā†’ { approved, ...argOverrides } // Single ref for both approval decisions and arg overrides const toolArgOverridesRef = useRef>>( @@ -103,6 +231,7 @@ export function ConversationView({ sendMessage, messages, status, + clearError, stop, regenerate, addToolApprovalResponse, @@ -110,6 +239,16 @@ export function ConversationView({ id: conversationId, resume: true, onFinish: () => { + autoRetryAttemptRef.current = 0; + setAutoRetryStatus({ + nextDelayMs: null, + attempt: 0, + exhausted: false, + }); + if (autoRetryTimeoutRef.current) { + window.clearTimeout(autoRetryTimeoutRef.current); + autoRetryTimeoutRef.current = null; + } toolArgOverridesRef.current = {}; pendingApprovalRequestsRef.current = []; readFetcher.submit(null, { @@ -117,6 +256,9 @@ export function ConversationView({ action: `/api/v1/conversation/${conversationId}/read`, }); }, + onError: () => { + scheduleAutoRetry(); + }, messages: history.map( (h) => ({ @@ -153,12 +295,146 @@ export function ConversationView({ // recorded approve/decline decision in toolArgOverridesRef. sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses, }); + regenerateRef.current = regenerate; + clearErrorRef.current = clearError; + const latestMessageIsUser = + messages.length > 0 && messages[messages.length - 1]?.role === "user"; + pendingUserTurnRef.current = latestMessageIsUser; useEffect(() => { - if (autoRegenerate && history.length === 1 && conversationStatus !== "running") { + return () => { + if (autoRetryTimeoutRef.current) { + window.clearTimeout(autoRetryTimeoutRef.current); + } + if (stalledTurnWatchdogRef.current) { + window.clearTimeout(stalledTurnWatchdogRef.current); + } + }; + }, []); + + useEffect(() => { + // Burst/provider recovery shim: once a reply actually starts streaming again, + // clear any pending retry timer but preserve the current attempt count so the + // user-facing banner still reflects the in-flight recovery state. + if ( + (status === "submitted" || status === "streaming") && + autoRetryTimeoutRef.current + ) { + window.clearTimeout(autoRetryTimeoutRef.current); + autoRetryTimeoutRef.current = null; + setAutoRetryStatus({ + nextDelayMs: null, + attempt: autoRetryAttemptRef.current, + exhausted: false, + }); + } + }, [status]); + + useEffect(() => { + // Any completed assistant turn resets the retry/remount state for the next + // pending user turn. + if (!latestMessageIsUser) { + autoRetryAttemptRef.current = 0; + if (autoRetryTimeoutRef.current) { + window.clearTimeout(autoRetryTimeoutRef.current); + autoRetryTimeoutRef.current = null; + } + if (stalledTurnWatchdogRef.current) { + window.clearTimeout(stalledTurnWatchdogRef.current); + stalledTurnWatchdogRef.current = null; + } + setAutoRetryStatus({ + nextDelayMs: null, + attempt: 0, + exhausted: false, + }); + } + }, [latestMessageIsUser]); + + useEffect(() => { + if ( + autoRegenerate && + initialHistoryHasPendingUserTurn && + conversationStatus !== "running" + ) { + const recoveryKey = `${conversationId}:${history[history.length - 1]?.id ?? "pending"}`; + if (autoRegeneratedConversationRef.current === recoveryKey) { + return; + } + autoRegeneratedConversationRef.current = recoveryKey; regenerate(); } - }, []); + }, [ + autoRegenerate, + conversationId, + conversationStatus, + history, + initialHistoryHasPendingUserTurn, + regenerate, + ]); + + useEffect(() => { + if ( + !enableBurstSafeFirstReplyRecovery || + !latestMessageIsUser || + status !== "error" + ) { + return; + } + + clearError(); + scheduleAutoRetry(); + }, [ + clearError, + enableBurstSafeFirstReplyRecovery, + latestMessageIsUser, + scheduleAutoRetry, + status, + ]); + + useEffect(() => { + if ( + !enableBurstSafeFirstReplyRecovery || + !latestMessageIsUser || + status === "submitted" || + status === "streaming" + ) { + if (stalledTurnWatchdogRef.current) { + window.clearTimeout(stalledTurnWatchdogRef.current); + stalledTurnWatchdogRef.current = null; + } + return; + } + + if (stalledTurnWatchdogRef.current || autoRetryTimeoutRef.current) { + return; + } + + // Transport resilience shim only: if a provider leaves us with a persisted + // trailing user turn and no active stream, revalidate/remount once for that + // turn so the existing chat flow can retry. This is not the normal chat UX. + stalledTurnWatchdogRef.current = window.setTimeout(() => { + stalledTurnWatchdogRef.current = null; + const recoveryKey = `${conversationId}:${history[history.length - 1]?.id ?? "pending"}`; + revalidator.revalidate(); + onForceRehydrate?.(recoveryKey); + }, 2000); + + return () => { + if (stalledTurnWatchdogRef.current) { + window.clearTimeout(stalledTurnWatchdogRef.current); + stalledTurnWatchdogRef.current = null; + } + }; + }, [ + clearError, + enableBurstSafeFirstReplyRecovery, + latestMessageIsUser, + onForceRehydrate, + revalidator, + scheduleAutoRetry, + status, + ]); // Measure scroll container and keep spacer in sync so any message can reach the top useEffect(() => { @@ -293,6 +569,33 @@ export function ConversationView({ {(status === "streaming" || status === "submitted" || keepSpacer) && (
)} + {enableBurstSafeFirstReplyRecovery && + latestMessageIsUser && + (autoRetryStatus.nextDelayMs !== null || autoRetryStatus.exhausted) && ( +
+ {autoRetryStatus.nextDelayMs !== null + ? `Provider is rate-limiting the latest reply. Retrying automatically in ${Math.ceil(autoRetryStatus.nextDelayMs / 1000)}s (attempt ${autoRetryStatus.attempt}).` + : "Provider is still rate-limiting the latest reply. You can retry without refreshing."} + {autoRetryStatus.exhausted && ( + + )} +
+ )}
diff --git a/apps/webapp/app/jobs/conversation/create-title.logic.ts b/apps/webapp/app/jobs/conversation/create-title.logic.ts index 57931980d..833aed42b 100644 --- a/apps/webapp/app/jobs/conversation/create-title.logic.ts +++ b/apps/webapp/app/jobs/conversation/create-title.logic.ts @@ -3,6 +3,7 @@ import { conversationTitlePrompt } from "~/trigger/conversation/prompt"; import { prisma } from "~/db.server"; import { logger } from "~/services/logger.service"; import { makeStructuredModelCall } from "~/lib/model.server"; +import { runWithBurstRetry } from "~/services/agent/burst-retry.server"; export interface CreateConversationTitlePayload { @@ -34,18 +35,23 @@ export async function processConversationTitleCreation( select: { workspaceId: true }, }); - const { object } = await makeStructuredModelCall( - TitleSchema, - [ - { - role: "user", - content: conversationTitlePrompt.replace("{{message}}", payload.message), - }, - ], - "medium", - "conversationTitle", - undefined, - conversation?.workspaceId ?? undefined, + const { object } = await runWithBurstRetry("conversation.title", () => + makeStructuredModelCall( + TitleSchema, + [ + { + role: "user", + content: conversationTitlePrompt.replace( + "{{message}}", + payload.message, + ), + }, + ], + "medium", + "conversationTitle", + undefined, + conversation?.workspaceId ?? undefined, + ), ); const title = object.title?.replace(/<[^>]*>/g, "").replace(/\s+/g, " ").trim().slice(0, 80) || ""; diff --git a/apps/webapp/app/jobs/ingest/preprocess-episode.logic.ts b/apps/webapp/app/jobs/ingest/preprocess-episode.logic.ts index 1f890ee54..c5e18986b 100644 --- a/apps/webapp/app/jobs/ingest/preprocess-episode.logic.ts +++ b/apps/webapp/app/jobs/ingest/preprocess-episode.logic.ts @@ -22,6 +22,7 @@ import { prisma } from "~/db.server"; import { type SessionCompactionPayload } from "~/jobs/session/session-compaction.logic"; import { getRecentEpisodes } from "~/services/vectorStorage.server"; import { type EpisodeEmbedding } from "@prisma/client"; +import { getBurstSafeBackgroundDelayMs } from "~/services/llm-provider.server"; export { IngestBodyRequest }; @@ -73,7 +74,10 @@ export async function processEpisodePreprocessing( // Callback function for enqueueing ingestion jobs (one per chunk) enqueueIngestEpisode?: (params: IngestEpisodePayload) => Promise, // Callback function for enqueueing session compaction (for conversations) - enqueueSessionCompaction?: (params: SessionCompactionPayload) => Promise, + enqueueSessionCompaction?: ( + params: SessionCompactionPayload, + delayMs?: number, + ) => Promise, ): Promise { try { logger.info(`Preprocessing episode for user ${payload.userId}`, { @@ -495,12 +499,13 @@ export async function processEpisodePreprocessing( }); try { + const delayMs = !document ? getBurstSafeBackgroundDelayMs() : 0; await enqueueSessionCompaction({ userId: payload.userId, sessionId, source: episodeBody.source, workspaceId: payload.workspaceId, - }); + }, delayMs); } catch (compactionError) { // Don't fail preprocessing if compaction enqueueing fails logger.warn(`Failed to enqueue session compaction`, { diff --git a/apps/webapp/app/jobs/spaces/aspect-persona-generation.ts b/apps/webapp/app/jobs/spaces/aspect-persona-generation.ts index eb4892264..65c4fc500 100644 --- a/apps/webapp/app/jobs/spaces/aspect-persona-generation.ts +++ b/apps/webapp/app/jobs/spaces/aspect-persona-generation.ts @@ -25,24 +25,33 @@ import { } from "@core/types"; import { ProviderFactory } from "@core/providers"; import { getActiveVoiceAspects } from "~/services/aspectStore.server"; +import { runWithBurstRetry } from "~/services/agent/burst-retry.server"; +import { isBurstSensitiveChatProvider } from "~/services/llm-provider.server"; -import { createAgent, resolveModelString } from "~/lib/model.server"; +import { makeModelCall } from "~/lib/model.server"; import { type ModelMessage } from "ai"; -import { type MessageListInput } from "@mastra/core/agent/message-list"; /** * Direct LLM call helper ā€” replaces batch for single/few requests. * Returns the text content from a single prompt. */ async function directLLMCall( - prompt: MessageListInput, + prompt: ModelMessage, label?: string, ): Promise { try { - const modelId = await resolveModelString("chat", "medium"); - const agent = createAgent(modelId); - const result = await agent.generate(prompt); - const text = result.text; + const text = await runWithBurstRetry( + `persona.${label ?? "direct-call"}`, + () => + makeModelCall( + false, + [prompt], + () => {}, + undefined, + "medium", + label ? `persona-${label}` : "persona-direct-call", + ) as Promise, + ); logger.info(`Direct LLM call completed${label ? ` [${label}]` : ""}`, { responseLength: text.length, preview: text.slice(0, 100), @@ -382,6 +391,60 @@ interface ChunkData { isLatest: boolean; } +function sortAspectData(aspectData: AspectData): AspectData { + return { + ...aspectData, + statements: [...aspectData.statements].sort( + (a, b) => b.createdAt.getTime() - a.createdAt.getTime(), + ), + episodes: [...aspectData.episodes].sort( + (a, b) => b.createdAt.getTime() - a.createdAt.getTime(), + ), + }; +} + +function buildSectionResult( + aspectData: AspectData, + content: string, +): PersonaSectionResult { + const sectionInfo = ASPECT_SECTION_MAP[aspectData.aspect]; + return { + aspect: aspectData.aspect, + title: sectionInfo.title, + content, + statementCount: aspectData.statements.length, + episodeCount: aspectData.episodes.length, + }; +} + +function extractResponseContent(response: unknown): string { + return typeof response === "string" + ? response + : ((response as { content?: string } | null)?.content ?? ""); +} + +function getSectionContent( + aspect: StatementAspect, + response: unknown, +): string | null { + const content = extractResponseContent(response); + + if (!content || content.includes("INSUFFICIENT_DATA")) { + logger.info(`${aspect} section returned INSUFFICIENT_DATA`); + return null; + } + + return content; +} + +function getSectionResult( + aspectData: AspectData, + response: unknown, +): PersonaSectionResult | null { + const content = getSectionContent(aspectData.aspect, response); + return content ? buildSectionResult(aspectData, content) : null; +} + /** * Fetch statements grouped by aspect with their provenance episodes. * Also fetches voice aspects from the Aspects Store and merges them @@ -817,6 +880,38 @@ async function generateSectionWithChunking( // Split into chunks const chunks = chunkAspectData(aspectData); + const burstSensitive = isBurstSensitiveChatProvider(); + + if (burstSensitive) { + const chunkSummaries: string[] = []; + + for (const chunk of chunks) { + const content = await directLLMCall( + buildChunkSummaryPrompt(aspect, chunk, userContext), + `chunk-${aspect}-${chunk.chunkIndex}`, + ); + + if (content && !content.includes("NO_PATTERNS")) { + chunkSummaries.push(content); + } + } + + if (chunkSummaries.length === 0) { + logger.info(`No patterns found in any chunk for ${aspect}`); + return "INSUFFICIENT_DATA"; + } + + if (chunkSummaries.length === 1) { + return chunkSummaries[0]; + } + + return ( + (await directLLMCall( + buildMergePrompt(aspect, chunkSummaries, userContext), + `merge-${aspect}`, + )) ?? chunkSummaries[0] + ); + } // Generate summary for each chunk via batch const chunkRequests = chunks.map((chunk) => ({ @@ -844,10 +939,7 @@ async function generateSectionWithChunking( for (const result of chunkBatch.results) { if (result.error || !result.response) continue; - const content = - typeof result.response === "string" - ? result.response - : result.response.content || ""; + const content = extractResponseContent(result.response); if (!content.includes("NO_PATTERNS")) { chunkSummaries.push(content); @@ -895,6 +987,68 @@ async function generateSectionWithChunking( : mergeResult.response.content || chunkSummaries[0]; } +/** + * Generate a single aspect section + */ +async function generateAspectSection( + aspectData: AspectData, + userContext: UserContext, +): Promise { + const { aspect, statements, episodes } = aspectData; + const sectionInfo = ASPECT_SECTION_MAP[aspect]; + + if (!sectionInfo) { + logger.warn(`No section mapping for aspect "${aspect}" ā€” skipping`); + return null; + } + + // Skip if insufficient data + if (statements.length < MIN_STATEMENTS_PER_SECTION) { + logger.info(`Skipping ${aspect} section - insufficient data`, { + statementCount: statements.length, + minRequired: MIN_STATEMENTS_PER_SECTION, + }); + return null; + } + + const prompt = buildAspectSectionPrompt(aspectData, userContext); + const burstSensitive = isBurstSensitiveChatProvider(); + + if (burstSensitive) { + const content = await directLLMCall(prompt, `section-${aspect}`); + return content ? getSectionResult(aspectData, content) : null; + } + + const batchRequest = { + customId: `persona-section-${aspect}-${Date.now()}`, + messages: [prompt], + systemPrompt: "", + }; + + const { batchId } = await createBatch({ + requests: [batchRequest], + outputSchema: SectionContentSchema, + maxRetries: 3, + timeoutMs: 1200000, + }); + + // Poll for completion + const batch = await pollBatchCompletion(batchId, 1200000); + + if (!batch.results || batch.results.length === 0) { + logger.warn(`No results for ${aspect} section`); + return null; + } + + const result = batch.results[0]; + if (result.error || !result.response) { + logger.warn(`Error generating ${aspect} section`, { error: result.error }); + return null; + } + + return getSectionResult(aspectData, result.response); +} + /** * Generate all aspect sections in parallel batches */ @@ -950,22 +1104,15 @@ async function generateAllAspectSections( // Run all sections in parallel - large (chunked) and small (single batch) concurrently const parallelTasks: Promise[] = []; + const burstSensitive = isBurstSensitiveChatProvider(); // Task for each large section (chunking handled internally) for (const aspectData of largeAspects) { parallelTasks.push( generateSectionWithChunking(aspectData, userContext).then((content) => { if (content && !content.includes("INSUFFICIENT_DATA")) { - const sectionInfo = ASPECT_SECTION_MAP[aspectData.aspect]; - return [ - { - aspect: aspectData.aspect, - title: sectionInfo.title, - content, - statementCount: aspectData.statements.length, - episodeCount: aspectData.episodes.length, - }, - ]; + const section = getSectionResult(aspectData, content); + return section ? [section] : []; } return []; }), @@ -974,82 +1121,81 @@ async function generateAllAspectSections( // Task for all small sections in a single batch if (smallAspects.length > 0) { - parallelTasks.push( - (async () => { - const sortedSmallAspects = smallAspects.map((aspectData) => ({ - ...aspectData, - statements: [...aspectData.statements].sort( - (a, b) => b.createdAt.getTime() - a.createdAt.getTime(), - ), - episodes: [...aspectData.episodes].sort( - (a, b) => b.createdAt.getTime() - a.createdAt.getTime(), - ), - })); - - const batchRequests = sortedSmallAspects.map((aspectData) => { - const prompt = buildAspectSectionPrompt(aspectData, userContext); - return { - customId: `persona-section-${aspectData.aspect}-${Date.now()}`, - messages: [prompt], - systemPrompt: "", - }; - }); - - logger.info( - `Generating ${batchRequests.length} small persona sections in batch`, - { - aspects: sortedSmallAspects.map((a) => a.aspect), - }, - ); - - const { batchId } = await createBatch({ - requests: batchRequests, - outputSchema: SectionContentSchema, - maxRetries: 3, - timeoutMs: 1200000, - }); + const sortedSmallAspects = smallAspects.map(sortAspectData); - const batch = await pollBatchCompletion(batchId, 1200000); - const results: PersonaSectionResult[] = []; + if (burstSensitive) { + parallelTasks.push( + (async () => { + const results: PersonaSectionResult[] = []; - if (batch.results && batch.results.length > 0) { - for (let i = 0; i < batch.results.length; i++) { - const result = batch.results[i]; - const aspectData = sortedSmallAspects[i]; - const sectionInfo = ASPECT_SECTION_MAP[aspectData.aspect]; + logger.info( + `Generating ${sortedSmallAspects.length} small persona sections sequentially for burst-sensitive provider`, + { + aspects: sortedSmallAspects.map((a) => a.aspect), + }, + ); - if (result.error || !result.response) { - logger.warn(`Error generating ${aspectData.aspect} section`, { - error: result.error, - }); - continue; + for (const aspectData of sortedSmallAspects) { + const section = await generateAspectSection(aspectData, userContext); + if (section) { + results.push(section); } + } - const content = - typeof result.response === "string" - ? result.response - : result.response.content || ""; - - if (content.includes("INSUFFICIENT_DATA")) { - logger.info( - `${aspectData.aspect} section returned INSUFFICIENT_DATA`, - ); - continue; + return results; + })(), + ); + } else { + parallelTasks.push( + (async () => { + const batchRequests = sortedSmallAspects.map((aspectData) => { + const prompt = buildAspectSectionPrompt(aspectData, userContext); + return { + customId: `persona-section-${aspectData.aspect}-${Date.now()}`, + messages: [prompt], + systemPrompt: "", + }; + }); + + logger.info( + `Generating ${batchRequests.length} small persona sections in batch`, + { + aspects: sortedSmallAspects.map((a) => a.aspect), + }, + ); + + const { batchId } = await createBatch({ + requests: batchRequests, + outputSchema: SectionContentSchema, + maxRetries: 3, + timeoutMs: 1200000, + }); + + const batch = await pollBatchCompletion(batchId, 1200000); + const results: PersonaSectionResult[] = []; + + if (batch.results && batch.results.length > 0) { + for (let i = 0; i < batch.results.length; i++) { + const result = batch.results[i]; + const aspectData = sortedSmallAspects[i]; + if (result.error || !result.response) { + logger.warn(`Error generating ${aspectData.aspect} section`, { + error: result.error, + }); + continue; + } + + const section = getSectionResult(aspectData, result.response); + if (section) { + results.push(section); + } } - - results.push({ - aspect: aspectData.aspect, - title: sectionInfo.title, - content, - statementCount: aspectData.statements.length, - episodeCount: aspectData.episodes.length, - }); } - } - return results; - })(), - ); + return results; + })(), + ); + } } // Wait for all tasks to complete in parallel diff --git a/apps/webapp/app/lib/ingest.server.ts b/apps/webapp/app/lib/ingest.server.ts index e2e43348b..22340ef61 100644 --- a/apps/webapp/app/lib/ingest.server.ts +++ b/apps/webapp/app/lib/ingest.server.ts @@ -18,6 +18,7 @@ export const addToQueue = async ( workspaceId: string, activityId?: string, ingestionQueueId?: string, + enqueueDelayMs?: number, ) => { const body = { ...rawBody, source: rawBody.source.toLowerCase() }; const sessionId = body.sessionId || crypto.randomUUID(); @@ -133,6 +134,7 @@ export const addToQueue = async ( queueId: queuePersist.id, }, rawBody.delay, + enqueueDelayMs, ); // Track feature usage diff --git a/apps/webapp/app/lib/model.server.ts b/apps/webapp/app/lib/model.server.ts index b9a0a8a49..034a5927c 100644 --- a/apps/webapp/app/lib/model.server.ts +++ b/apps/webapp/app/lib/model.server.ts @@ -1,4 +1,9 @@ -import { embed, type ModelMessage } from "ai"; +import { + embed, + generateText, + streamText, + type ModelMessage, +} from "ai"; import { type z } from "zod"; import { createOpenAI, @@ -117,21 +122,23 @@ export const getModel = (takeModel?: string) => { const modelId = getModelId(model); // Ollama: use direct AI SDK provider (needs custom URL) - if (provider === "ollama" || getDefaultChatProviderType() === "ollama") { + if (provider === "ollama") { + return getDirectModelInstance(model); + } + if (getDefaultChatProviderType() === "ollama") { const ollamaConfig = getProviderConfig("ollama"); const ollamaUrl = ollamaConfig.baseUrl; if (!ollamaUrl) { throw new Error("Ollama provider selected but no baseUrl configured."); } - if (!modelId) { - throw new Error("No chat model configured for Ollama."); - } - const ollama = createOllama({ baseURL: ollamaUrl }); - return ollama(modelId); + return createOllama({ baseURL: ollamaUrl })(modelId); } // Azure: use direct AI SDK provider (needs base URL + API key) - if (provider === "azure" || getDefaultChatProviderType() === "azure") { + if (provider === "azure") { + return getDirectModelInstance(model); + } + if (getDefaultChatProviderType() === "azure") { const azureConfig = getProviderConfig("azure"); const baseURL = azureConfig.baseUrl; if (!baseURL) { @@ -141,27 +148,13 @@ export const getModel = (takeModel?: string) => { if (!apiKey) { throw new Error("Azure provider selected but AZURE_API_KEY is not configured."); } - const azureClient = createAzure({ baseURL, apiKey }); - return azureClient(modelId); + return createAzure({ baseURL, apiKey })(modelId); } // OpenAI proxy: use direct AI SDK provider (needs custom base URL) const openaiConfig = getProviderConfig("openai"); if (provider === "openai" && openaiConfig.baseUrl) { - const openaiKey = resolveApiKey("openai"); - if (!openaiKey) { - throw new Error( - "OpenAI proxy configured but OPENAI_API_KEY is missing.", - ); - } - const openaiClient = createOpenAI({ - baseURL: openaiConfig.baseUrl, - apiKey: openaiKey, - }); - const apiMode = openaiConfig.apiMode ?? "responses"; - return apiMode === "chat_completions" - ? openaiClient.chat(modelId) - : openaiClient.responses(modelId); + return getDirectModelInstance(model); } // All other providers: use Mastra model router @@ -227,6 +220,77 @@ function logTokenUsage(prefix: string, model: string, tokenUsage: TokenUsage | u ); } +export function shouldBypassMastraAgent( + modelString: string, + resolvedBaseUrl?: string, +): boolean { + const provider = getProvider(modelString); + const providerConfig = getProviderConfig(provider); + const effectiveBaseUrl = resolvedBaseUrl ?? providerConfig.baseUrl; + + return ( + provider === "ollama" || + provider === "azure" || + (provider === "openai" && !!effectiveBaseUrl) + ); +} + +export function getDirectModelInstance( + modelString: string, + options?: { apiKey?: string; baseUrl?: string }, +) { + const provider = getProvider(modelString); + const modelId = getModelId(modelString); + + if (provider === "ollama") { + const baseURL = + options?.baseUrl ?? options?.apiKey ?? getProviderConfig("ollama").baseUrl; + if (!baseURL) { + throw new Error("Ollama provider selected but no baseUrl configured."); + } + return createOllama({ baseURL })(modelId); + } + + if (provider === "azure") { + const baseURL = options?.baseUrl ?? getProviderConfig("azure").baseUrl; + const apiKey = options?.apiKey ?? resolveApiKey("azure"); + if (!baseURL) { + throw new Error( + "Azure provider selected but AZURE_BASE_URL is not configured.", + ); + } + if (!apiKey) { + throw new Error( + "Azure provider selected but AZURE_API_KEY is not configured.", + ); + } + return createAzure({ baseURL, apiKey })(modelId); + } + + if (provider === "openai") { + const openaiConfig = getProviderConfig("openai"); + const baseURL = options?.baseUrl ?? openaiConfig.baseUrl; + const apiKey = options?.apiKey ?? resolveApiKey("openai"); + if (!baseURL) { + throw new Error( + "OpenAI direct model requested but no proxy baseUrl is configured.", + ); + } + if (!apiKey) { + throw new Error( + "OpenAI proxy configured but OPENAI_API_KEY is missing.", + ); + } + const openaiClient = createOpenAI({ baseURL, apiKey }); + const apiMode = openaiConfig.apiMode ?? "responses"; + return apiMode === "chat_completions" + ? openaiClient.chat(modelId) + : openaiClient.responses(modelId); + } + + throw new Error(`Direct model instance not supported for provider: ${provider}`); +} + // --------------------------------------------------------------------------- // Mastra Agent factory // --------------------------------------------------------------------------- @@ -328,6 +392,38 @@ export async function makeModelCall( ); const agentOptions = apiKey ? { apiKey, ...(baseUrl && { baseUrl }) } : undefined; + + if (shouldBypassMastraAgent(model, agentOptions?.baseUrl)) { + const modelInstance = getDirectModelInstance(model, agentOptions) as any; + + if (stream) { + return streamText({ + model: modelInstance, + messages, + ...options, + ...(providerOptions && { providerOptions }), + onFinish: async ({ text, usage }) => { + const tokenUsage = toTokenUsage(usage); + logTokenUsage(complexity.toUpperCase(), model, tokenUsage); + onFinish(text, model, tokenUsage); + }, + }); + } + + const result = await generateText({ + model: modelInstance, + messages, + ...options, + ...(providerOptions && { providerOptions }), + }); + + const tokenUsage = toTokenUsage(result.usage); + logTokenUsage(complexity.toUpperCase(), model, tokenUsage); + onFinish(result.text, model, tokenUsage); + + return result.text; + } + const agent = createAgent(model, undefined, undefined, agentOptions); if (stream) { @@ -378,14 +474,6 @@ function tryParseJsonFromText(raw: string): unknown | undefined { } } -function needsTolerantParsing(): boolean { - const openaiConfig = getProviderConfig("openai"); - const apiMode = openaiConfig.apiMode ?? "responses"; - const isProxyChatMode = apiMode === "chat_completions" && !!openaiConfig.baseUrl; - const isOllama = getDefaultChatProviderType() === "ollama"; - return isProxyChatMode || isOllama; -} - export async function makeStructuredModelCall( schema: T, messages: ModelMessage[], @@ -406,13 +494,13 @@ export async function makeStructuredModelCall( const agentApiOptions = apiKey ? { apiKey, ...(baseUrl && { baseUrl }) } : undefined; // Proxy/Ollama: manual JSON extraction (no structured output support) - if (needsTolerantParsing()) { + if (shouldBypassMastraAgent(model, agentApiOptions?.baseUrl)) { const { object, usage } = await structuredCallWithTolerantParsing( schema, messages, model, temperature, - apiKey, + agentApiOptions, ); const tokenUsage = toTokenUsage(usage); logTokenUsage(`Structured/${complexity.toUpperCase()}`, model, tokenUsage); @@ -459,19 +547,23 @@ async function structuredCallWithTolerantParsing( messages: ModelMessage[], modelString: string, temperature?: number, - apiKey?: string, + modelOptions?: { apiKey?: string; baseUrl?: string }, ): Promise<{ object: z.infer; usage: any }> { - const jsonPreamble = - "Return ONLY a single valid JSON object that matches the requested schema. " + - "Do not wrap it in Markdown fences. Do not include extra text. " + - "Include every required key; use null for nullable fields; use [] for empty arrays."; - - const agentOpts = apiKey ? { apiKey } : undefined; - const agent = createAgent(modelString, jsonPreamble, undefined, agentOpts); - - const textResult = await agent.generate(messages as any, { + const modelInstance = getDirectModelInstance(modelString, modelOptions) as any; + const textResult = await generateText({ + model: modelInstance, + messages: [ + { + role: "system", + content: + "Return ONLY a single valid JSON object that matches the requested schema. " + + "Do not wrap it in Markdown fences. Do not include extra text. " + + "Include every required key; use null for nullable fields; use [] for empty arrays.", + }, + ...messages, + ], ...(temperature !== undefined && { temperature }), - } as any); + }); const parsed = tryParseJsonFromText(textResult.text); const validated = parsed ? schema.safeParse(parsed) : undefined; @@ -480,18 +572,19 @@ async function structuredCallWithTolerantParsing( } // Repair attempt - const repairAgent = createAgent( - modelString, - "You are a JSON repair assistant. Convert the user's content into a single valid JSON object. " + - "Return ONLY the JSON object, with no Markdown fences and no extra text.", - undefined, - agentOpts, - ); - - const repairResult = await repairAgent.generate( - [{ role: "user", content: textResult.text }] as any, - { temperature: 0 } as any, - ); + const repairResult = await generateText({ + model: modelInstance, + messages: [ + { + role: "system", + content: + "You are a JSON repair assistant. Convert the user's content into a single valid JSON object. " + + "Return ONLY the JSON object, with no Markdown fences and no extra text.", + }, + { role: "user", content: textResult.text }, + ], + temperature: 0, + }); const repairedParsed = tryParseJsonFromText(repairResult.text); const repairedValidated = repairedParsed ? schema.safeParse(repairedParsed) : undefined; diff --git a/apps/webapp/app/lib/queue-adapter.server.ts b/apps/webapp/app/lib/queue-adapter.server.ts index dc2293543..257badc5b 100644 --- a/apps/webapp/app/lib/queue-adapter.server.ts +++ b/apps/webapp/app/lib/queue-adapter.server.ts @@ -35,8 +35,10 @@ export type QueueProvider = "trigger" | "bullmq"; export async function enqueuePreprocessEpisode( payload: IngestEpisodePayload, delay?: boolean, + delayMs?: number, ): Promise<{ id?: string; token?: string }> { const provider = env.QUEUE_PROVIDER as QueueProvider; + const resolvedDelayMs = delayMs ?? (delay ? 5 * 60 * 1000 : undefined); if (provider === "trigger") { const { preprocessTask } = @@ -45,7 +47,9 @@ export async function enqueuePreprocessEpisode( queue: "preprocessing-queue", concurrencyKey: payload.userId, tags: [payload.userId, payload.queueId], - delay: delay ? "5m" : undefined, + delay: resolvedDelayMs + ? `${Math.ceil(resolvedDelayMs / 1000)}s` + : undefined, }); return { id: handler.id, token: handler.publicAccessToken }; } else { @@ -55,7 +59,7 @@ export async function enqueuePreprocessEpisode( jobId: payload.queueId, attempts: 3, backoff: { type: "exponential", delay: 2000 }, - delay: delay ? 5 * 60 * 1000 : undefined, // 5 minutes in milliseconds + delay: resolvedDelayMs, }); return { id: job.id }; } @@ -94,13 +98,16 @@ export async function enqueueIngestEpisode( */ export async function enqueueCreateConversationTitle( payload: CreateConversationTitlePayload, + delayMs?: number, ): Promise<{ id?: string }> { const provider = env.QUEUE_PROVIDER as QueueProvider; if (provider === "trigger") { const { createConversationTitle } = await import("~/trigger/conversation/create-conversation-title"); - const handler = await createConversationTitle.trigger(payload); + const handler = await createConversationTitle.trigger(payload, { + ...(delayMs ? { delay: `${Math.ceil(delayMs / 1000)}s` } : {}), + }); return { id: handler.id }; } else { // BullMQ @@ -111,6 +118,7 @@ export async function enqueueCreateConversationTitle( { attempts: 3, backoff: { type: "exponential", delay: 2000 }, + ...(delayMs ? { delay: delayMs } : {}), }, ); return { id: job.id }; @@ -122,13 +130,16 @@ export async function enqueueCreateConversationTitle( */ export async function enqueueSessionCompaction( payload: SessionCompactionPayload, + delayMs?: number, ): Promise<{ id?: string }> { const provider = env.QUEUE_PROVIDER as QueueProvider; if (provider === "trigger") { const { triggerSessionCompaction } = await import("~/trigger/session/session-compaction"); - const handler = await triggerSessionCompaction(payload); + const handler = await triggerSessionCompaction(payload, { + ...(delayMs ? { delay: `${Math.ceil(delayMs / 1000)}s` } : {}), + }); return { id: handler.id }; } else { // BullMQ @@ -139,6 +150,7 @@ export async function enqueueSessionCompaction( { attempts: 3, backoff: { type: "exponential", delay: 2000 }, + ...(delayMs ? { delay: delayMs } : {}), }, ); return { id: job.id }; diff --git a/apps/webapp/app/routes/api.v1.conversation._index.tsx b/apps/webapp/app/routes/api.v1.conversation._index.tsx index f805c0535..8c41c2564 100644 --- a/apps/webapp/app/routes/api.v1.conversation._index.tsx +++ b/apps/webapp/app/routes/api.v1.conversation._index.tsx @@ -10,6 +10,7 @@ import { import { toRouterString } from "~/lib/model.server"; import { getDefaultChatModelId, + getBurstSafeBackgroundDelayMs, resolveModelConfig, } from "~/services/llm-provider.server"; import { UserTypeEnum } from "@core/types"; @@ -23,6 +24,7 @@ import { streamToUIResponse, drainAgentResult, } from "~/services/agent/mastra-stream.server"; +import { runWithBurstRetry } from "~/services/agent/burst-retry.server"; import { InputProcessor, type OutputProcessor, @@ -85,10 +87,11 @@ const { loader, action } = createHybridActionApiRoute( // ----------------------------------------------------------------------- if (!isAssistantApproval) { if (conversationHistory.length === 1 && incomingUserText) { + const delayMs = getBurstSafeBackgroundDelayMs(); await enqueueCreateConversationTitle({ conversationId: body.id, message: incomingUserText, - }); + }, delayMs); } const messageParts = normalizeParts(body.message?.parts); @@ -338,15 +341,17 @@ const { loader, action } = createHybridActionApiRoute( // Belt-and-suspenders: also fire if request.signal ever works request.signal.addEventListener("abort", cancelStream); - const stream = await agent.stream(modelMessages, { - toolsets: { core: tools }, - runId: body.id, - stopWhen: [stepCountIs(10)], - toolCallConcurrency: 1, - outputProcessors: [messageHistoryProcessor as OutputProcessor], - modelSettings: { temperature: 0.5 }, - abortSignal: abortController.signal, - }); + const stream = await runWithBurstRetry("conversation.stream", () => + agent.stream(modelMessages, { + toolsets: { core: tools }, + runId: body.id, + stopWhen: [stepCountIs(10)], + toolCallConcurrency: 1, + outputProcessors: [messageHistoryProcessor as OutputProcessor], + modelSettings: { temperature: 0.5 }, + abortSignal: abortController.signal, + }), + ); return streamToUIResponse(stream, cancelStream); }, diff --git a/apps/webapp/app/routes/home.conversation.$conversationId.tsx b/apps/webapp/app/routes/home.conversation.$conversationId.tsx index 486fc678d..e5840950f 100644 --- a/apps/webapp/app/routes/home.conversation.$conversationId.tsx +++ b/apps/webapp/app/routes/home.conversation.$conversationId.tsx @@ -13,7 +13,10 @@ import { deleteConversation, } from "~/services/conversation.server"; import { getIntegrationAccounts } from "~/services/integrationAccount.server"; -import { getAvailableModels } from "~/services/llm-provider.server"; +import { + getAvailableModels, + isBurstSensitiveChatProvider, +} from "~/services/llm-provider.server"; import { ConversationView } from "~/components/conversation"; import { useTypedLoaderData } from "remix-typedjson"; @@ -32,6 +35,7 @@ export async function loader({ params, request }: LoaderFunctionArgs) { getIntegrationAccounts(user.id, workspaceId), getAvailableModels(), ]); + const enableBurstSafeFirstReplyRecovery = isBurstSensitiveChatProvider(); const models = allModels .filter( @@ -46,7 +50,12 @@ export async function loader({ params, request }: LoaderFunctionArgs) { })); if (!conversation) { - return { conversation: null, integrationAccountMap: {}, models }; + return { + conversation: null, + integrationAccountMap: {}, + models, + enableBurstSafeFirstReplyRecovery, + }; } if (conversation.unread) { @@ -62,7 +71,13 @@ export async function loader({ params, request }: LoaderFunctionArgs) { } } - return { conversation, integrationAccountMap, integrationFrontendMap, models }; + return { + conversation, + integrationAccountMap, + integrationFrontendMap, + models, + enableBurstSafeFirstReplyRecovery, + }; } export async function action({ params, request }: ActionFunctionArgs) { @@ -72,7 +87,13 @@ export async function action({ params, request }: ActionFunctionArgs) { } export default function SingleConversation() { - const { conversation, integrationAccountMap, integrationFrontendMap, models } = + const { + conversation, + integrationAccountMap, + integrationFrontendMap, + models, + enableBurstSafeFirstReplyRecovery, + } = useTypedLoaderData(); const { conversationId } = useParams(); @@ -96,6 +117,7 @@ export default function SingleConversation() { conversationStatus={conversation.status} models={models} autoRegenerate + enableBurstSafeFirstReplyRecovery={enableBurstSafeFirstReplyRecovery} /> ); diff --git a/apps/webapp/app/routes/home.conversation.tsx b/apps/webapp/app/routes/home.conversation.tsx index 2b37ecafe..7c3cc5bf6 100644 --- a/apps/webapp/app/routes/home.conversation.tsx +++ b/apps/webapp/app/routes/home.conversation.tsx @@ -16,6 +16,9 @@ import { createConversation, CreateConversationSchema, } from "~/services/conversation.server"; +import { noStreamProcess } from "~/services/agent/no-stream-process"; +import { isBurstSensitiveChatProvider } from "~/services/llm-provider.server"; +import { logger } from "~/services/logger.service"; import { ResizablePanelGroup, ResizablePanel, @@ -71,6 +74,37 @@ export async function action({ request }: ActionFunctionArgs) { const conversationId = conversation?.conversationId; + const shouldPrecomputeFirstReply = + !!conversationId && + !submission.value.conversationId && + !submission.value.panelMode && + isBurstSensitiveChatProvider(); + + if (shouldPrecomputeFirstReply) { + // Keep the default client autoRegenerate flow for normal providers. Proxy/self-hosted + // setups can fail after the initial stream starts (late 429s), so we precompute here. + try { + await noStreamProcess( + { + id: conversationId, + modelId: submission.value.modelId, + message: { + parts: [{ text: submission.value.message, type: "text" }], + role: "user", + }, + source: submission.value.source ?? "core", + }, + userId, + workspace?.id as string, + ); + } catch (error) { + logger.warn("Burst-safe first reply precompute failed; falling back to client retry path", { + conversationId, + error: error instanceof Error ? error.message : String(error), + }); + } + } + if (submission.value.panelMode) { return json({ conversation, conversationId }); } diff --git a/apps/webapp/app/routes/home.tsx b/apps/webapp/app/routes/home.tsx index 71461f16a..b21df5f03 100644 --- a/apps/webapp/app/routes/home.tsx +++ b/apps/webapp/app/routes/home.tsx @@ -30,7 +30,10 @@ import { } from "~/components/chat-panel/chat-panel-context"; import React from "react"; import { getIntegrationAccounts } from "~/services/integrationAccount.server"; -import { getAvailableModels } from "~/services/llm-provider.server"; +import { + getAvailableModels, + isBurstSensitiveChatProvider, +} from "~/services/llm-provider.server"; import { type LLMModel } from "~/components/conversation"; import { tinykeys } from "tinykeys"; @@ -84,6 +87,7 @@ export const loader = async ({ request }: LoaderFunctionArgs) => { getIntegrationAccounts(user.id, workspace?.id as string), getAvailableModels(), ]); + const enableBurstSafeReplyRecovery = isBurstSensitiveChatProvider(); const models = allModels .filter( @@ -117,6 +121,7 @@ export const loader = async ({ request }: LoaderFunctionArgs) => { models, integrationAccountMap, integrationFrontendMap, + enableBurstSafeReplyRecovery, }, { headers: { @@ -136,6 +141,7 @@ function HomeInner({ needsButlerName, models, integrationAccountMap, + enableBurstSafeReplyRecovery, }: { conversationSources: any; workspace: any; @@ -145,6 +151,7 @@ function HomeInner({ needsButlerName: boolean; models: LLMModel[]; integrationAccountMap: Record; + enableBurstSafeReplyRecovery: boolean; }) { const { panelRef, closeChat, onPanelCollapse, chatOpen, toggleChat } = useChatPanel()!; @@ -226,6 +233,7 @@ function HomeInner({ onClose={closeChat} models={models} integrationAccountMap={integrationAccountMap} + enableBurstSafeReplyRecovery={enableBurstSafeReplyRecovery} /> )} @@ -238,8 +246,13 @@ function HomeInner({ } export default function Home() { - const { conversationSources, workspace, models, integrationAccountMap } = - useLoaderData() as any; + const { + conversationSources, + workspace, + models, + integrationAccountMap, + enableBurstSafeReplyRecovery, + } = useLoaderData() as any; const meta = (workspace?.metadata ?? {}) as Record; const needsButlerName = !meta.onboardingV2Complete; const accentColor = (meta.accentColor as string) || "#c87844"; @@ -257,6 +270,7 @@ export default function Home() { needsButlerName={needsButlerName} models={models} integrationAccountMap={integrationAccountMap} + enableBurstSafeReplyRecovery={enableBurstSafeReplyRecovery} /> diff --git a/apps/webapp/app/services/agent/burst-retry.server.ts b/apps/webapp/app/services/agent/burst-retry.server.ts new file mode 100644 index 000000000..68e784ac9 --- /dev/null +++ b/apps/webapp/app/services/agent/burst-retry.server.ts @@ -0,0 +1,58 @@ +import { logger } from "~/services/logger.service"; +import { isBurstSensitiveChatProvider } from "~/services/llm-provider.server"; + +const BURST_RETRY_DELAYS_MS = [5000, 15000]; + +function isBurstRetryableError(error: unknown): boolean { + if (!error || typeof error !== "object") return false; + + const candidate = error as { + statusCode?: number; + responseBody?: string; + cause?: unknown; + }; + + if (candidate.statusCode === 429) return true; + if ( + typeof candidate.responseBody === "string" && + candidate.responseBody.includes('"code":"1305"') + ) { + return true; + } + + return isBurstRetryableError(candidate.cause); +} + +export async function runWithBurstRetry( + label: string, + operation: () => Promise, +): Promise { + if (!isBurstSensitiveChatProvider()) { + return operation(); + } + + let attempt = 0; + + while (true) { + try { + return await operation(); + } catch (error) { + if ( + !isBurstRetryableError(error) || + attempt >= BURST_RETRY_DELAYS_MS.length + ) { + throw error; + } + + const delayMs = BURST_RETRY_DELAYS_MS[attempt]; + attempt += 1; + + logger.warn(`[burst-retry] ${label} hit 429/overload; retrying`, { + attempt, + delayMs, + }); + + await new Promise((resolve) => setTimeout(resolve, delayMs)); + } + } +} diff --git a/apps/webapp/app/services/agent/mastra-stream.server.ts b/apps/webapp/app/services/agent/mastra-stream.server.ts index 58d29a8d0..54e44572d 100644 --- a/apps/webapp/app/services/agent/mastra-stream.server.ts +++ b/apps/webapp/app/services/agent/mastra-stream.server.ts @@ -9,6 +9,7 @@ import { import { deductCredits } from "~/trigger/utils/utils"; import { logger } from "~/services/logger.service"; import { convertMastraChunkToAISDKv5 } from "@mastra/core/stream"; +import { getBurstSafeConversationIngestDelayMs } from "~/services/llm-provider.server"; /** * Builds assistant message parts from LLMStepResult[]. @@ -69,6 +70,7 @@ export async function saveConversationResult({ .map((p: any) => p.text); if (textParts.length > 0 && !incognito) { + const enqueueDelayMs = getBurstSafeConversationIngestDelayMs(); await addToQueue( { episodeBody: `${incomingUserText ?? ""}${textParts.join("\n")}`, @@ -79,6 +81,9 @@ export async function saveConversationResult({ }, userId, workspaceId, + undefined, + undefined, + enqueueDelayMs, ); } } diff --git a/apps/webapp/app/services/agent/no-stream-process.ts b/apps/webapp/app/services/agent/no-stream-process.ts index 154755577..18f13801c 100644 --- a/apps/webapp/app/services/agent/no-stream-process.ts +++ b/apps/webapp/app/services/agent/no-stream-process.ts @@ -6,7 +6,7 @@ import { setActiveStreamId, clearActiveStreamId, } from "../conversation.server"; -import { EpisodeType, UserTypeEnum } from "@core/types"; +import { UserTypeEnum } from "@core/types"; import { generateId, stepCountIs, JsonToSseTransformStream } from "ai"; import { Agent, convertMessages } from "@mastra/core/agent"; import type { OutputProcessor } from "@mastra/core/processors"; @@ -14,6 +14,7 @@ import { buildAgentContext } from "./context"; import { getMastra } from "./mastra"; import { getDefaultChatModelId, + getBurstSafeBackgroundDelayMs, resolveModelConfig, } from "~/services/llm-provider.server"; import { @@ -21,13 +22,16 @@ import { type DecisionContext, } from "~/services/agent/types/decision-agent"; import { type OrchestratorTools } from "~/services/agent/executors/base"; -import { createUIStreamWithApprovals } from "./mastra-stream.server"; +import { + createUIStreamWithApprovals, + saveConversationResult, +} from "./mastra-stream.server"; +import { runWithBurstRetry } from "./burst-retry.server"; import { getResumableStreamContext } from "~/bullmq/connection"; -import { deductCredits } from "~/trigger/utils/utils"; -import { addToQueue } from "~/lib/ingest.server"; interface NoStreamProcessBody { id: string; + modelId?: string; message?: { id?: string; parts: any[]; @@ -78,10 +82,11 @@ export async function noStreamProcess( if (conversationHistory.length === 1 && !isAssistantApproval) { const message = body.message?.parts[0].text; // Trigger conversation title task + const delayMs = getBurstSafeBackgroundDelayMs(); await enqueueCreateConversationTitle({ conversationId: body.id, message, - }); + }, delayMs); } const messageUserType = body.messageUserType ?? UserTypeEnum.User; @@ -93,7 +98,6 @@ export async function noStreamProcess( ) { const message = body.message?.parts[0].text; const messageParts = body.message?.parts; - await upsertConversationHistory( message.id ?? crypto.randomUUID(), messageParts, @@ -120,19 +124,30 @@ export async function noStreamProcess( if (!isAssistantApproval) { const id = body.message?.id; const userMessageId = id ?? generateId(); + const latestHistoryMessage = messages[messages.length - 1]; + const currentMessageParts = + body.message?.parts ?? [{ text: message, type: "text" }]; + const alreadyInHistory = + latestHistoryMessage?.role === "user" && + JSON.stringify(latestHistoryMessage.parts ?? []) === + JSON.stringify(currentMessageParts); finalMessages = [ ...messages, - { - parts: body.message?.parts ?? [{ text: message, type: "text" }], - role: "user", - id: userMessageId, - }, + ...(!alreadyInHistory + ? [ + { + parts: currentMessageParts, + role: "user", + id: userMessageId, + }, + ] + : []), ]; } else { finalMessages = body.messages as any; } - const modelString = getDefaultChatModelId(); + const modelString = body.modelId ?? getDefaultChatModelId(); const { modelConfig, isBYOK } = await resolveModelConfig( modelString, workspaceId, @@ -193,6 +208,7 @@ export async function noStreamProcess( // Capture final parts/text from outputProcessor for channel reply let capturedParts: any[] = []; let capturedText = ""; + let didPersistResult = false; const messageHistoryProcessor: OutputProcessor = { id: "message-history", @@ -207,25 +223,36 @@ export async function noStreamProcess( .filter((p: any) => p.type === "text") .map((p: any) => p.text) .join(""); + await saveConversationResult({ + parts: capturedParts, + conversationId: body.id, + incomingUserText: message, + incognito: conversation?.incognito ?? false, + userId, + workspaceId, + isBYOK, + }); + didPersistResult = true; return messages; }, }; let agentResult: any; try { - agentResult = await agent.generate(modelMessages, { - toolsets: { core: tools }, - stopWhen: [stepCountIs(10)], - modelSettings: { temperature: 0.5 }, - outputProcessors: [messageHistoryProcessor], - }); + agentResult = await runWithBurstRetry("conversation.generate", () => + agent.generate(modelMessages, { + toolsets: { core: tools }, + stopWhen: [stepCountIs(10)], + modelSettings: { temperature: 0.5 }, + outputProcessors: [messageHistoryProcessor], + }), + ); } catch (error) { await updateConversationStatus(body.id, "failed"); throw error; } // Build assistant parts from result.steps (handle Mastra payload wrapper) - const assistantMessageId = crypto.randomUUID(); const assistantParts: any[] = []; for (const step of agentResult.steps) { @@ -255,43 +282,24 @@ export async function noStreamProcess( } } - const assistantMessage = { - id: assistantMessageId, - role: "assistant", - parts: assistantParts, - }; - - try { - await upsertConversationHistory( - assistantMessageId, - assistantParts, - body.id, - UserTypeEnum.Agent, - false, - ); - - if (agentResult.text) { - await addToQueue( - { - episodeBody: `${message}${agentResult.text}`, - source: body.source, - referenceTime: new Date().toISOString(), - type: EpisodeType.CONVERSATION, - sessionId: body.id, - }, - userId, - workspaceId, - ); - } - - if (!isBYOK) { - await deductCredits(workspaceId, userId, "chatMessage", 1); - } - } finally { - await updateConversationStatus(body.id, "completed"); + if (!didPersistResult && assistantParts.length > 0) { + await saveConversationResult({ + parts: assistantParts, + conversationId: body.id, + incomingUserText: message, + incognito: conversation?.incognito ?? false, + userId, + workspaceId, + isBYOK, + }); } - return { ...assistantMessage, text: agentResult.text }; + return { + id: crypto.randomUUID(), + role: "assistant", + parts: capturedParts.length > 0 ? capturedParts : assistantParts, + text: capturedText || agentResult.text || "", + }; // const uiStream = createUIStreamWithApprovals(agentResult); // const sseStream = uiStream.pipeThrough(new JsonToSseTransformStream()); diff --git a/apps/webapp/app/services/llm-provider.server.ts b/apps/webapp/app/services/llm-provider.server.ts index d19431571..799a316ec 100644 --- a/apps/webapp/app/services/llm-provider.server.ts +++ b/apps/webapp/app/services/llm-provider.server.ts @@ -37,6 +37,8 @@ interface ProviderConfig { export type UseCase = "chat" | "memory" | "search"; export type ModelComplexity = "low" | "medium" | "high"; +const BURST_SAFE_BACKGROUND_DELAY_MS = 15_000; +const BURST_SAFE_CONVERSATION_INGEST_DELAY_MS = 45_000; // --------------------------------------------------------------------------- // Seeder @@ -73,6 +75,9 @@ function buildProviderConfig(providerType: string): Record { */ export async function ensureDefaultProviders(): Promise { const catalog = seedData as Record; + const protectedCustomModelIds = new Set( + [env.MODEL, env.EMBEDDING_MODEL].filter(Boolean) as string[], + ); for (const [providerType, providerData] of Object.entries(catalog)) { let provider = await prisma.lLMProvider.findFirst({ @@ -136,7 +141,11 @@ export async function ensureDefaultProviders(): Promise { } for (const existing of existingModels) { - if (!seedModelIds.has(existing.modelId) && !existing.isDeprecated) { + if ( + !seedModelIds.has(existing.modelId) && + !protectedCustomModelIds.has(existing.modelId) && + !existing.isDeprecated + ) { await prisma.lLMModel.update({ where: { id: existing.id }, data: { isDeprecated: true }, @@ -149,14 +158,19 @@ export async function ensureDefaultProviders(): Promise { // Dynamic model creation for env-specified models not in seed if (env.MODEL) { - const chatModelExists = await prisma.lLMModel.findFirst({ - where: { modelId: env.MODEL }, + const targetProvider = await prisma.lLMProvider.findFirst({ + where: { type: env.CHAT_PROVIDER, workspaceId: null }, }); - if (!chatModelExists) { - const targetProvider = await prisma.lLMProvider.findFirst({ - where: { type: env.CHAT_PROVIDER, workspaceId: null }, + if (targetProvider) { + const chatModel = await prisma.lLMModel.findFirst({ + where: { + providerId: targetProvider.id, + modelId: env.MODEL, + capabilities: { has: "chat" }, + }, }); - if (targetProvider) { + + if (!chatModel) { await prisma.lLMModel.create({ data: { providerId: targetProvider.id, @@ -170,21 +184,40 @@ export async function ensureDefaultProviders(): Promise { logger.info( `[LLM] Added custom chat model: ${env.MODEL} under ${env.CHAT_PROVIDER}`, ); + } else if (chatModel.isDeprecated) { + await prisma.lLMModel.update({ + where: { id: chatModel.id }, + data: { + label: env.MODEL, + complexity: chatModel.complexity, + supportsBatch: false, + capabilities: ["chat"], + isDeprecated: false, + isEnabled: true, + }, + }); + logger.info( + `[LLM] Restored custom chat model: ${env.MODEL} under ${env.CHAT_PROVIDER}`, + ); } } } const embeddingProvider = env.EMBEDDINGS_PROVIDER ?? "openai"; const embeddingModelId = env.EMBEDDING_MODEL || "text-embedding-3-small"; - const embeddingModelExists = await prisma.lLMModel.findFirst({ - where: { modelId: embeddingModelId, capabilities: { has: "embedding" } }, + const targetProvider = await prisma.lLMProvider.findFirst({ + where: { type: embeddingProvider, workspaceId: null }, }); - if (!embeddingModelExists) { - const targetProvider = await prisma.lLMProvider.findFirst({ - where: { type: embeddingProvider, workspaceId: null }, + if (targetProvider) { + const embeddingModel = await prisma.lLMModel.findFirst({ + where: { + providerId: targetProvider.id, + modelId: embeddingModelId, + capabilities: { has: "embedding" }, + }, }); - if (targetProvider) { - const dims = parseInt(env.EMBEDDING_MODEL_SIZE || "1024", 10); + const dims = parseInt(env.EMBEDDING_MODEL_SIZE || "1024", 10); + if (!embeddingModel) { await prisma.lLMModel.create({ data: { providerId: targetProvider.id, @@ -195,9 +228,25 @@ export async function ensureDefaultProviders(): Promise { capabilities: ["embedding"], dimensions: dims, }, + }); + logger.info( + `[LLM] Added custom embedding model: ${embeddingModelId} under ${embeddingProvider}`, + ); + } else if (embeddingModel.isDeprecated) { + await prisma.lLMModel.update({ + where: { id: embeddingModel.id }, + data: { + label: embeddingModelId, + complexity: embeddingModel.complexity, + supportsBatch: false, + capabilities: ["embedding"], + dimensions: dims, + isDeprecated: false, + isEnabled: true, + }, }); logger.info( - `[LLM] Added custom embedding model: ${embeddingModelId} under ${embeddingProvider}`, + `[LLM] Restored custom embedding model: ${embeddingModelId} under ${embeddingProvider}`, ); } } @@ -234,6 +283,64 @@ export function getProviderConfig(providerType: string): ProviderConfig { return {}; } +async function resolveProviderBaseUrl( + providerType: string, + workspaceId?: string | null, +): Promise { + if ( + providerType !== "azure" && + providerType !== "openai" && + providerType !== "ollama" + ) { + return undefined; + } + + const byokBaseUrl = workspaceId + ? await resolveWorkspaceProviderBaseUrl(workspaceId, providerType) + : null; + + const envBaseUrl = + providerType === "azure" + ? env.AZURE_BASE_URL + : providerType === "openai" + ? env.OPENAI_BASE_URL + : env.OLLAMA_URL; + + return byokBaseUrl ?? envBaseUrl; +} + +export function isBurstSensitiveChatProvider(): boolean { + return ( + env.CHAT_PROVIDER === "ollama" || + (env.CHAT_PROVIDER === "openai" && !!env.OPENAI_BASE_URL) + ); +} + +export function getBurstSafeBackgroundDelayMs(): number { + return isBurstSensitiveChatProvider() ? BURST_SAFE_BACKGROUND_DELAY_MS : 0; +} + +export function getBurstSafeConversationIngestDelayMs(): number { + return isBurstSensitiveChatProvider() + ? BURST_SAFE_CONVERSATION_INGEST_DELAY_MS + : 0; +} + +export function getBurstSafeBullmqConcurrency( + envKey: string, + defaultConcurrency: number, +): number { + if (process.env[envKey]) { + return defaultConcurrency; + } + + if (!isBurstSensitiveChatProvider()) { + return defaultConcurrency; + } + + return Math.min(defaultConcurrency, 1); +} + export async function getDefaultEmbeddingInfo(): Promise { const embeddingModelId = env.EMBEDDING_MODEL || "text-embedding-3-small"; const model = await prisma.lLMModel.findFirst({ @@ -263,8 +370,8 @@ export async function getEmbeddingDimensions(): Promise { * * Resolution order: * 1. workspace.metadata.modelConfig[useCase].modelId (explicit workspace override) - * 2. LLMModel with env.CHAT_PROVIDER + complexity (DB complexity routing) - * 3. env.MODEL (final fallback) + * 2. env.MODEL (server-level default) + * 3. LLMModel with env.CHAT_PROVIDER + complexity (DB complexity routing fallback) */ export async function getModelForUseCase( useCase: UseCase, @@ -286,7 +393,10 @@ export async function getModelForUseCase( if (modelId) return modelId; } - // 2. DB complexity routing via env.CHAT_PROVIDER + // 2. Server-level default model + if (env.MODEL) return env.MODEL; + + // 3. DB complexity routing via env.CHAT_PROVIDER const provider = await prisma.lLMProvider.findFirst({ where: { type: env.CHAT_PROVIDER, workspaceId: null }, }); @@ -303,7 +413,7 @@ export async function getModelForUseCase( if (model) return model.modelId; } - // 3. env fallback + // 4. env fallback return env.MODEL; } @@ -368,6 +478,10 @@ export async function getChatModels(workspaceId?: string) { } export async function getAvailableModels(workspaceId?: string) { + const defaultModelId = workspaceId + ? await getModelForUseCase("chat", workspaceId, "medium") + : getDefaultChatModelId(); + const providers = await getProviders(workspaceId); const providerIds = providers.map((p) => p.id); @@ -399,10 +513,13 @@ export async function getAvailableModels(workspaceId?: string) { include: { provider: true }, }); - return [...globalModels, ...customModels]; + return [...globalModels, ...customModels].map((model) => ({ + ...model, + isDefault: model.modelId === defaultModelId, + })); } - return prisma.lLMModel.findMany({ + const models = await prisma.lLMModel.findMany({ where: { providerId: { in: providerIds }, isEnabled: true, @@ -410,6 +527,11 @@ export async function getAvailableModels(workspaceId?: string) { }, include: { provider: true }, }); + + return models.map((model) => ({ + ...model, + isDefault: model.modelId === defaultModelId, + })); } // --------------------------------------------------------------------------- @@ -493,12 +615,13 @@ export async function resolveModelForWorkspace( providerType, ); - // For Azure, also resolve the base URL (BYOK stores it in baseUrl; env fallback) - if (providerType === "azure") { - const byokBaseUrl = workspaceId - ? await resolveWorkspaceProviderBaseUrl(workspaceId, "azure") - : null; - const baseUrl = byokBaseUrl ?? env.AZURE_BASE_URL; + // Providers with custom endpoints can resolve a workspace-scoped base URL. + if ( + providerType === "azure" || + providerType === "openai" || + providerType === "ollama" + ) { + const baseUrl = await resolveProviderBaseUrl(providerType, workspaceId); return { modelId, apiKey, isBYOK, baseUrl }; } @@ -512,7 +635,7 @@ export type OpenAICompatibleConfig = { headers?: Record; }; -export type ModelConfig = string | OpenAICompatibleConfig; +export type ModelConfig = string | OpenAICompatibleConfig | object; export interface ResolvedModelConfig { modelConfig: ModelConfig; @@ -523,15 +646,28 @@ export async function resolveModelConfig( modelString: string, workspaceId: string | null | undefined, ): Promise { - const { toRouterString, getProvider } = await import("~/lib/model.server"); + const { toRouterString, getProvider, getDirectModelInstance, shouldBypassMastraAgent } = await import( + "~/lib/model.server" + ); const providerType = getProvider(modelString); const { apiKey, isBYOK } = await resolveApiKeyForWorkspace( workspaceId, providerType, ); + const baseUrl = await resolveProviderBaseUrl(providerType, workspaceId); const routerString = toRouterString(modelString) as `${string}/${string}`; + if (shouldBypassMastraAgent(modelString, baseUrl)) { + return { + modelConfig: getDirectModelInstance(modelString, { + ...(apiKey && { apiKey }), + ...(baseUrl && { baseUrl }), + }) as object, + isBYOK, + }; + } + if (isBYOK && apiKey) { return { modelConfig: { id: routerString, apiKey }, isBYOK: true }; } diff --git a/apps/webapp/app/trigger/ingest/preprocess-episode.ts b/apps/webapp/app/trigger/ingest/preprocess-episode.ts index a6823414c..7058d36dc 100644 --- a/apps/webapp/app/trigger/ingest/preprocess-episode.ts +++ b/apps/webapp/app/trigger/ingest/preprocess-episode.ts @@ -29,11 +29,12 @@ export const preprocessTask = task({ }); }, // Callback to enqueue session compaction for conversations - async (compactionPayload) => { + async (compactionPayload, delayMs) => { await sessionCompactionTask.trigger(compactionPayload, { queue: "session-compaction-queue", concurrencyKey: compactionPayload.userId, tags: [compactionPayload.userId, compactionPayload.sessionId], + ...(delayMs ? { delay: `${Math.ceil(delayMs / 1000)}s` } : {}), }); }, );