fixes: 2229 - completePrompt doesn't work for bedrock. (#2230)

* fixes: 2229 - completePrompt doesn't work for bedrock.

* fix failing test
This commit is contained in:
Smartsheet-JB-Brown 2025-04-02 12:35:14 -07:00 • committed by GitHub
parent 1b25190947
commit 0409509002
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 79 additions and 38 deletions

View file

@ -3,10 +3,20 @@ import { ApiHandlerOptions } from "../../../shared/api"
// Mock the AWS SDK
jest.mock("@aws-sdk/client-bedrock-runtime", () => {
const mockResponse = {
output: {
message: {
content: [
{
text: "Test response",
},
],
},
},
}
const mockSend = jest.fn().mockImplementation(() => {
return Promise.resolve({
output: new TextEncoder().encode(JSON.stringify({ content: "Test response" })),
})
return Promise.resolve(mockResponse)
})
return {

View file

@ -399,14 +399,20 @@ describe("AwsBedrockHandler", () => {
})
})
//response.output.message.content[0].text
describe("completePrompt", () => {
it("should complete prompt successfully", async () => {
const mockResponse = {
output: new TextEncoder().encode(
JSON.stringify({
content: "Test response",
}),
),
output: {
message: {
content: [
{
text: "Test response",
},
],
},
},
}
const mockSend = jest.fn().mockResolvedValue(mockResponse)
@ -450,7 +456,9 @@ describe("AwsBedrockHandler", () => {
it("should handle invalid response format", async () => {
const mockResponse = {
output: new TextEncoder().encode("invalid json"),
output: {
message: {},
},
}
const mockSend = jest.fn().mockResolvedValue(mockResponse)
@ -464,9 +472,16 @@ describe("AwsBedrockHandler", () => {
it("should handle empty response", async () => {
const mockResponse = {
output: new TextEncoder().encode(JSON.stringify({})),
output: {
message: {
content: [
{
text: "",
},
],
},
},
}
const mockSend = jest.fn().mockResolvedValue(mockResponse)
handler["client"] = {
send: mockSend,
@ -486,11 +501,15 @@ describe("AwsBedrockHandler", () => {
})
const mockResponse = {
output: new TextEncoder().encode(
JSON.stringify({
content: "Test response",
}),
),
output: {
message: {
content: [
{
text: "Test response",
},
],
},
},
}
const mockSend = jest.fn().mockResolvedValue(mockResponse)
@ -519,11 +538,15 @@ describe("AwsBedrockHandler", () => {
})
const mockResponse = {
output: new TextEncoder().encode(
JSON.stringify({
content: "Test response",
}),
),
output: {
message: {
content: [
{
text: "Test response",
},
],
},
},
}
const mockSend = jest.fn().mockResolvedValue(mockResponse)
@ -552,13 +575,16 @@ describe("AwsBedrockHandler", () => {
})
const mockResponse = {
output: new TextEncoder().encode(
JSON.stringify({
content: "Test response",
}),
),
output: {
message: {
content: [
{
text: "Test response",
},
],
},
},
}
const mockSend = jest.fn().mockResolvedValue(mockResponse)
handler["client"] = {
send: mockSend,
@ -585,11 +611,15 @@ describe("AwsBedrockHandler", () => {
})
const mockResponse = {
output: new TextEncoder().encode(
JSON.stringify({
content: "Test response",
}),
),
output: {
message: {
content: [
{
text: "Test response",
},
],
},
},
}
const mockSend = jest.fn().mockResolvedValue(mockResponse)

View file

@ -665,13 +665,14 @@ Please check:
const command = new ConverseCommand(payload)
const response = await this.client.send(command)
if (response.output && response.output instanceof Uint8Array) {
if (
response?.output?.message?.content &&
response.output.message.content.length > 0 &&
response.output.message.content[0].text &&
response.output.message.content[0].text.trim().length > 0
) {
try {
const outputStr = new TextDecoder().decode(response.output)
const output = JSON.parse(outputStr)
if (output.content) {
return output.content
}
return response.output.message.content[0].text
} catch (parseError) {
logger.error("Failed to parse Bedrock response", {
ctx: "bedrock",