switch stop sequence in tool calls

Author: charleslien

backend/src/__tests__/run-agent-step-tools.test.ts

@@ -176,7 +176,9 @@ describe('runAgentStep - update_report tool', () => {
           message: 'Hi',
         },
         false
-      ) + getToolCallString('end_turn', {}, true)
+      ) +
+      '\n\n' +
+      getToolCallString('end_turn', {}, true)
 
     spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* () {
       yield mockResponse

backend/src/__tests__/xml-stream-parser.test.ts

@@ -3,7 +3,7 @@ import { CoreMessage } from 'ai'
 
 import { promptAiSdkStream } from './llm-apis/vercel-ai-sdk/ai-sdk'
 import { AgentTemplate } from './templates/types'
-import { globalStopSequences } from './tools/constants'
+import { globalStopSequence } from './tools/constants'
 
 export const getAgentStreamFromTemplate = (params: {
   clientSessionId: string
@@ -26,7 +26,7 @@ export const getAgentStreamFromTemplate = (params: {
     const options: Parameters<typeof promptAiSdkStream>[0] = {
       messages,
       model,
-      stopSequences: globalStopSequences,
+      stopSequences: [globalStopSequence],
       clientSessionId,
       fingerprintId,
       userInputId,

backend/src/prompt-agent-stream.ts

@@ -72,9 +72,14 @@ export const toolParams = Object.fromEntries(
 ) as Record<ToolName, string[]>
 
 function paramsSection(schema: z.ZodObject, endsAgentStep: boolean) {
-  const jsonSchema = z.toJSONSchema(
-    schema.extend({ [endsAgentStepParam]: z.literal(endsAgentStep) })
-  )
+  const schemaWithEndsAgentStepParam = endsAgentStep
+    ? schema.extend({
+        [endsAgentStepParam]: z
+          .literal(endsAgentStep)
+          .describe('Easp flag must be set to true'),
+      })
+    : schema
+  const jsonSchema = z.toJSONSchema(schemaWithEndsAgentStepParam)
   delete jsonSchema.description
   delete jsonSchema['$schema']
   const paramsDescription = Object.keys(jsonSchema.properties ?? {}).length

backend/src/tools.ts

@@ -46,10 +46,7 @@ import { handleUpdateSubgoal } from './handlers/update-subgoal'
 import { handleWebSearch } from './handlers/web-search'
 import { handleWriteFile } from './handlers/write-file'
 
-export const globalStopSequences = [
-  `${JSON.stringify(endsAgentStepParam)}:true`,
-  `${JSON.stringify(endsAgentStepParam)}: true`,
-]
+export const globalStopSequence = `${JSON.stringify(endsAgentStepParam)}`
 
 type Prettify<T> = { [K in keyof T]: T[K] } & {}
 

backend/src/tools/constants.ts

@@ -58,15 +58,17 @@ export function parseRawToolCall<T extends ToolName = ToolName>(
       codebuffToolDefs[validName].endsAgentStep
   }
 
-  const result = (
-    codebuffToolDefs[validName].parameters satisfies z.ZodObject as z.ZodObject
-  )
-    .extend({
-      [endsAgentStepParam]: z.literal(
-        codebuffToolDefs[validName].endsAgentStep
-      ),
-    })
-    .safeParse(processedParameters)
+  const paramsSchema = codebuffToolDefs[validName].endsAgentStep
+    ? (
+        codebuffToolDefs[validName]
+          .parameters satisfies z.ZodObject as z.ZodObject
+      ).extend({
+        [endsAgentStepParam]: z.literal(
+          codebuffToolDefs[validName].endsAgentStep
+        ),
+      })
+    : codebuffToolDefs[validName].parameters
+  const result = paramsSchema.safeParse(processedParameters)
 
   if (!result.success) {
     return {
@@ -81,7 +83,10 @@ export function parseRawToolCall<T extends ToolName = ToolName>(
     }
   }
 
-  delete result.data[endsAgentStepParam]
+  if (endsAgentStepParam in result.data) {
+    delete result.data[endsAgentStepParam]
+  }
+
   return {
     toolName: validName,
     args: result.data,

backend/src/tools/tool-executor.ts

@@ -1,3 +1,6 @@
+import { trackEvent } from '@codebuff/common/analytics'
+import { Model } from '@codebuff/common/constants'
+import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
 import {
   endsAgentStepParam,
   endToolTag,
@@ -37,7 +40,12 @@ export async function* processStreamWithTags(
       onTagEnd: (tagName: string, params: Record<string, any>) => void
     }
   >,
-  onError: (tagName: string, errorMessage: string) => void
+  onError: (tagName: string, errorMessage: string) => void,
+  loggerOptions?: {
+    userId?: string
+    model?: Model
+    agentName?: string
+  }
 ): AsyncGenerator<string> {
   let streamCompleted = false
   let buffer = ''
@@ -62,16 +70,49 @@ export async function* processStreamWithTags(
     try {
       parsedParams = JSON.parse(contents)
     } catch (error: any) {
-      onError('parse_error', error.message)
+      trackEvent(
+        AnalyticsEvent.MALFORMED_TOOL_CALL_JSON,
+        loggerOptions?.userId ?? '',
+        {
+          contents,
+          model: loggerOptions?.model,
+          agent: loggerOptions?.agentName,
+        }
+      )
+      const shortenedContents =
+        contents.length < 50
+          ? contents
+          : contents.slice(0, 20) + '...' + contents.slice(-20)
+      onError(
+        'parse_error',
+        `Invalid JSON: ${JSON.stringify(shortenedContents)}\nError: ${error.message}`
+      )
       return
     }
 
     const toolName = parsedParams[toolNameParam] as keyof typeof processors
     if (!processors[toolName]) {
+      trackEvent(
+        AnalyticsEvent.UNKNOWN_TOOL_CALL,
+        loggerOptions?.userId ?? '',
+        {
+          contents,
+          toolName,
+          model: loggerOptions?.model,
+          agent: loggerOptions?.agentName,
+        }
+      )
       onError(toolName, `Tool not found: ${toolName}`)
       return
     }
 
+    trackEvent(AnalyticsEvent.TOOL_USE, loggerOptions?.userId ?? '', {
+      toolName,
+      contents,
+      parsedParams,
+      model: loggerOptions?.model,
+      agent: loggerOptions?.agentName,
+    })
     delete parsedParams[toolNameParam]
 
     processors[toolName].onTagStart(toolName, {})

backend/src/xml-stream-parser.ts

@@ -25,11 +25,13 @@ export enum AnalyticsEvent {
   UPDATE_CODEBUFF_FAILED = 'cli.update_codebuff_failed',
 
   // Backend
-  USER_INPUT = 'backend.user_input',
   AGENT_STEP = 'backend.agent_step',
   CREDIT_GRANT = 'backend.credit_grant',
   CREDIT_CONSUMED = 'backend.credit_consumed',
+  MALFORMED_TOOL_CALL_JSON = 'backend.malformed_tool_call_json',
   TOOL_USE = 'backend.tool_use',
+  UNKNOWN_TOOL_CALL = 'backend.unknown_tool_call',
+  USER_INPUT = 'backend.user_input',
 
   // Web
   SIGNUP = 'web.signup',

common/src/constants/analytics-events.ts

@@ -2,7 +2,7 @@ import { ToolResultPart } from 'ai'
 import { closeXml } from '../util/xml'
 
 export const toolNameParam = 'codebuff_tool_name'
-export const endsAgentStepParam = 'codebuff_end_step'
+export const endsAgentStepParam = 'codebuff_easp'
 export const toolXmlName = 'codebuff_tool_call'
 export const startToolTag = `<${toolXmlName}>\n`
 export const endToolTag = `\n</${toolXmlName}>`
@@ -81,19 +81,14 @@ export const getToolCallString = (
   params: Record<string, any>,
   endsAgentStep: boolean
 ) => {
-  return [
-    startToolTag,
-    JSON.stringify(
-      {
-        [toolNameParam]: toolName,
-        ...params,
-        [endsAgentStepParam]: endsAgentStep,
-      },
-      null,
-      2
-    ),
-    endToolTag,
-  ].join('')
+  const obj: Record<string, any> = {
+    [toolNameParam]: toolName,
+    ...params,
+  }
+  if (endsAgentStep) {
+    obj[endsAgentStepParam] = true
+  }
+  return [startToolTag, JSON.stringify(obj, null, 2), endToolTag].join('')
 }
 
 export type StringToolResultPart = Omit<ToolResultPart, 'type'> & {