feat: widgets and ai chat
This commit is contained in:
Rotimi Best
@@ -0,0 +1,24 @@
|
||||
import {
|
||||
listChatConversations,
|
||||
getChatConversation,
|
||||
createChatConversation,
|
||||
saveChatMessages,
|
||||
deleteChatConversation,
|
||||
updateConversationTitle,
|
||||
getChatHistory,
|
||||
saveChatHistory,
|
||||
deleteChatHistory
|
||||
} from '@cio/db/queries/agent';
|
||||
|
||||
export {
|
||||
listChatConversations,
|
||||
getChatConversation,
|
||||
createChatConversation,
|
||||
saveChatMessages,
|
||||
deleteChatConversation,
|
||||
updateConversationTitle,
|
||||
// Legacy compat
|
||||
getChatHistory as fetchChatHistory,
|
||||
saveChatHistory as persistChatHistory,
|
||||
deleteChatHistory as clearChatHistory
|
||||
};
|
||||
@@ -0,0 +1,27 @@
|
||||
/**
|
||||
* Trims conversation message history to stay within token budget.
|
||||
*
|
||||
* Strategy:
|
||||
* - Always keep the first message (often contains the initial request/document context)
|
||||
* - Always keep the last N messages (recent context is most relevant)
|
||||
* - Drop middle messages when the array is too long
|
||||
*
|
||||
* This is a message-count heuristic, not a token counter. Each message averages
|
||||
* ~200-500 tokens, so capping at ~40 messages keeps us under ~20K tokens of
|
||||
* history, leaving room for the system prompt + document text + LLM response.
|
||||
*/
|
||||
|
||||
const MAX_MESSAGES = 40;
|
||||
const KEEP_RECENT = 30;
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
export function trimMessageHistory(messages: any[]): any[] {
|
||||
if (messages.length <= MAX_MESSAGES) {
|
||||
return messages;
|
||||
}
|
||||
|
||||
const firstMessage = messages[0];
|
||||
const recentMessages = messages.slice(-KEEP_RECENT);
|
||||
|
||||
return [firstMessage, ...recentMessages];
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
import { nanoid } from 'nanoid';
|
||||
import { AppError } from '@api/utils/errors';
|
||||
import {
|
||||
MAX_DOCUMENT_TEXT_LENGTH,
|
||||
MAX_AGENT_DOCUMENT_SIZE,
|
||||
DOCUMENT_REDIS_TTL,
|
||||
SUPPORTED_DOCUMENT_TYPES
|
||||
} from '@cio/ai-assistant';
|
||||
import type { DocumentUploadResult } from '@cio/ai-assistant';
|
||||
import { agentDocumentKey } from '@api/utils/redis/key-generators';
|
||||
import { trackAgentEvent, AgentEvent } from '@api/utils/tinybird';
|
||||
import type { RedisClient } from '@api/utils/redis/redis';
|
||||
import { createChatDocument, getChatDocument } from '@cio/db/queries/agent';
|
||||
import { generateFileKey } from '@api/utils/upload';
|
||||
import { uploadToS3 } from '@api/utils/s3';
|
||||
import { getStorageConfig } from '@api/config/storage';
|
||||
import { createAssetFromUploadService } from '@api/services/assets/assets';
|
||||
|
||||
/**
|
||||
* Parse an uploaded document, store extracted text in Redis (hot cache) and Postgres
|
||||
* (durable, scoped to a conversation). Supports PDF, DOCX, and PPTX files.
|
||||
*/
|
||||
export async function parseAndStoreDocument(
|
||||
file: File,
|
||||
orgId: string,
|
||||
userId: string,
|
||||
courseId: string,
|
||||
conversationId: string,
|
||||
redis: RedisClient
|
||||
): Promise<DocumentUploadResult> {
|
||||
const mimeType = file.type;
|
||||
|
||||
if (!SUPPORTED_DOCUMENT_TYPES.includes(mimeType as (typeof SUPPORTED_DOCUMENT_TYPES)[number])) {
|
||||
throw new AppError('Unsupported file type. Allowed: PDF, DOCX, PPTX', 'UNSUPPORTED_FILE_TYPE', 415);
|
||||
}
|
||||
|
||||
if (file.size > MAX_AGENT_DOCUMENT_SIZE) {
|
||||
throw new AppError(
|
||||
`File too large. Maximum size is ${MAX_AGENT_DOCUMENT_SIZE / (1024 * 1024)}MB`,
|
||||
'FILE_TOO_LARGE',
|
||||
413
|
||||
);
|
||||
}
|
||||
|
||||
const buffer = Buffer.from(await file.arrayBuffer());
|
||||
let extractedText: string;
|
||||
let pageCount: number | null = null;
|
||||
|
||||
switch (mimeType) {
|
||||
case 'application/pdf':
|
||||
({ text: extractedText, pageCount } = await extractPdfText(buffer));
|
||||
break;
|
||||
case 'application/vnd.openxmlformats-officedocument.wordprocessingml.document':
|
||||
extractedText = await extractDocxText(buffer);
|
||||
break;
|
||||
case 'application/vnd.openxmlformats-officedocument.presentationml.presentation':
|
||||
({ text: extractedText, pageCount } = await extractPptxText(buffer));
|
||||
break;
|
||||
default:
|
||||
throw new AppError('Unsupported file type', 'UNSUPPORTED_FILE_TYPE', 415);
|
||||
}
|
||||
|
||||
const truncated = extractedText.length > MAX_DOCUMENT_TEXT_LENGTH;
|
||||
|
||||
if (truncated) {
|
||||
extractedText = extractedText.slice(0, MAX_DOCUMENT_TEXT_LENGTH);
|
||||
}
|
||||
|
||||
const wordCount = extractedText.split(/\s+/).filter(Boolean).length;
|
||||
const textPreview = extractedText.slice(0, 200);
|
||||
const documentId = nanoid();
|
||||
|
||||
// Persist the original file to S3 and register it as an asset so it shows
|
||||
// up in the org's asset manager — same pipeline lessons use.
|
||||
const storageKey = generateFileKey(file.name);
|
||||
const storageConfig = getStorageConfig();
|
||||
const uploadResult = await uploadToS3({
|
||||
Bucket: storageConfig.bucketDocuments,
|
||||
Key: storageKey,
|
||||
Body: buffer,
|
||||
ContentType: mimeType
|
||||
});
|
||||
|
||||
if (!uploadResult.success) {
|
||||
throw new AppError(
|
||||
`Failed to upload document to storage: ${uploadResult.error ?? 'unknown error'}`,
|
||||
'DOCUMENT_STORAGE_FAILED',
|
||||
500
|
||||
);
|
||||
}
|
||||
|
||||
const asset = await createAssetFromUploadService(orgId, userId, {
|
||||
kind: 'document',
|
||||
provider: 'upload',
|
||||
storageProvider: 's3',
|
||||
storageKey,
|
||||
byteSize: file.size,
|
||||
mimeType,
|
||||
title: file.name,
|
||||
isExternal: false,
|
||||
metadata: { source: 'ai_chat', conversationId }
|
||||
});
|
||||
|
||||
await redis.set(
|
||||
agentDocumentKey(documentId),
|
||||
JSON.stringify({
|
||||
text: extractedText,
|
||||
fileName: file.name,
|
||||
mimeType,
|
||||
uploadedAt: new Date().toISOString()
|
||||
}),
|
||||
{ EX: DOCUMENT_REDIS_TTL }
|
||||
);
|
||||
|
||||
await createChatDocument({
|
||||
id: documentId,
|
||||
conversationId,
|
||||
courseId,
|
||||
userId,
|
||||
assetId: asset.id,
|
||||
fileName: file.name,
|
||||
mimeType,
|
||||
text: extractedText,
|
||||
wordCount,
|
||||
pageCount
|
||||
});
|
||||
|
||||
trackAgentEvent(AgentEvent.DOCUMENT_UPLOADED, {
|
||||
orgId,
|
||||
userId,
|
||||
courseId,
|
||||
mimeType,
|
||||
fileSize: file.size,
|
||||
wordCount,
|
||||
truncated
|
||||
});
|
||||
|
||||
return {
|
||||
documentId,
|
||||
fileName: file.name,
|
||||
mimeType,
|
||||
pageCount,
|
||||
wordCount,
|
||||
textPreview,
|
||||
truncated
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieve stored document text. Tries the Redis hot cache first; on miss
|
||||
* falls back to Postgres and rehydrates Redis for next time.
|
||||
*/
|
||||
export async function getDocumentText(documentId: string, redis: RedisClient): Promise<string | null> {
|
||||
const raw = await redis.get(agentDocumentKey(documentId));
|
||||
|
||||
if (raw) {
|
||||
const parsed = JSON.parse(raw) as { text: string };
|
||||
|
||||
return parsed.text;
|
||||
}
|
||||
|
||||
const record = await getChatDocument(documentId);
|
||||
|
||||
if (!record) return null;
|
||||
|
||||
await redis.set(
|
||||
agentDocumentKey(documentId),
|
||||
JSON.stringify({
|
||||
text: record.text,
|
||||
fileName: record.fileName,
|
||||
mimeType: record.mimeType,
|
||||
uploadedAt: record.createdAt
|
||||
}),
|
||||
{ EX: DOCUMENT_REDIS_TTL }
|
||||
);
|
||||
|
||||
return record.text;
|
||||
}
|
||||
|
||||
// ─── Extraction Helpers ──────────────────────────────────────────────────────
|
||||
|
||||
async function extractPdfText(buffer: Buffer): Promise<{ text: string; pageCount: number }> {
|
||||
const pdfParse = (await import('pdf-parse')).default;
|
||||
const result = await pdfParse(buffer);
|
||||
|
||||
return { text: result.text, pageCount: result.numpages };
|
||||
}
|
||||
|
||||
async function extractDocxText(buffer: Buffer): Promise<string> {
|
||||
const mammoth = await import('mammoth');
|
||||
const result = await mammoth.extractRawText({ buffer });
|
||||
|
||||
return result.value;
|
||||
}
|
||||
|
||||
async function extractPptxText(buffer: Buffer): Promise<{ text: string; pageCount: number }> {
|
||||
// pptx-parser may not be available yet — use a simpler approach
|
||||
// For now, try to use officegen or a basic XML extraction
|
||||
try {
|
||||
const JSZip = (await import('jszip')).default;
|
||||
const zip = await JSZip.loadAsync(buffer);
|
||||
|
||||
const slideTexts: string[] = [];
|
||||
const slideFiles = Object.keys(zip.files)
|
||||
.filter((name) => name.match(/^ppt\/slides\/slide\d+\.xml$/))
|
||||
.sort();
|
||||
|
||||
for (const slideFile of slideFiles) {
|
||||
const content = await zip.files[slideFile].async('text');
|
||||
// Extract text from XML by stripping tags
|
||||
const textContent = content
|
||||
.replace(/<a:t[^>]*>/g, '')
|
||||
.replace(/<\/a:t>/g, '\n')
|
||||
.replace(/<[^>]+>/g, '')
|
||||
.replace(/\n{3,}/g, '\n\n')
|
||||
.trim();
|
||||
|
||||
if (textContent) {
|
||||
slideTexts.push(textContent);
|
||||
}
|
||||
}
|
||||
|
||||
return { text: slideTexts.join('\n\n---\n\n'), pageCount: slideFiles.length };
|
||||
} catch {
|
||||
throw new AppError('Failed to parse PPTX file', 'PPTX_PARSE_ERROR', 422);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
function escapeRegExp(value: string): string {
|
||||
return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
|
||||
}
|
||||
|
||||
function stripLeadingLessonTitle(content: string, lessonTitle: string): string {
|
||||
const normalizedTitle = lessonTitle.trim();
|
||||
|
||||
if (!normalizedTitle) {
|
||||
return content;
|
||||
}
|
||||
|
||||
const escapedTitle = escapeRegExp(normalizedTitle);
|
||||
const patterns = [
|
||||
new RegExp(`^\\s*<h[1-6][^>]*>\\s*${escapedTitle}\\s*</h[1-6]>\\s*`, 'i'),
|
||||
new RegExp(
|
||||
`^\\s*<p[^>]*>\\s*(?:<strong[^>]*>|<b[^>]*>)?\\s*${escapedTitle}\\s*(?:</strong>|</b>)?\\s*</p>\\s*`,
|
||||
'i'
|
||||
)
|
||||
];
|
||||
|
||||
let nextContent = content;
|
||||
|
||||
for (const pattern of patterns) {
|
||||
nextContent = nextContent.replace(pattern, '');
|
||||
}
|
||||
|
||||
return nextContent;
|
||||
}
|
||||
|
||||
function normalizeHeadingLevels(content: string): string {
|
||||
return content
|
||||
.replace(/<h1(\s|>)/gi, '<h3$1')
|
||||
.replace(/<\/h1>/gi, '</h3>')
|
||||
.replace(/<h2(\s|>)/gi, '<h3$1')
|
||||
.replace(/<\/h2>/gi, '</h3>');
|
||||
}
|
||||
|
||||
export function normalizeAgentLessonContent(content: string, lessonTitle: string): string {
|
||||
let normalizedContent = content.trim();
|
||||
|
||||
normalizedContent = stripLeadingLessonTitle(normalizedContent, lessonTitle);
|
||||
normalizedContent = normalizeHeadingLevels(normalizedContent);
|
||||
|
||||
return normalizedContent.trim();
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
import { AppError } from '@api/utils/errors';
|
||||
|
||||
export type CurrentQuestionForPatch = {
|
||||
id: number | string;
|
||||
title: string;
|
||||
questionTypeId: number;
|
||||
points: number;
|
||||
order: number;
|
||||
settings?: Record<string, unknown>;
|
||||
options: Array<{
|
||||
id: number | string;
|
||||
label: string | null;
|
||||
isCorrect: boolean;
|
||||
settings?: Record<string, unknown>;
|
||||
}>;
|
||||
};
|
||||
|
||||
export type QuestionPatch = {
|
||||
id: number;
|
||||
question?: string;
|
||||
questionTypeId?: number;
|
||||
points?: number;
|
||||
order?: number;
|
||||
settings?: Record<string, unknown>;
|
||||
options?: Array<{
|
||||
id?: number;
|
||||
label?: string;
|
||||
isCorrect?: boolean;
|
||||
settings?: Record<string, unknown>;
|
||||
}>;
|
||||
};
|
||||
|
||||
export type MergedOption = {
|
||||
id?: number;
|
||||
label: string;
|
||||
isCorrect: boolean;
|
||||
settings?: Record<string, unknown>;
|
||||
};
|
||||
|
||||
export type MergedQuestion = {
|
||||
id: number;
|
||||
question: string;
|
||||
questionTypeId: number;
|
||||
points: number;
|
||||
order: number;
|
||||
settings?: Record<string, unknown>;
|
||||
options: MergedOption[];
|
||||
};
|
||||
|
||||
export function buildUpdatedQuestions(
|
||||
currentQuestions: CurrentQuestionForPatch[],
|
||||
patches: QuestionPatch[],
|
||||
exerciseId: string
|
||||
): MergedQuestion[] {
|
||||
const byId = new Map(currentQuestions.map((q) => [Number(q.id), q] as const));
|
||||
|
||||
return patches.map((patch) => {
|
||||
const current = byId.get(patch.id);
|
||||
if (!current) {
|
||||
throw new AppError(
|
||||
`Question ${patch.id} does not belong to exercise ${exerciseId}`,
|
||||
'RESOURCE_NOT_IN_COURSE',
|
||||
404
|
||||
);
|
||||
}
|
||||
|
||||
const currentOptions: MergedOption[] = (current.options || []).map((o) => ({
|
||||
id: Number(o.id),
|
||||
label: o.label || '',
|
||||
isCorrect: o.isCorrect,
|
||||
settings: o.settings ?? undefined
|
||||
}));
|
||||
|
||||
const mergedSettings =
|
||||
patch.settings === undefined ? current.settings : { ...(current.settings ?? {}), ...patch.settings };
|
||||
|
||||
let mergedOptions: MergedOption[];
|
||||
if (patch.options === undefined) {
|
||||
mergedOptions = currentOptions;
|
||||
} else {
|
||||
mergedOptions = patch.options.map((opt) => {
|
||||
if (opt.id != null) {
|
||||
const existingOpt = currentOptions.find((co) => co.id === opt.id);
|
||||
return {
|
||||
id: opt.id,
|
||||
label: opt.label ?? existingOpt?.label ?? '',
|
||||
isCorrect: opt.isCorrect ?? existingOpt?.isCorrect ?? false,
|
||||
settings: opt.settings ?? existingOpt?.settings
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
label: opt.label ?? '',
|
||||
isCorrect: opt.isCorrect ?? false,
|
||||
settings: opt.settings
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
id: patch.id,
|
||||
question: patch.question ?? current.title ?? '',
|
||||
questionTypeId: patch.questionTypeId ?? current.questionTypeId,
|
||||
points: patch.points ?? current.points,
|
||||
order: patch.order ?? current.order,
|
||||
settings: mergedSettings,
|
||||
options: mergedOptions
|
||||
};
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
import { generateText } from 'ai';
|
||||
import { AIProvider, type AIProviderConfig } from '@cio/ai-assistant';
|
||||
import { createOpenAI } from '@ai-sdk/openai';
|
||||
import { createAnthropic } from '@ai-sdk/anthropic';
|
||||
import { createGoogleGenerativeAI } from '@ai-sdk/google';
|
||||
|
||||
/**
|
||||
* Cheap models used specifically for title generation to minimize token cost.
|
||||
*/
|
||||
const TITLE_MODELS: Record<AIProvider, string> = {
|
||||
[AIProvider.OPENAI]: 'gpt-4o-mini',
|
||||
[AIProvider.ANTHROPIC]: 'claude-haiku-4-20250414',
|
||||
[AIProvider.GOOGLE]: 'gemini-2.5-flash'
|
||||
};
|
||||
|
||||
function createTitleModel(config: AIProviderConfig) {
|
||||
const modelName = TITLE_MODELS[config.provider];
|
||||
|
||||
switch (config.provider) {
|
||||
case AIProvider.OPENAI: {
|
||||
const openai = createOpenAI({ apiKey: config.apiKey });
|
||||
return openai(modelName);
|
||||
}
|
||||
case AIProvider.ANTHROPIC: {
|
||||
const anthropic = createAnthropic({ apiKey: config.apiKey });
|
||||
return anthropic(modelName);
|
||||
}
|
||||
case AIProvider.GOOGLE: {
|
||||
const google = createGoogleGenerativeAI({ apiKey: config.apiKey });
|
||||
return google(modelName);
|
||||
}
|
||||
default:
|
||||
throw new Error(`Unsupported AI provider: ${config.provider}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a short conversational title from the user's first message.
|
||||
* Uses the cheapest available model to minimize token cost.
|
||||
*/
|
||||
export async function generateConversationTitle(
|
||||
firstMessageText: string,
|
||||
providerConfig: AIProviderConfig
|
||||
): Promise<string> {
|
||||
const model = createTitleModel(providerConfig);
|
||||
|
||||
const { text } = await generateText({
|
||||
model,
|
||||
system:
|
||||
"Generate a short title (max 6 words) for a chat conversation based on the user's first message. Return ONLY the title, no quotes, no punctuation at the end. Be concise and descriptive.",
|
||||
prompt: firstMessageText,
|
||||
maxOutputTokens: 30
|
||||
});
|
||||
|
||||
return text.trim().slice(0, 80);
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
import { db } from '@cio/db';
|
||||
import * as schema from '@cio/db/schema';
|
||||
import { eq, and, gte, sql } from 'drizzle-orm';
|
||||
import { getActiveOrganizationPlan } from '@cio/db/queries/organization';
|
||||
import { AppError, ErrorCodes } from '@api/utils/errors';
|
||||
import type { TokenBalance, TokenUsage } from '@cio/ai-assistant';
|
||||
|
||||
// ─── Plan-Based Token Allowances ─────────────────────────────────────────────
|
||||
|
||||
const PLAN_TOKEN_ALLOWANCES: Record<string, number> = {
|
||||
BASIC: 50_000,
|
||||
EARLY_ADOPTER: 500_000,
|
||||
ENTERPRISE: 2_000_000
|
||||
};
|
||||
|
||||
/**
|
||||
* Get the monthly token allowance for an org based on its active plan.
|
||||
* Also checks if the plan payload has a custom `aiTokenAllowance` override.
|
||||
*/
|
||||
async function getPlanAllowance(orgId: string): Promise<{ planName: string; allowance: number }> {
|
||||
const activePlan = await getActiveOrganizationPlan(orgId);
|
||||
|
||||
if (!activePlan || !activePlan.planName) {
|
||||
return { planName: 'BASIC', allowance: PLAN_TOKEN_ALLOWANCES.BASIC };
|
||||
}
|
||||
|
||||
const planName = activePlan.planName;
|
||||
const payload = activePlan.payload as Record<string, unknown> | null;
|
||||
const customAllowance = payload?.aiTokenAllowance as number | undefined;
|
||||
const allowance = customAllowance ?? PLAN_TOKEN_ALLOWANCES[planName] ?? 0;
|
||||
|
||||
return { planName, allowance };
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the total tokens used by an org in the current month.
|
||||
*/
|
||||
async function getMonthlyUsage(orgId: string): Promise<number> {
|
||||
const startOfMonth = new Date();
|
||||
startOfMonth.setDate(1);
|
||||
startOfMonth.setHours(0, 0, 0, 0);
|
||||
|
||||
const [result] = await db
|
||||
.select({
|
||||
totalTokens: sql<number>`COALESCE(SUM(${schema.aiTokenUsage.promptTokens} + ${schema.aiTokenUsage.completionTokens}), 0)`
|
||||
})
|
||||
.from(schema.aiTokenUsage)
|
||||
.where(and(eq(schema.aiTokenUsage.orgId, orgId), gte(schema.aiTokenUsage.createdAt, startOfMonth.toISOString())));
|
||||
|
||||
return Number(result?.totalTokens ?? 0);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the credit balance for an org.
|
||||
*/
|
||||
async function getCreditBalance(orgId: string): Promise<number> {
|
||||
const [result] = await db
|
||||
.select({ balance: schema.aiCreditBalance.balance })
|
||||
.from(schema.aiCreditBalance)
|
||||
.where(eq(schema.aiCreditBalance.orgId, orgId))
|
||||
.limit(1);
|
||||
|
||||
return result?.balance ?? 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the full token balance for an org: used, allowance, credits, remaining.
|
||||
*/
|
||||
export async function getTokenBalance(orgId: string): Promise<TokenBalance> {
|
||||
const [{ planName, allowance }, monthlyUsage, creditBalance] = await Promise.all([
|
||||
getPlanAllowance(orgId),
|
||||
getMonthlyUsage(orgId),
|
||||
getCreditBalance(orgId)
|
||||
]);
|
||||
|
||||
const remainingAllowance = Math.max(0, allowance - monthlyUsage);
|
||||
const remaining = remainingAllowance + creditBalance;
|
||||
|
||||
return {
|
||||
used: monthlyUsage,
|
||||
allowance,
|
||||
creditBalance,
|
||||
remaining
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the org has sufficient tokens. Throws 402 if exhausted.
|
||||
*/
|
||||
export async function enforceTokenBalance(orgId: string): Promise<TokenBalance> {
|
||||
const balance = await getTokenBalance(orgId);
|
||||
|
||||
if (balance.remaining <= 0) {
|
||||
throw new AppError('Token limit reached', 'TOKEN_LIMIT_REACHED', 402);
|
||||
}
|
||||
|
||||
return balance;
|
||||
}
|
||||
|
||||
/**
|
||||
* Record token usage after an LLM call.
|
||||
*/
|
||||
export async function recordTokenUsage(
|
||||
orgId: string,
|
||||
userId: string,
|
||||
courseId: string,
|
||||
usage: TokenUsage,
|
||||
model: string
|
||||
): Promise<void> {
|
||||
await db.insert(schema.aiTokenUsage).values({
|
||||
orgId,
|
||||
userId,
|
||||
courseId,
|
||||
promptTokens: usage.promptTokens,
|
||||
completionTokens: usage.completionTokens,
|
||||
model
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if an org's plan supports a specific feature (e.g., document upload).
|
||||
* Returns the plan name so callers can decide what to gate.
|
||||
*/
|
||||
export async function getOrgPlanName(orgId: string): Promise<string> {
|
||||
const { planName } = await getPlanAllowance(orgId);
|
||||
|
||||
return planName;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if an org is on a paid plan (EARLY_ADOPTER or above).
|
||||
*/
|
||||
export async function isOrgOnPaidPlan(orgId: string): Promise<boolean> {
|
||||
const planName = await getOrgPlanName(orgId);
|
||||
|
||||
return planName !== 'BASIC';
|
||||
}
|
||||
|
||||
// ─── Usage History ───────────────────────────────────────────────────────────
|
||||
|
||||
export interface DailyUsage {
|
||||
date: string;
|
||||
tokens: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get daily token usage for the current month, grouped by day.
|
||||
*/
|
||||
export async function getUsageHistory(orgId: string): Promise<DailyUsage[]> {
|
||||
const startOfMonth = new Date();
|
||||
startOfMonth.setDate(1);
|
||||
startOfMonth.setHours(0, 0, 0, 0);
|
||||
|
||||
const rows = await db
|
||||
.select({
|
||||
date: sql<string>`DATE(${schema.aiTokenUsage.createdAt})`,
|
||||
tokens: sql<number>`SUM(${schema.aiTokenUsage.promptTokens} + ${schema.aiTokenUsage.completionTokens})`
|
||||
})
|
||||
.from(schema.aiTokenUsage)
|
||||
.where(and(eq(schema.aiTokenUsage.orgId, orgId), gte(schema.aiTokenUsage.createdAt, startOfMonth.toISOString())))
|
||||
.groupBy(sql`DATE(${schema.aiTokenUsage.createdAt})`)
|
||||
.orderBy(sql`DATE(${schema.aiTokenUsage.createdAt})`);
|
||||
|
||||
return rows.map((row) => ({
|
||||
date: String(row.date),
|
||||
tokens: Number(row.tokens)
|
||||
}));
|
||||
}
|
||||
|
||||
/**
|
||||
* Get full detailed usage stats for the usage endpoint.
|
||||
*/
|
||||
export async function getDetailedUsage(orgId: string) {
|
||||
const [balance, history] = await Promise.all([getTokenBalance(orgId), getUsageHistory(orgId)]);
|
||||
|
||||
return { ...balance, history };
|
||||
}
|
||||
|
||||
// ─── Credit Purchases ────────────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Add credits to an org's balance.
|
||||
* Creates the balance row if it doesn't exist (upsert).
|
||||
*/
|
||||
export async function addCredits(orgId: string, amount: number): Promise<number> {
|
||||
if (amount <= 0) {
|
||||
throw new AppError('Credit amount must be positive', 'INVALID_CREDIT_AMOUNT', 400);
|
||||
}
|
||||
|
||||
const existing = await db
|
||||
.select({ id: schema.aiCreditBalance.id, balance: schema.aiCreditBalance.balance })
|
||||
.from(schema.aiCreditBalance)
|
||||
.where(eq(schema.aiCreditBalance.orgId, orgId))
|
||||
.limit(1);
|
||||
|
||||
if (existing.length > 0) {
|
||||
const newBalance = existing[0].balance + amount;
|
||||
await db
|
||||
.update(schema.aiCreditBalance)
|
||||
.set({ balance: newBalance, updatedAt: new Date().toISOString() })
|
||||
.where(eq(schema.aiCreditBalance.orgId, orgId));
|
||||
|
||||
return newBalance;
|
||||
}
|
||||
|
||||
await db.insert(schema.aiCreditBalance).values({ orgId, balance: amount });
|
||||
|
||||
return amount;
|
||||
}
|
||||
Reference in New Issue
Block a user