From 295594099dfc1e6efa6669ba63ee301e09ca3359 Mon Sep 17 00:00:00 2001 From: PopDaph Date: Thu, 7 Sep 2023 09:27:48 +0200 Subject: [PATCH 01/12] Assistant: Lib function createAgentConfiguration --- front/lib/api/assistant/actions/retrieval.ts | 69 +++------ front/lib/api/assistant/agent.ts | 122 ++++++++++++++-- .../lib/models/assistant/actions/retrieval.ts | 25 ++-- front/types/assistant/actions/retrieval.ts | 24 +++- front/types/assistant/agent-utils.ts | 134 ++++++++++++++++++ 5 files changed, 292 insertions(+), 82 deletions(-) create mode 100644 front/types/assistant/agent-utils.ts diff --git a/front/lib/api/assistant/actions/retrieval.ts b/front/lib/api/assistant/actions/retrieval.ts index e7a77e7ecc45..b151f9b95af1 100644 --- a/front/lib/api/assistant/actions/retrieval.ts +++ b/front/lib/api/assistant/actions/retrieval.ts @@ -387,66 +387,31 @@ 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.name, + })); + + 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); } } - // Handle timestamp filtering. - if (params.relativeTimeFrame) { - config.DATASOURCE.filter.timestamp = { - gt: timeFrameFromNow(params.relativeTimeFrame), - }; - } - // Handle top k. config.DATASOURCE.top_k = params.topK; diff --git a/front/lib/api/assistant/agent.ts b/front/lib/api/assistant/agent.ts index 75dce5fe133c..0dd606ea7584 100644 --- a/front/lib/api/assistant/agent.ts +++ b/front/lib/api/assistant/agent.ts @@ -14,9 +14,24 @@ import { runGeneration, } from "@app/lib/api/assistant/generation"; import { Authenticator } from "@app/lib/auth"; +import { front_sequelize } from "@app/lib/databases"; +import { DataSource } from "@app/lib/models"; +import { + AgentDataSourceConfiguration, + AgentRetrievalConfiguration, +} from "@app/lib/models/assistant/actions/retrieval"; +import { + AgentConfiguration, + AgentGenerationConfiguration, +} from "@app/lib/models/assistant/agent"; import { Err, Ok, Result } from "@app/lib/result"; import { generateModelSId } from "@app/lib/utils"; import { isRetrievalConfiguration } from "@app/types/assistant/actions/retrieval"; +import { + isTemplatedQuery, + isTimeFrame, + RetrievalDataSourcesConfiguration, +} from "@app/types/assistant/actions/retrieval"; import { AgentActionConfigurationType, AgentActionSpecification, @@ -24,6 +39,7 @@ import { AgentConfigurationType, AgentGenerationConfigurationType, } from "@app/types/assistant/agent"; +import { _getAgentConfigurationType } from "@app/types/assistant/agent-utils"; import { AgentActionType, AgentMessageType, @@ -48,15 +64,103 @@ export async function createAgentConfiguration( action?: AgentActionConfigurationType; generation?: AgentGenerationConfigurationType; } -): Promise { - return { - sId: generateModelSId(), - name, - pictureUrl: pictureUrl ?? null, - status: "active", - action: action ?? null, - generation: generation ?? null, - }; +): Promise { + const owner = auth.workspace(); + + if (!owner) { + return; + } + + return await front_sequelize.transaction(async (t) => { + // Create AgentConfiguration + const agentConfigRow = await AgentConfiguration.create( + { + sId: generateModelSId(), + status: "active", + name: name, + pictureUrl: pictureUrl ?? null, + scope: "workspace", + workspaceId: owner.id, + }, + { transaction: t } + ); + + let agentGenerationConfigRow: AgentGenerationConfiguration | null = null; + let agentActionConfigRow: AgentRetrievalConfiguration | null = null; + const agentDataSourcesConfigRows: AgentDataSourceConfiguration[] = []; + + // Create AgentGenerationConfiguration + if (generation) { + agentGenerationConfigRow = await AgentGenerationConfiguration.create( + { + prompt: generation.prompt, + modelProvider: generation.model.providerId, + modelId: generation.model.modelId, + agentId: agentConfigRow.id, + }, + { transaction: t } + ); + } + + // Create AgentRetrievalConfiguration & associated AgentDataSourceConfiguration + if (action) { + const query = action.query; + const timeframe = action.relativeTimeFrame; + + agentActionConfigRow = await AgentRetrievalConfiguration.create( + { + query: isTemplatedQuery(query) ? "templated" : query, + queryTemplate: isTemplatedQuery(query) ? query.template : null, + relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, + relativeTimeFrameDuration: isTimeFrame(timeframe) + ? timeframe.duration + : null, + relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, + topK: action.topK, + agentId: agentConfigRow.id, + }, + { transaction: t } + ); + + if (!agentActionConfigRow) { + return; + } + + if (action.dataSources) { + let dsRow: AgentDataSourceConfiguration | null = null; + action.dataSources.map(async (d) => { + const ds = await DataSource.findOne({ + where: { + name: d.name, + workspaceId: d.workspaceId, + }, + }); + + if (ds && agentActionConfigRow) { + dsRow = await AgentDataSourceConfiguration.create( + { + dataSourceId: ds.id, + tagsIn: d.filter.tags?.in, + tagsNotIn: d.filter.tags?.not, + parentsIn: d.filter.parents?.in, + parentsNotIn: d.filter.parents?.not, + retrievalConfigurationId: agentActionConfigRow.id, + }, + { transaction: t } + ); + agentDataSourcesConfigRows.push(dsRow); + } + }); + } + } + + return _getAgentConfigurationType({ + agent: agentConfigRow, + action: agentActionConfigRow, + generation: agentGenerationConfigRow, + dataSources: agentDataSourcesConfigRows, + }); + }); } export async function updateAgentConfiguration( diff --git a/front/lib/models/assistant/actions/retrieval.ts b/front/lib/models/assistant/actions/retrieval.ts index f12372810306..98dcaac79346 100644 --- a/front/lib/models/assistant/actions/retrieval.ts +++ b/front/lib/models/assistant/actions/retrieval.ts @@ -117,9 +117,6 @@ export class AgentDataSourceConfiguration extends Model< declare createdAt: CreationOptional; declare updatedAt: CreationOptional; - declare timeframeDuration: number | null; - declare timeframeUnit: TimeframeUnit | null; - declare tagsIn: string[] | null; declare tagsNotIn: string[] | null; declare parentsIn: string[] | null; @@ -147,14 +144,6 @@ AgentDataSourceConfiguration.init( allowNull: false, defaultValue: DataTypes.NOW, }, - timeframeDuration: { - type: DataTypes.INTEGER, - allowNull: true, - }, - timeframeUnit: { - type: DataTypes.STRING, - allowNull: true, - }, tagsIn: { type: DataTypes.ARRAY(DataTypes.STRING), allowNull: true, @@ -178,12 +167,16 @@ AgentDataSourceConfiguration.init( hooks: { beforeValidate: (dataSourceConfig: AgentDataSourceConfiguration) => { if ( - (dataSourceConfig.timeframeDuration === null) !== - (dataSourceConfig.timeframeUnit === null) + (dataSourceConfig.tagsIn === null) !== + (dataSourceConfig.tagsNotIn === null) ) { - throw new Error( - "Timeframe duration/unit must be both set or both null" - ); + throw new Error("Tags must be both set or both null"); + } + if ( + (dataSourceConfig.parentsIn === null) !== + (dataSourceConfig.parentsNotIn === null) + ) { + throw new Error("Parents must be both set or both null"); } }, }, diff --git a/front/types/assistant/actions/retrieval.ts b/front/types/assistant/actions/retrieval.ts index 2544edeb7ee7..4d875d2ccc3e 100644 --- a/front/types/assistant/actions/retrieval.ts +++ b/front/types/assistant/actions/retrieval.ts @@ -17,9 +17,10 @@ export type DataSourceFilter = { parents: { in: string[]; not: string[] } | null; }; +// DataSources have a unique pair (name, workspaceId) export type DataSourceConfiguration = { - workspaceId: string; - dataSourceId: string; + workspaceId: ModelId; + name: string; filter: DataSourceFilter; }; @@ -30,6 +31,15 @@ export type DataSourceConfiguration = { export type TemplatedQuery = { template: string; }; +export function isTemplatedQuery(arg: RetrievalQuery): arg is TemplatedQuery { + return (arg as TemplatedQuery).template !== undefined; +} +export function isTimeFrame(arg: RetrievalTimeframe): arg is TimeFrame { + return ( + (arg as TimeFrame).duration !== undefined && + (arg as TimeFrame).unit !== undefined + ); +} // Retrieval specifies a list of data sources (with possible parent / tags filtering, possible "all" // data sources), a query ("auto" generated by the model "none", no query, `TemplatedQuery`, fixed @@ -39,13 +49,17 @@ export type TemplatedQuery = { // `query` and `relativeTimeFrame` will be used to generate the inputs specification for the model // in charge of generating the action inputs. The results will be used along with `topK` and // `dataSources` to query the data. +export type RetrievalTimeframe = "auto" | "none" | TimeFrame; +export type RetrievalQuery = "auto" | "none" | TemplatedQuery; +export type RetrievalDataSourcesConfiguration = DataSourceConfiguration[]; + export type RetrievalConfigurationType = { id: ModelId; type: "retrieval_configuration"; - dataSources: "all" | DataSourceConfiguration[]; - query: "auto" | "none" | TemplatedQuery; - relativeTimeFrame: "auto" | "none" | TimeFrame; + dataSources: RetrievalDataSourcesConfiguration; + query: RetrievalQuery; + relativeTimeFrame: RetrievalTimeframe; topK: number; // Dynamically decide to skip, if needed in the future diff --git a/front/types/assistant/agent-utils.ts b/front/types/assistant/agent-utils.ts new file mode 100644 index 000000000000..3a3abdc09a49 --- /dev/null +++ b/front/types/assistant/agent-utils.ts @@ -0,0 +1,134 @@ +import { DataSource } from "@app/lib/models"; +import { + AgentDataSourceConfiguration, + AgentRetrievalConfiguration, +} from "@app/lib/models/assistant/actions/retrieval"; +import { + AgentConfiguration, + AgentGenerationConfiguration, +} from "@app/lib/models/assistant/agent"; +import { + RetrievalDataSourcesConfiguration, + RetrievalQuery, + RetrievalTimeframe, +} from "@app/types/assistant/actions/retrieval"; +import { + AgentActionConfigurationType, + AgentConfigurationType, + AgentGenerationConfigurationType, +} from "@app/types/assistant/agent"; + +/** + * Builds the agent configuration type from the model + */ +export function _getAgentConfigurationType({ + agent, + action, + generation, + dataSources, +}: { + agent: AgentConfiguration; + action: AgentRetrievalConfiguration | null; + generation: AgentGenerationConfiguration | null; + dataSources: AgentDataSourceConfiguration[] | null; +}): AgentConfigurationType { + return { + sId: agent.sId, + name: agent.name, + pictureUrl: agent.pictureUrl, + status: agent.status, + action: action + ? _buildAgentActionConfigurationType(action, dataSources) + : null, + generation: generation + ? _buildAgentGenerationConfigurationType(generation) + : null, + }; +} + +/** + * Builds the agent action configuration type from the model + */ +export function _buildAgentActionConfigurationType( + action: AgentRetrievalConfiguration, + dataSourcesConfig: AgentDataSourceConfiguration[] | null +): AgentActionConfigurationType { + // Build Retrieval Timeframe + let timeframe: RetrievalTimeframe = "auto"; + if ( + action.relativeTimeFrame === "custom" && + action.relativeTimeFrameDuration && + action.relativeTimeFrameUnit + ) { + timeframe = { + duration: action.relativeTimeFrameDuration, + unit: action.relativeTimeFrameUnit, + }; + } else if (action.relativeTimeFrame === "none") { + timeframe = "none"; + } + + // Build Retrieval Query + let query: RetrievalQuery = "auto"; + if (action.query === "templated" && action.queryTemplate) { + query = { + template: action.queryTemplate, + }; + } else if (action.query === "none") { + query = "none"; + } + + // Build Retrieval DataSources + const retrievalDataSourcesConfig: RetrievalDataSourcesConfiguration = []; + let dataSource: DataSource | null = null; + + dataSourcesConfig?.forEach(async (dsConfig) => { + dataSource = await DataSource.findOne({ + where: { + id: dsConfig.dataSourceId, + }, + }); + + if (!dataSource) { + return; + } + retrievalDataSourcesConfig.push({ + name: dataSource.name, + workspaceId: dataSource.workspaceId, + filter: { + tags: + dsConfig.tagsIn && dsConfig.tagsNotIn + ? { in: dsConfig.tagsIn, not: dsConfig.tagsNotIn } + : null, + parents: + dsConfig.parentsIn && dsConfig.parentsNotIn + ? { in: dsConfig.parentsIn, not: dsConfig.parentsNotIn } + : null, + }, + }); + }); + + return { + id: action.id, + query: query, + relativeTimeFrame: timeframe, + topK: action.topK, + type: "retrieval_configuration", + dataSources: retrievalDataSourcesConfig, + }; +} + +/** + * Builds the agent generation configuration type from the model + */ +export function _buildAgentGenerationConfigurationType( + generation: AgentGenerationConfiguration +): AgentGenerationConfigurationType { + return { + prompt: generation.prompt, + model: { + providerId: generation.modelProvider, + modelId: generation.modelId, + }, + }; +} From 7fb27046ed50a18e3795e885f9275ea191136683 Mon Sep 17 00:00:00 2001 From: PopDaph Date: Thu, 7 Sep 2023 18:34:05 +0200 Subject: [PATCH 02/12] Fix agent datasource config filter types --- front/lib/api/assistant/actions/retrieval.ts | 15 +- front/lib/api/assistant/agent.ts | 179 +++++++++++++----- .../lib/models/assistant/actions/retrieval.ts | 6 + front/types/assistant/actions/retrieval.ts | 13 +- front/types/assistant/agent-utils.ts | 51 +++-- 5 files changed, 192 insertions(+), 72 deletions(-) diff --git a/front/lib/api/assistant/actions/retrieval.ts b/front/lib/api/assistant/actions/retrieval.ts index b151f9b95af1..56957bf3c6ab 100644 --- a/front/lib/api/assistant/actions/retrieval.ts +++ b/front/lib/api/assistant/actions/retrieval.ts @@ -17,7 +17,7 @@ import { Err, Ok, Result } from "@app/lib/result"; import { new_id } from "@app/lib/utils"; import logger from "@app/logger/logger"; import { - DataSourceConfiguration, + AgentDataSourceConfigurationType, isRetrievalConfiguration, RetrievalActionType, RetrievalConfigurationType, @@ -273,7 +273,7 @@ export type RetrievalParamsEvent = { created: number; configurationId: string; messageId: string; - dataSources: "all" | DataSourceConfiguration[]; + dataSources: AgentDataSourceConfigurationType[]; query: string | null; relativeTimeFrame: TimeFrame | null; topK: number; @@ -388,8 +388,8 @@ export async function* runRetrieval( // Handle data sources list and parents/tags filtering. config.DATASOURCE.data_sources = c.dataSources.map((d) => ({ - workspace_id: d.workspaceId, - data_source_id: d.name, + workspace_id: d.workspaceSId, + data_source_id: d.dataSourceName, })); for (const ds of c.dataSources) { @@ -412,6 +412,13 @@ export async function* runRetrieval( } } + // Handle timestamp filtering. + if (params.relativeTimeFrame) { + config.DATASOURCE.filter.timestamp = { + gt: timeFrameFromNow(params.relativeTimeFrame), + }; + } + // Handle top k. config.DATASOURCE.top_k = params.topK; diff --git a/front/lib/api/assistant/agent.ts b/front/lib/api/assistant/agent.ts index 0dd606ea7584..6a1edd278832 100644 --- a/front/lib/api/assistant/agent.ts +++ b/front/lib/api/assistant/agent.ts @@ -1,3 +1,5 @@ +import { Op, Transaction } from "sequelize"; + import { cloneBaseConfig, DustProdActionRegistry, @@ -14,8 +16,8 @@ import { runGeneration, } from "@app/lib/api/assistant/generation"; import { Authenticator } from "@app/lib/auth"; -import { front_sequelize } from "@app/lib/databases"; -import { DataSource } from "@app/lib/models"; +import { front_sequelize, ModelId } from "@app/lib/databases"; +import { DataSource, Workspace } from "@app/lib/models"; import { AgentDataSourceConfiguration, AgentRetrievalConfiguration, @@ -28,9 +30,10 @@ import { Err, Ok, Result } from "@app/lib/result"; import { generateModelSId } from "@app/lib/utils"; import { isRetrievalConfiguration } from "@app/types/assistant/actions/retrieval"; import { + AgentDataSourceConfigurationType, + DataSourceFilter, isTemplatedQuery, isTimeFrame, - RetrievalDataSourcesConfiguration, } from "@app/types/assistant/actions/retrieval"; import { AgentActionConfigurationType, @@ -48,9 +51,8 @@ import { } from "@app/types/assistant/conversation"; /** - * Agent configuration. + * Create a new Agent */ - export async function createAgentConfiguration( auth: Authenticator, { @@ -66,14 +68,18 @@ export async function createAgentConfiguration( } ): Promise { const owner = auth.workspace(); - if (!owner) { return; } return await front_sequelize.transaction(async (t) => { + let agentConfigRow: AgentConfiguration | null = null; + let agentGenerationConfigRow: AgentGenerationConfiguration | null = null; + let agentActionConfigRow: AgentRetrievalConfiguration | null = null; + let agentDataSourcesConfigRows: AgentDataSourceConfiguration[] = []; + // Create AgentConfiguration - const agentConfigRow = await AgentConfiguration.create( + agentConfigRow = await AgentConfiguration.create( { sId: generateModelSId(), status: "active", @@ -85,10 +91,6 @@ export async function createAgentConfiguration( { transaction: t } ); - let agentGenerationConfigRow: AgentGenerationConfiguration | null = null; - let agentActionConfigRow: AgentRetrievalConfiguration | null = null; - const agentDataSourcesConfigRows: AgentDataSourceConfiguration[] = []; - // Create AgentGenerationConfiguration if (generation) { agentGenerationConfigRow = await AgentGenerationConfiguration.create( @@ -102,11 +104,10 @@ export async function createAgentConfiguration( ); } - // Create AgentRetrievalConfiguration & associated AgentDataSourceConfiguration + // Create AgentRetrievalConfiguration if (action) { const query = action.query; const timeframe = action.relativeTimeFrame; - agentActionConfigRow = await AgentRetrievalConfiguration.create( { query: isTemplatedQuery(query) ? "templated" : query, @@ -121,40 +122,18 @@ export async function createAgentConfiguration( }, { transaction: t } ); + } - if (!agentActionConfigRow) { - return; - } - - if (action.dataSources) { - let dsRow: AgentDataSourceConfiguration | null = null; - action.dataSources.map(async (d) => { - const ds = await DataSource.findOne({ - where: { - name: d.name, - workspaceId: d.workspaceId, - }, - }); - - if (ds && agentActionConfigRow) { - dsRow = await AgentDataSourceConfiguration.create( - { - dataSourceId: ds.id, - tagsIn: d.filter.tags?.in, - tagsNotIn: d.filter.tags?.not, - parentsIn: d.filter.parents?.in, - parentsNotIn: d.filter.parents?.not, - retrievalConfigurationId: agentActionConfigRow.id, - }, - { transaction: t } - ); - agentDataSourcesConfigRows.push(dsRow); - } - }); - } + // Create AgentDataSourceConfiguration + if (agentActionConfigRow && action?.dataSources) { + agentDataSourcesConfigRows = await _createAgentDataSourcesConfigData( + t, + action.dataSources, + agentActionConfigRow.id + ); } - return _getAgentConfigurationType({ + return await _getAgentConfigurationType({ agent: agentConfigRow, action: agentActionConfigRow, generation: agentGenerationConfigRow, @@ -163,6 +142,118 @@ export async function createAgentConfiguration( }); } +/** + * Create the AgentDataSourceConfiguration rows in database. + * + * Knowing that a datasource is uniquely identified by its name and its workspaceId + * We need to fetch the dataSources from the database from that. + * We obvisously need to do as few queries as possible. + */ +async function _createAgentDataSourcesConfigData( + t: Transaction, + dataSourcesConfig: AgentDataSourceConfigurationType[], + agentActionId: number +): Promise { + // dsConfig contains this format: + // [ + // { workspaceSId: s1o1u1p, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, + // { workspaceSId: s1o1u1p, dataSourceName: "managed-slack", filter: { tags: null, parents: null } }, + // { workspaceSId: i2n2o2u, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, + // ] + + // First we get the list of workspaces because we need the mapping between workspaceSId and workspaceId + const workspaces = await Workspace.findAll({ + where: { + sId: dataSourcesConfig.map((dsConfig) => dsConfig.workspaceSId), + }, + attributes: ["id", "sId"], + }); + + // Now will want to group the datasource names by workspaceId to do only one query per workspace. + // We want this: + // [ + // { workspaceId: 1, dataSourceNames: [""managed-notion", "managed-slack"] }, + // { workspaceId: 2, dataSourceNames: ["managed-notion"] } + // ] + type _DsNamesPerWorkspaceIdType = { + workspaceId: number; + dataSourceNames: string[]; + }; + const dsNamesPerWorkspaceId = dataSourcesConfig.reduce( + ( + acc: _DsNamesPerWorkspaceIdType[], + curr: AgentDataSourceConfigurationType + ) => { + // First we need to get the workspaceId from the workspaceSId + const workspace = workspaces.find((w) => w.sId === curr.workspaceSId); + if (!workspace) { + throw new Error("Workspace not found"); + } + + // Find an existing entry for this workspaceId + const existingEntry: _DsNamesPerWorkspaceIdType | undefined = acc.find( + (entry: _DsNamesPerWorkspaceIdType) => + entry.workspaceId === workspace.id + ); + if (existingEntry) { + // Append dataSourceName to existing entry + existingEntry.dataSourceNames.push(curr.dataSourceName); + } else { + // Add a new entry for this workspaceId + acc.push({ + workspaceId: workspace.id, + dataSourceNames: [curr.dataSourceName], + }); + } + return acc; + }, + [] + ); + + // Then we get do one findAllQuery per workspaceId, in a Promise.all + const getDataSourcesQueries = dsNamesPerWorkspaceId.map( + ({ workspaceId, dataSourceNames }) => { + return DataSource.findAll({ + where: { + workspaceId, + name: { + [Op.in]: dataSourceNames, + }, + }, + }); + } + ); + const results = await Promise.all(getDataSourcesQueries); + const dataSources = results.flat(); + + const agentDataSourcesConfigRows: AgentDataSourceConfiguration[] = + await Promise.all( + dataSourcesConfig.map(async (dsConfig) => { + const dataSource = dataSources.find( + (ds) => + ds.name === dsConfig.dataSourceName && + ds.workspaceId === + workspaces.find((w) => w.sId === dsConfig.workspaceSId)?.id + ); + if (!dataSource) { + throw new Error("DataSource not found"); + } + return AgentDataSourceConfiguration.create( + { + dataSourceId: dataSource.id, + tagsIn: dsConfig.filter.tags?.in, + tagsNotIn: dsConfig.filter.tags?.not, + parentsIn: dsConfig.filter.parents?.in, + parentsNotIn: dsConfig.filter.parents?.not, + retrievalConfigurationId: agentActionId, + }, + { transaction: t } + ); + }) + ); + return agentDataSourcesConfigRows; +} + export async function updateAgentConfiguration( auth: Authenticator, configurationId: string, diff --git a/front/lib/models/assistant/actions/retrieval.ts b/front/lib/models/assistant/actions/retrieval.ts index 98dcaac79346..95adea5566fd 100644 --- a/front/lib/models/assistant/actions/retrieval.ts +++ b/front/lib/models/assistant/actions/retrieval.ts @@ -189,6 +189,12 @@ AgentRetrievalConfiguration.hasMany(AgentDataSourceConfiguration, { onDelete: "CASCADE", }); +// Data source <> Data source config +DataSource.hasMany(AgentDataSourceConfiguration, { + foreignKey: { name: "dataSourceId", allowNull: false }, + onDelete: "CASCADE", +}); + // Agent config <> Retrieval config AgentConfiguration.hasOne(AgentRetrievalConfiguration, { foreignKey: { name: "agentId", allowNull: true }, // null = no generation set for this Agent diff --git a/front/types/assistant/actions/retrieval.ts b/front/types/assistant/actions/retrieval.ts index 4d875d2ccc3e..eff6727df63b 100644 --- a/front/types/assistant/actions/retrieval.ts +++ b/front/types/assistant/actions/retrieval.ts @@ -17,10 +17,10 @@ export type DataSourceFilter = { parents: { in: string[]; not: string[] } | null; }; -// DataSources have a unique pair (name, workspaceId) -export type DataSourceConfiguration = { - workspaceId: ModelId; - name: string; +// This is used to talk with Dust Apps and Core, so it store external Ids. +export type AgentDataSourceConfigurationType = { + workspaceSId: string; // = Workspace.sId + dataSourceName: string; // = Datasource.name filter: DataSourceFilter; }; @@ -51,7 +51,8 @@ export function isTimeFrame(arg: RetrievalTimeframe): arg is TimeFrame { // `dataSources` to query the data. export type RetrievalTimeframe = "auto" | "none" | TimeFrame; export type RetrievalQuery = "auto" | "none" | TemplatedQuery; -export type RetrievalDataSourcesConfiguration = DataSourceConfiguration[]; +export type RetrievalDataSourcesConfiguration = + AgentDataSourceConfigurationType[]; export type RetrievalConfigurationType = { id: ModelId; @@ -102,7 +103,7 @@ export type RetrievalActionType = { id: ModelId; // AgentRetrieval. type: "retrieval_action"; params: { - dataSources: "all" | DataSourceConfiguration[]; + dataSources: "all" | AgentDataSourceConfigurationType[]; relativeTimeFrame: TimeFrame | null; query: string | null; topK: number; diff --git a/front/types/assistant/agent-utils.ts b/front/types/assistant/agent-utils.ts index 3a3abdc09a49..a41a0f910eeb 100644 --- a/front/types/assistant/agent-utils.ts +++ b/front/types/assistant/agent-utils.ts @@ -1,4 +1,6 @@ -import { DataSource } from "@app/lib/models"; +import { Op } from "sequelize"; + +import { DataSource, Workspace } from "@app/lib/models"; import { AgentDataSourceConfiguration, AgentRetrievalConfiguration, @@ -21,7 +23,7 @@ import { /** * Builds the agent configuration type from the model */ -export function _getAgentConfigurationType({ +export async function _getAgentConfigurationType({ agent, action, generation, @@ -31,14 +33,14 @@ export function _getAgentConfigurationType({ action: AgentRetrievalConfiguration | null; generation: AgentGenerationConfiguration | null; dataSources: AgentDataSourceConfiguration[] | null; -}): AgentConfigurationType { +}): Promise { return { sId: agent.sId, name: agent.name, pictureUrl: agent.pictureUrl, status: agent.status, action: action - ? _buildAgentActionConfigurationType(action, dataSources) + ? await _buildAgentActionConfigurationType(action, dataSources || []) : null, generation: generation ? _buildAgentGenerationConfigurationType(generation) @@ -49,10 +51,10 @@ export function _getAgentConfigurationType({ /** * Builds the agent action configuration type from the model */ -export function _buildAgentActionConfigurationType( +export async function _buildAgentActionConfigurationType( action: AgentRetrievalConfiguration, - dataSourcesConfig: AgentDataSourceConfiguration[] | null -): AgentActionConfigurationType { + dataSourcesConfig: AgentDataSourceConfiguration[] +): Promise { // Build Retrieval Timeframe let timeframe: RetrievalTimeframe = "auto"; if ( @@ -80,21 +82,34 @@ export function _buildAgentActionConfigurationType( // Build Retrieval DataSources const retrievalDataSourcesConfig: RetrievalDataSourcesConfiguration = []; - let dataSource: DataSource | null = null; - dataSourcesConfig?.forEach(async (dsConfig) => { - dataSource = await DataSource.findOne({ - where: { - id: dsConfig.dataSourceId, - }, - }); + const dataSourcesIds = dataSourcesConfig?.map((ds) => ds.dataSourceId); + const dataSources = await DataSource.findAll({ + where: { + id: { [Op.in]: dataSourcesIds }, + }, + }); + const workspaceIds = dataSources.map((ds) => ds.workspaceId); + const workspaces = await Workspace.findAll({ + where: { + id: { [Op.in]: workspaceIds }, + }, + }); + + let dataSource: DataSource | undefined; + let workspace: Workspace | undefined; - if (!dataSource) { - return; + dataSourcesConfig.forEach(async (dsConfig) => { + dataSource = dataSources.find((ds) => ds.id === dsConfig.dataSourceId); + workspace = workspaces.find((w) => w.id === dataSource?.workspaceId); + + if (!dataSource || !workspace) { + throw new Error("Could not find dataSource or workspace"); } + retrievalDataSourcesConfig.push({ - name: dataSource.name, - workspaceId: dataSource.workspaceId, + dataSourceName: dataSource.name, + workspaceSId: workspace.sId, filter: { tags: dsConfig.tagsIn && dsConfig.tagsNotIn From f27a5bdd0171b060ad28b61494bd1abe9851d795 Mon Sep 17 00:00:00 2001 From: PopDaph Date: Thu, 7 Sep 2023 19:04:06 +0200 Subject: [PATCH 03/12] Add get function --- front/lib/api/assistant/agent.ts | 48 +++++++++++++++++-- .../api/assistant/agent_utils.ts} | 0 2 files changed, 45 insertions(+), 3 deletions(-) rename front/{types/assistant/agent-utils.ts => lib/api/assistant/agent_utils.ts} (100%) diff --git a/front/lib/api/assistant/agent.ts b/front/lib/api/assistant/agent.ts index 6a1edd278832..32dfb586b1d6 100644 --- a/front/lib/api/assistant/agent.ts +++ b/front/lib/api/assistant/agent.ts @@ -10,13 +10,14 @@ import { RetrievalParamsEvent, runRetrieval, } from "@app/lib/api/assistant/actions/retrieval"; +import { _getAgentConfigurationType } from "@app/lib/api/assistant/agent_utils"; import { GenerationTokensEvent, renderConversationForModel, runGeneration, } from "@app/lib/api/assistant/generation"; import { Authenticator } from "@app/lib/auth"; -import { front_sequelize, ModelId } from "@app/lib/databases"; +import { front_sequelize } from "@app/lib/databases"; import { DataSource, Workspace } from "@app/lib/models"; import { AgentDataSourceConfiguration, @@ -31,7 +32,6 @@ import { generateModelSId } from "@app/lib/utils"; import { isRetrievalConfiguration } from "@app/types/assistant/actions/retrieval"; import { AgentDataSourceConfigurationType, - DataSourceFilter, isTemplatedQuery, isTimeFrame, } from "@app/types/assistant/actions/retrieval"; @@ -42,7 +42,6 @@ import { AgentConfigurationType, AgentGenerationConfigurationType, } from "@app/types/assistant/agent"; -import { _getAgentConfigurationType } from "@app/types/assistant/agent-utils"; import { AgentActionType, AgentMessageType, @@ -50,6 +49,49 @@ import { UserMessageType, } from "@app/types/assistant/conversation"; +/** + * Get an agent configuration from its name + */ +export async function getAgentConfiguration(auth: Authenticator, name: string) { + const owner = auth.workspace(); + if (!owner) { + return; + } + const agent = await AgentConfiguration.findOne({ + where: { + name: name, + workspaceId: owner.id, + }, + }); + const agentGeneration = await AgentGenerationConfiguration.findOne({ + where: { + agentId: agent?.id, + }, + }); + const agentAction = await AgentRetrievalConfiguration.findOne({ + where: { + agentId: agent?.id, + }, + }); + const agentDataSources = agentAction?.id + ? await AgentDataSourceConfiguration.findAll({ + where: { + retrievalConfigurationId: agentAction?.id, + }, + }) + : []; + + if (!agent) { + return; + } + return await _getAgentConfigurationType({ + agent: agent, + generation: agentGeneration, + action: agentAction, + dataSources: agentDataSources, + }); +} + /** * Create a new Agent */ diff --git a/front/types/assistant/agent-utils.ts b/front/lib/api/assistant/agent_utils.ts similarity index 100% rename from front/types/assistant/agent-utils.ts rename to front/lib/api/assistant/agent_utils.ts From 6adc64893602157eb70a74736051b1dfe1b6f623 Mon Sep 17 00:00:00 2001 From: PopDaph Date: Thu, 7 Sep 2023 19:19:44 +0200 Subject: [PATCH 04/12] cleanup --- front/lib/api/assistant/agent.ts | 6 +----- front/lib/api/assistant/agent_utils.ts | 2 ++ 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/front/lib/api/assistant/agent.ts b/front/lib/api/assistant/agent.ts index 32dfb586b1d6..bd6f3c2755bd 100644 --- a/front/lib/api/assistant/agent.ts +++ b/front/lib/api/assistant/agent.ts @@ -146,7 +146,7 @@ export async function createAgentConfiguration( ); } - // Create AgentRetrievalConfiguration + // Create AgentRetrievalConfiguration & AgentDataSourceConfiguration if (action) { const query = action.query; const timeframe = action.relativeTimeFrame; @@ -164,10 +164,6 @@ export async function createAgentConfiguration( }, { transaction: t } ); - } - - // Create AgentDataSourceConfiguration - if (agentActionConfigRow && action?.dataSources) { agentDataSourcesConfigRows = await _createAgentDataSourcesConfigData( t, action.dataSources, diff --git a/front/lib/api/assistant/agent_utils.ts b/front/lib/api/assistant/agent_utils.ts index a41a0f910eeb..bd30750d26da 100644 --- a/front/lib/api/assistant/agent_utils.ts +++ b/front/lib/api/assistant/agent_utils.ts @@ -88,12 +88,14 @@ export async function _buildAgentActionConfigurationType( where: { id: { [Op.in]: dataSourcesIds }, }, + attributes: ["id", "name", "workspaceId"], }); const workspaceIds = dataSources.map((ds) => ds.workspaceId); const workspaces = await Workspace.findAll({ where: { id: { [Op.in]: workspaceIds }, }, + attributes: ["id", "sId"], }); let dataSource: DataSource | undefined; From 3fdbe73583d3aa20e6ec017a8eb67ca7c80ee1dd Mon Sep 17 00:00:00 2001 From: PopDaph Date: Fri, 8 Sep 2023 11:17:19 +0200 Subject: [PATCH 05/12] Rework functions - split --- front/lib/api/assistant/actions/retrieval.ts | 15 +- front/lib/api/assistant/agent.ts | 291 ------------------ front/lib/api/assistant/agent/agent_create.ts | 283 +++++++++++++++++ .../{agent_utils.ts => agent/agent_get.ts} | 102 ++++-- front/lib/api/assistant/agent/agent_update.ts | 208 +++++++++++++ front/types/assistant/agent.ts | 10 +- front/types/assistant/conversation.ts | 4 +- 7 files changed, 582 insertions(+), 331 deletions(-) create mode 100644 front/lib/api/assistant/agent/agent_create.ts rename front/lib/api/assistant/{agent_utils.ts => agent/agent_get.ts} (69%) create mode 100644 front/lib/api/assistant/agent/agent_update.ts diff --git a/front/lib/api/assistant/actions/retrieval.ts b/front/lib/api/assistant/actions/retrieval.ts index 56957bf3c6ab..f73e1be1d286 100644 --- a/front/lib/api/assistant/actions/retrieval.ts +++ b/front/lib/api/assistant/actions/retrieval.ts @@ -26,7 +26,7 @@ import { } from "@app/types/assistant/actions/retrieval"; import { AgentActionSpecification, - AgentConfigurationType, + AgentFullConfigurationType, } from "@app/types/assistant/agent"; import { AgentMessageType, @@ -167,6 +167,7 @@ export async function retrievalActionSpecification( } return { + id: configuration.id, name: "search_data_sources", description: "Search the data sources specified by the user for information to answer their request." + @@ -315,7 +316,7 @@ export type RetrievalSuccessEvent = { // error is expected to be stored by the caller on the parent agent message. export async function* runRetrieval( auth: Authenticator, - configuration: AgentConfigurationType, + configuration: AgentFullConfigurationType, conversation: ConversationType, userMessage: UserMessageType, agentMessage: AgentMessageType @@ -348,7 +349,7 @@ export async function* runRetrieval( return yield { type: "retrieval_error", created: Date.now(), - configurationId: configuration.sId, + configurationId: configuration.agent.sId, messageId: agentMessage.sId, error: { code: "retrieval_parameters_generation_error", @@ -432,7 +433,7 @@ export async function* runRetrieval( return yield { type: "retrieval_error", created: Date.now(), - configurationId: configuration.sId, + configurationId: configuration.agent.sId, messageId: agentMessage.sId, error: { code: "retrieval_search_error", @@ -449,7 +450,7 @@ export async function* runRetrieval( return yield { type: "retrieval_error", created: Date.now(), - configurationId: configuration.sId, + configurationId: configuration.agent.sId, messageId: agentMessage.sId, error: { code: "retrieval_search_error", @@ -528,7 +529,7 @@ export async function* runRetrieval( yield { type: "retrieval_documents", created: Date.now(), - configurationId: configuration.sId, + configurationId: configuration.agent.sId, messageId: agentMessage.sId, documents, }; @@ -536,7 +537,7 @@ export async function* runRetrieval( yield { type: "retrieval_success", created: Date.now(), - configurationId: configuration.sId, + configurationId: configuration.agent.sId, messageId: agentMessage.sId, action: { id: action.id, diff --git a/front/lib/api/assistant/agent.ts b/front/lib/api/assistant/agent.ts index bd6f3c2755bd..7799d8dd7549 100644 --- a/front/lib/api/assistant/agent.ts +++ b/front/lib/api/assistant/agent.ts @@ -1,5 +1,3 @@ -import { Op, Transaction } from "sequelize"; - import { cloneBaseConfig, DustProdActionRegistry, @@ -10,37 +8,18 @@ import { RetrievalParamsEvent, runRetrieval, } from "@app/lib/api/assistant/actions/retrieval"; -import { _getAgentConfigurationType } from "@app/lib/api/assistant/agent_utils"; import { GenerationTokensEvent, renderConversationForModel, runGeneration, } from "@app/lib/api/assistant/generation"; import { Authenticator } from "@app/lib/auth"; -import { front_sequelize } from "@app/lib/databases"; -import { DataSource, Workspace } from "@app/lib/models"; -import { - AgentDataSourceConfiguration, - AgentRetrievalConfiguration, -} from "@app/lib/models/assistant/actions/retrieval"; -import { - AgentConfiguration, - AgentGenerationConfiguration, -} from "@app/lib/models/assistant/agent"; import { Err, Ok, Result } from "@app/lib/result"; import { generateModelSId } from "@app/lib/utils"; import { isRetrievalConfiguration } from "@app/types/assistant/actions/retrieval"; import { - AgentDataSourceConfigurationType, - isTemplatedQuery, - isTimeFrame, -} from "@app/types/assistant/actions/retrieval"; -import { - AgentActionConfigurationType, AgentActionSpecification, - AgentConfigurationStatus, AgentConfigurationType, - AgentGenerationConfigurationType, } from "@app/types/assistant/agent"; import { AgentActionType, @@ -49,276 +28,6 @@ import { UserMessageType, } from "@app/types/assistant/conversation"; -/** - * Get an agent configuration from its name - */ -export async function getAgentConfiguration(auth: Authenticator, name: string) { - const owner = auth.workspace(); - if (!owner) { - return; - } - const agent = await AgentConfiguration.findOne({ - where: { - name: name, - workspaceId: owner.id, - }, - }); - const agentGeneration = await AgentGenerationConfiguration.findOne({ - where: { - agentId: agent?.id, - }, - }); - const agentAction = await AgentRetrievalConfiguration.findOne({ - where: { - agentId: agent?.id, - }, - }); - const agentDataSources = agentAction?.id - ? await AgentDataSourceConfiguration.findAll({ - where: { - retrievalConfigurationId: agentAction?.id, - }, - }) - : []; - - if (!agent) { - return; - } - return await _getAgentConfigurationType({ - agent: agent, - generation: agentGeneration, - action: agentAction, - dataSources: agentDataSources, - }); -} - -/** - * Create a new Agent - */ -export async function createAgentConfiguration( - auth: Authenticator, - { - name, - pictureUrl, - action, - generation, - }: { - name: string; - pictureUrl?: string; - action?: AgentActionConfigurationType; - generation?: AgentGenerationConfigurationType; - } -): Promise { - const owner = auth.workspace(); - if (!owner) { - return; - } - - return await front_sequelize.transaction(async (t) => { - let agentConfigRow: AgentConfiguration | null = null; - let agentGenerationConfigRow: AgentGenerationConfiguration | null = null; - let agentActionConfigRow: AgentRetrievalConfiguration | null = null; - let agentDataSourcesConfigRows: AgentDataSourceConfiguration[] = []; - - // Create AgentConfiguration - agentConfigRow = await AgentConfiguration.create( - { - sId: generateModelSId(), - status: "active", - name: name, - pictureUrl: pictureUrl ?? null, - scope: "workspace", - workspaceId: owner.id, - }, - { transaction: t } - ); - - // Create AgentGenerationConfiguration - if (generation) { - agentGenerationConfigRow = await AgentGenerationConfiguration.create( - { - prompt: generation.prompt, - modelProvider: generation.model.providerId, - modelId: generation.model.modelId, - agentId: agentConfigRow.id, - }, - { transaction: t } - ); - } - - // Create AgentRetrievalConfiguration & AgentDataSourceConfiguration - if (action) { - const query = action.query; - const timeframe = action.relativeTimeFrame; - agentActionConfigRow = await AgentRetrievalConfiguration.create( - { - query: isTemplatedQuery(query) ? "templated" : query, - queryTemplate: isTemplatedQuery(query) ? query.template : null, - relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, - relativeTimeFrameDuration: isTimeFrame(timeframe) - ? timeframe.duration - : null, - relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, - topK: action.topK, - agentId: agentConfigRow.id, - }, - { transaction: t } - ); - agentDataSourcesConfigRows = await _createAgentDataSourcesConfigData( - t, - action.dataSources, - agentActionConfigRow.id - ); - } - - return await _getAgentConfigurationType({ - agent: agentConfigRow, - action: agentActionConfigRow, - generation: agentGenerationConfigRow, - dataSources: agentDataSourcesConfigRows, - }); - }); -} - -/** - * Create the AgentDataSourceConfiguration rows in database. - * - * Knowing that a datasource is uniquely identified by its name and its workspaceId - * We need to fetch the dataSources from the database from that. - * We obvisously need to do as few queries as possible. - */ -async function _createAgentDataSourcesConfigData( - t: Transaction, - dataSourcesConfig: AgentDataSourceConfigurationType[], - agentActionId: number -): Promise { - // dsConfig contains this format: - // [ - // { workspaceSId: s1o1u1p, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, - // { workspaceSId: s1o1u1p, dataSourceName: "managed-slack", filter: { tags: null, parents: null } }, - // { workspaceSId: i2n2o2u, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, - // ] - - // First we get the list of workspaces because we need the mapping between workspaceSId and workspaceId - const workspaces = await Workspace.findAll({ - where: { - sId: dataSourcesConfig.map((dsConfig) => dsConfig.workspaceSId), - }, - attributes: ["id", "sId"], - }); - - // Now will want to group the datasource names by workspaceId to do only one query per workspace. - // We want this: - // [ - // { workspaceId: 1, dataSourceNames: [""managed-notion", "managed-slack"] }, - // { workspaceId: 2, dataSourceNames: ["managed-notion"] } - // ] - type _DsNamesPerWorkspaceIdType = { - workspaceId: number; - dataSourceNames: string[]; - }; - const dsNamesPerWorkspaceId = dataSourcesConfig.reduce( - ( - acc: _DsNamesPerWorkspaceIdType[], - curr: AgentDataSourceConfigurationType - ) => { - // First we need to get the workspaceId from the workspaceSId - const workspace = workspaces.find((w) => w.sId === curr.workspaceSId); - if (!workspace) { - throw new Error("Workspace not found"); - } - - // Find an existing entry for this workspaceId - const existingEntry: _DsNamesPerWorkspaceIdType | undefined = acc.find( - (entry: _DsNamesPerWorkspaceIdType) => - entry.workspaceId === workspace.id - ); - if (existingEntry) { - // Append dataSourceName to existing entry - existingEntry.dataSourceNames.push(curr.dataSourceName); - } else { - // Add a new entry for this workspaceId - acc.push({ - workspaceId: workspace.id, - dataSourceNames: [curr.dataSourceName], - }); - } - return acc; - }, - [] - ); - - // Then we get do one findAllQuery per workspaceId, in a Promise.all - const getDataSourcesQueries = dsNamesPerWorkspaceId.map( - ({ workspaceId, dataSourceNames }) => { - return DataSource.findAll({ - where: { - workspaceId, - name: { - [Op.in]: dataSourceNames, - }, - }, - }); - } - ); - const results = await Promise.all(getDataSourcesQueries); - const dataSources = results.flat(); - - const agentDataSourcesConfigRows: AgentDataSourceConfiguration[] = - await Promise.all( - dataSourcesConfig.map(async (dsConfig) => { - const dataSource = dataSources.find( - (ds) => - ds.name === dsConfig.dataSourceName && - ds.workspaceId === - workspaces.find((w) => w.sId === dsConfig.workspaceSId)?.id - ); - if (!dataSource) { - throw new Error("DataSource not found"); - } - return AgentDataSourceConfiguration.create( - { - dataSourceId: dataSource.id, - tagsIn: dsConfig.filter.tags?.in, - tagsNotIn: dsConfig.filter.tags?.not, - parentsIn: dsConfig.filter.parents?.in, - parentsNotIn: dsConfig.filter.parents?.not, - retrievalConfigurationId: agentActionId, - }, - { transaction: t } - ); - }) - ); - return agentDataSourcesConfigRows; -} - -export async function updateAgentConfiguration( - auth: Authenticator, - configurationId: string, - { - name, - pictureUrl, - status, - action, - generation, - }: { - name: string; - pictureUrl?: string; - status: AgentConfigurationStatus; - action?: AgentActionConfigurationType; - generation?: AgentGenerationConfigurationType; - } -): Promise { - return { - sId: generateModelSId(), - name, - pictureUrl: pictureUrl ?? null, - status, - action: action ?? null, - generation: generation ?? null, - }; -} - /** * Action Inputs generation. */ diff --git a/front/lib/api/assistant/agent/agent_create.ts b/front/lib/api/assistant/agent/agent_create.ts new file mode 100644 index 000000000000..45458446ceec --- /dev/null +++ b/front/lib/api/assistant/agent/agent_create.ts @@ -0,0 +1,283 @@ +import { Op, Transaction } from "sequelize"; + +import { + _buildAgentActionConfigurationTypeFromModel, + _buildAgentConfigurationTypeFromModel, + _buildAgentGenerationConfigurationTypeFromModel, +} from "@app/lib/api/assistant/agent/agent_get"; +import { Authenticator } from "@app/lib/auth"; +import { front_sequelize } from "@app/lib/databases"; +import { DataSource, Workspace } from "@app/lib/models"; +import { + AgentDataSourceConfiguration, + AgentRetrievalConfiguration, +} from "@app/lib/models/assistant/actions/retrieval"; +import { + AgentConfiguration, + AgentGenerationConfiguration, +} from "@app/lib/models/assistant/agent"; +import { generateModelSId } from "@app/lib/utils"; +import { + AgentDataSourceConfigurationType, + isTemplatedQuery, + isTimeFrame, + RetrievalDataSourcesConfiguration, + RetrievalQuery, + RetrievalTimeframe, +} from "@app/types/assistant/actions/retrieval"; +import { + AgentActionConfigurationType, + AgentConfigurationStatus, + AgentConfigurationType, + AgentGenerationConfigurationType, +} from "@app/types/assistant/agent"; + +/** + * Create Agent Configuration + */ +export async function createAgentConfiguration( + auth: Authenticator, + { + name, + pictureUrl, + status, + }: { + name: string; + pictureUrl: string; + status: AgentConfigurationStatus; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error("Cannot create AgentConfiguration without workspace"); + } + + const agentConfig = await AgentConfiguration.create({ + sId: generateModelSId(), + status: status, + name: name, + pictureUrl: pictureUrl, + scope: "workspace", + workspaceId: owner.id, + }); + + return _buildAgentConfigurationTypeFromModel({ + agent: agentConfig, + }); +} + +/** + * Create Agent Generation Configuration + */ +export async function createAgentGenerationConfiguration( + auth: Authenticator, + agentSid: string, + { + prompt, + modelProvider, + modelId, + }: { + prompt: string; + modelProvider: string; + modelId: string; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Workspace not found" + ); + } + + const agentConfig = await AgentConfiguration.findOne({ + where: { + sId: agentSid, + }, + }); + if (!agentConfig) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Agent not found" + ); + } + + const generation = await AgentGenerationConfiguration.create({ + prompt: prompt, + modelProvider: modelProvider, + modelId: modelId, + agentId: agentConfig.id, + }); + + return _buildAgentGenerationConfigurationTypeFromModel(generation); +} + +/** + * Create Agent Action Configuration (Retrieval) + */ +export async function createAgentActionRetrievalConfiguration( + auth: Authenticator, + agentSid: string, + { + query, + timeframe, + topK, + dataSources, + }: { + query: RetrievalQuery; + timeframe: RetrievalTimeframe; + topK: number; + dataSources: RetrievalDataSourcesConfiguration; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error( + "Cannot create AgentActionConfiguration: Workspace not found" + ); + } + + const agentConfig = await AgentConfiguration.findOne({ + where: { + sId: agentSid, + }, + }); + if (!agentConfig) { + throw new Error("Cannot create AgentActionConfiguration: Agent not found"); + } + return await front_sequelize.transaction(async (t) => { + const agentActionConfigRow = await AgentRetrievalConfiguration.create( + { + query: isTemplatedQuery(query) ? "templated" : query, + queryTemplate: isTemplatedQuery(query) ? query.template : null, + relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, + relativeTimeFrameDuration: isTimeFrame(timeframe) + ? timeframe.duration + : null, + relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, + topK: topK, + agentId: agentConfig.id, + }, + { transaction: t } + ); + const agentDataSourcesConfigRows = await _createAgentDataSourcesConfigData( + t, + dataSources, + agentActionConfigRow.id + ); + return await _buildAgentActionConfigurationTypeFromModel( + agentActionConfigRow, + agentDataSourcesConfigRows + ); + }); +} + +/** + * Create the AgentDataSourceConfiguration rows in database. + * + * Knowing that a datasource is uniquely identified by its name and its workspaceId + * We need to fetch the dataSources from the database from that. + * We obvisously need to do as few queries as possible. + */ +export async function _createAgentDataSourcesConfigData( + t: Transaction, + dataSourcesConfig: AgentDataSourceConfigurationType[], + agentActionId: number +): Promise { + // dsConfig contains this format: + // [ + // { workspaceSId: s1o1u1p, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, + // { workspaceSId: s1o1u1p, dataSourceName: "managed-slack", filter: { tags: null, parents: null } }, + // { workspaceSId: i2n2o2u, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, + // ] + + // First we get the list of workspaces because we need the mapping between workspaceSId and workspaceId + const workspaces = await Workspace.findAll({ + where: { + sId: dataSourcesConfig.map((dsConfig) => dsConfig.workspaceSId), + }, + attributes: ["id", "sId"], + }); + + // Now will want to group the datasource names by workspaceId to do only one query per workspace. + // We want this: + // [ + // { workspaceId: 1, dataSourceNames: [""managed-notion", "managed-slack"] }, + // { workspaceId: 2, dataSourceNames: ["managed-notion"] } + // ] + type _DsNamesPerWorkspaceIdType = { + workspaceId: number; + dataSourceNames: string[]; + }; + const dsNamesPerWorkspaceId = dataSourcesConfig.reduce( + ( + acc: _DsNamesPerWorkspaceIdType[], + curr: AgentDataSourceConfigurationType + ) => { + // First we need to get the workspaceId from the workspaceSId + const workspace = workspaces.find((w) => w.sId === curr.workspaceSId); + if (!workspace) { + throw new Error("Workspace not found"); + } + + // Find an existing entry for this workspaceId + const existingEntry: _DsNamesPerWorkspaceIdType | undefined = acc.find( + (entry: _DsNamesPerWorkspaceIdType) => + entry.workspaceId === workspace.id + ); + if (existingEntry) { + // Append dataSourceName to existing entry + existingEntry.dataSourceNames.push(curr.dataSourceName); + } else { + // Add a new entry for this workspaceId + acc.push({ + workspaceId: workspace.id, + dataSourceNames: [curr.dataSourceName], + }); + } + return acc; + }, + [] + ); + + // Then we get do one findAllQuery per workspaceId, in a Promise.all + const getDataSourcesQueries = dsNamesPerWorkspaceId.map( + ({ workspaceId, dataSourceNames }) => { + return DataSource.findAll({ + where: { + workspaceId, + name: { + [Op.in]: dataSourceNames, + }, + }, + }); + } + ); + const results = await Promise.all(getDataSourcesQueries); + const dataSources = results.flat(); + + const agentDataSourcesConfigRows: AgentDataSourceConfiguration[] = + await Promise.all( + dataSourcesConfig.map(async (dsConfig) => { + const dataSource = dataSources.find( + (ds) => + ds.name === dsConfig.dataSourceName && + ds.workspaceId === + workspaces.find((w) => w.sId === dsConfig.workspaceSId)?.id + ); + if (!dataSource) { + throw new Error("DataSource not found"); + } + return AgentDataSourceConfiguration.create( + { + dataSourceId: dataSource.id, + tagsIn: dsConfig.filter.tags?.in, + tagsNotIn: dsConfig.filter.tags?.not, + parentsIn: dsConfig.filter.parents?.in, + parentsNotIn: dsConfig.filter.parents?.not, + retrievalConfigurationId: agentActionId, + }, + { transaction: t } + ); + }) + ); + return agentDataSourcesConfigRows; +} diff --git a/front/lib/api/assistant/agent_utils.ts b/front/lib/api/assistant/agent/agent_get.ts similarity index 69% rename from front/lib/api/assistant/agent_utils.ts rename to front/lib/api/assistant/agent/agent_get.ts index bd30750d26da..360d99ef0ba4 100644 --- a/front/lib/api/assistant/agent_utils.ts +++ b/front/lib/api/assistant/agent/agent_get.ts @@ -1,5 +1,6 @@ import { Op } from "sequelize"; +import { Authenticator } from "@app/lib/auth"; import { DataSource, Workspace } from "@app/lib/models"; import { AgentDataSourceConfiguration, @@ -17,41 +18,99 @@ import { import { AgentActionConfigurationType, AgentConfigurationType, + AgentFullConfigurationType as AgentFullConfigurationType, AgentGenerationConfigurationType, } from "@app/types/assistant/agent"; +/** + * Get an agent full configuration from its name + */ +export async function getAgent( + auth: Authenticator, + sId: string +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error("Cannot find Agent: no workspace"); + } + const agent = await AgentConfiguration.findOne({ + where: { + sId: sId, + workspaceId: owner.id, + }, + }); + if (!agent) { + throw new Error("Cannot find Agent: no workspace"); + } + const agentGeneration = await AgentGenerationConfiguration.findOne({ + where: { + agentId: agent.id, + }, + }); + const agentAction = await AgentRetrievalConfiguration.findOne({ + where: { + agentId: agent.id, + }, + }); + const agentDataSources = agentAction?.id + ? await AgentDataSourceConfiguration.findAll({ + where: { + retrievalConfigurationId: agentAction?.id, + }, + }) + : []; + + return { + agent: await _buildAgentConfigurationTypeFromModel({ agent }), + action: agentAction + ? await _buildAgentActionConfigurationTypeFromModel( + agentAction, + agentDataSources || [] + ) + : null, + generation: agentGeneration + ? _buildAgentGenerationConfigurationTypeFromModel(agentGeneration) + : null, + }; +} + /** * Builds the agent configuration type from the model */ -export async function _getAgentConfigurationType({ +export async function _buildAgentConfigurationTypeFromModel({ agent, - action, - generation, - dataSources, }: { agent: AgentConfiguration; - action: AgentRetrievalConfiguration | null; - generation: AgentGenerationConfiguration | null; - dataSources: AgentDataSourceConfiguration[] | null; }): Promise { return { + id: agent.id, sId: agent.sId, name: agent.name, pictureUrl: agent.pictureUrl, status: agent.status, - action: action - ? await _buildAgentActionConfigurationType(action, dataSources || []) - : null, - generation: generation - ? _buildAgentGenerationConfigurationType(generation) - : null, + }; +} + +/** + * Builds the agent generation configuration type from the model + */ +export function _buildAgentGenerationConfigurationTypeFromModel( + generation: AgentGenerationConfiguration +): AgentGenerationConfigurationType { + return { + id: generation.id, + prompt: generation.prompt, + model: { + providerId: generation.modelProvider, + modelId: generation.modelId, + }, }; } /** * Builds the agent action configuration type from the model */ -export async function _buildAgentActionConfigurationType( +export async function _buildAgentActionConfigurationTypeFromModel( action: AgentRetrievalConfiguration, dataSourcesConfig: AgentDataSourceConfiguration[] ): Promise { @@ -134,18 +193,3 @@ export async function _buildAgentActionConfigurationType( dataSources: retrievalDataSourcesConfig, }; } - -/** - * Builds the agent generation configuration type from the model - */ -export function _buildAgentGenerationConfigurationType( - generation: AgentGenerationConfiguration -): AgentGenerationConfigurationType { - return { - prompt: generation.prompt, - model: { - providerId: generation.modelProvider, - modelId: generation.modelId, - }, - }; -} diff --git a/front/lib/api/assistant/agent/agent_update.ts b/front/lib/api/assistant/agent/agent_update.ts new file mode 100644 index 000000000000..2b9eb20b98e8 --- /dev/null +++ b/front/lib/api/assistant/agent/agent_update.ts @@ -0,0 +1,208 @@ +import { + _buildAgentActionConfigurationTypeFromModel, + _buildAgentConfigurationTypeFromModel, + _buildAgentGenerationConfigurationTypeFromModel, +} from "@app/lib/api/assistant/agent/agent_get"; +import { Authenticator } from "@app/lib/auth"; +import { front_sequelize } from "@app/lib/databases"; +import { + AgentConfiguration, + AgentDataSourceConfiguration, + AgentGenerationConfiguration, + AgentRetrievalConfiguration, +} from "@app/lib/models"; +import { + isTemplatedQuery, + isTimeFrame, + RetrievalConfigurationType, + RetrievalDataSourcesConfiguration, + RetrievalQuery, + RetrievalTimeframe, +} from "@app/types/assistant/actions/retrieval"; +import { + AgentConfigurationStatus, + AgentConfigurationType, + AgentGenerationConfigurationType, +} from "@app/types/assistant/agent"; + +import { _createAgentDataSourcesConfigData } from "./agent_create"; + +/** + * Update Agent Configuration + */ +export async function updateAgentConfiguration( + auth: Authenticator, + agentSid: string, + { + name, + pictureUrl, + status, + }: { + name: string; + pictureUrl: string; + status: AgentConfigurationStatus; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Workspace not found" + ); + } + + const agentConfig = await AgentConfiguration.findOne({ + where: { + sId: agentSid, + }, + }); + if (!agentConfig) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Agent not found" + ); + } + + const updatedAgent = await agentConfig.update({ + name: name, + pictureUrl: pictureUrl, + status: status, + }); + + return _buildAgentConfigurationTypeFromModel({ + agent: updatedAgent, + }); +} + +/** + * Update Agent Generation Configuration + */ +export async function updateAgentGenerationConfiguration( + auth: Authenticator, + agentSid: string, + { + prompt, + modelProvider, + modelId, + }: { + prompt: string; + modelProvider: string; + modelId: string; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Workspace not found" + ); + } + + const agentConfig = await AgentConfiguration.findOne({ + where: { + sId: agentSid, + }, + }); + if (!agentConfig) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Agent not found" + ); + } + + const generation = await AgentGenerationConfiguration.findOne({ + where: { + agentId: agentConfig.id, + }, + }); + if (!generation) { + throw new Error( + "Cannot update AgentGenerationConfiguration: Config not found" + ); + } + + const updatedGeneration = await generation.update({ + prompt: prompt, + modelProvider: modelProvider, + modelId: modelId, + }); + + return _buildAgentGenerationConfigurationTypeFromModel(updatedGeneration); +} + +/** + * Update Agent Generation Configuration + */ +export async function updateAgentActionRetrievalConfiguration( + auth: Authenticator, + agentSid: string, + { + query, + timeframe, + topK, + dataSources, + }: { + query: RetrievalQuery; + timeframe: RetrievalTimeframe; + topK: number; + dataSources: RetrievalDataSourcesConfiguration; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error( + "Cannot create AgentActionConfiguration: Workspace not found" + ); + } + + const agentConfig = await AgentConfiguration.findOne({ + where: { + sId: agentSid, + }, + }); + if (!agentConfig) { + throw new Error("Cannot create AgentActionConfiguration: Agent not found"); + } + + const action = await AgentRetrievalConfiguration.findOne({ + where: { + agentId: agentConfig.id, + }, + }); + if (!action) { + throw new Error("Cannot update AgentActionConfiguration: Config not found"); + } + + // Updating both the Action and datasources in a single transaction + // So that we update both or none + return await front_sequelize.transaction(async (t) => { + // Update Action + const updatedAction = await action.update( + { + query: isTemplatedQuery(query) ? "templated" : query, + queryTemplate: isTemplatedQuery(query) ? query.template : null, + relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, + relativeTimeFrameDuration: isTimeFrame(timeframe) + ? timeframe.duration + : null, + relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, + topK: topK, + agentId: agentConfig.id, + }, + { transaction: t } + ); + + // Update datasources: we drop and create them all + await AgentDataSourceConfiguration.destroy({ + where: { + retrievalConfigurationId: action.id, + }, + }); + const agentDataSourcesConfigRows = await _createAgentDataSourcesConfigData( + t, + dataSources, + action.id + ); + + return _buildAgentActionConfigurationTypeFromModel( + updatedAction, + agentDataSourcesConfigRows + ); + }); +} diff --git a/front/types/assistant/agent.ts b/front/types/assistant/agent.ts index 457320076e1a..f25d0197c184 100644 --- a/front/types/assistant/agent.ts +++ b/front/types/assistant/agent.ts @@ -1,3 +1,4 @@ +import { ModelId } from "@app/lib/databases"; import { RetrievalConfigurationType } from "@app/types/assistant/actions/retrieval"; /** @@ -29,7 +30,9 @@ export type AgentActionConfigurationType = RetrievalConfigurationType; // ] // } // ``` + export type AgentActionSpecification = { + id: ModelId; name: string; description: string; inputs: { @@ -44,6 +47,7 @@ export type AgentActionSpecification = { */ export type AgentGenerationConfigurationType = { + id: ModelId; prompt: string; model: { providerId: string; @@ -59,16 +63,18 @@ export type AgentConfigurationStatus = "active" | "archived"; export type AgentConfigurationScope = "global" | "workspace"; export type AgentConfigurationType = { + id: ModelId; sId: string; status: AgentConfigurationStatus; - name: string; pictureUrl: string | null; +}; +export type AgentFullConfigurationType = { + agent: AgentConfigurationType; // If undefined, no action performed, otherwise the action is // performed (potentially NoOp eg autoSkip above). action: AgentActionConfigurationType | null; - // If undefined, no text generation. generation: AgentGenerationConfigurationType | null; }; diff --git a/front/types/assistant/conversation.ts b/front/types/assistant/conversation.ts index 14ff04663172..89b1d3cbd9a4 100644 --- a/front/types/assistant/conversation.ts +++ b/front/types/assistant/conversation.ts @@ -1,8 +1,8 @@ import { ModelId } from "@app/lib/databases"; +import { AgentFullConfigurationType } from "@app/types/assistant/agent"; import { UserType } from "@app/types/user"; import { RetrievalActionType } from "./actions/retrieval"; -import { AgentConfigurationType } from "./agent"; /** * Mentions @@ -91,7 +91,7 @@ export type AgentMessageType = { version: number; parentMessageId: string | null; - configuration: AgentConfigurationType; + configuration: AgentFullConfigurationType; status: AgentMessageStatus; action: AgentActionType | null; message: string | null; From b20d0257536ea680f8d0e6a625d09cf08274a0ca Mon Sep 17 00:00:00 2001 From: PopDaph Date: Fri, 8 Sep 2023 14:12:53 +0200 Subject: [PATCH 06/12] WIP rework everything --- front/lib/api/assistant/actions/retrieval.ts | 25 +- front/lib/api/assistant/agent.ts | 6 +- front/lib/api/assistant/agent/agent_create.ts | 283 -------- front/lib/api/assistant/agent/agent_get.ts | 195 ------ front/lib/api/assistant/agent/agent_update.ts | 208 ------ front/lib/api/assistant/configuration.ts | 618 ++++++++++++++++++ front/lib/api/assistant/conversation.ts | 111 +--- .../lib/models/assistant/actions/retrieval.ts | 192 +----- front/lib/models/assistant/agent.ts | 167 ----- front/lib/models/assistant/configuration.ts | 352 ++++++++++ front/lib/models/index.ts | 6 +- front/types/assistant/actions/retrieval.ts | 45 +- front/types/assistant/agent.ts | 48 +- front/types/assistant/configuration.ts | 66 ++ front/types/assistant/conversation.ts | 4 +- 15 files changed, 1085 insertions(+), 1241 deletions(-) delete mode 100644 front/lib/api/assistant/agent/agent_create.ts delete mode 100644 front/lib/api/assistant/agent/agent_get.ts delete mode 100644 front/lib/api/assistant/agent/agent_update.ts create mode 100644 front/lib/api/assistant/configuration.ts delete mode 100644 front/lib/models/assistant/agent.ts create mode 100644 front/lib/models/assistant/configuration.ts create mode 100644 front/types/assistant/configuration.ts diff --git a/front/lib/api/assistant/actions/retrieval.ts b/front/lib/api/assistant/actions/retrieval.ts index f73e1be1d286..6730bef1fc89 100644 --- a/front/lib/api/assistant/actions/retrieval.ts +++ b/front/lib/api/assistant/actions/retrieval.ts @@ -17,17 +17,17 @@ import { Err, Ok, Result } from "@app/lib/result"; import { new_id } from "@app/lib/utils"; import logger from "@app/logger/logger"; import { - AgentDataSourceConfigurationType, - isRetrievalConfiguration, RetrievalActionType, - RetrievalConfigurationType, RetrievalDocumentType, TimeFrame, } from "@app/types/assistant/actions/retrieval"; +import { AgentActionSpecification } from "@app/types/assistant/agent"; import { - AgentActionSpecification, - AgentFullConfigurationType, -} from "@app/types/assistant/agent"; + AgentConfigurationType, + AgentDataSourceConfigurationType, + isRetrievalConfiguration, + RetrievalConfigurationType, +} from "@app/types/assistant/configuration"; import { AgentMessageType, ConversationType, @@ -167,7 +167,6 @@ export async function retrievalActionSpecification( } return { - id: configuration.id, name: "search_data_sources", description: "Search the data sources specified by the user for information to answer their request." + @@ -316,7 +315,7 @@ export type RetrievalSuccessEvent = { // error is expected to be stored by the caller on the parent agent message. export async function* runRetrieval( auth: Authenticator, - configuration: AgentFullConfigurationType, + configuration: AgentConfigurationType, conversation: ConversationType, userMessage: UserMessageType, agentMessage: AgentMessageType @@ -349,7 +348,7 @@ export async function* runRetrieval( return yield { type: "retrieval_error", created: Date.now(), - configurationId: configuration.agent.sId, + configurationId: configuration.sId, messageId: agentMessage.sId, error: { code: "retrieval_parameters_generation_error", @@ -433,7 +432,7 @@ export async function* runRetrieval( return yield { type: "retrieval_error", created: Date.now(), - configurationId: configuration.agent.sId, + configurationId: configuration.sId, messageId: agentMessage.sId, error: { code: "retrieval_search_error", @@ -450,7 +449,7 @@ export async function* runRetrieval( return yield { type: "retrieval_error", created: Date.now(), - configurationId: configuration.agent.sId, + configurationId: configuration.sId, messageId: agentMessage.sId, error: { code: "retrieval_search_error", @@ -529,7 +528,7 @@ export async function* runRetrieval( yield { type: "retrieval_documents", created: Date.now(), - configurationId: configuration.agent.sId, + configurationId: configuration.sId, messageId: agentMessage.sId, documents, }; @@ -537,7 +536,7 @@ export async function* runRetrieval( yield { type: "retrieval_success", created: Date.now(), - configurationId: configuration.agent.sId, + configurationId: configuration.sId, messageId: agentMessage.sId, action: { id: action.id, diff --git a/front/lib/api/assistant/agent.ts b/front/lib/api/assistant/agent.ts index 7799d8dd7549..4d9891171f10 100644 --- a/front/lib/api/assistant/agent.ts +++ b/front/lib/api/assistant/agent.ts @@ -16,11 +16,11 @@ import { import { Authenticator } from "@app/lib/auth"; import { Err, Ok, Result } from "@app/lib/result"; import { generateModelSId } from "@app/lib/utils"; -import { isRetrievalConfiguration } from "@app/types/assistant/actions/retrieval"; +import { AgentActionSpecification } from "@app/types/assistant/agent"; import { - AgentActionSpecification, AgentConfigurationType, -} from "@app/types/assistant/agent"; + isRetrievalConfiguration, +} from "@app/types/assistant/configuration"; import { AgentActionType, AgentMessageType, diff --git a/front/lib/api/assistant/agent/agent_create.ts b/front/lib/api/assistant/agent/agent_create.ts deleted file mode 100644 index 45458446ceec..000000000000 --- a/front/lib/api/assistant/agent/agent_create.ts +++ /dev/null @@ -1,283 +0,0 @@ -import { Op, Transaction } from "sequelize"; - -import { - _buildAgentActionConfigurationTypeFromModel, - _buildAgentConfigurationTypeFromModel, - _buildAgentGenerationConfigurationTypeFromModel, -} from "@app/lib/api/assistant/agent/agent_get"; -import { Authenticator } from "@app/lib/auth"; -import { front_sequelize } from "@app/lib/databases"; -import { DataSource, Workspace } from "@app/lib/models"; -import { - AgentDataSourceConfiguration, - AgentRetrievalConfiguration, -} from "@app/lib/models/assistant/actions/retrieval"; -import { - AgentConfiguration, - AgentGenerationConfiguration, -} from "@app/lib/models/assistant/agent"; -import { generateModelSId } from "@app/lib/utils"; -import { - AgentDataSourceConfigurationType, - isTemplatedQuery, - isTimeFrame, - RetrievalDataSourcesConfiguration, - RetrievalQuery, - RetrievalTimeframe, -} from "@app/types/assistant/actions/retrieval"; -import { - AgentActionConfigurationType, - AgentConfigurationStatus, - AgentConfigurationType, - AgentGenerationConfigurationType, -} from "@app/types/assistant/agent"; - -/** - * Create Agent Configuration - */ -export async function createAgentConfiguration( - auth: Authenticator, - { - name, - pictureUrl, - status, - }: { - name: string; - pictureUrl: string; - status: AgentConfigurationStatus; - } -): Promise { - const owner = auth.workspace(); - if (!owner) { - throw new Error("Cannot create AgentConfiguration without workspace"); - } - - const agentConfig = await AgentConfiguration.create({ - sId: generateModelSId(), - status: status, - name: name, - pictureUrl: pictureUrl, - scope: "workspace", - workspaceId: owner.id, - }); - - return _buildAgentConfigurationTypeFromModel({ - agent: agentConfig, - }); -} - -/** - * Create Agent Generation Configuration - */ -export async function createAgentGenerationConfiguration( - auth: Authenticator, - agentSid: string, - { - prompt, - modelProvider, - modelId, - }: { - prompt: string; - modelProvider: string; - modelId: string; - } -): Promise { - const owner = auth.workspace(); - if (!owner) { - throw new Error( - "Cannot create AgentGenerationConfiguration: Workspace not found" - ); - } - - const agentConfig = await AgentConfiguration.findOne({ - where: { - sId: agentSid, - }, - }); - if (!agentConfig) { - throw new Error( - "Cannot create AgentGenerationConfiguration: Agent not found" - ); - } - - const generation = await AgentGenerationConfiguration.create({ - prompt: prompt, - modelProvider: modelProvider, - modelId: modelId, - agentId: agentConfig.id, - }); - - return _buildAgentGenerationConfigurationTypeFromModel(generation); -} - -/** - * Create Agent Action Configuration (Retrieval) - */ -export async function createAgentActionRetrievalConfiguration( - auth: Authenticator, - agentSid: string, - { - query, - timeframe, - topK, - dataSources, - }: { - query: RetrievalQuery; - timeframe: RetrievalTimeframe; - topK: number; - dataSources: RetrievalDataSourcesConfiguration; - } -): Promise { - const owner = auth.workspace(); - if (!owner) { - throw new Error( - "Cannot create AgentActionConfiguration: Workspace not found" - ); - } - - const agentConfig = await AgentConfiguration.findOne({ - where: { - sId: agentSid, - }, - }); - if (!agentConfig) { - throw new Error("Cannot create AgentActionConfiguration: Agent not found"); - } - return await front_sequelize.transaction(async (t) => { - const agentActionConfigRow = await AgentRetrievalConfiguration.create( - { - query: isTemplatedQuery(query) ? "templated" : query, - queryTemplate: isTemplatedQuery(query) ? query.template : null, - relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, - relativeTimeFrameDuration: isTimeFrame(timeframe) - ? timeframe.duration - : null, - relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, - topK: topK, - agentId: agentConfig.id, - }, - { transaction: t } - ); - const agentDataSourcesConfigRows = await _createAgentDataSourcesConfigData( - t, - dataSources, - agentActionConfigRow.id - ); - return await _buildAgentActionConfigurationTypeFromModel( - agentActionConfigRow, - agentDataSourcesConfigRows - ); - }); -} - -/** - * Create the AgentDataSourceConfiguration rows in database. - * - * Knowing that a datasource is uniquely identified by its name and its workspaceId - * We need to fetch the dataSources from the database from that. - * We obvisously need to do as few queries as possible. - */ -export async function _createAgentDataSourcesConfigData( - t: Transaction, - dataSourcesConfig: AgentDataSourceConfigurationType[], - agentActionId: number -): Promise { - // dsConfig contains this format: - // [ - // { workspaceSId: s1o1u1p, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, - // { workspaceSId: s1o1u1p, dataSourceName: "managed-slack", filter: { tags: null, parents: null } }, - // { workspaceSId: i2n2o2u, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, - // ] - - // First we get the list of workspaces because we need the mapping between workspaceSId and workspaceId - const workspaces = await Workspace.findAll({ - where: { - sId: dataSourcesConfig.map((dsConfig) => dsConfig.workspaceSId), - }, - attributes: ["id", "sId"], - }); - - // Now will want to group the datasource names by workspaceId to do only one query per workspace. - // We want this: - // [ - // { workspaceId: 1, dataSourceNames: [""managed-notion", "managed-slack"] }, - // { workspaceId: 2, dataSourceNames: ["managed-notion"] } - // ] - type _DsNamesPerWorkspaceIdType = { - workspaceId: number; - dataSourceNames: string[]; - }; - const dsNamesPerWorkspaceId = dataSourcesConfig.reduce( - ( - acc: _DsNamesPerWorkspaceIdType[], - curr: AgentDataSourceConfigurationType - ) => { - // First we need to get the workspaceId from the workspaceSId - const workspace = workspaces.find((w) => w.sId === curr.workspaceSId); - if (!workspace) { - throw new Error("Workspace not found"); - } - - // Find an existing entry for this workspaceId - const existingEntry: _DsNamesPerWorkspaceIdType | undefined = acc.find( - (entry: _DsNamesPerWorkspaceIdType) => - entry.workspaceId === workspace.id - ); - if (existingEntry) { - // Append dataSourceName to existing entry - existingEntry.dataSourceNames.push(curr.dataSourceName); - } else { - // Add a new entry for this workspaceId - acc.push({ - workspaceId: workspace.id, - dataSourceNames: [curr.dataSourceName], - }); - } - return acc; - }, - [] - ); - - // Then we get do one findAllQuery per workspaceId, in a Promise.all - const getDataSourcesQueries = dsNamesPerWorkspaceId.map( - ({ workspaceId, dataSourceNames }) => { - return DataSource.findAll({ - where: { - workspaceId, - name: { - [Op.in]: dataSourceNames, - }, - }, - }); - } - ); - const results = await Promise.all(getDataSourcesQueries); - const dataSources = results.flat(); - - const agentDataSourcesConfigRows: AgentDataSourceConfiguration[] = - await Promise.all( - dataSourcesConfig.map(async (dsConfig) => { - const dataSource = dataSources.find( - (ds) => - ds.name === dsConfig.dataSourceName && - ds.workspaceId === - workspaces.find((w) => w.sId === dsConfig.workspaceSId)?.id - ); - if (!dataSource) { - throw new Error("DataSource not found"); - } - return AgentDataSourceConfiguration.create( - { - dataSourceId: dataSource.id, - tagsIn: dsConfig.filter.tags?.in, - tagsNotIn: dsConfig.filter.tags?.not, - parentsIn: dsConfig.filter.parents?.in, - parentsNotIn: dsConfig.filter.parents?.not, - retrievalConfigurationId: agentActionId, - }, - { transaction: t } - ); - }) - ); - return agentDataSourcesConfigRows; -} diff --git a/front/lib/api/assistant/agent/agent_get.ts b/front/lib/api/assistant/agent/agent_get.ts deleted file mode 100644 index 360d99ef0ba4..000000000000 --- a/front/lib/api/assistant/agent/agent_get.ts +++ /dev/null @@ -1,195 +0,0 @@ -import { Op } from "sequelize"; - -import { Authenticator } from "@app/lib/auth"; -import { DataSource, Workspace } from "@app/lib/models"; -import { - AgentDataSourceConfiguration, - AgentRetrievalConfiguration, -} from "@app/lib/models/assistant/actions/retrieval"; -import { - AgentConfiguration, - AgentGenerationConfiguration, -} from "@app/lib/models/assistant/agent"; -import { - RetrievalDataSourcesConfiguration, - RetrievalQuery, - RetrievalTimeframe, -} from "@app/types/assistant/actions/retrieval"; -import { - AgentActionConfigurationType, - AgentConfigurationType, - AgentFullConfigurationType as AgentFullConfigurationType, - AgentGenerationConfigurationType, -} from "@app/types/assistant/agent"; - -/** - * Get an agent full configuration from its name - */ -export async function getAgent( - auth: Authenticator, - sId: string -): Promise { - const owner = auth.workspace(); - if (!owner) { - throw new Error("Cannot find Agent: no workspace"); - } - const agent = await AgentConfiguration.findOne({ - where: { - sId: sId, - workspaceId: owner.id, - }, - }); - if (!agent) { - throw new Error("Cannot find Agent: no workspace"); - } - const agentGeneration = await AgentGenerationConfiguration.findOne({ - where: { - agentId: agent.id, - }, - }); - const agentAction = await AgentRetrievalConfiguration.findOne({ - where: { - agentId: agent.id, - }, - }); - const agentDataSources = agentAction?.id - ? await AgentDataSourceConfiguration.findAll({ - where: { - retrievalConfigurationId: agentAction?.id, - }, - }) - : []; - - return { - agent: await _buildAgentConfigurationTypeFromModel({ agent }), - action: agentAction - ? await _buildAgentActionConfigurationTypeFromModel( - agentAction, - agentDataSources || [] - ) - : null, - generation: agentGeneration - ? _buildAgentGenerationConfigurationTypeFromModel(agentGeneration) - : null, - }; -} - -/** - * Builds the agent configuration type from the model - */ -export async function _buildAgentConfigurationTypeFromModel({ - agent, -}: { - agent: AgentConfiguration; -}): Promise { - return { - id: agent.id, - sId: agent.sId, - name: agent.name, - pictureUrl: agent.pictureUrl, - status: agent.status, - }; -} - -/** - * Builds the agent generation configuration type from the model - */ -export function _buildAgentGenerationConfigurationTypeFromModel( - generation: AgentGenerationConfiguration -): AgentGenerationConfigurationType { - return { - id: generation.id, - prompt: generation.prompt, - model: { - providerId: generation.modelProvider, - modelId: generation.modelId, - }, - }; -} - -/** - * Builds the agent action configuration type from the model - */ -export async function _buildAgentActionConfigurationTypeFromModel( - action: AgentRetrievalConfiguration, - dataSourcesConfig: AgentDataSourceConfiguration[] -): Promise { - // Build Retrieval Timeframe - let timeframe: RetrievalTimeframe = "auto"; - if ( - action.relativeTimeFrame === "custom" && - action.relativeTimeFrameDuration && - action.relativeTimeFrameUnit - ) { - timeframe = { - duration: action.relativeTimeFrameDuration, - unit: action.relativeTimeFrameUnit, - }; - } else if (action.relativeTimeFrame === "none") { - timeframe = "none"; - } - - // Build Retrieval Query - let query: RetrievalQuery = "auto"; - if (action.query === "templated" && action.queryTemplate) { - query = { - template: action.queryTemplate, - }; - } else if (action.query === "none") { - query = "none"; - } - - // Build Retrieval DataSources - const retrievalDataSourcesConfig: RetrievalDataSourcesConfiguration = []; - - const dataSourcesIds = dataSourcesConfig?.map((ds) => ds.dataSourceId); - const dataSources = await DataSource.findAll({ - where: { - id: { [Op.in]: dataSourcesIds }, - }, - attributes: ["id", "name", "workspaceId"], - }); - const workspaceIds = dataSources.map((ds) => ds.workspaceId); - const workspaces = await Workspace.findAll({ - where: { - id: { [Op.in]: workspaceIds }, - }, - attributes: ["id", "sId"], - }); - - let dataSource: DataSource | undefined; - let workspace: Workspace | undefined; - - dataSourcesConfig.forEach(async (dsConfig) => { - dataSource = dataSources.find((ds) => ds.id === dsConfig.dataSourceId); - workspace = workspaces.find((w) => w.id === dataSource?.workspaceId); - - if (!dataSource || !workspace) { - throw new Error("Could not find dataSource or workspace"); - } - - retrievalDataSourcesConfig.push({ - dataSourceName: dataSource.name, - workspaceSId: workspace.sId, - filter: { - tags: - dsConfig.tagsIn && dsConfig.tagsNotIn - ? { in: dsConfig.tagsIn, not: dsConfig.tagsNotIn } - : null, - parents: - dsConfig.parentsIn && dsConfig.parentsNotIn - ? { in: dsConfig.parentsIn, not: dsConfig.parentsNotIn } - : null, - }, - }); - }); - - return { - id: action.id, - query: query, - relativeTimeFrame: timeframe, - topK: action.topK, - type: "retrieval_configuration", - dataSources: retrievalDataSourcesConfig, - }; -} diff --git a/front/lib/api/assistant/agent/agent_update.ts b/front/lib/api/assistant/agent/agent_update.ts deleted file mode 100644 index 2b9eb20b98e8..000000000000 --- a/front/lib/api/assistant/agent/agent_update.ts +++ /dev/null @@ -1,208 +0,0 @@ -import { - _buildAgentActionConfigurationTypeFromModel, - _buildAgentConfigurationTypeFromModel, - _buildAgentGenerationConfigurationTypeFromModel, -} from "@app/lib/api/assistant/agent/agent_get"; -import { Authenticator } from "@app/lib/auth"; -import { front_sequelize } from "@app/lib/databases"; -import { - AgentConfiguration, - AgentDataSourceConfiguration, - AgentGenerationConfiguration, - AgentRetrievalConfiguration, -} from "@app/lib/models"; -import { - isTemplatedQuery, - isTimeFrame, - RetrievalConfigurationType, - RetrievalDataSourcesConfiguration, - RetrievalQuery, - RetrievalTimeframe, -} from "@app/types/assistant/actions/retrieval"; -import { - AgentConfigurationStatus, - AgentConfigurationType, - AgentGenerationConfigurationType, -} from "@app/types/assistant/agent"; - -import { _createAgentDataSourcesConfigData } from "./agent_create"; - -/** - * Update Agent Configuration - */ -export async function updateAgentConfiguration( - auth: Authenticator, - agentSid: string, - { - name, - pictureUrl, - status, - }: { - name: string; - pictureUrl: string; - status: AgentConfigurationStatus; - } -): Promise { - const owner = auth.workspace(); - if (!owner) { - throw new Error( - "Cannot create AgentGenerationConfiguration: Workspace not found" - ); - } - - const agentConfig = await AgentConfiguration.findOne({ - where: { - sId: agentSid, - }, - }); - if (!agentConfig) { - throw new Error( - "Cannot create AgentGenerationConfiguration: Agent not found" - ); - } - - const updatedAgent = await agentConfig.update({ - name: name, - pictureUrl: pictureUrl, - status: status, - }); - - return _buildAgentConfigurationTypeFromModel({ - agent: updatedAgent, - }); -} - -/** - * Update Agent Generation Configuration - */ -export async function updateAgentGenerationConfiguration( - auth: Authenticator, - agentSid: string, - { - prompt, - modelProvider, - modelId, - }: { - prompt: string; - modelProvider: string; - modelId: string; - } -): Promise { - const owner = auth.workspace(); - if (!owner) { - throw new Error( - "Cannot create AgentGenerationConfiguration: Workspace not found" - ); - } - - const agentConfig = await AgentConfiguration.findOne({ - where: { - sId: agentSid, - }, - }); - if (!agentConfig) { - throw new Error( - "Cannot create AgentGenerationConfiguration: Agent not found" - ); - } - - const generation = await AgentGenerationConfiguration.findOne({ - where: { - agentId: agentConfig.id, - }, - }); - if (!generation) { - throw new Error( - "Cannot update AgentGenerationConfiguration: Config not found" - ); - } - - const updatedGeneration = await generation.update({ - prompt: prompt, - modelProvider: modelProvider, - modelId: modelId, - }); - - return _buildAgentGenerationConfigurationTypeFromModel(updatedGeneration); -} - -/** - * Update Agent Generation Configuration - */ -export async function updateAgentActionRetrievalConfiguration( - auth: Authenticator, - agentSid: string, - { - query, - timeframe, - topK, - dataSources, - }: { - query: RetrievalQuery; - timeframe: RetrievalTimeframe; - topK: number; - dataSources: RetrievalDataSourcesConfiguration; - } -): Promise { - const owner = auth.workspace(); - if (!owner) { - throw new Error( - "Cannot create AgentActionConfiguration: Workspace not found" - ); - } - - const agentConfig = await AgentConfiguration.findOne({ - where: { - sId: agentSid, - }, - }); - if (!agentConfig) { - throw new Error("Cannot create AgentActionConfiguration: Agent not found"); - } - - const action = await AgentRetrievalConfiguration.findOne({ - where: { - agentId: agentConfig.id, - }, - }); - if (!action) { - throw new Error("Cannot update AgentActionConfiguration: Config not found"); - } - - // Updating both the Action and datasources in a single transaction - // So that we update both or none - return await front_sequelize.transaction(async (t) => { - // Update Action - const updatedAction = await action.update( - { - query: isTemplatedQuery(query) ? "templated" : query, - queryTemplate: isTemplatedQuery(query) ? query.template : null, - relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, - relativeTimeFrameDuration: isTimeFrame(timeframe) - ? timeframe.duration - : null, - relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, - topK: topK, - agentId: agentConfig.id, - }, - { transaction: t } - ); - - // Update datasources: we drop and create them all - await AgentDataSourceConfiguration.destroy({ - where: { - retrievalConfigurationId: action.id, - }, - }); - const agentDataSourcesConfigRows = await _createAgentDataSourcesConfigData( - t, - dataSources, - action.id - ); - - return _buildAgentActionConfigurationTypeFromModel( - updatedAction, - agentDataSourcesConfigRows - ); - }); -} diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts new file mode 100644 index 000000000000..a6a74f62c3e1 --- /dev/null +++ b/front/lib/api/assistant/configuration.ts @@ -0,0 +1,618 @@ +import { Op, Transaction } from "sequelize"; + +import { Authenticator } from "@app/lib/auth"; +import { front_sequelize } from "@app/lib/databases"; +import { + AgentConfiguration, + AgentDataSourceConfiguration, + AgentGenerationConfiguration, + AgentRetrievalConfiguration, + DataSource, + Workspace, +} from "@app/lib/models"; +import { generateModelSId } from "@app/lib/utils"; +import { + isTemplatedQuery, + isTimeFrame, + RetrievalQuery, + RetrievalTimeframe, +} from "@app/types/assistant/actions/retrieval"; +import { + AgentActionConfigurationType, + AgentConfigurationStatus, + AgentConfigurationType, + AgentDataSourceConfigurationType, +} from "@app/types/assistant/configuration"; + +/** + * Get an agent configuration + */ +export async function getAgentConfiguration( + auth: Authenticator, + agentId: string +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error("Cannot find Agent: no workspace"); + } + const agent = await AgentConfiguration.findOne({ + where: { + sId: agentId, + workspaceId: owner.id, + }, + }); + if (!agent) { + throw new Error("Cannot find Agent: no workspace"); + } + + const generation = agent.generationId + ? await AgentGenerationConfiguration.findOne({ + where: { + id: agent.generationId, + }, + }) + : null; + + const action = agent.retrievalId + ? await AgentRetrievalConfiguration.findOne({ + where: { + id: agent.retrievalId, + }, + }) + : null; + const datasources = action?.id + ? await AgentDataSourceConfiguration.findAll({ + where: { + retrievalConfigurationId: action.id, + }, + }) + : []; + + return { + sId: agent.sId, + name: agent.name, + pictureUrl: agent.pictureUrl, + status: agent.status, + action: action ? await _agentActionType(action, datasources) : null, + generation: generation + ? { + id: generation.id, + prompt: generation.prompt, + model: { + providerId: generation.providerId, + modelId: generation.modelId, + }, + } + : null, + }; +} + +/** + * Create Agent Configuration + */ +export async function createAgentConfiguration( + auth: Authenticator, + { + name, + pictureUrl, + status, + generation, + action, + }: { + name: string; + pictureUrl: string; + status: AgentConfigurationStatus; + generation: { + prompt: string; + model: { + providerId: string; + modelId: string; + }; + } | null; + action: { + type: string; + query: RetrievalQuery; + timeframe: RetrievalTimeframe; + topK: number; + dataSources: AgentDataSourceConfigurationType[]; + } | null; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error("Cannot create AgentConfiguration without workspace"); + } + + return await front_sequelize.transaction(async (t) => { + let genConfig: AgentGenerationConfiguration | null = null; + let retrievalConfig: AgentRetrievalConfiguration | null = null; + let dataSourcesConfig: AgentDataSourceConfiguration[] = []; + + // Create Generation config + if (generation) { + const { prompt, model } = generation; + genConfig = await AgentGenerationConfiguration.create({ + prompt: prompt, + providerId: model.providerId, + modelId: model.modelId, + }); + } + + // Create Retrieval & Datasources configs + if (action && action.type === "retrieval_configuration") { + const { query, timeframe, topK, dataSources } = action; + retrievalConfig = await AgentRetrievalConfiguration.create( + { + query: isTemplatedQuery(query) ? "templated" : query, + queryTemplate: isTemplatedQuery(query) ? query.template : null, + relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, + relativeTimeFrameDuration: isTimeFrame(timeframe) + ? timeframe.duration + : null, + relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, + topK: topK, + }, + { transaction: t } + ); + dataSourcesConfig = await _createAgentDataSourcesConfigData( + t, + dataSources, + retrievalConfig.id + ); + } + + // Create Agent config + const agentConfig = await AgentConfiguration.create({ + sId: generateModelSId(), + status: status, + name: name, + pictureUrl: pictureUrl, + scope: "workspace", + workspaceId: owner.id, + generationId: genConfig?.id ?? null, + retrievalId: retrievalConfig?.id ?? null, + }); + + return { + sId: agentConfig.sId, + name: agentConfig.name, + pictureUrl: agentConfig.pictureUrl, + status: agentConfig.status, + action: retrievalConfig + ? await _agentActionType(retrievalConfig, dataSourcesConfig) + : null, + generation: genConfig + ? { + id: genConfig.id, + prompt: genConfig.prompt, + model: { + providerId: genConfig.providerId, + modelId: genConfig.modelId, + }, + } + : null, + }; + }); +} + +/** + * Update Agent Generation Configuration + */ +export async function updateAgentGenerationConfiguration( + auth: Authenticator, + agentId: string, + { + name, + pictureUrl, + status, + generation, + }: { + name: string; + pictureUrl: string; + status: AgentConfigurationStatus; + generation: { + prompt: string; + model: { + providerId: string; + modelId: string; + }; + } | null; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Workspace not found" + ); + } + const agentConfig = await AgentConfiguration.findOne({ + where: { + sId: agentId, + }, + }); + if (!agentConfig) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Agent not found" + ); + } + const existingGeneration = agentConfig.generationId + ? await AgentGenerationConfiguration.findOne({ + where: { + id: agentConfig.generationId, + }, + }) + : null; + + const existingRetrivalConfig = agentConfig.retrievalId + ? await AgentRetrievalConfiguration.findOne({ + where: { + id: agentConfig.retrievalId, + }, + }) + : null; + + const existingDataSourcesConfig = existingRetrivalConfig?.id + ? await AgentDataSourceConfiguration.findAll({ + where: { + retrievalConfigurationId: existingRetrivalConfig.id, + }, + }) + : []; + + return await front_sequelize.transaction(async (t) => { + // Upserting Agent Config + const updatedAgentConfig = await agentConfig.update( + { + name: name, + pictureUrl: pictureUrl, + status: status, + }, + { transaction: t } + ); + + // Upserting Generation Config + let upsertedGenerationConfig: AgentGenerationConfiguration | null = null; + if (generation) { + const { prompt, model } = generation; + if (existingGeneration) { + upsertedGenerationConfig = await existingGeneration.update( + { + prompt: prompt, + providerId: model.providerId, + modelId: model.modelId, + }, + { transaction: t } + ); + } else { + upsertedGenerationConfig = await AgentGenerationConfiguration.create( + { + prompt: prompt, + providerId: model.providerId, + modelId: model.modelId, + }, + { transaction: t } + ); + } + } else if (existingGeneration) { + await existingGeneration.destroy(); + } + + return { + sId: updatedAgentConfig.sId, + name: updatedAgentConfig.name, + pictureUrl: updatedAgentConfig.pictureUrl, + status: updatedAgentConfig.status, + action: existingRetrivalConfig + ? await _agentActionType( + existingRetrivalConfig, + existingDataSourcesConfig + ) + : null, + generation: + generation && upsertedGenerationConfig + ? { + id: upsertedGenerationConfig.id, + prompt: upsertedGenerationConfig.prompt, + model: { + providerId: upsertedGenerationConfig.providerId, + modelId: upsertedGenerationConfig.modelId, + }, + } + : null, + }; + }); +} + +/** + * Update Agent Retrieval Configuration + * This will destroy and recreate the retrieval config + */ +export async function updateAgentRetrievalConfiguration( + auth: Authenticator, + agentId: string, + { + query, + timeframe, + topK, + dataSources, + }: { + query: RetrievalQuery; + timeframe: RetrievalTimeframe; + topK: number; + dataSources: AgentDataSourceConfigurationType[]; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Workspace not found" + ); + } + const agentConfig = await AgentConfiguration.findOne({ + where: { + sId: agentId, + }, + }); + if (!agentConfig) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Agent not found" + ); + } + const generationConfig = agentConfig.generationId + ? await AgentGenerationConfiguration.findOne({ + where: { + id: agentConfig.generationId, + }, + }) + : null; + + return await front_sequelize.transaction(async (t) => { + if (agentConfig.retrievalId) { + const existingRetrivalConfig = await AgentRetrievalConfiguration.findOne({ + where: { + id: agentConfig.retrievalId, + }, + }); + if (existingRetrivalConfig) { + await existingRetrivalConfig.destroy(); // That will destroy the dataSourcesConfig too + } + } + + const newRetrievalConfig = await AgentRetrievalConfiguration.create( + { + query: isTemplatedQuery(query) ? "templated" : query, + queryTemplate: isTemplatedQuery(query) ? query.template : null, + relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, + relativeTimeFrameDuration: isTimeFrame(timeframe) + ? timeframe.duration + : null, + relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, + topK: topK, + }, + { transaction: t } + ); + const dataSourcesConfig = await _createAgentDataSourcesConfigData( + t, + dataSources, + newRetrievalConfig.id + ); + + return { + sId: agentConfig.sId, + name: agentConfig.name, + pictureUrl: agentConfig.pictureUrl, + status: agentConfig.status, + action: newRetrievalConfig + ? await _agentActionType(newRetrievalConfig, dataSourcesConfig) + : null, + generation: generationConfig + ? { + id: generationConfig.id, + prompt: generationConfig.prompt, + model: { + providerId: generationConfig.providerId, + modelId: generationConfig.modelId, + }, + } + : null, + }; + }); +} + +/** + * Builds the agent action configuration type from the model + */ +export async function _agentActionType( + action: AgentRetrievalConfiguration, + dataSourcesConfig: AgentDataSourceConfiguration[] +): Promise { + // Build Retrieval Timeframe + let timeframe: RetrievalTimeframe = "auto"; + if ( + action.relativeTimeFrame === "custom" && + action.relativeTimeFrameDuration && + action.relativeTimeFrameUnit + ) { + timeframe = { + duration: action.relativeTimeFrameDuration, + unit: action.relativeTimeFrameUnit, + }; + } else if (action.relativeTimeFrame === "none") { + timeframe = "none"; + } + + // Build Retrieval Query + let query: RetrievalQuery = "auto"; + if (action.query === "templated" && action.queryTemplate) { + query = { + template: action.queryTemplate, + }; + } else if (action.query === "none") { + query = "none"; + } + + // Build Retrieval DataSources + const dataSourcesIds = dataSourcesConfig?.map((ds) => ds.dataSourceId); + const dataSources = await DataSource.findAll({ + where: { + id: { [Op.in]: dataSourcesIds }, + }, + attributes: ["id", "name", "workspaceId"], + }); + const workspaceIds = dataSources.map((ds) => ds.workspaceId); + const workspaces = await Workspace.findAll({ + where: { + id: { [Op.in]: workspaceIds }, + }, + attributes: ["id", "sId"], + }); + + let dataSource: DataSource | undefined; + let workspace: Workspace | undefined; + const dataSourcesConfigType: AgentDataSourceConfigurationType[] = []; + + dataSourcesConfig.forEach(async (dsConfig) => { + dataSource = dataSources.find((ds) => ds.id === dsConfig.dataSourceId); + workspace = workspaces.find((w) => w.id === dataSource?.workspaceId); + + if (!dataSource || !workspace) { + throw new Error("Could not find dataSource or workspace"); + } + + dataSourcesConfigType.push({ + dataSourceName: dataSource.name, + workspaceSId: workspace.sId, + filter: { + tags: + dsConfig.tagsIn && dsConfig.tagsNotIn + ? { in: dsConfig.tagsIn, not: dsConfig.tagsNotIn } + : null, + parents: + dsConfig.parentsIn && dsConfig.parentsNotIn + ? { in: dsConfig.parentsIn, not: dsConfig.parentsNotIn } + : null, + }, + }); + }); + + return { + id: action.id, + type: "retrieval_configuration", + query: query, + relativeTimeFrame: timeframe, + topK: action.topK, + dataSources: dataSourcesConfigType, + }; +} + +/** + * Create the AgentDataSourceConfiguration rows in database. + * + * Knowing that a datasource is uniquely identified by its name and its workspaceId + * We need to fetch the dataSources from the database from that. + * We obvisously need to do as few queries as possible. + */ +export async function _createAgentDataSourcesConfigData( + t: Transaction, + dataSourcesConfig: AgentDataSourceConfigurationType[], + agentActionId: number +): Promise { + // dsConfig contains this format: + // [ + // { workspaceSId: s1o1u1p, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, + // { workspaceSId: s1o1u1p, dataSourceName: "managed-slack", filter: { tags: null, parents: null } }, + // { workspaceSId: i2n2o2u, dataSourceName: "managed-notion", filter: { tags: null, parents: null } }, + // ] + + // First we get the list of workspaces because we need the mapping between workspaceSId and workspaceId + const workspaces = await Workspace.findAll({ + where: { + sId: dataSourcesConfig.map((dsConfig) => dsConfig.workspaceSId), + }, + attributes: ["id", "sId"], + }); + + // Now will want to group the datasource names by workspaceId to do only one query per workspace. + // We want this: + // [ + // { workspaceId: 1, dataSourceNames: [""managed-notion", "managed-slack"] }, + // { workspaceId: 2, dataSourceNames: ["managed-notion"] } + // ] + type _DsNamesPerWorkspaceIdType = { + workspaceId: number; + dataSourceNames: string[]; + }; + const dsNamesPerWorkspaceId = dataSourcesConfig.reduce( + ( + acc: _DsNamesPerWorkspaceIdType[], + curr: AgentDataSourceConfigurationType + ) => { + // First we need to get the workspaceId from the workspaceSId + const workspace = workspaces.find((w) => w.sId === curr.workspaceSId); + if (!workspace) { + throw new Error("Workspace not found"); + } + + // Find an existing entry for this workspaceId + const existingEntry: _DsNamesPerWorkspaceIdType | undefined = acc.find( + (entry: _DsNamesPerWorkspaceIdType) => + entry.workspaceId === workspace.id + ); + if (existingEntry) { + // Append dataSourceName to existing entry + existingEntry.dataSourceNames.push(curr.dataSourceName); + } else { + // Add a new entry for this workspaceId + acc.push({ + workspaceId: workspace.id, + dataSourceNames: [curr.dataSourceName], + }); + } + return acc; + }, + [] + ); + + // Then we get do one findAllQuery per workspaceId, in a Promise.all + const getDataSourcesQueries = dsNamesPerWorkspaceId.map( + ({ workspaceId, dataSourceNames }) => { + return DataSource.findAll({ + where: { + workspaceId, + name: { + [Op.in]: dataSourceNames, + }, + }, + }); + } + ); + const results = await Promise.all(getDataSourcesQueries); + const dataSources = results.flat(); + + const agentDataSourcesConfigRows: AgentDataSourceConfiguration[] = + await Promise.all( + dataSourcesConfig.map(async (dsConfig) => { + const dataSource = dataSources.find( + (ds) => + ds.name === dsConfig.dataSourceName && + ds.workspaceId === + workspaces.find((w) => w.sId === dsConfig.workspaceSId)?.id + ); + if (!dataSource) { + throw new Error("DataSource not found"); + } + return AgentDataSourceConfiguration.create( + { + dataSourceId: dataSource.id, + tagsIn: dsConfig.filter.tags?.in, + tagsNotIn: dsConfig.filter.tags?.not, + parentsIn: dsConfig.filter.parents?.in, + parentsNotIn: dsConfig.filter.parents?.not, + retrievalConfigurationId: agentActionId, + }, + { transaction: t } + ); + }) + ); + return agentDataSourcesConfigRows; +} diff --git a/front/lib/api/assistant/conversation.ts b/front/lib/api/assistant/conversation.ts index b89e253484d3..1dd3da5926ef 100644 --- a/front/lib/api/assistant/conversation.ts +++ b/front/lib/api/assistant/conversation.ts @@ -75,92 +75,31 @@ export async function* postUserMessage( where: { conversationId: conversation.id, }, - transaction: t, - })) ?? -1) + 1; - - const m = await Message.create( - { - sId: generateModelSId(), - rank: nextMessageRank++, - conversationId: conversation.id, - parentId: null, - userMessageId: ( - await UserMessage.create( - { - message: message, - userContextUsername: context.username, - userContextTimezone: context.timezone, - userContextFullName: context.fullName, - userContextEmail: context.email, - userContextProfilePictureUrl: context.profilePictureUrl, - userId: user ? user.id : null, - }, - { transaction: t } - ) - ).id, - }, - { - transaction: t, - } - ); - - const userMessage: UserMessageType = { - id: m.id, - sId: m.sId, - type: "user_message", - visibility: "visible", - version: 0, - user: user, - mentions: mentions, - message: message, - context: context, - }; - - const agentMessages: AgentMessageType[] = []; - const agentMessageRows: AgentMessage[] = []; - - // for each assistant mention, create an "empty" agent message - for (const mention of mentions) { - if (isAgentMention(mention)) { - const agentMessageRow = await AgentMessage.create( - {}, - { transaction: t } - ); - const m = await Message.create( - { - sId: generateModelSId(), - rank: nextMessageRank++, - conversationId: conversation.id, - parentId: userMessage.id, - agentMessageId: agentMessageRow.id, - }, - { - transaction: t, - } - ); - agentMessageRows.push(agentMessageRow); - agentMessages.push({ - id: m.id, - sId: m.sId, - type: "agent_message", - visibility: "visible", - version: 0, - parentMessageId: userMessage.sId, - status: "created", - action: null, - message: null, - feedbacks: [], - error: null, - configuration: { - sId: mention.configurationId, - status: "active", - name: "foo", // TODO - pictureUrl: null, // TODO - action: null, // TODO - generation: null, // TODO - }, - }); - } + { + transaction: t, + } + ); + agentMessages.push({ + id: agentMessageRow.id, + sId: agentMessageRow.sId, + type: "agent_message", + visibility: "visible", + version: 0, + parentMessageId: userMessage.sId, + status: "created", + action: null, + message: null, + feedbacks: [], + error: null, + configuration: { + sId: m.configurationId, + status: "active", + name: "foo", // TODO + pictureUrl: null, // TODO + action: null, // TODO + generation: null, // TODO + }, + }); } return { userMessage, agentMessages, agentMessageRows }; diff --git a/front/lib/models/assistant/actions/retrieval.ts b/front/lib/models/assistant/actions/retrieval.ts index 95adea5566fd..6394c461465c 100644 --- a/front/lib/models/assistant/actions/retrieval.ts +++ b/front/lib/models/assistant/actions/retrieval.ts @@ -8,199 +8,9 @@ import { } from "sequelize"; import { front_sequelize } from "@app/lib/databases"; -import { AgentConfiguration } from "@app/lib/models/assistant/agent"; -import { DataSource } from "@app/lib/models/data_source"; +import { AgentRetrievalConfiguration } from "@app/lib/models/assistant/configuration"; import { TimeframeUnit } from "@app/types/assistant/actions/retrieval"; -/** - * Action Retrieval configuration - */ -export class AgentRetrievalConfiguration extends Model< - InferAttributes, - InferCreationAttributes -> { - declare id: CreationOptional; - declare createdAt: CreationOptional; - declare updatedAt: CreationOptional; - - declare query: "auto" | "none" | "templated"; - declare queryTemplate: string | null; - declare relativeTimeFrame: "auto" | "none" | "custom"; - declare relativeTimeFrameDuration: number | null; - declare relativeTimeFrameUnit: TimeframeUnit | null; - declare topK: number; - - declare agentId: ForeignKey; -} -AgentRetrievalConfiguration.init( - { - id: { - type: DataTypes.INTEGER, - autoIncrement: true, - primaryKey: true, - }, - createdAt: { - type: DataTypes.DATE, - allowNull: false, - defaultValue: DataTypes.NOW, - }, - updatedAt: { - type: DataTypes.DATE, - allowNull: false, - defaultValue: DataTypes.NOW, - }, - query: { - type: DataTypes.STRING, - allowNull: false, - defaultValue: "auto", - }, - queryTemplate: { - type: DataTypes.TEXT, - allowNull: true, - }, - relativeTimeFrame: { - type: DataTypes.STRING, - allowNull: false, - defaultValue: "auto", - }, - relativeTimeFrameDuration: { - type: DataTypes.INTEGER, - allowNull: true, - }, - relativeTimeFrameUnit: { - type: DataTypes.STRING, - allowNull: true, - }, - topK: { - type: DataTypes.INTEGER, - allowNull: false, - }, - }, - { - modelName: "agent_retrieval_configuration", - sequelize: front_sequelize, - hooks: { - beforeValidate: (retrieval: AgentRetrievalConfiguration) => { - // Validation for templated Query - if (retrieval.query == "templated") { - if (retrieval.queryTemplate === null) { - throw new Error("Must set a template for templated query"); - } - } else if (retrieval.queryTemplate !== null) { - throw new Error("Can't set a template without templated query"); - } - - // Validation for Timeframe - if (retrieval.relativeTimeFrame == "custom") { - if ( - retrieval.relativeTimeFrameDuration === null || - retrieval.relativeTimeFrameUnit === null - ) { - throw new Error( - "Custom relative time frame must have a duration and unit set" - ); - } - } - }, - }, - } -); - -/** - * Configuration of Datasources used for Retrieval Action. - */ -export class AgentDataSourceConfiguration extends Model< - InferAttributes, - InferCreationAttributes -> { - declare id: CreationOptional; - declare createdAt: CreationOptional; - declare updatedAt: CreationOptional; - - declare tagsIn: string[] | null; - declare tagsNotIn: string[] | null; - declare parentsIn: string[] | null; - declare parentsNotIn: string[] | null; - - declare dataSourceId: ForeignKey; - declare retrievalConfigurationId: ForeignKey< - AgentRetrievalConfiguration["id"] - >; -} -AgentDataSourceConfiguration.init( - { - id: { - type: DataTypes.INTEGER, - autoIncrement: true, - primaryKey: true, - }, - createdAt: { - type: DataTypes.DATE, - allowNull: false, - defaultValue: DataTypes.NOW, - }, - updatedAt: { - type: DataTypes.DATE, - allowNull: false, - defaultValue: DataTypes.NOW, - }, - tagsIn: { - type: DataTypes.ARRAY(DataTypes.STRING), - allowNull: true, - }, - tagsNotIn: { - type: DataTypes.ARRAY(DataTypes.STRING), - allowNull: true, - }, - parentsIn: { - type: DataTypes.ARRAY(DataTypes.STRING), - allowNull: true, - }, - parentsNotIn: { - type: DataTypes.ARRAY(DataTypes.STRING), - allowNull: true, - }, - }, - { - modelName: "agent_data_source_configuration", - sequelize: front_sequelize, - hooks: { - beforeValidate: (dataSourceConfig: AgentDataSourceConfiguration) => { - if ( - (dataSourceConfig.tagsIn === null) !== - (dataSourceConfig.tagsNotIn === null) - ) { - throw new Error("Tags must be both set or both null"); - } - if ( - (dataSourceConfig.parentsIn === null) !== - (dataSourceConfig.parentsNotIn === null) - ) { - throw new Error("Parents must be both set or both null"); - } - }, - }, - } -); - -// Retrieval config <> data source config -AgentRetrievalConfiguration.hasMany(AgentDataSourceConfiguration, { - foreignKey: { name: "retrievalId", allowNull: false }, - onDelete: "CASCADE", -}); - -// Data source <> Data source config -DataSource.hasMany(AgentDataSourceConfiguration, { - foreignKey: { name: "dataSourceId", allowNull: false }, - onDelete: "CASCADE", -}); - -// Agent config <> Retrieval config -AgentConfiguration.hasOne(AgentRetrievalConfiguration, { - foreignKey: { name: "agentId", allowNull: true }, // null = no generation set for this Agent - onDelete: "CASCADE", -}); - /** * Retrieval Action */ diff --git a/front/lib/models/assistant/agent.ts b/front/lib/models/assistant/agent.ts deleted file mode 100644 index c10393cceac4..000000000000 --- a/front/lib/models/assistant/agent.ts +++ /dev/null @@ -1,167 +0,0 @@ -import { - CreationOptional, - DataTypes, - ForeignKey, - InferAttributes, - InferCreationAttributes, - Model, -} from "sequelize"; - -import { front_sequelize } from "@app/lib/databases"; -import { AgentRetrievalConfiguration } from "@app/lib/models/assistant/actions/retrieval"; -import { Workspace } from "@app/lib/models/workspace"; -import { - AgentConfigurationScope, - AgentConfigurationStatus, -} from "@app/types/assistant/agent"; - -/** - * Agent configuration - */ -export class AgentConfiguration extends Model< - InferAttributes, - InferCreationAttributes -> { - declare id: CreationOptional; - declare createdAt: CreationOptional; - declare updatedAt: CreationOptional; - - declare sId: string; - declare status: AgentConfigurationStatus; - declare name: string; - declare pictureUrl: string | null; - - declare scope: AgentConfigurationScope; - declare workspaceId: ForeignKey | null; // null = it's a global agent - - declare model: ForeignKey | null; -} -AgentConfiguration.init( - { - id: { - type: DataTypes.INTEGER, - autoIncrement: true, - primaryKey: true, - }, - createdAt: { - type: DataTypes.DATE, - allowNull: false, - defaultValue: DataTypes.NOW, - }, - updatedAt: { - type: DataTypes.DATE, - allowNull: false, - defaultValue: DataTypes.NOW, - }, - sId: { - type: DataTypes.STRING, - allowNull: false, - unique: true, - }, - status: { - type: DataTypes.STRING, - allowNull: false, - defaultValue: "active", - }, - name: { - type: DataTypes.TEXT, - allowNull: false, - }, - pictureUrl: { - type: DataTypes.TEXT, - allowNull: true, - }, - scope: { - type: DataTypes.STRING, - allowNull: false, - defaultValue: "workspace", - }, - }, - { - modelName: "agent_configuration", - sequelize: front_sequelize, - indexes: [ - { fields: ["workspaceId"] }, - // Unique name per workspace. - // Note that on PostgreSQL a unique constraint on multiple columns will treat NULL - // as distinct from any other value, so we can create twice the same name if at least - // one of the workspaceId is null. We're okay with it. - { fields: ["workspaceId", "name", "scope"], unique: true }, - { fields: ["sId"], unique: true }, - ], - hooks: { - beforeValidate: (agent: AgentConfiguration) => { - if (agent.scope !== "workspace" && agent.workspaceId) { - throw new Error("Workspace id must be null for global agent"); - } else if (agent.scope === "workspace" && !agent.workspaceId) { - throw new Error("Workspace id must be set for non-global agent"); - } - }, - }, - } -); - -/** - * Configuration of Agent generation. - */ -export class AgentGenerationConfiguration extends Model< - InferAttributes, - InferCreationAttributes -> { - declare id: CreationOptional; - declare createdAt: CreationOptional; - declare updatedAt: CreationOptional; - - declare prompt: string; - declare modelProvider: string; - declare modelId: string; - - declare agentId: ForeignKey; -} -AgentGenerationConfiguration.init( - { - id: { - type: DataTypes.INTEGER, - autoIncrement: true, - primaryKey: true, - }, - createdAt: { - type: DataTypes.DATE, - allowNull: false, - defaultValue: DataTypes.NOW, - }, - updatedAt: { - type: DataTypes.DATE, - allowNull: false, - defaultValue: DataTypes.NOW, - }, - prompt: { - type: DataTypes.TEXT, - allowNull: false, - }, - modelProvider: { - type: DataTypes.STRING, - allowNull: false, - }, - modelId: { - type: DataTypes.STRING, - allowNull: false, - }, - }, - { - modelName: "agent_generation_configuration", - sequelize: front_sequelize, - } -); - -// Workspace <> Agent config -Workspace.hasMany(AgentConfiguration, { - foreignKey: { name: "workspaceId", allowNull: true }, // null = global Agent - onDelete: "CASCADE", -}); - -// Agent config <> Generation config -AgentConfiguration.hasOne(AgentGenerationConfiguration, { - foreignKey: { name: "agentId", allowNull: false }, // null = no retrieval action set for this Agent - onDelete: "CASCADE", -}); diff --git a/front/lib/models/assistant/configuration.ts b/front/lib/models/assistant/configuration.ts new file mode 100644 index 000000000000..018fc1683830 --- /dev/null +++ b/front/lib/models/assistant/configuration.ts @@ -0,0 +1,352 @@ +import { + CreationOptional, + DataTypes, + ForeignKey, + InferAttributes, + InferCreationAttributes, + Model, +} from "sequelize"; + +import { front_sequelize } from "@app/lib/databases"; +import { DataSource } from "@app/lib/models/data_source"; +import { Workspace } from "@app/lib/models/workspace"; +import { TimeframeUnit } from "@app/types/assistant/actions/retrieval"; +import { + AgentConfigurationScope, + AgentConfigurationStatus, +} from "@app/types/assistant/configuration"; + +/** + * Agent configuration + */ +export class AgentConfiguration extends Model< + InferAttributes, + InferCreationAttributes +> { + declare id: CreationOptional; + declare createdAt: CreationOptional; + declare updatedAt: CreationOptional; + + declare sId: string; + declare status: AgentConfigurationStatus; + declare name: string; + declare pictureUrl: string | null; + declare scope: AgentConfigurationScope; + + declare workspaceId: ForeignKey | null; // null = it's a global agent + declare generationId: ForeignKey | null; + declare retrievalId: ForeignKey | null; +} +AgentConfiguration.init( + { + id: { + type: DataTypes.INTEGER, + autoIncrement: true, + primaryKey: true, + }, + createdAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + }, + updatedAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + }, + sId: { + type: DataTypes.STRING, + allowNull: false, + unique: true, + }, + status: { + type: DataTypes.STRING, + allowNull: false, + defaultValue: "active", + }, + name: { + type: DataTypes.TEXT, + allowNull: false, + }, + pictureUrl: { + type: DataTypes.TEXT, + allowNull: true, + }, + scope: { + type: DataTypes.STRING, + allowNull: false, + defaultValue: "workspace", + }, + }, + { + modelName: "agent_configuration", + sequelize: front_sequelize, + indexes: [ + { fields: ["workspaceId"] }, + // Unique name per workspace. + // Note that on PostgreSQL a unique constraint on multiple columns will treat NULL + // as distinct from any other value, so we can create twice the same name if at least + // one of the workspaceId is null. We're okay with it. + { fields: ["workspaceId", "name", "scope"], unique: true }, + { fields: ["sId"], unique: true }, + ], + hooks: { + beforeValidate: (agent: AgentConfiguration) => { + if (agent.scope !== "workspace" && agent.workspaceId) { + throw new Error("Workspace id must be null for global agent"); + } else if (agent.scope === "workspace" && !agent.workspaceId) { + throw new Error("Workspace id must be set for non-global agent"); + } + }, + }, + } +); + +/** + * Configuration of Agent generation. + */ +export class AgentGenerationConfiguration extends Model< + InferAttributes, + InferCreationAttributes +> { + declare id: CreationOptional; + declare createdAt: CreationOptional; + declare updatedAt: CreationOptional; + + declare prompt: string; + declare providerId: string; + declare modelId: string; +} +AgentGenerationConfiguration.init( + { + id: { + type: DataTypes.INTEGER, + autoIncrement: true, + primaryKey: true, + }, + createdAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + }, + updatedAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + }, + prompt: { + type: DataTypes.TEXT, + allowNull: false, + }, + providerId: { + type: DataTypes.STRING, + allowNull: false, + }, + modelId: { + type: DataTypes.STRING, + allowNull: false, + }, + }, + { + modelName: "agent_generation_configuration", + sequelize: front_sequelize, + } +); + +/** + * Action Retrieval configuration + */ +export class AgentRetrievalConfiguration extends Model< + InferAttributes, + InferCreationAttributes +> { + declare id: CreationOptional; + declare createdAt: CreationOptional; + declare updatedAt: CreationOptional; + + declare query: "auto" | "none" | "templated"; + declare queryTemplate: string | null; + declare relativeTimeFrame: "auto" | "none" | "custom"; + declare relativeTimeFrameDuration: number | null; + declare relativeTimeFrameUnit: TimeframeUnit | null; + declare topK: number; +} +AgentRetrievalConfiguration.init( + { + id: { + type: DataTypes.INTEGER, + autoIncrement: true, + primaryKey: true, + }, + createdAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + }, + updatedAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + }, + query: { + type: DataTypes.STRING, + allowNull: false, + defaultValue: "auto", + }, + queryTemplate: { + type: DataTypes.TEXT, + allowNull: true, + }, + relativeTimeFrame: { + type: DataTypes.STRING, + allowNull: false, + defaultValue: "auto", + }, + relativeTimeFrameDuration: { + type: DataTypes.INTEGER, + allowNull: true, + }, + relativeTimeFrameUnit: { + type: DataTypes.STRING, + allowNull: true, + }, + topK: { + type: DataTypes.INTEGER, + allowNull: false, + }, + }, + { + modelName: "agent_retrieval_configuration", + sequelize: front_sequelize, + hooks: { + beforeValidate: (retrieval: AgentRetrievalConfiguration) => { + // Validation for templated Query + if (retrieval.query == "templated") { + if (retrieval.queryTemplate === null) { + throw new Error("Must set a template for templated query"); + } + } else if (retrieval.queryTemplate !== null) { + throw new Error("Can't set a template without templated query"); + } + + // Validation for Timeframe + if (retrieval.relativeTimeFrame == "custom") { + if ( + retrieval.relativeTimeFrameDuration === null || + retrieval.relativeTimeFrameUnit === null + ) { + throw new Error( + "Custom relative time frame must have a duration and unit set" + ); + } + } + }, + }, + } +); + +/** + * Configuration of Datasources used for Retrieval Action. + */ +export class AgentDataSourceConfiguration extends Model< + InferAttributes, + InferCreationAttributes +> { + declare id: CreationOptional; + declare createdAt: CreationOptional; + declare updatedAt: CreationOptional; + + declare tagsIn: string[] | null; + declare tagsNotIn: string[] | null; + declare parentsIn: string[] | null; + declare parentsNotIn: string[] | null; + + declare dataSourceId: ForeignKey; + declare retrievalConfigurationId: ForeignKey< + AgentRetrievalConfiguration["id"] + >; +} +AgentDataSourceConfiguration.init( + { + id: { + type: DataTypes.INTEGER, + autoIncrement: true, + primaryKey: true, + }, + createdAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + }, + updatedAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + }, + tagsIn: { + type: DataTypes.ARRAY(DataTypes.STRING), + allowNull: true, + }, + tagsNotIn: { + type: DataTypes.ARRAY(DataTypes.STRING), + allowNull: true, + }, + parentsIn: { + type: DataTypes.ARRAY(DataTypes.STRING), + allowNull: true, + }, + parentsNotIn: { + type: DataTypes.ARRAY(DataTypes.STRING), + allowNull: true, + }, + }, + { + modelName: "agent_data_source_configuration", + sequelize: front_sequelize, + hooks: { + beforeValidate: (dataSourceConfig: AgentDataSourceConfiguration) => { + if ( + (dataSourceConfig.tagsIn === null) !== + (dataSourceConfig.tagsNotIn === null) + ) { + throw new Error("Tags must be both set or both null"); + } + if ( + (dataSourceConfig.parentsIn === null) !== + (dataSourceConfig.parentsNotIn === null) + ) { + throw new Error("Parents must be both set or both null"); + } + }, + }, + } +); + +// Agent config <> Workspace +Workspace.hasMany(AgentConfiguration, { + foreignKey: { name: "workspaceId", allowNull: true }, // null = global Agent + onDelete: "CASCADE", +}); + +// Agent config <> Generation config +AgentConfiguration.hasOne(AgentGenerationConfiguration, { + foreignKey: { name: "generationId", allowNull: true }, // null = no generation set for this Agent + onDelete: "CASCADE", +}); +// Agent config <> Retrieval config +AgentConfiguration.hasOne(AgentRetrievalConfiguration, { + foreignKey: { name: "retrievalId", allowNull: true }, // null = no retrieval action set for this Agent + onDelete: "CASCADE", +}); + +// Retrieval config <> Data source config +AgentRetrievalConfiguration.hasMany(AgentDataSourceConfiguration, { + foreignKey: { name: "retrievalId", allowNull: false }, + onDelete: "CASCADE", +}); + +// Data source config <> Data source +DataSource.hasMany(AgentDataSourceConfiguration, { + foreignKey: { name: "dataSourceId", allowNull: false }, + onDelete: "CASCADE", +}); diff --git a/front/lib/models/index.ts b/front/lib/models/index.ts index 0bc9edbbbdc0..d15b932065a2 100644 --- a/front/lib/models/index.ts +++ b/front/lib/models/index.ts @@ -1,15 +1,15 @@ import { App, Clone, Dataset, Provider, Run } from "@app/lib/models/apps"; import { - AgentDataSourceConfiguration, AgentRetrievalAction, - AgentRetrievalConfiguration, RetrievalDocument, RetrievalDocumentChunk, } from "@app/lib/models/assistant/actions/retrieval"; import { AgentConfiguration, + AgentDataSourceConfiguration, AgentGenerationConfiguration, -} from "@app/lib/models/assistant/agent"; + AgentRetrievalConfiguration, +} from "@app/lib/models/assistant/configuration"; import { AgentMessage, Conversation, diff --git a/front/types/assistant/actions/retrieval.ts b/front/types/assistant/actions/retrieval.ts index eff6727df63b..47a8138a6120 100644 --- a/front/types/assistant/actions/retrieval.ts +++ b/front/types/assistant/actions/retrieval.ts @@ -1,9 +1,5 @@ -/** - * Data Source configuration - */ - import { ModelId } from "@app/lib/databases"; -import { AgentActionConfigurationType } from "@app/types/assistant/agent"; +import { AgentDataSourceConfigurationType } from "@app/types/assistant/configuration"; import { AgentActionType } from "@app/types/assistant/conversation"; export type TimeframeUnit = "hour" | "day" | "week" | "month" | "year"; @@ -12,22 +8,6 @@ export type TimeFrame = { unit: TimeframeUnit; }; -export type DataSourceFilter = { - tags: { in: string[]; not: string[] } | null; - parents: { in: string[]; not: string[] } | null; -}; - -// This is used to talk with Dust Apps and Core, so it store external Ids. -export type AgentDataSourceConfigurationType = { - workspaceSId: string; // = Workspace.sId - dataSourceName: string; // = Datasource.name - filter: DataSourceFilter; -}; - -/** - * Retrieval configuration - */ - export type TemplatedQuery = { template: string; }; @@ -51,27 +31,6 @@ export function isTimeFrame(arg: RetrievalTimeframe): arg is TimeFrame { // `dataSources` to query the data. export type RetrievalTimeframe = "auto" | "none" | TimeFrame; export type RetrievalQuery = "auto" | "none" | TemplatedQuery; -export type RetrievalDataSourcesConfiguration = - AgentDataSourceConfigurationType[]; - -export type RetrievalConfigurationType = { - id: ModelId; - - type: "retrieval_configuration"; - dataSources: RetrievalDataSourcesConfiguration; - query: RetrievalQuery; - relativeTimeFrame: RetrievalTimeframe; - topK: number; - - // Dynamically decide to skip, if needed in the future - // autoSkip: boolean; -}; - -export function isRetrievalConfiguration( - arg: AgentActionConfigurationType | null -): arg is RetrievalConfigurationType { - return arg !== null && arg.type && arg.type === "retrieval_configuration"; -} /** * Retrieval action @@ -103,7 +62,7 @@ export type RetrievalActionType = { id: ModelId; // AgentRetrieval. type: "retrieval_action"; params: { - dataSources: "all" | AgentDataSourceConfigurationType[]; + dataSources: AgentDataSourceConfigurationType[]; relativeTimeFrame: TimeFrame | null; query: string | null; topK: number; diff --git a/front/types/assistant/agent.ts b/front/types/assistant/agent.ts index f25d0197c184..4eb7cd6c0040 100644 --- a/front/types/assistant/agent.ts +++ b/front/types/assistant/agent.ts @@ -1,15 +1,7 @@ -import { ModelId } from "@app/lib/databases"; -import { RetrievalConfigurationType } from "@app/types/assistant/actions/retrieval"; - /** - * Agent Action configuration + * Agent Action */ -// New AgentActionConfigurationType checklist: -// - Add the type to the union type below -// - Add model rendering support in `renderConversationForModel` -export type AgentActionConfigurationType = RetrievalConfigurationType; - // Each AgentActionConfigurationType is capable of generating this type at runtime to specify which // inputs should be generated by the model. As an example, to run the retrieval action for which the // `relativeTimeFrame` has been specified in the configuration but for which the `query` is "auto", @@ -32,7 +24,6 @@ export type AgentActionConfigurationType = RetrievalConfigurationType; // ``` export type AgentActionSpecification = { - id: ModelId; name: string; description: string; inputs: { @@ -41,40 +32,3 @@ export type AgentActionSpecification = { type: "string" | "number" | "boolean"; }[]; }; - -/** - * Agent Message configuration - */ - -export type AgentGenerationConfigurationType = { - id: ModelId; - prompt: string; - model: { - providerId: string; - modelId: string; - }; -}; - -/** - * Agent configuration - */ - -export type AgentConfigurationStatus = "active" | "archived"; -export type AgentConfigurationScope = "global" | "workspace"; - -export type AgentConfigurationType = { - id: ModelId; - sId: string; - status: AgentConfigurationStatus; - name: string; - pictureUrl: string | null; -}; - -export type AgentFullConfigurationType = { - agent: AgentConfigurationType; - // If undefined, no action performed, otherwise the action is - // performed (potentially NoOp eg autoSkip above). - action: AgentActionConfigurationType | null; - // If undefined, no text generation. - generation: AgentGenerationConfigurationType | null; -}; diff --git a/front/types/assistant/configuration.ts b/front/types/assistant/configuration.ts new file mode 100644 index 000000000000..62d825a62beb --- /dev/null +++ b/front/types/assistant/configuration.ts @@ -0,0 +1,66 @@ +import { ModelId } from "@app/lib/databases"; +import { + RetrievalQuery, + RetrievalTimeframe, +} from "@app/types/assistant/actions/retrieval"; + +/** + * Agent config + */ +export type AgentConfigurationStatus = "active" | "archived"; +export type AgentConfigurationScope = "global" | "workspace"; +export type AgentConfigurationType = { + sId: string; + status: AgentConfigurationStatus; + name: string; + pictureUrl: string | null; + action: AgentActionConfigurationType | null; // If undefined, no action performed + generation: AgentGenerationConfigurationType | null; // If undefined, no text generation. +}; + +/** + * Generation config + */ +export type AgentGenerationConfigurationType = { + id: ModelId; + prompt: string; + model: { + providerId: string; + modelId: string; + }; +}; + +/** + * Action > Retrieval + */ +export type AgentActionConfigurationType = RetrievalConfigurationType; + +/** + * Retrieval Action config + */ +export type RetrievalConfigurationType = { + id: ModelId; + + type: "retrieval_configuration"; + dataSources: AgentDataSourceConfigurationType[]; + query: RetrievalQuery; + relativeTimeFrame: RetrievalTimeframe; + topK: number; +}; +export function isRetrievalConfiguration( + arg: AgentActionConfigurationType | null +): arg is RetrievalConfigurationType { + return arg !== null && arg.type && arg.type === "retrieval_configuration"; +} + +/** + * Datasources config for Retrieval Action + */ +export type AgentDataSourceConfigurationType = { + workspaceSId: string; // need sId to talk with Core (external id) + dataSourceName: string; // need Datasource.name to talk with Core (external id) + filter: { + tags: { in: string[]; not: string[] } | null; + parents: { in: string[]; not: string[] } | null; + }; +}; diff --git a/front/types/assistant/conversation.ts b/front/types/assistant/conversation.ts index 89b1d3cbd9a4..342d1f8a246f 100644 --- a/front/types/assistant/conversation.ts +++ b/front/types/assistant/conversation.ts @@ -1,5 +1,5 @@ import { ModelId } from "@app/lib/databases"; -import { AgentFullConfigurationType } from "@app/types/assistant/agent"; +import { AgentConfigurationType } from "@app/types/assistant/configuration"; import { UserType } from "@app/types/user"; import { RetrievalActionType } from "./actions/retrieval"; @@ -91,7 +91,7 @@ export type AgentMessageType = { version: number; parentMessageId: string | null; - configuration: AgentFullConfigurationType; + configuration: AgentConfigurationType; status: AgentMessageStatus; action: AgentActionType | null; message: string | null; From cc3bd444f5495ebc25129b13df4d9894e48b1371 Mon Sep 17 00:00:00 2001 From: PopDaph Date: Fri, 8 Sep 2023 14:52:26 +0200 Subject: [PATCH 07/12] Fix bad rebase --- front/lib/api/assistant/conversation.ts | 111 ++++++++++++++++++------ 1 file changed, 86 insertions(+), 25 deletions(-) diff --git a/front/lib/api/assistant/conversation.ts b/front/lib/api/assistant/conversation.ts index 1dd3da5926ef..b89e253484d3 100644 --- a/front/lib/api/assistant/conversation.ts +++ b/front/lib/api/assistant/conversation.ts @@ -75,31 +75,92 @@ export async function* postUserMessage( where: { conversationId: conversation.id, }, - { - transaction: t, - } - ); - agentMessages.push({ - id: agentMessageRow.id, - sId: agentMessageRow.sId, - type: "agent_message", - visibility: "visible", - version: 0, - parentMessageId: userMessage.sId, - status: "created", - action: null, - message: null, - feedbacks: [], - error: null, - configuration: { - sId: m.configurationId, - status: "active", - name: "foo", // TODO - pictureUrl: null, // TODO - action: null, // TODO - generation: null, // TODO - }, - }); + transaction: t, + })) ?? -1) + 1; + + const m = await Message.create( + { + sId: generateModelSId(), + rank: nextMessageRank++, + conversationId: conversation.id, + parentId: null, + userMessageId: ( + await UserMessage.create( + { + message: message, + userContextUsername: context.username, + userContextTimezone: context.timezone, + userContextFullName: context.fullName, + userContextEmail: context.email, + userContextProfilePictureUrl: context.profilePictureUrl, + userId: user ? user.id : null, + }, + { transaction: t } + ) + ).id, + }, + { + transaction: t, + } + ); + + const userMessage: UserMessageType = { + id: m.id, + sId: m.sId, + type: "user_message", + visibility: "visible", + version: 0, + user: user, + mentions: mentions, + message: message, + context: context, + }; + + const agentMessages: AgentMessageType[] = []; + const agentMessageRows: AgentMessage[] = []; + + // for each assistant mention, create an "empty" agent message + for (const mention of mentions) { + if (isAgentMention(mention)) { + const agentMessageRow = await AgentMessage.create( + {}, + { transaction: t } + ); + const m = await Message.create( + { + sId: generateModelSId(), + rank: nextMessageRank++, + conversationId: conversation.id, + parentId: userMessage.id, + agentMessageId: agentMessageRow.id, + }, + { + transaction: t, + } + ); + agentMessageRows.push(agentMessageRow); + agentMessages.push({ + id: m.id, + sId: m.sId, + type: "agent_message", + visibility: "visible", + version: 0, + parentMessageId: userMessage.sId, + status: "created", + action: null, + message: null, + feedbacks: [], + error: null, + configuration: { + sId: mention.configurationId, + status: "active", + name: "foo", // TODO + pictureUrl: null, // TODO + action: null, // TODO + generation: null, // TODO + }, + }); + } } return { userMessage, agentMessages, agentMessageRows }; From 75297c17607dac3b8d67e3f0315c7bffe3d38d0b Mon Sep 17 00:00:00 2001 From: PopDaph Date: Fri, 8 Sep 2023 15:34:31 +0200 Subject: [PATCH 08/12] WIP --- front/lib/api/assistant/configuration.ts | 37 +++++++++++---------- front/lib/models/assistant/configuration.ts | 10 +++--- 2 files changed, 24 insertions(+), 23 deletions(-) diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts index a6a74f62c3e1..0d0a896b90ad 100644 --- a/front/lib/api/assistant/configuration.ts +++ b/front/lib/api/assistant/configuration.ts @@ -63,7 +63,7 @@ export async function getAgentConfiguration( const datasources = action?.id ? await AgentDataSourceConfiguration.findAll({ where: { - retrievalConfigurationId: action.id, + retrievalId: action.id, }, }) : []; @@ -162,16 +162,19 @@ export async function createAgentConfiguration( } // Create Agent config - const agentConfig = await AgentConfiguration.create({ - sId: generateModelSId(), - status: status, - name: name, - pictureUrl: pictureUrl, - scope: "workspace", - workspaceId: owner.id, - generationId: genConfig?.id ?? null, - retrievalId: retrievalConfig?.id ?? null, - }); + const agentConfig = await AgentConfiguration.create( + { + sId: generateModelSId(), + status: status, + name: name, + pictureUrl: pictureUrl, + scope: "workspace", + workspaceId: owner.id, + generationId: genConfig?.id ?? null, + retrievalId: retrievalConfig?.id ?? null, + }, + { transaction: t } + ); return { sId: agentConfig.sId, @@ -196,9 +199,9 @@ export async function createAgentConfiguration( } /** - * Update Agent Generation Configuration + * Update Agent Configuration */ -export async function updateAgentGenerationConfiguration( +export async function updateAgentConfiguration( auth: Authenticator, agentId: string, { @@ -254,13 +257,13 @@ export async function updateAgentGenerationConfiguration( const existingDataSourcesConfig = existingRetrivalConfig?.id ? await AgentDataSourceConfiguration.findAll({ where: { - retrievalConfigurationId: existingRetrivalConfig.id, + retrievalId: existingRetrivalConfig.id, }, }) : []; return await front_sequelize.transaction(async (t) => { - // Upserting Agent Config + // Updating Agent Config const updatedAgentConfig = await agentConfig.update( { name: name, @@ -515,7 +518,7 @@ export async function _agentActionType( export async function _createAgentDataSourcesConfigData( t: Transaction, dataSourcesConfig: AgentDataSourceConfigurationType[], - agentActionId: number + retrievalConfigId: number ): Promise { // dsConfig contains this format: // [ @@ -608,7 +611,7 @@ export async function _createAgentDataSourcesConfigData( tagsNotIn: dsConfig.filter.tags?.not, parentsIn: dsConfig.filter.parents?.in, parentsNotIn: dsConfig.filter.parents?.not, - retrievalConfigurationId: agentActionId, + retrievalId: retrievalConfigId, }, { transaction: t } ); diff --git a/front/lib/models/assistant/configuration.ts b/front/lib/models/assistant/configuration.ts index 018fc1683830..77dd98b608ab 100644 --- a/front/lib/models/assistant/configuration.ts +++ b/front/lib/models/assistant/configuration.ts @@ -262,9 +262,7 @@ export class AgentDataSourceConfiguration extends Model< declare parentsNotIn: string[] | null; declare dataSourceId: ForeignKey; - declare retrievalConfigurationId: ForeignKey< - AgentRetrievalConfiguration["id"] - >; + declare retrievalId: ForeignKey; } AgentDataSourceConfiguration.init( { @@ -329,18 +327,18 @@ Workspace.hasMany(AgentConfiguration, { }); // Agent config <> Generation config -AgentConfiguration.hasOne(AgentGenerationConfiguration, { +AgentGenerationConfiguration.hasOne(AgentConfiguration, { foreignKey: { name: "generationId", allowNull: true }, // null = no generation set for this Agent onDelete: "CASCADE", }); // Agent config <> Retrieval config -AgentConfiguration.hasOne(AgentRetrievalConfiguration, { +AgentRetrievalConfiguration.hasOne(AgentConfiguration, { foreignKey: { name: "retrievalId", allowNull: true }, // null = no retrieval action set for this Agent onDelete: "CASCADE", }); // Retrieval config <> Data source config -AgentRetrievalConfiguration.hasMany(AgentDataSourceConfiguration, { +AgentRetrievalConfiguration.hasOne(AgentDataSourceConfiguration, { foreignKey: { name: "retrievalId", allowNull: false }, onDelete: "CASCADE", }); From 31202d986e47c2cee89771b988b527e1dfe32681 Mon Sep 17 00:00:00 2001 From: PopDaph Date: Fri, 8 Sep 2023 15:39:57 +0200 Subject: [PATCH 09/12] WIP --- front/admin/db.ts | 3 +-- front/lib/api/assistant/actions/retrieval.ts | 4 ++-- front/lib/api/assistant/configuration.ts | 16 ++++++++-------- front/types/assistant/configuration.ts | 4 ++-- 4 files changed, 13 insertions(+), 14 deletions(-) diff --git a/front/admin/db.ts b/front/admin/db.ts index 29890f3a3d0b..477b50f3b030 100644 --- a/front/admin/db.ts +++ b/front/admin/db.ts @@ -56,11 +56,10 @@ async function main() { await ExtractedEvent.sync({ alter: true }); await DocumentTrackerChangeSuggestion.sync({ alter: true }); - await AgentConfiguration.sync({ alter: true }); await AgentGenerationConfiguration.sync({ alter: true }); - await AgentRetrievalConfiguration.sync({ alter: true }); await AgentDataSourceConfiguration.sync({ alter: true }); + await AgentConfiguration.sync({ alter: true }); await AgentRetrievalAction.sync({ alter: true }); await RetrievalDocument.sync({ alter: true }); await RetrievalDocumentChunk.sync({ alter: true }); diff --git a/front/lib/api/assistant/actions/retrieval.ts b/front/lib/api/assistant/actions/retrieval.ts index 6730bef1fc89..bbe5d7762cad 100644 --- a/front/lib/api/assistant/actions/retrieval.ts +++ b/front/lib/api/assistant/actions/retrieval.ts @@ -388,8 +388,8 @@ export async function* runRetrieval( // Handle data sources list and parents/tags filtering. config.DATASOURCE.data_sources = c.dataSources.map((d) => ({ - workspace_id: d.workspaceSId, - data_source_id: d.dataSourceName, + workspace_id: d.workspaceId, + data_source_id: d.dataSourceId, })); for (const ds of c.dataSources) { diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts index 0d0a896b90ad..1550a9d34a72 100644 --- a/front/lib/api/assistant/configuration.ts +++ b/front/lib/api/assistant/configuration.ts @@ -483,8 +483,8 @@ export async function _agentActionType( } dataSourcesConfigType.push({ - dataSourceName: dataSource.name, - workspaceSId: workspace.sId, + dataSourceId: dataSource.name, + workspaceId: workspace.sId, filter: { tags: dsConfig.tagsIn && dsConfig.tagsNotIn @@ -530,7 +530,7 @@ export async function _createAgentDataSourcesConfigData( // First we get the list of workspaces because we need the mapping between workspaceSId and workspaceId const workspaces = await Workspace.findAll({ where: { - sId: dataSourcesConfig.map((dsConfig) => dsConfig.workspaceSId), + sId: dataSourcesConfig.map((dsConfig) => dsConfig.workspaceId), }, attributes: ["id", "sId"], }); @@ -551,7 +551,7 @@ export async function _createAgentDataSourcesConfigData( curr: AgentDataSourceConfigurationType ) => { // First we need to get the workspaceId from the workspaceSId - const workspace = workspaces.find((w) => w.sId === curr.workspaceSId); + const workspace = workspaces.find((w) => w.sId === curr.workspaceId); if (!workspace) { throw new Error("Workspace not found"); } @@ -563,12 +563,12 @@ export async function _createAgentDataSourcesConfigData( ); if (existingEntry) { // Append dataSourceName to existing entry - existingEntry.dataSourceNames.push(curr.dataSourceName); + existingEntry.dataSourceNames.push(curr.dataSourceId); } else { // Add a new entry for this workspaceId acc.push({ workspaceId: workspace.id, - dataSourceNames: [curr.dataSourceName], + dataSourceNames: [curr.dataSourceId], }); } return acc; @@ -597,9 +597,9 @@ export async function _createAgentDataSourcesConfigData( dataSourcesConfig.map(async (dsConfig) => { const dataSource = dataSources.find( (ds) => - ds.name === dsConfig.dataSourceName && + ds.name === dsConfig.dataSourceId && ds.workspaceId === - workspaces.find((w) => w.sId === dsConfig.workspaceSId)?.id + workspaces.find((w) => w.sId === dsConfig.workspaceId)?.id ); if (!dataSource) { throw new Error("DataSource not found"); diff --git a/front/types/assistant/configuration.ts b/front/types/assistant/configuration.ts index 62d825a62beb..33c925671fc3 100644 --- a/front/types/assistant/configuration.ts +++ b/front/types/assistant/configuration.ts @@ -57,8 +57,8 @@ export function isRetrievalConfiguration( * Datasources config for Retrieval Action */ export type AgentDataSourceConfigurationType = { - workspaceSId: string; // need sId to talk with Core (external id) - dataSourceName: string; // need Datasource.name to talk with Core (external id) + workspaceId: string; // need sId to talk with Core (external id) + dataSourceId: string; // need Datasource.name to talk with Core (external id) filter: { tags: { in: string[]; not: string[] } | null; parents: { in: string[]; not: string[] } | null; From b49e22d8cd523ed33d6862db51986f04927ad4ce Mon Sep 17 00:00:00 2001 From: PopDaph Date: Fri, 8 Sep 2023 15:45:08 +0200 Subject: [PATCH 10/12] Move types --- front/lib/api/assistant/actions/retrieval.ts | 6 +-- front/lib/api/assistant/agent.ts | 6 +-- front/lib/api/assistant/configuration.ts | 36 +++++++------- front/lib/api/assistant/generation.ts | 2 +- front/lib/models/assistant/configuration.ts | 16 +++--- front/types/assistant/actions/retrieval.ts | 32 +++++++++++- front/types/assistant/agent.ts | 21 -------- front/types/assistant/configuration.ts | 51 ++++++++------------ 8 files changed, 86 insertions(+), 84 deletions(-) diff --git a/front/lib/api/assistant/actions/retrieval.ts b/front/lib/api/assistant/actions/retrieval.ts index bbe5d7762cad..0c65fa0f468b 100644 --- a/front/lib/api/assistant/actions/retrieval.ts +++ b/front/lib/api/assistant/actions/retrieval.ts @@ -21,13 +21,13 @@ import { RetrievalDocumentType, TimeFrame, } from "@app/types/assistant/actions/retrieval"; -import { AgentActionSpecification } from "@app/types/assistant/agent"; import { - AgentConfigurationType, AgentDataSourceConfigurationType, isRetrievalConfiguration, RetrievalConfigurationType, -} from "@app/types/assistant/configuration"; +} from "@app/types/assistant/actions/retrieval"; +import { AgentActionSpecification } from "@app/types/assistant/agent"; +import { AgentConfigurationType } from "@app/types/assistant/configuration"; import { AgentMessageType, ConversationType, diff --git a/front/lib/api/assistant/agent.ts b/front/lib/api/assistant/agent.ts index 4d9891171f10..702f0bb1cb61 100644 --- a/front/lib/api/assistant/agent.ts +++ b/front/lib/api/assistant/agent.ts @@ -16,11 +16,9 @@ import { import { Authenticator } from "@app/lib/auth"; import { Err, Ok, Result } from "@app/lib/result"; import { generateModelSId } from "@app/lib/utils"; +import { isRetrievalConfiguration } from "@app/types/assistant/actions/retrieval"; import { AgentActionSpecification } from "@app/types/assistant/agent"; -import { - AgentConfigurationType, - isRetrievalConfiguration, -} from "@app/types/assistant/configuration"; +import { AgentConfigurationType } from "@app/types/assistant/configuration"; import { AgentActionType, AgentMessageType, diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts index 1550a9d34a72..1406d06e1bab 100644 --- a/front/lib/api/assistant/configuration.ts +++ b/front/lib/api/assistant/configuration.ts @@ -12,6 +12,7 @@ import { } from "@app/lib/models"; import { generateModelSId } from "@app/lib/utils"; import { + AgentDataSourceConfigurationType, isTemplatedQuery, isTimeFrame, RetrievalQuery, @@ -21,7 +22,6 @@ import { AgentActionConfigurationType, AgentConfigurationStatus, AgentConfigurationType, - AgentDataSourceConfigurationType, } from "@app/types/assistant/configuration"; /** @@ -45,25 +45,25 @@ export async function getAgentConfiguration( throw new Error("Cannot find Agent: no workspace"); } - const generation = agent.generationId + const generation = agent.generationConfigId ? await AgentGenerationConfiguration.findOne({ where: { - id: agent.generationId, + id: agent.generationConfigId, }, }) : null; - const action = agent.retrievalId + const action = agent.retrievalConfigId ? await AgentRetrievalConfiguration.findOne({ where: { - id: agent.retrievalId, + id: agent.retrievalConfigId, }, }) : null; const datasources = action?.id ? await AgentDataSourceConfiguration.findAll({ where: { - retrievalId: action.id, + retrievalConfigId: action.id, }, }) : []; @@ -170,8 +170,8 @@ export async function createAgentConfiguration( pictureUrl: pictureUrl, scope: "workspace", workspaceId: owner.id, - generationId: genConfig?.id ?? null, - retrievalId: retrievalConfig?.id ?? null, + generationConfigId: genConfig?.id ?? null, + retrievalConfigId: retrievalConfig?.id ?? null, }, { transaction: t } ); @@ -238,18 +238,18 @@ export async function updateAgentConfiguration( "Cannot create AgentGenerationConfiguration: Agent not found" ); } - const existingGeneration = agentConfig.generationId + const existingGeneration = agentConfig.generationConfigId ? await AgentGenerationConfiguration.findOne({ where: { - id: agentConfig.generationId, + id: agentConfig.generationConfigId, }, }) : null; - const existingRetrivalConfig = agentConfig.retrievalId + const existingRetrivalConfig = agentConfig.retrievalConfigId ? await AgentRetrievalConfiguration.findOne({ where: { - id: agentConfig.retrievalId, + id: agentConfig.retrievalConfigId, }, }) : null; @@ -257,7 +257,7 @@ export async function updateAgentConfiguration( const existingDataSourcesConfig = existingRetrivalConfig?.id ? await AgentDataSourceConfiguration.findAll({ where: { - retrievalId: existingRetrivalConfig.id, + retrievalConfigId: existingRetrivalConfig.id, }, }) : []; @@ -361,19 +361,19 @@ export async function updateAgentRetrievalConfiguration( "Cannot create AgentGenerationConfiguration: Agent not found" ); } - const generationConfig = agentConfig.generationId + const generationConfig = agentConfig.generationConfigId ? await AgentGenerationConfiguration.findOne({ where: { - id: agentConfig.generationId, + id: agentConfig.generationConfigId, }, }) : null; return await front_sequelize.transaction(async (t) => { - if (agentConfig.retrievalId) { + if (agentConfig.retrievalConfigId) { const existingRetrivalConfig = await AgentRetrievalConfiguration.findOne({ where: { - id: agentConfig.retrievalId, + id: agentConfig.retrievalConfigId, }, }); if (existingRetrivalConfig) { @@ -611,7 +611,7 @@ export async function _createAgentDataSourcesConfigData( tagsNotIn: dsConfig.filter.tags?.not, parentsIn: dsConfig.filter.parents?.in, parentsNotIn: dsConfig.filter.parents?.not, - retrievalId: retrievalConfigId, + retrievalConfigId: retrievalConfigId, }, { transaction: t } ); diff --git a/front/lib/api/assistant/generation.ts b/front/lib/api/assistant/generation.ts index 70bdb71160ee..c4a3c6c3a2da 100644 --- a/front/lib/api/assistant/generation.ts +++ b/front/lib/api/assistant/generation.ts @@ -9,7 +9,7 @@ import { CoreAPI } from "@app/lib/core_api"; import { Err, Ok, Result } from "@app/lib/result"; import logger from "@app/logger/logger"; import { isRetrievalActionType } from "@app/types/assistant/actions/retrieval"; -import { AgentConfigurationType } from "@app/types/assistant/agent"; +import { AgentConfigurationType } from "@app/types/assistant/configuration"; import { AgentMessageType, ConversationType, diff --git a/front/lib/models/assistant/configuration.ts b/front/lib/models/assistant/configuration.ts index 77dd98b608ab..4ce9a40c7c72 100644 --- a/front/lib/models/assistant/configuration.ts +++ b/front/lib/models/assistant/configuration.ts @@ -34,8 +34,12 @@ export class AgentConfiguration extends Model< declare scope: AgentConfigurationScope; declare workspaceId: ForeignKey | null; // null = it's a global agent - declare generationId: ForeignKey | null; - declare retrievalId: ForeignKey | null; + declare generationConfigId: ForeignKey< + AgentGenerationConfiguration["id"] + > | null; + declare retrievalConfigId: ForeignKey< + AgentRetrievalConfiguration["id"] + > | null; } AgentConfiguration.init( { @@ -262,7 +266,7 @@ export class AgentDataSourceConfiguration extends Model< declare parentsNotIn: string[] | null; declare dataSourceId: ForeignKey; - declare retrievalId: ForeignKey; + declare retrievalConfigId: ForeignKey; } AgentDataSourceConfiguration.init( { @@ -328,18 +332,18 @@ Workspace.hasMany(AgentConfiguration, { // Agent config <> Generation config AgentGenerationConfiguration.hasOne(AgentConfiguration, { - foreignKey: { name: "generationId", allowNull: true }, // null = no generation set for this Agent + foreignKey: { name: "generationConfigId", allowNull: true }, // null = no generation set for this Agent onDelete: "CASCADE", }); // Agent config <> Retrieval config AgentRetrievalConfiguration.hasOne(AgentConfiguration, { - foreignKey: { name: "retrievalId", allowNull: true }, // null = no retrieval action set for this Agent + foreignKey: { name: "retrievalConfigId", allowNull: true }, // null = no retrieval action set for this Agent onDelete: "CASCADE", }); // Retrieval config <> Data source config AgentRetrievalConfiguration.hasOne(AgentDataSourceConfiguration, { - foreignKey: { name: "retrievalId", allowNull: false }, + foreignKey: { name: "retrievalConfigId", allowNull: false }, onDelete: "CASCADE", }); diff --git a/front/types/assistant/actions/retrieval.ts b/front/types/assistant/actions/retrieval.ts index 47a8138a6120..42d0935d920d 100644 --- a/front/types/assistant/actions/retrieval.ts +++ b/front/types/assistant/actions/retrieval.ts @@ -1,5 +1,5 @@ import { ModelId } from "@app/lib/databases"; -import { AgentDataSourceConfigurationType } from "@app/types/assistant/configuration"; +import { AgentActionConfigurationType } from "@app/types/assistant/configuration"; import { AgentActionType } from "@app/types/assistant/conversation"; export type TimeframeUnit = "hour" | "day" | "week" | "month" | "year"; @@ -21,6 +21,36 @@ export function isTimeFrame(arg: RetrievalTimeframe): arg is TimeFrame { ); } +/** + * Retrieval Action config + */ +export type RetrievalConfigurationType = { + id: ModelId; + + type: "retrieval_configuration"; + dataSources: AgentDataSourceConfigurationType[]; + query: RetrievalQuery; + relativeTimeFrame: RetrievalTimeframe; + topK: number; +}; +export function isRetrievalConfiguration( + arg: AgentActionConfigurationType | null +): arg is RetrievalConfigurationType { + return arg !== null && arg.type && arg.type === "retrieval_configuration"; +} + +/** + * Datasources config for Retrieval Action + */ +export type AgentDataSourceConfigurationType = { + workspaceId: string; // need sId to talk with Core (external id) + dataSourceId: string; // need Datasource.name to talk with Core (external id) + filter: { + tags: { in: string[]; not: string[] } | null; + parents: { in: string[]; not: string[] } | null; + }; +}; + // Retrieval specifies a list of data sources (with possible parent / tags filtering, possible "all" // data sources), a query ("auto" generated by the model "none", no query, `TemplatedQuery`, fixed // query), a relative time frame ("auto" generated by the model, "none" no time filtering diff --git a/front/types/assistant/agent.ts b/front/types/assistant/agent.ts index 4eb7cd6c0040..0ea0e0096821 100644 --- a/front/types/assistant/agent.ts +++ b/front/types/assistant/agent.ts @@ -2,27 +2,6 @@ * Agent Action */ -// Each AgentActionConfigurationType is capable of generating this type at runtime to specify which -// inputs should be generated by the model. As an example, to run the retrieval action for which the -// `relativeTimeFrame` has been specified in the configuration but for which the `query` is "auto", -// it would generate: -// -// ``` -// { inputs: [{ name: "query", description: "...", type: "string" }] -// ``` -// -// The params generator model for this action would be tasked to generate that query. If the -// retrieval configuration sets `relativeTimeFrame` to "auto" as well we would get: -// -// ``` -// { -// inputs: [ -// { name: "query", description: "...", type: "string" }, -// { name: "relativeTimeFrame", description: "...", type: "string" }, -// ] -// } -// ``` - export type AgentActionSpecification = { name: string; description: string; diff --git a/front/types/assistant/configuration.ts b/front/types/assistant/configuration.ts index 33c925671fc3..002bc69371cd 100644 --- a/front/types/assistant/configuration.ts +++ b/front/types/assistant/configuration.ts @@ -1,5 +1,6 @@ import { ModelId } from "@app/lib/databases"; import { + RetrievalConfigurationType, RetrievalQuery, RetrievalTimeframe, } from "@app/types/assistant/actions/retrieval"; @@ -33,34 +34,24 @@ export type AgentGenerationConfigurationType = { /** * Action > Retrieval */ +// Each AgentActionConfigurationType is capable of generating this type at runtime to specify which +// inputs should be generated by the model. As an example, to run the retrieval action for which the +// `relativeTimeFrame` has been specified in the configuration but for which the `query` is "auto", +// it would generate: +// +// ``` +// { inputs: [{ name: "query", description: "...", type: "string" }] +// ``` +// +// The params generator model for this action would be tasked to generate that query. If the +// retrieval configuration sets `relativeTimeFrame` to "auto" as well we would get: +// +// ``` +// { +// inputs: [ +// { name: "query", description: "...", type: "string" }, +// { name: "relativeTimeFrame", description: "...", type: "string" }, +// ] +// } +// ``` export type AgentActionConfigurationType = RetrievalConfigurationType; - -/** - * Retrieval Action config - */ -export type RetrievalConfigurationType = { - id: ModelId; - - type: "retrieval_configuration"; - dataSources: AgentDataSourceConfigurationType[]; - query: RetrievalQuery; - relativeTimeFrame: RetrievalTimeframe; - topK: number; -}; -export function isRetrievalConfiguration( - arg: AgentActionConfigurationType | null -): arg is RetrievalConfigurationType { - return arg !== null && arg.type && arg.type === "retrieval_configuration"; -} - -/** - * Datasources config for Retrieval Action - */ -export type AgentDataSourceConfigurationType = { - workspaceId: string; // need sId to talk with Core (external id) - dataSourceId: string; // need Datasource.name to talk with Core (external id) - filter: { - tags: { in: string[]; not: string[] } | null; - parents: { in: string[]; not: string[] } | null; - }; -}; From a99a376e374f6623c5951b4bb5e702661f0b5063 Mon Sep 17 00:00:00 2001 From: PopDaph Date: Fri, 8 Sep 2023 16:38:29 +0200 Subject: [PATCH 11/12] Split create / update --- front/lib/api/assistant/configuration.ts | 424 ++++++++++++++------ front/lib/models/assistant/configuration.ts | 2 - 2 files changed, 293 insertions(+), 133 deletions(-) diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts index 1406d06e1bab..24da2c104670 100644 --- a/front/lib/api/assistant/configuration.ts +++ b/front/lib/api/assistant/configuration.ts @@ -1,7 +1,7 @@ import { Op, Transaction } from "sequelize"; import { Authenticator } from "@app/lib/auth"; -import { front_sequelize } from "@app/lib/databases"; +import { front_sequelize, ModelId } from "@app/lib/databases"; import { AgentConfiguration, AgentDataSourceConfiguration, @@ -22,6 +22,7 @@ import { AgentActionConfigurationType, AgentConfigurationStatus, AgentConfigurationType, + AgentGenerationConfigurationType, } from "@app/types/assistant/configuration"; /** @@ -45,7 +46,7 @@ export async function getAgentConfiguration( throw new Error("Cannot find Agent: no workspace"); } - const generation = agent.generationConfigId + const generationConfig = agent.generationConfigId ? await AgentGenerationConfiguration.findOne({ where: { id: agent.generationConfigId, @@ -53,17 +54,17 @@ export async function getAgentConfiguration( }) : null; - const action = agent.retrievalConfigId + const actionConfig = agent.retrievalConfigId ? await AgentRetrievalConfiguration.findOne({ where: { id: agent.retrievalConfigId, }, }) : null; - const datasources = action?.id + const dataSourcesConfig = actionConfig?.id ? await AgentDataSourceConfiguration.findAll({ where: { - retrievalConfigId: action.id, + retrievalConfigId: actionConfig.id, }, }) : []; @@ -73,14 +74,19 @@ export async function getAgentConfiguration( name: agent.name, pictureUrl: agent.pictureUrl, status: agent.status, - action: action ? await _agentActionType(action, datasources) : null, - generation: generation + action: actionConfig + ? await renderAgentActionConfigurationType( + actionConfig, + dataSourcesConfig + ) + : null, + generation: generationConfig ? { - id: generation.id, - prompt: generation.prompt, + id: generationConfig.id, + prompt: generationConfig.prompt, model: { - providerId: generation.providerId, - modelId: generation.modelId, + providerId: generationConfig.providerId, + modelId: generationConfig.modelId, }, } : null, @@ -96,26 +102,14 @@ export async function createAgentConfiguration( name, pictureUrl, status, - generation, - action, + generationConfigId, + actionConfigId, }: { name: string; pictureUrl: string; status: AgentConfigurationStatus; - generation: { - prompt: string; - model: { - providerId: string; - modelId: string; - }; - } | null; - action: { - type: string; - query: RetrievalQuery; - timeframe: RetrievalTimeframe; - topK: number; - dataSources: AgentDataSourceConfigurationType[]; - } | null; + generationConfigId: ModelId | null; + actionConfigId: ModelId | null; } ): Promise { const owner = auth.workspace(); @@ -123,78 +117,158 @@ export async function createAgentConfiguration( throw new Error("Cannot create AgentConfiguration without workspace"); } + // Create Agent config + const agentConfig = await AgentConfiguration.create({ + sId: generateModelSId(), + status: status, + name: name, + pictureUrl: pictureUrl, + scope: "workspace", + workspaceId: owner.id, + generationConfigId: generationConfigId, + retrievalConfigId: actionConfigId, + }); + + // Return the config with Generation and Action if any + const generationConfig = agentConfig.generationConfigId + ? await AgentGenerationConfiguration.findOne({ + where: { + id: agentConfig.generationConfigId, + }, + }) + : null; + const actionConfig = agentConfig.retrievalConfigId + ? await AgentRetrievalConfiguration.findOne({ + where: { + id: agentConfig.retrievalConfigId, + }, + }) + : null; + const dataSourcesConfig = actionConfig?.id + ? await AgentDataSourceConfiguration.findAll({ + where: { + retrievalConfigId: actionConfig.id, + }, + }) + : []; + + return { + sId: agentConfig.sId, + name: agentConfig.name, + pictureUrl: agentConfig.pictureUrl, + status: agentConfig.status, + action: actionConfig + ? await renderAgentActionConfigurationType( + actionConfig, + dataSourcesConfig + ) + : null, + generation: generationConfig + ? { + id: generationConfig.id, + prompt: generationConfig.prompt, + model: { + providerId: generationConfig.providerId, + modelId: generationConfig.modelId, + }, + } + : null, + }; +} + +/** + * Create Agent Generation Configuration + */ +export async function createAgentGenerationConfiguration( + auth: Authenticator, + { + prompt, + model, + }: { + prompt: string; + model: { + providerId: string; + modelId: string; + }; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error("Cannot create AgentConfiguration without workspace"); + } + + const genConfig = await AgentGenerationConfiguration.create({ + prompt: prompt, + providerId: model.providerId, + modelId: model.modelId, + }); + + return { + id: genConfig.id, + prompt: genConfig.prompt, + model: { + providerId: genConfig.providerId, + modelId: genConfig.modelId, + }, + }; +} + +/** + * Create Agent RetrievalConfiguration + */ +export async function createAgentActionConfiguration( + auth: Authenticator, + { + type, + query, + timeframe, + topK, + dataSources, + }: { + type: string; + query: RetrievalQuery; + timeframe: RetrievalTimeframe; + topK: number; + dataSources: AgentDataSourceConfigurationType[]; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error("Cannot create AgentConfiguration without workspace"); + } + return await front_sequelize.transaction(async (t) => { - let genConfig: AgentGenerationConfiguration | null = null; let retrievalConfig: AgentRetrievalConfiguration | null = null; let dataSourcesConfig: AgentDataSourceConfiguration[] = []; - // Create Generation config - if (generation) { - const { prompt, model } = generation; - genConfig = await AgentGenerationConfiguration.create({ - prompt: prompt, - providerId: model.providerId, - modelId: model.modelId, - }); + if (type !== "retrieval_configuration") { + throw new Error("Unkown Agent action type"); } // Create Retrieval & Datasources configs - if (action && action.type === "retrieval_configuration") { - const { query, timeframe, topK, dataSources } = action; - retrievalConfig = await AgentRetrievalConfiguration.create( - { - query: isTemplatedQuery(query) ? "templated" : query, - queryTemplate: isTemplatedQuery(query) ? query.template : null, - relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, - relativeTimeFrameDuration: isTimeFrame(timeframe) - ? timeframe.duration - : null, - relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, - topK: topK, - }, - { transaction: t } - ); - dataSourcesConfig = await _createAgentDataSourcesConfigData( - t, - dataSources, - retrievalConfig.id - ); - } - - // Create Agent config - const agentConfig = await AgentConfiguration.create( + retrievalConfig = await AgentRetrievalConfiguration.create( { - sId: generateModelSId(), - status: status, - name: name, - pictureUrl: pictureUrl, - scope: "workspace", - workspaceId: owner.id, - generationConfigId: genConfig?.id ?? null, - retrievalConfigId: retrievalConfig?.id ?? null, + query: isTemplatedQuery(query) ? "templated" : query, + queryTemplate: isTemplatedQuery(query) ? query.template : null, + relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, + relativeTimeFrameDuration: isTimeFrame(timeframe) + ? timeframe.duration + : null, + relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, + topK: topK, }, { transaction: t } ); + dataSourcesConfig = await _createAgentDataSourcesConfigData( + t, + dataSources, + retrievalConfig.id + ); - return { - sId: agentConfig.sId, - name: agentConfig.name, - pictureUrl: agentConfig.pictureUrl, - status: agentConfig.status, - action: retrievalConfig - ? await _agentActionType(retrievalConfig, dataSourcesConfig) - : null, - generation: genConfig - ? { - id: genConfig.id, - prompt: genConfig.prompt, - model: { - providerId: genConfig.providerId, - modelId: genConfig.modelId, - }, - } - : null, - }; + return await renderAgentActionConfigurationType( + retrievalConfig, + dataSourcesConfig + ); }); } @@ -202,6 +276,94 @@ export async function createAgentConfiguration( * Update Agent Configuration */ export async function updateAgentConfiguration( + auth: Authenticator, + agentId: string, + { + name, + pictureUrl, + status, + }: { + name: string; + pictureUrl: string; + status: AgentConfigurationStatus; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Workspace not found" + ); + } + const agentConfig = await AgentConfiguration.findOne({ + where: { + sId: agentId, + }, + }); + if (!agentConfig) { + throw new Error( + "Cannot create AgentGenerationConfiguration: Agent not found" + ); + } + // Updating Agent Config + const updatedAgentConfig = await agentConfig.update({ + name: name, + pictureUrl: pictureUrl, + status: status, + }); + + // Return the config with Generation and Action if any + const existingGeneration = agentConfig.generationConfigId + ? await AgentGenerationConfiguration.findOne({ + where: { + id: agentConfig.generationConfigId, + }, + }) + : null; + + const existingRetrivalConfig = agentConfig.retrievalConfigId + ? await AgentRetrievalConfiguration.findOne({ + where: { + id: agentConfig.retrievalConfigId, + }, + }) + : null; + + const existingDataSourcesConfig = existingRetrivalConfig?.id + ? await AgentDataSourceConfiguration.findAll({ + where: { + retrievalConfigId: existingRetrivalConfig.id, + }, + }) + : []; + + return { + sId: updatedAgentConfig.sId, + name: updatedAgentConfig.name, + pictureUrl: updatedAgentConfig.pictureUrl, + status: updatedAgentConfig.status, + action: existingRetrivalConfig + ? await renderAgentActionConfigurationType( + existingRetrivalConfig, + existingDataSourcesConfig + ) + : null, + generation: existingGeneration + ? { + id: existingGeneration.id, + prompt: existingGeneration.prompt, + model: { + providerId: existingGeneration.providerId, + modelId: existingGeneration.modelId, + }, + } + : null, + }; +} + +/** + * Update Agent Configuration + */ +export async function updateAgentGenerationConfiguration( auth: Authenticator, agentId: string, { @@ -306,7 +468,7 @@ export async function updateAgentConfiguration( pictureUrl: updatedAgentConfig.pictureUrl, status: updatedAgentConfig.status, action: existingRetrivalConfig - ? await _agentActionType( + ? await renderAgentActionConfigurationType( existingRetrivalConfig, existingDataSourcesConfig ) @@ -330,27 +492,33 @@ export async function updateAgentConfiguration( * Update Agent Retrieval Configuration * This will destroy and recreate the retrieval config */ -export async function updateAgentRetrievalConfiguration( +export async function updateAgentActionConfiguration( auth: Authenticator, agentId: string, { + type, query, timeframe, topK, dataSources, }: { + type: string; query: RetrievalQuery; timeframe: RetrievalTimeframe; topK: number; dataSources: AgentDataSourceConfigurationType[]; } -): Promise { +): Promise { const owner = auth.workspace(); if (!owner) { throw new Error( "Cannot create AgentGenerationConfiguration: Workspace not found" ); } + if (type !== "retrieval_configuration") { + throw new Error("Unkown Agent action type"); + } + const agentConfig = await AgentConfiguration.findOne({ where: { sId: agentId, @@ -361,27 +529,25 @@ export async function updateAgentRetrievalConfiguration( "Cannot create AgentGenerationConfiguration: Agent not found" ); } - const generationConfig = agentConfig.generationConfigId - ? await AgentGenerationConfiguration.findOne({ - where: { - id: agentConfig.generationConfigId, - }, - }) - : null; - return await front_sequelize.transaction(async (t) => { - if (agentConfig.retrievalConfigId) { - const existingRetrivalConfig = await AgentRetrievalConfiguration.findOne({ - where: { - id: agentConfig.retrievalConfigId, - }, - }); - if (existingRetrivalConfig) { - await existingRetrivalConfig.destroy(); // That will destroy the dataSourcesConfig too - } - } + if (!agentConfig.retrievalConfigId) { + throw new Error( + "Cannot update Agent Retrieval Configuration: Agent has no retrieval config" + ); + } + const existingRetrievalConfig = await AgentRetrievalConfiguration.findOne({ + where: { + id: agentConfig.retrievalConfigId, + }, + }); + if (!existingRetrievalConfig) { + throw new Error( + "Cannot update Agent Retrieval Configuration: Retrieval config not found" + ); + } - const newRetrievalConfig = await AgentRetrievalConfiguration.create( + return await front_sequelize.transaction(async (t) => { + const updatedRetrievalConfig = await existingRetrievalConfig.update( { query: isTemplatedQuery(query) ? "templated" : query, queryTemplate: isTemplatedQuery(query) ? query.template : null, @@ -394,38 +560,32 @@ export async function updateAgentRetrievalConfiguration( }, { transaction: t } ); + + // Destroy existing dataSources config + await AgentDataSourceConfiguration.destroy({ + where: { + retrievalConfigId: existingRetrievalConfig.id, + }, + }); + + // Create new dataSources config const dataSourcesConfig = await _createAgentDataSourcesConfigData( t, dataSources, - newRetrievalConfig.id + updatedRetrievalConfig.id ); - return { - sId: agentConfig.sId, - name: agentConfig.name, - pictureUrl: agentConfig.pictureUrl, - status: agentConfig.status, - action: newRetrievalConfig - ? await _agentActionType(newRetrievalConfig, dataSourcesConfig) - : null, - generation: generationConfig - ? { - id: generationConfig.id, - prompt: generationConfig.prompt, - model: { - providerId: generationConfig.providerId, - modelId: generationConfig.modelId, - }, - } - : null, - }; + return await renderAgentActionConfigurationType( + updatedRetrievalConfig, + dataSourcesConfig + ); }); } /** * Builds the agent action configuration type from the model */ -export async function _agentActionType( +async function renderAgentActionConfigurationType( action: AgentRetrievalConfiguration, dataSourcesConfig: AgentDataSourceConfiguration[] ): Promise { @@ -553,7 +713,9 @@ export async function _createAgentDataSourcesConfigData( // First we need to get the workspaceId from the workspaceSId const workspace = workspaces.find((w) => w.sId === curr.workspaceId); if (!workspace) { - throw new Error("Workspace not found"); + throw new Error( + "Can't create Datasources config for retrieval: Workspace not found" + ); } // Find an existing entry for this workspaceId diff --git a/front/lib/models/assistant/configuration.ts b/front/lib/models/assistant/configuration.ts index 4ce9a40c7c72..0cf66ac3de67 100644 --- a/front/lib/models/assistant/configuration.ts +++ b/front/lib/models/assistant/configuration.ts @@ -333,12 +333,10 @@ Workspace.hasMany(AgentConfiguration, { // Agent config <> Generation config AgentGenerationConfiguration.hasOne(AgentConfiguration, { foreignKey: { name: "generationConfigId", allowNull: true }, // null = no generation set for this Agent - onDelete: "CASCADE", }); // Agent config <> Retrieval config AgentRetrievalConfiguration.hasOne(AgentConfiguration, { foreignKey: { name: "retrievalConfigId", allowNull: true }, // null = no retrieval action set for this Agent - onDelete: "CASCADE", }); // Retrieval config <> Data source config From 32afccc0a419644be115421b01e9d0b179f40168 Mon Sep 17 00:00:00 2001 From: PopDaph Date: Fri, 8 Sep 2023 17:18:22 +0200 Subject: [PATCH 12/12] wip --- front/lib/api/assistant/configuration.ts | 395 +++++++++-------------- 1 file changed, 148 insertions(+), 247 deletions(-) diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts index 24da2c104670..f106c5c69698 100644 --- a/front/lib/api/assistant/configuration.ts +++ b/front/lib/api/assistant/configuration.ts @@ -34,7 +34,7 @@ export async function getAgentConfiguration( ): Promise { const owner = auth.workspace(); if (!owner) { - throw new Error("Cannot find Agent: no workspace"); + throw new Error("Cannot find AgentConfiguration: no workspace."); } const agent = await AgentConfiguration.findOne({ where: { @@ -43,7 +43,7 @@ export async function getAgentConfiguration( }, }); if (!agent) { - throw new Error("Cannot find Agent: no workspace"); + throw new Error("Cannot find AgentConfiguration."); } const generationConfig = agent.generationConfigId @@ -102,19 +102,19 @@ export async function createAgentConfiguration( name, pictureUrl, status, - generationConfigId, - actionConfigId, + generationConfig, + actionConfig, }: { name: string; pictureUrl: string; status: AgentConfigurationStatus; - generationConfigId: ModelId | null; - actionConfigId: ModelId | null; + generationConfig: AgentGenerationConfigurationType | null; + actionConfig: AgentActionConfigurationType | null; } ): Promise { const owner = auth.workspace(); if (!owner) { - throw new Error("Cannot create AgentConfiguration without workspace"); + throw new Error("Cannot create AgentConfiguration: Workspace not found."); } // Create Agent config @@ -125,153 +125,20 @@ export async function createAgentConfiguration( pictureUrl: pictureUrl, scope: "workspace", workspaceId: owner.id, - generationConfigId: generationConfigId, - retrievalConfigId: actionConfigId, + generationConfigId: generationConfig?.id || null, + retrievalConfigId: actionConfig?.id || null, }); - // Return the config with Generation and Action if any - const generationConfig = agentConfig.generationConfigId - ? await AgentGenerationConfiguration.findOne({ - where: { - id: agentConfig.generationConfigId, - }, - }) - : null; - const actionConfig = agentConfig.retrievalConfigId - ? await AgentRetrievalConfiguration.findOne({ - where: { - id: agentConfig.retrievalConfigId, - }, - }) - : null; - const dataSourcesConfig = actionConfig?.id - ? await AgentDataSourceConfiguration.findAll({ - where: { - retrievalConfigId: actionConfig.id, - }, - }) - : []; - return { sId: agentConfig.sId, name: agentConfig.name, pictureUrl: agentConfig.pictureUrl, status: agentConfig.status, - action: actionConfig - ? await renderAgentActionConfigurationType( - actionConfig, - dataSourcesConfig - ) - : null, - generation: generationConfig - ? { - id: generationConfig.id, - prompt: generationConfig.prompt, - model: { - providerId: generationConfig.providerId, - modelId: generationConfig.modelId, - }, - } - : null, - }; -} - -/** - * Create Agent Generation Configuration - */ -export async function createAgentGenerationConfiguration( - auth: Authenticator, - { - prompt, - model, - }: { - prompt: string; - model: { - providerId: string; - modelId: string; - }; - } -): Promise { - const owner = auth.workspace(); - if (!owner) { - throw new Error("Cannot create AgentConfiguration without workspace"); - } - - const genConfig = await AgentGenerationConfiguration.create({ - prompt: prompt, - providerId: model.providerId, - modelId: model.modelId, - }); - - return { - id: genConfig.id, - prompt: genConfig.prompt, - model: { - providerId: genConfig.providerId, - modelId: genConfig.modelId, - }, + action: actionConfig, + generation: generationConfig, }; } -/** - * Create Agent RetrievalConfiguration - */ -export async function createAgentActionConfiguration( - auth: Authenticator, - { - type, - query, - timeframe, - topK, - dataSources, - }: { - type: string; - query: RetrievalQuery; - timeframe: RetrievalTimeframe; - topK: number; - dataSources: AgentDataSourceConfigurationType[]; - } -): Promise { - const owner = auth.workspace(); - if (!owner) { - throw new Error("Cannot create AgentConfiguration without workspace"); - } - - return await front_sequelize.transaction(async (t) => { - let retrievalConfig: AgentRetrievalConfiguration | null = null; - let dataSourcesConfig: AgentDataSourceConfiguration[] = []; - - if (type !== "retrieval_configuration") { - throw new Error("Unkown Agent action type"); - } - - // Create Retrieval & Datasources configs - retrievalConfig = await AgentRetrievalConfiguration.create( - { - query: isTemplatedQuery(query) ? "templated" : query, - queryTemplate: isTemplatedQuery(query) ? query.template : null, - relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, - relativeTimeFrameDuration: isTimeFrame(timeframe) - ? timeframe.duration - : null, - relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, - topK: topK, - }, - { transaction: t } - ); - dataSourcesConfig = await _createAgentDataSourcesConfigData( - t, - dataSources, - retrievalConfig.id - ); - - return await renderAgentActionConfigurationType( - retrievalConfig, - dataSourcesConfig - ); - }); -} - /** * Update Agent Configuration */ @@ -291,7 +158,7 @@ export async function updateAgentConfiguration( const owner = auth.workspace(); if (!owner) { throw new Error( - "Cannot create AgentGenerationConfiguration: Workspace not found" + "Cannot create AgentGenerationConfiguration: Workspace not found."); ); } const agentConfig = await AgentConfiguration.findOne({ @@ -301,7 +168,7 @@ export async function updateAgentConfiguration( }); if (!agentConfig) { throw new Error( - "Cannot create AgentGenerationConfiguration: Agent not found" + "Cannot create AgentGenerationConfiguration: Agent not found."); ); } // Updating Agent Config @@ -361,33 +228,63 @@ export async function updateAgentConfiguration( } /** - * Update Agent Configuration + * Create Agent Generation Configuration + */ +export async function createAgentGenerationConfiguration( + auth: Authenticator, + { + prompt, + model, + }: { + prompt: string; + model: { + providerId: string; + modelId: string; + }; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error("Cannot create AgentGenerationConfiguration: Workspace not found."); + } + + const genConfig = await AgentGenerationConfiguration.create({ + prompt: prompt, + providerId: model.providerId, + modelId: model.modelId, + }); + + return { + id: genConfig.id, + prompt: genConfig.prompt, + model: { + providerId: genConfig.providerId, + modelId: genConfig.modelId, + }, + }; +} + +/** + * Update Agent Generation Configuration */ export async function updateAgentGenerationConfiguration( auth: Authenticator, agentId: string, { - name, - pictureUrl, - status, - generation, + prompt, + model, }: { - name: string; - pictureUrl: string; - status: AgentConfigurationStatus; - generation: { - prompt: string; - model: { - providerId: string; - modelId: string; - }; - } | null; + prompt: string; + model: { + providerId: string; + modelId: string; + }; } -): Promise { +): Promise { const owner = auth.workspace(); if (!owner) { throw new Error( - "Cannot create AgentGenerationConfiguration: Workspace not found" + "Cannot update AgentGenerationConfiguration: Workspace not found."); ); } const agentConfig = await AgentConfiguration.findOne({ @@ -397,94 +294,98 @@ export async function updateAgentGenerationConfiguration( }); if (!agentConfig) { throw new Error( - "Cannot create AgentGenerationConfiguration: Agent not found" + "Cannot update AgentGenerationConfiguration: Agent not found."); ); } - const existingGeneration = agentConfig.generationConfigId - ? await AgentGenerationConfiguration.findOne({ - where: { - id: agentConfig.generationConfigId, - }, - }) - : null; - const existingRetrivalConfig = agentConfig.retrievalConfigId - ? await AgentRetrievalConfiguration.findOne({ - where: { - id: agentConfig.retrievalConfigId, - }, - }) - : null; + if (!agentConfig.generationConfigId) { + throw new Error( + "Cannot update AgentGenerationConfiguration: Agent has no config."); + ); + } + const existingGenConfig = await AgentGenerationConfiguration.findOne({ + where: { + id: agentConfig.generationConfigId, + }, + }); + if (!existingGenConfig) { + throw new Error( + "Cannot update AgentGenerationConfiguration: config not found."); + ); + } - const existingDataSourcesConfig = existingRetrivalConfig?.id - ? await AgentDataSourceConfiguration.findAll({ - where: { - retrievalConfigId: existingRetrivalConfig.id, - }, - }) - : []; + const updatedGenConfig = await existingGenConfig.update({ + prompt: prompt, + providerId: model.providerId, + modelId: model.modelId, + }); + + return { + id: updatedGenConfig.id, + prompt: updatedGenConfig.prompt, + model: { + providerId: updatedGenConfig.providerId, + modelId: updatedGenConfig.modelId, + }, + }; +} + +/** + * Create Agent RetrievalConfiguration + */ +export async function createAgentActionConfiguration( + auth: Authenticator, + { + type, + query, + timeframe, + topK, + dataSources, + }: { + type: string; + query: RetrievalQuery; + timeframe: RetrievalTimeframe; + topK: number; + dataSources: AgentDataSourceConfigurationType[]; + } +): Promise { + const owner = auth.workspace(); + if (!owner) { + throw new Error("Cannot create AgentActionConfiguration: no workspace"); + } return await front_sequelize.transaction(async (t) => { - // Updating Agent Config - const updatedAgentConfig = await agentConfig.update( + let retrievalConfig: AgentRetrievalConfiguration | null = null; + let dataSourcesConfig: AgentDataSourceConfiguration[] = []; + + if (type !== "retrieval_configuration") { + throw new Error("Cannot create AgentActionConfiguration: unknow type"); + } + + // Create Retrieval & Datasources configs + retrievalConfig = await AgentRetrievalConfiguration.create( { - name: name, - pictureUrl: pictureUrl, - status: status, + query: isTemplatedQuery(query) ? "templated" : query, + queryTemplate: isTemplatedQuery(query) ? query.template : null, + relativeTimeFrame: isTimeFrame(timeframe) ? "custom" : timeframe, + relativeTimeFrameDuration: isTimeFrame(timeframe) + ? timeframe.duration + : null, + relativeTimeFrameUnit: isTimeFrame(timeframe) ? timeframe.unit : null, + topK: topK, }, { transaction: t } ); + dataSourcesConfig = await _createAgentDataSourcesConfigData( + t, + dataSources, + retrievalConfig.id + ); - // Upserting Generation Config - let upsertedGenerationConfig: AgentGenerationConfiguration | null = null; - if (generation) { - const { prompt, model } = generation; - if (existingGeneration) { - upsertedGenerationConfig = await existingGeneration.update( - { - prompt: prompt, - providerId: model.providerId, - modelId: model.modelId, - }, - { transaction: t } - ); - } else { - upsertedGenerationConfig = await AgentGenerationConfiguration.create( - { - prompt: prompt, - providerId: model.providerId, - modelId: model.modelId, - }, - { transaction: t } - ); - } - } else if (existingGeneration) { - await existingGeneration.destroy(); - } - - return { - sId: updatedAgentConfig.sId, - name: updatedAgentConfig.name, - pictureUrl: updatedAgentConfig.pictureUrl, - status: updatedAgentConfig.status, - action: existingRetrivalConfig - ? await renderAgentActionConfigurationType( - existingRetrivalConfig, - existingDataSourcesConfig - ) - : null, - generation: - generation && upsertedGenerationConfig - ? { - id: upsertedGenerationConfig.id, - prompt: upsertedGenerationConfig.prompt, - model: { - providerId: upsertedGenerationConfig.providerId, - modelId: upsertedGenerationConfig.modelId, - }, - } - : null, - }; + return await renderAgentActionConfigurationType( + retrievalConfig, + dataSourcesConfig + ); }); } @@ -512,7 +413,7 @@ export async function updateAgentActionConfiguration( const owner = auth.workspace(); if (!owner) { throw new Error( - "Cannot create AgentGenerationConfiguration: Workspace not found" + "Cannot update AgentActionConfiguration: Workspace not found."); ); } if (type !== "retrieval_configuration") { @@ -526,13 +427,13 @@ export async function updateAgentActionConfiguration( }); if (!agentConfig) { throw new Error( - "Cannot create AgentGenerationConfiguration: Agent not found" + "Cannot update AgentActionConfiguration: Agent not found."); ); } if (!agentConfig.retrievalConfigId) { throw new Error( - "Cannot update Agent Retrieval Configuration: Agent has no retrieval config" + "Cannot update AgentActionConfiguration: Agent has no retrieval config."); ); } const existingRetrievalConfig = await AgentRetrievalConfiguration.findOne({ @@ -542,7 +443,7 @@ export async function updateAgentActionConfiguration( }); if (!existingRetrievalConfig) { throw new Error( - "Cannot update Agent Retrieval Configuration: Retrieval config not found" + "Cannot update AgentActionConfiguration: config not found."); ); } @@ -639,7 +540,7 @@ async function renderAgentActionConfigurationType( workspace = workspaces.find((w) => w.id === dataSource?.workspaceId); if (!dataSource || !workspace) { - throw new Error("Could not find dataSource or workspace"); + throw new Error("Can't render Agent Retrieval dataSources: not found."); } dataSourcesConfigType.push({ @@ -675,7 +576,7 @@ async function renderAgentActionConfigurationType( * We need to fetch the dataSources from the database from that. * We obvisously need to do as few queries as possible. */ -export async function _createAgentDataSourcesConfigData( +async function _createAgentDataSourcesConfigData( t: Transaction, dataSourcesConfig: AgentDataSourceConfigurationType[], retrievalConfigId: number @@ -764,7 +665,7 @@ export async function _createAgentDataSourcesConfigData( workspaces.find((w) => w.sId === dsConfig.workspaceId)?.id ); if (!dataSource) { - throw new Error("DataSource not found"); + throw new Error("Can't create AgentDataSourcesConfig: datasource not found."); } return AgentDataSourceConfiguration.create( {