Skip to content

Commit

Permalink
fix(chat): support non-streaming chat completion requests (#5565)
Browse files Browse the repository at this point in the history
Follow up on #5508
CLOSE
https://linear.app/sourcegraph/issue/CODY-3741/oai-o1-hits-parsing-error

This change adds support for non-streaming requests, particularly for
new OpenAI models that do not support streaming. It also fixes a "Parse
Error: JS exception" issue that occurred when handling non-streaming
responses in the streaming client.

- Remove HACK for handling non-streaming requests in the stream method
in nodeClient.ts that's resulting in parsing error
- Implement a new `_fetchWithCallbacks` method in the
`SourcegraphCompletionsClient` class to handle non-streaming requests
- Update the `SourcegraphNodeCompletionsClient` and
`SourcegraphBrowserCompletionsClient` classes to use the new
`_fetchWithCallbacks` method when `params.stream` is `false`
- This change ensures that non-streaming requests, such as those for the
new OpenAI models that do not support streaming, are handled correctly
and avoid issues like "Parse Error: JS exception" that can occur when
the response is not parsed correctly in the streaming client

The `SourcegraphCompletionsClient` class now has a new abstract method
`_fetchWithCallbacks` that must be implemented by subclasses

## Test plan

<!-- Required. See
https://docs-legacy.sourcegraph.com/dev/background-information/testing_principles.
-->

Verify that chat requests using non-streaming models (e.g., OpenAI O1,
OpenAI O1-mini) work correctly without any parsing errors. Example: ask
Cody about cody repo with "cody what are the latest updates in this
repo?"

### After

First try


![image](https://github.com/user-attachments/assets/24fbd1b0-185d-4b07-90be-385fdd5da975)

Second try


![image](https://github.com/user-attachments/assets/becdc7f6-18bb-4fec-8fcc-8263072680f8)

### Before

First try


![image](https://github.com/user-attachments/assets/8975cb5a-b605-4e0a-ac2c-dab053124920)

## Changelog

<!-- OPTIONAL; info at
https://www.notion.so/sourcegraph/Writing-a-changelog-entry-dd997f411d524caabf0d8d38a24a878c
-->

feat(chat): support non-streaming requests

---------

Co-authored-by: Valery Bugakov <skymk1@gmail.com>
  • Loading branch information
abeatrix and valerybugakov authored Sep 13, 2024
1 parent ccab0ed commit 24192e4
Show file tree
Hide file tree
Showing 4 changed files with 143 additions and 17 deletions.
49 changes: 48 additions & 1 deletion lib/shared/src/sourcegraph-api/completions/browserClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import { addClientInfoParams } from '../client-name-version'
import { CompletionsResponseBuilder } from './CompletionsResponseBuilder'
import { type CompletionRequestParameters, SourcegraphCompletionsClient } from './client'
import { parseCompletionJSON } from './parse'
import type { CompletionCallbacks, CompletionParameters, Event } from './types'
import type { CompletionCallbacks, CompletionParameters, CompletionResponse, Event } from './types'
import { getSerializedParams } from './utils'

declare const WorkerGlobalScope: never
Expand Down Expand Up @@ -115,6 +115,53 @@ export class SourcegraphBrowserCompletionsClient extends SourcegraphCompletionsC
console.error(error)
})
}

protected async _fetchWithCallbacks(
params: CompletionParameters,
requestParams: CompletionRequestParameters,
cb: CompletionCallbacks,
signal?: AbortSignal
): Promise<void> {
const { url, serializedParams } = await this.prepareRequest(params, requestParams)
const headersInstance = new Headers({
'Content-Type': 'application/json; charset=utf-8',
...this.config.customHeaders,
...requestParams.customHeaders,
})
addCustomUserAgent(headersInstance)
if (this.config.accessToken) {
headersInstance.set('Authorization', `token ${this.config.accessToken}`)
}
if (new URLSearchParams(globalThis.location.search).get('trace')) {
headersInstance.set('X-Sourcegraph-Should-Trace', 'true')
}
try {
const response = await fetch(url.toString(), {
method: 'POST',
headers: headersInstance,
body: JSON.stringify(serializedParams),
signal,
})
if (!response.ok) {
const errorMessage = await response.text()
throw new Error(
errorMessage.length === 0
? `Request failed with status code ${response.status}`
: errorMessage
)
}
const data = (await response.json()) as CompletionResponse
if (data?.completion) {
cb.onChange(data.completion)
cb.onComplete()
} else {
throw new Error('Unexpected response format')
}
} catch (error) {
console.error(error)
cb.onError(error instanceof Error ? error : new Error(`${error}`))
}
}
}

if (isRunningInWebWorker) {
Expand Down
32 changes: 30 additions & 2 deletions lib/shared/src/sourcegraph-api/completions/client.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import type { Span } from '@opentelemetry/api'
import { addClientInfoParams, getSerializedParams } from '../..'
import type { ClientConfigurationWithAccessToken } from '../../configuration'

import { useCustomChatClient } from '../../llm-providers'
import { recordErrorToSpan } from '../../tracing'
import type {
Expand All @@ -9,6 +9,7 @@ import type {
CompletionParameters,
CompletionResponse,
Event,
SerializedCompletionParameters,
} from './types'

export interface CompletionLogger {
Expand Down Expand Up @@ -45,6 +46,8 @@ export type CompletionsClientConfig = Pick<
export abstract class SourcegraphCompletionsClient {
private errorEncountered = false

protected readonly isTemperatureZero = process.env.CODY_TEMPERATURE_ZERO === 'true'

constructor(
protected config: CompletionsClientConfig,
protected logger?: CompletionLogger
Expand Down Expand Up @@ -88,6 +91,27 @@ export abstract class SourcegraphCompletionsClient {
}
}

protected async prepareRequest(
params: CompletionParameters,
requestParams: CompletionRequestParameters
): Promise<{ url: URL; serializedParams: SerializedCompletionParameters }> {
const { apiVersion } = requestParams
const serializedParams = await getSerializedParams(params)
const url = new URL(this.completionsEndpoint)
if (apiVersion >= 1) {
url.searchParams.append('api-version', '' + apiVersion)
}
addClientInfoParams(url.searchParams)
return { url, serializedParams }
}

protected abstract _fetchWithCallbacks(
params: CompletionParameters,
requestParams: CompletionRequestParameters,
cb: CompletionCallbacks,
signal?: AbortSignal
): Promise<void>

protected abstract _streamWithCallbacks(
params: CompletionParameters,
requestParams: CompletionRequestParameters,
Expand Down Expand Up @@ -144,7 +168,11 @@ export abstract class SourcegraphCompletionsClient {
})

if (!isNonSourcegraphProvider) {
await this._streamWithCallbacks(params, requestParams, callbacks, signal)
if (params.stream === false) {
await this._fetchWithCallbacks(params, requestParams, callbacks, signal)
} else {
await this._streamWithCallbacks(params, requestParams, callbacks, signal)
}
}

for (let i = 0; ; i++) {
Expand Down
1 change: 1 addition & 0 deletions vscode/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ This is a log of all notable changes to Cody for VS Code. [Unreleased] changes a

- The [new OpenAI models (OpenAI O1 & OpenAI O1-mini)](https://sourcegraph.com/blog/openai-o1-for-cody) are now available to selected Cody Pro users for early access. [pull/5508](https://github.com/sourcegraph/cody/pull/5508)
- Cody Pro users can join the waitlist for the new models by clicking the "Join Waitlist" button. [pull/5508](https://github.com/sourcegraph/cody/pull/5508)
- Chat: Support non-streaming requests. [pull/5565](https://github.com/sourcegraph/cody/pull/5565)

### Fixed

Expand Down
78 changes: 64 additions & 14 deletions vscode/src/completions/nodeClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import {
type CompletionCallbacks,
type CompletionParameters,
type CompletionRequestParameters,
type CompletionResponse,
NetworkError,
RateLimitError,
SourcegraphCompletionsClient,
Expand All @@ -29,8 +30,6 @@ import {
} from '@sourcegraph/cody-shared'
import { CompletionsResponseBuilder } from '@sourcegraph/cody-shared/src/sourcegraph-api/completions/CompletionsResponseBuilder'

const isTemperatureZero = process.env.CODY_TEMPERATURE_ZERO === 'true'

export class SourcegraphNodeCompletionsClient extends SourcegraphCompletionsClient {
protected _streamWithCallbacks(
params: CompletionParameters,
Expand All @@ -56,7 +55,7 @@ export class SourcegraphNodeCompletionsClient extends SourcegraphCompletionsClie
model: params.model,
})

if (isTemperatureZero) {
if (this.isTemperatureZero) {
params = {
...params,
temperature: 0,
Expand Down Expand Up @@ -203,17 +202,6 @@ export class SourcegraphNodeCompletionsClient extends SourcegraphCompletionsClie
bufferText += str
bufferBin = buf

// HACK: Handles non-stream request.
// TODO: Implement a function to make and process non-stream requests.
if (params.stream === false) {
const json = JSON.parse(bufferText)
if (json?.completion) {
cb.onChange(json.completion)
cb.onComplete()
return
}
}

const parseResult = parseEvents(builder, bufferText)
if (isError(parseResult)) {
logError(
Expand Down Expand Up @@ -286,6 +274,68 @@ export class SourcegraphNodeCompletionsClient extends SourcegraphCompletionsClie
onAbort(signal, () => request.destroy())
})
}

protected async _fetchWithCallbacks(
params: CompletionParameters,
requestParams: CompletionRequestParameters,
cb: CompletionCallbacks,
signal?: AbortSignal
): Promise<void> {
const { url, serializedParams } = await this.prepareRequest(params, requestParams)
const log = this.logger?.startCompletion(params, url.toString())
return tracer.startActiveSpan(`POST ${url.toString()}`, async span => {
span.setAttributes({
fast: params.fast,
maxTokensToSample: params.maxTokensToSample,
temperature: this.isTemperatureZero ? 0 : params.temperature,
topK: params.topK,
topP: params.topP,
model: params.model,
})
try {
const response = await fetch(url.toString(), {
method: 'POST',
headers: {
'Content-Type': 'application/json',
'Accept-Encoding': 'gzip;q=0',
...(this.config.accessToken
? { Authorization: `token ${this.config.accessToken}` }
: null),
...(customUserAgent ? { 'User-Agent': customUserAgent } : null),
...this.config.customHeaders,
...requestParams.customHeaders,
...getTraceparentHeaders(),
},
body: JSON.stringify(serializedParams),
signal,
})
if (!response.ok) {
const errorMessage = await response.text()
throw new NetworkError(
{
url: url.toString(),
status: response.status,
statusText: response.statusText,
},
errorMessage,
getActiveTraceAndSpanId()?.traceId
)
}
const json = (await response.json()) as CompletionResponse
if (typeof json?.completion === 'string') {
cb.onChange(json.completion)
cb.onComplete()
return
}
throw new Error('Unexpected response format')
} catch (error) {
const errorObject = error instanceof Error ? error : new Error(`${error}`)
log?.onError(errorObject.message, error)
recordErrorToSpan(span, errorObject)
cb.onError(errorObject)
}
})
}
}

function getHeader(value: string | undefined | string[]): string | undefined {
Expand Down

0 comments on commit 24192e4

Please sign in to comment.