graphile-llm 0.8.0 → 0.9.0

package/__tests__/graphile-llm.test.js

@@ -36,19 +36,19 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-const path_1 = require("path");
 const ollama_1 = __importDefault(require("@agentic-kit/ollama"));
-const graphile_test_1 = require("graphile-test");
 const graphile_connection_filter_1 = require("graphile-connection-filter");
+const pgvector_1 = require("graphile-search/adapters/pgvector");
 const vector_codec_1 = require("graphile-search/codecs/vector-codec");
 const plugin_1 = require("graphile-search/plugin");
-const pgvector_1 = require("graphile-search/adapters/pgvector");
+const graphile_test_1 = require("graphile-test");
+const path_1 = require("path");
+const chat_1 = require("../../src/chat");
+const embedder_1 = require("../../src/embedder");
 const llm_module_plugin_1 = require("../../src/plugins/llm-module-plugin");
-const text_search_plugin_1 = require("../../src/plugins/text-search-plugin");
-const text_mutation_plugin_1 = require("../../src/plugins/text-mutation-plugin");
 const rag_plugin_1 = require("../../src/plugins/rag-plugin");
-const embedder_1 = require("../../src/embedder");
-const chat_1 = require("../../src/chat");
+const text_mutation_plugin_1 = require("../../src/plugins/text-mutation-plugin");
+const text_search_plugin_1 = require("../../src/plugins/text-search-plugin");
 // ─── @agentic-kit/ollama client ─────────────────────────────────────────────
 const ollamaClient = new ollama_1.default('http://localhost:11434');
 async function ensureNomicModel() {
@@ -68,14 +68,14 @@ describe('Embedder abstraction', () => {
             const embedder = (0, embedder_1.buildEmbedder)({
                 provider: 'ollama',
                 model: 'nomic-embed-text',
-                baseUrl: 'http://localhost:11434',
+                baseUrl: 'http://localhost:11434'
             });
             expect(embedder).not.toBeNull();
             expect(typeof embedder).toBe('function');
         });
         it('returns null for unknown provider', () => {
             const embedder = (0, embedder_1.buildEmbedder)({
-                provider: 'unknown-provider',
+                provider: 'unknown-provider'
             });
             expect(embedder).toBeNull();
         });
@@ -90,7 +90,7 @@ describe('Embedder abstraction', () => {
             const moduleData = {
                 embedding_provider: 'ollama',
                 embedding_model: 'nomic-embed-text',
-                embedding_base_url: 'http://localhost:11434',
+                embedding_base_url: 'http://localhost:11434'
             };
             const embedder = (0, embedder_1.buildEmbedderFromModule)(moduleData);
             expect(embedder).not.toBeNull();
@@ -98,7 +98,7 @@ describe('Embedder abstraction', () => {
         });
         it('returns null for unsupported provider in module data', () => {
             const moduleData = {
-                embedding_provider: 'unsupported',
+                embedding_provider: 'unsupported'
             };
             const embedder = (0, embedder_1.buildEmbedderFromModule)(moduleData);
             expect(embedder).toBeNull();
@@ -121,7 +121,7 @@ describe('Embedder abstraction', () => {
                 ...originalEnv,
                 EMBEDDER_PROVIDER: 'ollama',
                 EMBEDDER_MODEL: 'nomic-embed-text',
-                EMBEDDER_BASE_URL: 'http://localhost:11434',
+                EMBEDDER_BASE_URL: 'http://localhost:11434'
             };
             const embedder = (0, embedder_1.buildEmbedderFromEnv)();
             expect(embedder).not.toBeNull();
@@ -139,7 +139,7 @@ describe('graphile-llm schema enrichment', () => {
     let query;
     beforeAll(async () => {
         const unifiedPlugin = (0, plugin_1.createUnifiedSearchPlugin)({
-            adapters: [(0, pgvector_1.createPgvectorAdapter)()],
+            adapters: [(0, pgvector_1.createPgvectorAdapter)()]
         });
         const testPreset = {
             extends: [(0, graphile_connection_filter_1.ConnectionFilterPreset)()],
@@ -152,18 +152,18 @@ describe('graphile-llm schema enrichment', () => {
                     defaultEmbedder: {
                         provider: 'ollama',
                         model: 'nomic-embed-text',
-                        baseUrl: 'http://localhost:11434',
-                    },
+                        baseUrl: 'http://localhost:11434'
+                    }
                 }),
                 (0, text_search_plugin_1.createLlmTextSearchPlugin)(),
-                (0, text_mutation_plugin_1.createLlmTextMutationPlugin)(),
-            ],
+                (0, text_mutation_plugin_1.createLlmTextMutationPlugin)()
+            ]
         };
         const connections = await (0, graphile_test_1.getConnections)({
             schemas: ['llm_test'],
             preset: testPreset,
             useRoot: true,
-            authRole: 'postgres',
+            authRole: 'postgres'
         }, [graphile_test_1.seed.sqlfile([(0, path_1.join)(__dirname, './setup.sql')])]);
         db = connections.db;
         teardown = connections.teardown;
@@ -292,33 +292,36 @@ describe('graphile-llm with real Ollama embedding', () => {
         const embedder = (0, embedder_1.buildEmbedder)({
             provider: 'ollama',
             model: 'nomic-embed-text',
-            baseUrl: 'http://localhost:11434',
+            baseUrl: 'http://localhost:11434'
         });
         expect(embedder).not.toBeNull();
-        const vector = await embedder('Machine learning is transforming AI');
+        const result = await embedder('Machine learning is transforming AI');
         // nomic-embed-text produces 768-dimensional vectors
-        expect(Array.isArray(vector)).toBe(true);
-        expect(vector.length).toBe(768);
+        expect(Array.isArray(result.embedding)).toBe(true);
+        expect(result.embedding.length).toBe(768);
+        expect(result.promptTokens).toBeGreaterThan(0);
         // All elements should be numbers
-        for (const v of vector) {
+        for (const v of result.embedding) {
             expect(typeof v).toBe('number');
             expect(Number.isFinite(v)).toBe(true);
         }
         // Vector should not be all zeros
-        const magnitude = Math.sqrt(vector.reduce((sum, v) => sum + v * v, 0));
+        const magnitude = Math.sqrt(result.embedding.reduce((sum, v) => sum + v * v, 0));
         expect(magnitude).toBeGreaterThan(0);
     });
     it('should produce different vectors for semantically different text', async () => {
         const embedder = (0, embedder_1.buildEmbedder)({
             provider: 'ollama',
             model: 'nomic-embed-text',
-            baseUrl: 'http://localhost:11434',
+            baseUrl: 'http://localhost:11434'
         });
         expect(embedder).not.toBeNull();
-        const [vecA, vecB] = await Promise.all([
+        const [resultA, resultB] = await Promise.all([
             embedder('Artificial intelligence and machine learning'),
-            embedder('Cooking recipes for Italian pasta dishes'),
+            embedder('Cooking recipes for Italian pasta dishes')
         ]);
+        const vecA = resultA.embedding;
+        const vecB = resultB.embedding;
         expect(vecA.length).toBe(768);
         expect(vecB.length).toBe(768);
         // Compute cosine similarity
@@ -338,13 +341,15 @@ describe('graphile-llm with real Ollama embedding', () => {
         const embedder = (0, embedder_1.buildEmbedder)({
             provider: 'ollama',
             model: 'nomic-embed-text',
-            baseUrl: 'http://localhost:11434',
+            baseUrl: 'http://localhost:11434'
         });
         expect(embedder).not.toBeNull();
-        const [vecA, vecB] = await Promise.all([
+        const [resultA, resultB] = await Promise.all([
             embedder('Machine learning and artificial intelligence'),
-            embedder('AI and ML are subfields of computer science'),
+            embedder('AI and ML are subfields of computer science')
         ]);
+        const vecA = resultA.embedding;
+        const vecB = resultB.embedding;
         expect(vecA.length).toBe(768);
         expect(vecB.length).toBe(768);
         // Compute cosine similarity
@@ -360,11 +365,14 @@ describe('graphile-llm with real Ollama embedding', () => {
         // Semantically similar texts should have high similarity
         expect(cosineSimilarity).toBeGreaterThan(0.5);
     });
-    it('should produce embeddings via @agentic-kit/ollama OllamaClient directly', async () => {
-        const vector = await ollamaClient.generateEmbedding('Testing the agentic-kit Ollama client directly', 'nomic-embed-text');
-        expect(Array.isArray(vector)).toBe(true);
-        expect(vector.length).toBe(768);
-        for (const v of vector) {
+    it('should produce embeddings with token count via @agentic-kit/ollama OllamaClient directly', async () => {
+        const result = await ollamaClient.generateEmbedding('Testing the agentic-kit Ollama client directly', 'nomic-embed-text');
+        expect(result).toHaveProperty('embedding');
+        expect(result).toHaveProperty('promptTokens');
+        expect(Array.isArray(result.embedding)).toBe(true);
+        expect(result.embedding.length).toBe(768);
+        expect(result.promptTokens).toBeGreaterThan(0);
+        for (const v of result.embedding) {
             expect(typeof v).toBe('number');
             expect(Number.isFinite(v)).toBe(true);
         }
@@ -379,14 +387,14 @@ describe('Chat completion abstraction', () => {
             const chat = (0, chat_1.buildChatCompleter)({
                 provider: 'ollama',
                 model: 'llama3',
-                baseUrl: 'http://localhost:11434',
+                baseUrl: 'http://localhost:11434'
             });
             expect(chat).not.toBeNull();
             expect(typeof chat).toBe('function');
         });
         it('returns null for unknown provider', () => {
             const chat = (0, chat_1.buildChatCompleter)({
-                provider: 'unknown-provider',
+                provider: 'unknown-provider'
             });
             expect(chat).toBeNull();
         });
@@ -402,7 +410,7 @@ describe('Chat completion abstraction', () => {
                 embedding_provider: 'ollama',
                 chat_provider: 'ollama',
                 chat_model: 'llama3',
-                chat_base_url: 'http://localhost:11434',
+                chat_base_url: 'http://localhost:11434'
             };
             const chat = (0, chat_1.buildChatCompleterFromModule)(moduleData);
             expect(chat).not.toBeNull();
@@ -410,7 +418,7 @@ describe('Chat completion abstraction', () => {
         });
         it('returns null when chat_provider is not set', () => {
             const moduleData = {
-                embedding_provider: 'ollama',
+                embedding_provider: 'ollama'
             };
             const chat = (0, chat_1.buildChatCompleterFromModule)(moduleData);
             expect(chat).toBeNull();
@@ -433,7 +441,7 @@ describe('Chat completion abstraction', () => {
                 ...originalEnv,
                 CHAT_PROVIDER: 'ollama',
                 CHAT_MODEL: 'llama3',
-                CHAT_BASE_URL: 'http://localhost:11434',
+                CHAT_BASE_URL: 'http://localhost:11434'
             };
             const chat = (0, chat_1.buildChatCompleterFromEnv)();
             expect(chat).not.toBeNull();
@@ -472,10 +480,10 @@ function makeTestSmartTagsPlugin(tagsByTable) {
                             Object.assign(c.extensions.tags, tags);
                         }
                         return _;
-                    },
-                },
-            },
-        },
+                    }
+                }
+            }
+        }
     };
 }
 describe('RAG plugin schema enrichment', () => {
@@ -484,7 +492,7 @@ describe('RAG plugin schema enrichment', () => {
     let query;
     beforeAll(async () => {
         const unifiedPlugin = (0, plugin_1.createUnifiedSearchPlugin)({
-            adapters: [(0, pgvector_1.createPgvectorAdapter)()],
+            adapters: [(0, pgvector_1.createPgvectorAdapter)()]
         });
         const smartTagsPlugin = makeTestSmartTagsPlugin({
             articles: {
@@ -493,16 +501,22 @@ describe('RAG plugin schema enrichment', () => {
                     parentFk: 'parent_id',
                     parentPk: 'id',
                     embeddingField: 'embedding',
-                    contentField: 'content',
-                },
-            },
+                    contentField: 'content'
+                }
+            }
         });
-        // Mock embedder that returns a fixed 3-dim vector
-        const mockEmbedder = async (_text) => [1, 0, 0];
-        // Mock chat completer that returns a canned response
+        // Mock embedder that returns a fixed 3-dim vector with token count
+        const mockEmbedder = async (_text) => ({
+            embedding: [1, 0, 0],
+            promptTokens: 5
+        });
+        // Mock chat completer that returns a canned response with usage
         const mockChatCompleter = async (messages) => {
             const userMessage = messages.find((m) => m.role === 'user');
-            return `Mock answer for: ${userMessage?.content || 'unknown'}`;
+            return {
+                content: `Mock answer for: ${userMessage?.content || 'unknown'}`,
+                usage: { input: 10, output: 15, reasoning: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 25 }
+            };
         };
         const testPreset = {
             extends: [(0, graphile_connection_filter_1.ConnectionFilterPreset)()],
@@ -514,13 +528,13 @@ describe('RAG plugin schema enrichment', () => {
                     defaultEmbedder: {
                         provider: 'ollama',
                         model: 'nomic-embed-text',
-                        baseUrl: 'http://localhost:11434',
-                    },
+                        baseUrl: 'http://localhost:11434'
+                    }
                 }),
                 (0, text_search_plugin_1.createLlmTextSearchPlugin)(),
                 (0, text_mutation_plugin_1.createLlmTextMutationPlugin)(),
-                (0, rag_plugin_1.createLlmRagPlugin)(),
-            ],
+                (0, rag_plugin_1.createLlmRagPlugin)()
+            ]
         };
         // Override the embedder and chat completer on the build context
         // by wrapping the LlmModulePlugin's build hook
@@ -533,20 +547,20 @@ describe('RAG plugin schema enrichment', () => {
                     build(build) {
                         return build.extend(build, {
                             llmEmbedder: mockEmbedder,
-                            llmChatCompleter: mockChatCompleter,
+                            llmChatCompleter: mockChatCompleter
                         }, 'TestOverridePlugin overriding embedder and chat completer');
-                    },
-                },
-            },
+                    }
+                }
+            }
         };
         const connections = await (0, graphile_test_1.getConnections)({
             schemas: ['llm_test'],
             preset: {
                 ...testPreset,
-                plugins: [...testPreset.plugins, overridePlugin],
+                plugins: [...testPreset.plugins, overridePlugin]
             },
             useRoot: true,
-            authRole: 'postgres',
+            authRole: 'postgres'
         }, [graphile_test_1.seed.sqlfile([(0, path_1.join)(__dirname, './setup.sql')])]);
         db = connections.db;
         teardown = connections.teardown;
@@ -670,7 +684,7 @@ describe('GraphileLlmPreset toggles', () => {
     it('enableRag=false excludes RAG plugin (no ragQuery field)', async () => {
         const { GraphileLlmPreset } = await Promise.resolve().then(() => __importStar(require('../../src/preset')));
         const preset = GraphileLlmPreset({
-            enableRag: false,
+            enableRag: false
         });
         const pluginNames = preset.plugins.map((p) => p.name);
         expect(pluginNames).not.toContain('LlmRagPlugin');
@@ -678,7 +692,7 @@ describe('GraphileLlmPreset toggles', () => {
     it('enableRag=true includes RAG plugin', async () => {
         const { GraphileLlmPreset } = await Promise.resolve().then(() => __importStar(require('../../src/preset')));
         const preset = GraphileLlmPreset({
-            enableRag: true,
+            enableRag: true
         });
         const pluginNames = preset.plugins.map((p) => p.name);
         expect(pluginNames).toContain('LlmRagPlugin');
@@ -686,7 +700,7 @@ describe('GraphileLlmPreset toggles', () => {
     it('enableTextSearch=false excludes text search plugin', async () => {
         const { GraphileLlmPreset } = await Promise.resolve().then(() => __importStar(require('../../src/preset')));
         const preset = GraphileLlmPreset({
-            enableTextSearch: false,
+            enableTextSearch: false
         });
         const pluginNames = preset.plugins.map((p) => p.name);
         expect(pluginNames).not.toContain('LlmTextSearchPlugin');
@@ -696,7 +710,7 @@ describe('GraphileLlmPreset toggles', () => {
     it('enableTextMutations=false excludes text mutation plugin', async () => {
         const { GraphileLlmPreset } = await Promise.resolve().then(() => __importStar(require('../../src/preset')));
         const preset = GraphileLlmPreset({
-            enableTextMutations: false,
+            enableTextMutations: false
         });
         const pluginNames = preset.plugins.map((p) => p.name);
         expect(pluginNames).not.toContain('LlmTextMutationPlugin');
@@ -706,7 +720,7 @@ describe('GraphileLlmPreset toggles', () => {
         const preset = GraphileLlmPreset({
             enableTextSearch: false,
             enableTextMutations: false,
-            enableRag: false,
+            enableRag: false
         });
         const pluginNames = preset.plugins.map((p) => p.name);
         expect(pluginNames).toEqual(['LlmModulePlugin']);

package/chat.js

@@ -12,45 +12,55 @@
  *   2. The preset's `defaultChatCompleter` option (fallback for dev/testing)
  *   3. Environment variables (CHAT_PROVIDER, CHAT_MODEL, CHAT_BASE_URL)
  */
-var __importDefault = (this && this.__importDefault) || function (mod) {
-    return (mod && mod.__esModule) ? mod : { "default": mod };
-};
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.buildChatCompleter = buildChatCompleter;
 exports.buildChatCompleterFromModule = buildChatCompleterFromModule;
 exports.buildChatCompleterFromEnv = buildChatCompleterFromEnv;
-const ollama_1 = __importDefault(require("@agentic-kit/ollama"));
+const ollama_1 = require("@agentic-kit/ollama");
 const env_1 = require("./env");
 // ─── Built-in Providers ─────────────────────────────────────────────────────
 /**
  * Create an Ollama-based chat completion function.
  *
- * Uses OllamaClient.generate() with a messages array, which internally
- * routes to the /api/chat endpoint.
+ * Uses OllamaAdapter.stream() to get both response content and real token
+ * usage counts from the provider (prompt_eval_count, eval_count).
  */
 function createOllamaChatCompleter(baseUrl = 'http://localhost:11434', model = 'llama3') {
-    const client = new ollama_1.default(baseUrl);
+    const adapter = new ollama_1.OllamaAdapter(baseUrl);
     return async (messages, options) => {
-        // Build the input for OllamaClient.generate() in chat mode
-        const input = {
-            model,
-            messages: messages.filter((m) => m.role !== 'system'),
-        };
-        // Extract system message if present
         const systemMsg = messages.find((m) => m.role === 'system');
-        if (systemMsg) {
-            input.system = systemMsg.content;
-        }
-        if (options?.temperature !== undefined) {
-            input.temperature = options.temperature;
-        }
-        const startTime = Date.now();
-        const response = await client.generate(input);
-        const latencyMs = Date.now() - startTime;
-        // Token count logging (metering deferred to billing system)
-        console.log(`[graphile-llm] Chat completion: model=${model}, latency=${latencyMs}ms, ` +
-            `messages=${messages.length}`);
-        return response;
+        const nonSystem = messages.filter((m) => m.role !== 'system');
+        const modelDesc = adapter.createModel(model, {
+            maxOutputTokens: options?.maxTokens
+        });
+        const context = {
+            systemPrompt: systemMsg?.content,
+            messages: nonSystem.map((m) => ({
+                role: m.role,
+                content: m.content,
+                timestamp: Date.now()
+            }))
+        };
+        const stream = adapter.stream(modelDesc, context, {
+            temperature: options?.temperature,
+            maxTokens: options?.maxTokens
+        });
+        const result = await stream.result();
+        const content = result.content
+            .filter((block) => block.type === 'text')
+            .map((block) => block.text)
+            .join('');
+        return {
+            content,
+            usage: {
+                input: result.usage.input,
+                output: result.usage.output,
+                reasoning: result.usage.reasoning,
+                cacheRead: result.usage.cacheRead,
+                cacheWrite: result.usage.cacheWrite,
+                totalTokens: result.usage.totalTokens
+            }
+        };
     };
 }
 // ─── Chat Completer Construction ────────────────────────────────────────────
@@ -81,7 +91,7 @@ function buildChatCompleterFromModule(data) {
     return buildChatCompleter({
         provider: data.chat_provider,
         model: data.chat_model,
-        baseUrl: data.chat_base_url,
+        baseUrl: data.chat_base_url
     });
 }
 /**

package/config-cache.js

@@ -56,7 +56,7 @@ const INFERENCE_LOG_MODULE_SQL = `
 const billingCache = new graphile_cache_1.ModuleConfigCache({
     name: 'billing-config',
     ttlMs: 5 * 60 * 1000, // 5 minutes
-    max: 50,
+    max: 50
 });
 // ─── Resolution Functions ───────────────────────────────────────────────────
 /**
@@ -77,7 +77,7 @@ async function resolveInferenceLogConfig(pgClient, databaseId) {
             return null;
         return {
             schema: row.schema,
-            tableName: row.table_name,
+            tableName: row.table_name
         };
     }
     catch {
@@ -101,7 +101,7 @@ async function resolveBillingConfig(pgClient, databaseId) {
             privateSchema: row.private_schema,
             recordUsageFunction: row.record_usage_function,
             // The check_billing_quota function name follows the inflection pattern
-            checkBillingQuotaFunction: 'check_billing_quota',
+            checkBillingQuotaFunction: 'check_billing_quota'
         };
     }
     catch {
@@ -123,7 +123,7 @@ async function getLlmBillingConfig(pgClient, databaseId) {
         return cached;
     const [billing, inferenceLog] = await Promise.all([
         resolveBillingConfig(pgClient, databaseId),
-        resolveInferenceLogConfig(pgClient, databaseId),
+        resolveInferenceLogConfig(pgClient, databaseId)
     ]);
     const entry = { billing, inferenceLog };
     billingCache.set(databaseId, entry);

package/embedder.js

@@ -22,6 +22,8 @@ const env_1 = require("./env");
 // ─── Built-in Providers ─────────────────────────────────────────────────────
 /**
  * Create an Ollama-based embedder function.
+ *
+ * Uses the /api/embed endpoint which returns prompt_eval_count (real token count).
  */
 function createOllamaEmbedder(baseUrl = 'http://localhost:11434', model = 'nomic-embed-text') {
     const client = new ollama_1.default(baseUrl);
@@ -55,7 +57,7 @@ function buildEmbedderFromModule(data) {
     return buildEmbedder({
         provider: data.embedding_provider,
         model: data.embedding_model,
-        baseUrl: data.embedding_base_url,
+        baseUrl: data.embedding_base_url
     });
 }
 /**

package/env.js

@@ -21,13 +21,13 @@ const LLM_DEFAULTS = {
     embedding: {
         provider: 'ollama',
         model: 'nomic-embed-text',
-        baseUrl: 'http://localhost:11434',
+        baseUrl: 'http://localhost:11434'
     },
     chat: {
         provider: 'ollama',
         model: 'llama3',
-        baseUrl: 'http://localhost:11434',
-    },
+        baseUrl: 'http://localhost:11434'
+    }
 };
 // ─── Resolution ─────────────────────────────────────────────────────────────
 /**
@@ -41,12 +41,12 @@ function getLlmEnvOptions() {
         embedding: {
             provider: process.env.EMBEDDER_PROVIDER ?? LLM_DEFAULTS.embedding.provider,
             model: process.env.EMBEDDER_MODEL ?? LLM_DEFAULTS.embedding.model,
-            baseUrl: process.env.EMBEDDER_BASE_URL ?? LLM_DEFAULTS.embedding.baseUrl,
+            baseUrl: process.env.EMBEDDER_BASE_URL ?? LLM_DEFAULTS.embedding.baseUrl
         },
         chat: {
             provider: process.env.CHAT_PROVIDER ?? LLM_DEFAULTS.chat.provider,
             model: process.env.CHAT_MODEL ?? LLM_DEFAULTS.chat.model,
-            baseUrl: process.env.CHAT_BASE_URL ?? LLM_DEFAULTS.chat.baseUrl,
-        },
+            baseUrl: process.env.CHAT_BASE_URL ?? LLM_DEFAULTS.chat.baseUrl
+        }
     };
 }