From 5fda8c85652dd3ad9b3473ce1060a37220e33188 Mon Sep 17 00:00:00 2001 From: Anas Barghouthy Date: Wed, 4 Sep 2024 11:37:06 +0300 Subject: [PATCH] feat: tested proxy usage limits --- docker-compose.dev.yml | 23 +++++++++++ .../src/lib/ai/providers/anthropic/index.ts | 2 +- packages/react-ui/vite.config.ts | 4 +- .../1724777988812-CreateProxyConfig.ts | 2 + .../ee/project-plan/project-plan.entity.ts | 5 +++ .../project/usage/project-usage-service.ts | 32 +++++++++++++++ .../api/src/app/proxy/proxy-config-service.ts | 3 -- .../server/api/src/app/proxy/proxy-module.ts | 40 +++++++++++++++++-- packages/shared/src/lib/project/project.ts | 1 + 9 files changed, 102 insertions(+), 10 deletions(-) create mode 100644 docker-compose.dev.yml rename packages/server/api/src/app/database/migration/{common => postgres}/1724777988812-CreateProxyConfig.ts (89%) diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml new file mode 100644 index 0000000000..9a17eb75ae --- /dev/null +++ b/docker-compose.dev.yml @@ -0,0 +1,23 @@ +services: + db: + image: postgres:14.4 + environment: + POSTGRES_DB: activepieces + POSTGRES_USER: postgres + POSTGRES_PASSWORD: A79Vm5D4p2VQHOp2gd5 + + volumes: + - postgres_data:/var/lib/postgresql/data + ports: + - "5432:5432" + + redis: + image: redis:7.0.7 + volumes: + - redis_data:/data + ports: + - "6379:6379" + +volumes: + postgres_data: + redis_data: \ No newline at end of file diff --git a/packages/pieces/community/common/src/lib/ai/providers/anthropic/index.ts b/packages/pieces/community/common/src/lib/ai/providers/anthropic/index.ts index 9f415fb482..aa6e837efb 100644 --- a/packages/pieces/community/common/src/lib/ai/providers/anthropic/index.ts +++ b/packages/pieces/community/common/src/lib/ai/providers/anthropic/index.ts @@ -11,7 +11,7 @@ export const anthropic = ({ apiKey: engineToken, baseURL: proxyUrl, defaultHeaders: { - 'X-AP-TOTAL-USAGE-BODY-PATH': 'usage.total_tokens', + 'X-AP-TOTAL-USAGE-BODY-PATH': 'usage.output_tokens+usage.input_tokens', 'Authorization': `Bearer ${engineToken}`, }, }) diff --git a/packages/react-ui/vite.config.ts b/packages/react-ui/vite.config.ts index daa52d75a8..505dffcc31 100644 --- a/packages/react-ui/vite.config.ts +++ b/packages/react-ui/vite.config.ts @@ -13,12 +13,12 @@ export default defineConfig({ server: { proxy: { '/api': { - target: 'http://localhost:3000', + target: 'http://127.0.0.1:3000', secure: false, changeOrigin: true, rewrite: (path) => path.replace(/^\/api/, ''), headers: { - Host: 'localhost:4200', + Host: '127.0.0.1:4200', }, ws: true, }, diff --git a/packages/server/api/src/app/database/migration/common/1724777988812-CreateProxyConfig.ts b/packages/server/api/src/app/database/migration/postgres/1724777988812-CreateProxyConfig.ts similarity index 89% rename from packages/server/api/src/app/database/migration/common/1724777988812-CreateProxyConfig.ts rename to packages/server/api/src/app/database/migration/postgres/1724777988812-CreateProxyConfig.ts index 304323fc70..ee928c5095 100644 --- a/packages/server/api/src/app/database/migration/common/1724777988812-CreateProxyConfig.ts +++ b/packages/server/api/src/app/database/migration/postgres/1724777988812-CreateProxyConfig.ts @@ -21,6 +21,7 @@ export class CreateProxyConfig1724777988812 implements MigrationInterface { ALTER TABLE "proxy_config" ADD CONSTRAINT "fk_proxy_config_platform_id" FOREIGN KEY ("platformId") REFERENCES "platform"("id") ON DELETE NO ACTION ON UPDATE NO ACTION `); + await queryRunner.query(`ALTER TABLE "project_plan" ADD COLUMN "aiTokens" integer DEFAULT 1000;`); } public async down(queryRunner: QueryRunner): Promise { @@ -33,6 +34,7 @@ export class CreateProxyConfig1724777988812 implements MigrationInterface { await queryRunner.query(` DROP TABLE "proxy_config" `); + await queryRunner.query(`ALTER TABLE "project_plan" DROP COLUMN "aiTokens";`); } } diff --git a/packages/server/api/src/app/ee/project-plan/project-plan.entity.ts b/packages/server/api/src/app/ee/project-plan/project-plan.entity.ts index e6a238c188..7629e20d5c 100644 --- a/packages/server/api/src/app/ee/project-plan/project-plan.entity.ts +++ b/packages/server/api/src/app/ee/project-plan/project-plan.entity.ts @@ -39,6 +39,11 @@ export const ProjectPlanEntity = new EntitySchema({ tasks: { type: Number, }, + aiTokens: { + type: Number, + default: 1000, + nullable: true + }, }, indices: [ { diff --git a/packages/server/api/src/app/project/usage/project-usage-service.ts b/packages/server/api/src/app/project/usage/project-usage-service.ts index 260c4fae26..827c98056e 100644 --- a/packages/server/api/src/app/project/usage/project-usage-service.ts +++ b/packages/server/api/src/app/project/usage/project-usage-service.ts @@ -18,6 +18,8 @@ export const projectUsageService = { increaseTasks, getCurrentingStartPeriod, getCurrentingEndPeriod, + increaseAITokens, + getAITokensUsage, } function getCurrentingStartPeriod(datetime: string): string { @@ -57,6 +59,36 @@ async function getTasksUsage(projectId: string, startBillingPeriod: string): Pro return Number(value) || 0 } + +async function increaseAITokens(projectId: string, incrementBy: number): Promise { + const project = await projectService.getOneOrThrow(projectId) + const startBillingPeriod = getCurrentingStartPeriod(project.created) + return incrementOrCreateAITokensRedisRecord(projectId, startBillingPeriod, incrementBy) +} + +async function incrementOrCreateAITokensRedisRecord(projectId: string, startBillingPeriod: string, incrementBy: number): Promise { + const environment = system.get(SharedSystemProp.ENVIRONMENT) + if (environment === ApEnvironment.TESTING) { + return 0 + } + const key = constructAITokensUsageKey(projectId, startBillingPeriod) + return getRedisConnection().incrby(key, incrementBy) +} + +async function getAITokensUsage(projectId: string, startBillingPeriod: string): Promise { + const environment = system.get(SharedSystemProp.ENVIRONMENT) + if (environment === ApEnvironment.TESTING) { + return 0 + } + const key = constructAITokensUsageKey(projectId, startBillingPeriod) + const value = await getRedisConnection().get(key) + return Number(value) || 0 +} + function constructUsageKey(projectId: string, startBillingPeriod: string): string { return `project-usage:${projectId}:${startBillingPeriod}` } + +function constructAITokensUsageKey(projectId: string, startBillingPeriod: string): string { + return `project-ai-tokens-usage:${projectId}:${startBillingPeriod}` +} \ No newline at end of file diff --git a/packages/server/api/src/app/proxy/proxy-config-service.ts b/packages/server/api/src/app/proxy/proxy-config-service.ts index 1526e258a5..f712a50843 100644 --- a/packages/server/api/src/app/proxy/proxy-config-service.ts +++ b/packages/server/api/src/app/proxy/proxy-config-service.ts @@ -10,9 +10,6 @@ export const proxyConfigService = { platformId, provider, }: GetParams): Promise { - console.log("XXXXXXXXXX provider", provider) - console.log("XXXXXXXXXX provider", provider) - console.log("XXXXXXXXXX platformId", platformId) const config = await repo().findOneBy({ platformId, provider, diff --git a/packages/server/api/src/app/proxy/proxy-module.ts b/packages/server/api/src/app/proxy/proxy-module.ts index a4d740eee7..cdbee2aa31 100644 --- a/packages/server/api/src/app/proxy/proxy-module.ts +++ b/packages/server/api/src/app/proxy/proxy-module.ts @@ -1,10 +1,9 @@ -import { PrincipalType } from '@activepieces/shared' +import { isNil, PrincipalType } from '@activepieces/shared' import { FastifyPluginAsyncTypebox, FastifyPluginCallbackTypebox, Type } from '@fastify/type-provider-typebox' -import axios, { AxiosRequestConfig } from 'axios' -import { StatusCodes } from 'http-status-codes' +import { projectLimitsService } from 'packages/server/api/src/app/ee/project-plan/project-plan.service' import { projectService } from 'packages/server/api/src/app/project/project-service' +import { projectUsageService } from 'packages/server/api/src/app/project/usage/project-usage-service' import { proxyConfigService } from './proxy-config-service' -import https from 'https' export const proxyModule: FastifyPluginAsyncTypebox = async (app) => { await app.register(projectProxyController, { prefix: '/v1/proxy' }) @@ -33,6 +32,10 @@ export const projectProxyController: FastifyPluginCallbackTypebox = ( try { const projectId = request.principal.projectId + const projectPlan = await projectLimitsService.getPlanByProjectId(projectId) + + const planTokens = projectPlan?.aiTokens + const platformId = await projectService.getPlatformId(projectId) const provider = request.params.provider @@ -48,6 +51,10 @@ export const projectProxyController: FastifyPluginCallbackTypebox = ( const requestHeaders = structuredClone({ ...request.headers }) + const tokensUsagePath = requestHeaders['x-ap-total-usage-body-path'] + + console.log('XXXXXXXXXX Request Headers', requestHeaders) + delete requestHeaders['authorization'] delete requestHeaders['Authorization'] delete requestHeaders['content-length'] @@ -74,6 +81,18 @@ export const projectProxyController: FastifyPluginCallbackTypebox = ( const data = await response.json() + const usage = projectPlan && tokensUsagePath ? calculateUsage(data, tokensUsagePath) : null + + const totalTokensUsage = usage ? await projectUsageService.increaseAITokens(projectId, usage) : null + + if (isNil(totalTokensUsage)) { + throw new Error(`Failed to calculate usage from response headers ${tokensUsagePath}`) + } + + if (!isNil(planTokens) && totalTokensUsage > planTokens) { + throw new Error(`You have exceeded your plan limit of ${planTokens} tokens`) + } + await reply .code(response.status) .send(data); @@ -84,4 +103,17 @@ export const projectProxyController: FastifyPluginCallbackTypebox = ( }) done() +} + +const calculateUsage = (body: any, usagePath: string | string[]): number => { + const fields = typeof usagePath === "string" ? usagePath.split('+').map(field => field.trim()) : usagePath + + return fields.reduce((acc, field) => { + const fieldPath = field.split('.') + const value = fieldPath.reduce((acc, field) => acc[field], body) + if (typeof value !== 'number') { + return acc + } + return acc + value + }, 0) } \ No newline at end of file diff --git a/packages/shared/src/lib/project/project.ts b/packages/shared/src/lib/project/project.ts index dcbd1e4f12..44675d82ab 100755 --- a/packages/shared/src/lib/project/project.ts +++ b/packages/shared/src/lib/project/project.ts @@ -49,6 +49,7 @@ export const ProjectPlan = Type.Object({ connections: Type.Number(), teamMembers: Type.Number(), tasks: Type.Number(), + aiTokens: Type.Optional(Type.Number()), }) export type ProjectPlan = Static