mirror of
https://github.com/rowboatlabs/rowboat.git
synced 2026-05-09 07:12:39 +02:00
refactor usage quota service as a policy
This commit is contained in:
parent
5b6d592d09
commit
831356a155
9 changed files with 37 additions and 37 deletions
|
|
@ -14,9 +14,9 @@ import { USE_BILLING } from "../lib/feature_flags";
|
|||
import { WithStringId } from "../lib/types/types";
|
||||
import { getEditAgentInstructionsResponse } from "../lib/copilot/copilot";
|
||||
import { container } from "@/di/container";
|
||||
import { IUsageQuotaPolicyService } from "@/src/application/services/usage-quota-policy.service.interface";
|
||||
import { IUsageQuotaPolicy } from "@/src/application/policies/usage-quota.policy.interface";
|
||||
|
||||
const usageQuotaPolicyService = container.resolve<IUsageQuotaPolicyService>('usageQuotaPolicyService');
|
||||
const usageQuotaPolicy = container.resolve<IUsageQuotaPolicy>('usageQuotaPolicy');
|
||||
|
||||
export async function getCopilotResponseStream(
|
||||
projectId: string,
|
||||
|
|
@ -28,7 +28,7 @@ export async function getCopilotResponseStream(
|
|||
streamId: string;
|
||||
} | { billingError: string }> {
|
||||
await projectAuthCheck(projectId);
|
||||
await usageQuotaPolicyService.assertAndConsume(projectId);
|
||||
await usageQuotaPolicy.assertAndConsume(projectId);
|
||||
|
||||
// Check billing authorization
|
||||
const authResponse = await authorizeUserAction({
|
||||
|
|
@ -39,7 +39,7 @@ export async function getCopilotResponseStream(
|
|||
return { billingError: authResponse.error || 'Billing error' };
|
||||
}
|
||||
|
||||
await usageQuotaPolicyService.assertAndConsume(projectId);
|
||||
await usageQuotaPolicy.assertAndConsume(projectId);
|
||||
|
||||
// prepare request
|
||||
const request: z.infer<typeof CopilotAPIRequest> = {
|
||||
|
|
@ -71,7 +71,7 @@ export async function getCopilotAgentInstructions(
|
|||
agentName: string,
|
||||
): Promise<string | { billingError: string }> {
|
||||
await projectAuthCheck(projectId);
|
||||
await usageQuotaPolicyService.assertAndConsume(projectId);
|
||||
await usageQuotaPolicy.assertAndConsume(projectId);
|
||||
|
||||
// Check billing authorization
|
||||
const authResponse = await authorizeUserAction({
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import { authorize, getCustomerIdForProject, logUsage } from "@/app/lib/billing"
|
|||
import { USE_BILLING } from "@/app/lib/feature_flags";
|
||||
import { getResponse } from "@/app/lib/agents";
|
||||
import { Message, AssistantMessage, AssistantMessageWithToolCalls, ToolMessage } from "@/app/lib/types/types";
|
||||
import { IUsageQuotaPolicyService } from "@/src/application/services/usage-quota-policy.service.interface";
|
||||
import { IUsageQuotaPolicy } from "@/src/application/policies/usage-quota.policy.interface";
|
||||
import { container } from "@/di/container";
|
||||
|
||||
function convert(messages: z.infer<typeof apiV1.ChatMessage>[]): z.infer<typeof Message>[] {
|
||||
|
|
@ -125,8 +125,8 @@ export async function POST(
|
|||
}
|
||||
|
||||
// assert and consume quota
|
||||
const usageQuotaPolicyService = container.resolve<IUsageQuotaPolicyService>('usageQuotaPolicyService');
|
||||
await usageQuotaPolicyService.assertAndConsume(session.projectId);
|
||||
const usageQuotaPolicy = container.resolve<IUsageQuotaPolicy>('usageQuotaPolicy');
|
||||
await usageQuotaPolicy.assertAndConsume(session.projectId);
|
||||
|
||||
// parse and validate the request body
|
||||
let body;
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import { CreateCachedTurnUseCase } from "@/src/application/use-cases/conversatio
|
|||
import { FetchCachedTurnUseCase } from "@/src/application/use-cases/conversations/fetch-cached-turn.use-case";
|
||||
import { CreateCachedTurnController } from "@/src/interface-adapters/controllers/conversations/create-cached-turn.controller";
|
||||
import { RunTurnController } from "@/src/interface-adapters/controllers/conversations/run-turn.controller";
|
||||
import { RedisUsageQuotaPolicyService } from "@/src/infrastructure/services/redis.usage-quota-policy.service";
|
||||
import { RedisUsageQuotaPolicy } from "@/src/infrastructure/policies/redis.usage-quota.policy";
|
||||
|
||||
export const container = createContainer({
|
||||
injectionMode: InjectionMode.PROXY,
|
||||
|
|
@ -20,7 +20,7 @@ container.register({
|
|||
// services
|
||||
// ---
|
||||
cacheService: asClass(RedisCacheService).singleton(),
|
||||
usageQuotaPolicyService: asClass(RedisUsageQuotaPolicyService).singleton(),
|
||||
usageQuotaPolicy: asClass(RedisUsageQuotaPolicy).singleton(),
|
||||
|
||||
// conversations
|
||||
// ---
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
export interface IUsageQuotaPolicyService {
|
||||
export interface IUsageQuotaPolicy {
|
||||
// this method will throw a QuotaExceededError if the quota is exceeded
|
||||
assertAndConsume(projectId: string): Promise<void>;
|
||||
}
|
||||
|
|
@ -5,7 +5,7 @@ import { z } from "zod";
|
|||
import { nanoid } from 'nanoid';
|
||||
import { ICacheService } from '@/src/application/services/cache.service.interface';
|
||||
import { CachedTurnRequest, Turn } from '@/src/entities/models/turn';
|
||||
import { IUsageQuotaPolicyService } from '../../services/usage-quota-policy.service.interface';
|
||||
import { IUsageQuotaPolicy } from '../../policies/usage-quota.policy.interface';
|
||||
|
||||
const inputSchema = z.object({
|
||||
caller: z.enum(["user", "api"]),
|
||||
|
|
@ -22,20 +22,20 @@ export interface ICreateCachedTurnUseCase {
|
|||
export class CreateCachedTurnUseCase implements ICreateCachedTurnUseCase {
|
||||
private readonly cacheService: ICacheService;
|
||||
private readonly conversationsRepository: IConversationsRepository;
|
||||
private readonly usageQuotaPolicyService: IUsageQuotaPolicyService;
|
||||
private readonly usageQuotaPolicy: IUsageQuotaPolicy;
|
||||
|
||||
constructor({
|
||||
cacheService,
|
||||
conversationsRepository,
|
||||
usageQuotaPolicyService,
|
||||
usageQuotaPolicy,
|
||||
}: {
|
||||
cacheService: ICacheService,
|
||||
conversationsRepository: IConversationsRepository,
|
||||
usageQuotaPolicyService: IUsageQuotaPolicyService,
|
||||
usageQuotaPolicy: IUsageQuotaPolicy,
|
||||
}) {
|
||||
this.cacheService = cacheService;
|
||||
this.conversationsRepository = conversationsRepository;
|
||||
this.usageQuotaPolicyService = usageQuotaPolicyService;
|
||||
this.usageQuotaPolicy = usageQuotaPolicy;
|
||||
}
|
||||
|
||||
async execute(data: z.infer<typeof inputSchema>): Promise<{ key: string }> {
|
||||
|
|
@ -49,7 +49,7 @@ export class CreateCachedTurnUseCase implements ICreateCachedTurnUseCase {
|
|||
const { projectId } = conversation;
|
||||
|
||||
// assert and consume quota
|
||||
await this.usageQuotaPolicyService.assertAndConsume(projectId);
|
||||
await this.usageQuotaPolicy.assertAndConsume(projectId);
|
||||
|
||||
// if caller is a user, ensure they are a member of project
|
||||
if (data.caller === "user") {
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import { IConversationsRepository } from "@/src/application/repositories/convers
|
|||
import { z } from "zod";
|
||||
import { Conversation } from "@/src/entities/models/conversation";
|
||||
import { Workflow } from "@/app/lib/types/workflow_types";
|
||||
import { IUsageQuotaPolicyService } from '../../services/usage-quota-policy.service.interface';
|
||||
import { IUsageQuotaPolicy } from '../../policies/usage-quota.policy.interface';
|
||||
|
||||
const inputSchema = z.object({
|
||||
caller: z.enum(["user", "api"]),
|
||||
|
|
@ -21,17 +21,17 @@ export interface ICreateConversationUseCase {
|
|||
|
||||
export class CreateConversationUseCase implements ICreateConversationUseCase {
|
||||
private readonly conversationsRepository: IConversationsRepository;
|
||||
private readonly usageQuotaPolicyService: IUsageQuotaPolicyService;
|
||||
private readonly usageQuotaPolicy: IUsageQuotaPolicy;
|
||||
|
||||
constructor({
|
||||
conversationsRepository,
|
||||
usageQuotaPolicyService,
|
||||
usageQuotaPolicy,
|
||||
}: {
|
||||
conversationsRepository: IConversationsRepository,
|
||||
usageQuotaPolicyService: IUsageQuotaPolicyService,
|
||||
usageQuotaPolicy: IUsageQuotaPolicy,
|
||||
}) {
|
||||
this.conversationsRepository = conversationsRepository;
|
||||
this.usageQuotaPolicyService = usageQuotaPolicyService;
|
||||
this.usageQuotaPolicy = usageQuotaPolicy;
|
||||
}
|
||||
|
||||
async execute(data: z.infer<typeof inputSchema>): Promise<z.infer<typeof Conversation>> {
|
||||
|
|
@ -40,7 +40,7 @@ export class CreateConversationUseCase implements ICreateConversationUseCase {
|
|||
let workflow = data.workflow;
|
||||
|
||||
// assert and consume quota
|
||||
await this.usageQuotaPolicyService.assertAndConsume(projectId);
|
||||
await this.usageQuotaPolicy.assertAndConsume(projectId);
|
||||
|
||||
// if caller is a user, ensure they are a member of project
|
||||
if (caller === "user") {
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import { IConversationsRepository } from "@/src/application/repositories/convers
|
|||
import { z } from "zod";
|
||||
import { ICacheService } from '@/src/application/services/cache.service.interface';
|
||||
import { CachedTurnRequest, Turn } from '@/src/entities/models/turn';
|
||||
import { IUsageQuotaPolicyService } from '../../services/usage-quota-policy.service.interface';
|
||||
import { IUsageQuotaPolicy } from '../../policies/usage-quota.policy.interface';
|
||||
|
||||
const inputSchema = z.object({
|
||||
caller: z.enum(["user", "api"]),
|
||||
|
|
@ -20,20 +20,20 @@ export interface IFetchCachedTurnUseCase {
|
|||
export class FetchCachedTurnUseCase implements IFetchCachedTurnUseCase {
|
||||
private readonly cacheService: ICacheService;
|
||||
private readonly conversationsRepository: IConversationsRepository;
|
||||
private readonly usageQuotaPolicyService: IUsageQuotaPolicyService;
|
||||
private readonly usageQuotaPolicy: IUsageQuotaPolicy;
|
||||
|
||||
constructor({
|
||||
cacheService,
|
||||
conversationsRepository,
|
||||
usageQuotaPolicyService,
|
||||
usageQuotaPolicy,
|
||||
}: {
|
||||
cacheService: ICacheService,
|
||||
conversationsRepository: IConversationsRepository,
|
||||
usageQuotaPolicyService: IUsageQuotaPolicyService,
|
||||
usageQuotaPolicy: IUsageQuotaPolicy,
|
||||
}) {
|
||||
this.cacheService = cacheService;
|
||||
this.conversationsRepository = conversationsRepository;
|
||||
this.usageQuotaPolicyService = usageQuotaPolicyService;
|
||||
this.usageQuotaPolicy = usageQuotaPolicy;
|
||||
}
|
||||
|
||||
async execute(data: z.infer<typeof inputSchema>): Promise<z.infer<typeof CachedTurnRequest>> {
|
||||
|
|
@ -56,7 +56,7 @@ export class FetchCachedTurnUseCase implements IFetchCachedTurnUseCase {
|
|||
const { projectId } = conversation;
|
||||
|
||||
// assert and consume quota
|
||||
await this.usageQuotaPolicyService.assertAndConsume(projectId);
|
||||
await this.usageQuotaPolicy.assertAndConsume(projectId);
|
||||
|
||||
// if caller is a user, ensure they are a member of project
|
||||
if (data.caller === "user") {
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import { IConversationsRepository } from "@/src/application/repositories/convers
|
|||
import { streamResponse } from "@/app/lib/agents";
|
||||
import { z } from "zod";
|
||||
import { Message } from "@/app/lib/types/types";
|
||||
import { IUsageQuotaPolicyService } from '../../services/usage-quota-policy.service.interface';
|
||||
import { IUsageQuotaPolicy } from '../../policies/usage-quota.policy.interface';
|
||||
|
||||
const inputSchema = z.object({
|
||||
caller: z.enum(["user", "api"]),
|
||||
|
|
@ -24,17 +24,17 @@ export interface IRunConversationTurnUseCase {
|
|||
|
||||
export class RunConversationTurnUseCase implements IRunConversationTurnUseCase {
|
||||
private readonly conversationsRepository: IConversationsRepository;
|
||||
private readonly usageQuotaPolicyService: IUsageQuotaPolicyService;
|
||||
private readonly usageQuotaPolicy: IUsageQuotaPolicy;
|
||||
|
||||
constructor({
|
||||
conversationsRepository,
|
||||
usageQuotaPolicyService,
|
||||
usageQuotaPolicy,
|
||||
}: {
|
||||
conversationsRepository: IConversationsRepository,
|
||||
usageQuotaPolicyService: IUsageQuotaPolicyService,
|
||||
usageQuotaPolicy: IUsageQuotaPolicy,
|
||||
}) {
|
||||
this.conversationsRepository = conversationsRepository;
|
||||
this.usageQuotaPolicyService = usageQuotaPolicyService;
|
||||
this.usageQuotaPolicy = usageQuotaPolicy;
|
||||
}
|
||||
|
||||
async *execute(data: z.infer<typeof inputSchema>): AsyncGenerator<z.infer<typeof TurnEvent>, void, unknown> {
|
||||
|
|
@ -48,7 +48,7 @@ export class RunConversationTurnUseCase implements IRunConversationTurnUseCase {
|
|||
const { id: conversationId, projectId } = conversation;
|
||||
|
||||
// assert and consume quota
|
||||
await this.usageQuotaPolicyService.assertAndConsume(projectId);
|
||||
await this.usageQuotaPolicy.assertAndConsume(projectId);
|
||||
|
||||
// if caller is a user, ensure they are a member of project
|
||||
if (data.caller === "user") {
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
import { IUsageQuotaPolicyService } from "@/src/application/services/usage-quota-policy.service.interface";
|
||||
import { IUsageQuotaPolicy } from "@/src/application/policies/usage-quota.policy.interface";
|
||||
import { redisClient } from "@/app/lib/redis";
|
||||
import { QuotaExceededError } from "@/src/entities/errors/common";
|
||||
|
||||
const MAX_QUERIES_PER_MINUTE = Number(process.env.MAX_QUERIES_PER_MINUTE) || 0;
|
||||
|
||||
export class RedisUsageQuotaPolicyService implements IUsageQuotaPolicyService {
|
||||
export class RedisUsageQuotaPolicy implements IUsageQuotaPolicy {
|
||||
async assertAndConsume(projectId: string): Promise<void> {
|
||||
if (MAX_QUERIES_PER_MINUTE === 0) {
|
||||
return;
|
||||
Loading…
Add table
Add a link
Reference in a new issue