Skip to content

Add Ollama LLM Client Support #62

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SumonRayy wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Add Ollama LLM Client Support #62

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
fc643d6
feat: add Ollama support and configuration
SumonRayy May 27, 2025
a53e3aa
feat: enhanced ollama MySQL prompt
SumonRayy May 27, 2025
8730944
feat: add support for YugabyteDB and MySQL in Ollama prompts
SumonRayy May 27, 2025
ed880c0
feat: extend Ollama support for Clickhouse and MongoDB
SumonRayy May 27, 2025
2fb0cdd
fix: update README to reflect Ollama support changes
SumonRayy May 27, 2025
6073839
feat: add Ollama support for PostgreSQL, MySQL, Clickhouse, MongoDB, …
SumonRayy May 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,10 +35,10 @@
## Supported LLM Clients
- OpenAI (Any chat completion model)
- Google Gemini (Any chat completion model)
- Ollama (Any chat completion model)

## Planned to be supported LLM Clients
- Anthropic (Claude 3.5 Sonnet)
- Ollama (Any chat completion model)

## Tech Stack

Expand Down Expand Up @@ -84,7 +84,7 @@ Read ([SETUP](https://github.com/bhaskarblur/neobase-ai-dba/blob/main/SETUP.md))

## Contributing

We welcome contributions! Here’s how you can help:
We welcome contributions! Here's how you can help:

1. Fork the repository.
2. Create a new branch (`git checkout -b feature/your-feature`).
Expand Down
8 changes: 7 additions & 1 deletion backend/.env.example
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ NEOBASE_REDIS_PORT=6379
NEOBASE_REDIS_USERNAME=neobase
NEOBASE_REDIS_PASSWORD=default

DEFAULT_LLM_CLIENT=openai # openai, gemini
DEFAULT_LLM_CLIENT=ollama # openai, gemini, ollama
# OpenAI API Key
OPENAI_API_KEY=<openai-api-key> # Your OpenAI Api Key
OPENAI_MODEL=gpt-4o # OpenAI Model
Expand All @@ -33,6 +33,12 @@ GEMINI_MODEL=gemini-2.0-flash # Gemini Model
GEMINI_MAX_COMPLETION_TOKENS=30000 # Example: 30000
GEMINI_TEMPERATURE=1 # 0-2

# Ollama API Keys
OLLAMA_API_KEY=http://localhost:11434 # your ollama url:port
OLLAMA_MODEL=llama3:latest
OLLAMA_MAX_COMPLETION_TOKENS=30000
OLLAMA_TEMPERATURE=1

# Example DB for Development Environment
EXAMPLE_DB_TYPE=
EXAMPLE_DB_HOST=
Expand Down
12 changes: 12 additions & 0 deletions backend/config/env_values.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,12 @@ type Environment struct {
GeminiModel string
GeminiMaxCompletionTokens int
GeminiTemperature float64

// Ollama configs
OllamaAPIKey string
OllamaModel string
OllamaMaxCompletionTokens int
OllamaTemperature float64
}

var Env Environment
Expand Down Expand Up @@ -113,6 +119,12 @@ func LoadEnv() error {
Env.GeminiMaxCompletionTokens = getIntEnvWithDefault("GEMINI_MAX_COMPLETION_TOKENS", constants.GeminiMaxCompletionTokens)
Env.GeminiTemperature = getFloatEnvWithDefault("GEMINI_TEMPERATURE", constants.GeminiTemperature)

// Ollama configs
Env.OllamaAPIKey = getRequiredEnv("OLLAMA_API_KEY", "")
Env.OllamaModel = getEnvWithDefault("OLLAMA_MODEL", constants.OllamaModel)
Env.OllamaMaxCompletionTokens = getIntEnvWithDefault("OLLAMA_MAX_COMPLETION_TOKENS", constants.OllamaMaxCompletionTokens)
Env.OllamaTemperature = getFloatEnvWithDefault("OLLAMA_TEMPERATURE", constants.OllamaTemperature)

return validateConfig()
}

Expand Down
5 changes: 1 addition & 4 deletions backend/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,11 +42,9 @@ require (
github.com/go-playground/universal-translator v0.18.1 // indirect
github.com/go-playground/validator/v10 v10.23.0 // indirect
github.com/goccy/go-json v0.10.4 // indirect
github.com/google/go-cmp v0.7.0 // indirect
github.com/google/s2a-go v0.1.9 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
github.com/googleapis/gax-go/v2 v2.14.1 // indirect
github.com/gorilla/websocket v1.5.3 // indirect
github.com/hashicorp/go-version v1.7.0 // indirect
github.com/jackc/pgpassfile v1.0.0 // indirect
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
Expand Down Expand Up @@ -87,7 +85,7 @@ require (
)

require (
github.com/ClickHouse/clickhouse-go/v2 v2.32.2
github.com/ClickHouse/clickhouse-go/v2 v2.32.2 // indirect
github.com/gin-contrib/cors v1.7.3
github.com/go-sql-driver/mysql v1.9.0
github.com/golang/snappy v0.0.4 // indirect
Expand All @@ -103,7 +101,6 @@ require (
golang.org/x/sync v0.11.0 // indirect
golang.org/x/text v0.22.0 // indirect
google.golang.org/api v0.223.0
google.golang.org/genai v0.4.0
gorm.io/driver/clickhouse v0.6.1
gorm.io/driver/mysql v1.5.7
)
10 changes: 1 addition & 9 deletions backend/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
github.com/ClickHouse/ch-go v0.65.1 h1:SLuxmLl5Mjj44/XbINsK2HFvzqup0s6rwKLFH347ZhU=
github.com/ClickHouse/ch-go v0.65.1/go.mod h1:bsodgURwmrkvkBe5jw1qnGDgyITsYErfONKAHn05nv4=
github.com/ClickHouse/clickhouse-go v1.5.4 h1:cKjXeYLNWVJIx2J1K6H2CqyRmfwVJVY1OV1coaaFcI0=
github.com/ClickHouse/clickhouse-go/v2 v2.32.2 h1:Y8fAXt0CpLhqNXMLlSddg+cMfAr7zHBWqXLpih6ozCY=
github.com/ClickHouse/clickhouse-go/v2 v2.32.2/go.mod h1:/vE8N/+9pozLkIiTMWbNUGviccDv/czEGS1KACvpXIk=
github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA=
Expand Down Expand Up @@ -95,8 +94,6 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.4 h1:XYIDZApgAnrN1c855gT
github.com/googleapis/enterprise-certificate-proxy v0.3.4/go.mod h1:YKe7cfqYXjKGpGvmSg28/fFvhNzinZQm8DGnaburhGA=
github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrkurSS/Q=
github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA=
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/hashicorp/go-version v1.7.0 h1:5tqGy27NaOTB8yJKUZELlFAS/LTKJkrmONwQKeRZfjY=
github.com/hashicorp/go-version v1.7.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA=
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
Expand All @@ -105,8 +102,6 @@ github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
github.com/jackc/pgx/v5 v5.5.5 h1:amBjrZVmksIdNjxGW/IiIMzxMKZFelXbUoPNb+8sjQw=
github.com/jackc/pgx/v5 v5.5.5/go.mod h1:ez9gk+OAat140fv9ErkZDYFWmXLfV+++K0uAOiwgm1A=
github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk=
github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
Expand All @@ -120,8 +115,6 @@ github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHm
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA=
github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
github.com/klauspost/compress v1.17.11 h1:In6xLpyWOi1+C7tXUUWv2ot1QvBjxevKAaI6IXrJmUc=
github.com/klauspost/compress v1.17.11/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0=
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
Expand Down Expand Up @@ -194,6 +187,7 @@ github.com/xdg-go/scram v1.1.2/go.mod h1:RT/sEzTbU5y00aCK8UOx6R7YryM0iF1N2MOmC3k
github.com/xdg-go/stringprep v1.0.3/go.mod h1:W3f5j4i+9rC0kuIEJL0ky1VpHXQU3ocBgklLGvcBnW8=
github.com/xdg-go/stringprep v1.0.4 h1:XLI/Ng3O1Atzq0oBs3TWm+5ZVgkq2aqdlvP9JtoZ6c8=
github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM=
github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d/go.mod h1:rHwXgn7JulP+udvsHwJoVG1YGAP6VLg4y9I5dyZdqmA=
github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 h1:ilQV1hzziu+LLM3zUTJ0trRztfwgjqKnBWNtSRkbmwM=
Expand Down Expand Up @@ -285,8 +279,6 @@ golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8T
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/api v0.223.0 h1:JUTaWEriXmEy5AhvdMgksGGPEFsYfUKaPEYXd4c3Wvc=
google.golang.org/api v0.223.0/go.mod h1:C+RS7Z+dDwds2b+zoAk5hN/eSfsiCn0UDrYof/M4d2M=
google.golang.org/genai v0.4.0 h1:nHLpFvp1i2nUGQ8CjIQ8j/6d3H79Echt1jiNLb9myDk=
google.golang.org/genai v0.4.0/go.mod h1:yPyKKBezIg2rqZziLhHQ5CD62HWr7sLDLc2PDzdrNVs=
google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576 h1:CkkIfIt50+lT6NHAVoRYEyAvQGFM7xEwXUUywFvEb3Q=
google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576/go.mod h1:1R3kvZ1dtP3+4p4d3G8uJ8rFk/fWlScl38vanWACI08=
google.golang.org/genproto/googleapis/rpc v0.0.0-20250219182151-9fdb1cabc7b2 h1:DMTIbak9GhdaSxEjvVzAeNZvyc03I61duqNbnm3SU0M=
Expand Down
32 changes: 32 additions & 0 deletions backend/internal/constants/llms.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package constants
const (
OpenAI = "openai"
Gemini = "gemini"
Ollama = "ollama"
)

func GetLLMResponseSchema(provider string, dbType string) interface{} {
Expand Down Expand Up @@ -37,7 +38,23 @@ func GetLLMResponseSchema(provider string, dbType string) interface{} {
default:
return GeminiPostgresLLMResponseSchema
}
case Ollama:
switch dbType {
case DatabaseTypePostgreSQL:
return OllamaPostgreSQLLLMResponseSchema
case DatabaseTypeYugabyteDB:
return OllamaYugabyteDBLLMResponseSchema
case DatabaseTypeMySQL:
return OllamaMySQLLLMResponseSchema
case DatabaseTypeClickhouse:
return OllamaClickhouseLLMResponseSchema
case DatabaseTypeMongoDB:
return OllamaMongoDBLLMResponseSchema
default:
return OllamaPostgreSQLLLMResponseSchema
}
}

return ""
}

Expand Down Expand Up @@ -74,6 +91,21 @@ func GetSystemPrompt(provider string, dbType string) string {
default:
return GeminiPostgreSQLPrompt // Default to PostgreSQL
}
case Ollama:
switch dbType {
case DatabaseTypePostgreSQL:
return OllamaPostgreSQLPrompt
case DatabaseTypeYugabyteDB:
return OllamaYugabyteDBPrompt
case DatabaseTypeMySQL:
return OllamaMySQLPrompt
case DatabaseTypeClickhouse:
return OllamaClickhousePrompt
case DatabaseTypeMongoDB:
return OllamaMongoDBPrompt
default:
return OllamaPostgreSQLPrompt
}
}
return ""
}
Loading