cleanup pg_graphql #1 (#7673)

## Context
This PR removes workspace-query-runner/builder in preparation for fully
deprecating pg_graphql

next steps: Remove from the setup and make a command to remove comments
on schema/tables related to pg_graphql
This commit is contained in:
Weiko
2024-10-14 14:19:13 +02:00
committed by GitHub
parent a64635a9db
commit efba3b14be
34 changed files with 138 additions and 2790 deletions

View File

@ -1,30 +0,0 @@
import { Module } from '@nestjs/common';
import { TypeOrmModule } from '@nestjs/typeorm';
import { WorkspaceDataSourceModule } from 'src/engine/workspace-datasource/workspace-datasource.module';
import { UserModule } from 'src/engine/core-modules/user/user.module';
import { AISQLQueryResolver } from 'src/engine/core-modules/ai-sql-query/ai-sql-query.resolver';
import { AISQLQueryService } from 'src/engine/core-modules/ai-sql-query/ai-sql-query.service';
import { FeatureFlagEntity } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
import { WorkspaceQueryRunnerModule } from 'src/engine/api/graphql/workspace-query-runner/workspace-query-runner.module';
import { LLMChatModelModule } from 'src/engine/core-modules/llm-chat-model/llm-chat-model.module';
import { EnvironmentModule } from 'src/engine/core-modules/environment/environment.module';
import { LLMTracingModule } from 'src/engine/core-modules/llm-tracing/llm-tracing.module';
import { ObjectMetadataModule } from 'src/engine/metadata-modules/object-metadata/object-metadata.module';
import { WorkspaceSyncMetadataModule } from 'src/engine/workspace-manager/workspace-sync-metadata/workspace-sync-metadata.module';
@Module({
imports: [
WorkspaceDataSourceModule,
WorkspaceQueryRunnerModule,
UserModule,
TypeOrmModule.forFeature([FeatureFlagEntity], 'core'),
LLMChatModelModule,
LLMTracingModule,
EnvironmentModule,
ObjectMetadataModule,
WorkspaceSyncMetadataModule,
],
exports: [],
providers: [AISQLQueryResolver, AISQLQueryService],
})
export class AISQLQueryModule {}

View File

@ -1,14 +0,0 @@
import { PromptTemplate } from '@langchain/core/prompts';
export const sqlGenerationPromptTemplate = PromptTemplate.fromTemplate<{
llmOutputJsonSchema: string;
sqlCreateTableStatements: string;
userQuestion: string;
}>(`Always respond following this JSON Schema: {llmOutputJsonSchema}
Based on the table schema below, write a PostgreSQL query that would answer the user's question. All column names must be enclosed in double quotes.
{sqlCreateTableStatements}
Question: {userQuestion}
SQL Query:`);

View File

@ -1,63 +0,0 @@
import { ForbiddenException, UseGuards } from '@nestjs/common';
import { Args, ArgsType, Field, Query, Resolver } from '@nestjs/graphql';
import { InjectRepository } from '@nestjs/typeorm';
import { Repository } from 'typeorm';
import { AISQLQueryService } from 'src/engine/core-modules/ai-sql-query/ai-sql-query.service';
import { AISQLQueryResult } from 'src/engine/core-modules/ai-sql-query/dtos/ai-sql-query-result.dto';
import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
import { FeatureFlagEntity } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
import { User } from 'src/engine/core-modules/user/user.entity';
import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
import { AuthUser } from 'src/engine/decorators/auth/auth-user.decorator';
import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
import { UserAuthGuard } from 'src/engine/guards/user-auth.guard';
import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
@ArgsType()
class GetAISQLQueryArgs {
@Field(() => String)
text: string;
}
@UseGuards(WorkspaceAuthGuard, UserAuthGuard)
@Resolver(() => AISQLQueryResult)
export class AISQLQueryResolver {
constructor(
private readonly aiSqlQueryService: AISQLQueryService,
@InjectRepository(FeatureFlagEntity, 'core')
private readonly featureFlagRepository: Repository<FeatureFlagEntity>,
) {}
@Query(() => AISQLQueryResult)
async getAISQLQuery(
@AuthWorkspace() { id: workspaceId }: Workspace,
@AuthUser() user: User,
@Args() { text }: GetAISQLQueryArgs,
) {
const isCopilotEnabledFeatureFlag =
await this.featureFlagRepository.findOneBy({
workspaceId,
key: FeatureFlagKey.IsCopilotEnabled,
value: true,
});
if (!isCopilotEnabledFeatureFlag?.value) {
throw new ForbiddenException(
`${FeatureFlagKey.IsCopilotEnabled} feature flag is disabled`,
);
}
const traceMetadata = {
userId: user.id,
userEmail: user.email,
};
return this.aiSqlQueryService.generateAndExecute(
workspaceId,
text,
traceMetadata,
);
}
}

View File

@ -1,250 +0,0 @@
import { Injectable, Logger } from '@nestjs/common';
import { StructuredOutputParser } from '@langchain/core/output_parsers';
import { RunnableSequence } from '@langchain/core/runnables';
import groupBy from 'lodash.groupby';
import { DataSource, QueryFailedError } from 'typeorm';
import { PostgresConnectionOptions } from 'typeorm/driver/postgres/PostgresConnectionOptions';
import { z } from 'zod';
import { zodToJsonSchema } from 'zod-to-json-schema';
import { WorkspaceQueryRunnerService } from 'src/engine/api/graphql/workspace-query-runner/workspace-query-runner.service';
import { sqlGenerationPromptTemplate } from 'src/engine/core-modules/ai-sql-query/ai-sql-query.prompt-templates';
import { AISQLQueryResult } from 'src/engine/core-modules/ai-sql-query/dtos/ai-sql-query-result.dto';
import { LLMChatModelService } from 'src/engine/core-modules/llm-chat-model/llm-chat-model.service';
import { LLMTracingService } from 'src/engine/core-modules/llm-tracing/llm-tracing.service';
import { DEFAULT_LABEL_IDENTIFIER_FIELD_NAME } from 'src/engine/metadata-modules/object-metadata/object-metadata.constants';
import { ObjectMetadataEntity } from 'src/engine/metadata-modules/object-metadata/object-metadata.entity';
import { WorkspaceDataSourceService } from 'src/engine/workspace-datasource/workspace-datasource.service';
import { StandardObjectFactory } from 'src/engine/workspace-manager/workspace-sync-metadata/factories/standard-object.factory';
@Injectable()
export class AISQLQueryService {
private readonly logger = new Logger(AISQLQueryService.name);
constructor(
private readonly workspaceDataSourceService: WorkspaceDataSourceService,
private readonly workspaceQueryRunnerService: WorkspaceQueryRunnerService,
private readonly llmChatModelService: LLMChatModelService,
private readonly llmTracingService: LLMTracingService,
private readonly standardObjectFactory: StandardObjectFactory,
) {}
private getLabelIdentifierName(
objectMetadata: ObjectMetadataEntity,
_dataSourceId,
_workspaceId,
_workspaceFeatureFlagsMap,
): string | undefined {
const customObjectLabelIdentifierFieldMetadata = objectMetadata.fields.find(
(fieldMetadata) =>
fieldMetadata.id === objectMetadata.labelIdentifierFieldMetadataId,
);
/* const standardObjectMetadataCollection = this.standardObjectFactory.create(
standardObjectMetadataDefinitions,
{ workspaceId, dataSourceId },
workspaceFeatureFlagsMap,
);
const standardObjectLabelIdentifierFieldMetadata =
standardObjectMetadataCollection
.find(
(standardObjectMetadata) =>
standardObjectMetadata.nameSingular === objectMetadata.nameSingular,
)
?.fields.find(
(field: PartialFieldMetadata) =>
field.name === DEFAULT_LABEL_IDENTIFIER_FIELD_NAME,
) as PartialFieldMetadata; */
const labelIdentifierFieldMetadata =
customObjectLabelIdentifierFieldMetadata; /*??
standardObjectLabelIdentifierFieldMetadata*/
return (
labelIdentifierFieldMetadata?.name ?? DEFAULT_LABEL_IDENTIFIER_FIELD_NAME
);
}
private async getColInfosByTableName(dataSource: DataSource) {
const { schema } = dataSource.options as PostgresConnectionOptions;
// From LangChain sql_utils.ts
const sqlQuery = `SELECT
t.table_name,
c.*
FROM
information_schema.tables t
JOIN information_schema.columns c
ON t.table_name = c.table_name
WHERE
t.table_schema = '${schema}'
AND c.table_schema = '${schema}'
ORDER BY
t.table_name,
c.ordinal_position;`;
const colInfos = await dataSource.query<
{
table_name: string;
column_name: string;
data_type: string | undefined;
is_nullable: 'YES' | 'NO';
}[]
>(sqlQuery);
return groupBy(colInfos, (colInfo) => colInfo.table_name);
}
private getCreateTableStatement(tableName: string, colInfos: any[]) {
return `CREATE TABLE ${tableName} (\n ${colInfos
.map(
(colInfo) =>
`${colInfo.column_name} ${colInfo.data_type} ${
colInfo.is_nullable === 'YES' ? '' : 'NOT NULL'
}`,
)
.join(', ')});`;
}
private getRelationDescriptions() {
// TODO - Construct sentences like the following:
// investorId: a foreign key referencing the person table, indicating the investor who owns this portfolio company.
return '';
}
private getTableDescription(tableName: string, colInfos: any[]) {
return [
this.getCreateTableStatement(tableName, colInfos),
this.getRelationDescriptions(),
].join('\n');
}
private async getWorkspaceSchemaDescription(
dataSource: DataSource,
): Promise<string> {
const colInfoByTableName = await this.getColInfosByTableName(dataSource);
return Object.entries(colInfoByTableName)
.map(([tableName, colInfos]) =>
this.getTableDescription(tableName, colInfos),
)
.join('\n\n');
}
private async generateWithDataSource(
workspaceId: string,
workspaceDataSource: DataSource,
userQuestion: string,
traceMetadata: Record<string, string> = {},
) {
const workspaceSchemaName =
this.workspaceDataSourceService.getSchemaName(workspaceId);
workspaceDataSource.setOptions({
schema: workspaceSchemaName,
});
const workspaceSchemaDescription =
await this.getWorkspaceSchemaDescription(workspaceDataSource);
const llmOutputSchema = z.object({
sqlQuery: z.string(),
});
const llmOutputJsonSchema = JSON.stringify(
zodToJsonSchema(llmOutputSchema),
);
const structuredOutputParser =
StructuredOutputParser.fromZodSchema(llmOutputSchema);
const sqlQueryGeneratorChain = RunnableSequence.from([
sqlGenerationPromptTemplate,
this.llmChatModelService.getJSONChatModel(),
structuredOutputParser,
]);
const metadata = {
workspaceId,
...traceMetadata,
};
const tracingCallbackHandler =
this.llmTracingService.getCallbackHandler(metadata);
const { sqlQuery } = await sqlQueryGeneratorChain.invoke(
{
llmOutputJsonSchema,
sqlCreateTableStatements: workspaceSchemaDescription,
userQuestion,
},
{
callbacks: [tracingCallbackHandler],
},
);
return sqlQuery;
}
async generate(
workspaceId: string,
userQuestion: string,
traceMetadata: Record<string, string> = {},
) {
const workspaceDataSource =
await this.workspaceDataSourceService.connectToWorkspaceDataSource(
workspaceId,
);
return this.generateWithDataSource(
workspaceId,
workspaceDataSource,
userQuestion,
traceMetadata,
);
}
async generateAndExecute(
workspaceId: string,
userQuestion: string,
traceMetadata: Record<string, string> = {},
): Promise<AISQLQueryResult> {
const workspaceDataSource =
await this.workspaceDataSourceService.connectToWorkspaceDataSource(
workspaceId,
);
const sqlQuery = await this.generateWithDataSource(
workspaceId,
workspaceDataSource,
userQuestion,
traceMetadata,
);
try {
const sqlQueryResult: Record<string, any>[] =
await this.workspaceQueryRunnerService.executeSQL(
workspaceDataSource,
workspaceId,
sqlQuery,
);
return {
sqlQuery,
sqlQueryResult: JSON.stringify(sqlQueryResult),
};
} catch (error) {
if (error instanceof QueryFailedError) {
return {
sqlQuery,
queryFailedErrorMessage: error.message,
};
}
this.logger.error(error.message, error.stack);
return {
sqlQuery,
};
}
}
}

View File

@ -1,17 +0,0 @@
import { Field, ObjectType } from '@nestjs/graphql';
import { IsOptional } from 'class-validator';
@ObjectType('AISQLQueryResult')
export class AISQLQueryResult {
@Field(() => String)
sqlQuery: string;
@Field(() => String, { nullable: true })
@IsOptional()
sqlQueryResult?: string;
@Field(() => String, { nullable: true })
@IsOptional()
queryFailedErrorMessage?: string;
}

View File

@ -3,7 +3,6 @@ import { HttpAdapterHost } from '@nestjs/core';
import { EventEmitterModule } from '@nestjs/event-emitter';
import { ActorModule } from 'src/engine/core-modules/actor/actor.module';
import { AISQLQueryModule } from 'src/engine/core-modules/ai-sql-query/ai-sql-query.module';
import { AppTokenModule } from 'src/engine/core-modules/app-token/app-token.module';
import { AuthModule } from 'src/engine/core-modules/auth/auth.module';
import { BillingModule } from 'src/engine/core-modules/billing/billing.module';
@ -62,7 +61,6 @@ import { FileModule } from './file/file.module';
UserModule,
WorkspaceModule,
WorkspaceInvitationModule,
AISQLQueryModule,
PostgresCredentialsModule,
WorkflowTriggerApiModule,
WorkspaceEventEmitterModule,

View File

@ -1,11 +0,0 @@
import { Module } from '@nestjs/common';
import { DuplicateService } from 'src/engine/core-modules/duplicate/duplicate.service';
import { WorkspaceDataSourceModule } from 'src/engine/workspace-datasource/workspace-datasource.module';
@Module({
imports: [WorkspaceDataSourceModule],
exports: [DuplicateService],
providers: [DuplicateService],
})
export class DuplicateModule {}

View File

@ -1,97 +0,0 @@
import { Injectable } from '@nestjs/common';
import {
Record as IRecord,
Record,
} from 'src/engine/api/graphql/workspace-query-builder/interfaces/record.interface';
import { ObjectMetadataInterface } from 'src/engine/metadata-modules/field-metadata/interfaces/object-metadata.interface';
import { settings } from 'src/engine/constants/settings';
import { DUPLICATE_CRITERIA_COLLECTION } from 'src/engine/core-modules/duplicate/constants/duplicate-criteria.constants';
import { computeObjectTargetTable } from 'src/engine/utils/compute-object-target-table.util';
import { WorkspaceDataSourceService } from 'src/engine/workspace-datasource/workspace-datasource.service';
@Injectable()
export class DuplicateService {
constructor(
private readonly workspaceDataSourceService: WorkspaceDataSourceService,
) {}
async findExistingRecords(
recordIds: (string | number)[],
objectMetadata: ObjectMetadataInterface,
workspaceId: string,
) {
const dataSourceSchema =
this.workspaceDataSourceService.getSchemaName(workspaceId);
const results = await this.workspaceDataSourceService.executeRawQuery(
`
SELECT
*
FROM
${dataSourceSchema}."${computeObjectTargetTable(
objectMetadata,
)}" p
WHERE
p."id" IN (${recordIds
.map((_, index) => `$${index + 1}`)
.join(', ')})
`,
recordIds,
workspaceId,
);
return results as IRecord[];
}
buildDuplicateConditionForGraphQL(
objectMetadata: ObjectMetadataInterface,
argsData?: Partial<Record>,
filteringByExistingRecordId?: string,
) {
if (!argsData) {
return;
}
const criteriaCollection =
this.getApplicableDuplicateCriteriaCollection(objectMetadata);
const criteriaWithMatchingArgs = criteriaCollection.filter((criteria) =>
criteria.columnNames.every((columnName) => {
const value = argsData[columnName] as string | undefined;
return (
!!value && value.length >= settings.minLengthOfStringForDuplicateCheck
);
}),
);
const filterCriteria = criteriaWithMatchingArgs.map((criteria) =>
Object.fromEntries(
criteria.columnNames.map((columnName) => [
columnName,
{ eq: argsData[columnName] },
]),
),
);
return {
// when filtering by an existing record, we need to filter that explicit record out
...(filteringByExistingRecordId && {
id: { neq: filteringByExistingRecordId },
}),
// keep condition as "or" to get results by more duplicate criteria
or: filterCriteria,
};
}
private getApplicableDuplicateCriteriaCollection(
objectMetadataItem: ObjectMetadataInterface,
) {
return DUPLICATE_CRITERIA_COLLECTION.filter(
(duplicateCriteria) =>
duplicateCriteria.objectName === objectMetadataItem.nameSingular,
);
}
}