Files
RAG-Lambdas/query.js

53 lines
1.9 KiB
JavaScript

import { BedrockRuntimeClient, InvokeModelCommand } from "@aws-sdk/client-bedrock-runtime";
import { Pinecone } from "@pinecone-database/pinecone";
const bedrock = new BedrockRuntimeClient({});
const pc = new Pinecone({ apiKey: process.env.PINECONE_API_KEY });
export const handler = async (event) => {
const { question } = JSON.parse(event.body);
// 1. Embed user query
const embedResponse = await bedrock.send(new InvokeModelCommand({
modelId: "amazon.titan-embed-text-v1",
contentType: "application/json",
accept: "application/json",
body: JSON.stringify({ inputText: question })
}));
const { embedding } = JSON.parse(new TextDecoder().decode(embedResponse.body));
// 2. Query Vector DB for relevant context
const index = pc.index(process.env.PINECONE_INDEX_NAME);
const queryResponse = await index.query({
vector: embedding,
topK: 3,
includeMetadata: true
});
const context = queryResponse.matches.map(match => match.metadata.text).join("\n\n");
const systemPrompt = `Use the following context to answer the question. If you don't know, say you don't know.\n\nContext:\n${context}`;
const llmResponse = await bedrock.send(new InvokeModelCommand({
modelId: "anthropic.claude-3-haiku-20240307-v1:0", // Fast & Cost-effective
contentType: "application/json",
accept: "application/json",
body: JSON.stringify({
anthropic_version: "bedrock-2023-05-31",
max_tokens: 500,
messages: [
{ role: "user", content: `${systemPrompt}\n\nQuestion: ${question}` }
]
})
}));
const result = JSON.parse(new TextDecoder().decode(llmResponse.body));
const answer = result.content[0].text;
return {
statusCode: 200,
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ answer, contextUsed: queryResponse.matches.length })
};
};