fix(cody-gateway): getAPIURL before transformBody (#63406)

<!-- 💡 To write a useful PR description, make sure that your description
covers:
- WHAT this PR is changing:
    - How was it PREVIOUSLY.
    - How it will be from NOW on.
- WHY this PR is needed.
- CONTEXT, i.e. to which initiative, project or RFC it belongs.

The structure of the description doesn't matter as much as covering
these points, so use
your best judgement based on your context.
Learn how to write good pull request description:
https://www.notion.so/sourcegraph/Write-a-good-pull-request-description-610a7fd3e613496eb76f450db5a49b6e?pvs=4
-->

Fix an issue where the requestBody is used after transformBody has been
executed.


## Test plan

<!-- All pull requests REQUIRE a test plan:
https://docs-legacy.sourcegraph.com/dev/background-information/testing_principles
-->

1. Start Cody Gateway locally
2. Start SG local dev instance
3. Connect SG local dev instance to your local Cody Gateway instance
4. Set Gemini Flash as your chatModel
5. Connect Cody to your local dev instance
6. Ask Cody a question and verify you are getting a response


![image](https://github.com/sourcegraph/sourcegraph/assets/68532117/fbce22f9-8531-4f6e-8eb7-5c6b26e0a9fa)



## Changelog

<!--
1. Ensure your pull request title is formatted as: $type($domain): $what
8. Add bullet list items for each additional detail you want to cover
(see example below)
9. You can edit this after the pull request was merged, as long as
release shipping it hasn't been promoted to the public.
10. For more information, please see this how-to
https://www.notion.so/sourcegraph/Writing-a-changelog-entry-dd997f411d524caabf0d8d38a24a878c?

Audience: TS/CSE > Customers > Teammates (in that order).

Cheat sheet: $type = chore|fix|feat $domain:
source|search|ci|release|plg|cody|local|...
-->

<!--
Example:

Title: fix(search): parse quotes with the appropriate context
Changelog section:

## Changelog

- When a quote is used with regexp pattern type, then ...
- Refactored underlying code.
-->

Author: abeatrix

cmd/cody-gateway/internal/httpapi/completions/anthropic.go

@@ -166,7 +166,7 @@ func (a *AnthropicHandlerMethods) transformRequest(r *http.Request) {
 	r.Header.Set("anthropic-version", "2023-01-01")
 }
 
-func (a *AnthropicHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody anthropicRequest, r io.Reader) (promptUsage, completionUsage usageStats) {
+func (a *AnthropicHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody anthropicRequest, r io.Reader, isStreamRequest bool) (promptUsage, completionUsage usageStats) {
 	var err error
 
 	// Setting a default -1 value so that in case of errors the tokenizer computed tokens don't impact the data
@@ -183,7 +183,7 @@ func (a *AnthropicHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBo
 
 	// Try to parse the request we saw, if it was non-streaming, we can simply parse
 	// it as JSON.
-	if !reqBody.Stream {
+	if !isStreamRequest {
 		var res anthropicResponse
 		if err := json.NewDecoder(r).Decode(&res); err != nil {
 			logger.Error("failed to parse Anthropic response as JSON", log.Error(err))

cmd/cody-gateway/internal/httpapi/completions/anthropicmessages.go

@@ -215,7 +215,7 @@ func (a *AnthropicMessagesHandlerMethods) transformRequest(r *http.Request) {
 	r.Header.Set("anthropic-version", "2023-06-01")
 }
 
-func (a *AnthropicMessagesHandlerMethods) parseResponseAndUsage(logger log.Logger, body anthropicMessagesRequest, r io.Reader) (promptUsage, completionUsage usageStats) {
+func (a *AnthropicMessagesHandlerMethods) parseResponseAndUsage(logger log.Logger, body anthropicMessagesRequest, r io.Reader, isStreamRequest bool) (promptUsage, completionUsage usageStats) {
 	// First, extract prompt usage details from the request.
 	for _, m := range body.Messages {
 		promptUsage.characters += len(m.Content)
@@ -232,7 +232,7 @@ func (a *AnthropicMessagesHandlerMethods) parseResponseAndUsage(logger log.Logge
 
 	// Try to parse the request we saw, if it was non-streaming, we can simply parse
 	// it as JSON.
-	if !body.ShouldStream() {
+	if !isStreamRequest {
 		var res anthropicMessagesNonStreamingResponse
 		if err := json.NewDecoder(r).Decode(&res); err != nil {
 			logger.Error("failed to parse Anthropic response as JSON", log.Error(err))

cmd/cody-gateway/internal/httpapi/completions/fireworks.go

@@ -148,13 +148,13 @@ func (f *FireworksHandlerMethods) transformRequest(r *http.Request) {
 	r.Header.Set("Authorization", "Bearer "+f.config.AccessToken)
 }
 
-func (f *FireworksHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody fireworksRequest, r io.Reader) (promptUsage, completionUsage usageStats) {
+func (f *FireworksHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody fireworksRequest, r io.Reader, isStreamRequest bool) (promptUsage, completionUsage usageStats) {
 	// First, extract prompt usage details from the request.
 	promptUsage.characters = len(reqBody.Prompt)
 
 	// Try to parse the request we saw, if it was non-streaming, we can simply parse
 	// it as JSON.
-	if !reqBody.Stream {
+	if !isStreamRequest {
 		var res fireworksResponse
 		if err := json.NewDecoder(r).Decode(&res); err != nil {
 			logger.Error("failed to parse fireworks response as JSON", log.Error(err))

cmd/cody-gateway/internal/httpapi/completions/fireworks_test.go

@@ -18,7 +18,7 @@ func TestFireworksRequestGetTokenCount(t *testing.T) {
 		req := fireworksRequest{Stream: true}
 		r := strings.NewReader(fireworksStreamingResponse)
 		handler := &FireworksHandlerMethods{}
-		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r)
+		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r, true)
 
 		assert.Equal(t, 79, promptUsage.tokens)
 		assert.Equal(t, 30, completionUsage.tokens)
@@ -28,7 +28,7 @@ func TestFireworksRequestGetTokenCount(t *testing.T) {
 		req := fireworksRequest{Stream: false}
 		r := strings.NewReader(fireworksNonStreamingResponse)
 		handler := &FireworksHandlerMethods{}
-		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r)
+		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r, false)
 
 		assert.Equal(t, 79, promptUsage.tokens)
 		assert.Equal(t, 30, completionUsage.tokens)

cmd/cody-gateway/internal/httpapi/completions/google.go

@@ -64,10 +64,9 @@ func (g *GoogleHandlerMethods) getAPIURL(feature codygateway.Feature, req google
 	rpc := "generateContent"
 	sseSuffix := ""
 	// If we're streaming, we need to use the stream endpoint.
-	if feature == codygateway.FeatureChatCompletions || req.ShouldStream() {
+	if req.ShouldStream() {
 		rpc = "streamGenerateContent"
 		sseSuffix = "&alt=sse"
-		req.Stream = true
 	}
 	return fmt.Sprintf("https://generativelanguage.googleapis.com/v1beta/models/%s:%s?key=%s%s", req.Model, rpc, g.config.AccessToken, sseSuffix)
 }
@@ -108,12 +107,12 @@ func (o *GoogleHandlerMethods) transformRequest(r *http.Request) {
 	r.Header.Set("Content-Type", "application/json")
 }
 
-func (*GoogleHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody googleRequest, r io.Reader) (promptUsage, completionUsage usageStats) {
+func (*GoogleHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody googleRequest, r io.Reader, isStreamRequest bool) (promptUsage, completionUsage usageStats) {
 	// First, extract prompt usage details from the request.
 	promptUsage.characters = len(reqBody.BuildPrompt())
 	// Try to parse the request we saw, if it was non-streaming, we can simply parse
 	// it as JSON.
-	if !reqBody.Stream && !reqBody.ShouldStream() {
+	if !isStreamRequest {
 		var res googleResponse
 		if err := json.NewDecoder(r).Decode(&res); err != nil {
 			logger.Error("failed to parse Google response as JSON", log.Error(err))
@@ -135,10 +134,10 @@ func (*GoogleHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody go
 	if err != nil {
 		logger.Error("failed to decode Google streaming response", log.Error(err))
 	}
+	promptUsage.tokens, completionUsage.tokens = promptTokens, completionTokens
 	if completionUsage.tokens == -1 || promptUsage.tokens == -1 {
 		logger.Warn("did not extract token counts from Google streaming response", log.Int("prompt-tokens", promptUsage.tokens), log.Int("completion-tokens", completionUsage.tokens))
 	}
-	promptUsage.tokens, completionUsage.tokens = promptTokens, completionTokens
 	return promptUsage, completionUsage
 }
 

cmd/cody-gateway/internal/httpapi/completions/google_test.go

@@ -17,7 +17,7 @@ func TestGoogleRequestGetTokenCount(t *testing.T) {
 		req := googleRequest{Stream: true}
 		r := strings.NewReader(googleStreamingResponse)
 		handler := &GoogleHandlerMethods{}
-		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r)
+		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r, true)
 
 		assert.Equal(t, 21, promptUsage.tokens)
 		assert.Equal(t, 87, completionUsage.tokens)
@@ -27,7 +27,7 @@ func TestGoogleRequestGetTokenCount(t *testing.T) {
 		req := googleRequest{Stream: false}
 		r := strings.NewReader(googleNonStreamingResponse)
 		handler := &GoogleHandlerMethods{}
-		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r)
+		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r, false)
 
 		assert.Equal(t, 59, promptUsage.tokens)
 		assert.Equal(t, 54, completionUsage.tokens)

cmd/cody-gateway/internal/httpapi/completions/openai.go

@@ -153,7 +153,7 @@ func (o *OpenAIHandlerMethods) transformRequest(r *http.Request) {
 	}
 }
 
-func (*OpenAIHandlerMethods) parseResponseAndUsage(logger log.Logger, body openaiRequest, r io.Reader) (promptUsage, completionUsage usageStats) {
+func (*OpenAIHandlerMethods) parseResponseAndUsage(logger log.Logger, body openaiRequest, r io.Reader, isStreamRequest bool) (promptUsage, completionUsage usageStats) {
 	// First, extract prompt usage details from the request.
 	for _, m := range body.Messages {
 		promptUsage.characters += len(m.Content)
@@ -164,7 +164,7 @@ func (*OpenAIHandlerMethods) parseResponseAndUsage(logger log.Logger, body opena
 	completionUsage.tokenizerTokens = -1
 	// Try to parse the request we saw, if it was non-streaming, we can simply parse
 	// it as JSON.
-	if !body.Stream {
+	if !isStreamRequest {
 		var res openaiResponse
 		if err := json.NewDecoder(r).Decode(&res); err != nil {
 			logger.Error("failed to parse OpenAI response as JSON", log.Error(err))

cmd/cody-gateway/internal/httpapi/completions/openai_test.go

@@ -15,7 +15,7 @@ func TestOpenAIRequestGetTokenCount(t *testing.T) {
 		req := openaiRequest{Stream: true}
 		r := strings.NewReader(openaiStreamingResponse)
 		handler := &OpenAIHandlerMethods{}
-		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r)
+		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r, true)
 
 		assert.Equal(t, 427, promptUsage.tokens)
 		assert.Equal(t, 12, completionUsage.tokens)
@@ -25,7 +25,7 @@ func TestOpenAIRequestGetTokenCount(t *testing.T) {
 		req := openaiRequest{Stream: false}
 		r := strings.NewReader(openaiNonStreamingResponse)
 		handler := &OpenAIHandlerMethods{}
-		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r)
+		promptUsage, completionUsage := handler.parseResponseAndUsage(logger, req, r, false)
 
 		assert.Equal(t, 12, promptUsage.tokens)
 		assert.Equal(t, 9, completionUsage.tokens)

cmd/cody-gateway/internal/httpapi/completions/upstream.go

@@ -104,7 +104,7 @@ type upstreamHandlerMethods[ReqT UpstreamRequest] interface {
 	//
 	// If data is unavailable, implementations should set relevant usage fields
 	// to -1 as a sentinel value.
-	parseResponseAndUsage(log.Logger, ReqT, io.Reader) (promptUsage, completionUsage usageStats)
+	parseResponseAndUsage(log.Logger, ReqT, io.Reader, bool) (promptUsage, completionUsage usageStats)
 }
 
 type UpstreamRequest interface {
@@ -273,6 +273,13 @@ func makeUpstreamHandler[ReqT UpstreamRequest](
 			}
 		}
 
+		// Get the URL to call for the upstream provider before we transform the request.
+		upstreamURL := methods.getAPIURL(feature, body)
+
+		// Store the shouldStream value in case it changes during the transformation.
+		// Example: We remove it for Google requests.
+		shouldStream := body.ShouldStream()
+
 		// identifier that can be provided to upstream for abuse detection
 		// has the format '$ACTOR_ID:$SG_ACTOR_ID'. The latter is anonymized
 		// (specific per-instance)
@@ -288,7 +295,6 @@ func makeUpstreamHandler[ReqT UpstreamRequest](
 		}
 
 		// Create a new request to send upstream, making sure we retain the same context.
-		upstreamURL := methods.getAPIURL(feature, body)
 		req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(upstreamPayload))
 		if err != nil {
 			response.JSONError(logger, w, http.StatusInternalServerError, errors.Wrap(err, "failed to create request"))
@@ -359,7 +365,7 @@ func makeUpstreamHandler[ReqT UpstreamRequest](
 			o.Feature = feature
 			o.UpstreamLatency = upstreamLatency
 			o.Provider = upstreamName
-			o.Stream = body.ShouldStream()
+			o.Stream = shouldStream
 
 			err := eventLogger.LogEvent(
 				ctx,
@@ -461,7 +467,7 @@ func makeUpstreamHandler[ReqT UpstreamRequest](
 		// if this is a streaming request, we want to flush ourselves instead of leaving that to the http.Server
 		// (so events are sent to the client as soon as possible)
 		var responseWriter io.Writer = w
-		if config.AutoFlushStreamingResponses && body.ShouldStream() {
+		if config.AutoFlushStreamingResponses && shouldStream {
 			if fw, err := response.NewAutoFlushingWriter(w); err == nil {
 				responseWriter = fw
 			} else {
@@ -475,7 +481,7 @@ func makeUpstreamHandler[ReqT UpstreamRequest](
 
 		if upstreamStatusCode >= 200 && upstreamStatusCode < 300 {
 			// Pass reader to response transformer to capture token counts.
-			promptUsage, completionUsage = methods.parseResponseAndUsage(logger, body, &responseBuf)
+			promptUsage, completionUsage = methods.parseResponseAndUsage(logger, body, &responseBuf, shouldStream)
 		} else if upstreamStatusCode >= 500 {
 			logger.Error("error from upstream",
 				log.Int("status_code", upstreamStatusCode))

cmd/frontend/internal/dotcom/productsubscription/codygateway_dotcom_user.go

@@ -334,6 +334,9 @@ var allCodeCompletionModels = slices.Concat([]string{"anthropic/" + anthropic.Cl
 	"anthropic/claude-instant-1.2-cyan",
 	"anthropic/claude-instant-1.2",
 	"google/" + google.Gemini15Flash,
+	"google/" + google.Gemini15FlashLatest,
+	"google/" + google.GeminiPro,
+	"google/" + google.GeminiProLatest,
 	"fireworks/starcoder",
 	"fireworks/" + fireworks.Llama213bCode,
 	"fireworks/" + fireworks.StarcoderTwo15b,