feat(adapters): extends Ollama embedding options

Ref: #176 Signed-off-by: Tomas Dvorak <toomas2d@gmail.com>
i-am-bee · Dec 13, 2024 · d3c9364 · d3c9364
1 parent ca17a18
commit d3c9364
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git 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<OllamaChatLLMOutput> {
     return extractModelMeta(model);
   }
 
-  // eslint-disable-next-line unused-imports/no-unused-vars
-  async embed(input: BaseMessage[][], options?: EmbeddingOptions): Promise<EmbeddingOutput> {
+  async embed(
+    input: BaseMessage[][],
+    options: OllamaEmbeddingOptions = {},
+  ): Promise<EmbeddingOutput> {
     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
@@ -58,6 +58,11 @@ interface Input {
   cache?: LLMCache<OllamaLLMOutput>;
 }
 
+export interface OllamaEmbeddingOptions extends EmbeddingOptions {
+  options?: Partial<Parameters>;
+  truncate?: boolean;
+}
+
 export class OllamaLLMOutput extends BaseLLMOutput {
   public readonly results: GenerateResponse[];
 
@@ -187,9 +192,13 @@ export class OllamaLLM extends LLM<OllamaLLMOutput> {
     return extractModelMeta(model);
   }
 
-  // eslint-disable-next-line unused-imports/no-unused-vars
-  async embed(input: LLMInput[], options?: EmbeddingOptions): Promise<EmbeddingOutput> {
-    const response = await this.client.embed({ model: this.modelId, input: input });
+  async embed(input: LLMInput[], options: OllamaEmbeddingOptions = {}): Promise<EmbeddingOutput> {
+    const response = await this.client.embed({
+      model: this.modelId,
+      input: input,
+      options: options?.options,
+      truncate: options?.truncate,
+    });
     return { embeddings: response.embeddings };
   }