@nextsparkjs/plugin-langchain 0.1.0-beta.1

package/lib/streaming.ts

@@ -0,0 +1,232 @@
+/**
+ * Streaming Service
+ *
+ * Provides token-by-token streaming via AsyncGenerator.
+ * Integrates with LangChain's streamEvents() method.
+ */
+
+import { BaseMessage, HumanMessage, AIMessage } from '@langchain/core/messages'
+import type { AgentContext, SessionConfig } from '../types/langchain.types'
+import { dbMemoryStore } from './db-memory-store'
+import { tokenTracker } from './token-tracker'
+import { tracer } from './tracer'
+import { createTracingCallbacks } from './tracer-callbacks'
+
+// Stream chunk types
+export type StreamChunk =
+    | { type: 'token'; content: string }
+    | { type: 'done'; fullContent: string; agentUsed?: string; tokenUsage?: TokenUsage }
+    | { type: 'error'; error: string }
+    | { type: 'tool_start'; toolName: string }
+    | { type: 'tool_end'; toolName: string; result: unknown }
+
+interface TokenUsage {
+    inputTokens: number
+    outputTokens: number
+    totalTokens: number
+}
+
+export interface StreamChatOptions {
+    sessionId?: string
+    sessionConfig?: SessionConfig
+    agentName?: string
+    onToken?: (token: string) => void
+    onToolCall?: (name: string, input: unknown) => void
+    signal?: AbortSignal  // For cancellation
+}
+
+interface AgentConfig {
+    modelConfig?: {
+        provider?: string
+        model?: string
+    }
+}
+
+/**
+ * Stream chat with an agent
+ *
+ * Uses LangChain's streamEvents() for token-by-token streaming.
+ * Handles memory persistence and token tracking.
+ */
+export async function* streamChat(
+    agent: { streamEvents: Function; invoke: Function },
+    input: string,
+    context: AgentContext,
+    config: AgentConfig,
+    options: StreamChatOptions = {}
+): AsyncGenerator<StreamChunk, void, unknown> {
+    const { sessionId, sessionConfig, agentName, signal } = options
+
+    let fullContent = ''
+    let tokenUsage: TokenUsage = { inputTokens: 0, outputTokens: 0, totalTokens: 0 }
+
+    // Start trace if context exists
+    const traceContext = await tracer.startTrace(
+        { userId: context.userId, teamId: context.teamId },
+        agentName || 'StreamAgent',
+        input,
+        { sessionId }
+    )
+
+    // Create tracing callbacks if trace was started
+    const tracingHandler = traceContext
+        ? createTracingCallbacks({ userId: context.userId, teamId: context.teamId }, traceContext.traceId)
+        : null
+    const tracingCallbacks = tracingHandler ? [tracingHandler] : []
+
+    try {
+        // Get history if session exists
+        let history: BaseMessage[] = []
+        if (sessionId) {
+            history = await dbMemoryStore.getMessages(sessionId, context)
+        }
+
+        // Create input with history
+        const messages = [...history, new HumanMessage(input)]
+
+        // Stream events from LangChain with tracing callbacks
+        const stream = agent.streamEvents(
+            { messages },
+            {
+                version: 'v2',
+                callbacks: tracingCallbacks,
+            }
+        )
+
+        for await (const event of stream) {
+            // Check for cancellation
+            if (signal?.aborted) {
+                yield { type: 'error', error: 'Stream cancelled by user' }
+                return
+            }
+
+            // Handle different event types
+            if (event.event === 'on_chat_model_stream') {
+                const token = event.data?.chunk?.content
+                if (token && typeof token === 'string') {
+                    fullContent += token
+                    yield { type: 'token', content: token }
+                }
+            }
+
+            if (event.event === 'on_tool_start') {
+                yield {
+                    type: 'tool_start',
+                    toolName: event.name || 'unknown',
+                }
+            }
+
+            if (event.event === 'on_tool_end') {
+                yield {
+                    type: 'tool_end',
+                    toolName: event.name || 'unknown',
+                    result: event.data?.output,
+                }
+            }
+
+            // Capture token usage from LLM events
+            if (event.event === 'on_llm_end') {
+                const usage = event.data?.output?.usage_metadata ||
+                              event.data?.output?.llmOutput?.tokenUsage
+                if (usage) {
+                    tokenUsage = {
+                        inputTokens: usage.input_tokens || usage.promptTokens || 0,
+                        outputTokens: usage.output_tokens || usage.completionTokens || 0,
+                        totalTokens: usage.total_tokens || usage.totalTokens || 0,
+                    }
+                }
+            }
+        }
+
+        // Persist to memory if session exists
+        if (sessionId && fullContent) {
+            await dbMemoryStore.addMessages(
+                sessionId,
+                [new HumanMessage(input), new AIMessage(fullContent)],
+                context,
+                sessionConfig
+            )
+        }
+
+        // Track token usage
+        if (tokenUsage.totalTokens > 0) {
+            await tokenTracker.trackUsage({
+                context,
+                sessionId,
+                provider: config.modelConfig?.provider || 'unknown',
+                model: config.modelConfig?.model || 'unknown',
+                usage: tokenUsage,
+                agentName,
+            })
+        }
+
+        // End trace successfully if started
+        if (traceContext) {
+            // Flush pending operations and get call counts from tracing handler
+            await tracingHandler?.flush()
+            const counts = tracingHandler?.getCounts() || { llmCalls: 0, toolCalls: 0 }
+
+            await tracer.endTrace(
+                { userId: context.userId, teamId: context.teamId },
+                traceContext.traceId,
+                {
+                    output: fullContent,
+                    tokens: tokenUsage.totalTokens > 0
+                        ? {
+                              input: tokenUsage.inputTokens,
+                              output: tokenUsage.outputTokens,
+                              total: tokenUsage.totalTokens,
+                          }
+                        : undefined,
+                    llmCalls: counts.llmCalls,
+                    toolCalls: counts.toolCalls,
+                }
+            )
+        }
+
+        yield {
+            type: 'done',
+            fullContent,
+            agentUsed: agentName,
+            tokenUsage,
+        }
+
+    } catch (error) {
+        // End trace with error if started
+        if (traceContext) {
+            // Flush pending operations and get call counts from tracing handler (even on error)
+            await tracingHandler?.flush()
+            const counts = tracingHandler?.getCounts() || { llmCalls: 0, toolCalls: 0 }
+
+            await tracer.endTrace(
+                { userId: context.userId, teamId: context.teamId },
+                traceContext.traceId,
+                {
+                    error: error instanceof Error ? error : new Error(String(error)),
+                    llmCalls: counts.llmCalls,
+                    toolCalls: counts.toolCalls,
+                }
+            )
+        }
+
+        const message = error instanceof Error ? error.message : 'Unknown streaming error'
+        yield { type: 'error', error: message }
+    }
+}
+
+/**
+ * Create SSE encoder for streaming responses
+ */
+export function createSSEEncoder() {
+    const encoder = new TextEncoder()
+
+    return {
+        encode(chunk: StreamChunk): Uint8Array {
+            return encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`)
+        },
+
+        encodeDone(): Uint8Array {
+            return encoder.encode('data: [DONE]\n\n')
+        },
+    }
+}

package/lib/token-tracker.ts

@@ -0,0 +1,298 @@
+/**
+ * Token Tracker Service
+ *
+ * Tracks token usage and calculates costs per request.
+ * Integrates with LangChain callbacks for automatic tracking.
+ */
+
+import { mutateWithRLS, queryWithRLS } from '@nextsparkjs/core/lib/db'
+import type { AgentContext } from '../types/langchain.types'
+
+interface TokenUsage {
+    inputTokens: number
+    outputTokens: number
+    totalTokens: number
+}
+
+interface TrackUsageParams {
+    context: AgentContext
+    sessionId?: string
+    provider: string
+    model: string
+    usage: TokenUsage
+    agentName?: string
+    metadata?: Record<string, unknown>
+}
+
+interface UsageStats {
+    totalTokens: number
+    totalCost: number
+    inputTokens: number
+    outputTokens: number
+    requestCount: number
+    byModel: Record<string, { tokens: number; cost: number }>
+}
+
+type Period = 'today' | '7d' | '30d' | 'all'
+
+// Pricing per 1M tokens (USD)
+const DEFAULT_PRICING: Record<string, { input: number; output: number }> = {
+    'gpt-4o': { input: 5.00, output: 15.00 },
+    'gpt-4o-mini': { input: 0.15, output: 0.60 },
+    'gpt-4-turbo': { input: 10.00, output: 30.00 },
+    'gpt-3.5-turbo': { input: 0.50, output: 1.50 },
+    'claude-3-5-sonnet': { input: 3.00, output: 15.00 },
+    'claude-3-opus': { input: 15.00, output: 75.00 },
+    'claude-3-haiku': { input: 0.25, output: 1.25 },
+    // Ollama models are free (local)
+    'ollama/*': { input: 0, output: 0 },
+}
+
+export const tokenTracker = {
+    /**
+     * Calculate cost for token usage
+     */
+    calculateCost(model: string, usage: TokenUsage, customPricing?: typeof DEFAULT_PRICING): {
+        inputCost: number
+        outputCost: number
+        totalCost: number
+    } {
+        const pricing = customPricing || DEFAULT_PRICING
+
+        // Check for exact match first, then wildcard
+        let modelPricing = pricing[model]
+        if (!modelPricing) {
+            // Check for provider wildcard (e.g., 'ollama/*')
+            const provider = model.split('/')[0] || model.split('-')[0]
+            modelPricing = pricing[`${provider}/*`] || { input: 0, output: 0 }
+        }
+
+        const inputCost = (usage.inputTokens / 1_000_000) * modelPricing.input
+        const outputCost = (usage.outputTokens / 1_000_000) * modelPricing.output
+
+        return {
+            inputCost,
+            outputCost,
+            totalCost: inputCost + outputCost,
+        }
+    },
+
+    /**
+     * Track token usage for a request
+     */
+    async trackUsage(params: TrackUsageParams): Promise<void> {
+        const { context, sessionId, provider, model, usage, agentName, metadata } = params
+        const { userId, teamId } = context
+
+        const costs = this.calculateCost(model, usage)
+
+        await mutateWithRLS(
+            `INSERT INTO public."langchain_token_usage"
+             (id, "userId", "teamId", "sessionId", provider, model,
+              "inputTokens", "outputTokens", "totalTokens",
+              "inputCost", "outputCost", "totalCost",
+              "agentName", metadata)
+             VALUES (gen_random_uuid()::text, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)`,
+            [
+                userId, teamId, sessionId || null, provider, model,
+                usage.inputTokens, usage.outputTokens, usage.totalTokens,
+                costs.inputCost, costs.outputCost, costs.totalCost,
+                agentName || null, JSON.stringify(metadata || {})
+            ],
+            userId
+        )
+    },
+
+    /**
+     * Get usage statistics for a user
+     */
+    async getUsage(context: AgentContext, period: Period = '30d'): Promise<UsageStats> {
+        const { userId, teamId } = context
+
+        const periodClause = this.getPeriodClause(period)
+
+        const result = await queryWithRLS<{
+            totalTokens: string
+            totalCost: string
+            inputTokens: string
+            outputTokens: string
+            requestCount: string
+            model: string
+            modelTokens: string
+            modelCost: string
+        }>(
+            `SELECT
+                SUM("totalTokens")::text as "totalTokens",
+                SUM("totalCost")::text as "totalCost",
+                SUM("inputTokens")::text as "inputTokens",
+                SUM("outputTokens")::text as "outputTokens",
+                COUNT(*)::text as "requestCount",
+                model,
+                SUM("totalTokens")::text as "modelTokens",
+                SUM("totalCost")::text as "modelCost"
+             FROM public."langchain_token_usage"
+             WHERE "userId" = $1 AND "teamId" = $2 ${periodClause}
+             GROUP BY model`,
+            [userId, teamId],
+            userId
+        )
+
+        if (!result.length) {
+            return {
+                totalTokens: 0,
+                totalCost: 0,
+                inputTokens: 0,
+                outputTokens: 0,
+                requestCount: 0,
+                byModel: {},
+            }
+        }
+
+        // Aggregate results
+        const byModel: Record<string, { tokens: number; cost: number }> = {}
+        let totalTokens = 0
+        let totalCost = 0
+        let inputTokens = 0
+        let outputTokens = 0
+        let requestCount = 0
+
+        for (const row of result) {
+            byModel[row.model] = {
+                tokens: parseInt(row.modelTokens, 10),
+                cost: parseFloat(row.modelCost),
+            }
+            totalTokens += parseInt(row.totalTokens || '0', 10)
+            totalCost += parseFloat(row.totalCost || '0')
+            inputTokens += parseInt(row.inputTokens || '0', 10)
+            outputTokens += parseInt(row.outputTokens || '0', 10)
+            requestCount += parseInt(row.requestCount || '0', 10)
+        }
+
+        return {
+            totalTokens,
+            totalCost,
+            inputTokens,
+            outputTokens,
+            requestCount,
+            byModel,
+        }
+    },
+
+    /**
+     * Get daily usage for charts
+     */
+    async getDailyUsage(context: AgentContext, days: number = 30): Promise<Array<{
+        date: string
+        tokens: number
+        cost: number
+        requests: number
+    }>> {
+        const { userId, teamId } = context
+
+        // Validate days parameter to prevent SQL injection (must be positive integer)
+        const safeDays = Math.max(1, Math.min(365, Math.floor(Number(days) || 30)))
+
+        const result = await queryWithRLS<{
+            date: string
+            tokens: string
+            cost: string
+            requests: string
+        }>(
+            `SELECT
+                DATE("createdAt")::text as date,
+                SUM("totalTokens")::text as tokens,
+                SUM("totalCost")::text as cost,
+                COUNT(*)::text as requests
+             FROM public."langchain_token_usage"
+             WHERE "userId" = $1 AND "teamId" = $2
+               AND "createdAt" >= now() - ($3 || ' days')::interval
+             GROUP BY DATE("createdAt")
+             ORDER BY date DESC`,
+            [userId, teamId, safeDays.toString()],
+            userId
+        )
+
+        return result.map(row => ({
+            date: row.date,
+            tokens: parseInt(row.tokens, 10),
+            cost: parseFloat(row.cost),
+            requests: parseInt(row.requests, 10),
+        }))
+    },
+
+    /**
+     * Get team usage (admin only)
+     */
+    async getTeamUsage(teamId: string, period: Period = '30d'): Promise<UsageStats & {
+        byUser: Record<string, { tokens: number; cost: number }>
+    }> {
+        const periodClause = this.getPeriodClause(period)
+
+        // This bypasses RLS - caller must verify admin permissions
+        const result = await queryWithRLS<{
+            userId: string
+            totalTokens: string
+            totalCost: string
+            inputTokens: string
+            outputTokens: string
+            requestCount: string
+        }>(
+            `SELECT
+                "userId",
+                SUM("totalTokens")::text as "totalTokens",
+                SUM("totalCost")::text as "totalCost",
+                SUM("inputTokens")::text as "inputTokens",
+                SUM("outputTokens")::text as "outputTokens",
+                COUNT(*)::text as "requestCount"
+             FROM public."langchain_token_usage"
+             WHERE "teamId" = $1 ${periodClause}
+             GROUP BY "userId"`,
+            [teamId],
+            'admin' // Use admin context for team-wide queries
+        )
+
+        const byUser: Record<string, { tokens: number; cost: number }> = {}
+        let totalTokens = 0
+        let totalCost = 0
+        let inputTokens = 0
+        let outputTokens = 0
+        let requestCount = 0
+
+        for (const row of result) {
+            byUser[row.userId] = {
+                tokens: parseInt(row.totalTokens, 10),
+                cost: parseFloat(row.totalCost),
+            }
+            totalTokens += parseInt(row.totalTokens || '0', 10)
+            totalCost += parseFloat(row.totalCost || '0')
+            inputTokens += parseInt(row.inputTokens || '0', 10)
+            outputTokens += parseInt(row.outputTokens || '0', 10)
+            requestCount += parseInt(row.requestCount || '0', 10)
+        }
+
+        return {
+            totalTokens,
+            totalCost,
+            inputTokens,
+            outputTokens,
+            requestCount,
+            byModel: {}, // Not grouped by model for team view
+            byUser,
+        }
+    },
+
+    // Helper: Get SQL period clause
+    getPeriodClause(period: Period): string {
+        switch (period) {
+            case 'today':
+                return `AND "createdAt" >= CURRENT_DATE`
+            case '7d':
+                return `AND "createdAt" >= now() - interval '7 days'`
+            case '30d':
+                return `AND "createdAt" >= now() - interval '30 days'`
+            case 'all':
+            default:
+                return ''
+        }
+    },
+}