gooseek/services/llm-svc/src/lib/models/providers/ollama/ollamaEmbedding.ts

import { Ollama } from 'ollama';
import BaseEmbedding from '../../base/embedding.js';
import type { Chunk } from '../../../types.js';

type OllamaConfig = {
  model: string;
  baseURL?: string;
};

class OllamaEmbedding extends BaseEmbedding<OllamaConfig> {
  ollamaClient: Ollama;

  constructor(protected config: OllamaConfig) {
    super(config);
    this.ollamaClient = new Ollama({
      host: this.config.baseURL ?? 'http://localhost:11434',
    });
  }

  async embedText(texts: string[]): Promise<number[][]> {
    const response = await this.ollamaClient.embed({
      input: texts,
      model: this.config.model,
    });
    return response.embeddings;
  }

  async embedChunks(chunks: Chunk[]): Promise<number[][]> {
    const response = await this.ollamaClient.embed({
      input: chunks.map((c) => c.content),
      model: this.config.model,
    });
    return response.embeddings;
  }
}

export default OllamaEmbedding;