From fadf736f7034d5b2b685fdca46dac3d127f24e21 Mon Sep 17 00:00:00 2001
From: n4ze3m <n4ze3m@gmail.com>
Date: Thu, 23 May 2024 11:22:46 +0530
Subject: [PATCH] refactor: Parse keepAlive value in OllamaEmbeddingsPageAssist

---
 src/models/ChatOllama.ts      |  3 ++-
 src/models/OllamaEmbedding.ts |  7 ++++---
 src/models/utils/ollama.ts    | 10 +++++++++-
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/models/ChatOllama.ts b/src/models/ChatOllama.ts
index 69138bf..2a4c28b 100644
--- a/src/models/ChatOllama.ts
+++ b/src/models/ChatOllama.ts
@@ -15,6 +15,7 @@ import type { StringWithAutocomplete } from "@langchain/core/utils/types";
 import {
     createOllamaChatStream,
     createOllamaGenerateStream,
+    parseKeepAlive,
     type OllamaInput,
     type OllamaMessage,
 } from "./utils/ollama";
@@ -112,7 +113,7 @@ export class ChatOllama
         this.baseUrl = fields.baseUrl?.endsWith("/")
             ? fields.baseUrl.slice(0, -1)
             : fields.baseUrl ?? this.baseUrl;
-        this.keepAlive = fields.keepAlive ?? this.keepAlive;
+        this.keepAlive = parseKeepAlive(fields.keepAlive) ?? this.keepAlive;
         this.embeddingOnly = fields.embeddingOnly;
         this.f16KV = fields.f16KV;
         this.frequencyPenalty = fields.frequencyPenalty;
diff --git a/src/models/OllamaEmbedding.ts b/src/models/OllamaEmbedding.ts
index 2d57ef2..eadec2d 100644
--- a/src/models/OllamaEmbedding.ts
+++ b/src/models/OllamaEmbedding.ts
@@ -1,12 +1,13 @@
 import { Embeddings, EmbeddingsParams } from "@langchain/core/embeddings"
 import type { StringWithAutocomplete } from "@langchain/core/utils/types"
+import { parseKeepAlive } from "./utils/ollama"
 
 export interface OllamaInput {
   embeddingOnly?: boolean
   f16KV?: boolean
   frequencyPenalty?: number
   headers?: Record<string, string>
-  keepAlive?: string
+  keepAlive?: any
   logitsAll?: boolean
   lowVram?: boolean
   mainGpu?: number
@@ -98,7 +99,7 @@ interface OllamaEmbeddingsParams extends EmbeddingsParams {
   headers?: Record<string, string>
 
   /** Defaults to "5m" */
-  keepAlive?: string
+  keepAlive?: any
 
   /** Advanced Ollama API request parameters in camelCase, see
    * https://github.com/jmorganca/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values
@@ -138,7 +139,7 @@ export class OllamaEmbeddingsPageAssist extends Embeddings {
     }
 
     if (params?.keepAlive) {
-      this.keepAlive = params.keepAlive
+      this.keepAlive = parseKeepAlive(params.keepAlive)
     }
 
     if (params?.requestOptions) {
diff --git a/src/models/utils/ollama.ts b/src/models/utils/ollama.ts
index d3524fb..55ce05e 100644
--- a/src/models/utils/ollama.ts
+++ b/src/models/utils/ollama.ts
@@ -7,7 +7,7 @@ export interface OllamaInput {
   f16KV?: boolean;
   frequencyPenalty?: number;
   headers?: Record<string, string>;
-  keepAlive?: string;
+  keepAlive?: any;
   logitsAll?: boolean;
   lowVram?: boolean;
   mainGpu?: number;
@@ -198,4 +198,12 @@ export async function* createOllamaChatStream(
   options: OllamaCallOptions
 ): AsyncGenerator<OllamaChatGenerationChunk> {
   yield* createOllamaStream(`${baseUrl}/api/chat`, params, options);
+}
+
+
+export const parseKeepAlive = (keepAlive: any) => {
+  if (keepAlive === "-1") {
+    return -1
+  }
+  return keepAlive
 }
\ No newline at end of file