Skip to content

Commit

Permalink
feat: add Ollama backend
Browse files Browse the repository at this point in the history
Signed-off-by: Kay Yan <[email protected]>
  • Loading branch information
yankay committed May 17, 2024
1 parent cc890df commit 56e7949
Show file tree
Hide file tree
Showing 8 changed files with 115 additions and 7 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/build_container.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ on:
- "**.md"

env:
GO_VERSION: "~1.21"
GO_VERSION: "~1.22"
IMAGE_NAME: "k8sgpt"
defaults:
run:
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/release.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ jobs:
- name: Set up Go
uses: actions/setup-go@0c52d547c9bc32b1aa3301fd7a9cb496313a4491 # v5
with:
go-version: '1.21'
go-version: '1.22'
- name: Download Syft
uses: anchore/sbom-action/download-syft@7ccf588e3cf3cc2611714c2eeae48550fbc17552 # v0.15.11
- name: Run GoReleaser
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/test.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ on:
- main

env:
GO_VERSION: "~1.21"
GO_VERSION: "~1.22"

jobs:
build:
Expand Down
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -308,6 +308,7 @@ Active:
Unused:
> openai
> localai
> ollama
> azureopenai
> cohere
> amazonbedrock
Expand Down
3 changes: 2 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
module github.com/k8sgpt-ai/k8sgpt

go 1.21
go 1.22

require (
github.com/aquasecurity/trivy-operator v0.17.1
github.com/fatih/color v1.16.0
github.com/kedacore/keda/v2 v2.11.2
github.com/magiconair/properties v1.8.7
github.com/mittwald/go-helm-client v0.12.5
github.com/ollama/ollama v0.1.33
github.com/sashabaranov/go-openai v1.23.0
github.com/schollz/progressbar/v3 v3.14.2
github.com/spf13/cobra v1.8.0
Expand Down
6 changes: 4 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -1812,8 +1812,8 @@ github.com/klauspost/compress v1.17.4/go.mod h1:/dCuZOvVtNoHsyb+cuJD3itjs3NbnF6K
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
github.com/klauspost/cpuid/v2 v2.0.10/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c=
github.com/klauspost/cpuid/v2 v2.0.12/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c=
github.com/klauspost/cpuid/v2 v2.2.3 h1:sxCkb+qR91z4vsqw4vGGZlDgPz3G7gjaLyK3V8y70BU=
github.com/klauspost/cpuid/v2 v2.2.3/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZXnvk=
github.com/klauspost/cpuid/v2 v2.2.4/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
github.com/kolo/xmlrpc v0.0.0-20220921171641-a4b6fa1dd06b h1:udzkj9S/zlT5X367kqJis0QP7YMxobob6zhzq6Yre00=
github.com/kolo/xmlrpc v0.0.0-20220921171641-a4b6fa1dd06b/go.mod h1:pcaDhQK0/NJZEvtCO0qQPPropqV0sJOJ6YW7X+9kRwM=
github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
Expand Down Expand Up @@ -1928,6 +1928,8 @@ github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f h1:KUppIJq7/+
github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec=
github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
github.com/ollama/ollama v0.1.33 h1:Kb2bK4YxLJrRPNOaznwIYUoRYnT7AFOZioJKjreCM6E=
github.com/ollama/ollama v0.1.33/go.mod h1:u9Bo9/pxhGe2YiL1I/ePNRTH0Ik5U3B2C/i2EYp1lZk=
github.com/onsi/ginkgo/v2 v2.13.1 h1:LNGfMbR2OVGBfXjvRZIZ2YCTQdGKtPLvuI1rMCCj3OU=
github.com/onsi/ginkgo/v2 v2.13.1/go.mod h1:XStQ8QcGwLyF4HdfcZB8SFOS/MWCgDuXMSBe6zrvLgM=
github.com/onsi/gomega v1.30.0 h1:hvMK7xYz4D3HapigLTeGdId/NcfQx1VHMJc60ew99+8=
Expand Down
4 changes: 3 additions & 1 deletion pkg/ai/iai.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ var (
&OpenAIClient{},
&AzureAIClient{},
&LocalAIClient{},
&OllamaClient{},
&NoOpAIClient{},
&CohereClient{},
&AmazonBedRockClient{},
Expand All @@ -34,6 +35,7 @@ var (
Backends = []string{
openAIClientName,
localAIClientName,
ollamaClientName,
azureAIClientName,
cohereAIClientName,
amazonbedrockAIClientName,
Expand Down Expand Up @@ -164,7 +166,7 @@ func (p *AIProvider) GetCompartmentId() string {
return p.CompartmentId
}

var passwordlessProviders = []string{"localai", "amazonsagemaker", "amazonbedrock", "googlevertexai", "oci"}
var passwordlessProviders = []string{"localai", "ollama", "amazonsagemaker", "amazonbedrock", "googlevertexai", "oci"}

func NeedPassword(backend string) bool {
for _, b := range passwordlessProviders {
Expand Down
102 changes: 102 additions & 0 deletions pkg/ai/ollama.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
/*
Copyright 2023 The K8sGPT Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package ai

import (
"context"
"errors"
"net/http"
"net/url"

ollama "github.com/ollama/ollama/api"
)

const ollamaClientName = "ollama"

type OllamaClient struct {
nopCloser

client *ollama.Client
model string
temperature float32
topP float32
}

const (
defaultBaseURL = "http://localhost:11434"
defaultModel = "llama3"
)

func (c *OllamaClient) Configure(config IAIConfig) error {
baseURL := config.GetBaseURL()
if baseURL == "" {
baseURL = defaultBaseURL
}
baseClientURL, err := url.Parse(baseURL)
if err != nil {
return err
}

proxyEndpoint := config.GetProxyEndpoint()
httpClient := http.DefaultClient
if proxyEndpoint != "" {
proxyUrl, err := url.Parse(proxyEndpoint)
if err != nil {
return err
}
transport := &http.Transport{
Proxy: http.ProxyURL(proxyUrl),
}

httpClient = &http.Client{
Transport: transport,
}
}

c.client = ollama.NewClient(baseClientURL, httpClient)
if c.client == nil {
return errors.New("error creating Ollama client")
}
c.model = config.GetModel()
if c.model == "" {
c.model = defaultModel
}
c.temperature = config.GetTemperature()
c.topP = config.GetTopP()
return nil
}
func (c *OllamaClient) GetCompletion(ctx context.Context, prompt string) (string, error) {
req := &ollama.GenerateRequest{
Model: c.model,
Prompt: prompt,
Stream: new(bool),
Options: map[string]interface{}{
"temperature": c.temperature,
"top_p": c.topP,
},
}
completion := ""
respFunc := func(resp ollama.GenerateResponse) error {
completion = resp.Response
return nil
}
err := c.client.Generate(ctx, req, respFunc)
if err != nil {
return "", err
}
return completion, nil
}
func (a *OllamaClient) GetName() string {
return ollamaClientName
}

0 comments on commit 56e7949

Please sign in to comment.