feat(self-host): ollama support (#1219)

Author: mogery

SELF_HOST.md

@@ -51,6 +51,11 @@ USE_DB_AUTHENTICATION=false
 # Provide your OpenAI API key here to enable AI features
 # OPENAI_API_KEY=
 
+# Experimental: Use Ollama
+# OPENAI_API_KEY=ollama
+# OPENAI_BASE_URL=http://localhost:11434/v1
+# MODEL_NAME=deepseek-r1:7b
+
 ## === Proxy ===
 # PROXY_SERVER can be a full URL (e.g. http://0.1.2.3:1234) or just an IP and port combo (e.g. 0.1.2.3:1234)
 # Do not uncomment PROXY_USERNAME and PROXY_PASSWORD if your proxy is unauthenticated

apps/api/src/lib/LLM-extraction/index.ts

@@ -67,7 +67,7 @@ export async function generateCompletions(
 
 export async function generateBasicCompletion(prompt: string) {
   const openai = new OpenAI();
-  const model = "gpt-4o";
+  const model = process.env.MODEL_NAME || "gpt-4o";
 
   const completion = await openai.chat.completions.create({
     temperature: 0,

apps/api/src/lib/extract/completions.ts

@@ -27,7 +27,7 @@
 // export async function generateBasicCompletion(prompt: string) {
 //   const openai = new OpenAI();
 //   const model: TiktokenModel =
-//     (process.env.MODEL_NAME as TiktokenModel) ?? "gpt-4o-mini";
+//     (process.env.MODEL_NAME as TiktokenModel) || "gpt-4o-mini";
 
 //   const completion = await openai.chat.completions.create({
 //     model,
@@ -48,7 +48,7 @@
 // }> {
 //   const openai = new OpenAI();
 //   const model: TiktokenModel =
-//     (process.env.MODEL_NAME as TiktokenModel) ?? "gpt-4o-mini";
+//     (process.env.MODEL_NAME as TiktokenModel) || "gpt-4o-mini";
 
 //   let extractionContent = pagesContent;
 //   let numTokens = 0;
@@ -109,7 +109,7 @@
 //       : { type: "json_object" },
 //   });
 
-//   if (jsonCompletion.choices[0].message.refusal !== null) {
+//   if (jsonCompletion.choices[0].message.refusal !== null && jsonCompletion.choices[0].message.refusal !== undefined) {
 //     throw new LLMRefusalError(jsonCompletion.choices[0].message.refusal);
 //   }
 

apps/api/src/lib/extract/completions/analyzeSchemaAndPrompt.ts

@@ -37,11 +37,11 @@ export async function analyzeSchemaAndPrompt(
       "isMultiEntity was true, but no multiEntityKeys",
     );
 
-  const model = "gpt-4o";
+  const model = process.env.MODEL_NAME || "gpt-4o";
 
   const openai = new OpenAI();
   const result = await openai.beta.chat.completions.parse({
-    model: model,
+    model,
     messages: [
       {
         role: "system",

apps/api/src/lib/extract/index/pinecone.ts

@@ -28,7 +28,7 @@ async function getEmbedding(text: string) {
   });
 
   const embedding = await openai.embeddings.create({
-    model: "text-embedding-3-small",
+    model: process.env.EMBEDDING_MODEL_NAME || "text-embedding-3-small",
     input: text,
     encoding_format: "float",
   });

apps/api/src/lib/extract/usage/llm-cost.ts

@@ -26,7 +26,7 @@ export function estimateTotalCost(tokenUsage: TokenUsage[]): number {
 export function estimateCost(tokenUsage: TokenUsage): number {
   let totalCost = 0;
   try {
-    let model = tokenUsage.model ?? process.env.MODEL_NAME ?? "gpt-4o-mini";
+    let model = tokenUsage.model ?? (process.env.MODEL_NAME || "gpt-4o-mini");
     const pricing = modelPrices[model] as ModelPricing;
 
     if (!pricing) {

apps/api/src/lib/generate-llmstxt/generate-llmstxt-service.ts

@@ -132,7 +132,7 @@ export async function performGenerateLlmsTxt(options: GenerateLLMsTextServiceOpt
           _logger.debug(`Generating description for ${document.metadata?.url}`);
 
           const completion = await openai.beta.chat.completions.parse({
-            model: "gpt-4o-mini",
+            model: process.env.MODEL_NAME || "gpt-4o-mini",
             messages: [
               {
                 role: "user", 

apps/api/src/lib/llm/generate.ts

@@ -124,7 +124,7 @@ export async function generateOpenAICompletions(
   markdown?: string,
   previousWarning?: string,
   isExtractEndpoint?: boolean,
-  model: TiktokenModel = (process.env.MODEL_NAME as TiktokenModel) ??
+  model: TiktokenModel = (process.env.MODEL_NAME as TiktokenModel) ||
     "gpt-4o-mini",
 ): Promise<{
   extract: any;
@@ -151,15 +151,21 @@ export async function generateOpenAICompletions(
 
   // count number of tokens
   let numTokens = 0;
-  const encoder = encoding_for_model(model as TiktokenModel);
   try {
     // Encode the message into tokens
-    const tokens = encoder.encode(markdown);
-
-    // Return the number of tokens
-    numTokens = tokens.length;
+    const encoder = encoding_for_model(model as TiktokenModel);
+    
+    try {
+      const tokens = encoder.encode(markdown);
+      numTokens = tokens.length;
+    } catch (e) {
+      throw e;
+    } finally {
+      // Free the encoder resources after use
+      encoder.free();
+    }
   } catch (error) {
-    logger.warn("Calculating num tokens of string failed", { error, markdown });
+    logger.warn("Calculating num tokens of string failed", { error });
 
     markdown = markdown.slice(0, maxTokensSafe * modifier);
 
@@ -168,9 +174,6 @@ export async function generateOpenAICompletions(
       maxTokensSafe +
       ") we support.";
     warning = previousWarning === undefined ? w : w + " " + previousWarning;
-  } finally {
-    // Free the encoder resources after use
-    encoder.free();
   }
 
   if (numTokens > maxTokensSafe) {
@@ -247,7 +250,7 @@ export async function generateOpenAICompletions(
       : { type: "json_object" },
   });
 
-  if (jsonCompletion.choices[0].message.refusal !== null) {
+  if (jsonCompletion.choices[0].message.refusal !== null && jsonCompletion.choices[0].message.refusal !== undefined) {
     throw new LLMRefusalError(jsonCompletion.choices[0].message.refusal);
   }
 
@@ -351,7 +354,7 @@ export async function generateSchemaFromPrompt(prompt: string): Promise<any> {
   for (const temp of temperatures) {
     try {
       const result = await openai.beta.chat.completions.parse({
-        model: "gpt-4o",
+        model: process.env.MODEL_NAME || "gpt-4o",
         temperature: temp,
         messages: [
           {
@@ -392,7 +395,7 @@ Return a valid JSON schema object with properties that would capture the informa
         },
       });
 
-      if (result.choices[0].message.refusal !== null) {
+      if (result.choices[0].message.refusal !== null && result.choices[0].message.refusal !== undefined) {
         throw new Error("LLM refused to generate schema");
       }