Modified function signatures (#9652)

Co-authored-by: Lucas <lucas@dust.tt>
dust-tt · Dec 30, 2024 · 36c648d · 36c648d
1 parent 247e19e
commit 36c648d
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 69 deletions.
diff --git a/front/lib/resources/agent_message_feedback_resource.ts b/front/lib/resources/agent_message_feedback_resource.ts
@@ -1,4 +1,9 @@
-import type { AgentConfigurationType, Result, UserType } from "@dust-tt/types";
+import type {
+  AgentConfigurationType,
+  Result,
+  UserType,
+  WorkspaceType,
+} from "@dust-tt/types";
 import { Err, Ok } from "@dust-tt/types";
 import type { Attributes, ModelStatic } from "sequelize";
 import type { CreationAttributes, Transaction } from "sequelize";
@@ -8,7 +13,6 @@ import type { AgentMessageFeedbackDirection } from "@app/lib/api/assistant/conve
 import type { Authenticator } from "@app/lib/auth";
 import type { AgentMessage } from "@app/lib/models/assistant/conversation";
 import { AgentMessageFeedback } from "@app/lib/models/assistant/conversation";
-import type { Workspace } from "@app/lib/models/workspace";
 import { BaseResource } from "@app/lib/resources/base_resource";
 import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types";
 import { UserResource } from "@app/lib/resources/user_resource";
@@ -107,7 +111,7 @@ export class AgentMessageFeedbackResource extends BaseResource<AgentMessageFeedb
     startDate,
     endDate,
   }: {
-    workspace: Workspace;
+    workspace: WorkspaceType;
     startDate: Date;
     endDate: Date;
   }): Promise<AgentMessageFeedbackResource[]> {

diff --git a/front/lib/workspace_usage.ts b/front/lib/workspace_usage.ts
@@ -1,4 +1,4 @@
-import type { ModelId } from "@dust-tt/types";
+import type { ModelId, WorkspaceType } from "@dust-tt/types";
 import { stringify } from "csv-stringify/sync";
 import { format } from "date-fns/format";
 import { Op, QueryTypes, Sequelize } from "sequelize";
@@ -10,7 +10,6 @@ import {
   Message,
   UserMessage,
 } from "@app/lib/models/assistant/conversation";
-import { Workspace } from "@app/lib/models/workspace";
 import { AgentMessageFeedbackResource } from "@app/lib/resources/agent_message_feedback_resource";
 import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { frontSequelize } from "@app/lib/resources/storage";
@@ -87,8 +86,9 @@ interface FeedbackQueryResult {
 export async function unsafeGetUsageData(
   startDate: Date,
   endDate: Date,
-  wId: string
+  workspace: WorkspaceType
 ): Promise<string> {
+  const wId = workspace.sId;
   const results = await frontSequelize.query<WorkspaceUsageQueryResult>(
     `
       SELECT
@@ -163,14 +163,8 @@ export async function unsafeGetUsageData(
 export async function getMessageUsageData(
   startDate: Date,
   endDate: Date,
-  workspaceId: string
+  workspace: WorkspaceType
 ): Promise<string> {
-  const workspace = await Workspace.findOne({
-    where: { sId: workspaceId },
-  });
-  if (!workspace) {
-    throw new Error(`Workspace not found for sId: ${workspaceId}`);
-  }
   const wId = workspace.id;
   const results = await frontSequelize.query<MessageUsageQueryResult>(
     `
@@ -229,14 +223,8 @@ export async function getMessageUsageData(
 export async function getUserUsageData(
   startDate: Date,
   endDate: Date,
-  workspaceId: string
+  workspace: WorkspaceType
 ): Promise<string> {
-  const workspace = await Workspace.findOne({
-    where: { sId: workspaceId },
-  });
-  if (!workspace) {
-    throw new Error(`Workspace not found for sId: ${workspaceId}`);
-  }
   const wId = workspace.id;
   const userMessages = await Message.findAll({
     attributes: [
@@ -320,14 +308,8 @@ export async function getUserUsageData(
 export async function getBuildersUsageData(
   startDate: Date,
   endDate: Date,
-  workspaceId: string
+  workspace: WorkspaceType
 ): Promise<string> {
-  const workspace = await Workspace.findOne({
-    where: { sId: workspaceId },
-  });
-  if (!workspace) {
-    throw new Error(`Workspace not found for sId: ${workspaceId}`);
-  }
   const wId = workspace.id;
   const agentConfigurations = await AgentConfiguration.findAll({
     attributes: [
@@ -401,14 +383,8 @@ export async function getBuildersUsageData(
 export async function getAssistantsUsageData(
   startDate: Date,
   endDate: Date,
-  workspaceId: string
+  workspace: WorkspaceType
 ): Promise<string> {
-  const workspace = await Workspace.findOne({
-    where: { sId: workspaceId },
-  });
-  if (!workspace) {
-    throw new Error(`Workspace not found for sId: ${workspaceId}`);
-  }
   const wId = workspace.id;
   const mentions = await frontSequelize.query<AgentUsageQueryResult>(
     `
@@ -462,15 +438,8 @@ export async function getAssistantsUsageData(
 export async function getFeedbacksUsageData(
   startDate: Date,
   endDate: Date,
-  workspaceId: string
+  workspace: WorkspaceType
 ): Promise<string> {
-  const workspace = await Workspace.findOne({
-    where: { sId: workspaceId },
-  });
-  if (!workspace) {
-    throw new Error(`Workspace not found for sId: ${workspaceId}`);
-  }
-
   const feedbacks =
     await AgentMessageFeedbackResource.listByWorkspaceAndDateRange({
       workspace: workspace,

diff --git a/front/pages/api/v1/w/[wId]/usage.ts b/front/pages/api/v1/w/[wId]/usage.ts
@@ -65,7 +65,7 @@ async function handler(
       const csvData = await unsafeGetUsageData(
         new Date(query.start_date),
         query.end_date ? new Date(query.end_date) : new Date(),
-        owner.sId
+        owner
       );
       res.setHeader("Content-Type", "text/csv");
       res.status(200).send(csvData);

diff --git a/front/pages/api/v1/w/[wId]/workspace-usage.ts b/front/pages/api/v1/w/[wId]/workspace-usage.ts
@@ -3,6 +3,7 @@ import type {
   UsageTableType,
 } from "@dust-tt/client";
 import { GetWorkspaceUsageRequestSchema } from "@dust-tt/client";
+import type { WorkspaceType } from "@dust-tt/types";
 import { assertNever } from "@dust-tt/types";
 import { endOfMonth } from "date-fns/endOfMonth";
 import JSZip from "jszip";
@@ -128,7 +129,7 @@ async function handler(
         table: query.table,
         start: startDate,
         end: endDate,
-        workspaceId: owner.sId,
+        workspace: owner,
       });
       const zip = new JSZip();
       const csvSuffix = startDate
@@ -200,38 +201,38 @@ async function fetchUsageData({
   table,
   start,
   end,
-  workspaceId,
+  workspace,
 }: {
   table: UsageTableType;
   start: Date;
   end: Date;
-  workspaceId: string;
+  workspace: WorkspaceType;
 }): Promise<Partial<Record<UsageTableType, string>>> {
   switch (table) {
     case "users":
-      return { users: await getUserUsageData(start, end, workspaceId) };
+      return { users: await getUserUsageData(start, end, workspace) };
     case "assistant_messages":
       return {
-        assistant_messages: await getMessageUsageData(start, end, workspaceId),
+        assistant_messages: await getMessageUsageData(start, end, workspace),
       };
     case "builders":
-      return { builders: await getBuildersUsageData(start, end, workspaceId) };
+      return { builders: await getBuildersUsageData(start, end, workspace) };
     case "assistants":
       return {
-        assistants: await getAssistantsUsageData(start, end, workspaceId),
+        assistants: await getAssistantsUsageData(start, end, workspace),
       };
     case "feedbacks":
       return {
-        feedbacks: await getFeedbacksUsageData(start, end, workspaceId),
+        feedbacks: await getFeedbacksUsageData(start, end, workspace),
       };
     case "all":
       const [users, assistant_messages, builders, assistants, feedbacks] =
         await Promise.all([
-          getUserUsageData(start, end, workspaceId),
-          getMessageUsageData(start, end, workspaceId),
-          getBuildersUsageData(start, end, workspaceId),
-          getAssistantsUsageData(start, end, workspaceId),
-          getFeedbacksUsageData(start, end, workspaceId),
+          getUserUsageData(start, end, workspace),
+          getMessageUsageData(start, end, workspace),
+          getBuildersUsageData(start, end, workspace),
+          getAssistantsUsageData(start, end, workspace),
+          getFeedbacksUsageData(start, end, workspace),
         ]);
       return { users, assistant_messages, builders, assistants, feedbacks };
     default:

diff --git a/front/pages/api/w/[wId]/workspace-usage.ts b/front/pages/api/w/[wId]/workspace-usage.ts
@@ -1,3 +1,4 @@
+import type { WorkspaceType } from "@dust-tt/types";
 import { assertNever } from "@dust-tt/types";
 import { endOfMonth } from "date-fns/endOfMonth";
 import { isLeft } from "fp-ts/lib/Either";
@@ -102,7 +103,7 @@ async function handler(
         table: query.table,
         start: startDate,
         end: endDate,
-        workspaceId: owner.sId,
+        workspace: owner,
       });
       if (query.table === "all") {
         const zip = new JSZip();
@@ -173,36 +174,36 @@ async function fetchUsageData({
   table,
   start,
   end,
-  workspaceId,
+  workspace,
 }: {
   table: usageTableType;
   start: Date;
   end: Date;
-  workspaceId: string;
+  workspace: WorkspaceType;
 }): Promise<Partial<Record<usageTableType, string>>> {
   switch (table) {
     case "users":
-      return { users: await getUserUsageData(start, end, workspaceId) };
+      return { users: await getUserUsageData(start, end, workspace) };
     case "assistant_messages":
-      return { mentions: await getMessageUsageData(start, end, workspaceId) };
+      return { mentions: await getMessageUsageData(start, end, workspace) };
     case "builders":
-      return { builders: await getBuildersUsageData(start, end, workspaceId) };
+      return { builders: await getBuildersUsageData(start, end, workspace) };
     case "feedbacks":
       return {
-        feedbacks: await getFeedbacksUsageData(start, end, workspaceId),
+        feedbacks: await getFeedbacksUsageData(start, end, workspace),
       };
     case "assistants":
       return {
-        assistants: await getAssistantsUsageData(start, end, workspaceId),
+        assistants: await getAssistantsUsageData(start, end, workspace),
       };
     case "all":
       const [users, assistant_messages, builders, assistants, feedbacks] =
         await Promise.all([
-          getUserUsageData(start, end, workspaceId),
-          getMessageUsageData(start, end, workspaceId),
-          getBuildersUsageData(start, end, workspaceId),
-          getAssistantsUsageData(start, end, workspaceId),
-          getFeedbacksUsageData(start, end, workspaceId),
+          getUserUsageData(start, end, workspace),
+          getMessageUsageData(start, end, workspace),
+          getBuildersUsageData(start, end, workspace),
+          getAssistantsUsageData(start, end, workspace),
+          getFeedbacksUsageData(start, end, workspace),
         ]);
       return { users, assistant_messages, builders, assistants, feedbacks };
     default: