Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mrubens merged 2 commits into from
Apr 2, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3bf3f38
fixes: 2229 - completePrompt doesn't work for bedrock.
Smartsheet-JB-Brown Apr 2, 2025
9e6377d
fix failing test
Smartsheet-JB-Brown Apr 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions src/api/providers/__tests__/bedrock-custom-arn.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
88 changes: 59 additions & 29 deletions src/api/providers/__tests__/bedrock.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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,
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand Down
13 changes: 7 additions & 6 deletions src/api/providers/bedrock.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Loading