Add custom fetch to embedder for Ollama (#5180)
Refactor ollama timeout to be shared. Add custom fetch to embedder for ollama as well
This commit is contained in:
parent
dc0bdf112b
commit
4e3bcfc616
@ -33,7 +33,7 @@ class OllamaAILLM {
|
||||
this.client = new Ollama({
|
||||
host: this.basePath,
|
||||
headers: headers,
|
||||
fetch: this.#applyFetch(),
|
||||
fetch: OllamaAILLM.applyOllamaFetch(),
|
||||
});
|
||||
this.embedder = embedder ?? new NativeEmbedder();
|
||||
this.defaultTemp = 0.7;
|
||||
@ -132,7 +132,7 @@ class OllamaAILLM {
|
||||
* for machines which run responses very slowly.
|
||||
* @returns {Function} The custom fetch function.
|
||||
*/
|
||||
#applyFetch() {
|
||||
static applyOllamaFetch() {
|
||||
try {
|
||||
if (!("OLLAMA_RESPONSE_TIMEOUT" in process.env)) return fetch;
|
||||
const { Agent } = require("undici");
|
||||
@ -140,7 +140,7 @@ class OllamaAILLM {
|
||||
let timeout = process.env.OLLAMA_RESPONSE_TIMEOUT;
|
||||
|
||||
if (!timeout || isNaN(Number(timeout)) || Number(timeout) <= 5 * 60_000) {
|
||||
this.#log(
|
||||
OllamaAILLM.#slog(
|
||||
"Timeout option was not set, is not a number, or is less than 5 minutes in ms - falling back to default",
|
||||
{ timeout }
|
||||
);
|
||||
@ -155,10 +155,13 @@ class OllamaAILLM {
|
||||
};
|
||||
|
||||
const humanDiff = moment.duration(timeout).humanize();
|
||||
this.#log(`Applying custom fetch w/timeout of ${humanDiff}.`);
|
||||
OllamaAILLM.#slog(`Applying custom fetch w/timeout of ${humanDiff}.`);
|
||||
return noTimeoutFetch;
|
||||
} catch (error) {
|
||||
this.#log("Error applying custom fetch - using default fetch", error);
|
||||
OllamaAILLM.#slog(
|
||||
"Error applying custom fetch - using default fetch",
|
||||
error
|
||||
);
|
||||
return fetch;
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
const { maximumChunkLength } = require("../../helpers");
|
||||
const { Ollama } = require("ollama");
|
||||
const { OllamaAILLM } = require("../../AiProviders/ollama");
|
||||
|
||||
class OllamaEmbedder {
|
||||
constructor() {
|
||||
@ -16,10 +17,15 @@ class OllamaEmbedder {
|
||||
: 1;
|
||||
this.embeddingMaxChunkLength = maximumChunkLength();
|
||||
this.authToken = process.env.OLLAMA_AUTH_TOKEN;
|
||||
|
||||
const headers = this.authToken
|
||||
? { Authorization: `Bearer ${this.authToken}` }
|
||||
: {};
|
||||
this.client = new Ollama({ host: this.basePath, headers });
|
||||
this.client = new Ollama({
|
||||
host: this.basePath,
|
||||
headers,
|
||||
fetch: OllamaAILLM.applyOllamaFetch(),
|
||||
});
|
||||
this.log(
|
||||
`initialized with model ${this.model} at ${this.basePath}. Batch size: ${this.maxConcurrentChunks}, num_ctx: ${this.embeddingMaxChunkLength}`
|
||||
);
|
||||
|
||||
@ -28,7 +28,7 @@ class OllamaProvider extends InheritMultiple([Provider, UnTooled]) {
|
||||
this._client = new Ollama({
|
||||
host: basePath,
|
||||
headers: headers,
|
||||
fetch: this.#applyFetch(),
|
||||
fetch: OllamaAILLM.applyOllamaFetch(),
|
||||
});
|
||||
this.model = model;
|
||||
this.verbose = true;
|
||||
@ -526,46 +526,6 @@ class OllamaProvider extends InheritMultiple([Provider, UnTooled]) {
|
||||
getCost(_usage) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply a custom fetch function to the Ollama client.
|
||||
* This is useful when we want to bypass the default 5m timeout for global fetch
|
||||
* for machines which run responses very slowly.
|
||||
* @returns {Function} The custom fetch function.
|
||||
*/
|
||||
#applyFetch() {
|
||||
try {
|
||||
if (!("OLLAMA_RESPONSE_TIMEOUT" in process.env)) return fetch;
|
||||
const { Agent } = require("undici");
|
||||
const moment = require("moment");
|
||||
let timeout = process.env.OLLAMA_RESPONSE_TIMEOUT;
|
||||
|
||||
if (!timeout || isNaN(Number(timeout)) || Number(timeout) <= 5 * 60_000) {
|
||||
this.providerLog(
|
||||
"Timeout option was not set, is not a number, or is less than 5 minutes in ms - falling back to default",
|
||||
{ timeout }
|
||||
);
|
||||
return fetch;
|
||||
} else timeout = Number(timeout);
|
||||
|
||||
const noTimeoutFetch = (input, init = {}) => {
|
||||
return fetch(input, {
|
||||
...init,
|
||||
dispatcher: new Agent({ headersTimeout: timeout }),
|
||||
});
|
||||
};
|
||||
|
||||
const humanDiff = moment.duration(timeout).humanize();
|
||||
this.providerLog(`Applying custom fetch w/timeout of ${humanDiff}.`);
|
||||
return noTimeoutFetch;
|
||||
} catch (error) {
|
||||
this.providerLog(
|
||||
"Error applying custom fetch - using default fetch",
|
||||
error
|
||||
);
|
||||
return fetch;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = OllamaProvider;
|
||||
|
||||
Loading…
Reference in New Issue
Block a user