Integrate VSCode Cody with server-side Context API (#4840)

This PR adds the Client-side integration of
https://github.com/sourcegraph/sourcegraph/pull/63789 for VSCode:
- GraphQL queries, types and client methods
- ContextAPIClient class mapping VSCode objects to GraphQL
- wiring of ContextAPIClient to existing code

Currently, for all added backend calls:

- all remote RPC are guarded by a feature flag
- results of RPCs are not used by integrating call
- calls are not awaited (to avoid slowing down interactions)

Related to AI-128.

## Test plan

- tested locally -> works with local and sourcegraph.com backend

lib/shared/src/experimentation/FeatureFlagProvider.ts

@@ -75,6 +75,9 @@ export enum FeatureFlag {
 
     /** Use embeddings to provide enhanced context. */
     CodyEnhancedContexUseEmbeddings = 'cody-enhanced-context-use-embeddings',
+
+    /** Whether to use server-side Context API. */
+    CodyServerSideContextAPI = 'cody-server-side-context-api-enabled',
 }
 
 const ONE_HOUR = 60 * 60 * 1000

lib/shared/src/sourcegraph-api/graphql/client.ts

@@ -13,6 +13,7 @@ import { addTraceparent, wrapInActiveSpan } from '../../tracing'
 import { isError } from '../../utils'
 import { DOTCOM_URL, isDotCom } from '../environments'
 import {
+    CHAT_INTENT_QUERY,
     CONTEXT_FILTERS_QUERY,
     CONTEXT_SEARCH_QUERY,
     CURRENT_SITE_CODY_CONFIG_FEATURES,
@@ -37,6 +38,8 @@ import {
     LOG_EVENT_MUTATION,
     LOG_EVENT_MUTATION_DEPRECATED,
     PACKAGE_LIST_QUERY,
+    RANK_CONTEXT_QUERY,
+    RECORD_CONTEXT_QUERY,
     RECORD_TELEMETRY_EVENTS_MUTATION,
     REPOSITORY_IDS_QUERY,
     REPOSITORY_ID_QUERY,
@@ -316,6 +319,23 @@ interface SearchAttributionResponse {
 
 type LogEventResponse = unknown
 
+interface ChatIntentResponse {
+    chatIntent: {
+        intent: string
+        score: number
+    }
+}
+
+type RecordContextResponse = unknown
+
+interface RankContextResponse {
+    rankContext: {
+        ranker: string
+        used: number[]
+        ignored: number[]
+    }
+}
+
 interface ContextSearchResponse {
     getCodyContext: {
         blob: {
@@ -344,6 +364,20 @@ export interface EmbeddingsSearchResult {
     content: string
 }
 
+export interface ChatIntentResult {
+    intent: string
+    score: number
+}
+
+/**
+ * Experimental API.
+ */
+export interface InputContextItem {
+    content: string
+    retriever: string
+    score?: number
+}
+
 export interface ContextSearchResult {
     repoName: string
     commit: string
@@ -847,6 +881,52 @@ export class SourcegraphGraphQLAPIClient {
         return isError(result) ? null : result
     }
 
+    /** Experimental API */
+    public async chatIntent(interactionID: string, query: string): Promise<ChatIntentResult | Error> {
+        const response = await this.fetchSourcegraphAPI<APIResponse<ChatIntentResponse>>(
+            CHAT_INTENT_QUERY,
+            {
+                query: query,
+                interactionId: interactionID,
+            }
+        )
+        return extractDataOrError(response, data => data.chatIntent)
+    }
+
+    /** Experimental API */
+    public async recordContext(
+        interactionID: string,
+        used: InputContextItem[],
+        ignored: InputContextItem[]
+    ): Promise<RecordContextResponse | Error> {
+        const response = await this.fetchSourcegraphAPI<APIResponse<RecordContextResponse>>(
+            RECORD_CONTEXT_QUERY,
+            {
+                interactionId: interactionID,
+                usedContextItems: used,
+                ignoredContextItems: ignored,
+            }
+        )
+        return extractDataOrError(response, data => data)
+    }
+
+    /** Experimental API */
+    public async rankContext(
+        interactionID: string,
+        query: string,
+        context: InputContextItem[]
+    ): Promise<RankContextResponse | Error> {
+        const response = await this.fetchSourcegraphAPI<APIResponse<RankContextResponse>>(
+            RANK_CONTEXT_QUERY,
+            {
+                interactionId: interactionID,
+                query,
+                contextItems: context,
+            }
+        )
+        return extractDataOrError(response, data => data)
+    }
+
     public async contextSearch(
         repoIDs: string[],
         query: string

lib/shared/src/sourcegraph-api/graphql/queries.ts

@@ -194,6 +194,30 @@ query Repositories($names: [String!]!, $first: Int!) {
   }
 `
 
+export const CHAT_INTENT_QUERY = `
+query ChatIntent($query: String!, $interactionId: String!) {
+    chatIntent(query: $query, interactionId: $interactionId) {
+        intent
+        score
+    }
+}`
+
+export const RECORD_CONTEXT_QUERY = `
+query RecordContext($interactionId: String!, $usedContextItems: [InputContextItem!]!, $ignoredContextItems: [InputContextItem!]!) {
+    recordContext(interactionId: $interactionId, usedContextItems: $usedContextItems, ignoredContextItems: $ignoredContextItems) {
+        alwaysNil
+    }
+}`
+
+export const RANK_CONTEXT_QUERY = `
+query RankContext($interactionId: String!, $query: String!, $contextItems: [InputContextItem!]!) {
+    rankContext(interactionId: $interactionId, query:$query, contextItems: $contextItems) {
+        ranker
+        used
+        ignored
+    }
+}`
+
 export const CONTEXT_SEARCH_QUERY = `
 query GetCodyContext($repos: [ID!]!, $query: String!, $codeResultsCount: Int!, $textResultsCount: Int!) {
 	getCodyContext(repos: $repos, query: $query, codeResultsCount: $codeResultsCount, textResultsCount: $textResultsCount) {

vscode/src/chat/chat-view/ChatController.ts

@@ -92,6 +92,7 @@ import { TestSupport } from '../../test-support'
 import type { MessageErrorType } from '../MessageProvider'
 import { startClientStateBroadcaster } from '../clientStateBroadcaster'
 import { getChatContextItemsForMention } from '../context/chatContext'
+import type { ContextAPIClient } from '../context/contextAPIClient'
 import type {
     ChatSubmitType,
     ConfigurationSubsetForWebview,
@@ -121,6 +122,7 @@ interface ChatControllerOptions {
     models: Model[]
     guardrails: Guardrails
     startTokenReceiver?: typeof startTokenReceiver
+    contextAPIClient: ContextAPIClient | null
 }
 
 export interface ChatSession {
@@ -169,11 +171,13 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
     private readonly remoteSearch: RemoteSearch | null
     private readonly repoPicker: RemoteRepoPicker | null
     private readonly startTokenReceiver: typeof startTokenReceiver | undefined
+    private readonly contextAPIClient: ContextAPIClient | null
 
     private contextFilesQueryCancellation?: vscode.CancellationTokenSource
     private allMentionProvidersMetadataQueryCancellation?: vscode.CancellationTokenSource
 
     private disposables: vscode.Disposable[] = []
+
     public dispose(): void {
         vscode.Disposable.from(...this.disposables).dispose()
         this.disposables = []
@@ -191,6 +195,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
         guardrails,
         enterpriseContext,
         startTokenReceiver,
+        contextAPIClient,
     }: ChatControllerOptions) {
         this.extensionUri = extensionUri
         this.authProvider = authProvider
@@ -206,6 +211,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
 
         this.guardrails = guardrails
         this.startTokenReceiver = startTokenReceiver
+        this.contextAPIClient = contextAPIClient
 
         if (TestSupport.instance) {
             TestSupport.instance.chatPanelProvider.set(this)
@@ -655,6 +661,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                 abortSignal.throwIfAborted()
 
                 this.postEmptyMessageInProgress()
+                this.contextAPIClient?.detectChatIntent(requestID, inputText.toString())
 
                 // Add user's current selection as context for chat messages.
                 const selectionContext = source === 'chat' ? await getContextFileFromSelection() : []
@@ -697,8 +704,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                                         chatModel: this.chatModel,
                                     })
                                   : inputText
-
-                              return getEnhancedContext({
+                              const context = getEnhancedContext({
                                   strategy: contextStrategy,
                                   editor: this.editor,
                                   input: { text: rewrite, mentions },
@@ -710,6 +716,11 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                                   },
                                   contextRanking: this.contextRanking,
                               })
+                              // add a callback, but return the original context
+                              context.then(c =>
+                                  this.contextAPIClient?.rankContext(requestID, inputText.toString(), c)
+                              )
+                              return context
                           }
                         : undefined,
                     command !== undefined
@@ -743,7 +754,12 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                 }
 
                 try {
-                    const prompt = await this.buildPrompt(prompter, abortSignal, sendTelemetry)
+                    const prompt = await this.buildPrompt(
+                        prompter,
+                        abortSignal,
+                        requestID,
+                        sendTelemetry
+                    )
                     abortSignal.throwIfAborted()
                     this.streamAssistantResponse(requestID, prompt, span, firstTokenSpan, abortSignal)
                 } catch (error) {
@@ -1117,6 +1133,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
     private async buildPrompt(
         prompter: DefaultPrompter,
         abortSignal: AbortSignal,
+        requestID: string,
         sendTelemetry?: (contextSummary: any, privateContextSummary?: any) => void
     ): Promise<Message[]> {
         const { prompt, context } = await prompter.makePrompt(
@@ -1128,6 +1145,8 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
         // Update UI based on prompt construction
         // Includes the excluded context items to display in the UI
         this.chatModel.setLastMessageContext([...context.used, ...context.ignored])
+        // this is not awaited, so we kick the call off but don't block on it returning
+        this.contextAPIClient?.recordContext(requestID, context.used, context.ignored)
 
         if (sendTelemetry) {
             // Create a summary of how many code snippets of each context source are being

vscode/src/chat/chat-view/ChatsController.ts

@@ -28,6 +28,7 @@ import { getConfiguration } from '../../configuration'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { ContextRankingController } from '../../local-context/context-ranking'
 import type { AuthProvider } from '../../services/AuthProvider'
+import type { ContextAPIClient } from '../context/contextAPIClient'
 import {
     ChatController,
     type ChatSession,
@@ -69,7 +70,8 @@ export class ChatsController implements vscode.Disposable {
         private readonly localEmbeddings: LocalEmbeddingsController | null,
         private readonly contextRanking: ContextRankingController | null,
         private readonly symf: SymfRunner | null,
-        private readonly guardrails: Guardrails
+        private readonly guardrails: Guardrails,
+        private readonly contextAPIClient: ContextAPIClient | null
     ) {
         logDebug('ChatsController:constructor', 'init')
         this.panel = this.createChatController()
@@ -441,6 +443,7 @@ export class ChatsController implements vscode.Disposable {
             models,
             guardrails: this.guardrails,
             startTokenReceiver: this.options.startTokenReceiver,
+            contextAPIClient: this.contextAPIClient,
         })
     }
 

vscode/src/chat/context/contextAPIClient.ts

@@ -0,0 +1,59 @@
+import {
+    type ContextItem,
+    FeatureFlag,
+    type FeatureFlagProvider,
+    type SourcegraphGraphQLAPIClient,
+    isError,
+    logError,
+} from '@sourcegraph/cody-shared'
+import type { InputContextItem } from '@sourcegraph/cody-shared/src/sourcegraph-api/graphql/client'
+
+function toInput(input: ContextItem[]): InputContextItem[] {
+    return input
+        .map(i =>
+            !i || !i.content
+                ? null
+                : {
+                      content: i.content,
+                      retriever: i.source || '',
+                  }
+        )
+        .filter(i => i !== null) as InputContextItem[]
+}
+
+export class ContextAPIClient {
+    constructor(
+        private readonly apiClient: SourcegraphGraphQLAPIClient,
+        private readonly featureFlagProvider: FeatureFlagProvider
+    ) {}
+
+    public async detectChatIntent(interactionID: string, query: string) {
+        if (!(await this.isServerSideContextAPIEnabled())) {
+            return
+        }
+        return this.apiClient.chatIntent(interactionID, query)
+    }
+
+    public async rankContext(interactionID: string, query: string, context: ContextItem[]) {
+        if (!(await this.isServerSideContextAPIEnabled())) {
+            return
+        }
+        const res = await this.apiClient.rankContext(interactionID, query, toInput(context))
+        if (isError(res)) {
+            logError('rankContext', 'ranking result', res)
+            return res
+        }
+        return { used: res.rankContext.used, ignored: res.rankContext.ignored }
+    }
+
+    public async recordContext(interactionID: string, used: ContextItem[], ignored: ContextItem[]) {
+        if (!(await this.isServerSideContextAPIEnabled())) {
+            return
+        }
+        await this.apiClient.recordContext(interactionID, toInput(used), toInput(ignored))
+    }
+
+    private async isServerSideContextAPIEnabled() {
+        return await this.featureFlagProvider.evaluateFeatureFlag(FeatureFlag.CodyServerSideContextAPI)
+    }
+}

vscode/src/external-services.ts

@@ -7,10 +7,12 @@ import {
     type Guardrails,
     type SourcegraphCompletionsClient,
     SourcegraphGuardrailsClient,
+    featureFlagProvider,
     graphqlClient,
     isError,
 } from '@sourcegraph/cody-shared'
 
+import { ContextAPIClient } from './chat/context/contextAPIClient'
 import { createClient as createCodeCompletionsClient } from './completions/client'
 import type { ConfigWatcher } from './configwatcher'
 import type { PlatformContext } from './extension.common'
@@ -29,6 +31,7 @@ interface ExternalServices {
     contextRanking: ContextRankingController | undefined
     localEmbeddings: LocalEmbeddingsController | undefined
     symfRunner: SymfRunner | undefined
+    contextAPIClient: ContextAPIClient | undefined
 
     /** Update configuration for all of the services in this interface. */
     onConfigurationChange: (newConfig: ExternalServicesConfiguration) => void
@@ -86,6 +89,8 @@ export async function configureExternalServices(
 
     const guardrails = new SourcegraphGuardrailsClient(graphqlClient)
 
+    const contextAPIClient = new ContextAPIClient(graphqlClient, featureFlagProvider)
+
     return {
         chatClient,
         completionsClient,
@@ -94,6 +99,7 @@ export async function configureExternalServices(
         localEmbeddings,
         contextRanking,
         symfRunner,
+        contextAPIClient,
         onConfigurationChange: newConfig => {
             sentryService?.onConfigurationChange(newConfig)
             openTelemetryService?.onConfigurationChange(newConfig)