From 2b9a1ae3167a855a1dbe33cb8d4a832b3aa98829 Mon Sep 17 00:00:00 2001
From: Sarge Joseph <github.com@toilemaitre.com>
Date: Thu, 13 Mar 2025 15:59:27 -0400
Subject: [PATCH] Add an Ollama baseUrl in chatflow configuration for Follow-up
 prompts (#4169)

Add abn Ollama baseUrl in chatflow configuration for Follow-up prompts
---
 packages/components/src/Interface.ts                      | 1 +
 packages/components/src/followUpPrompts.ts                | 8 ++++++--
 packages/ui/src/ui-component/extended/FollowUpPrompts.jsx | 8 ++++++++
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/packages/components/src/Interface.ts b/packages/components/src/Interface.ts
index 68979ca66..194e83582 100644
--- a/packages/components/src/Interface.ts
+++ b/packages/components/src/Interface.ts
@@ -436,6 +436,7 @@ export type FollowUpPromptProviderConfig = {
     [key in FollowUpPromptProvider]: {
         credentialId: string
         modelName: string
+        baseUrl: string
         prompt: string
         temperature: string
     }
diff --git a/packages/components/src/followUpPrompts.ts b/packages/components/src/followUpPrompts.ts
index b888a87d9..ecfcfe825 100644
--- a/packages/components/src/followUpPrompts.ts
+++ b/packages/components/src/followUpPrompts.ts
@@ -8,7 +8,7 @@ import { z } from 'zod'
 import { PromptTemplate } from '@langchain/core/prompts'
 import { StructuredOutputParser } from '@langchain/core/output_parsers'
 import { ChatGroq } from '@langchain/groq'
-import ollama from 'ollama'
+import { Ollama } from 'ollama'
 
 const FollowUpPromptType = z
     .object({
@@ -122,7 +122,11 @@ export const generateFollowUpPrompts = async (
                 return structuredResponse
             }
             case FollowUpPromptProvider.OLLAMA: {
-                const response = await ollama.chat({
+                const ollamaClient = new Ollama({
+                    host: providerConfig.baseUrl || 'http://127.0.0.1:11434'
+                })
+
+                const response = await ollamaClient.chat({
                     model: providerConfig.modelName,
                     messages: [
                         {
diff --git a/packages/ui/src/ui-component/extended/FollowUpPrompts.jsx b/packages/ui/src/ui-component/extended/FollowUpPrompts.jsx
index 44da9bc61..c27a917a5 100644
--- a/packages/ui/src/ui-component/extended/FollowUpPrompts.jsx
+++ b/packages/ui/src/ui-component/extended/FollowUpPrompts.jsx
@@ -269,6 +269,14 @@ const followUpPromptsOptions = {
         name: FollowUpPromptProviders.OLLAMA,
         icon: ollamaIcon,
         inputs: [
+            {
+                label: 'Base URL',
+                name: 'baseUrl',
+                type: 'string',
+                placeholder: 'http://127.0.0.1:11434',
+                description: 'Base URL of your Ollama instance',
+                default: 'http://127.0.0.1:11434'
+            },
             {
                 label: 'Model Name',
                 name: 'modelName',