Add shared library for APIs/chunking

query.js

@@ -1,52 +1,89 @@
-import { BedrockRuntimeClient, InvokeModelCommand } from "@aws-sdk/client-bedrock-runtime";
 import { Pinecone } from "@pinecone-database/pinecone";
+import { embedTexts, rerank } from "./lib/voyage.js";
+import { generate } from "./lib/anthropic.js";
+import { computeGroundedness } from "./lib/scoring.js";
 
-const bedrock = new BedrockRuntimeClient({});
 const pc = new Pinecone({ apiKey: process.env.PINECONE_API_KEY });
+const TOP_K_SEARCH = 20;
+const TOP_K_RERANK = 5;
 
 export const handler = async (event) => {
-    const { question } = JSON.parse(event.body);
+    const { question } = JSON.parse(event.body || "{}");
+    if (!question || typeof question !== "string") {
+        return {
+            statusCode: 400,
+            headers: { "Content-Type": "application/json" },
+            body: JSON.stringify({ error: "question is required" }),
+        };
+    }
 
-    // 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 [queryEmbedding] = await embedTexts([question], "query");
     const index = pc.index(process.env.PINECONE_INDEX_NAME);
     const queryResponse = await index.query({
-        vector: embedding,
-        topK: 3,
-        includeMetadata: true
+        vector: queryEmbedding,
+        topK: TOP_K_SEARCH,
+        includeMetadata: true,
     });
 
-    const context = queryResponse.matches.map(match => match.metadata.text).join("\n\n");
+    const candidates = (queryResponse.matches || [])
+        .filter((m) => m.metadata && m.metadata.text)
+        .map((m) => ({
+            id: m.id,
+            text: m.metadata.text,
+            source: m.metadata.source || m.id,
+            chunkIndex: m.metadata.chunkIndex,
+        }));
 
-    const systemPrompt = `Use the following context to answer the question. If you don't know, say you don't know.\n\nContext:\n${context}`;
+    if (candidates.length === 0) {
+        return {
+            statusCode: 200,
+            headers: { "Content-Type": "application/json" },
+            body: JSON.stringify({
+                answer: "No sources found. Ingest documents first.",
+                citations: [],
+                groundednessScore: null,
+                followUpQuestions: [],
+                contextUsed: 0,
+            }),
+        };
+    }
 
-    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 ranked = await rerank(question, candidates);
+    const topChunks = ranked.slice(0, TOP_K_RERANK).map((r) => r.document);
+
+    const sources = topChunks.map((chunk, i) => ({
+        docIndex: i + 1,
+        name: chunk.source,
+        text: chunk.text,
+        id: chunk.id,
     }));
 
-    const result = JSON.parse(new TextDecoder().decode(llmResponse.body));
-    const answer = result.content[0].text;
+    const { answer, citedSourceIndices, followUpQuestions } = await generate(question, sources);
+
+    const citedChunks = citedSourceIndices
+        .map((idx) => sources.find((s) => s.docIndex === idx))
+        .filter(Boolean);
+
+    const groundednessScore = await computeGroundedness(
+        answer,
+        citedChunks.map((c) => c.text)
+    );
+
+    const citations = citedChunks.map((c) => ({
+        sourceIndex: c.docIndex,
+        source: c.name,
+        chunkText: c.text,
+    }));
 
     return {
         statusCode: 200,
         headers: { "Content-Type": "application/json" },
-        body: JSON.stringify({ answer, contextUsed: queryResponse.matches.length })
+        body: JSON.stringify({
+            answer,
+            citations,
+            groundednessScore,
+            followUpQuestions,
+            contextUsed: topChunks.length,
+        }),
     };
 };