diff --git a/front/lib/api/assistant/actions/retrieval.ts b/front/lib/api/assistant/actions/retrieval.ts index e7a77e7ecc45..04f709681398 100644 --- a/front/lib/api/assistant/actions/retrieval.ts +++ b/front/lib/api/assistant/actions/retrieval.ts @@ -387,56 +387,28 @@ export async function* runRetrieval( ); // Handle data sources list and parents/tags filtering. - if (c.dataSources === "all") { - const prodCredentials = await prodAPICredentialsForOwner(owner); - const api = new DustAPI(prodCredentials); + config.DATASOURCE.data_sources = c.dataSources.map((d) => ({ + workspace_id: d.workspaceId, + data_source_id: d.dataSourceId, + })); + + for (const ds of c.dataSources) { + if (ds.filter.tags) { + if (!config.DATASOURCE.filter.tags) { + config.DATASOURCE.filter.tags = { in: [], not: [] }; + } - const dsRes = await api.getDataSources(prodCredentials.workspaceId); - if (dsRes.isErr()) { - return yield { - type: "retrieval_error", - created: Date.now(), - configurationId: configuration.sId, - messageId: agentMessage.sId, - error: { - code: "retrieval_data_sources_error", - message: `Error retrieving workspace data sources: ${dsRes.error.message}`, - }, - }; + config.DATASOURCE.filter.tags.in.push(...ds.filter.tags.in); + config.DATASOURCE.filter.tags.not.push(...ds.filter.tags.not); } - const ds = dsRes.value.filter((d) => d.assistantDefaultSelected); - - config.DATASOURCE.data_sources = ds.map((d) => { - return { - workspace_id: prodCredentials.workspaceId, - data_source_id: d.name, - }; - }); - } else { - config.DATASOURCE.data_sources = c.dataSources.map((d) => ({ - workspace_id: d.workspaceId, - data_source_id: d.dataSourceId, - })); - - for (const ds of c.dataSources) { - if (ds.filter.tags) { - if (!config.DATASOURCE.filter.tags) { - config.DATASOURCE.filter.tags = { in: [], not: [] }; - } - - config.DATASOURCE.filter.tags.in.push(...ds.filter.tags.in); - config.DATASOURCE.filter.tags.not.push(...ds.filter.tags.not); + if (ds.filter.parents) { + if (!config.DATASOURCE.filter.parents) { + config.DATASOURCE.filter.parents = { in: [], not: [] }; } - if (ds.filter.parents) { - if (!config.DATASOURCE.filter.parents) { - config.DATASOURCE.filter.parents = { in: [], not: [] }; - } - - config.DATASOURCE.filter.parents.in.push(...ds.filter.parents.in); - config.DATASOURCE.filter.parents.not.push(...ds.filter.parents.not); - } + config.DATASOURCE.filter.parents.in.push(...ds.filter.parents.in); + config.DATASOURCE.filter.parents.not.push(...ds.filter.parents.not); } } diff --git a/front/types/assistant/actions/retrieval.ts b/front/types/assistant/actions/retrieval.ts index 2544edeb7ee7..90c9a7ae498a 100644 --- a/front/types/assistant/actions/retrieval.ts +++ b/front/types/assistant/actions/retrieval.ts @@ -43,7 +43,7 @@ export type RetrievalConfigurationType = { id: ModelId; type: "retrieval_configuration"; - dataSources: "all" | DataSourceConfiguration[]; + dataSources: DataSourceConfiguration[]; query: "auto" | "none" | TemplatedQuery; relativeTimeFrame: "auto" | "none" | TimeFrame; topK: number;