Skip to content

Commit 27ef6d3

Browse files
brendan-kellambrendan-kellam
committed
Add chat history to LLM input
1 parent c349add commit 27ef6d3

File tree

4 files changed

+152
-28
lines changed

4 files changed

+152
-28
lines changed

‎packages/web/src/app/api/(server)/chat/route.ts‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { env } from "@/env.mjs";
33
import { saveChatMessages, updateChatName } from "@/features/chat/actions";
44
import { createAgentStream } from "@/features/chat/agent";
55
import { SBChatMessage } from "@/features/chat/types";
6-
import { getConfiguredModelProviderInfo } from "@/features/chat/utils";
6+
import { getAnswerPartFromAssistantMessage, getConfiguredModelProviderInfo } from "@/features/chat/utils";
77
import { ErrorCode } from "@/lib/errorCodes";
88
import { schemaValidationError, serviceErrorResponse } from "@/lib/serviceError";
99
import { isServiceError } from "@/lib/utils";
@@ -14,7 +14,7 @@ import { createOpenAI, OpenAIResponsesProviderOptions } from "@ai-sdk/openai";
1414
import { LanguageModelV2 } from "@ai-sdk/provider";
1515
import { OrgRole } from "@sourcebot/db";
1616
import { createLogger } from "@sourcebot/logger";
17-
import { convertToModelMessages, createUIMessageStream, createUIMessageStreamResponse, generateText, JSONValue, StreamTextResult, UIMessageStreamOptions, UIMessageStreamWriter } from "ai";
17+
import { createUIMessageStream, createUIMessageStreamResponse, generateText, JSONValue, ModelMessage, StreamTextResult, UIMessageStreamOptions, UIMessageStreamWriter } from "ai";
1818
import { StatusCodes } from "http-status-codes";
1919
import { z } from "zod";
2020

@@ -86,20 +86,40 @@ const chatHandler = ({ messages, id, selectedRepos }: { messages: SBChatMessage[
8686
) {
8787
const content = messages[0].parts[0].text;
8888

89-
logger.debug("Generating chat title...");
9089
const title = await generateChatTitle(content, model);
9190
if (title) {
92-
logger.debug("Chat title generated:", title);
9391
updateChatName({
9492
chatId: id,
9593
name: title,
9694
}, domain);
9795
}
9896
else {
99-
logger.debug("Failed to generate chat title.");
97+
logger.error("Failed to generate chat title.");
10098
}
10199
}
102100

101+
// Extract user messages and assistant answers.
102+
// We will use this as the context we carry between messages.
103+
const messageHistory =
104+
messages.map((message): ModelMessage | undefined => {
105+
if (message.role === 'user') {
106+
return {
107+
role: 'user',
108+
content: message.parts[0].type === 'text' ? message.parts[0].text : '',
109+
};
110+
}
111+
112+
if (message.role === 'assistant') {
113+
const answerPart = getAnswerPartFromAssistantMessage(message, false);
114+
if (answerPart) {
115+
return {
116+
role: 'assistant',
117+
content: [answerPart]
118+
}
119+
}
120+
}
121+
}).filter(message => message !== undefined);
122+
103123
try {
104124
const stream = createUIMessageStream<SBChatMessage>({
105125
execute: async ({ writer }) => {
@@ -113,8 +133,7 @@ const chatHandler = ({ messages, id, selectedRepos }: { messages: SBChatMessage[
113133
model,
114134
providerOptions,
115135
headers,
116-
// @todo: we will need to incorporate the previous messages into the prompt.
117-
inputMessages: convertToModelMessages([latestMessage]),
136+
inputMessages: messageHistory,
118137
inputSources: sources,
119138
selectedRepos,
120139
onWriteSource: (source) => {

‎packages/web/src/features/chat/components/chatThread/chatThreadListItem.tsx‎

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,20 +7,20 @@ import { ResizablePanel, ResizablePanelGroup } from '@/components/ui/resizable';
77
import { Separator } from '@/components/ui/separator';
88
import { Skeleton } from '@/components/ui/skeleton';
99
import { cn } from '@/lib/utils';
10-
import { Brain, CheckCircle, ChevronDown, ChevronRight, Clock, Copy, Cpu, InfoIcon, Loader2, ThumbsDown, ThumbsUp, Zap } from 'lucide-react';
10+
import { Brain, CheckCircle, ChevronDown, ChevronRight, Clock, Cpu, InfoIcon, Loader2, Zap } from 'lucide-react';
1111
import { CSSProperties, forwardRef, useCallback, useEffect, useMemo, useRef, useState } from 'react';
1212
import scrollIntoView from 'scroll-into-view-if-needed';
1313
import { ANSWER_TAG } from '../../constants';
1414
import { Reference, referenceSchema, SBChatMessage, SBChatMessageMetadata, Source } from "../../types";
1515
import { useExtractReferences } from '../../useExtractReferences';
16-
import { groupMessageIntoSteps } from '../../utils';
16+
import { getAnswerPartFromAssistantMessage, groupMessageIntoSteps } from '../../utils';
17+
import { AnswerCard } from './answerCard';
1718
import { MarkdownRenderer, REFERENCE_PAYLOAD_ATTRIBUTE } from './markdownRenderer';
1819
import { ReferencedSourcesListView } from './referencedSourcesListView';
1920
import { FindSymbolDefinitionsToolComponent } from './tools/findSymbolDefinitionsToolComponent';
2021
import { FindSymbolReferencesToolComponent } from './tools/findSymbolReferencesToolComponent';
2122
import { ReadFilesToolComponent } from './tools/readFilesToolComponent';
2223
import { SearchCodeToolComponent } from './tools/searchCodeToolComponent';
23-
import { AnswerCard } from './answerCard';
2424

2525
interface ChatThreadListItemProps {
2626
userMessage: SBChatMessage;
@@ -58,21 +58,12 @@ export const ChatThreadListItem = forwardRef<HTMLDivElement, ChatThreadListItemP
5858
}, [assistantMessage?.metadata]);
5959

6060
const answerPart = useMemo(() => {
61-
const lastTextPart = assistantMessage?.parts
62-
.findLast((part) => part.type === 'text')
63-
64-
if (lastTextPart?.text.startsWith(ANSWER_TAG)) {
65-
return lastTextPart;
66-
}
67-
68-
// If the agent did not include the answer tag, then fallback to using the last text part.
69-
// Only do this when we are no longer streaming since the agent may still be thinking.
70-
if (!isStreaming && lastTextPart) {
71-
return lastTextPart;
61+
if (!assistantMessage) {
62+
return undefined;
7263
}
7364

74-
return undefined;
75-
}, [assistantMessage?.parts, isStreaming]);
65+
return getAnswerPartFromAssistantMessage(assistantMessage, isStreaming);
66+
}, [assistantMessage, isStreaming]);
7667

7768

7869
const thinkingSteps = useMemo(() => {

‎packages/web/src/features/chat/utils.test.ts‎

Lines changed: 98 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { expect, test, vi } from 'vitest'
2-
import { fileReferenceToString, groupMessageIntoSteps, sourceCodeChunksToModelOutput, sourceCodeToModelOutput } from './utils'
3-
import { FILE_REFERENCE_REGEX } from './constants';
4-
import { SBChatMessagePart } from './types';
2+
import { fileReferenceToString, getAnswerPartFromAssistantMessage, groupMessageIntoSteps, sourceCodeChunksToModelOutput, sourceCodeToModelOutput } from './utils'
3+
import { FILE_REFERENCE_REGEX, ANSWER_TAG } from './constants';
4+
import { SBChatMessage, SBChatMessagePart } from './types';
55

66
// Mock the env module
77
vi.mock('@/env.mjs', () => ({
@@ -288,3 +288,98 @@ test('groupMessageIntoSteps returns a single group when there is no step-start p
288288
]
289289
]);
290290
});
291+
292+
test('getAnswerPartFromAssistantMessage returns text part when it starts with ANSWER_TAG while not streaming', () => {
293+
const message: SBChatMessage = {
294+
role: 'assistant',
295+
parts: [
296+
{
297+
type: 'text',
298+
text: 'Some initial text'
299+
},
300+
{
301+
type: 'text',
302+
text: `${ANSWER_TAG}This is the answer to your question.`
303+
}
304+
]
305+
} as SBChatMessage;
306+
307+
const result = getAnswerPartFromAssistantMessage(message, false);
308+
309+
expect(result).toEqual({
310+
type: 'text',
311+
text: `${ANSWER_TAG}This is the answer to your question.`
312+
});
313+
});
314+
315+
test('getAnswerPartFromAssistantMessage returns text part when it starts with ANSWER_TAG while streaming', () => {
316+
const message: SBChatMessage = {
317+
role: 'assistant',
318+
parts: [
319+
{
320+
type: 'text',
321+
text: 'Some initial text'
322+
},
323+
{
324+
type: 'text',
325+
text: `${ANSWER_TAG}This is the answer to your question.`
326+
}
327+
]
328+
} as SBChatMessage;
329+
330+
const result = getAnswerPartFromAssistantMessage(message, true);
331+
332+
expect(result).toEqual({
333+
type: 'text',
334+
text: `${ANSWER_TAG}This is the answer to your question.`
335+
});
336+
});
337+
338+
test('getAnswerPartFromAssistantMessage returns last text part as fallback when not streaming and no ANSWER_TAG', () => {
339+
const message: SBChatMessage = {
340+
role: 'assistant',
341+
parts: [
342+
{
343+
type: 'text',
344+
text: 'First text part'
345+
},
346+
{
347+
type: 'tool-call',
348+
id: 'call-1',
349+
name: 'search',
350+
args: {}
351+
},
352+
{
353+
type: 'text',
354+
text: 'This is the last text part without answer tag'
355+
}
356+
]
357+
} as SBChatMessage;
358+
359+
const result = getAnswerPartFromAssistantMessage(message, false);
360+
361+
expect(result).toEqual({
362+
type: 'text',
363+
text: 'This is the last text part without answer tag'
364+
});
365+
});
366+
367+
test('getAnswerPartFromAssistantMessage returns undefined when streaming and no ANSWER_TAG', () => {
368+
const message: SBChatMessage = {
369+
role: 'assistant',
370+
parts: [
371+
{
372+
type: 'text',
373+
text: 'Some text without answer tag'
374+
},
375+
{
376+
type: 'text',
377+
text: 'Another text part'
378+
}
379+
]
380+
} as SBChatMessage;
381+
382+
const result = getAnswerPartFromAssistantMessage(message, true);
383+
384+
expect(result).toBeUndefined();
385+
});

‎packages/web/src/features/chat/utils.ts‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { env } from "@/env.mjs"
2-
import { CreateUIMessage, UIMessagePart } from "ai"
2+
import { CreateUIMessage, TextUIPart, UIMessagePart } from "ai"
33
import { Descendant, Editor, Point, Range, Transforms } from "slate"
4-
import { FILE_REFERENCE_PREFIX, FILE_REFERENCE_REGEX } from "./constants"
4+
import { ANSWER_TAG, FILE_REFERENCE_PREFIX, FILE_REFERENCE_REGEX } from "./constants"
55
import { CustomEditor, CustomText, FileReference, FileSource, MentionData, MentionElement, ModelProviderInfo, ParagraphElement, SBChatMessage, SBChatMessagePart, SBChatMessageToolTypes, Source } from "./types"
66

77
export const insertMention = (editor: CustomEditor, data: MentionData, target?: Range | null) => {
@@ -355,3 +355,22 @@ export const groupMessageIntoSteps = (parts: SBChatMessagePart[]) => {
355355

356356
return steps;
357357
}
358+
359+
// Attempts to find the part of the assistant's message
360+
// that contains the answer.
361+
export const getAnswerPartFromAssistantMessage = (message: SBChatMessage, isStreaming: boolean): TextUIPart | undefined => {
362+
const lastTextPart = message.parts
363+
.findLast((part) => part.type === 'text')
364+
365+
if (lastTextPart?.text.startsWith(ANSWER_TAG)) {
366+
return lastTextPart;
367+
}
368+
369+
// If the agent did not include the answer tag, then fallback to using the last text part.
370+
// Only do this when we are no longer streaming since the agent may still be thinking.
371+
if (!isStreaming && lastTextPart) {
372+
return lastTextPart;
373+
}
374+
375+
return undefined;
376+
}

0 commit comments

Comments
 (0)