From 3d4fd569930f979f242fb6c31270084859bba73d Mon Sep 17 00:00:00 2001 From: Kyle Lierer Date: Tue, 2 Sep 2025 20:41:00 -0500 Subject: [PATCH] Changed how thread lookups are provided. --- seed/data/default/agents.json | 34 ++++++++--------------------- seed/data/default/capabilities.json | 2 +- seed/data/default/collections.json | 2 +- seed/data/default/prompts.json | 2 +- src/controllers/prediction.ts | 10 ++++++--- src/database/procedures/quota.ts | 5 ++++- src/utils/ai/system.ts | 14 ++++++++---- 7 files changed, 33 insertions(+), 36 deletions(-) diff --git a/seed/data/default/agents.json b/seed/data/default/agents.json index 66b6772..d57dc3f 100644 --- a/seed/data/default/agents.json +++ b/seed/data/default/agents.json @@ -5,9 +5,7 @@ "logicalCollection": "67200e060a342cc6ff42ddee", "description": "Generates the information given the node visit history and graph", "reasoning": null, - "capabilities": [ - "68b6542f873e1debe91487bc" - ], + "capabilities": ["68b6542f873e1debe91487bc"], "memoryEnabled": false, "subscriptionFilter": null, "outputFilter": null, @@ -21,9 +19,7 @@ "logicalCollection": "67200e060a342cc6ff42ddee", "description": "This agent takes in an existing graph and a request to modify it and returns a list of nodes to upsert, a list of nodes to delete, a list of edges to add, and a list of edges to delete. Good for making complex modifications to existing graphs.", "reasoning": null, - "capabilities": [ - "66fed2ad135939f6d3130fa6" - ], + "capabilities": ["66fed2ad135939f6d3130fa6"], "memoryEnabled": false, "subscriptionFilter": null, "outputFilter": null, @@ -88,9 +84,7 @@ "logicalCollection": "67200e060a342cc6ff42ddee", "description": "Chooses which node to go to next given the node visit history and graph", "reasoning": null, - "capabilities": [ - "66bf8c22a9430d3fa139143a" - ], + "capabilities": ["66bf8c22a9430d3fa139143a"], "memoryEnabled": false, "subscriptionFilter": null, "outputFilter": null, @@ -104,9 +98,7 @@ "logicalCollection": "67200e060a342cc6ff42ddee", "description": "Analyzes data received by the Inquiry and provides either a conversation summary or a summary for all possible answers to a particular question.", "reasoning": null, - "capabilities": [ - "669bee8d6773961cdf4604ae" - ], + "capabilities": ["669bee8d6773961cdf4604ae"], "memoryEnabled": false, "subscriptionFilter": null, "outputFilter": null, @@ -120,9 +112,7 @@ "logicalCollection": "67200e060a342cc6ff42ddee", "description": "Generates the question to ask next given the node visit history and graph", "reasoning": null, - "capabilities": [ - "66c8a9f8d28a5133dd801ab7" - ], + "capabilities": ["66c8a9f8d28a5133dd801ab7"], "memoryEnabled": false, "subscriptionFilter": null, "outputFilter": null, @@ -136,9 +126,7 @@ "logicalCollection": "67200e060a342cc6ff42ddee", "description": "For a given graph, creates synthetic node visit data", "reasoning": null, - "capabilities": [ - "66c0040e525628d5c5861ed8" - ], + "capabilities": ["66c0040e525628d5c5861ed8"], "memoryEnabled": false, "subscriptionFilter": null, "outputFilter": null, @@ -157,9 +145,7 @@ "variablePassThrough": true, "_id": "6858291e0b814e997e60176d" }, - "capabilities": [ - "66c25d95cbd9c2075c1e5f4e" - ], + "capabilities": ["66c25d95cbd9c2075c1e5f4e"], "memoryEnabled": false, "subscriptionFilter": null, "outputFilter": null, @@ -173,9 +159,7 @@ "logicalCollection": "67200e060a342cc6ff42ddee", "description": "Generates the information given the node visit history and graph", "reasoning": null, - "capabilities": [ - "66c95576ae782803372bd9f8" - ], + "capabilities": ["66c95576ae782803372bd9f8"], "memoryEnabled": false, "subscriptionFilter": null, "outputFilter": null, @@ -183,4 +167,4 @@ "updatedAt": "2025-08-01T05:57:15.215Z", "__v": 0 } -] \ No newline at end of file +] diff --git a/seed/data/default/capabilities.json b/seed/data/default/capabilities.json index b750eba..b8f0bbf 100644 --- a/seed/data/default/capabilities.json +++ b/seed/data/default/capabilities.json @@ -499,4 +499,4 @@ "updatedAt": "2025-09-02T02:19:27.488Z", "__v": 0 } -] \ No newline at end of file +] diff --git a/seed/data/default/collections.json b/seed/data/default/collections.json index ab69b0d..906aec8 100644 --- a/seed/data/default/collections.json +++ b/seed/data/default/collections.json @@ -13,4 +13,4 @@ "updatedAt": "2024-10-28T22:19:50.860Z", "__v": 0 } -] \ No newline at end of file +] diff --git a/seed/data/default/prompts.json b/seed/data/default/prompts.json index f3fbe4c..4ed64fd 100644 --- a/seed/data/default/prompts.json +++ b/seed/data/default/prompts.json @@ -368,4 +368,4 @@ "updatedAt": "2025-08-23T22:52:30.746Z", "__v": 0 } -] \ No newline at end of file +] diff --git a/src/controllers/prediction.ts b/src/controllers/prediction.ts index 891fbd3..44c403b 100644 --- a/src/controllers/prediction.ts +++ b/src/controllers/prediction.ts @@ -344,7 +344,7 @@ export async function generatePrediction({ auth, subscriptionId, agentId, - inquiryId, + threadLookups, variables, attachments, integrationId, @@ -353,7 +353,7 @@ export async function generatePrediction({ auth?: { sub?: string }; subscriptionId: string; agentId: string; - inquiryId?: string; + threadLookups?: Map; variables: Record; attachments: Array; integrationId?: string; @@ -372,7 +372,7 @@ export async function generatePrediction({ const agent = await getAgent(agentId); if (!agent) throw new Error(`No agent found for ID: ${agentId}`); - const thread = await findOrCreateThread(subscriptionId, inquiryId); + const thread = await findOrCreateThread(subscriptionId, threadLookups); await addToThread(thread, auth?.sub, variables, true); log.debug({ @@ -942,5 +942,9 @@ export async function generatePredictionWithInquiry({ attachments, integrationId, correlationId, + threadLookups: new Map([ + ['inquiryId', inquiryId], + ['userId', auth?.sub ?? ''], + ]), }); } diff --git a/src/database/procedures/quota.ts b/src/database/procedures/quota.ts index 8b7b36e..b08b6cb 100644 --- a/src/database/procedures/quota.ts +++ b/src/database/procedures/quota.ts @@ -65,7 +65,10 @@ export async function calculateUserTokenUsage(userId: string): Promise<{ return tokenTotals; } catch (error) { - log.error({ error, userId }, `Error calculating token usage for user ${userId}`); + log.error( + { error, userId }, + `Error calculating token usage for user ${userId}`, + ); throw new Error( `Failed to calculate token usage for user ${userId}: ${error instanceof Error ? error.message : 'Unknown error'}`, ); diff --git a/src/utils/ai/system.ts b/src/utils/ai/system.ts index c408051..b2cfed9 100644 --- a/src/utils/ai/system.ts +++ b/src/utils/ai/system.ts @@ -50,15 +50,21 @@ export async function getCapability( return await Capability.findOne({ alias }).populate('prompts'); } +/** + * Finds or creates a thread for a subscription. + * @param subscriptionId {string} - The ID of the subscription. + * @param threadLookups {Map} - A map of values that can be used to look up the thread (e.g., userId, inquiryId). + * @returns {Promise} - The found or created thread. + */ export async function findOrCreateThread( subscriptionId: string, - inquiryId?: string, + threadLookups?: Map, ) { - const updateData: { $setOnInsert: { messages: []; inquiryId?: string } } = { + const updateData = { $setOnInsert: { messages: [] }, }; - if (inquiryId) { - updateData.$setOnInsert.inquiryId = inquiryId; + for (const [key, value] of threadLookups?.entries() || []) { + updateData.$setOnInsert[key] = value; } return await Thread.findOneAndUpdate({ subscriptionId }, updateData, {