kyd-shared-badge 0.3.14 → 0.3.16

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "kyd-shared-badge",
-  "version": "0.3.14",
+  "version": "0.3.16",
   "private": false,
   "main": "./src/index.ts",
   "module": "./src/index.ts",

package/src/lib/context.ts

@@ -8,46 +8,51 @@ import { S3Client, GetObjectCommand } from '@aws-sdk/client-s3';
 import { GraphInsightsPayload } from '../types';
 
 const region = process.env.AWS_REGION!;
-const usersTableName = process.env.USER_TABLE_NAME || `UserTable${(process.env.NEXT_PUBLIC_STAGE || process.env.STAGE || 'dev').replace(/^./, c => c.toUpperCase())}`;
+const usersTableName = process.env.USER_TABLE_NAME
 const doc = DynamoDBDocumentClient.from(new DynamoDBClient({ region }));
 const s3 = new S3Client({ region });
 const providerDataBucket = process.env.PROVIDER_DATA_BUCKET_NAME!;
 const badgeTableName = process.env.BADGE_TABLE_NAME!;
 
 export async function aggregateUserData(userId: string, enterpriseMode?: boolean, companyId?: string | null): Promise<any> {
-  // Minimal: load user record. Optionally, in enterprise mode, you could also fetch link overlay.
-  const userResp = await doc.send(new GetCommand({ TableName: usersTableName, Key: { userId } }));
+  let userResp;
+  if (enterpriseMode) { 
+    const linkTableName = process.env.USER_COMPANY_LINK_TABLE_NAME
+    userResp = await doc.send(new GetCommand({ TableName: linkTableName, Key: { userId, companyId } }));
+  } else {
+    userResp = await doc.send(new GetCommand({ TableName: usersTableName, Key: { userId } }));
+  }
   const user = userResp.Item || {};
-  const merged = { ...user }
+
   if (enterpriseMode && companyId) {
     // Optionally merge enterprise overlay later; keeping simple for MVP.
-    merged.companyId = companyId;
+    user.companyId = companyId;
   }
 
   // For each provider_* entry, fetch filtered_data_key JSON from S3 and attach under <provider>_analysis
-  const providerKeys = Object.keys(merged).filter(k => k.startsWith('provider_'));
+  const providerKeys = Object.keys(user).filter(k => k.startsWith('provider_'));
   const fetches = providerKeys.map(async (key) => {
     try {
-      const providerData = merged[key];
+      const providerData = user[key];
       if (!providerData || typeof providerData !== 'object') return;
       const filteredKey = providerData['filtered_data_key'];
       if (!filteredKey || typeof filteredKey !== 'string') return;
       const obj = await getJsonFromS3(providerDataBucket, filteredKey);
       const providerName = key.replace(/^provider_/, '');
       const analysisKey = `${providerName}_analysis`;
-      merged[analysisKey] = obj;
+      user[analysisKey] = obj;
     } catch (e) {
       console.error('Failed to load provider filtered data', e);
       // Non-fatal; annotate error
       try {
         const providerName = key.replace(/^provider_/, '');
-        merged[`${providerName}_analysis`] = { error: 'Failed to load provider filtered data' };
+        user[`${providerName}_analysis`] = { error: 'Failed to load provider filtered data' };
       } catch {}
     }
   });
   await Promise.all(fetches);
 
-  return merged;
+  return user;
 }
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -150,4 +155,9 @@ async function getJsonFromS3(bucket: string, key: string): Promise<any> {
 export async function getReportGraphData(badgeId: string): Promise<GraphInsightsPayload | undefined> {
   const badge = await doc.send(new GetCommand({ TableName: badgeTableName, Key: { badgeId } }));
   return badge.Item?.assessmentResult?.graph_insights as GraphInsightsPayload | undefined;
+}
+
+export async function getBadgeUserId(badgeId: string): Promise<string> {
+  const badge = await doc.send(new GetCommand({ TableName: badgeTableName, Key: { badgeId } }));
+  return badge.Item?.userId as string;
 }

package/src/lib/routes.ts

@@ -6,7 +6,13 @@ import { streamText } from 'ai';
 import { bedrock } from '@ai-sdk/amazon-bedrock';
 
 import { getHistory, putMessage } from './chat-store';
-import { aggregateUserData, cleanDeveloperProfile, buildAllContextPrompt, getReportGraphData } from './context';
+import { 
+  aggregateUserData, 
+  cleanDeveloperProfile, 
+  buildAllContextPrompt, 
+  getReportGraphData,
+  getBadgeUserId
+} from './context';
 import { checkAndConsumeToken } from './rate-limit';
 
 import { createSession } from './chat-store';
@@ -29,7 +35,8 @@ export async function chatStreamRoute(req: NextRequest, userId: string, companyI
 
     await putMessage({ sessionId, role: 'user', content });
 
-    const aggregated = await aggregateUserData(userId, !!companyId, companyId);
+    const badgeUserId = await getBadgeUserId(badgeId);
+    const aggregated = await aggregateUserData(badgeUserId, !!companyId, companyId);
     const cleaned = cleanDeveloperProfile(aggregated);
     const graphData = await getReportGraphData(badgeId);
     const system = buildAllContextPrompt(cleaned, graphData, { concise: true });