Skip to content

Fix/follow up context #1082

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Oct 13, 2025
Merged

Fix/follow up context #1082

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
8ea8b73
fix: XYNE-159 fixed followup context for attached files
Himanshvarma Oct 10, 2025
7fc3a69
fix: XYNE-159 fixed ai comments
Himanshvarma Oct 10, 2025
5c38e8a
fix: XYNE-159 resolved ai comments
Himanshvarma Oct 10, 2025
3d7f351
Merge branch 'main' into fix/followUpContext
Himanshvarma Oct 13, 2025
88bd2eb
fix: XYNE-159 removed error message context gathering
Himanshvarma Oct 13, 2025
da87ff4
Merge branch 'main' of https://github.com/xynehq/xyne into fix/follow…
Himanshvarma Oct 13, 2025
6d4b37f
Update chat.ts
Himanshvarma Oct 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 107 additions & 54 deletions server/api/chat/agents.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,6 @@ import {
} from "@xyne/vespa-ts/types"
import { APIError } from "openai"
import { insertChatTrace } from "@/db/chatTrace"
import type { AttachmentMetadata } from "@/shared/types"
import { storeAttachmentMetadata } from "@/db/attachment"
import { parseAttachmentMetadata } from "@/utils/parseAttachment"
import { isCuid } from "@paralleldrive/cuid2"
Expand Down Expand Up @@ -3755,32 +3754,6 @@ export const AgentMessageApi = async (c: Context) => {
})
}

const filteredMessages = messages
.slice(0, messages.length - 1)
.filter(
(msg) =>
!(msg.messageRole === MessageRole.Assistant && !msg.message),
)

// Check for follow-up context carry-forward
const lastIdx = filteredMessages.length - 1
const workingSet = collectFollowupContext(filteredMessages, lastIdx)

const hasCarriedContext =
workingSet.fileIds.length > 0 ||
workingSet.attachmentFileIds.length > 0
if (hasCarriedContext) {
fileIds = Array.from(new Set([...fileIds, ...workingSet.fileIds]))
imageAttachmentFileIds = Array.from(
new Set([
...imageAttachmentFileIds,
...workingSet.attachmentFileIds,
]),
)
loggerWithChild({ email: email }).info(
`Carried forward context from follow-up: ${JSON.stringify(workingSet)}`,
)
}
if (
(fileIds && fileIds?.length > 0) ||
(imageAttachmentFileIds && imageAttachmentFileIds?.length > 0)
Expand Down Expand Up @@ -4163,6 +4136,7 @@ export const AgentMessageApi = async (c: Context) => {
offset: 0,
}
let parsed = {
isFollowUp: false,
answer: "",
queryRewrite: "",
temporalDirection: null,
Expand Down Expand Up @@ -4317,6 +4291,19 @@ export const AgentMessageApi = async (c: Context) => {
parsed.queryRewrite,
)
conversationSpan.end()
let classification: TemporalClassifier & QueryRouterResponse =
{
direction: parsed.temporalDirection,
type: parsed.type as QueryType,
filterQuery: parsed.filter_query,
isFollowUp: parsed.isFollowUp,
filters: {
...(parsed?.filters ?? {}),
apps: parsed.filters?.apps || [],
entities: parsed.filters?.entities as any,
intent: parsed.intent || {},
},
}

if (parsed.answer === null || parsed.answer === "") {
const ragSpan = streamSpan.startSpan("rag_processing")
Expand All @@ -4332,37 +4319,103 @@ export const AgentMessageApi = async (c: Context) => {
"There was no need for a query rewrite and there was no answer in the conversation, applying RAG",
)
}
const classification: TemporalClassifier & QueryRouterResponse =
{
direction: parsed.temporalDirection,
type: parsed.type as QueryType,
filterQuery: parsed.filter_query,
filters: {
...(parsed?.filters ?? {}),
apps: parsed.filters?.apps || [],
entities: parsed.filters?.entities as any,
intent: parsed.intent || {},
},
}

Logger.info(
loggerWithChild({ email: email }).info(
`Classifying the query as:, ${JSON.stringify(classification)}`,
)
const understandSpan = ragSpan.startSpan("understand_message")
const iterator = UnderstandMessageAndAnswer(
email,
ctx,
userMetadata,
message,
classification,
limitedMessages,
0.5,
userRequestsReasoning,
understandSpan,
agentPromptForLLM,
actualModelId,
pathExtractedInfo,

ragSpan.setAttribute(
"isFollowUp",
classification.isFollowUp ?? false,
)
const understandSpan = ragSpan.startSpan("understand_message")

let iterator:
| AsyncIterableIterator<
ConverseResponse & {
citation?: { index: number; item: any }
imageCitation?: ImageCitation
}
>
| undefined = undefined

if (messages.length < 2) {
classification.isFollowUp = false // First message or not enough history to be a follow-up
} else if (classification.isFollowUp) {
// Use the NEW classification that already contains:
// - Updated filters (with proper offset calculation)
// - Preserved app/entity from previous query
// - Updated count/pagination info
// - All the smart follow-up logic from the LLM

const filteredMessages = messages
.filter(
(msg) => !msg?.errorMessage,
)
.filter(
(msg) =>
!(msg.messageRole === MessageRole.Assistant && !msg.message),
)

// Check for follow-up context carry-forward
const workingSet = collectFollowupContext(filteredMessages);

const hasCarriedContext =
workingSet.fileIds.length > 0 ||
workingSet.attachmentFileIds.length > 0;
if (hasCarriedContext) {
fileIds = workingSet.fileIds;
imageAttachmentFileIds = workingSet.attachmentFileIds;
loggerWithChild({ email: email }).info(
`Carried forward context from follow-up: ${JSON.stringify(workingSet)}`,
);
}

if (fileIds && fileIds.length > 0 || imageAttachmentFileIds && imageAttachmentFileIds.length > 0) {
loggerWithChild({ email: email }).info(
`Follow-up query with file context detected. Using file-based context with NEW classification: ${JSON.stringify(classification)}, FileIds: ${JSON.stringify([fileIds, imageAttachmentFileIds])}`,
)
iterator = UnderstandMessageAndAnswerForGivenContext(
email,
ctx,
userMetadata,
message,
0.5,
fileIds as string[],
userRequestsReasoning,
understandSpan,
undefined,
imageAttachmentFileIds as string[],
agentPromptForLLM,
undefined,
actualModelId || config.defaultBestModel,
)
} else {
loggerWithChild({ email: email }).info(
`Follow-up query detected.`,
)
// Use the new classification directly - it already has all the smart follow-up logic
// No need to reuse old classification, the LLM has generated an updated one
}
}

// If no iterator was set above (non-file-context scenario), use the regular flow with the new classification
if(!iterator) {
iterator = UnderstandMessageAndAnswer(
email,
ctx,
userMetadata,
message,
classification,
limitedMessages,
0.5,
userRequestsReasoning,
understandSpan,
agentPromptForLLM,
actualModelId,
pathExtractedInfo,
)
}
stream.writeSSE({
event: ChatSSEvents.Start,
data: "",
Expand Down
112 changes: 52 additions & 60 deletions server/api/chat/chat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1972,23 +1972,14 @@ async function* generateAnswerFromGivenContext(
combinedSearchResponse.push(...results.root.children)
}
} else {
const collectionFileIds = fileIds.filter(
(fid) => fid.startsWith("clf-") || fid.startsWith("att_"),
)
const nonCollectionFileIds = fileIds.filter(
(fid) => !fid.startsWith("clf-") && !fid.startsWith("att_"),
)
if (nonCollectionFileIds && nonCollectionFileIds.length > 0) {
results = await searchVespaInFiles(
builtUserQuery,
email,
nonCollectionFileIds,
{
const collectionFileIds = fileIds.filter((fid) => fid.startsWith("clf-") || fid.startsWith("att_"))
const nonCollectionFileIds = fileIds.filter((fid) => !fid.startsWith("clf-") && !fid.startsWith("att"))
const attachmentFileIds = fileIds.filter((fid) => fid.startsWith("attf_"))
if(nonCollectionFileIds && nonCollectionFileIds.length > 0) {
results = await searchVespaInFiles(builtUserQuery, email, nonCollectionFileIds, {
limit: fileIds?.length,
alpha: userAlpha,
rankProfile: SearchModes.AttachmentRank,
},
)
})
if (results.root.children) {
combinedSearchResponse.push(...results.root.children)
}
Expand All @@ -2003,6 +1994,16 @@ async function* generateAnswerFromGivenContext(
combinedSearchResponse.push(...results.root.children)
}
}
if(attachmentFileIds && attachmentFileIds.length > 0) {
results = await searchVespaInFiles(builtUserQuery, email, attachmentFileIds, {
limit: fileIds?.length,
alpha: userAlpha,
rankProfile: SearchModes.attachmentRank,
})
if (results.root.children) {
combinedSearchResponse.push(...results.root.children)
}
}
}

// Apply intelligent chunk selection based on document relevance and chunk scores
Expand Down Expand Up @@ -4468,32 +4469,6 @@ export const MessageApi = async (c: Context) => {
})
}

let filteredMessages = messages
.slice(0, messages.length - 1)
.filter(
(msg) =>
!(msg.messageRole === MessageRole.Assistant && !msg.message),
)

// Check for follow-up context carry-forward
const lastIdx = filteredMessages.length - 1
const workingSet = collectFollowupContext(filteredMessages, lastIdx)

const hasCarriedContext =
workingSet.fileIds.length > 0 ||
workingSet.attachmentFileIds.length > 0
if (hasCarriedContext) {
fileIds = Array.from(new Set([...fileIds, ...workingSet.fileIds]))
imageAttachmentFileIds = Array.from(
new Set([
...imageAttachmentFileIds,
...workingSet.attachmentFileIds,
]),
)
loggerWithChild({ email: email }).info(
`Carried forward context from follow-up: ${JSON.stringify(workingSet)}`,
)
}
if (
(fileIds && fileIds?.length > 0) ||
(imageAttachmentFileIds && imageAttachmentFileIds?.length > 0)
Expand Down Expand Up @@ -4755,9 +4730,16 @@ export const MessageApi = async (c: Context) => {
streamSpan.end()
rootSpan.end()
} else {
filteredMessages = filteredMessages.filter(
(msg) => !msg?.errorMessage,
)
const filteredMessages = messages
.slice(0, messages.length - 1)
.filter(
(msg) => !msg?.errorMessage,
)
.filter(
(msg) =>
!(msg.messageRole === MessageRole.Assistant && !msg.message),
)

loggerWithChild({ email: email }).info(
"Checking if answer is in the conversation or a mandatory query rewrite is needed before RAG",
)
Expand Down Expand Up @@ -5232,34 +5214,44 @@ export const MessageApi = async (c: Context) => {
// - Updated count/pagination info
// - All the smart follow-up logic from the LLM

// Only check for fileIds if we need file context
const lastUserMessage = messages[messages.length - 3] // Assistant is at -2, last user is at -3
const parsedMessage =
selectMessageSchema.safeParse(lastUserMessage)

if (parsedMessage.error) {
loggerWithChild({ email: email }).error(
`Error while parsing last user message for file context check`,
const filteredMessages = messages
.filter(
(msg) => !msg?.errorMessage,
)
} else if (
parsedMessage.success &&
Array.isArray(parsedMessage.data.fileIds) &&
parsedMessage.data.fileIds.length // If the message contains fileIds then the follow up is must for @file
) {
.filter(
(msg) =>
!(msg.messageRole === MessageRole.Assistant && !msg.message),
)

// Check for follow-up context carry-forward
const workingSet = collectFollowupContext(filteredMessages);

const hasCarriedContext =
workingSet.fileIds.length > 0 ||
workingSet.attachmentFileIds.length > 0;
if (hasCarriedContext) {
fileIds = workingSet.fileIds;
imageAttachmentFileIds = workingSet.attachmentFileIds;
loggerWithChild({ email: email }).info(
`Follow-up query with file context detected. Using file-based context with NEW classification: ${JSON.stringify(classification)}, FileIds: ${JSON.stringify(parsedMessage.data.fileIds)}`,
`Carried forward context from follow-up: ${JSON.stringify(workingSet)}`,
);
}

if (fileIds && fileIds.length > 0 || imageAttachmentFileIds && imageAttachmentFileIds.length > 0) {
loggerWithChild({ email: email }).info(
`Follow-up query with file context detected. Using file-based context with NEW classification: ${JSON.stringify(classification)}, FileIds: ${JSON.stringify([fileIds, imageAttachmentFileIds])}`,
)
iterator = UnderstandMessageAndAnswerForGivenContext(
email,
ctx,
userMetadata,
message,
0.5,
parsedMessage.data.fileIds as string[],
fileIds as string[],
userRequestsReasoning,
understandSpan,
undefined,
undefined,
imageAttachmentFileIds as string[],
agentPromptValue,
undefined,
actualModelId || config.defaultBestModel,
Expand Down
Loading
Loading