diff --git a/frontend/src/components/EnhancedReasoning.tsx b/frontend/src/components/EnhancedReasoning.tsx index 4c6acea6f..e6aeda8a2 100644 --- a/frontend/src/components/EnhancedReasoning.tsx +++ b/frontend/src/components/EnhancedReasoning.tsx @@ -339,7 +339,7 @@ const ReasoningStepComponent: React.FC<{ // Get the display text - show actual message for step 1, summaries for steps 2 and 3 const getDisplayText = () => { if (isIteration) { - return `Attempt ${step.iterationNumber}` + return `Turn ${step.iterationNumber}` } // For initial messages, always show full content diff --git a/server/ai/modelConfig.ts b/server/ai/modelConfig.ts index abc8aae29..0b0154269 100644 --- a/server/ai/modelConfig.ts +++ b/server/ai/modelConfig.ts @@ -665,6 +665,53 @@ export const MODEL_CONFIGURATIONS: Record = { }, } +const DEFAULT_MAX_INPUT_TOKENS = 128_000 + +const MODEL_MAX_INPUT_TOKEN_OVERRIDES: Partial> = { + [Models.Claude_3_5_Haiku]: 200_000, + [Models.Claude_3_5_Sonnet]: 200_000, + [Models.Claude_3_5_SonnetV2]: 200_000, + [Models.Claude_3_7_Sonnet]: 200_000, + [Models.Claude_Opus_4]: 200_000, + [Models.Claude_Sonnet_4]: 200_000, + [Models.Amazon_Nova_Micro]: 300_000, + [Models.Amazon_Nova_Lite]: 300_000, + [Models.Amazon_Nova_Pro]: 300_000, + [Models.Gpt_4]: 8_192, + [Models.Gpt_4o]: 128_000, + [Models.Gpt_4o_mini]: 128_000, + [Models.o3_Deep_Research]: 200_000, + [Models.o4_Mini_Deep_Research]: 200_000, + [Models.Gemini_2_5_Flash]: 1_000_000, + [Models.Gemini_2_0_Flash_Thinking]: 1_000_000, + [Models.Vertex_Claude_Sonnet_4]: 200_000, + [Models.Vertex_Gemini_2_5_Pro]: 1_000_000, + [Models.Vertex_Gemini_2_5_Flash]: 1_000_000, + [Models.Vertex_Gemini_3_Pro]: 1_000_000, + [Models.Vertex_Gemini_3_Flash]: 1_000_000, +} + +for (const [model, maxInputTokens] of Object.entries( + MODEL_MAX_INPUT_TOKEN_OVERRIDES, +)) { + const entry = MODEL_CONFIGURATIONS[model as Models] + if (entry) { + entry.maxInputTokens = maxInputTokens + } +} + +export const getModelMaxInputTokens = ( + modelId?: Models | string | null, +): number => { + if (!modelId) { + return DEFAULT_MAX_INPUT_TOKENS + } + return ( + MODEL_CONFIGURATIONS[modelId as Models]?.maxInputTokens ?? + DEFAULT_MAX_INPUT_TOKENS + ) +} + // Model display name mappings - using the new enum-based approach export const MODEL_DISPLAY_NAMES: Record = { // Build from ModelDisplayNames enum diff --git a/server/ai/prompts.ts b/server/ai/prompts.ts index 9a798ce2c..b9c789fd5 100644 --- a/server/ai/prompts.ts +++ b/server/ai/prompts.ts @@ -2616,6 +2616,7 @@ Your goal is to capture not only directly matching documents but also those that - Offer **related background**, **context**, **examples**, or **clarifying information**. 4. **Prioritize quality** — prefer documents that are specific, factual, and contribute distinct value. 5. **Output** — Return only the indexes of the most relevant and complementary contexts. +6. **Honor agent prompt** — if you see "This is the system prompt of agent:", analyse it for instructions related to selection, ranking and filtering of source documents and treat it as binding and follow it strictly while selecting. ### Input - Query: "${query}" @@ -2712,4 +2713,4 @@ User question: ${query} Schema (all tables in the same database): ${schema} -` \ No newline at end of file +` diff --git a/server/api/chat/agent-schemas.ts b/server/api/chat/agent-schemas.ts index 254fb839d..9404187b7 100644 --- a/server/api/chat/agent-schemas.ts +++ b/server/api/chat/agent-schemas.ts @@ -164,6 +164,7 @@ export interface AgentRunContext { currentSubTask: string | null // Active substep ID userContext: string agentPrompt?: string + dedicatedAgentSystemPrompt?: string // Clarification tracking clarifications: Clarification[] diff --git a/server/api/chat/final-answer-synthesis.ts b/server/api/chat/final-answer-synthesis.ts new file mode 100644 index 000000000..9e3fef81a --- /dev/null +++ b/server/api/chat/final-answer-synthesis.ts @@ -0,0 +1,1301 @@ +import { getModelMaxInputTokens } from "@/ai/modelConfig" +import { getProviderByModel, jsonParseLLMOutput } from "@/ai/provider" +import { Models, type ModelParams } from "@/ai/types" +import config from "@/config" +import { getLogger, getLoggerWithChild } from "@/logger" +import { AgentReasoningStepType } from "@/shared/types" +import { Subsystem } from "@/types" +import { ConversationRole } from "@aws-sdk/client-bedrock-runtime" +import type { Message } from "@aws-sdk/client-bedrock-runtime" +import { z } from "zod" +import type { AgentRunContext, PlanState, SubTask } from "./agent-schemas" +import { + buildAgentSystemPromptContextBlock, + formatFragmentWithMetadata, + formatFragmentsWithMetadata, +} from "./message-agents-metadata" +import type { FragmentImageReference, MinimalAgentFragment } from "./types" + +const { defaultBestModel, IMAGE_CONTEXT_CONFIG } = config + +const Logger = getLogger(Subsystem.Chat) +const loggerWithChild = getLoggerWithChild(Subsystem.Chat) + +const IMAGE_TOKEN_ESTIMATE = 1_844 +const FINAL_OUTPUT_HEADROOM_RATIO = 0.15 +const FALLBACK_OUTPUT_TOKENS = 1_500 +const PREVIEW_TEXT_LENGTH = 320 +const MAX_SECTION_COUNT = 5 +const MAPPER_CONCURRENCY = 4 +const SECTION_CONCURRENCY = 4 + +export type FinalSynthesisExecutionResult = { + textLength: number + totalImagesAvailable: number + imagesProvided: number + estimatedCostUsd: number + mode: "single" | "sectional" +} + +type SynthesisModeSelection = { + mode: "single" | "sectional" + maxInputTokens: number + safeInputBudget: number + estimatedInputTokens: number +} + +type FinalSection = { + sectionId: number + title: string + objective: string +} + +type SectionAnswerResult = { + sectionId: number + title: string + body: string +} + +type FragmentPreviewRecord = { + fragmentIndex: number + docId: string + title?: string + app?: string + entity?: string + timestamp?: string + previewText: string +} + +type FragmentAssignmentBatch = { + fragmentIndex: number + fragment: MinimalAgentFragment +} + +type SelectedImagesResult = { + selected: string[] + total: number + dropped: string[] + userAttachmentCount: number +} + +type SectionMappingEnvelope = { + sections?: Record +} + +const SectionPlanSchema = z.object({ + sections: z + .array( + z.object({ + sectionId: z.number().int().positive().optional(), + title: z.string().trim().min(1), + objective: z.string().trim().min(1), + }), + ) + .min(1) + .max(MAX_SECTION_COUNT), +}) + +const SectionMappingSchema = z.object({ + sections: z + .record(z.string(), z.array(z.number().int().positive())) + .default({}), +}) + +function normalizeWhitespace(value: string): string { + return value.replace(/\s+/g, " ").trim() +} + +function truncateText(value: string, maxLength: number): string { + if (value.length <= maxLength) return value + return `${value.slice(0, Math.max(1, maxLength - 1))}…` +} + +export function estimateTextTokens(text: string): number { + return Math.ceil(text.length / 4) +} + +function estimateImageTokens(imageCount: number): number { + return imageCount * IMAGE_TOKEN_ESTIMATE +} + +function estimatePromptTokens( + systemPrompt: string, + userMessage: string, + imageCount = 0, +): number { + return ( + estimateTextTokens(systemPrompt) + + estimateTextTokens(userMessage) + + estimateImageTokens(imageCount) + ) +} + +function estimateSafeInputBudget( + modelId: Models, + maxOutputTokens?: number, +): { maxInputTokens: number; safeInputBudget: number } { + const maxInputTokens = getModelMaxInputTokens(modelId) + const reservedOutputTokens = Math.ceil( + (maxOutputTokens ?? FALLBACK_OUTPUT_TOKENS) * + (1 + FINAL_OUTPUT_HEADROOM_RATIO), + ) + const safeInputBudget = Math.max(1_024, maxInputTokens - reservedOutputTokens) + return { maxInputTokens, safeInputBudget } +} + +export function formatPlanForPrompt(plan: PlanState | null): string { + if (!plan) return "" + const lines = [`Goal: ${plan.goal}`] + plan.subTasks.forEach((task, idx) => { + const icon = + task.status === "completed" + ? "✓" + : task.status === "in_progress" + ? "→" + : task.status === "failed" + ? "✗" + : task.status === "blocked" + ? "!" + : "○" + const baseLine = `${idx + 1}. [${icon}] ${task.description}` + const detailParts: string[] = [] + if (task.result) detailParts.push(`Result: ${task.result}`) + if (task.toolsRequired?.length) { + detailParts.push(`Tools: ${task.toolsRequired.join(", ")}`) + } + lines.push( + detailParts.length > 0 ? `${baseLine}\n ${detailParts.join(" | ")}` : baseLine, + ) + }) + return lines.join("\n") +} + +export function formatClarificationsForPrompt( + clarifications: AgentRunContext["clarifications"], +): string { + if (!clarifications?.length) return "" + return clarifications + .map( + (clarification, idx) => + `${idx + 1}. Q: ${clarification.question}\n A: ${clarification.answer}`, + ) + .join("\n") +} + +export function buildSharedFinalAnswerContext( + context: AgentRunContext, +): string { + const agentSystemPromptBlock = buildAgentSystemPromptContextBlock( + context.dedicatedAgentSystemPrompt, + ) + const agentSystemPromptSection = agentSystemPromptBlock + ? `Agent System Prompt Context:\n${agentSystemPromptBlock}` + : "" + const planSection = formatPlanForPrompt(context.plan) + const clarificationSection = formatClarificationsForPrompt( + context.clarifications, + ) + const workspaceSection = context.userContext?.trim() + ? `Workspace Context:\n${context.userContext}` + : "" + + return [ + `User Question:\n${context.message.text}`, + agentSystemPromptSection, + planSection ? `Execution Plan Snapshot:\n${planSection}` : "", + clarificationSection + ? `Clarifications Resolved:\n${clarificationSection}` + : "", + workspaceSection, + ] + .filter(Boolean) + .join("\n\n") +} + +export function buildBaseFinalAnswerSystemPrompt( + mode: "final" | "section" = "final", +): string { + const mission = + mode === "final" + ? "- Deliver the user's final answer using the conversation, plan snapshot, clarifications, workspace context, context fragments, and supplied images; never plan or call tools." + : "- Deliver only the assigned answer section using the conversation, plan snapshot, clarifications, workspace context, mapped context fragments, and supplied images. Other sections are being generated in parallel and a final ordered answer will be assembled later; never attempt to write the full final answer." + const sectionRules = + mode === "section" + ? ` + +### Section Constraints +- Write only the requested section body for the assigned section. +- Do not add a global introduction, conclusion, or next-step sentence that assumes other sections are already visible. +- Do not repeat section headings for other sections. +- Treat the provided section list as context only; answer exclusively for the assigned section. +`.trim() + : "" + + return ` +### Mission +${mission} + +### Evidence Intake +- Prioritize the highest-signal fragments, but pull any supporting fragment that improves accuracy. +- Only draw on context that directly answers the user's question; ignore unrelated fragments even if they were retrieved earlier. +- Treat delegated-agent outputs as citeable fragments; reference them like any other context entry. +- Describe evidence gaps plainly before concluding; never guess. +- Extract actionable details from provided images and cite them via their fragment indices. +- Respect user-imposed constraints using fragment metadata (any metadata field). If compliant evidence is missing, state that clearly. +- If "This is the system prompt of agent:" is present, analyse for instructions relevant for answering and strictly bind by them . + +### Response Construction +- Lead with the conclusion, then stack proof underneath. +- Organize output into tight sections (e.g., **Summary**, **Proof**, **Next Steps** when relevant); omit empty sections. +- Never mention internal tooling, planning logs, or this synthesis process. + +### Constraint Handling +- When the user asks for an action the system cannot execute (e.g., sending an email), deliver the closest actionable substitute (draft, checklist, explicit next steps) inside the answer. +- Pair the substitute with a concise explanation of the limitation and the manual action the user must take. + +### File & Chunk Formatting (CRITICAL) +- Each file starts with a header line exactly like: + index {docId} {file context begins here...} +- \`docId\` is a unique identifier for that file (e.g., 0, 1, 2, etc.). +- Inside the file context, text is split into chunks. +- Each chunk might begin with a bracketed numeric index, e.g.: [0], [1], [2], etc. +- This is the chunk index within that file, if it exists. + +### Guidelines for Response +1. Data Interpretation: + - Use ONLY the provided files and their chunks as your knowledge base. + - Treat every file header \`index {docId} ...\` as the start of a new document. + - Treat every bracketed number like [0], [1], [2] as the authoritative chunk index within that document. + - If dates exist, interpret them relative to the user's timezone when paraphrasing. +2. Response Structure: + - Start with the most relevant facts from the chunks across files. + - Keep order chronological when it helps comprehension. + - Every factual statement MUST cite the exact chunk it came from using the format: + K[docId_chunkIndex] + where: + - \`docId\` is taken from the file header line ("index {docId} ..."). + - \`chunkIndex\` is the bracketed number prefixed on that chunk within the same file. + - Examples: + - Single citation: "X is true K[12_3]." + - Two citations in one sentence (from different files or chunks): "X K[12_3] and Y K[7_0]." + - Use at most 1-2 citations per sentence; NEVER add more than 2 for one sentence. +3. Citation Rules (DOCUMENT+CHUNK LEVEL ONLY): + - ALWAYS cite at the chunk level with the K[docId_chunkIndex] format. + - Every chunk level citation must start with the K prefix eg. K[12_3] K[7_0] correct, but K[12_3] [7_0] is incorrect. + - Place the citation immediately after the relevant claim. + - Do NOT group indices inside one set of brackets (WRONG: "K[12_3,7_1]"). + - If a sentence draws on two distinct chunks (possibly from different files), include two separate citations inline, e.g., "... K[12_3] ... K[7_1]". + - Only cite information that appears verbatim or is directly inferable from the cited chunk. + - If you cannot ground a claim to a specific chunk, do not make the claim. +4. Quality Assurance: + - Cross-check across multiple chunks/files when available and briefly note inconsistencies if they exist. + - Keep tone professional and concise. + - Acknowledge gaps if the provided chunks don't contain enough detail. + +### Tone & Delivery +- Answer with confident, declarative, verb-first sentences that use concrete nouns. +- Highlight key deliverables using **bold** labels or short lists; keep wording razor-concise. +- Ask one targeted follow-up question only if missing info blocks action. + +### Tool Spotlighting +- Reference critical tool outputs explicitly, e.g., "**Slack Search:** Ops escalated the RCA at 09:42 [2]." +- Explain why each highlighted tool mattered so reviewers see coverage breadth. +- When multiple tools contribute, show the sequence, e.g., "**Vespa Search:** context -> **Sheet Lookup:** metrics." +${sectionRules ? `\n\n${sectionRules}` : ""} + +### Finish +- Close with a single sentence confirming completion or the next action you recommend. +`.trim() +} + +export function buildFinalSynthesisPayload( + context: AgentRunContext, + fragmentsLimit = Math.max(12, context.allFragments.length || 1), +): { systemPrompt: string; userMessage: string } { + const sharedContext = buildSharedFinalAnswerContext(context) + const formattedFragments = formatFragmentsWithMetadata( + context.allFragments, + fragmentsLimit, + ) + const fragmentsSection = formattedFragments + ? `Context Fragments:\n${formattedFragments}` + : "" + + return { + systemPrompt: buildBaseFinalAnswerSystemPrompt("final"), + userMessage: [sharedContext, fragmentsSection].filter(Boolean).join("\n\n"), + } +} + +function formatSectionPlanOverview(sections: FinalSection[]): string { + return sections + .map( + (section) => + `${section.sectionId}. ${section.title}\n Objective: ${section.objective}`, + ) + .join("\n") +} + +function buildPlannerSystemPrompt(): string { + return ` +You are planning a final answer for a large evidence set. + +Return JSON only in the shape: +{ + "sections": [ + { "sectionId": 1, "title": "string", "objective": "string" } + ] +} + +Rules: +- Produce between 2 and ${MAX_SECTION_COUNT} sections when possible. +- Keep sections ordered exactly as the final answer should appear. +- Use concise, user-facing titles. +- Objectives should state what each section must accomplish. +- Do not include a catch-all section unless needed. +- Do not mention internal processing, tools, or token limits. +`.trim() +} + +function buildMapperSystemPrompt(): string { + return ` +You are mapping evidence fragments to pre-planned answer sections. + +Return JSON only in the shape: +{ + "sections": { + "1": [3, 7], + "2": [1] + } +} + +Rules: +- Keys are section ids. +- Values are fragment indexes from the provided batch. +- A fragment may belong to multiple sections when directly relevant. +- Omit fragments that do not help any section. +- Omit section ids with no fragments from this batch. +- Never invent fragment indexes. +`.trim() +} + +function formatSectionFragments( + entries: FragmentAssignmentBatch[], +): string { + return entries + .map((entry) => formatFragmentWithMetadata(entry.fragment, entry.fragmentIndex - 1)) + .join("\n\n") +} + +function findTimestamp(fragment: MinimalAgentFragment): string | undefined { + const source = fragment.source ?? {} + return ( + source.closedAt || + source.resolvedAt || + source.createdAt || + undefined + ) +} + +function buildFragmentPreviewRecord( + fragment: MinimalAgentFragment, + fragmentIndex: number, +): FragmentPreviewRecord { + const source = fragment.source ?? {} + const previewText = truncateText( + normalizeWhitespace(fragment.content ?? ""), + PREVIEW_TEXT_LENGTH, + ) + + return { + fragmentIndex, + docId: source.docId || fragment.id, + title: source.title || source.page_title || undefined, + app: source.app ? String(source.app) : undefined, + entity: source.entity ? String(source.entity) : undefined, + timestamp: findTimestamp(fragment), + previewText, + } +} + +function formatPreviewRecord(preview: FragmentPreviewRecord): string { + const meta = [ + `fragmentIndex: ${preview.fragmentIndex}`, + `docId: ${preview.docId}`, + preview.title ? `title: ${preview.title}` : "", + preview.app ? `app: ${preview.app}` : "", + preview.entity ? `entity: ${preview.entity}` : "", + preview.timestamp ? `timestamp: ${preview.timestamp}` : "", + ] + .filter(Boolean) + .join(" | ") + + return `${meta}\npreviewText: ${preview.previewText}` +} + +function buildPreviewOmissionSummary(previews: FragmentPreviewRecord[]): string { + if (previews.length === 0) return "" + const counts = new Map() + for (const preview of previews) { + const key = preview.app || "unknown" + counts.set(key, (counts.get(key) ?? 0) + 1) + } + const summary = Array.from(counts.entries()) + .sort((a, b) => b[1] - a[1]) + .map(([key, count]) => `${key}: ${count}`) + .join(", ") + return `Additional fragment previews omitted due to budget: ${previews.length} (${summary}).` +} + +function buildPreviewTextWithinBudget( + previews: FragmentPreviewRecord[], + budgetTokens: number, +): { includedText: string; omittedSummary: string } { + if (previews.length === 0) { + return { includedText: "None.", omittedSummary: "" } + } + + const included: string[] = [] + let usedTokens = 0 + let cutoff = previews.length + + for (let index = 0; index < previews.length; index++) { + const previewText = formatPreviewRecord(previews[index]) + const previewTokens = estimateTextTokens(`${previewText}\n\n`) + if (included.length > 0 && usedTokens + previewTokens > budgetTokens) { + cutoff = index + break + } + if (included.length === 0 && previewTokens > budgetTokens) { + included.push(previewText) + cutoff = index + 1 + usedTokens += previewTokens + break + } + included.push(previewText) + usedTokens += previewTokens + } + + return { + includedText: included.join("\n\n"), + omittedSummary: buildPreviewOmissionSummary(previews.slice(cutoff)), + } +} + +function buildFragmentBatchesWithinBudget( + entries: FragmentAssignmentBatch[], + baseTokens: number, + budgetTokens: number, +): FragmentAssignmentBatch[][] { + if (entries.length === 0) return [] + + const batches: FragmentAssignmentBatch[][] = [] + let currentBatch: FragmentAssignmentBatch[] = [] + let currentTokens = baseTokens + + for (const entry of entries) { + const itemTokens = estimateTextTokens( + `${formatFragmentWithMetadata(entry.fragment, entry.fragmentIndex - 1)}\n\n`, + ) + const wouldOverflow = + currentBatch.length > 0 && currentTokens + itemTokens > budgetTokens + + if (wouldOverflow) { + batches.push(currentBatch) + currentBatch = [] + currentTokens = baseTokens + } + + currentBatch.push(entry) + currentTokens += itemTokens + } + + if (currentBatch.length > 0) { + batches.push(currentBatch) + } + + return batches +} + +function normalizeSectionPlan(data: z.infer): FinalSection[] { + return data.sections.slice(0, MAX_SECTION_COUNT).map((section, index) => ({ + sectionId: index + 1, + title: normalizeWhitespace(section.title), + objective: normalizeWhitespace(section.objective), + })) +} + +function normalizeSectionAssignments( + raw: SectionMappingEnvelope, + sections: FinalSection[], +): Map> { + const validSectionIds = new Set(sections.map((section) => section.sectionId)) + const merged = new Map>() + + for (const [key, indexes] of Object.entries(raw.sections ?? {})) { + const sectionId = Number(key) + if (!validSectionIds.has(sectionId)) continue + const target = merged.get(sectionId) ?? new Set() + for (const index of indexes) { + if (Number.isInteger(index) && index > 0) { + target.add(index) + } + } + if (target.size > 0) { + merged.set(sectionId, target) + } + } + + return merged +} + +function mergeSectionAssignments( + assignments: Array>>, + sections: FinalSection[], +): Map { + const merged = new Map>() + for (const section of sections) { + merged.set(section.sectionId, new Set()) + } + for (const batchAssignment of assignments) { + for (const [sectionId, indexes] of batchAssignment.entries()) { + const target = merged.get(sectionId) ?? new Set() + for (const index of indexes) { + target.add(index) + } + merged.set(sectionId, target) + } + } + + return new Map( + Array.from(merged.entries()).map(([sectionId, indexes]) => [ + sectionId, + Array.from(indexes).sort((a, b) => a - b), + ]), + ) +} + +async function runWithConcurrency( + items: T[], + limit: number, + worker: (item: T, index: number) => Promise, +): Promise { + if (items.length === 0) return [] + const results = new Array(items.length) + let cursor = 0 + + const runWorker = async () => { + while (true) { + const current = cursor + cursor += 1 + if (current >= items.length) { + return + } + results[current] = await worker(items[current], current) + } + } + + const concurrency = Math.max(1, Math.min(limit, items.length)) + await Promise.all(Array.from({ length: concurrency }, () => runWorker())) + return results +} + +function buildDefaultSectionPlan(): FinalSection[] { + return [ + { + sectionId: 1, + title: "Answer", + objective: "Provide the best complete answer using the mapped evidence.", + }, + ] +} + +function createSelectedImagesResult( + images: FragmentImageReference[], + turnCount: number, +): SelectedImagesResult { + const total = images.length + if (!IMAGE_CONTEXT_CONFIG.enabled || total === 0) { + return { selected: [], total, dropped: [], userAttachmentCount: 0 } + } + + const attachments = images.filter((img) => img.isUserAttachment) + const nonAttachments = images + .filter((img) => !img.isUserAttachment) + .sort((a, b) => { + const ageA = turnCount - a.addedAtTurn + const ageB = turnCount - b.addedAtTurn + return ageA - ageB + }) + + const prioritized = [...attachments, ...nonAttachments] + const uniqueNames: string[] = [] + const seen = new Set() + for (const image of prioritized) { + if (seen.has(image.fileName)) continue + seen.add(image.fileName) + uniqueNames.push(image.fileName) + } + + let selected = uniqueNames + let dropped: string[] = [] + + if ( + IMAGE_CONTEXT_CONFIG.maxImagesPerCall > 0 && + uniqueNames.length > IMAGE_CONTEXT_CONFIG.maxImagesPerCall + ) { + selected = uniqueNames.slice(0, IMAGE_CONTEXT_CONFIG.maxImagesPerCall) + dropped = uniqueNames.slice(IMAGE_CONTEXT_CONFIG.maxImagesPerCall) + } + + return { + selected, + total, + dropped, + userAttachmentCount: attachments.length, + } +} + +function selectImagesForFinalSynthesis( + context: AgentRunContext, +): SelectedImagesResult { + return createSelectedImagesResult(context.allImages, context.turnCount) +} + +function selectImagesForFragmentIds( + context: AgentRunContext, + fragmentIds: Set, +): SelectedImagesResult { + const images = context.allImages.filter((image) => + fragmentIds.has(image.sourceFragmentId), + ) + return createSelectedImagesResult(images, context.turnCount) +} + +function selectMappedEntriesWithinBudget( + entries: FragmentAssignmentBatch[], + baseTokens: number, + budgetTokens: number, +): { selected: FragmentAssignmentBatch[]; trimmedCount: number } { + if (entries.length === 0) { + return { selected: [], trimmedCount: 0 } + } + + const selected: FragmentAssignmentBatch[] = [] + let usedTokens = baseTokens + + for (const entry of entries) { + const entryTokens = estimateTextTokens( + `${formatFragmentWithMetadata(entry.fragment, entry.fragmentIndex - 1)}\n\n`, + ) + if (selected.length > 0 && usedTokens + entryTokens > budgetTokens) { + break + } + if (selected.length === 0 && usedTokens + entryTokens > budgetTokens) { + selected.push(entry) + break + } + selected.push(entry) + usedTokens += entryTokens + } + + return { + selected, + trimmedCount: Math.max(entries.length - selected.length, 0), + } +} + +function buildSectionAnswerPayload( + context: AgentRunContext, + sections: FinalSection[], + section: FinalSection, + entries: FragmentAssignmentBatch[], + imageFileNames: string[], +): { systemPrompt: string; userMessage: string; imageFileNames: string[] } { + const sharedContext = buildSharedFinalAnswerContext(context) + const sectionOverview = formatSectionPlanOverview(sections) + const fragmentsText = formatSectionFragments(entries) + const fragmentsSection = fragmentsText + ? `Context Fragments For This Section:\n${fragmentsText}` + : "Context Fragments For This Section:\nNone." + + const userMessage = [ + sharedContext, + `All Planned Sections (generated in parallel; a final ordered answer will be assembled later):\n${sectionOverview}`, + `Assigned Section:\n${section.sectionId}. ${section.title}\nObjective: ${section.objective}`, + [ + "Section Instructions:", + "- Write only this section.", + "- Do not write the full final answer.", + "- Other sections are being generated in parallel.", + "- A final ordered answer will be assembled afterwards.", + "- Avoid intro or outro language that assumes the whole answer is already visible.", + "- Use the provided global fragment indexes exactly as shown for citations.", + ].join("\n"), + fragmentsSection, + ] + .filter(Boolean) + .join("\n\n") + + return { + systemPrompt: buildBaseFinalAnswerSystemPrompt("section"), + userMessage, + imageFileNames, + } +} + +function decideSynthesisMode( + context: AgentRunContext, + modelId: Models, + imageSelection: SelectedImagesResult, +): SynthesisModeSelection { + const payload = buildFinalSynthesisPayload( + context, + Math.max(12, context.allFragments.length || 1), + ) + const { maxInputTokens, safeInputBudget } = estimateSafeInputBudget( + modelId, + context.maxOutputTokens, + ) + const estimatedInputTokens = estimatePromptTokens( + payload.systemPrompt, + payload.userMessage, + imageSelection.selected.length, + ) + + return { + mode: estimatedInputTokens > safeInputBudget ? "sectional" : "single", + maxInputTokens, + safeInputBudget, + estimatedInputTokens, + } +} + +async function planSections( + context: AgentRunContext, + providerModelId: Models, + safeInputBudget: number, +): Promise<{ sections: FinalSection[]; estimatedCostUsd: number }> { + const previews = context.allFragments.map((fragment, index) => + buildFragmentPreviewRecord(fragment, index + 1), + ) + const sharedContext = buildSharedFinalAnswerContext(context) + const plannerUserIntro = [ + sharedContext, + "Create the final answer section plan using the fragment previews below.", + `Return at most ${MAX_SECTION_COUNT} sections.`, + "Fragment Previews:", + ].join("\n\n") + const baseTokens = + estimatePromptTokens( + buildPlannerSystemPrompt(), + plannerUserIntro, + 0, + ) + 256 + const previewBudget = Math.max(512, safeInputBudget - baseTokens) + const { includedText, omittedSummary } = buildPreviewTextWithinBudget( + previews, + previewBudget, + ) + const userMessage = [plannerUserIntro, includedText, omittedSummary] + .filter(Boolean) + .join("\n\n") + + const response = await getProviderByModel(providerModelId).converse( + [ + { + role: ConversationRole.USER, + content: [{ text: userMessage }], + }, + ], + { + modelId: providerModelId, + json: true, + stream: false, + temperature: 0, + max_new_tokens: 800, + systemPrompt: buildPlannerSystemPrompt(), + }, + ) + + const parsed = SectionPlanSchema.safeParse( + jsonParseLLMOutput(response.text ?? ""), + ) + + if (!parsed.success) { + Logger.warn( + { + issues: parsed.error.issues, + response: response.text, + chatId: context.chat.externalId, + }, + "[FinalAnswerSynthesis] Invalid planner output; falling back to default section plan.", + ) + return { + sections: buildDefaultSectionPlan(), + estimatedCostUsd: response.cost ?? 0, + } + } + + return { + sections: normalizeSectionPlan(parsed.data), + estimatedCostUsd: response.cost ?? 0, + } +} + +async function mapFragmentsToSections( + context: AgentRunContext, + sections: FinalSection[], + providerModelId: Models, + safeInputBudget: number, +): Promise<{ assignments: Map; estimatedCostUsd: number }> { + const sharedContext = buildSharedFinalAnswerContext(context) + const sectionOverview = formatSectionPlanOverview(sections) + const batchIntro = [ + sharedContext, + `Sections:\n${sectionOverview}`, + "Map the following fragments to the relevant section ids.", + "Fragments:", + ].join("\n\n") + const baseTokens = + estimatePromptTokens(buildMapperSystemPrompt(), batchIntro, 0) + 256 + const entries = context.allFragments.map((fragment, index) => ({ + fragmentIndex: index + 1, + fragment, + })) + const batches = buildFragmentBatchesWithinBudget( + entries, + baseTokens, + safeInputBudget, + ) + + if (batches.length === 0) { + return { assignments: new Map(), estimatedCostUsd: 0 } + } + + const batchResults = await runWithConcurrency( + batches, + MAPPER_CONCURRENCY, + async (batch) => { + const fragmentsText = formatSectionFragments(batch) + const userMessage = [batchIntro, fragmentsText].join("\n\n") + try { + const response = await getProviderByModel(providerModelId).converse( + [ + { + role: ConversationRole.USER, + content: [{ text: userMessage }], + }, + ], + { + modelId: providerModelId, + json: true, + stream: false, + temperature: 0, + max_new_tokens: 1_000, + systemPrompt: buildMapperSystemPrompt(), + }, + ) + + const parsed = SectionMappingSchema.safeParse( + jsonParseLLMOutput(response.text ?? ""), + ) + if (!parsed.success) { + Logger.warn( + { + issues: parsed.error.issues, + response: response.text, + chatId: context.chat.externalId, + }, + "[FinalAnswerSynthesis] Invalid mapper output for batch; skipping batch.", + ) + return { + assignments: new Map>(), + estimatedCostUsd: response.cost ?? 0, + } + } + + return { + assignments: normalizeSectionAssignments(parsed.data, sections), + estimatedCostUsd: response.cost ?? 0, + } + } catch (error) { + Logger.warn( + { + err: error instanceof Error ? error.message : String(error), + chatId: context.chat.externalId, + }, + "[FinalAnswerSynthesis] Mapper batch failed; skipping batch.", + ) + return { + assignments: new Map>(), + estimatedCostUsd: 0, + } + } + }, + ) + + return { + assignments: mergeSectionAssignments( + batchResults.map((result) => result.assignments), + sections, + ), + estimatedCostUsd: batchResults.reduce( + (sum, result) => sum + result.estimatedCostUsd, + 0, + ), + } +} + +function buildDefaultAssignments( + context: AgentRunContext, + sections: FinalSection[], +): Map { + const indexes = context.allFragments.map((_, index) => index + 1) + return new Map( + sections.map((section) => [section.sectionId, [...indexes]]), + ) +} + +async function synthesizeSingleAnswer( + context: AgentRunContext, + modelId: Models, + imageSelection: SelectedImagesResult, +): Promise { + const streamAnswer = context.runtime?.streamAnswerText + if (!streamAnswer) { + throw new Error("Streaming channel unavailable. Cannot deliver final answer.") + } + + const { systemPrompt, userMessage } = buildFinalSynthesisPayload(context) + const provider = getProviderByModel(modelId) + let streamedCharacters = 0 + let estimatedCostUsd = 0 + + const iterator = provider.converseStream( + [ + { + role: ConversationRole.USER, + content: [ + { + text: `${userMessage}\n\nSynthesize the final answer using the evidence above.`, + }, + ], + }, + ], + { + modelId, + systemPrompt, + stream: true, + temperature: 0.2, + max_new_tokens: context.maxOutputTokens ?? FALLBACK_OUTPUT_TOKENS, + imageFileNames: imageSelection.selected, + }, + ) + + for await (const chunk of iterator) { + if (chunk.text) { + streamedCharacters += chunk.text.length + context.finalSynthesis.streamedText += chunk.text + await streamAnswer(chunk.text) + } + const chunkCost = chunk.metadata?.cost + if (typeof chunkCost === "number" && !Number.isNaN(chunkCost)) { + estimatedCostUsd += chunkCost + } + } + + return { + textLength: streamedCharacters, + totalImagesAvailable: imageSelection.total, + imagesProvided: imageSelection.selected.length, + estimatedCostUsd, + mode: "single", + } +} + +async function synthesizeSection( + context: AgentRunContext, + sections: FinalSection[], + section: FinalSection, + mappedIndexes: number[], + providerModelId: Models, + safeInputBudget: number, +): Promise<{ + result: SectionAnswerResult | null + estimatedCostUsd: number + imageFileNames: string[] +}> { + const orderedEntries = mappedIndexes + .map((index) => ({ + fragmentIndex: index, + fragment: context.allFragments[index - 1], + })) + .filter((entry) => !!entry.fragment) as FragmentAssignmentBatch[] + + if (orderedEntries.length === 0) { + return { result: null, estimatedCostUsd: 0, imageFileNames: [] } + } + + const fragmentIds = new Set(orderedEntries.map((entry) => entry.fragment.id)) + const imageSelection = selectImagesForFragmentIds(context, fragmentIds) + + const emptyPayload = buildSectionAnswerPayload( + context, + sections, + section, + [], + imageSelection.selected, + ) + const baseTokens = + estimatePromptTokens( + emptyPayload.systemPrompt, + emptyPayload.userMessage, + imageSelection.selected.length, + ) + 128 + const { selected, trimmedCount } = selectMappedEntriesWithinBudget( + orderedEntries, + baseTokens, + safeInputBudget, + ) + + if (trimmedCount > 0) { + loggerWithChild({ email: context.user.email }).info( + { + chatId: context.chat.externalId, + sectionId: section.sectionId, + trimmedCount, + }, + "[FinalAnswerSynthesis] Trimmed mapped fragments to fit section input budget.", + ) + } + + const payload = buildSectionAnswerPayload( + context, + sections, + section, + selected, + imageSelection.selected, + ) + const sectionMaxTokens = Math.min( + context.maxOutputTokens ?? FALLBACK_OUTPUT_TOKENS, + Math.max( + 250, + Math.ceil( + (context.maxOutputTokens ?? FALLBACK_OUTPUT_TOKENS) / + Math.max(sections.length, 1), + ), + ), + ) + + try { + const response = await getProviderByModel(providerModelId).converse( + [ + { + role: ConversationRole.USER, + content: [{ text: payload.userMessage }], + }, + ], + { + modelId: providerModelId, + stream: false, + temperature: 0.2, + max_new_tokens: sectionMaxTokens, + systemPrompt: payload.systemPrompt, + imageFileNames: payload.imageFileNames, + }, + ) + + const body = response.text?.trim() ?? "" + return { + result: body + ? { + sectionId: section.sectionId, + title: section.title, + body, + } + : null, + estimatedCostUsd: response.cost ?? 0, + imageFileNames: payload.imageFileNames, + } + } catch (error) { + Logger.warn( + { + err: error instanceof Error ? error.message : String(error), + chatId: context.chat.externalId, + sectionId: section.sectionId, + }, + "[FinalAnswerSynthesis] Section synthesis failed; omitting section.", + ) + return { + result: null, + estimatedCostUsd: 0, + imageFileNames: payload.imageFileNames, + } + } +} + +function assembleSectionAnswers(results: SectionAnswerResult[]): string { + return results + .sort((a, b) => a.sectionId - b.sectionId) + .map((result) => `**${result.title}**\n${result.body}`) + .join("\n\n") + .trim() +} + +async function synthesizeSectionalAnswer( + context: AgentRunContext, + modelId: Models, + imageSelection: SelectedImagesResult, + safeInputBudget: number, +): Promise { + let estimatedCostUsd = 0 + + await context.runtime?.emitReasoning?.({ + text: `Final synthesis exceeded the model input budget. Switching to sectional synthesis across ${context.allFragments.length} fragments.`, + step: { type: AgentReasoningStepType.LogMessage }, + }) + + const planned = await planSections(context, modelId, safeInputBudget) + estimatedCostUsd += planned.estimatedCostUsd + let sections = planned.sections + + const mapped = await mapFragmentsToSections( + context, + sections, + modelId, + safeInputBudget, + ) + estimatedCostUsd += mapped.estimatedCostUsd + + let assignments = mapped.assignments + const hasAssignments = Array.from(assignments.values()).some( + (indexes) => indexes.length > 0, + ) + + if (!hasAssignments) { + sections = buildDefaultSectionPlan() + assignments = buildDefaultAssignments(context, sections) + } + + const sectionResults = await runWithConcurrency( + sections, + SECTION_CONCURRENCY, + async (section) => + synthesizeSection( + context, + sections, + section, + assignments.get(section.sectionId) ?? [], + modelId, + safeInputBudget, + ), + ) + + estimatedCostUsd += sectionResults.reduce( + (sum, result) => sum + result.estimatedCostUsd, + 0, + ) + + let assembledText = assembleSectionAnswers( + sectionResults + .map((result) => result.result) + .filter((result): result is SectionAnswerResult => !!result), + ) + + if (!assembledText) { + const fallbackSections = buildDefaultSectionPlan() + const fallbackAssignments = buildDefaultAssignments(context, fallbackSections) + const fallback = await synthesizeSection( + context, + fallbackSections, + fallbackSections[0], + fallbackAssignments.get(1) ?? [], + modelId, + safeInputBudget, + ) + estimatedCostUsd += fallback.estimatedCostUsd + assembledText = assembleSectionAnswers( + fallback.result ? [fallback.result] : [], + ) + if (fallback.result) { + sectionResults.push(fallback) + } + } + + if (!assembledText) { + throw new Error("Sectional final synthesis produced no answer text.") + } + + const uniqueImagesProvided = new Set() + for (const result of sectionResults) { + for (const imageName of result.imageFileNames) { + uniqueImagesProvided.add(imageName) + } + } + + context.finalSynthesis.streamedText = assembledText + await context.runtime?.streamAnswerText?.(assembledText) + + return { + textLength: assembledText.length, + totalImagesAvailable: imageSelection.total, + imagesProvided: uniqueImagesProvided.size, + estimatedCostUsd, + mode: "sectional", + } +} + +export async function executeFinalSynthesis( + context: AgentRunContext, +): Promise { + const modelId = + (context.modelId as Models) || + (defaultBestModel as Models) || + Models.Gpt_4o + const imageSelection = selectImagesForFinalSynthesis(context) + const modeSelection = decideSynthesisMode(context, modelId, imageSelection) + + loggerWithChild({ email: context.user.email }).debug( + { + chatId: context.chat.externalId, + mode: modeSelection.mode, + maxInputTokens: modeSelection.maxInputTokens, + safeInputBudget: modeSelection.safeInputBudget, + estimatedInputTokens: modeSelection.estimatedInputTokens, + fragmentsCount: context.allFragments.length, + selectedImages: imageSelection.selected, + droppedImages: imageSelection.dropped, + userAttachmentCount: imageSelection.userAttachmentCount, + }, + "[FinalAnswerSynthesis] Selected final synthesis mode.", + ) + + if (imageSelection.dropped.length > 0) { + loggerWithChild({ email: context.user.email }).info( + { + chatId: context.chat.externalId, + droppedCount: imageSelection.dropped.length, + limit: IMAGE_CONTEXT_CONFIG.maxImagesPerCall, + totalImages: imageSelection.total, + }, + "[FinalAnswerSynthesis] Image limit enforced for single-shot selection.", + ) + } + + return modeSelection.mode === "single" + ? synthesizeSingleAnswer(context, modelId, imageSelection) + : synthesizeSectionalAnswer( + context, + modelId, + imageSelection, + modeSelection.safeInputBudget, + ) +} + +export const __finalAnswerSynthesisInternals = { + buildFragmentPreviewRecord, + buildSectionAnswerPayload, + decideSynthesisMode, + selectImagesForFragmentIds, +} diff --git a/server/api/chat/message-agents-metadata.ts b/server/api/chat/message-agents-metadata.ts new file mode 100644 index 000000000..22d40cf59 --- /dev/null +++ b/server/api/chat/message-agents-metadata.ts @@ -0,0 +1,469 @@ +import { ConversationRole } from "@aws-sdk/client-bedrock-runtime" +import type { Message } from "@aws-sdk/client-bedrock-runtime" +import type { MinimalAgentFragment } from "./types" + +const METADATA_VALUE_MAX_LENGTH = 220 +const METADATA_TERM_MAX_LENGTH = 96 +const AGENT_SYSTEM_PROMPT_MAX_LENGTH = 6000 +const AGENT_SYSTEM_PROMPT_LABEL = "This is the system prompt of agent:" +const METADATA_QUERY_TRIGGER = + /\b(only|strictly|exclusively|from|source|sources|document|documents|doc|docs|file|files|app|entity|exclude|excluding|except|without|not)\b/i +const METADATA_STRICT_TRIGGER = /\b(only|strictly|exclusively|just)\b/i +const METADATA_TARGETED_TRIGGER = + /\b(only|strictly|exclusively|source|sources|document|documents|doc|docs|file|files|app|entity|exclude|excluding|except|without|not)\b/i +const NON_SIGNAL_TERMS = new Set([ + "the", + "a", + "an", + "document", + "documents", + "doc", + "docs", + "file", + "files", + "source", + "sources", + "app", + "entity", + "metadata", + "chunk", + "chunks", + "context", + "data", + "result", + "results", +]) + +type MetadataQueryConstraints = { + includeTerms: string[] + excludeTerms: string[] + strict: boolean +} + +type RankedMetadataCandidate = { + fragment: MinimalAgentFragment + includeScore: number + excludeScore: number + score: number + compliant: boolean +} + +function truncateValue(value: string, maxLength = 160): string { + if (value.length <= maxLength) return value + return `${value.slice(0, maxLength - 1)}…` +} + +function normalizeWhitespace(value: string): string { + return value.replace(/\s+/g, " ").trim() +} + +export function sanitizeAgentSystemPromptSnapshot( + prompt: string | undefined +): string | undefined { + if (!prompt || typeof prompt !== "string") { + return undefined + } + const normalized = normalizeWhitespace(prompt) + if (!normalized) { + return undefined + } + return truncateValue(normalized, AGENT_SYSTEM_PROMPT_MAX_LENGTH) +} + +export function buildAgentSystemPromptContextBlock( + prompt: string | undefined +): string | undefined { + const snapshot = sanitizeAgentSystemPromptSnapshot(prompt) + if (!snapshot) { + return undefined + } + return `${AGENT_SYSTEM_PROMPT_LABEL} + +${snapshot} +` +} + +function hasAgentSystemPromptBlock(messages: Message[]): boolean { + return messages.some((message) => { + const content = (message as any)?.content + if (!Array.isArray(content)) return false + return content.some((entry: any) => { + return ( + typeof entry?.text === "string" && + entry.text.includes(AGENT_SYSTEM_PROMPT_LABEL) + ) + }) + }) +} + +export function withAgentSystemPromptMessage( + messages: Message[], + prompt: string | undefined +): Message[] { + const block = buildAgentSystemPromptContextBlock(prompt) + if (!block || hasAgentSystemPromptBlock(messages)) { + return messages + } + return [ + ...messages, + { + role: ConversationRole.USER, + content: [{ text: block }], + }, + ] +} + +function normalizeMetadataValue(value: unknown): string | null { + if (value === undefined || value === null) return null + let serialized = "" + if (typeof value === "string") { + serialized = value + } else if (typeof value === "number" || typeof value === "boolean") { + serialized = String(value) + } else if (Array.isArray(value)) { + serialized = value + .map((entry) => (entry === undefined || entry === null ? "" : String(entry))) + .filter(Boolean) + .join(", ") + } else { + try { + serialized = JSON.stringify(value) + } catch { + serialized = String(value) + } + } + const normalized = normalizeWhitespace(serialized) + if (!normalized) return null + return truncateValue(normalized, METADATA_VALUE_MAX_LENGTH) +} + +function collectFragmentMetadataEntries( + fragment: MinimalAgentFragment +): Array<[string, string]> { + const source = (fragment.source || {}) as Record + const preferredOrder = [ + "title", + "page_title", + "app", + "entity", + "docId", + "url", + "threadId", + "itemId", + "clId", + "parentThreadId", + "createdAt", + "resolvedAt", + "closedAt", + "status", + "ticketNumber", + ] + const keys = Object.keys(source) + const orderedKeys = [ + ...preferredOrder.filter((key) => key in source), + ...keys.filter((key) => !preferredOrder.includes(key)).sort(), + ] + + const entries: Array<[string, string]> = [] + for (const key of orderedKeys) { + const normalized = normalizeMetadataValue(source[key]) + if (!normalized) continue + entries.push([key, normalized]) + } + + const fragmentId = normalizeMetadataValue(fragment.id) + if (fragmentId) { + entries.push(["fragmentId", fragmentId]) + } + return entries +} + +function buildFragmentMetadataSearchText(fragment: MinimalAgentFragment): string { + const pairs = collectFragmentMetadataEntries(fragment) + const metadataText = pairs.map(([key, value]) => `${key}: ${value}`).join(" | ") + const confidenceText = + typeof fragment.confidence === "number" && Number.isFinite(fragment.confidence) + ? ` | confidence: ${fragment.confidence.toFixed(3)}` + : "" + return `${metadataText}${confidenceText}`.toLowerCase() +} + +export function formatFragmentWithMetadata( + fragment: MinimalAgentFragment, + index: number +): string { + const metadataEntries = collectFragmentMetadataEntries(fragment) + if (typeof fragment.confidence === "number" && Number.isFinite(fragment.confidence)) { + metadataEntries.push(["confidence", fragment.confidence.toFixed(3)]) + } + const metadataBlock = metadataEntries.length + ? metadataEntries.map(([key, value]) => `- ${key}: ${value}`).join("\n") + : "- unavailable" + const content = fragment.content?.trim() || "No content." + return `index ${index + 1} {file context begins here...} +Metadata: +${metadataBlock} +Content: +${content}` +} + +export function formatFragmentsWithMetadata( + fragments: MinimalAgentFragment[], + maxFragments?: number +): string { + if (!fragments || fragments.length === 0) { + return "" + } + const limit = + typeof maxFragments === "number" + ? Math.max(0, Math.min(maxFragments, fragments.length)) + : fragments.length + if (limit === 0) { + return "" + } + return fragments + .slice(0, limit) + .map((fragment, index) => formatFragmentWithMetadata(fragment, index)) + .join("\n\n") +} + +function splitConstraintCandidates(raw: string): string[] { + return raw + .split(/,|;|\band\b|\bor\b/gi) + .map((part) => normalizeWhitespace(part)) + .filter(Boolean) +} + +function normalizeConstraintTerm(raw: string): string | null { + let normalized = normalizeWhitespace(raw.toLowerCase()) + if (!normalized) return null + normalized = normalized + .replace(/^[\s"'`]+|[\s"'`]+$/g, "") + .replace( + /\b(?:documents?|docs?|files?|sources?|metadata|records?|items?)\b/g, + " " + ) + normalized = normalizeWhitespace(normalized).replace(/[.?!,:;]+$/g, "") + if (!normalized || normalized.length < 2) return null + if (NON_SIGNAL_TERMS.has(normalized)) return null + if ( + /^(last|this|next)\s+(day|week|month|quarter|year)$/i.test(normalized) || + /^(today|yesterday|tomorrow)$/i.test(normalized) || + /^\d{4}$/.test(normalized) + ) { + return null + } + if (normalized.split(/\s+/).length > 8) return null + return truncateValue(normalized, METADATA_TERM_MAX_LENGTH) +} + +function addConstraintTerms( + target: Set, + rawValue: string +): void { + splitConstraintCandidates(rawValue).forEach((candidate) => { + const normalized = normalizeConstraintTerm(candidate) + if (normalized) { + target.add(normalized) + } + }) +} + +export function extractMetadataConstraintsFromUserMessage( + userMessage: string +): MetadataQueryConstraints { + const includeTerms = new Set() + const excludeTerms = new Set() + const normalizedMessage = normalizeWhitespace(userMessage) + if (!normalizedMessage) { + return { includeTerms: [], excludeTerms: [], strict: false } + } + + const hasConstraintSignal = METADATA_QUERY_TRIGGER.test(normalizedMessage) + const hasTargetedSignal = METADATA_TARGETED_TRIGGER.test(normalizedMessage) + const strict = METADATA_STRICT_TRIGGER.test(normalizedMessage) + + if (hasTargetedSignal || strict) { + const quotedPattern = /"([^"]{2,180})"|'([^']{2,180})'/g + let quotedMatch: RegExpExecArray | null + while ((quotedMatch = quotedPattern.exec(normalizedMessage)) !== null) { + addConstraintTerms(includeTerms, quotedMatch[1] || quotedMatch[2]) + } + } + + if (hasConstraintSignal) { + const includePatterns: RegExp[] = [ + /\bonly\s+from\s+([^,.!?;\n]{2,180})/gi, + /\b(?:source|sources|app|entity|document|documents|doc|docs|file|files)\s*(?:is|are|=|:)?\s*([^,.!?;\n]{2,180})/gi, + ] + if (hasTargetedSignal || includeTerms.size > 0) { + includePatterns.push(/\bfrom\s+([^,.!?;\n]{2,180})/gi) + } + const excludePatterns = [ + /\bnot\s+from\s+([^,.!?;\n]{2,180})/gi, + /\b(?:exclude|excluding|except|without)\s+([^,.!?;\n]{2,180})/gi, + ] + + for (const pattern of includePatterns) { + let match: RegExpExecArray | null + while ((match = pattern.exec(normalizedMessage)) !== null) { + addConstraintTerms(includeTerms, match[1]) + } + } + for (const pattern of excludePatterns) { + let match: RegExpExecArray | null + while ((match = pattern.exec(normalizedMessage)) !== null) { + addConstraintTerms(excludeTerms, match[1]) + } + } + } + + for (const excluded of excludeTerms) { + includeTerms.delete(excluded) + } + + return { + includeTerms: Array.from(includeTerms), + excludeTerms: Array.from(excludeTerms), + strict, + } +} + +function computeTermMatchScore(metadataText: string, term: string): number { + if (!term || !metadataText) return 0 + if (metadataText.includes(term)) { + return term.includes(" ") ? 3 : 2 + } + const tokens = term + .split(/\s+/) + .map((token) => token.trim()) + .filter((token) => token.length >= 3 && !NON_SIGNAL_TERMS.has(token)) + if (tokens.length === 0) return 0 + let tokenHits = 0 + for (const token of tokens) { + if (metadataText.includes(token)) { + tokenHits++ + } + } + if (tokenHits === tokens.length) return 2 + if (tokenHits >= Math.max(1, Math.ceil(tokens.length / 2))) return 1 + return 0 +} + +export function rankFragmentsByMetadataConstraints( + fragments: MinimalAgentFragment[], + constraints: MetadataQueryConstraints +): { + rankedCandidates: RankedMetadataCandidate[] + hasConstraints: boolean + hasCompliantCandidates: boolean +} { + const hasConstraints = + constraints.includeTerms.length > 0 || constraints.excludeTerms.length > 0 + const scored: RankedMetadataCandidate[] = fragments.map((fragment) => { + const metadataText = buildFragmentMetadataSearchText(fragment) + const includeScore = constraints.includeTerms.reduce((total, term) => { + return total + computeTermMatchScore(metadataText, term) + }, 0) + const excludeScore = constraints.excludeTerms.reduce((total, term) => { + return total + computeTermMatchScore(metadataText, term) + }, 0) + const includeCompliant = + constraints.includeTerms.length === 0 || includeScore > 0 + const excludeCompliant = excludeScore === 0 + const compliant = includeCompliant && excludeCompliant + let score = + includeScore * 3 - + excludeScore * 4 + + (fragment.confidence || 0) + if (constraints.strict && !compliant) { + score -= 100 + } + return { + fragment, + includeScore, + excludeScore, + score, + compliant, + } + }) + + if (!hasConstraints) { + return { + rankedCandidates: scored, + hasConstraints: false, + hasCompliantCandidates: false, + } + } + + scored.sort((a, b) => { + if (a.compliant !== b.compliant) { + return a.compliant ? -1 : 1 + } + if (a.score !== b.score) { + return b.score - a.score + } + return (b.fragment.confidence || 0) - (a.fragment.confidence || 0) + }) + + const compliantCandidates = scored.filter((candidate) => candidate.compliant) + if (constraints.strict && compliantCandidates.length > 0) { + return { + rankedCandidates: compliantCandidates, + hasConstraints: true, + hasCompliantCandidates: true, + } + } + + return { + rankedCandidates: scored, + hasConstraints: true, + hasCompliantCandidates: compliantCandidates.length > 0, + } +} + +export function enforceMetadataConstraintsOnSelection( + selected: MinimalAgentFragment[], + rankedCandidates: RankedMetadataCandidate[], + constraints: MetadataQueryConstraints +): MinimalAgentFragment[] { + const hasConstraints = + constraints.includeTerms.length > 0 || constraints.excludeTerms.length > 0 + if (!hasConstraints || selected.length === 0) { + return selected + } + + const candidateById = new Map( + rankedCandidates.map((candidate) => [candidate.fragment.id, candidate]) + ) + const compliantSelected = selected.filter((fragment) => { + return candidateById.get(fragment.id)?.compliant + }) + if (constraints.strict) { + return compliantSelected + } + if (compliantSelected.length > 0) { + return compliantSelected + } + const compliantFallback = rankedCandidates + .filter((candidate) => candidate.compliant) + .slice(0, Math.min(3, rankedCandidates.length)) + .map((candidate) => candidate.fragment) + if (compliantFallback.length === 0) { + return selected + } + const merged = [...compliantFallback, ...selected] + const seen = new Set() + return merged.filter((fragment) => { + if (seen.has(fragment.id)) return false + seen.add(fragment.id) + return true + }) +} + +export const __messageAgentsMetadataInternals = { + formatFragmentWithMetadata, + formatFragmentsWithMetadata, + extractMetadataConstraintsFromUserMessage, + rankFragmentsByMetadataConstraints, +} diff --git a/server/api/chat/message-agents.ts b/server/api/chat/message-agents.ts index d09ef5803..7158e0125 100644 --- a/server/api/chat/message-agents.ts +++ b/server/api/chat/message-agents.ts @@ -65,6 +65,7 @@ import { getProviderByModel, jsonParseLLMOutput, } from "@/ai/provider" +import { extractBestDocumentsPrompt } from "@/ai/prompts" import { answerContextMap, answerContextMapFromFragments, userContext } from "@/ai/context" import { ConversationRole } from "@aws-sdk/client-bedrock-runtime" import type { Message } from "@aws-sdk/client-bedrock-runtime" @@ -102,7 +103,7 @@ import { parseAttachmentMetadata } from "@/utils/parseAttachment" import { db } from "@/db/client" import { insertChat, updateChatByExternalIdWithAuth } from "@/db/chat" import { insertChatTrace } from "@/db/chatTrace" -import { insertMessage } from "@/db/message" +import { getChatMessagesWithAuth, insertMessage } from "@/db/message" import { storeAttachmentMetadata } from "@/db/attachment" import { ChatType, @@ -137,6 +138,11 @@ import { parseAgentAppIntegrations } from "./tools/utils" import { buildAgentPromptAddendum } from "./agentPromptCreation" import { getConnectorById } from "@/db/connector" import { getToolsByConnectorId } from "@/db/tool" +import { + executeFinalSynthesis, + formatClarificationsForPrompt, + formatPlanForPrompt, +} from "./final-answer-synthesis" import { buildMCPJAFTools, type FinalToolsList, @@ -157,6 +163,17 @@ import { parseMessageText } from "./chat" import { getUserPersonalizationByEmail } from "@/db/personalization" import { getChunkCountPerDoc } from "./chunk-selection" import { getPrecomputedDbContextIfNeeded } from "@/lib/databaseContext" +import { + enforceMetadataConstraintsOnSelection, + extractMetadataConstraintsFromUserMessage, + formatFragmentWithMetadata, + rankFragmentsByMetadataConstraints, + sanitizeAgentSystemPromptSnapshot, + withAgentSystemPromptMessage, +} from "./message-agents-metadata" + +export { __messageAgentsMetadataInternals } from "./message-agents-metadata" +export { buildFinalSynthesisPayload } from "./final-answer-synthesis" const { defaultBestModel, @@ -319,6 +336,81 @@ function truncateValue(value: string, maxLength = 160): string { return `${value.slice(0, maxLength - 1)}…` } +function normalizeUserMessageForHistory(message: SelectMessage): string { + const fileIds = Array.isArray(message?.fileIds) ? message.fileIds : [] + if ( + message.messageRole !== MessageRole.User || + !fileIds.length || + !message.message.startsWith("[{") + ) { + return message.message + } + + try { + const parsed = JSON.parse(message.message) + if (!Array.isArray(parsed)) { + return message.message + } + return parsed + .map((item) => { + if (item?.type === "text") { + return `${item?.value ?? ""} ` + } + if (item?.type === "pill") { + const title = item?.value?.title ?? "Unknown file" + return ` ` + } + if (item?.type === "link") { + return " " + } + return "" + }) + .join("") + .trim() + } catch { + return message.message + } +} + +function buildConversationHistoryForAgentRun( + history: SelectMessage[], +): { + jafHistory: JAFMessage[] + llmHistory: Message[] +} { + const filtered = history + .filter((msg) => !msg?.errorMessage) + .filter( + (msg) => !(msg.messageRole === MessageRole.Assistant && !msg.message), + ) + .filter( + (msg) => + msg.messageRole === MessageRole.User || + msg.messageRole === MessageRole.Assistant, + ) + + const toText = (msg: SelectMessage) => normalizeUserMessageForHistory(msg) + + return { + jafHistory: filtered.map((msg) => ({ + role: msg.messageRole === MessageRole.Assistant ? "assistant" : "user", + content: toText(msg), + })), + llmHistory: filtered.map((msg) => ({ + role: + msg.messageRole === MessageRole.Assistant + ? ConversationRole.ASSISTANT + : ConversationRole.USER, + content: [{ text: toText(msg) }], + })), + } +} + +export const __messageAgentsHistoryInternals = { + normalizeUserMessageForHistory, + buildConversationHistoryForAgentRun, +} + const RECENT_IMAGE_WINDOW = 2 function mergeFragmentLists( @@ -634,33 +726,6 @@ function getMetadataValue( return undefined } -function formatPlanForPrompt(plan: PlanState | null): string { - if (!plan) return "" - const lines = [`Goal: ${plan.goal}`] - plan.subTasks.forEach((task, idx) => { - const icon = - task.status === "completed" - ? "✓" - : task.status === "in_progress" - ? "→" - : task.status === "failed" - ? "✗" - : task.status === "blocked" - ? "!" - : "○" - const baseLine = `${idx + 1}. [${icon}] ${task.description}` - const detailParts: string[] = [] - if (task.result) detailParts.push(`Result: ${task.result}`) - if (task.toolsRequired?.length) { - detailParts.push(`Tools: ${task.toolsRequired.join(", ")}`) - } - lines.push( - detailParts.length > 0 ? `${baseLine}\n ${detailParts.join(" | ")}` : baseLine - ) - }) - return lines.join("\n") -} - function selectActiveSubTaskId(plan: PlanState | null): string | null { if (!plan || !Array.isArray(plan.subTasks) || plan.subTasks.length === 0) { return null @@ -772,168 +837,6 @@ function advancePlanAfterTool( } } -function formatClarificationsForPrompt( - clarifications: AgentRunContext["clarifications"] -): string { - if (!clarifications?.length) return "" - const formatted = clarifications - .map( - (clarification, idx) => - `${idx + 1}. Q: ${clarification.question}\n A: ${clarification.answer}` - ) - .join("\n") - return formatted -} - -function buildFinalSynthesisPayload( - context: AgentRunContext, - fragmentsLimit = Math.max(12, context.allFragments.length || 1) -): { systemPrompt: string; userMessage: string } { - const fragments = context.allFragments - const fragmentsSection = answerContextMapFromFragments(fragments, fragmentsLimit) - const planSection = formatPlanForPrompt(context.plan) - const clarificationSection = formatClarificationsForPrompt(context.clarifications) - const workspaceSection = context.userContext?.trim() - ? `Workspace Context:\n${context.userContext}` - : "" - - const parts = [ - `User Question:\n${context.message.text}`, - planSection ? `Execution Plan Snapshot:\n${planSection}` : "", - clarificationSection ? `Clarifications Resolved:\n${clarificationSection}` : "", - workspaceSection, - fragmentsSection, - ].filter(Boolean) - - const userMessage = parts.join("\n\n") - - const systemPrompt = ` -### Mission -- Deliver the user's final answer using the conversation, plan snapshot, clarifications, workspace context, context fragments, and supplied images; never plan or call tools. - -### Evidence Intake -- Prioritize the highest-signal fragments, but pull any supporting fragment that improves accuracy. -- Only draw on context that directly answers the user's question; ignore unrelated fragments even if they were retrieved earlier. -- Treat delegated-agent outputs as citeable fragments; reference them like any other context entry. -- Describe evidence gaps plainly before concluding; never guess. -- Extract actionable details from provided images and cite them via their fragment indices. - -### Response Construction -- Lead with the conclusion, then stack proof underneath. -- Organize output into tight sections (e.g., **Summary**, **Proof**, **Next Steps** when relevant); omit empty sections. -- Never mention internal tooling, planning logs, or this synthesis process. - -### Constraint Handling -- When the user asks for an action the system cannot execute (e.g., sending an email), deliver the closest actionable substitute (draft, checklist, explicit next steps) inside the answer. -- Pair the substitute with a concise explanation of the limitation and the manual action the user must take. - -### File & Chunk Formatting (CRITICAL) -- Each file starts with a header line exactly like: - index {docId} {file context begins here...} -- \`docId\` is a unique identifier for that file (e.g., 0, 1, 2, etc.). -- Inside the file context, text is split into chunks. -- Each chunk might begin with a bracketed numeric index, e.g.: [0], [1], [2], etc. -- This is the chunk index within that file, if it exists. - -### Guidelines for Response -1. Data Interpretation: - - Use ONLY the provided files and their chunks as your knowledge base. - - Treat every file header \`index {docId} ...\` as the start of a new document. - - Treat every bracketed number like [0], [1], [2] as the authoritative chunk index within that document. - - If dates exist, interpret them relative to the user's timezone when paraphrasing. -2. Response Structure: - - Start with the most relevant facts from the chunks across files. - - Keep order chronological when it helps comprehension. - - Every factual statement MUST cite the exact chunk it came from using the format: - K[docId_chunkIndex] - where: - - \`docId\` is taken from the file header line ("index {docId} ..."). - - \`chunkIndex\` is the bracketed number prefixed on that chunk within the same file. - - Examples: - - Single citation: "X is true K[12_3]." - - Two citations in one sentence (from different files or chunks): "X K[12_3] and Y K[7_0]." - - Use at most 1-2 citations per sentence; NEVER add more than 2 for one sentence. -3. Citation Rules (DOCUMENT+CHUNK LEVEL ONLY): - - ALWAYS cite at the chunk level with the K[docId_chunkIndex] format. - - Every chunk level citation must start with the K prefix eg. K[12_3] K[7_0] correct, but K[12_3] [7_0] is incorrect. - - Place the citation immediately after the relevant claim. - - Do NOT group indices inside one set of brackets (WRONG: "K[12_3,7_1]"). - - If a sentence draws on two distinct chunks (possibly from different files), include two separate citations inline, e.g., "... K[12_3] ... K[7_1]". - - Only cite information that appears verbatim or is directly inferable from the cited chunk. - - If you cannot ground a claim to a specific chunk, do not make the claim. -4. Quality Assurance: - - Cross-check across multiple chunks/files when available and briefly note inconsistencies if they exist. - - Keep tone professional and concise. - - Acknowledge gaps if the provided chunks don't contain enough detail. - -### Tone & Delivery -- Answer with confident, declarative, verb-first sentences that use concrete nouns. -- Highlight key deliverables using **bold** labels or short lists; keep wording razor-concise. -- Ask one targeted follow-up question only if missing info blocks action. - -### Tool Spotlighting -- Reference critical tool outputs explicitly, e.g., "**Slack Search:** Ops escalated the RCA at 09:42 [2]." -- Explain why each highlighted tool mattered so reviewers see coverage breadth. -- When multiple tools contribute, show the sequence, e.g., "**Vespa Search:** context -> **Sheet Lookup:** metrics." - -### Finish -- Close with a single sentence confirming completion or the next action you recommend. -`.trim() - - return { systemPrompt, userMessage } -} - -function selectImagesForFinalSynthesis( - context: AgentRunContext -): { - selected: string[] - total: number - dropped: string[] - userAttachmentCount: number -} { - const images = context.allImages - const total = images.length - if (!IMAGE_CONTEXT_CONFIG.enabled || total === 0) { - return { selected: [], total, dropped: [], userAttachmentCount: 0 } - } - - const attachments = images.filter((img) => img.isUserAttachment) - const nonAttachments = images - .filter((img) => !img.isUserAttachment) - .sort((a, b) => { - const ageA = context.turnCount - a.addedAtTurn - const ageB = context.turnCount - b.addedAtTurn - return ageA - ageB - }) - - const prioritized = [...attachments, ...nonAttachments] - const uniqueNames: string[] = [] - const seen = new Set() - for (const image of prioritized) { - if (seen.has(image.fileName)) continue - seen.add(image.fileName) - uniqueNames.push(image.fileName) - } - - let selected = uniqueNames - let dropped: string[] = [] - - if ( - IMAGE_CONTEXT_CONFIG.maxImagesPerCall > 0 && - uniqueNames.length > IMAGE_CONTEXT_CONFIG.maxImagesPerCall - ) { - selected = uniqueNames.slice(0, IMAGE_CONTEXT_CONFIG.maxImagesPerCall) - dropped = uniqueNames.slice(IMAGE_CONTEXT_CONFIG.maxImagesPerCall) - } - - return { - selected, - total, - dropped, - userAttachmentCount: attachments.length, - } -} - function buildAttachmentToolMessage( fragments: MinimalAgentFragment[], summary: string @@ -964,6 +867,7 @@ function initializeAgentContext( options?: { userContext?: string agentPrompt?: string + dedicatedAgentSystemPrompt?: string workspaceNumericId?: number chatId?: number stopController?: AbortController @@ -1003,6 +907,7 @@ function initializeAgentContext( currentSubTask: null, userContext: options?.userContext ?? "", agentPrompt: options?.agentPrompt, + dedicatedAgentSystemPrompt: options?.dedicatedAgentSystemPrompt, clarifications: [], ambiguityResolved: false, toolCallHistory: [], @@ -1202,6 +1107,7 @@ type ChatBootstrapParams = { type ChatBootstrapResult = { chat: SelectChat userMessage: SelectMessage + conversationHistory: SelectMessage[] attachmentError?: Error } @@ -1254,7 +1160,12 @@ async function ensureChatAndPersistUserMessage( } } - return { chat, userMessage, attachmentError: attachmentError ?? undefined } + return { + chat, + userMessage, + conversationHistory: [], + attachmentError: attachmentError ?? undefined, + } } const chat = await updateChatByExternalIdWithAuth( @@ -1263,6 +1174,11 @@ async function ensureChatAndPersistUserMessage( String(params.email), {} ) + const conversationHistory = await getChatMessagesWithAuth( + tx, + String(incomingChatId), + String(params.email), + ) const messageInsert = { chatId: chat.id, @@ -1290,7 +1206,12 @@ async function ensureChatAndPersistUserMessage( } } - return { chat, userMessage, attachmentError: attachmentError ?? undefined } + return { + chat, + userMessage, + conversationHistory, + attachmentError: attachmentError ?? undefined, + } }) } @@ -1741,7 +1662,7 @@ export async function afterToolExecutionHook( } // LOG: Context extraction results - loggerWithChild({ email: context.user.email }).debug( + loggerWithChild({ email: context.user.email }).info( { toolName, totalContextsExtracted: contexts.length, @@ -1756,7 +1677,7 @@ export async function afterToolExecutionHook( ) // LOG: Filtering results - loggerWithChild({ email: context.user.email }).debug( + loggerWithChild({ email: context.user.email }).info( { toolName, totalContexts: contexts.length, @@ -1775,14 +1696,21 @@ export async function afterToolExecutionHook( { toolName } ) - const contextStrings = filteredContexts.map( - (v: MinimalAgentFragment) => { - context.seenDocuments.add(v.id) - return ` - title: ${v.source.title}\n - content: ${v.content}\n - ` - } + const metadataConstraints = extractMetadataConstraintsFromUserMessage(userMessage) + const { + rankedCandidates, + hasConstraints: hasMetadataConstraints, + hasCompliantCandidates, + } = rankFragmentsByMetadataConstraints(filteredContexts, metadataConstraints) + const rankingCandidates = rankedCandidates.map((candidate) => candidate.fragment) + const strictNoCompliantCandidates = + hasMetadataConstraints && + metadataConstraints.strict && + !hasCompliantCandidates + + const contextStrings = rankingCandidates.map( + (fragment: MinimalAgentFragment, index: number) => + formatFragmentWithMetadata(fragment, index) ) // LOG: Prepared context strings for ranking @@ -1791,11 +1719,40 @@ export async function afterToolExecutionHook( toolName, contextStringsCount: contextStrings.length, userMessage: userMessage.slice(0, 200), + hasMetadataConstraints, + metadataIncludeTerms: metadataConstraints.includeTerms, + metadataExcludeTerms: metadataConstraints.excludeTerms, + metadataStrict: metadataConstraints.strict, + compliantCandidateCount: rankedCandidates.filter((v) => v.compliant).length, contextStringsSample: contextStrings.slice(0, 2).map(s => s.slice(0, 150)), }, "[afterToolExecutionHook] Prepared context strings for select_best_documents" ) + if (hasMetadataConstraints) { + await streamReasoningStep( + reasoningEmitter, + strictNoCompliantCandidates + ? "Strict metadata constraints detected and no compliant documents were found." + : hasCompliantCandidates + ? "Applied metadata constraints from the user request before ranking." + : "Detected metadata constraints in the user request but found no clearly compliant metadata matches.", + { + toolName, + detail: [ + metadataConstraints.includeTerms.length + ? `include=${metadataConstraints.includeTerms.join("|")}` + : "", + metadataConstraints.excludeTerms.length + ? `exclude=${metadataConstraints.excludeTerms.join("|")}` + : "", + ] + .filter(Boolean) + .join(" "), + } + ) + } + try { // LOG: Calling extractBestDocumentIndexes const rankingModelId = (context.modelId as Models) || config.defaultBestModel @@ -1815,6 +1772,29 @@ export async function afterToolExecutionHook( selectionSpan.setAttribute("context_count", contextStrings.length) let bestDocIndexes: number[] = [] try { + const rankingMessages = withAgentSystemPromptMessage( + messagesWithNoErrResponse, + context.dedicatedAgentSystemPrompt + ) + const rankingSystemPrompt = extractBestDocumentsPrompt( + userMessage, + contextStrings + ) + loggerWithChild({ email: context.user.email }).debug( + { + toolName, + modelId: rankingModelId, + rankingSystemPrompt, + rankingMessages, + }, + "[afterToolExecutionHook][select_best_documents] FINAL LLM PAYLOAD" + ) + selectionSpan.setAttribute( + "has_agent_system_prompt_snapshot", + !!sanitizeAgentSystemPromptSnapshot( + context.dedicatedAgentSystemPrompt + ) + ) bestDocIndexes = await extractBestDocumentIndexes( userMessage, contextStrings, @@ -1823,7 +1803,7 @@ export async function afterToolExecutionHook( json: false, stream: false, }, - messagesWithNoErrResponse + rankingMessages ) selectionSpan.setAttribute("selected_count", bestDocIndexes.length) } catch (error) { @@ -1846,117 +1826,168 @@ export async function afterToolExecutionHook( ) if (bestDocIndexes.length > 0) { - const selectedDocs: MinimalAgentFragment[] = [] + let selectedDocs: MinimalAgentFragment[] = [] bestDocIndexes.forEach((idx) => { - if (idx >= 1 && idx <= filteredContexts.length) { - const doc: MinimalAgentFragment = filteredContexts[idx - 1] + if (idx >= 1 && idx <= rankingCandidates.length) { + const doc: MinimalAgentFragment = rankingCandidates[idx - 1] selectedDocs.push(doc) } }) - // LOG: Document selection results - loggerWithChild({ email: context.user.email }).debug( - { - toolName, - selectedDocsCount: selectedDocs.length, - selectedDocIds: selectedDocs.map(d => d.id), - selectedDocTitles: selectedDocs.map(d => d.source.title || "untitled"), - }, - "[afterToolExecutionHook][select_best_documents] Selected documents after filtering" + selectedDocs = enforceMetadataConstraintsOnSelection( + selectedDocs, + rankedCandidates, + metadataConstraints ) - await streamReasoningStep( - reasoningEmitter, - `Filtered down to ${selectedDocs.length} best document${ - selectedDocs.length === 1 ? "" : "s" - } for analysis.`, - { toolName, detail: selectedDocs.map((doc) => doc.source.title || doc.id).join(", ") } - ) + if (selectedDocs.length === 0) { + const fallbackWhenSelectionEmpty = + strictNoCompliantCandidates + ? [] + : hasMetadataConstraints && + metadataConstraints.strict && + hasCompliantCandidates + ? rankedCandidates + .filter((candidate) => candidate.compliant) + .map((candidate) => candidate.fragment) + : rankingCandidates + loggerWithChild({ email: context.user.email }).debug( + { + toolName, + fallbackContextsCount: fallbackWhenSelectionEmpty.length, + }, + "[afterToolExecutionHook] No contexts survived selection enforcement; using metadata-aware fallback" + ) + addToolFragments(fallbackWhenSelectionEmpty) + } else { + // LOG: Document selection results + loggerWithChild({ email: context.user.email }).debug( + { + toolName, + selectedDocsCount: selectedDocs.length, + selectedDocIds: selectedDocs.map((d) => d.id), + selectedDocTitles: selectedDocs.map( + (d) => d.source.title || "untitled" + ), + }, + "[afterToolExecutionHook][select_best_documents] Selected documents after filtering" + ) - let fragmentsForResult = selectedDocs + await streamReasoningStep( + reasoningEmitter, + `Filtered down to ${selectedDocs.length} best document${ + selectedDocs.length === 1 ? "" : "s" + } for analysis.`, + { + toolName, + detail: selectedDocs + .map((doc) => doc.source.title || doc.id) + .join(", "), + } + ) - if (IMAGE_CONTEXT_CONFIG.enabled && selectedDocs.length > 0) { - const vespaLikeResults = selectedDocs.map((doc, idx) => ({ - id: doc.id, - relevance: 0, - fields: { docId: doc.source.docId }, - })) as unknown as VespaSearchResult[] + let fragmentsForResult = selectedDocs - const combinedContext = selectedDocs - .map((doc) => doc.content) - .join("\n") + if (IMAGE_CONTEXT_CONFIG.enabled && selectedDocs.length > 0) { + const vespaLikeResults = selectedDocs.map((doc, idx) => ({ + id: doc.id, + relevance: 0, + fields: { docId: doc.source.docId }, + })) as unknown as VespaSearchResult[] - const { imageFileNames: extractedImages } = extractImageFileNames( - combinedContext, - vespaLikeResults - ) + const combinedContext = selectedDocs + .map((doc) => doc.content) + .join("\n") - if (extractedImages.length > 0) { - const turnForImages = Math.max( - context.turnCount ?? MIN_TURN_NUMBER, - MIN_TURN_NUMBER - ) - const imageSpan = getTracer("chat").startSpan( - "tool_image_extraction" - ) - imageSpan.setAttribute("tool_name", toolName) - imageSpan.setAttribute("turn_number", turnForImages) - imageSpan.setAttribute("extracted_count", extractedImages.length) - imageSpan.setAttribute( - "image_names_preview", - extractedImages.slice(0, 5).join(",") + const { imageFileNames: extractedImages } = extractImageFileNames( + combinedContext, + vespaLikeResults ) - // LOG: Before attaching images - loggerWithChild({ email: context.user.email }).debug( - { - toolName, - extractedImagesCount: extractedImages.length, - extractedImages, - turnForImages, - selectedDocsCount: selectedDocs.length, - }, - "[afterToolExecutionHook] Extracted images from fragments - preparing to attach" - ) + if (extractedImages.length > 0) { + const turnForImages = Math.max( + context.turnCount ?? MIN_TURN_NUMBER, + MIN_TURN_NUMBER + ) + const imageSpan = getTracer("chat").startSpan( + "tool_image_extraction" + ) + imageSpan.setAttribute("tool_name", toolName) + imageSpan.setAttribute("turn_number", turnForImages) + imageSpan.setAttribute("extracted_count", extractedImages.length) + imageSpan.setAttribute( + "image_names_preview", + extractedImages.slice(0, 5).join(",") + ) - fragmentsForResult = attachImagesToFragments(selectedDocs, extractedImages, { - turnNumber: turnForImages, - sourceToolName: toolName, - isUserAttachment: false, - }) - imageSpan.setAttribute( - "fragments_with_images", - fragmentsForResult.filter( - (fragment) => fragment.images && fragment.images.length > 0 - ).length - ) - imageSpan.end() + // LOG: Before attaching images + loggerWithChild({ email: context.user.email }).debug( + { + toolName, + extractedImagesCount: extractedImages.length, + extractedImages, + turnForImages, + selectedDocsCount: selectedDocs.length, + }, + "[afterToolExecutionHook] Extracted images from fragments - preparing to attach" + ) - // LOG: After attaching images - loggerWithChild({ email: context.user.email }).debug( - { - toolName, - attachedImagesCount: extractedImages.length, - fragmentsWithImages: fragmentsForResult.filter(f => f.images && f.images.length > 0).length, - }, - `[afterToolExecutionHook] Tracked ${extractedImages.length} image reference${ - extractedImages.length === 1 ? "" : "s" - } from ${toolName} on turn ${turnForImages}` - ) + fragmentsForResult = attachImagesToFragments( + selectedDocs, + extractedImages, + { + turnNumber: turnForImages, + sourceToolName: toolName, + isUserAttachment: false, + } + ) + imageSpan.setAttribute( + "fragments_with_images", + fragmentsForResult.filter( + (fragment) => fragment.images && fragment.images.length > 0 + ).length + ) + imageSpan.end() + + // LOG: After attaching images + loggerWithChild({ email: context.user.email }).debug( + { + toolName, + attachedImagesCount: extractedImages.length, + fragmentsWithImages: fragmentsForResult.filter( + (f) => f.images && f.images.length > 0 + ).length, + }, + `[afterToolExecutionHook] Tracked ${extractedImages.length} image reference${ + extractedImages.length === 1 ? "" : "s" + } from ${toolName} on turn ${turnForImages}` + ) + } } - } - addToolFragments(fragmentsForResult) + addToolFragments(fragmentsForResult) + } } else { + const metadataFilteredFallback = + strictNoCompliantCandidates + ? [] + : hasMetadataConstraints && metadataConstraints.strict && hasCompliantCandidates + ? rankedCandidates + .filter((candidate) => candidate.compliant) + .map((candidate) => candidate.fragment) + : rankingCandidates loggerWithChild({ email: context.user.email }).debug( { toolName, filteredContextsCount: filteredContexts.length, + fallbackContextsCount: metadataFilteredFallback.length, }, - "[afterToolExecutionHook] Document ranking returned no results; retaining all filtered contexts" + strictNoCompliantCandidates + ? "[afterToolExecutionHook] Document ranking returned no results and strict metadata constraints had no compliant contexts" + : "[afterToolExecutionHook] Document ranking returned no results; retaining all filtered contexts" ) - addToolFragments(filteredContexts) + addToolFragments(metadataFilteredFallback) } } catch (error) { // LOG: Error in document ranking @@ -1970,10 +2001,20 @@ export async function afterToolExecutionHook( ) await streamReasoningStep( reasoningEmitter, - "Context ranking failed, retaining all retrieved documents.", + strictNoCompliantCandidates + ? "Context ranking failed and strict metadata constraints had no compliant documents; returning no contexts." + : "Context ranking failed, retaining all retrieved documents.", { toolName } ) - addToolFragments(filteredContexts) + const fallbackContexts = + strictNoCompliantCandidates + ? [] + : hasMetadataConstraints && metadataConstraints.strict && hasCompliantCandidates + ? rankedCandidates + .filter((candidate) => candidate.compliant) + .map((candidate) => candidate.fragment) + : rankingCandidates + addToolFragments(fallbackContexts) } } } @@ -3070,37 +3111,14 @@ function createFinalSynthesisTool(): Tool { ) } - const streamAnswer = mutableContext.runtime?.streamAnswerText - if (!streamAnswer) { + if (!mutableContext.runtime?.streamAnswerText) { return ToolResponse.error( "EXECUTION_FAILED", "Streaming channel unavailable. Cannot deliver final answer." ) } - const { selected, total, dropped, userAttachmentCount } = - selectImagesForFinalSynthesis(context) - loggerWithChild({ email: context.user.email }).debug( - { - chatId: context.chat.externalId, - selectedImages: selected, - totalImages: total, - droppedImages: dropped, - userAttachmentCount, - }, - "[MessageAgents][FinalSynthesis] Image payload" - ) - - const { systemPrompt, userMessage } = buildFinalSynthesisPayload(context) const fragmentsCount = context.allFragments.length - loggerWithChild({ email: context.user.email }).debug( - { - chatId: context.chat.externalId, - finalSynthesisSystemPrompt: systemPrompt, - finalSynthesisUserMessage: userMessage, - }, - "[MessageAgents][FinalSynthesis] Full context payload" - ) mutableContext.finalSynthesis.requested = true mutableContext.finalSynthesis.suppressAssistantStreaming = true @@ -3108,93 +3126,19 @@ function createFinalSynthesisTool(): Tool { mutableContext.finalSynthesis.streamedText = "" await mutableContext.runtime?.emitReasoning?.({ - text: `Initiating final synthesis with ${fragmentsCount} context fragments and ${selected.length}/${total} images (${userAttachmentCount} user attachments).`, + text: `Initiating final synthesis with ${fragmentsCount} context fragments.`, step: { type: AgentReasoningStepType.LogMessage }, }) - const logger = loggerWithChild({ email: context.user.email }) - if (dropped.length > 0) { - logger.info( - { - droppedCount: dropped.length, - limit: IMAGE_CONTEXT_CONFIG.maxImagesPerCall, - totalImages: total, - }, - "Final synthesis image limit enforced; dropped oldest references." - ) - } - - const modelId = - (context.modelId as Models) || (defaultBestModel as Models) || Models.Gpt_4o - const modelParams: ModelParams = { - modelId, - systemPrompt, - stream: true, - temperature: 0.2, - max_new_tokens: context.maxOutputTokens ?? 1500, - imageFileNames: selected, - } - - const finalUserPrompt = `${userMessage}\n\nSynthesize the final answer using the evidence above.` - const messages: Message[] = [ - { - role: ConversationRole.USER, - content: [{ text: finalUserPrompt }], - }, - ] - Logger.debug( - { - email: context.user.email, - chatId: context.chat.externalId, - fragmentsCount: context.allFragments.length, - planPresent: !!context.plan, - clarificationsCount: context.clarifications.length, - toolOutputsThisTurn: context.currentTurnArtifacts.toolOutputs.length, - imageNames: selected, - }, - "[MessageAgents][FinalSynthesis] Context summary for synthesis call" - ) - - Logger.debug({ - email: context.user.email, - chatId: context.chat.externalId, - modelId, - systemPrompt, - messagesCount: messages.length, - imagesProvided: selected.length, - }, "[MessageAgents][FinalSynthesis] LLM call parameters") - - const provider = getProviderByModel(modelId) - let streamedCharacters = 0 - let estimatedCostUsd = 0 - try { - const iterator = provider.converseStream(messages, modelParams) - for await (const chunk of iterator) { - if (chunk.text) { - streamedCharacters += chunk.text.length - context.finalSynthesis.streamedText += chunk.text - await streamAnswer(chunk.text) - } - const chunkCost = chunk.metadata?.cost - if (typeof chunkCost === "number" && !Number.isNaN(chunkCost)) { - estimatedCostUsd += chunkCost - } - } - + const synthesisResult = await executeFinalSynthesis(mutableContext) context.finalSynthesis.completed = true - loggerWithChild({ email: context.user.email }).debug( - { - chatId: context.chat.externalId, - streamedCharacters, - estimatedCostUsd, - imagesProvided: selected, - }, - "[MessageAgents][FinalSynthesis] LLM call completed" - ) await context.runtime?.emitReasoning?.({ - text: "Final synthesis completed and streamed to the user.", + text: + synthesisResult.mode === "sectional" + ? "Sectional final synthesis completed and delivered to the user." + : "Final synthesis completed and streamed to the user.", step: { type: AgentReasoningStepType.LogMessage }, }) @@ -3203,20 +3147,20 @@ function createFinalSynthesisTool(): Tool { result: "Final answer streamed to user.", streamed: true, metadata: { - textLength: streamedCharacters, - totalImagesAvailable: total, - imagesProvided: selected.length, + textLength: synthesisResult.textLength, + totalImagesAvailable: synthesisResult.totalImagesAvailable, + imagesProvided: synthesisResult.imagesProvided, }, }, { - estimatedCostUsd, + estimatedCostUsd: synthesisResult.estimatedCostUsd, } ) } catch (error) { context.finalSynthesis.suppressAssistantStreaming = false context.finalSynthesis.requested = false context.finalSynthesis.completed = false - logger.error( + loggerWithChild({ email: context.user.email }).error( { err: error instanceof Error ? error.message : String(error) }, "Final synthesis tool failed." ) @@ -3735,6 +3679,7 @@ export async function MessageAgents(c: Context): Promise { let lastPersistedMessageId = 0 let lastPersistedMessageExternalId = "" let attachmentStorageError: Error | null = null + let previousConversationHistory: SelectMessage[] = [] try { const bootstrap = await ensureChatAndPersistUserMessage({ @@ -3752,6 +3697,7 @@ export async function MessageAgents(c: Context): Promise { lastPersistedMessageId = bootstrap.userMessage.id as number lastPersistedMessageExternalId = String(bootstrap.userMessage.externalId) attachmentStorageError = bootstrap.attachmentError ?? null + previousConversationHistory = bootstrap.conversationHistory ?? [] const chatAgentId = chatRecord.agentId ? String(chatRecord.agentId) : undefined @@ -3779,6 +3725,10 @@ export async function MessageAgents(c: Context): Promise { }) } rootSpan.setAttribute("chatId", String(chatRecord.externalId)) + rootSpan.setAttribute( + "conversation_history_count", + previousConversationHistory.length, + ) if ( resolvedAgentId && @@ -3803,6 +3753,10 @@ export async function MessageAgents(c: Context): Promise { } const hasExplicitAgent = Boolean(resolvedAgentId && agentPromptForLLM) + const dedicatedAgentSystemPrompt = + typeof agentRecord?.prompt === "string" && agentRecord.prompt.trim().length > 0 + ? agentRecord.prompt.trim() + : undefined const delegationEnabled = !hasExplicitAgent return streamSSE(c, async (stream) => { @@ -3879,6 +3833,7 @@ export async function MessageAgents(c: Context): Promise { userContext: userCtxString, workspaceNumericId: workspace.id, agentPrompt: agentPromptForLLM, + dedicatedAgentSystemPrompt, chatId: chatRecord.id as number, stopController, modelId: agenticModelId, @@ -4259,10 +4214,14 @@ export async function MessageAgents(c: Context): Promise { // Initialize run state const runId = generateRunId() const traceId = generateTraceId() + const { jafHistory, llmHistory } = buildConversationHistoryForAgentRun( + previousConversationHistory + ) const initialMessages: JAFMessage[] = [ + ...jafHistory, { role: "user", - content: message, // Use actual user message + content: message, }, ...initialSyntheticMessages, ] @@ -4279,6 +4238,8 @@ export async function MessageAgents(c: Context): Promise { jafStreamingSpan.setAttribute("chat_external_id", chatRecord.externalId) jafStreamingSpan.setAttribute("run_id", runId) jafStreamingSpan.setAttribute("trace_id", traceId) + jafStreamingSpan.setAttribute("history_message_count", jafHistory.length) + jafStreamingSpan.setAttribute("history_seeded", jafHistory.length > 0) let turnSpan: Span | undefined const endTurnSpan = () => { if (turnSpan) { @@ -4296,6 +4257,7 @@ export async function MessageAgents(c: Context): Promise { } const messagesWithNoErrResponse: Message[] = [ + ...llmHistory, { role: ConversationRole.USER, content: [{ text: message }], @@ -5736,6 +5698,10 @@ async function runDelegatedAgentWithMessageAgents( } const agentPromptForLLM = JSON.stringify(agentRecord) + const dedicatedAgentSystemPrompt = + typeof agentRecord.prompt === "string" && agentRecord.prompt.trim().length > 0 + ? agentRecord.prompt.trim() + : undefined const userCtxString = userContext(userAndWorkspace) const userTimezone = user.timeZone || "Asia/Kolkata" const dateForAI = getDateForAI({ userTimeZone: userTimezone }) @@ -5762,6 +5728,7 @@ async function runDelegatedAgentWithMessageAgents( { userContext: userCtxString, agentPrompt: agentPromptForLLM, + dedicatedAgentSystemPrompt, workspaceNumericId: workspace.id, stopSignal: params.stopSignal, modelId: delegateModelId, @@ -5808,14 +5775,15 @@ async function runDelegatedAgentWithMessageAgents( const gatheredFragmentsKeys = new Set() - const instructions = () => - buildAgentInstructions( + const instructions = () => { + return buildAgentInstructions( agentContext, allTools.map((tool) => tool.schema.name), dateForAI, agentPromptForLLM, false ) + } const jafAgent: JAFAgent = { name: "xyne-delegate", diff --git a/server/logger/index.ts b/server/logger/index.ts index 119615143..6c1ad93a5 100644 --- a/server/logger/index.ts +++ b/server/logger/index.ts @@ -15,6 +15,16 @@ import { object } from "zod" export { Subsystem } const { JwtPayloadKey } = config +const defaultLogLevel = "info" + +const isValidLogLevel = (level: string) => level in levels.values + +const getConfiguredLogLevel = () => { + const configuredLevel = + process.env.NODE_LOG_LEVEL ?? process.env.LOG_LEVEL ?? defaultLogLevel + + return isValidLogLevel(configuredLevel) ? configuredLevel : defaultLogLevel +} const humanize = (times: string[]) => { const [delimiter, separator] = [",", "."] @@ -35,6 +45,7 @@ const time = (start: number) => { export const getLogger = (loggerType: Subsystem) => { const isProduction = process.env.NODE_ENV === "production" + const level = getConfiguredLogLevel() if (isProduction) { const destination = pino.destination(1) // stdout @@ -42,6 +53,7 @@ export const getLogger = (loggerType: Subsystem) => { return pino( { name: loggerType, + level, timestamp: false, formatters: { level: (label) => ({ level: label }), @@ -75,6 +87,7 @@ export const getLogger = (loggerType: Subsystem) => { // Dev logger return pino({ name: loggerType, + level, transport: { target: "pino-pretty", options: { diff --git a/server/shared/types.ts b/server/shared/types.ts index 325e87dc1..34fadab1c 100644 --- a/server/shared/types.ts +++ b/server/shared/types.ts @@ -827,6 +827,7 @@ export interface ModelConfiguration { websearch: boolean deepResearch: boolean description: string + maxInputTokens?: number } export const getDocumentSchema = z.object({ docId: z.string().min(1), diff --git a/server/tests/finalAnswerSynthesis.test.ts b/server/tests/finalAnswerSynthesis.test.ts new file mode 100644 index 000000000..c2bca6fb6 --- /dev/null +++ b/server/tests/finalAnswerSynthesis.test.ts @@ -0,0 +1,183 @@ +import { describe, expect, test } from "bun:test" +import type { AgentRunContext } from "@/api/chat/agent-schemas" +import { + __finalAnswerSynthesisInternals, + buildFinalSynthesisPayload, +} from "@/api/chat/final-answer-synthesis" +import type { MinimalAgentFragment } from "@/api/chat/types" +import { Models } from "@/ai/types" +import { Apps } from "@xyne/vespa-ts/types" + +const baseFragment: MinimalAgentFragment = { + id: "doc-1", + content: "Quarterly ARR grew 12% and pipeline coverage improved.", + source: { + docId: "doc-1", + title: "ARR Summary", + url: "https://example.com/doc-1", + app: Apps.KnowledgeBase, + entity: "file" as any, + }, + confidence: 0.9, +} + +const createMockContext = (): AgentRunContext => ({ + user: { + email: "tester@example.com", + workspaceId: "workspace", + id: "user-1", + }, + chat: { + externalId: "chat-1", + metadata: {}, + }, + message: { + text: "How is ARR tracking?", + attachments: [], + timestamp: new Date().toISOString(), + }, + modelId: Models.Gpt_4o, + plan: null, + currentSubTask: null, + userContext: "", + agentPrompt: undefined, + dedicatedAgentSystemPrompt: undefined, + clarifications: [], + ambiguityResolved: true, + toolCallHistory: [], + seenDocuments: new Set(), + allFragments: [], + turnFragments: new Map(), + allImages: [], + imagesByTurn: new Map(), + recentImages: [], + currentTurnArtifacts: { + fragments: [], + expectations: [], + toolOutputs: [], + images: [], + }, + turnCount: 1, + totalLatency: 0, + totalCost: 0, + tokenUsage: { input: 0, output: 0 }, + availableAgents: [], + usedAgents: [], + enabledTools: new Set(), + delegationEnabled: true, + failedTools: new Map(), + retryCount: 0, + maxRetries: 3, + review: { + lastReviewTurn: null, + reviewFrequency: 5, + outstandingAnomalies: [], + clarificationQuestions: [], + lastReviewResult: null, + lockedByFinalSynthesis: false, + lockedAtTurn: null, + }, + decisions: [], + finalSynthesis: { + requested: false, + completed: false, + suppressAssistantStreaming: false, + streamedText: "", + ackReceived: false, + }, + stopRequested: false, +}) + +describe("final-answer-synthesis", () => { + test("builds deterministic fragment previews from raw fragment content", () => { + const preview = __finalAnswerSynthesisInternals.buildFragmentPreviewRecord( + { + ...baseFragment, + content: + " Quarterly ARR grew 12% and pipeline coverage improved.\n\nCustomers expanded seats. ", + source: { + ...baseFragment.source, + createdAt: "2026-03-08T09:00:00.000Z", + }, + }, + 7, + ) + + expect(preview.fragmentIndex).toBe(7) + expect(preview.docId).toBe("doc-1") + expect(preview.title).toBe("ARR Summary") + expect(preview.app).toBe(String(Apps.KnowledgeBase)) + expect(preview.previewText).toBe( + "Quarterly ARR grew 12% and pipeline coverage improved. Customers expanded seats.", + ) + expect(preview.timestamp).toBe("2026-03-08T09:00:00.000Z") + }) + + test("section payload keeps shared context and adds section-only instructions", () => { + const context = createMockContext() + context.dedicatedAgentSystemPrompt = + "You are an enterprise agent. Always use verified workspace evidence." + context.allFragments = [baseFragment] + + const payload = __finalAnswerSynthesisInternals.buildSectionAnswerPayload( + context, + [ + { + sectionId: 1, + title: "Summary", + objective: "Summarize the ARR status.", + }, + { + sectionId: 2, + title: "Evidence", + objective: "Provide supporting evidence.", + }, + ], + { + sectionId: 2, + title: "Evidence", + objective: "Provide supporting evidence.", + }, + [{ fragmentIndex: 3, fragment: baseFragment }], + ["3_doc-1_0"], + ) + + expect(payload.systemPrompt).toContain("Deliver only the assigned answer section") + expect(payload.userMessage).toContain( + "All Planned Sections (generated in parallel; a final ordered answer will be assembled later):", + ) + expect(payload.userMessage).toContain("Assigned Section:\n2. Evidence") + expect(payload.userMessage).toContain("Write only this section.") + expect(payload.userMessage).toContain("index 3 {file context begins here...}") + expect(payload.userMessage).toContain("Agent System Prompt Context:") + expect(payload.imageFileNames).toEqual(["3_doc-1_0"]) + }) + + test("switches to sectional mode when full final payload exceeds the model input budget", () => { + const context = createMockContext() + context.modelId = Models.Gpt_4 + context.allFragments = [ + { + ...baseFragment, + content: "A".repeat(40_000), + }, + ] + + const payload = buildFinalSynthesisPayload(context) + expect(payload.userMessage.length).toBeGreaterThan(0) + + const decision = __finalAnswerSynthesisInternals.decideSynthesisMode( + context, + Models.Gpt_4, + { + selected: [], + total: 0, + dropped: [], + userAttachmentCount: 0, + }, + ) + + expect(decision.mode).toBe("sectional") + expect(decision.estimatedInputTokens).toBeGreaterThan(decision.safeInputBudget) + }) +}) diff --git a/server/tests/messageAgentsFragments.test.ts b/server/tests/messageAgentsFragments.test.ts index 94ac196f4..809e8aa69 100644 --- a/server/tests/messageAgentsFragments.test.ts +++ b/server/tests/messageAgentsFragments.test.ts @@ -1,8 +1,11 @@ import { describe, expect, test } from "bun:test" import type { AgentRunContext } from "@/api/chat/agent-schemas" import { + __messageAgentsHistoryInternals, + __messageAgentsMetadataInternals, afterToolExecutionHook, buildDelegatedAgentFragments, + buildFinalSynthesisPayload, buildReviewPromptFromContext, } from "@/api/chat/message-agents" import { getRecentImagesFromContext } from "@/api/chat/runContextUtils" @@ -43,6 +46,7 @@ const createMockContext = (): AgentRunContext => ({ currentSubTask: null, userContext: "", agentPrompt: undefined, + dedicatedAgentSystemPrompt: undefined, clarifications: [], ambiguityResolved: true, toolCallHistory: [], @@ -146,6 +150,53 @@ describe("message-agents context tracking", () => { ) }) + test("afterToolExecutionHook enforces strict metadata constraints when no compliant docs exist", async () => { + const context = createMockContext() + const nonCompliantFragment: MinimalAgentFragment = { + ...baseFragment, + source: { + ...baseFragment.source, + title: "General Notes", + }, + } + + await afterToolExecutionHook( + "searchGlobal", + { + status: "success", + metadata: { + contexts: [nonCompliantFragment], + }, + data: { + result: "Found documents.", + }, + }, + { + toolCall: { id: "call-2" } as any, + args: { query: "notes" }, + state: { + context, + messages: [], + runId: createRunId("run-2"), + traceId: createTraceId("trace-2"), + currentAgentName: "xyne-agent", + turnCount: 1, + }, + agentName: "xyne-agent", + executionTime: 10, + status: "success", + }, + 'Answer only from source "Q4 Planning".', + [], + new Set(), + undefined, + context.turnCount + ) + + expect(context.allFragments).toHaveLength(0) + expect(context.currentTurnArtifacts.fragments).toHaveLength(0) + }) + test("buildReviewPromptFromContext includes plan, expectations, and image metadata", () => { const context = createMockContext() context.plan = { @@ -211,6 +262,48 @@ describe("message-agents context tracking", () => { expect(fragments[0].content).toContain("Delegate") }) + test("buildConversationHistoryForAgentRun normalizes context JSON and filters invalid turns", () => { + const { buildConversationHistoryForAgentRun } = + __messageAgentsHistoryInternals + + const history = [ + { + messageRole: "user", + message: '[{"type":"text","value":"Summarize"},{"type":"pill","value":{"title":"Q4 Plan"}}]', + fileIds: ["clf-1"], + errorMessage: "", + }, + { + messageRole: "assistant", + message: "Sure, sharing summary.", + fileIds: [], + errorMessage: "", + }, + { + messageRole: "assistant", + message: "", + fileIds: [], + errorMessage: "", + }, + { + messageRole: "user", + message: "bad turn", + fileIds: [], + errorMessage: "timeout", + }, + ] as any + + const { jafHistory, llmHistory } = buildConversationHistoryForAgentRun(history) + + expect(jafHistory).toHaveLength(2) + expect(jafHistory[0].role).toBe("user") + expect(jafHistory[0].content).toContain('User referred a file with title "Q4 Plan"') + expect(jafHistory[1].role).toBe("assistant") + expect(llmHistory).toHaveLength(2) + expect((llmHistory[0] as any).role).toBe("user") + expect((llmHistory[1] as any).role).toBe("assistant") + }) + test("getRecentImagesFromContext prioritizes attachments and last two turns", () => { const context = createMockContext() context.currentTurnArtifacts.images.push({ @@ -266,4 +359,74 @@ describe("message-agents context tracking", () => { expect(contexts[0].source.title).toContain("Connector 123") expect(contexts[0].content).toContain("MCP response text") }) + + test("metadata constraints are inferred and ranked generically from user request", () => { + const constraints = + __messageAgentsMetadataInternals.extractMetadataConstraintsFromUserMessage( + 'Answer only from source "Q4 Planning" and exclude "Legacy Notes".' + ) + + expect(constraints.strict).toBe(true) + expect(constraints.includeTerms).toContain("q4 planning") + expect(constraints.excludeTerms).toContain("legacy notes") + + const matchingFragment: MinimalAgentFragment = { + ...baseFragment, + id: "doc-2", + source: { + ...baseFragment.source, + title: "Q4 Planning", + }, + } + const excludedFragment: MinimalAgentFragment = { + ...baseFragment, + id: "doc-3", + source: { + ...baseFragment.source, + title: "Legacy Notes", + }, + } + + const ranked = __messageAgentsMetadataInternals.rankFragmentsByMetadataConstraints( + [excludedFragment, matchingFragment], + constraints + ) + expect(ranked.hasConstraints).toBe(true) + expect(ranked.hasCompliantCandidates).toBe(true) + expect(ranked.rankedCandidates[0].fragment.id).toBe("doc-2") + }) + + test("final synthesis payload includes metadata-enriched fragment context", () => { + const context = createMockContext() + context.dedicatedAgentSystemPrompt = + "You are an enterprise agent. Always use verified workspace evidence." + context.allFragments = [ + { + ...baseFragment, + source: { + ...baseFragment.source, + page_title: "Quarterly Planning Sheet", + status: "Open", + }, + }, + ] + + const payload = buildFinalSynthesisPayload(context) + expect(payload.userMessage).toContain("Agent System Prompt Context:") + expect(payload.userMessage).toContain("This is the system prompt of agent:") + expect(payload.userMessage).toContain("") + expect(payload.userMessage).toContain( + "You are an enterprise agent. Always use verified workspace evidence." + ) + expect(payload.userMessage).not.toContain( + "You are Xyne, an enterprise search assistant with agentic capabilities." + ) + expect(payload.userMessage).toContain("") + expect(payload.userMessage).toContain("Context Fragments:") + expect(payload.userMessage).toContain("index 1 {file context begins here...}") + expect(payload.userMessage).toContain("- title: ARR Summary") + expect(payload.userMessage).toContain("- page_title: Quarterly Planning Sheet") + expect(payload.userMessage).toContain("Content:") + expect(payload.userMessage).toContain("Quarterly ARR grew 12%") + }) })