front: CoreAPISearchFilter initial plumbing

dust-tt · Jul 25, 2024 · 69749b2 · 69749b2
1 parent 6d2ab76
commit 69749b2
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 27 deletions.
diff --git a/front/pages/api/registry/[type]/lookup.ts b/front/pages/api/registry/[type]/lookup.ts
@@ -1,3 +1,4 @@
+import type { CoreAPISearchFilter } from "@dust-tt/types";
 import type { NextApiRequest, NextApiResponse } from "next";
 
 import { DataSource } from "@app/lib/models/data_source";
@@ -9,6 +10,7 @@ const { DUST_REGISTRY_SECRET } = process.env;
 type LookupDataSourceResponseBody = {
   project_id: number;
   data_source_id: string;
+  view_filter: CoreAPISearchFilter | null;
 };
 
 /**
@@ -56,6 +58,8 @@ async function handler(
     return;
   }
 
+  // TODO(GROUPS_INFRA): Add x-dust-group-ids header retrieval + checks
+
   const dustWorkspaceId = req.headers["x-dust-workspace-id"] as string;
 
   switch (req.method) {
@@ -98,17 +102,25 @@ async function handler(
             return;
           }
 
+          // TODO(GROUPS_INFRA):
+          // - Implement view_filter return when a data source view is looked up.
+          // - If data_source_ids is of the form `dsv_...` then it's a data source view
+          //   and we pull the view_filter to return it below
+          // - otherwise it's data source and the view_filter is null
+          // - Obviously this is where we check based on the x-dust-group-ids header that we
+          //   have access to the data source or data source view
+
           res.status(200).json({
             project_id: parseInt(dataSource.dustAPIProjectId),
             data_source_id: req.query.data_source_id,
+            view_filter: null,
           });
           return;
 
         default:
           res.status(405).end();
           return;
       }
-      return;
 
     default:
       res.status(405).end();

diff --git a/front/pages/api/v1/w/[wId]/data_sources/[name]/search.ts b/front/pages/api/v1/w/[wId]/data_sources/[name]/search.ts
@@ -246,16 +246,16 @@ async function handler(
           target_document_tokens: query.target_document_tokens,
           filter: {
             tags: {
-              in: query.tags_in,
-              not: query.tags_not,
+              in: query.tags_in ?? null,
+              not: query.tags_not ?? null,
             },
             parents: {
-              in: query.parents_in,
-              not: query.parents_not,
+              in: query.parents_in ?? null,
+              not: query.parents_not ?? null,
             },
             timestamp: {
-              gt: query.timestamp_gt,
-              lt: query.timestamp_lt,
+              gt: query.timestamp_gt ?? null,
+              lt: query.timestamp_lt ?? null,
             },
           },
           credentials,

diff --git a/front/pages/api/w/[wId]/data_sources/[name]/search.ts b/front/pages/api/w/[wId]/data_sources/[name]/search.ts
@@ -137,16 +137,16 @@ export async function handleSearchDataSource({
       target_document_tokens: searchQuery.target_document_tokens,
       filter: {
         tags: {
-          in: searchQuery.tags_in,
-          not: searchQuery.tags_not,
+          in: searchQuery.tags_in ?? null,
+          not: searchQuery.tags_not ?? null,
         },
         parents: {
-          in: searchQuery.parents_in,
-          not: searchQuery.parents_not,
+          in: searchQuery.parents_in ?? null,
+          not: searchQuery.parents_not ?? null,
         },
         timestamp: {
-          gt: searchQuery.timestamp_gt,
-          lt: searchQuery.timestamp_lt,
+          gt: searchQuery.timestamp_gt ?? null,
+          lt: searchQuery.timestamp_lt ?? null,
         },
       },
       credentials: credentials,

diff --git a/types/src/front/lib/core_api.ts b/types/src/front/lib/core_api.ts
@@ -140,6 +140,21 @@ export type CoreAPIQueryResult = {
   value: Record<string, unknown>;
 };
 
+export type CoreAPISearchFilter = {
+  tags: {
+    in: string[] | null;
+    not: string[] | null;
+  } | null;
+  parents: {
+    in: string[] | null;
+    not: string[] | null;
+  } | null;
+  timestamp: {
+    gt: number | null;
+    lt: number | null;
+  } | null;
+};
+
 export class CoreAPI {
   _url: string;
   declare _logger: LoggerInterface;
@@ -277,6 +292,9 @@ export class CoreAPI {
     credentials,
     secrets,
   }: CoreAPICreateRunParams): Promise<CoreAPIResponse<{ run: CoreAPIRun }>> {
+    // TODO(GROUPS_INFRA): use the auth as argument of that method instead of `runAsWorkspaceId`
+    // and pass both X-Dust-Workspace-Id and X-Dust-Group-Ids.
+
     const response = await this._fetchWithError(
       `${this._url}/projects/${encodeURIComponent(projectId)}/runs`,
       {
@@ -318,6 +336,9 @@ export class CoreAPI {
       dustRunId: Promise<string>;
     }>
   > {
+    // TODO(GROUPS_INFRA): use the auth as argument of that method instead of `runAsWorkspaceId`
+    // and pass both X-Dust-Workspace-Id and X-Dust-Group-Ids.
+
     const res = await this._fetchWithError(
       `${this._url}/projects/${projectId}/runs/stream`,
       {
@@ -614,20 +635,7 @@ export class CoreAPI {
     payload: {
       query: string;
       topK: number;
-      filter?: {
-        tags: {
-          in?: string[] | null;
-          not?: string[] | null;
-        };
-        parents?: {
-          in?: string[] | null;
-          not?: string[] | null;
-        };
-        timestamp?: {
-          gt?: number | null;
-          lt?: number | null;
-        };
-      } | null;
+      filter?: CoreAPISearchFilter | null;
       fullText: boolean;
       credentials: { [key: string]: string };
       target_document_tokens?: number | null;