diff --git a/src/adapters/ollama/chat.ts b/src/adapters/ollama/chat.ts index e51400c6..40d359f7 100644 --- a/src/adapters/ollama/chat.ts +++ b/src/adapters/ollama/chat.ts @@ -17,7 +17,6 @@ import { AsyncStream, BaseLLMTokenizeOutput, - EmbeddingOptions, EmbeddingOutput, ExecutionOptions, GenerateOptions, @@ -42,6 +41,7 @@ import { retrieveVersion, } from "@/adapters/ollama/shared.js"; import { getEnv } from "@/internals/env.js"; +import { OllamaEmbeddingOptions } from "@/adapters/ollama/llm.js"; export class OllamaChatLLMOutput extends ChatLLMOutput { public readonly results: ChatResponse[]; @@ -160,11 +160,15 @@ export class OllamaChatLLM extends ChatLLM { return extractModelMeta(model); } - // eslint-disable-next-line unused-imports/no-unused-vars - async embed(input: BaseMessage[][], options?: EmbeddingOptions): Promise { + async embed( + input: BaseMessage[][], + options: OllamaEmbeddingOptions = {}, + ): Promise { const response = await this.client.embed({ model: this.modelId, input: input.flatMap((messages) => messages).flatMap((msg) => msg.text), + options: options?.options, + truncate: options?.truncate, }); return { embeddings: response.embeddings }; } diff --git a/src/adapters/ollama/llm.ts b/src/adapters/ollama/llm.ts index ce8704a9..47ca5ecd 100644 --- a/src/adapters/ollama/llm.ts +++ b/src/adapters/ollama/llm.ts @@ -58,6 +58,11 @@ interface Input { cache?: LLMCache; } +export interface OllamaEmbeddingOptions extends EmbeddingOptions { + options?: Partial; + truncate?: boolean; +} + export class OllamaLLMOutput extends BaseLLMOutput { public readonly results: GenerateResponse[]; @@ -187,9 +192,13 @@ export class OllamaLLM extends LLM { return extractModelMeta(model); } - // eslint-disable-next-line unused-imports/no-unused-vars - async embed(input: LLMInput[], options?: EmbeddingOptions): Promise { - const response = await this.client.embed({ model: this.modelId, input: input }); + async embed(input: LLMInput[], options: OllamaEmbeddingOptions = {}): Promise { + const response = await this.client.embed({ + model: this.modelId, + input: input, + options: options?.options, + truncate: options?.truncate, + }); return { embeddings: response.embeddings }; }