trace calls to completion providers

vscode/src/completions/index.ts

@@ -14,7 +14,7 @@ import { getCurrentDocContext } from './document'
 import { History } from './history'
 import * as CompletionLogger from './logger'
 import { detectMultiline } from './multiline'
-import { Provider, ProviderConfig, ProviderOptions } from './providers/provider'
+import { CompletionProviderTracer, Provider, ProviderConfig, ProviderOptions } from './providers/provider'
 import { RequestManager } from './request-manager'
 import { sharedPostProcess } from './shared-post-process'
 import { ProvideInlineCompletionItemsTracer, ProvideInlineCompletionsItemTraceData } from './tracer'
@@ -359,7 +359,8 @@ export class CodyCompletionItemProvider implements vscode.InlineCompletionItemPr
             prefix,
             completers,
             contextResult.context,
-            abortController.signal
+            abortController.signal,
+            tracer ? createCompletionProviderTracer(tracer) : undefined
         )
 
         // Shared post-processing logic
@@ -487,6 +488,13 @@ function createTracerForInvocation(tracer: ProvideInlineCompletionItemsTracer):
     }
 }
 
+function createCompletionProviderTracer(tracer: SingleInvocationTracer): CompletionProviderTracer {
+    return {
+        params: data => tracer({ completionProviderCallParams: data }),
+        result: data => tracer({ completionProviderCallResult: data }),
+    }
+}
+
 function delay(milliseconds: number): Promise<void> {
     return new Promise<void>(resolve => setTimeout(resolve, milliseconds))
 }

vscode/src/completions/providers/anthropic.ts

@@ -17,7 +17,7 @@ import {
 } from '../text-processing'
 import { batchCompletions, messagesToText } from '../utils'
 
-import { Provider, ProviderConfig, ProviderOptions } from './provider'
+import { CompletionProviderTracer, Provider, ProviderConfig, ProviderOptions } from './provider'
 
 const CHARS_PER_TOKEN = 4
 
@@ -137,7 +137,11 @@ export class AnthropicProvider extends Provider {
         return completion.trimEnd()
     }
 
-    public async generateCompletions(abortSignal: AbortSignal, snippets: ReferenceSnippet[]): Promise<Completion[]> {
+    public async generateCompletions(
+        abortSignal: AbortSignal,
+        snippets: ReferenceSnippet[],
+        tracer?: CompletionProviderTracer
+    ): Promise<Completion[]> {
         // Create prompt
         const { messages: prompt } = this.createPrompt(snippets)
         if (prompt.length > this.promptChars) {
@@ -157,6 +161,7 @@ export class AnthropicProvider extends Provider {
                   maxTokensToSample: Math.min(50, this.responseTokens),
                   stopSequences: [anthropic.HUMAN_PROMPT, CLOSING_CODE_TAG, '\n\n'],
               }
+        tracer?.params(args)
 
         // Issue request
         const responses = await batchCompletions(this.completionsClient, args, this.options.n, abortSignal)
@@ -172,14 +177,16 @@ export class AnthropicProvider extends Provider {
             return [
                 {
                     prefix: this.options.prefix,
-                    messages: prompt,
                     content,
                     stopReason: resp.stopReason,
                 },
             ]
         })
 
-        return ret.flat()
+        const completions = ret.flat()
+        tracer?.result({ rawResponses: responses, completions })
+
+        return completions
     }
 }
 

vscode/src/completions/providers/provider.ts

@@ -1,3 +1,5 @@
+import { CompletionParameters } from '@sourcegraph/cody-shared/src/sourcegraph-api/completions/types'
+
 import { Completion } from '..'
 import { ReferenceSnippet } from '../context'
 
@@ -50,5 +52,28 @@ export interface ProviderOptions {
 export abstract class Provider {
     constructor(public readonly options: Readonly<ProviderOptions>) {}
 
-    public abstract generateCompletions(abortSignal: AbortSignal, snippets: ReferenceSnippet[]): Promise<Completion[]>
+    public abstract generateCompletions(
+        abortSignal: AbortSignal,
+        snippets: ReferenceSnippet[],
+        tracer?: CompletionProviderTracer
+    ): Promise<Completion[]>
+}
+
+/**
+ * Tracer for {@link Provider}.
+ */
+export interface CompletionProviderTracer {
+    /** Called with the params passed to the LLM. */
+    params(params: CompletionParameters): void
+
+    /** Called with the result from the LLM. */
+    result(data: CompletionProviderTracerResultData): void
+}
+
+export interface CompletionProviderTracerResultData {
+    /** The raw response from the LLM. */
+    rawResponses: unknown
+
+    /** The post-processed completions that are returned by the provider. */
+    completions: Completion[]
 }

vscode/src/completions/request-manager.ts

@@ -2,10 +2,11 @@ import { Completion } from '.'
 import { CompletionsCache } from './cache'
 import { ReferenceSnippet } from './context'
 import { logCompletionEvent } from './logger'
-import { Provider } from './providers/provider'
+import { CompletionProviderTracer, Provider } from './providers/provider'
 
 interface Request {
     prefix: string
+    tracer?: CompletionProviderTracer
     resolve(completions: Completion[]): void
     reject(error: Error): void
 }
@@ -27,7 +28,8 @@ export class RequestManager {
         prefix: string,
         providers: Provider[],
         context: ReferenceSnippet[],
-        signal: AbortSignal
+        signal: AbortSignal,
+        tracer?: CompletionProviderTracer
     ): Promise<Completion[]> {
         let resolve: Request['resolve'] = () => {}
         let reject: Request['reject'] = () => {}
@@ -40,6 +42,7 @@ export class RequestManager {
             prefix,
             resolve,
             reject,
+            tracer,
         }
         this.startRequest(request, documentUri, logId, providers, context, signal)
 
@@ -61,7 +64,9 @@ export class RequestManager {
 
         this.addRequest(documentUri, request)
 
-        Promise.all(providers.map(c => c.generateCompletions(networkRequestAbortController.signal, context)))
+        Promise.all(
+            providers.map(c => c.generateCompletions(networkRequestAbortController.signal, context, request.tracer))
+        )
             .then(res => res.flat())
             .then(completions => {
                 // Add the completed results to the cache, even if the request

vscode/src/completions/tracer/index.ts

@@ -1,7 +1,9 @@
 import * as vscode from 'vscode'
 
+import { CompletionParameters } from '@sourcegraph/cody-shared/src/sourcegraph-api/completions/types'
+
 import { GetContextResult } from '../context'
-import { Provider } from '../providers/provider'
+import { CompletionProviderTracerResultData, Provider } from '../providers/provider'
 
 /**
  * Traces invocations of {@link CodyCompletionItemProvider.provideInlineCompletionItems}.
@@ -27,6 +29,13 @@ export interface ProvideInlineCompletionsItemTraceData {
         context: vscode.InlineCompletionContext
     }
     completers?: Provider['options'][]
+
+    /**
+     * @todo Make this support recording more than 1 call to a completion provider.
+     */
+    completionProviderCallParams?: CompletionParameters
+    completionProviderCallResult?: CompletionProviderTracerResultData
+
     context?: GetContextResult
     result?: vscode.InlineCompletionList
     cacheHit?: boolean

vscode/src/completions/tracer/traceView.ts

@@ -108,6 +108,19 @@ ${
               )
               .join('\n\n')
 }
+`,
+        data?.completionProviderCallParams &&
+            `
+## Completion provider calls
+
+${codeDetailsWithSummary('Params', JSON.stringify(data.completionProviderCallParams, null, 2))}
+
+${
+    data.completionProviderCallResult
+        ? codeDetailsWithSummary('Result', JSON.stringify(data.completionProviderCallResult, null, 2))
+        : '_Loading result..._'
+}
+
 `,
         data?.result &&
             `
@@ -130,7 +143,7 @@ ${markdownCodeBlock(data.error)}
         `
 ## Advanced tools
 
-${codeDetailsWithSummary('JSON for dataset', jsonForDataset(data), 'start')}
+${codeDetailsWithSummary('JSON for dataset', jsonForDataset(data))}
 
 `,
     ]
@@ -141,13 +154,23 @@ ${codeDetailsWithSummary('JSON for dataset', jsonForDataset(data), 'start')}
     return renderMarkdown(markdownSource, { noDomPurify: true })
 }
 
-function codeDetailsWithSummary(title: string, value: string, anchor: 'start' | 'end', excerptLength = 50): string {
-    const excerpt = anchor === 'start' ? value.slice(0, excerptLength) : value.slice(-excerptLength)
+function codeDetailsWithSummary(
+    title: string,
+    value: string,
+    anchor: 'start' | 'end' | 'none' = 'none',
+    excerptLength = 50
+): string {
+    const excerpt =
+        anchor === 'start' ? value.slice(0, excerptLength) : anchor === 'end' ? value.slice(-excerptLength) : null
+    const excerptMarkdown =
+        excerpt !== null
+            ? `: <code>${anchor === 'end' ? '⋯' : ''}${withVisibleWhitespace(excerpt)
+                  .replace(/</g, '&lt;')
+                  .replace(/>/g, '&gt;')}${anchor === 'start' ? '⋯' : ''}</code>`
+            : ''
     return `
 <details>
-<summary>${title}: <code>${anchor === 'end' ? '⋯' : ''}${withVisibleWhitespace(excerpt)
-        .replace(/</g, '&lt;')
-        .replace(/>/g, '&gt;')}${anchor === 'start' ? '⋯' : ''}</code></summary>
+<summary>${title}${excerptMarkdown}</summary>
 
 ${markdownCodeBlock(value)}