Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions packages/ollama/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand All @@ -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>`

Expand All @@ -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
Expand Down
25 changes: 20 additions & 5 deletions packages/ollama/__tests__/ollama.live.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
});

Expand Down
28 changes: 22 additions & 6 deletions packages/ollama/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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[] = [];
Expand Down Expand Up @@ -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>;
Expand Down Expand Up @@ -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);
Expand Down
Loading