Merge pull request #11 from constructive-io/feat/embedding-token-counts

feat(ollama): return real token counts from embedding endpoint

Author: pyramation

packages/ollama/README.md

@@ -50,9 +50,10 @@ await client.generate(
 // Pull a model to local cache
 await client.pullModel('mistral');
 
-// Generate embeddings
-const embedding = await client.generateEmbedding('Compute embeddings');
-console.log('Embedding vector length:', embedding.length);
+// Generate embeddings (with token count from /api/embed)
+const result = await client.generateEmbedding('Compute embeddings');
+console.log('Embedding vector length:', result.embedding.length);
+console.log('Prompt tokens:', result.promptTokens);
 
 // Delete a pulled model when done
 await client.deleteModel('mistral');
@@ -64,7 +65,7 @@ await client.deleteModel('mistral');
 - `.listModels(): Promise<string[]>`
 - `.showModel(model: string): Promise<{ capabilities?: string[] } | null>`
 - `.generate(input: GenerateInput, onChunk?: (chunk: string) => void): Promise<string | void>`
-- `.generateEmbedding(text: string, model?: string): Promise<number[]>` — defaults to `nomic-embed-text`
+- `.generateEmbedding(text: string, model?: string): Promise<EmbeddingResult>` — returns `{ embedding: number[], promptTokens: number }`, defaults to `nomic-embed-text`
 - `.pullModel(model: string): Promise<void>`
 - `.deleteModel(model: string): Promise<void>`
 
@@ -75,6 +76,10 @@ import { OllamaAdapter } from '@agentic-kit/ollama';
 
 const provider = new OllamaAdapter('http://localhost:11434');
 const model = provider.createModel('llama3');
+
+// Embeddings with real token counts
+const result = await provider.embed('Compute embeddings', 'nomic-embed-text');
+console.log(result.embedding.length, result.promptTokens);
 ```
 
 ## Local Live Tests

packages/ollama/__tests__/ollama.live.test.ts

@@ -205,13 +205,28 @@ describeExtended('Ollama live extended', () => {
     expect(output.trim().toLowerCase()).toContain('marble');
   });
 
-  itWithEmbeddings('generates local embeddings when an embed model is installed', async () => {
+  itWithEmbeddings('generates local embeddings with token count via /api/embed', async () => {
     const client = new OllamaClient(baseUrl);
-    const embedding = await client.generateEmbedding('hello world', embedModel);
+    const result = await client.generateEmbedding('hello world', embedModel);
+
+    expect(result).toHaveProperty('embedding');
+    expect(result).toHaveProperty('promptTokens');
+    expect(Array.isArray(result.embedding)).toBe(true);
+    expect(result.embedding.length).toBeGreaterThan(0);
+    expect(result.embedding.every((value) => Number.isFinite(value))).toBe(true);
+    expect(result.promptTokens).toBeGreaterThan(0);
+  });
+
+  itWithEmbeddings('OllamaAdapter.embed() returns embedding with token count', async () => {
+    const { OllamaAdapter } = require('../src/index');
+    const adapter = new OllamaAdapter(baseUrl);
+    const result = await adapter.embed('hello world', embedModel);
 
-    expect(Array.isArray(embedding)).toBe(true);
-    expect(embedding.length).toBeGreaterThan(0);
-    expect(embedding.every((value) => Number.isFinite(value))).toBe(true);
+    expect(result).toHaveProperty('embedding');
+    expect(result).toHaveProperty('promptTokens');
+    expect(Array.isArray(result.embedding)).toBe(true);
+    expect(result.embedding.length).toBeGreaterThan(0);
+    expect(result.promptTokens).toBeGreaterThan(0);
   });
 });
 

packages/ollama/src/index.ts

@@ -246,8 +246,17 @@ interface OllamaChatLine {
   response?: string;
 }
 
-interface OllamaEmbeddingResponse {
+interface OllamaEmbedResponse {
+  model: string;
+  embeddings: number[][];
+  total_duration?: number;
+  load_duration?: number;
+  prompt_eval_count?: number;
+}
+
+export interface EmbeddingResult {
   embedding: number[];
+  promptTokens: number;
 }
 
 export const OLLAMA_MODELS: ModelDescriptor[] = [];
@@ -297,18 +306,21 @@ export class OllamaClient {
     }
   }
 
-  async generateEmbedding(text: string, model = 'nomic-embed-text'): Promise<number[]> {
-    const response = await fetch(`${this.baseUrl}/api/embeddings`, {
+  async generateEmbedding(text: string, model = 'nomic-embed-text'): Promise<EmbeddingResult> {
+    const response = await fetch(`${this.baseUrl}/api/embed`, {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
-      body: JSON.stringify({ model, prompt: text }),
+      body: JSON.stringify({ model, input: text }),
     });
     if (!response.ok) {
       throw new Error(`generateEmbedding failed: ${response.status} ${response.statusText}`);
     }
 
-    const payload = (await response.json()) as OllamaEmbeddingResponse;
-    return payload.embedding;
+    const payload = (await response.json()) as OllamaEmbedResponse;
+    return {
+      embedding: payload.embeddings[0],
+      promptTokens: payload.prompt_eval_count ?? 0,
+    };
   }
 
   async generate(input: GenerateInput): Promise<string>;
@@ -385,6 +397,10 @@ export class OllamaAdapter {
     return this.client.listModels();
   }
 
+  async embed(text: string, model = 'nomic-embed-text'): Promise<EmbeddingResult> {
+    return this.client.generateEmbedding(text, model);
+  }
+
   stream(model: ModelDescriptor, context: Context, options?: StreamOptions): AssistantMessageEventStream {
     const stream = new DefaultAssistantMessageEventStream();
     const output = createAssistantMessage(model);