graphile-llm 0.7.3 → 0.9.0

package/esm/metering.js

@@ -0,0 +1,352 @@
+/**
+ * metering — Billing-aware wrappers for embedder and chat functions
+ *
+ * Wraps EmbedderFunction and ChatFunction with:
+ *   1. Pre-check: `check_billing_quota(meter_slug, entity_id, estimated_amount)`
+ *   2. Execute the underlying function
+ *   3. Post-record: `record_usage(meter_slug, entity_id, actual_amount)`
+ *
+ * When the quota check fails, the wrapper returns null (graceful degradation)
+ * instead of throwing, so the search pipeline can fall back to text-only.
+ *
+ * Token counts:
+ *   - Chat: real provider counts via ChatResult.usage (from OllamaAdapter.stream())
+ *   - Embedding: real provider counts via EmbeddingResult.promptTokens (from /api/embed)
+ *
+ * The billing functions live in the tenant database and are called via the
+ * Graphile `withPgClient` callback. Function locations (schema, names) are
+ * resolved from `billing_module` metaschema and cached by `config-cache.ts`.
+ */
+// ─── Billing SQL Helpers ────────────────────────────────────────────────────
+/**
+ * Check if the entity has sufficient quota for the requested amount.
+ * Returns true if the call is allowed, false if quota is exceeded.
+ *
+ * Gracefully returns true if the billing function doesn't exist or errors —
+ * metering is opt-in, so missing infrastructure means "allow".
+ */
+async function checkQuota(pgClient, billing, entityId, meterSlug, amount) {
+    try {
+        const sql = `SELECT "${billing.privateSchema}"."${billing.checkBillingQuotaFunction}"($1, $2::uuid, $3) AS allowed`;
+        const result = await pgClient.query(sql, [meterSlug, entityId, amount]);
+        return result.rows[0]?.allowed !== false;
+    }
+    catch (e) {
+        const message = e instanceof Error ? e.message : String(e);
+        console.warn(`[graphile-llm] check_billing_quota failed (allowing): ${message}`);
+        return true;
+    }
+}
+/**
+ * Record usage after a successful call.
+ * Gracefully skips if the billing function doesn't exist or errors.
+ */
+async function recordUsage(pgClient, billing, entityId, meterSlug, amount, metadata) {
+    try {
+        const sql = `SELECT "${billing.privateSchema}"."${billing.recordUsageFunction}"($1, $2::uuid, $3, $4::jsonb)`;
+        await pgClient.query(sql, [meterSlug, entityId, amount, JSON.stringify(metadata)]);
+    }
+    catch (e) {
+        const message = e instanceof Error ? e.message : String(e);
+        console.warn(`[graphile-llm] record_usage failed (non-fatal): ${message}`);
+    }
+}
+/**
+ * Write a row to the usage_log_inference table.
+ * Gracefully skips if the inference_log_module is not provisioned.
+ *
+ * TODO: Also write to child (generated) database when dual-write is needed.
+ */
+export async function logInferenceUsage(ctx, entry) {
+    if (!ctx.inferenceLog)
+        return;
+    const { schema, tableName } = ctx.inferenceLog;
+    const sql = `INSERT INTO "${schema}"."${tableName}" (
+    database_id, entity_id, actor_id,
+    model, provider, service, operation,
+    input_tokens, output_tokens, total_tokens,
+    cache_read_tokens, cache_write_tokens,
+    latency_ms, rag_enabled, chunks_retrieved,
+    embedding_model, embedding_latency_ms,
+    status, error_type, raw_usage
+  ) VALUES (
+    $1, $2, $3,
+    $4, $5, $6, $7,
+    $8, $9, $10,
+    $11, $12,
+    $13, $14, $15,
+    $16, $17,
+    $18, $19, $20
+  )`;
+    try {
+        await ctx.withPgClient(ctx.pgSettings, async (pgClient) => {
+            await pgClient.query(sql, [
+                entry.databaseId, entry.entityId, entry.actorId,
+                entry.model, entry.provider, entry.service, entry.operation,
+                entry.inputTokens, entry.outputTokens, entry.totalTokens,
+                entry.cacheReadTokens, entry.cacheWriteTokens,
+                entry.latencyMs, entry.ragEnabled, entry.chunksRetrieved,
+                entry.embeddingModel, entry.embeddingLatencyMs,
+                entry.status, entry.errorType,
+                entry.rawUsage ? JSON.stringify(entry.rawUsage) : null
+            ]);
+        });
+    }
+    catch (e) {
+        const message = e instanceof Error ? e.message : String(e);
+        console.warn(`[graphile-llm] inference log INSERT failed (non-fatal): ${message}`);
+    }
+}
+// ─── Metered Embedder ───────────────────────────────────────────────────────
+/**
+ * Wrap an embedder with billing quota check + usage recording.
+ *
+ * The returned MeterResult contains `quotaExceeded: true` when the pre-check
+ * fails, enabling the caller to fall back to text-only search.
+ */
+export async function meteredEmbed(embedder, text, ctx, options = {}) {
+    const startTime = Date.now();
+    // No billing context → just embed without metering
+    if (!ctx) {
+        const { embedding } = await embedder(text);
+        return {
+            result: embedding,
+            metered: false,
+            quotaExceeded: false,
+            latencyMs: Date.now() - startTime
+        };
+    }
+    const meterSlug = options.embeddingMeterSlug;
+    if (!meterSlug) {
+        const { embedding } = await embedder(text);
+        return {
+            result: embedding,
+            metered: false,
+            quotaExceeded: false,
+            latencyMs: Date.now() - startTime
+        };
+    }
+    if (options.skipMetering) {
+        const { embedding } = await embedder(text);
+        return {
+            result: embedding,
+            metered: false,
+            quotaExceeded: false,
+            latencyMs: Date.now() - startTime
+        };
+    }
+    // Pre-check: can this entity afford this call?
+    let allowed = true;
+    try {
+        await ctx.withPgClient(ctx.pgSettings, async (pgClient) => {
+            allowed = await checkQuota(pgClient, ctx.billing, ctx.entityId, meterSlug, 1);
+        });
+    }
+    catch {
+        allowed = true;
+    }
+    if (!allowed) {
+        logInferenceUsage(ctx, {
+            databaseId: ctx.databaseId,
+            entityId: ctx.entityId,
+            actorId: ctx.actorId,
+            model: options.embeddingModel ?? meterSlug,
+            provider: options.provider ?? null,
+            service: 'embedding',
+            operation: 'create',
+            inputTokens: 0,
+            outputTokens: 0,
+            totalTokens: 0,
+            cacheReadTokens: null,
+            cacheWriteTokens: null,
+            latencyMs: Date.now() - startTime,
+            ragEnabled: false,
+            chunksRetrieved: null,
+            embeddingModel: options.embeddingModel ?? null,
+            embeddingLatencyMs: null,
+            status: 'quota_exceeded',
+            errorType: null,
+            rawUsage: null
+        }).catch(() => { });
+        return {
+            result: null,
+            metered: true,
+            quotaExceeded: true,
+            latencyMs: Date.now() - startTime
+        };
+    }
+    // Execute embedding — real token count from provider via EmbeddingResult
+    const { embedding, promptTokens } = await embedder(text);
+    const latencyMs = Date.now() - startTime;
+    ctx.withPgClient(ctx.pgSettings, async (pgClient) => {
+        await recordUsage(pgClient, ctx.billing, ctx.entityId, meterSlug, promptTokens, {
+            request_id: ctx.requestId,
+            input_chars: text.length,
+            prompt_tokens: promptTokens,
+            dims: embedding.length,
+            latency_ms: latencyMs
+        });
+    }).catch(() => { });
+    // Log to inference usage table
+    logInferenceUsage(ctx, {
+        databaseId: ctx.databaseId,
+        entityId: ctx.entityId,
+        actorId: ctx.actorId,
+        model: options.embeddingModel ?? meterSlug,
+        provider: options.provider ?? null,
+        service: 'embedding',
+        operation: 'create',
+        inputTokens: promptTokens,
+        outputTokens: 0,
+        totalTokens: promptTokens,
+        cacheReadTokens: null,
+        cacheWriteTokens: null,
+        latencyMs,
+        ragEnabled: false,
+        chunksRetrieved: null,
+        embeddingModel: options.embeddingModel ?? null,
+        embeddingLatencyMs: latencyMs,
+        status: 'success',
+        errorType: null,
+        rawUsage: { prompt_tokens: promptTokens }
+    }).catch(() => { });
+    return {
+        result: embedding,
+        metered: true,
+        quotaExceeded: false,
+        latencyMs
+    };
+}
+// ─── Metered Chat ───────────────────────────────────────────────────────────
+/**
+ * Wrap a chat completion call with billing quota check + usage recording.
+ */
+export async function meteredChat(chat, messages, ctx, chatOptions, meteringOptions = {}) {
+    const startTime = Date.now();
+    if (!ctx) {
+        const chatResult = await chat(messages, chatOptions);
+        return {
+            result: chatResult.content,
+            metered: false,
+            quotaExceeded: false,
+            latencyMs: Date.now() - startTime
+        };
+    }
+    const meterSlug = meteringOptions.chatMeterSlug;
+    if (!meterSlug) {
+        const chatResult = await chat(messages, chatOptions);
+        return {
+            result: chatResult.content,
+            metered: false,
+            quotaExceeded: false,
+            latencyMs: Date.now() - startTime
+        };
+    }
+    if (meteringOptions.skipMetering) {
+        const chatResult = await chat(messages, chatOptions);
+        return {
+            result: chatResult.content,
+            metered: false,
+            quotaExceeded: false,
+            latencyMs: Date.now() - startTime
+        };
+    }
+    // Pre-check: can this entity afford this call?
+    let allowed = true;
+    try {
+        await ctx.withPgClient(ctx.pgSettings, async (pgClient) => {
+            allowed = await checkQuota(pgClient, ctx.billing, ctx.entityId, meterSlug, 1);
+        });
+    }
+    catch {
+        allowed = true;
+    }
+    if (!allowed) {
+        const estimatedInputTokens = Math.ceil(messages.reduce((sum, m) => sum + m.content.length, 0) / 4);
+        logInferenceUsage(ctx, {
+            databaseId: ctx.databaseId,
+            entityId: ctx.entityId,
+            actorId: ctx.actorId,
+            model: meteringOptions.chatModel ?? meterSlug,
+            provider: meteringOptions.provider ?? null,
+            service: 'llm',
+            operation: 'chat',
+            inputTokens: estimatedInputTokens,
+            outputTokens: 0,
+            totalTokens: estimatedInputTokens,
+            cacheReadTokens: null,
+            cacheWriteTokens: null,
+            latencyMs: Date.now() - startTime,
+            ragEnabled: false,
+            chunksRetrieved: null,
+            embeddingModel: null,
+            embeddingLatencyMs: null,
+            status: 'quota_exceeded',
+            errorType: null,
+            rawUsage: null
+        }).catch(() => { });
+        return {
+            result: null,
+            metered: true,
+            quotaExceeded: true,
+            latencyMs: Date.now() - startTime
+        };
+    }
+    // Execute chat completion — returns real token usage from provider
+    const chatResult = await chat(messages, chatOptions);
+    const latencyMs = Date.now() - startTime;
+    const usage = chatResult.usage;
+    ctx.withPgClient(ctx.pgSettings, async (pgClient) => {
+        await recordUsage(pgClient, ctx.billing, ctx.entityId, meterSlug, usage.totalTokens, {
+            request_id: ctx.requestId,
+            input_tokens: usage.input,
+            output_tokens: usage.output,
+            cache_read_tokens: usage.cacheRead,
+            cache_write_tokens: usage.cacheWrite,
+            messages_count: messages.length,
+            latency_ms: latencyMs
+        });
+    }).catch(() => { });
+    // Log to inference usage table with real provider token counts
+    logInferenceUsage(ctx, {
+        databaseId: ctx.databaseId,
+        entityId: ctx.entityId,
+        actorId: ctx.actorId,
+        model: meteringOptions.chatModel ?? meterSlug,
+        provider: meteringOptions.provider ?? null,
+        service: 'llm',
+        operation: 'chat',
+        inputTokens: usage.input,
+        outputTokens: usage.output,
+        totalTokens: usage.totalTokens,
+        cacheReadTokens: usage.cacheRead || null,
+        cacheWriteTokens: usage.cacheWrite || null,
+        latencyMs,
+        ragEnabled: false,
+        chunksRetrieved: null,
+        embeddingModel: null,
+        embeddingLatencyMs: null,
+        status: 'success',
+        errorType: null,
+        rawUsage: { reasoning: usage.reasoning }
+    }).catch(() => { });
+    return {
+        result: chatResult.content,
+        metered: true,
+        quotaExceeded: false,
+        latencyMs
+    };
+}
+// ─── Error Types ────────────────────────────────────────────────────────────
+export class QuotaExceededError extends Error {
+    code = 'QUOTA_EXCEEDED';
+    meterSlug;
+    entityId;
+    constructor(meterSlug, entityId) {
+        super(`LLM quota exceeded for meter '${meterSlug}' on entity '${entityId}'. ` +
+            'Upgrade your plan or wait for the next billing period.');
+        this.name = 'QuotaExceededError';
+        this.meterSlug = meterSlug;
+        this.entityId = entityId;
+    }
+}

package/esm/plugins/agent-discovery-plugin.d.ts

@@ -0,0 +1,29 @@
+/**
+ * Agent Discovery
+ *
+ * Discovers agent tables by querying the agent_chat_module config table
+ * at runtime. The module stores schema_id, table names, and table IDs
+ * when provisioned — no smart tags needed.
+ *
+ * Results are cached per-database with a TTL so the REST middleware
+ * doesn't hit the database on every request.
+ */
+import { Pool } from 'pg';
+export interface AgentTableInfo {
+    /** The PostgreSQL schema name (e.g. 'agent_public') */
+    schemaName: string;
+    /** The table name (e.g. 'agent_thread') */
+    tableName: string;
+}
+export interface AgentDiscovery {
+    thread: AgentTableInfo | null;
+    message: AgentTableInfo | null;
+    task: AgentTableInfo | null;
+}
+/** Clear all cached discovery results (for testing) */
+export declare function clearAgentDiscoveryCache(): void;
+/**
+ * Look up agent table info for a database, querying the module config table.
+ * Results are cached per-database with a 60s TTL.
+ */
+export declare function getAgentDiscovery(pool: Pool, dbname: string): Promise<AgentDiscovery | null>;

package/esm/plugins/agent-discovery-plugin.js

@@ -0,0 +1,65 @@
+/**
+ * Agent Discovery
+ *
+ * Discovers agent tables by querying the agent_chat_module config table
+ * at runtime. The module stores schema_id, table names, and table IDs
+ * when provisioned — no smart tags needed.
+ *
+ * Results are cached per-database with a TTL so the REST middleware
+ * doesn't hit the database on every request.
+ */
+import { ModuleConfigCache } from 'graphile-cache';
+// ─── Cache ──────────────────────────────────────────────────────────────────
+const agentDiscoveryCache = new ModuleConfigCache({
+    name: 'agent-discovery',
+    ttlMs: 60_000
+});
+/** Clear all cached discovery results (for testing) */
+export function clearAgentDiscoveryCache() {
+    agentDiscoveryCache.clear();
+}
+// ─── Discovery Query ────────────────────────────────────────────────────────
+const DISCOVERY_SQL = `
+  SELECT
+    s.schema_name,
+    acm.thread_table_name,
+    acm.message_table_name,
+    acm.task_table_name
+  FROM metaschema_modules_public.agent_chat_module acm
+  JOIN metaschema_public.schema s ON s.id = acm.schema_id
+  LIMIT 1
+`;
+/**
+ * Look up agent table info for a database, querying the module config table.
+ * Results are cached per-database with a 60s TTL.
+ */
+export async function getAgentDiscovery(pool, dbname) {
+    const cached = agentDiscoveryCache.get(dbname);
+    if (cached !== undefined) {
+        return cached;
+    }
+    let discovery = null;
+    try {
+        const { rows } = await pool.query(DISCOVERY_SQL);
+        if (rows.length > 0) {
+            const row = rows[0];
+            const schemaName = row.schema_name;
+            discovery = {
+                thread: row.thread_table_name
+                    ? { schemaName, tableName: row.thread_table_name }
+                    : null,
+                message: row.message_table_name
+                    ? { schemaName, tableName: row.message_table_name }
+                    : null,
+                task: row.task_table_name
+                    ? { schemaName, tableName: row.task_table_name }
+                    : null
+            };
+        }
+    }
+    catch {
+        // Module table doesn't exist in this database — not provisioned
+    }
+    agentDiscoveryCache.set(dbname, discovery);
+    return discovery;
+}

package/esm/plugins/llm-module-plugin.d.ts

@@ -2,7 +2,8 @@
  * LlmModulePlugin
  *
  * Detects and loads the `llm_module` configuration from `services_public.api_modules`.
- * Makes the resolved embedder available to other plugins via the build context.
+ * Makes the resolved embedder and chat completer available to other plugins
+ * via the build context.
  *
  * This plugin is the foundation that enables per-database LLM configuration.
  * When an API has an `llm_module` configured, the embedder is resolved and
@@ -14,9 +15,13 @@
  *   2. `defaultEmbedder` from preset options (dev/testing fallback)
  *   3. Environment variables (EMBEDDER_PROVIDER, EMBEDDER_MODEL, EMBEDDER_BASE_URL)
  *   4. null — LLM features are disabled
+ *
+ * This plugin is intentionally pure — no billing or metering logic.
+ * The optional LlmMeteringPlugin wraps the embedder with billing integration
+ * if loaded (it runs after this plugin and before the consumer plugins).
  */
 import type { GraphileConfig } from 'graphile-config';
-import type { EmbedderFunction, ChatFunction, GraphileLlmOptions } from '../types';
+import type { ChatFunction, EmbedderFunction, GraphileLlmOptions } from '../types';
 declare global {
     namespace GraphileBuild {
         interface Build {
@@ -24,6 +29,10 @@ declare global {
             llmEmbedder: EmbedderFunction | null;
             /** The resolved chat completion function, or null if not configured */
             llmChatCompleter: ChatFunction | null;
+            /** The embedding model name (used as billing meter slug) */
+            llmEmbeddingModel: string | null;
+            /** The chat model name (used as billing meter slug) */
+            llmChatModel: string | null;
         }
     }
     namespace GraphileConfig {

package/esm/plugins/llm-module-plugin.js

@@ -2,7 +2,8 @@
  * LlmModulePlugin
  *
  * Detects and loads the `llm_module` configuration from `services_public.api_modules`.
- * Makes the resolved embedder available to other plugins via the build context.
+ * Makes the resolved embedder and chat completer available to other plugins
+ * via the build context.
  *
  * This plugin is the foundation that enables per-database LLM configuration.
  * When an API has an `llm_module` configured, the embedder is resolved and
@@ -14,9 +15,14 @@
  *   2. `defaultEmbedder` from preset options (dev/testing fallback)
  *   3. Environment variables (EMBEDDER_PROVIDER, EMBEDDER_MODEL, EMBEDDER_BASE_URL)
  *   4. null — LLM features are disabled
+ *
+ * This plugin is intentionally pure — no billing or metering logic.
+ * The optional LlmMeteringPlugin wraps the embedder with billing integration
+ * if loaded (it runs after this plugin and before the consumer plugins).
  */
-import { buildEmbedder, buildEmbedderFromEnv } from '../embedder';
 import { buildChatCompleter, buildChatCompleterFromEnv } from '../chat';
+import { buildEmbedder, buildEmbedderFromEnv } from '../embedder';
+import { getLlmEnvOptions } from '../env';
 /**
  * Creates the LlmModulePlugin with the given options.
  */
@@ -24,7 +30,7 @@ export function createLlmModulePlugin(options = {}) {
     const { defaultEmbedder, defaultChatCompleter } = options;
     return {
         name: 'LlmModulePlugin',
-        version: '0.1.0',
+        version: '0.2.0',
         description: 'Resolves LLM embedder and chat completer configuration and makes them available to other plugins',
         schema: {
             hooks: {
@@ -74,9 +80,11 @@ export function createLlmModulePlugin(options = {}) {
                     return build.extend(build, {
                         llmEmbedder: embedder,
                         llmChatCompleter: chat,
-                    }, 'LlmModulePlugin adding llmEmbedder and llmChatCompleter to build');
-                },
-            },
-        },
+                        llmEmbeddingModel: defaultEmbedder?.model ?? getLlmEnvOptions().embedding.model,
+                        llmChatModel: defaultChatCompleter?.model ?? getLlmEnvOptions().chat.model
+                    }, 'LlmModulePlugin adding llmEmbedder, llmChatCompleter, and model names to build');
+                }
+            }
+        }
     };
 }

package/esm/plugins/metering-plugin.d.ts

@@ -0,0 +1,42 @@
+/**
+ * LlmMeteringPlugin
+ *
+ * Opt-in billing integration for graphile-llm. Completely separate from the
+ * pure LLM plugins (text-search, text-mutation, rag).
+ *
+ * **How it works:**
+ * 1. At schema build time, replaces `build.llmEmbedder` with a metered wrapper
+ *    that has the same `(text: string) => Promise<number[]>` signature
+ * 2. At request time, wraps every root query/mutation resolver to set up a
+ *    request-scoped MeteringContext via AsyncLocalStorage
+ * 3. When the embedder is called (by any plugin), the wrapper checks
+ *    AsyncLocalStorage for a metering context and if found, calls
+ *    check_billing_quota before and record_usage after
+ * 4. If quota is exceeded, the wrapper returns null — the calling plugin sees
+ *    null and handles it (search falls back to text-only, mutations throw)
+ *
+ * The pure plugins never import metering, config-cache, or billing types.
+ * They call the embedder and handle null results — that's it.
+ *
+ * **Entity ID resolution:**
+ * The billing `entity_id` is resolved via a configurable callback.
+ * Default: reads `jwt.claims.user_id` from pgSettings. Override via
+ * `metering.resolveEntityId` in GraphileLlmPreset options.
+ *
+ * **Graceful behavior:**
+ * - billing_module not provisioned → embedder passes through unmetered
+ * - entity_id not available → embedder passes through unmetered
+ * - check_billing_quota throws → call is allowed (billing is opt-in)
+ * - record_usage throws → call succeeds, recording silently skipped
+ * - quota exceeded → embedder returns null
+ */
+import type { GraphileConfig } from 'graphile-config';
+import type { MeteringConfig } from '../types';
+declare global {
+    namespace GraphileConfig {
+        interface Plugins {
+            LlmMeteringPlugin: true;
+        }
+    }
+}
+export declare function createLlmMeteringPlugin(meteringConfig?: MeteringConfig): GraphileConfig.Plugin;