@tuturuuu/ai 0.0.10

package/src/embeddings/metered.ts

@@ -0,0 +1,283 @@
+import { google } from '@ai-sdk/google';
+import { checkAiCredits } from '@tuturuuu/ai/credits/check-credits';
+import {
+  commitMeteredEmbeddingCredits,
+  releaseMeteredEmbeddingCredits,
+  reserveMeteredEmbeddingCredits,
+} from '@tuturuuu/ai/credits/reservations';
+import { embed } from 'ai';
+
+export const GEMINI_EMBEDDING_2_MODEL_ID = 'gemini-embedding-2';
+export const GEMINI_EMBEDDING_2_GATEWAY_MODEL_ID = 'google/gemini-embedding-2';
+export const GEMINI_EMBEDDING_2_DIMENSIONS = 3072;
+
+const DEFAULT_COUNT_TOKENS_ENDPOINT =
+  'https://generativelanguage.googleapis.com/v1beta/models';
+
+export type MeteredEmbeddingTaskType =
+  | 'CLASSIFICATION'
+  | 'CLUSTERING'
+  | 'RETRIEVAL_DOCUMENT'
+  | 'RETRIEVAL_QUERY'
+  | 'SEMANTIC_SIMILARITY';
+
+type MeteredEmbeddingSkipReason =
+  | 'credit_check_failed'
+  | 'credits_exhausted'
+  | 'embedding_generation_failed'
+  | 'empty_value'
+  | 'invalid_embedding_shape'
+  | 'missing_google_api_key'
+  | 'missing_scope'
+  | 'reservation_commit_failed'
+  | 'reservation_failed'
+  | 'token_count_failed';
+
+export type MeteredEmbeddingResult =
+  | {
+      creditsDeducted: number;
+      embedding: number[];
+      inputTokens: number;
+      modelId: string;
+      ok: true;
+      skipped?: false;
+    }
+  | {
+      errorCode?: string | null;
+      ok: false;
+      reason: MeteredEmbeddingSkipReason | string;
+      skipped: true;
+      structuralDisable?: boolean;
+    };
+
+export type MeteredTextEmbeddingParams = {
+  metadata?: Record<string, unknown>;
+  source: string;
+  taskType: MeteredEmbeddingTaskType;
+  userId?: string | null;
+  value: string;
+  wsId?: string | null;
+};
+
+function isStructuralCreditError(errorCode?: string | null) {
+  return [
+    'CREDIT_CHECK_FAILED',
+    'FEATURE_NOT_ALLOWED',
+    'MODEL_DISABLED',
+    'MODEL_NOT_ALLOWED',
+    'MODEL_PRICING_UNAVAILABLE',
+    'NO_ALLOCATION',
+    'NO_BALANCE',
+  ].includes(errorCode ?? '');
+}
+
+export function shouldDisableMemoryForMeteringReason(
+  result: MeteredEmbeddingResult
+) {
+  if (result.ok) return false;
+  return (
+    result.structuralDisable ||
+    ['missing_google_api_key', 'missing_scope'].includes(result.reason)
+  );
+}
+
+function googleApiKey() {
+  return process.env.GOOGLE_GENERATIVE_AI_API_KEY?.trim() ?? '';
+}
+
+async function countGeminiEmbeddingTokens({
+  apiKey,
+  modelId,
+  value,
+}: {
+  apiKey: string;
+  modelId: string;
+  value: string;
+}) {
+  const endpoint = (
+    process.env.GOOGLE_GENERATIVE_AI_COUNT_TOKENS_ENDPOINT ??
+    DEFAULT_COUNT_TOKENS_ENDPOINT
+  ).replace(/\/+$/u, '');
+  const url = new URL(`${endpoint}/${modelId}:countTokens`);
+  url.searchParams.set('key', apiKey);
+
+  const response = await fetch(url, {
+    body: JSON.stringify({
+      contents: [{ parts: [{ text: value }] }],
+    }),
+    headers: { 'Content-Type': 'application/json' },
+    method: 'POST',
+  });
+
+  if (!response.ok) {
+    throw new Error(`CountTokens failed with HTTP ${response.status}`);
+  }
+
+  const json = (await response.json()) as {
+    totalTokens?: number;
+    total_tokens?: number;
+  };
+  const totalTokens = Number(json.totalTokens ?? json.total_tokens ?? 0);
+  if (!Number.isFinite(totalTokens) || totalTokens <= 0) {
+    throw new Error('CountTokens returned no billable tokens');
+  }
+
+  return Math.ceil(totalTokens);
+}
+
+function skipped(
+  reason: MeteredEmbeddingSkipReason | string,
+  options?: { errorCode?: string | null; structuralDisable?: boolean }
+): MeteredEmbeddingResult {
+  return {
+    ok: false,
+    reason,
+    skipped: true,
+    ...(options?.errorCode ? { errorCode: options.errorCode } : {}),
+    ...(options?.structuralDisable
+      ? { structuralDisable: options.structuralDisable }
+      : {}),
+  };
+}
+
+export async function createMeteredTextEmbedding({
+  metadata,
+  source,
+  taskType,
+  userId,
+  value,
+  wsId,
+}: MeteredTextEmbeddingParams): Promise<MeteredEmbeddingResult> {
+  const text = value.trim();
+  if (!text) return skipped('empty_value');
+  if (!wsId || !userId) {
+    return skipped('missing_scope', { structuralDisable: true });
+  }
+
+  const apiKey = googleApiKey();
+  if (!apiKey) {
+    return skipped('missing_google_api_key', { structuralDisable: true });
+  }
+
+  const creditCheck = await checkAiCredits(
+    wsId,
+    GEMINI_EMBEDDING_2_GATEWAY_MODEL_ID,
+    'embeddings',
+    { userId }
+  );
+  if (!creditCheck.allowed) {
+    return skipped(
+      creditCheck.errorCode === 'CREDITS_EXHAUSTED'
+        ? 'credits_exhausted'
+        : 'credit_check_failed',
+      {
+        errorCode: creditCheck.errorCode,
+        structuralDisable: isStructuralCreditError(creditCheck.errorCode),
+      }
+    );
+  }
+
+  let inputTokens: number;
+  try {
+    inputTokens = await countGeminiEmbeddingTokens({
+      apiKey,
+      modelId: GEMINI_EMBEDDING_2_MODEL_ID,
+      value: text,
+    });
+  } catch {
+    return skipped('token_count_failed', { structuralDisable: false });
+  }
+
+  const reservation = await reserveMeteredEmbeddingCredits({
+    inputTokens,
+    metadata: {
+      ...(metadata ?? {}),
+      dimensions: GEMINI_EMBEDDING_2_DIMENSIONS,
+      inputTokens,
+      modelId: GEMINI_EMBEDDING_2_GATEWAY_MODEL_ID,
+      source,
+      taskType,
+      userId,
+      wsId,
+    },
+    modelId: GEMINI_EMBEDDING_2_GATEWAY_MODEL_ID,
+    userId,
+    wsId,
+  });
+
+  if (!reservation.success || !reservation.reservationId) {
+    return skipped(
+      reservation.errorCode === 'INSUFFICIENT_CREDITS'
+        ? 'credits_exhausted'
+        : 'reservation_failed',
+      {
+        errorCode: reservation.errorCode,
+        structuralDisable: isStructuralCreditError(reservation.errorCode),
+      }
+    );
+  }
+
+  try {
+    const result = await embed({
+      model: google.embedding(GEMINI_EMBEDDING_2_MODEL_ID),
+      providerOptions: {
+        google: {
+          outputDimensionality: GEMINI_EMBEDDING_2_DIMENSIONS,
+          taskType,
+        },
+      },
+      value: text,
+    });
+
+    if (
+      !Array.isArray(result.embedding) ||
+      result.embedding.length !== GEMINI_EMBEDDING_2_DIMENSIONS
+    ) {
+      await releaseMeteredEmbeddingCredits(reservation.reservationId, {
+        reason: 'invalid_embedding_shape',
+        source,
+        userId,
+        wsId,
+      });
+      return skipped('invalid_embedding_shape', { structuralDisable: false });
+    }
+
+    const commit = await commitMeteredEmbeddingCredits(
+      reservation.reservationId,
+      {
+        actualInputTokens: result.usage.tokens ?? inputTokens,
+        source,
+        userId,
+        wsId,
+      }
+    );
+    if (!commit.success) {
+      await releaseMeteredEmbeddingCredits(reservation.reservationId, {
+        reason: 'reservation_commit_failed',
+        source,
+        userId,
+        wsId,
+      });
+      return skipped('reservation_commit_failed', {
+        errorCode: commit.errorCode,
+        structuralDisable: false,
+      });
+    }
+
+    return {
+      creditsDeducted: commit.creditsDeducted,
+      embedding: result.embedding,
+      inputTokens,
+      modelId: GEMINI_EMBEDDING_2_GATEWAY_MODEL_ID,
+      ok: true,
+    };
+  } catch {
+    await releaseMeteredEmbeddingCredits(reservation.reservationId, {
+      reason: 'embedding_generation_failed',
+      source,
+      userId,
+      wsId,
+    });
+    return skipped('embedding_generation_failed');
+  }
+}

package/src/executions/route.ts

@@ -0,0 +1,137 @@
+import { createAdminClient } from '@tuturuuu/supabase/next/server';
+import type { WorkspaceAIExecution } from '@tuturuuu/types';
+import { ROOT_WORKSPACE_ID } from '@tuturuuu/utils/constants';
+import { type NextRequest, NextResponse } from 'next/server';
+import { validateApiKeyHash } from '../api-key-hash';
+
+export function createPOST({
+  getLast30DaysStats,
+  getAllTimeStats,
+}: {
+  getLast30DaysStats: (wsId: string) => Promise<{
+    summary: any;
+    dailyStats: any[];
+    modelStats: any[];
+  }>;
+  getAllTimeStats: (wsId: string) => Promise<{
+    summary: any;
+    dailyStats: any[];
+    modelStats: any[];
+  }>;
+}) {
+  // Higher-order function that returns the actual request handler
+  return async function handler(req: NextRequest): Promise<Response> {
+    const sbAdmin = await createAdminClient();
+
+    const {
+      accessKey,
+      query,
+      configs = {
+        wsId: ROOT_WORKSPACE_ID,
+      },
+    } = (await req.json()) as {
+      accessKey?: {
+        id: string;
+        value: string;
+      };
+      query?: {
+        q?: string;
+        page?: string;
+        pageSize?: string;
+      };
+      configs?: {
+        wsId: string;
+      };
+    };
+
+    try {
+      if (!accessKey?.id || !accessKey?.value) {
+        console.error('Missing accessId or accessKey');
+        return new Response('Missing accessId or accessKey', { status: 400 });
+      }
+
+      const { data: apiKeyData, error: apiKeyError } = await sbAdmin
+        .from('workspace_api_keys')
+        .select('id, scopes, key_hash')
+        .eq('ws_id', configs.wsId)
+        .eq('id', accessKey.id)
+        .single();
+
+      const isValidAccessKey =
+        !apiKeyError &&
+        !!apiKeyData?.key_hash &&
+        (await validateApiKeyHash(accessKey.value, apiKeyData.key_hash));
+
+      if (!isValidAccessKey) {
+        console.error('Invalid accessId or accessKey');
+        return new Response('Invalid accessId or accessKey', { status: 400 });
+      }
+
+      const [executionData, analyticsData, allTimeStats] = await Promise.all([
+        getData(configs.wsId, query || {}),
+        getLast30DaysStats(configs.wsId),
+        getAllTimeStats(configs.wsId),
+      ]);
+
+      return NextResponse.json({
+        executionData,
+        analyticsData,
+        allTimeStats,
+      });
+    } catch (error) {
+      if (error instanceof Error) {
+        console.log(error.message);
+        return NextResponse.json(
+          {
+            message: `## Edge API Failure\nCould not complete the request. Please view the **Stack trace** below.\n\`\`\`bash\n${error instanceof Error ? error.stack : 'Unknown error'}`,
+          },
+          {
+            status: 500,
+          }
+        );
+      }
+      console.log(error);
+      return new Response('Internal Server Error', { status: 500 });
+    }
+  };
+}
+
+async function getData(
+  wsId: string,
+  {
+    q,
+    page = '1',
+    pageSize = '10',
+    retry = true,
+  }: {
+    q?: string | undefined;
+    page?: string | undefined;
+    pageSize?: string | undefined;
+    retry?: boolean | undefined;
+  }
+) {
+  const sbAdmin = await createAdminClient();
+  const queryBuilder = sbAdmin
+    .from('workspace_ai_executions')
+    .select('*', { count: 'exact' })
+    .eq('ws_id', wsId)
+    .order('created_at', { ascending: false });
+
+  if (page && pageSize) {
+    const parsedPage = parseInt(page, 10);
+    const parsedSize = parseInt(pageSize, 10);
+    const start = (parsedPage - 1) * parsedSize;
+    const end = parsedPage * parsedSize;
+    queryBuilder.range(start, end).limit(parsedSize);
+  }
+
+  const { data, error, count } = await queryBuilder;
+  if (error) {
+    if (!retry) throw error;
+    return getData(wsId, { q, pageSize, retry: false });
+  }
+  return { data, count } as {
+    data: WorkspaceAIExecution[];
+    count: number;
+  };
+}