switch all instances of CoreMessage to CodebuffMessage

Author: charleslien

backend/src/__tests__/request-files-prompt.test.ts

@@ -14,8 +14,8 @@ import * as OriginalRequestFilesPromptModule from '../find-files/request-files-p
 import * as geminiWithFallbacksModule from '../llm-apis/gemini-with-fallbacks'
 
 import type { CostMode } from '@codebuff/common/constants'
+import type { CodebuffMessage } from '@codebuff/common/types/message'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
-import type { CoreMessage } from 'ai'
 import type { Mock } from 'bun:test'
 
 // Restore module-level mocks using bunMockFn for the mock implementations
@@ -65,7 +65,9 @@ bunMockFn.module('@codebuff/bigquery', () => ({
 }))
 
 describe('requestRelevantFiles', () => {
-  const mockMessages: CoreMessage[] = [{ role: 'user', content: 'test prompt' }]
+  const mockMessages: CodebuffMessage[] = [
+    { role: 'user', content: 'test prompt' },
+  ]
   const mockSystem = 'test system'
   const mockFileContext: ProjectFileContext = {
     projectRoot: '/test/project',

backend/src/fast-rewrite.ts

@@ -11,7 +11,7 @@ import { promptRelaceAI } from './llm-apis/relace-api'
 import { promptAiSdk } from './llm-apis/vercel-ai-sdk/ai-sdk'
 import { logger } from './util/logger'
 
-import type { CoreMessage } from 'ai'
+import type { CodebuffMessage } from '@codebuff/common/types/message'
 
 export async function fastRewrite(
   initialContent: string,
@@ -135,7 +135,7 @@ export const shouldAddFilePlaceholders = async (
   filePath: string,
   oldContent: string,
   rewrittenNewContent: string,
-  messageHistory: CoreMessage[],
+  messageHistory: CodebuffMessage[],
   fullResponse: string,
   userId: string | undefined,
   clientSessionId: string,

backend/src/find-files/check-new-files-necessary.ts

@@ -2,17 +2,17 @@ import { models } from '@codebuff/common/constants'
 import { closeXml } from '@codebuff/common/util/xml'
 
 import { promptFlashWithFallbacks } from '../llm-apis/gemini-with-fallbacks'
-import { getCoreMessagesSubset } from '../util/messages'
+import { getMessagesSubset } from '../util/messages'
 
 import type { System } from '../llm-apis/claude'
-import type { CoreMessage } from 'ai'
+import type { CodebuffMessage } from '@codebuff/common/types/message'
 
 const systemIntro = `
 You are assisting the user with their software project, in the application Codebuff. Codebuff is a coding agent that helps developers write code or perform utility tasks.
 `.trim()
 
 export const checkNewFilesNecessary = async (
-  messages: CoreMessage[],
+  messages: CodebuffMessage[],
   system: System,
   clientSessionId: string,
   fingerprintId: string,
@@ -56,7 +56,7 @@ Answer with just 'YES' if reading new files is helpful, or 'NO' if the current f
   const response = await promptFlashWithFallbacks(
     [
       { role: 'system', content: systemWithCodebuffInfo },
-      ...getCoreMessagesSubset(
+      ...getMessagesSubset(
         [...messages, { role: 'user', content: prompt }],
         bufferTokens,
       ),

backend/src/find-files/request-files-prompt.ts

@@ -19,8 +19,8 @@ import { promptAiSdk } from '../llm-apis/vercel-ai-sdk/ai-sdk'
 import { logger } from '../util/logger'
 import {
   castAssistantMessage,
-  coreMessagesWithSystem,
-  getCoreMessagesSubset,
+  messagesWithSystem,
+  getMessagesSubset,
 } from '../util/messages'
 import { getRequestContext } from '../websockets/request-context'
 
@@ -30,8 +30,8 @@ import type {
   GetExpandedFileContextForTrainingTrace,
   GetRelevantFilesTrace,
 } from '@codebuff/bigquery'
+import type { CodebuffMessage } from '@codebuff/common/types/message'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
-import type { CoreMessage } from 'ai'
 
 const NUMBER_OF_EXAMPLE_FILES = 100
 const MAX_FILES_PER_REQUEST = 30
@@ -124,7 +124,7 @@ export async function requestRelevantFiles(
     messages,
     system,
   }: {
-    messages: CoreMessage[]
+    messages: CodebuffMessage[]
     system: string | Array<TextBlock>
   },
   fileContext: ProjectFileContext,
@@ -260,7 +260,7 @@ export async function requestRelevantFilesForTraining(
     messages,
     system,
   }: {
-    messages: CoreMessage[]
+    messages: CodebuffMessage[]
     system: string | Array<TextBlock>
   },
   fileContext: ProjectFileContext,
@@ -341,7 +341,7 @@ async function getRelevantFiles(
     messages,
     system,
   }: {
-    messages: CoreMessage[]
+    messages: CodebuffMessage[]
     system: string | Array<TextBlock>
   },
   userPrompt: string,
@@ -355,7 +355,7 @@ async function getRelevantFiles(
   modelId?: FinetunedVertexModel,
 ) {
   const bufferTokens = 100_000
-  const messagesWithPrompt = getCoreMessagesSubset(
+  const messagesWithPrompt = getMessagesSubset(
     [
       ...messages,
       {
@@ -366,12 +366,12 @@ async function getRelevantFiles(
     bufferTokens,
   )
   const start = performance.now()
-  let coreMessages = coreMessagesWithSystem(messagesWithPrompt, system)
+  let codebuffMessages = messagesWithSystem(messagesWithPrompt, system)
 
   // Converts assistant messages to user messages for finetuned model
-  coreMessages = coreMessages
+  codebuffMessages = codebuffMessages
     .map((msg, i) => {
-      if (msg.role === 'assistant' && i !== coreMessages.length - 1) {
+      if (msg.role === 'assistant' && i !== codebuffMessages.length - 1) {
         return castAssistantMessage(msg)
       } else {
         return msg
@@ -380,7 +380,7 @@ async function getRelevantFiles(
     .filter((msg) => msg !== null)
   const finetunedModel = modelId ?? finetunedVertexModels.ft_filepicker_010
 
-  let response = await promptFlashWithFallbacks(coreMessages, {
+  let response = await promptFlashWithFallbacks(codebuffMessages, {
     clientSessionId,
     userInputId,
     model: models.gemini2flash,
@@ -424,7 +424,7 @@ async function getRelevantFilesForTraining(
     messages,
     system,
   }: {
-    messages: CoreMessage[]
+    messages: CodebuffMessage[]
     system: string | Array<TextBlock>
   },
   userPrompt: string,
@@ -437,7 +437,7 @@ async function getRelevantFilesForTraining(
   repoId: string | undefined,
 ) {
   const bufferTokens = 100_000
-  const messagesWithPrompt = getCoreMessagesSubset(
+  const messagesWithPrompt = getMessagesSubset(
     [
       ...messages,
       {
@@ -449,7 +449,7 @@ async function getRelevantFilesForTraining(
   )
   const start = performance.now()
   let response = await promptAiSdk({
-    messages: coreMessagesWithSystem(messagesWithPrompt, system),
+    messages: messagesWithSystem(messagesWithPrompt, system),
     clientSessionId,
     fingerprintId,
     userInputId,

backend/src/get-file-reading-updates.ts

@@ -13,8 +13,8 @@ import {
 import { countTokensJson } from './util/token-counter'
 import { requestFiles } from './websockets/websocket-action'
 
+import type { CodebuffMessage } from '@codebuff/common/types/message'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
-import type { CoreMessage } from 'ai'
 import type { WebSocket } from 'ws'
 
 const getInitialFiles = (fileContext: ProjectFileContext) => {
@@ -39,7 +39,7 @@ const getInitialFiles = (fileContext: ProjectFileContext) => {
 
 export async function getFileReadingUpdates(
   ws: WebSocket,
-  messages: CoreMessage[],
+  messages: CodebuffMessage[],
   fileContext: ProjectFileContext,
   options: {
     requestedFiles?: string[]

backend/src/llm-apis/check-for-loop.ts

@@ -8,7 +8,7 @@ import type {
   FinetunedVertexModel,
   GeminiModel,
 } from '@codebuff/common/constants'
-import type { CoreMessage } from 'ai'
+import type { CodebuffMessage } from '@codebuff/common/types/message'
 
 /**
  * Prompts a Gemini model with fallback logic.
@@ -36,7 +36,7 @@ import type { CoreMessage } from 'ai'
  * @throws If all API calls (primary and fallbacks) fail.
  */
 export async function promptFlashWithFallbacks(
-  messages: CoreMessage[],
+  messages: CodebuffMessage[],
   options: {
     clientSessionId: string
     fingerprintId: string

backend/src/llm-apis/gemini-with-fallbacks.ts

@@ -21,10 +21,8 @@ import { stripNullCharsFromObject } from '../util/object'
 import { SWITCHBOARD } from '../websockets/server'
 import { sendAction } from '../websockets/websocket-action'
 
-import type { OpenAIMessage } from './openai-api'
 import type { ClientState } from '../websockets/switchboard'
-import type { Message } from '@codebuff/common/types/message'
-import type { CoreMessage } from 'ai'
+import type { CodebuffMessage } from '@codebuff/common/types/message'
 
 export const PROFIT_MARGIN = 0.3
 
@@ -332,7 +330,7 @@ type InsertMessageParams = {
   fingerprintId: string
   userInputId: string
   model: string
-  request: Message[] | OpenAIMessage[] | CoreMessage[]
+  request: CodebuffMessage[]
   response: string
   inputTokens: number
   outputTokens: number
@@ -532,7 +530,7 @@ export const saveMessage = async (value: {
   fingerprintId: string
   userInputId: string
   model: string
-  request: Message[] | OpenAIMessage[] | CoreMessage[]
+  request: CodebuffMessage[]
   response: string
   inputTokens: number
   outputTokens: number