From 85756e77c8e585fe36d8e77c10e74986f9fced35 Mon Sep 17 00:00:00 2001
From: Hemang Kandwal <hemangkandwal@gmail.com>
Date: Fri, 23 Aug 2024 13:01:10 +0530
Subject: [PATCH] feat(core): fix ollama integration inside docker

---
 core/src/common/utils.go  |  6 ++++++
 core/src/env/env.go       |  2 ++
 core/src/llm/llm.go       | 27 +++++++++++++++++++++++----
 core/src/router/router.go |  3 ++-
 core/src/router/task.go   |  6 ------
 5 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/core/src/common/utils.go b/core/src/common/utils.go
index 8867622..932249a 100644
--- a/core/src/common/utils.go
+++ b/core/src/common/utils.go
@@ -2,6 +2,7 @@ package common
 
 import (
 	"fmt"
+	"os"
 	"regexp"
 	"strings"
 
@@ -75,3 +76,8 @@ func JoinWithQuotes(arr []string) string {
 
 	return strings.Join(quotedStrings, ", ")
 }
+
+func IsRunningInsideDocker() bool {
+	_, err := os.Stat("/.dockerenv")
+	return !os.IsNotExist(err)
+}
diff --git a/core/src/env/env.go b/core/src/env/env.go
index c775b9b..b2b8715 100644
--- a/core/src/env/env.go
+++ b/core/src/env/env.go
@@ -10,6 +10,8 @@ import (
 )
 
 var IsDevelopment = os.Getenv("ENVIRONMENT") == "dev"
+var OllamaHost = os.Getenv("WHODB_OLLAMA_HOST")
+var OllamaPort = os.Getenv("WHODB_OLLAMA_PORT")
 
 type DatabaseCredentials struct {
 	Hostname string            `json:"host"`
diff --git a/core/src/llm/llm.go b/core/src/llm/llm.go
index cf8cfb7..37081aa 100644
--- a/core/src/llm/llm.go
+++ b/core/src/llm/llm.go
@@ -9,6 +9,9 @@ import (
 	"io"
 	"net/http"
 	"strings"
+
+	"github.com/clidey/whodb/core/src/common"
+	"github.com/clidey/whodb/core/src/env"
 )
 
 type completionRequest struct {
@@ -23,8 +26,6 @@ type completionResponse struct {
 	Done      bool   `json:"done"`
 }
 
-const ollamaLocalEndpoint = "http://localhost:11434/api"
-
 type LLMType string
 
 const (
@@ -54,7 +55,7 @@ func (c *LLMClient) Complete(prompt string, model LLMModel, receiverChan *chan s
 	var url string
 	switch c.Type {
 	case Ollama_LLMType:
-		url = fmt.Sprintf("%v/generate", ollamaLocalEndpoint)
+		url = fmt.Sprintf("%v/generate", getOllamaEndpoint())
 	}
 
 	resp, err := http.Post(url, "application/json", bytes.NewBuffer(requestBody))
@@ -96,7 +97,7 @@ func (c *LLMClient) GetSupportedModels() ([]string, error) {
 	var url string
 	switch c.Type {
 	case Ollama_LLMType:
-		url = fmt.Sprintf("%v/tags", ollamaLocalEndpoint)
+		url = fmt.Sprintf("%v/tags", getOllamaEndpoint())
 	}
 
 	resp, err := http.Get(url)
@@ -143,3 +144,21 @@ func Instance(llmType LLMType) *LLMClient {
 	llmInstance[llmType] = instance
 	return instance
 }
+
+func getOllamaEndpoint() string {
+	host := "localhost"
+	port := "11434"
+
+	if common.IsRunningInsideDocker() {
+		host = "host.docker.internal"
+	}
+
+	if env.OllamaHost != "" {
+		host = env.OllamaHost
+	}
+	if env.OllamaPort != "" {
+		port = env.OllamaPort
+	}
+
+	return fmt.Sprintf("http://%v:%v/api", host, port)
+}
diff --git a/core/src/router/router.go b/core/src/router/router.go
index 274d0f4..b604ab7 100644
--- a/core/src/router/router.go
+++ b/core/src/router/router.go
@@ -11,6 +11,7 @@ import (
 	"github.com/99designs/gqlgen/graphql/handler/transport"
 	"github.com/clidey/whodb/core/graph"
 	"github.com/clidey/whodb/core/src/auth"
+	"github.com/clidey/whodb/core/src/common"
 	"github.com/clidey/whodb/core/src/log"
 	"github.com/go-chi/chi/v5"
 	"github.com/go-chi/chi/v5/middleware"
@@ -69,7 +70,7 @@ func InitializeRouter(staticFiles embed.FS) {
 	log.Logger.Infof("http://0.0.0.0:%s", port)
 	log.Logger.Info("Explore and enjoy working with your databases!")
 
-	if !isDocker() {
+	if !common.IsRunningInsideDocker() {
 		openBrowser(fmt.Sprintf("http://localhost:%v", port))
 	}
 
diff --git a/core/src/router/task.go b/core/src/router/task.go
index cc585dc..45bbb32 100644
--- a/core/src/router/task.go
+++ b/core/src/router/task.go
@@ -1,7 +1,6 @@
 package router
 
 import (
-	"os"
 	"os/exec"
 	"runtime"
 
@@ -24,8 +23,3 @@ func openBrowser(url string) {
 		log.Logger.Warnf("Failed to open browser: %v\n", err)
 	}
 }
-
-func isDocker() bool {
-	_, err := os.Stat("/.dockerenv")
-	return !os.IsNotExist(err)
-}