Add support for Cody API v1 (#3551)

sourcegraph · Mar 27, 2024 · bf57e36 · bf57e36
1 parent 3830c3a
commit bf57e36
Show file tree

Hide file tree

Showing 24 changed files with 246 additions and 433 deletions.
diff --git a/...bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/AuthStatus.kt b/...bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/AuthStatus.kt
@@ -12,6 +12,7 @@ data class AuthStatus(
   val requiresVerifiedEmail: Boolean? = null,
   val siteHasCodyEnabled: Boolean? = null,
   val siteVersion: String? = null,
+  val codyApiVersion: Int? = null,
   val configOverwrites: CodyLLMSiteConfiguration? = null,
   val showNetworkError: Boolean? = null,
   val primaryEmail: String? = null,

diff --git a/agent/src/index.test.ts b/agent/src/index.test.ts
@@ -292,7 +292,7 @@ describe('Agent', () => {
                 })
             )
             expect(reply2.messages.at(-1)?.text).toMatchInlineSnapshot(
-                '" You told me your name is Lars Monsen."',
+                `" You told me your name is Lars Monsen."`,
                 explainPollyError
             )
         }, 30_000)
@@ -950,22 +950,22 @@ describe('Agent', () => {
                 'TestClass.ts',
                 obtained =>
                     expect(obtained).toMatchInlineSnapshot(`
-                  "const foo = 42
+                      "const foo = 42
 
-                  export class TestClass {
-                      constructor(private shouldGreet: boolean) {}
+                      export class TestClass {
+                          constructor(private shouldGreet: boolean) {}
 
-                          /**
-                       * Logs a greeting message if the shouldGreet flag is true.
-                       */
-                  public functionName() {
-                          if (this.shouldGreet) {
-                              console.log(/* CURSOR */ 'Hello World!')
+                              /**
+                           * Logs a greeting message if the shouldGreet flag is true.
+                           */
+                      public functionName() {
+                              if (this.shouldGreet) {
+                                  console.log(/* CURSOR */ 'Hello World!')
+                              }
                           }
                       }
-                  }
-                  "
-                `)
+                      "
+                    `)
             )
 
             checkDocumentCommand(
@@ -974,23 +974,23 @@ describe('Agent', () => {
                 'TestLogger.ts',
                 obtained =>
                     expect(obtained).toMatchInlineSnapshot(`
-                  "const foo = 42
-                  export const TestLogger = {
-                      startLogging: () => {
-                          // Do some stuff
-
-                                  /**
-                           * Logs a message to indicate when logging activity begins.
-                           */
-                  function recordLog() {
-                              console.log(/* CURSOR */ 'Recording the log')
-                          }
-
-                          recordLog()
-                      },
-                  }
-                  "
-                `)
+                      "const foo = 42
+                      export const TestLogger = {
+                          startLogging: () => {
+                              // Do some stuff
+
+                                      /**
+                               * Logs a message to indicate when logging activity begins.
+                               */
+                      function recordLog() {
+                                  console.log(/* CURSOR */ 'Recording the log')
+                              }
+
+                              recordLog()
+                          },
+                      }
+                      "
+                    `)
             )
 
             checkDocumentCommand(
@@ -999,28 +999,28 @@ describe('Agent', () => {
                 'example.test.ts',
                 obtained =>
                     expect(obtained).toMatchInlineSnapshot(`
-                  "import { expect } from 'vitest'
-                  import { it } from 'vitest'
-                  import { describe } from 'vitest'
-
-                  describe('test block', () => {
-                      it('does 1', () => {
-                          expect(true).toBe(true)
-                      })
-
-                      it('does 2', () => {
-                          expect(true).toBe(true)
-                      })
-
-                      it('does something else', () => {
-                          // This line will error due to incorrect usage of \`performance.now\`
-                                  // Record start time for performance measurement
-                          const startTime = performance.now();
-                  const startTime = performance.now(/* CURSOR */)
+                      "import { expect } from 'vitest'
+                      import { it } from 'vitest'
+                      import { describe } from 'vitest'
+
+                      describe('test block', () => {
+                          it('does 1', () => {
+                              expect(true).toBe(true)
+                          })
+
+                          it('does 2', () => {
+                              expect(true).toBe(true)
+                          })
+
+                          it('does something else', () => {
+                              // This line will error due to incorrect usage of \`performance.now\`
+                                      // Record start time for performance measurement
+                              const startTime = performance.now();
+                      const startTime = performance.now(/* CURSOR */)
+                          })
                       })
-                  })
-                  "
-                `)
+                      "
+                    `)
             )
         })
     })
@@ -1349,32 +1349,32 @@ describe('Agent', () => {
             'example.test.ts',
             obtained =>
                 expect(obtained).toMatchInlineSnapshot(`
-              "import { expect } from 'vitest'
-              import { it } from 'vitest'
-              import { describe } from 'vitest'
+                  "import { expect } from 'vitest'
+                  import { it } from 'vitest'
+                  import { describe } from 'vitest'
 
-              /**
-               * Test block that contains 3 test cases:
-               * - Does test 1
-               * - Does test 2
-               * - Does something else (has incorrect usage of \`performance.now\`)
-               */
-              describe('test block', () => {
-                  it('does 1', () => {
-                      expect(true).toBe(true)
-                  })
+                  /**
+                   * Test block that contains 3 test cases:
+                   * - Does test 1
+                   * - Does test 2
+                   * - Does something else (has incorrect usage of \`performance.now\`)
+                   */
+                  describe('test block', () => {
+                      it('does 1', () => {
+                          expect(true).toBe(true)
+                      })
 
-                  it('does 2', () => {
-                      expect(true).toBe(true)
-                  })
+                      it('does 2', () => {
+                          expect(true).toBe(true)
+                      })
 
-                  it('does something else', () => {
-                      // This line will error due to incorrect usage of \`performance.now\`
-                      const startTime = performance.now(/* CURSOR */)
+                      it('does something else', () => {
+                          // This line will error due to incorrect usage of \`performance.now\`
+                          const startTime = performance.now(/* CURSOR */)
+                      })
                   })
-              })
-              "
-            `)
+                  "
+                `)
         )
 
         // NOTE(olafurpg) disabled on Windows because the multi-repo keyword

diff --git a/lib/shared/src/auth/types.ts b/lib/shared/src/auth/types.ts
@@ -15,6 +15,7 @@ export interface AuthStatus {
     requiresVerifiedEmail: boolean
     siteHasCodyEnabled: boolean
     siteVersion: string
+    codyApiVersion: number
     configOverwrites?: CodyLLMSiteConfiguration
     showNetworkError?: boolean
     primaryEmail: string

diff --git a/lib/shared/src/chat/chat.ts b/lib/shared/src/chat/chat.ts
@@ -1,18 +1,11 @@
-import { dotcomTokenToGatewayToken } from '../auth/tokens'
 import type { AuthStatus } from '../auth/types'
-import type { ConfigurationWithAccessToken } from '../configuration'
-import { supportsFastPath } from '../models/utils'
 import { ANSWER_TOKENS } from '../prompt/constants'
 import type { Message } from '../sourcegraph-api'
-import type {
-    CompletionLogger,
-    SourcegraphCompletionsClient,
-} from '../sourcegraph-api/completions/client'
+import type { SourcegraphCompletionsClient } from '../sourcegraph-api/completions/client'
 import type {
     CompletionGeneratorValue,
     CompletionParameters,
 } from '../sourcegraph-api/completions/types'
-import { createFastPathClient } from './fast-path-client'
 
 type ChatParameters = Omit<CompletionParameters, 'messages'>
 
@@ -24,41 +17,25 @@ const DEFAULT_CHAT_COMPLETION_PARAMETERS: ChatParameters = {
 }
 
 export class ChatClient {
-    private fastPathAccessToken?: string
-
     constructor(
         private completions: SourcegraphCompletionsClient,
-        config: Pick<ConfigurationWithAccessToken, 'accessToken'>,
-        private getAuthStatus: () => Pick<AuthStatus, 'userCanUpgrade' | 'isDotCom' | 'endpoint'>,
-        private completionLogger: CompletionLogger
-    ) {
-        this.onConfigurationChange(config)
-    }
-
-    public onConfigurationChange(newConfig: Pick<ConfigurationWithAccessToken, 'accessToken'>): void {
-        const isNode = typeof process !== 'undefined'
-        this.fastPathAccessToken =
-            newConfig.accessToken &&
-            // Require the upstream to be dotcom
-            this.getAuthStatus().isDotCom &&
-            // The fast path client only supports Node.js style response streams
-            isNode
-                ? dotcomTokenToGatewayToken(newConfig.accessToken)
-                : undefined
-    }
+        private getAuthStatus: () => Pick<
+            AuthStatus,
+            'userCanUpgrade' | 'isDotCom' | 'endpoint' | 'codyApiVersion'
+        >
+    ) {}
 
     public chat(
         messages: Message[],
         params: Partial<ChatParameters>,
         abortSignal?: AbortSignal
     ): AsyncGenerator<CompletionGeneratorValue> {
-        const useFastPath =
-            this.fastPathAccessToken !== undefined && params.model && supportsFastPath(params.model)
-
+        const authStatus = this.getAuthStatus()
+        const useApiV1 = authStatus.codyApiVersion >= 1 && params.model?.includes('claude-3')
         const isLastMessageFromHuman = messages.length > 0 && messages.at(-1)!.speaker === 'human'
 
         const augmentedMessages =
-            params?.model?.startsWith('fireworks/') || useFastPath
+            params?.model?.startsWith('fireworks/') || useApiV1
                 ? sanitizeMessages(messages)
                 : isLastMessageFromHuman
                   ? messages.concat([{ speaker: 'assistant' }])
@@ -77,17 +54,11 @@ export class ChatClient {
             messages: messagesToSend,
         }
 
-        if (useFastPath) {
-            return createFastPathClient(
-                completionParams,
-                this.getAuthStatus(),
-                this.fastPathAccessToken!,
-                abortSignal,
-                this.completionLogger
-            )
-        }
-
-        return this.completions.stream(completionParams, abortSignal)
+        return this.completions.stream(
+            completionParams,
+            useApiV1 ? authStatus.codyApiVersion : 0,
+            abortSignal
+        )
     }
 }