e2e: add agent-tool test

internal/extproc/translator/openai_awsbedrock.go

@@ -588,6 +588,8 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) ResponseBody(respHeaders
 		if toolCall := o.bedrockToolUseToOpenAICalls(output.ToolUse); toolCall != nil {
 			choice.Message.ToolCalls = []openai.ChatCompletionMessageToolCallParam{*toolCall}
 		}
+		// TODO: merge the choice message with the
+		// update dont append
 		openAIResp.Choices = append(openAIResp.Choices, choice)
 	}
 

tests/extproc/real_providers_test.go

@@ -181,34 +181,106 @@
 			})
 		}
 	})
-
-	t.Run("Bedrock calls tool get_weather function", func(t *testing.T) {
-		cc.maybeSkip(t, requiredCredentialAWS)
+
+	t.Run("Bedrock uses tool in response", func(t *testing.T) {
 		client := openai.NewClient(option.WithBaseURL(listenerAddress + "/v1/"))
-		require.Eventually(t, func() bool {
-			chatCompletion, err := client.Chat.Completions.New(context.Background(), openai.ChatCompletionNewParams{
-				Messages: openai.F([]openai.ChatCompletionMessageParamUnion{
-					openai.UserMessage("What is the weather like in Paris today?"),
-				}),
-				Tools: openai.F([]openai.ChatCompletionToolParam{
-					{
-						Type: openai.F(openai.ChatCompletionToolTypeFunction),
-						Function: openai.F(openai.FunctionDefinitionParam{
-							Name:        openai.String("get_weather"),
-							Description: openai.String("Get weather at the given location"),
-							Parameters: openai.F(openai.FunctionParameters{
-								"type": "object",
-								"properties": map[string]interface{}{
-									"location": map[string]string{
-										"type": "string",
-									},
-								},
-								"required": []string{"location"},
-							}),
+		for _, tc := range []struct {
+			testCaseName,
+			modelName string
+		}{
+			{testCaseName: "aws-bedrock", modelName: "us.anthropic.claude-3-5-sonnet-20240620-v1:0"}, // This will go to "aws-bedrock" using credentials file.
+		} {
+			t.Run(tc.modelName, func(t *testing.T) {
+				require.Eventually(t, func() bool {
+					// Step 1: Initial tool call request
+					question := "What is the weather in New York City?"
+					params := openai.ChatCompletionNewParams{
+						Messages: openai.F([]openai.ChatCompletionMessageParamUnion{
+							openai.UserMessage(question),
 						}),
-					},
-				}),
-				Model: openai.F("us.anthropic.claude-3-5-sonnet-20240620-v1:0"),
+						Tools: openai.F([]openai.ChatCompletionToolParam{
+							{
+								Type: openai.F(openai.ChatCompletionToolTypeFunction),
+								Function: openai.F(openai.FunctionDefinitionParam{
+									Name:        openai.String("get_weather"),
+									Description: openai.String("Get weather at the given location"),
+									Parameters: openai.F(openai.FunctionParameters{
+										"type": "object",
+										"properties": map[string]interface{}{
+											"location": map[string]string{
+												"type": "string",
+											},
+										},
+										"required": []string{"location"},
+									}),
+								}),
+							},
+						}),
+						// TODO: check if we should seed.
+						Seed:  openai.Int(0),
+						Model: openai.F(tc.modelName),
+					}
+					completion, err := client.Chat.Completions.New(context.Background(), params)
+					if err != nil {
+						t.Logf("error: %v", err)
+						return false
+					}
+					// Step 2: Verify tool call
+					// TODO: remove after test done
+					returnsToolCall := false
+					for _, choice := range completion.Choices {
+						t.Logf("choice content: %s", choice.Message.Content)
+						t.Logf("finish reason: %s", choice.FinishReason)
+						t.Logf("choice toolcall: %v", choice.Message.ToolCalls)
+						if choice.FinishReason == openai.ChatCompletionChoicesFinishReasonToolCalls {
+							returnsToolCall = true
+						}
+					}
+					if returnsToolCall == false {
+						t.Logf("Tool call not returned")
+						return false
+					}
+					toolCalls := completion.Choices[0].Message.ToolCalls
+					if len(toolCalls) == 0 {
+						t.Logf("Expected tool call from completion result but got none")
+						return false
+					}
+					// Step 3: Simulate the tool returning a response, add the tool response to the params, and check the second response
+					params.Messages.Value = append(params.Messages.Value, completion.Choices[0].Message)
+					getWeatherCalled := false
+					for _, toolCall := range toolCalls {
+						if toolCall.Function.Name == "get_weather" {
+							getWeatherCalled = true
+							// Extract the location from the function call arguments
+							var args map[string]interface{}
+							if argErr := json.Unmarshal([]byte(toolCall.Function.Arguments), &args); argErr != nil {
+								panic(argErr)
+							}
+							location := args["location"].(string)
+							if location != "New York City" {
+								t.Logf("Expected location to be New York City but got %s", location)
+							}
+							// Simulate getting weather data
+							weatherData := "Sunny, 25°C"
+							params.Messages.Value = append(params.Messages.Value, openai.ToolMessage(toolCall.ID, weatherData))
+						}
+					}
+					if getWeatherCalled == false {
+						t.Logf("get_weather tool not specified in chat completion response")
+						return false
+					}
+
+					secondChatCompletion, err := client.Chat.Completions.New(context.Background(), params)
+					if err != nil {
+						t.Logf("error during second response: %v", err)
+						return false
+					}
+
+					// Step 4: Verify that the second response is correct
+					completionResult := secondChatCompletion.Choices[0].Message.Content
+					t.Logf("content of completion response using tool: %s", secondChatCompletion.Choices[0].Message.Content)
+					return completionResult == "The weather in Paris is currently sunny and 25°C."
+				}, 30*time.Second, 2*time.Second)
 			})
 			if err != nil {
 				t.Logf("error: %v", err)
@@ -224,7 +296,7 @@
 				}
 			}
 			return returnsToolCall
 		}, 30*time.Second, 2*time.Second)
 	})
 }
 
@@ -252,17 +324,17 @@
 }
 
 // maybeSkip skips the test if the required credentials are not set.
 func (c credentialsContext) maybeSkip(t *testing.T, required requiredCredential) {
 	if required&requiredCredentialOpenAI != 0 && !c.openAIValid {
 		t.Skip("skipping test as OpenAI API key is not set in TEST_OPENAI_API_KEY")
 	}
 	if required&requiredCredentialAWS != 0 && !c.awsValid {
 		t.Skip("skipping test as AWS credentials are not set in TEST_AWS_ACCESS_KEY_ID and TEST_AWS_SECRET_ACCESS_KEY")
 	}
 }
 
 // requireNewCredentialsContext creates a new credential context for the tests from the environment variables.
 func requireNewCredentialsContext(t *testing.T) (ctx credentialsContext) {
 	// Set up credential file for OpenAI.
 	openAIAPIKey := os.Getenv("TEST_OPENAI_API_KEY")
 
@@ -297,4 +369,4 @@
 		openAIAPIKeyFilePath: openAIAPIKeyFilePath,
 		awsFilePath:          awsFilePath,
 	}
 }