feat: tested proxy usage limits

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:

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}`,
     },
   })

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,
       },

packages/server/api/src/app/database/migration/common/1724777988812-CreateProxyConfig.ts → 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<void> {
@@ -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";`);
     }
 
 }

packages/server/api/src/app/ee/project-plan/project-plan.entity.ts

@@ -39,6 +39,11 @@ export const ProjectPlanEntity = new EntitySchema<ProjectPlanSchema>({
         tasks: {
             type: Number,
         },
+        aiTokens: {
+            type: Number,
+            default: 1000,
+            nullable: true
+        },
     },
     indices: [
         {

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<number> {
+    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<number> {
+    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<number> {
+    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}`
+}

packages/server/api/src/app/proxy/proxy-config-service.ts

@@ -10,9 +10,6 @@ export const proxyConfigService = {
     platformId,
     provider,
   }: GetParams): Promise<ProxyConfig> {
-    console.log("XXXXXXXXXX provider", provider)
-    console.log("XXXXXXXXXX provider", provider)
-    console.log("XXXXXXXXXX platformId", platformId)
     const config = await repo().findOneBy({
       platformId,
       provider,

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)
 }

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<typeof ProjectPlan>