feat: tool support Ollama with official client

go.mod

@@ -1,6 +1,6 @@
 module github.com/flanksource/incident-commander
 
-go 1.23.0
+go 1.24.0
 
 toolchain go1.24.2
 
@@ -59,6 +59,7 @@ require (
 	github.com/henvic/httpretty v0.1.4
 	github.com/jackc/pgerrcode v0.0.0-20240316143900-6e2875d9b438
 	github.com/jenkins-x/go-scm v1.14.17
+	github.com/ollama/ollama v0.6.6
 	github.com/redis/go-redis/v9 v9.5.1
 	github.com/samber/lo v1.49.1
 	github.com/samber/oops v1.17.0
@@ -76,7 +77,7 @@ require (
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.27.0
 	go.opentelemetry.io/otel/sdk v1.35.0
 	go.opentelemetry.io/otel/trace v1.35.0
-	golang.org/x/exp v0.0.0-20241108190413-2d47ceb2692f
+	golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa
 	google.golang.org/genai v1.1.0
 	gorm.io/hints v1.1.2
 	k8s.io/api v0.31.1
@@ -135,7 +136,7 @@ require (
 	github.com/cyphar/filepath-securejoin v0.2.4 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/distribution/reference v0.5.0 // indirect
-	github.com/dlclark/regexp2 v1.11.0 // indirect
+	github.com/dlclark/regexp2 v1.11.4 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/eko/gocache/lib/v4 v4.1.6 // indirect
 	github.com/eko/gocache/store/go_cache/v4 v4.2.2 // indirect
@@ -263,10 +264,10 @@ require (
 	go.opentelemetry.io/otel/metric v1.35.0 // indirect
 	go.opentelemetry.io/proto/otlp v1.3.1 // indirect
 	gocloud.dev v0.40.0 // indirect
-	golang.org/x/mod v0.22.0 // indirect
+	golang.org/x/mod v0.23.0 // indirect
 	golang.org/x/sync v0.12.0 // indirect
 	golang.org/x/term v0.30.0 // indirect
-	golang.org/x/tools v0.28.0 // indirect
+	golang.org/x/tools v0.30.0 // indirect
 	golang.org/x/xerrors v0.0.0-20240716161551-93cc26a95ae9 // indirect
 	gomodules.xyz/jsonpatch/v2 v2.4.0 // indirect
 	gonum.org/v1/gonum v0.15.1 // indirect
@@ -353,7 +354,7 @@ require (
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
 	go.opencensus.io v0.24.0 // indirect
 	golang.org/x/crypto v0.36.0
-	golang.org/x/net v0.33.0 // indirect
+	golang.org/x/net v0.35.0 // indirect
 	golang.org/x/oauth2 v0.28.0 // indirect
 	golang.org/x/sys v0.31.0 // indirect
 	golang.org/x/text v0.23.0

go.sum

@@ -247,8 +247,8 @@ github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/r
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/distribution/reference v0.5.0 h1:/FUIFXtfc/x2gpa5/VGfiGLuOIdYa1t65IKK2OFGvA0=
 github.com/distribution/reference v0.5.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E=
-github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI=
-github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
+github.com/dlclark/regexp2 v1.11.4 h1:rPYF9/LECdNymJufQKmri9gV604RvvABwgOA8un7yAo=
+github.com/dlclark/regexp2 v1.11.4/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/dnaeon/go-vcr v1.1.0/go.mod h1:M7tiix8f0r6mKKJ3Yq/kqU1OYf3MnfmBWVbPx/yU9ko=
 github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
@@ -737,6 +737,8 @@ github.com/ohler55/ojg v1.25.0 h1:sDwc4u4zex65Uz5Nm7O1QwDKTT+YRcpeZQTy1pffRkw=
 github.com/ohler55/ojg v1.25.0/go.mod h1:gQhDVpQLqrmnd2eqGAvJtn+NfKoYJbe/A4Sj3/Vro4o=
 github.com/oklog/ulid/v2 v2.1.0 h1:+9lhoxAP56we25tyYETBBY1YLA2SaoLvUFgrP2miPJU=
 github.com/oklog/ulid/v2 v2.1.0/go.mod h1:rcEKHmBBKfef9DhnvX7y1HZBYxjXb0cP5ExxNsTT1QQ=
+github.com/ollama/ollama v0.6.6 h1:rnCQTSTiRD3Dsvd35dh2j2YB9DlQMFQR/y3XOhWZOmI=
+github.com/ollama/ollama v0.6.6/go.mod h1:pGgtoNyc9DdM6oZI6yMfI6jTk2Eh4c36c2GpfQCH7PY=
 github.com/onsi/ginkgo/v2 v2.22.2 h1:/3X8Panh8/WwhU/3Ssa6rCKqPLuAkVY2I0RoyDLySlU=
 github.com/onsi/ginkgo/v2 v2.22.2/go.mod h1:oeMosUL+8LtarXBHu/c0bx2D/K9zyQ6uX3cTyztHwsk=
 github.com/onsi/gomega v1.36.2 h1:koNYke6TVk6ZmnyHrCXba/T/MoLBXFjeC1PtvYgw0A8=
@@ -1052,8 +1054,8 @@ golang.org/x/exp v0.0.0-20191227195350-da58074b4299/go.mod h1:2RIsYlXP63K8oxa1u0
 golang.org/x/exp v0.0.0-20200119233911-0405dc783f0a/go.mod h1:2RIsYlXP63K8oxa1u096TMicItID8zy7Y6sNkU49FU4=
 golang.org/x/exp v0.0.0-20200207192155-f17229e696bd/go.mod h1:J/WKrq2StrnmMY6+EHIKF9dgMWnmCNThgcyBT1FY9mM=
 golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6/go.mod h1:3jZMyOhIsHpP37uCMkUooju7aAi5cS1Q23tOzKc+0MU=
-golang.org/x/exp v0.0.0-20241108190413-2d47ceb2692f h1:XdNn9LlyWAhLVp6P/i8QYBW+hlyhrhei9uErw2B5GJo=
-golang.org/x/exp v0.0.0-20241108190413-2d47ceb2692f/go.mod h1:D5SMRVC3C2/4+F/DB1wZsLRnSNimn2Sp/NPsCrsv8ak=
+golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa h1:t2QcU6V556bFjYgu4L6C+6VrCPyJZ+eyRsABUPs1mz4=
+golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa/go.mod h1:BHOTPb3L19zxehTsLoJXVaTktb06DFgmdW6Wb9s8jqk=
 golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js=
 golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
@@ -1079,8 +1081,8 @@ golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/mod v0.14.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
-golang.org/x/mod v0.22.0 h1:D4nJWe9zXqHOmWqj4VMOJhvzj7bEZg4wEYa759z1pH4=
-golang.org/x/mod v0.22.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY=
+golang.org/x/mod v0.23.0 h1:Zb7khfcRGKk+kqfxFaP5tZqCnDZMjC5VtUBs87Hr6QM=
+golang.org/x/mod v0.23.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20181114220301-adae6a3d119a/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -1129,8 +1131,8 @@ golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
 golang.org/x/net v0.14.0/go.mod h1:PpSgVXXLK0OxS0F31C1/tv6XNguvCrnXIDrFMspZIUI=
 golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
 golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY=
-golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
-golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
+golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8=
+golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
 golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
@@ -1315,8 +1317,8 @@ golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc
 golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
 golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
 golang.org/x/tools v0.17.0/go.mod h1:xsh6VxdV005rRVaS6SSAf9oiAqljS7UZUacMZ8Bnsps=
-golang.org/x/tools v0.28.0 h1:WuB6qZ4RPCQo5aP3WdKZS7i595EdWqWR8vqJTlwTVK8=
-golang.org/x/tools v0.28.0/go.mod h1:dcIOrVd3mfQKTgrDVQHqCPMWy6lnhfhtX3hLXYVLfRw=
+golang.org/x/tools v0.30.0 h1:BgcpHewrV5AUp2G9MebG4XPFI1E2W41zU1SaqVA9vJY=
+golang.org/x/tools v0.30.0/go.mod h1:c347cR/OJfw5TI+GfX7RUPNMdDRRbjvYTS0jPyvsVtY=
 golang.org/x/xerrors v0.0.0-20190410155217-1f06c39b4373/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20190513163551-3ee3066db522/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

llm/cost.go

@@ -431,6 +431,10 @@ func populateOpenAIModels() {
 // It considers tiered pricing if defined for the model.
 // cacheReadTokens and cacheWriteTokens are optional.
 func CalculateCost(provider api.LLMBackend, modelID string, genInfo GenerationInfo) (float64, error) {
+	if provider == api.LLMBackendOllama {
+		return 0, nil
+	}
+
 	modelInfo, err := GetModelInfo(provider, modelID)
 	if err != nil {
 		return 0, err

llm/llm.go

@@ -8,7 +8,6 @@ import (
 	"github.com/samber/lo"
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/llms/anthropic"
-	"github.com/tmc/langchaingo/llms/ollama"
 	"github.com/tmc/langchaingo/llms/openai"
 	"google.golang.org/genai"
 
@@ -79,6 +78,10 @@ func PromptWithHistory(ctx dutyctx.Context, config Config, history []llms.Messag
 		// Do nothing
 		// NOTE: Handled by the wrapper
 
+	case api.LLMBackendOllama:
+		// Do nothing
+		// NOTE: Handled by the wrapper
+
 	default:
 		const forceToolUsePrompt = `You MUST use the %s tool to extract the diagnosis information. 
 	Do not provide any other response format. 
@@ -161,6 +164,14 @@ func calculateGenerationInfo(llmBackend api.LLMBackend, model string, resp *llms
 				if outputTokens, ok := choice.GenerationInfo["OutputTokens"]; ok {
 					genInfo.OutputTokens += int(outputTokens.(int32))
 				}
+
+			case api.LLMBackendOllama:
+				if inputTokens, ok := choice.GenerationInfo["PromptTokens"]; ok {
+					genInfo.InputTokens += inputTokens.(int)
+				}
+				if outputTokens, ok := choice.GenerationInfo["CompletionTokens"]; ok {
+					genInfo.OutputTokens += outputTokens.(int)
+				}
 			}
 
 			cost, err := CalculateCost(llmBackend, model, genInfo)
@@ -226,19 +237,12 @@ func getLLMModel(ctx dutyctx.Context, config Config) (llms.Model, error) {
 		return openaiLLM, nil
 
 	case api.LLMBackendOllama:
-		var opts []ollama.Option
-		if config.APIURL != "" {
-			opts = append(opts, ollama.WithServerURL(config.APIURL))
-		}
-		if config.Model != "" {
-			opts = append(opts, ollama.WithModel(config.Model))
-		}
-
-		openaiLLM, err := ollama.New(opts...)
+		client, err := NewOllamaClient(config.Model, config.APIURL)
 		if err != nil {
-			return nil, fmt.Errorf("failed to created ollama llm: %w", err)
+			return nil, fmt.Errorf("failed to create ollama client: %w", err)
 		}
-		return openaiLLM, nil
+
+		return client, nil
 
 	case api.LLMBackendAnthropic:
 		var opts []anthropic.Option

llm/ollama.go

@@ -0,0 +1,124 @@
+package llm
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strings"
+
+	"github.com/ollama/ollama/api"
+	"github.com/tmc/langchaingo/llms"
+)
+
+type OllamaClient struct {
+	*api.Client
+	model string
+}
+
+func NewOllamaClient(model, baseURL string) (*OllamaClient, error) {
+	parsedURL, err := url.Parse(baseURL)
+	if err != nil {
+		return nil, err
+	}
+
+	c := api.NewClient(parsedURL, &http.Client{})
+	return &OllamaClient{Client: c, model: model}, nil
+}
+
+type Params struct {
+	Type       string   `json:"type"`
+	Defs       any      `json:"$defs,omitempty"`
+	Items      any      `json:"items,omitempty"`
+	Required   []string `json:"required"`
+	Properties map[string]struct {
+		Type        api.PropertyType `json:"type"`
+		Items       any              `json:"items,omitempty"`
+		Description string           `json:"description"`
+		Enum        []any            `json:"enum,omitempty"`
+	} `json:"properties"`
+}
+
+// GenerateContent asks the model to generate content from a sequence of
+// messages. It's the most general interface for multi-modal LLMs that support
+// chat-like interactions.
+func (c *OllamaClient) GenerateContent(ctx context.Context, messages []llms.MessageContent, options ...llms.CallOption) (*llms.ContentResponse, error) {
+	// https://ollama.com/blog/tool-support
+	tools := []api.Tool{
+		{
+			Type: "function",
+			Function: api.ToolFunction{
+				Name:        "report_diagnosis",
+				Description: "Report the diagnosis information for the affected resource",
+				Parameters: Params{
+					Type: "object",
+					Properties: map[string]struct {
+						Type        api.PropertyType `json:"type"`
+						Items       any              `json:"items,omitempty"`
+						Description string           `json:"description"`
+						Enum        []any            `json:"enum,omitempty"`
+					}{
+						"headline": {
+							Type:        api.PropertyType([]string{"string"}),
+							Description: "Headline that clearly mentions the affected resource & the issue. Feel free to add emojis. Keep it short and concise.",
+						},
+						"summary": {
+							Type:        api.PropertyType([]string{"string"}),
+							Description: "Summary of the issue in markdown. Use bullet points if needed.",
+						},
+						"recommended_fix": {
+							Type:        api.PropertyType([]string{"string"}),
+							Description: "Short and concise recommended fix for the issue in markdown. Use bullet points if needed.",
+						},
+					},
+					Required: []string{"headline", "summary", "recommended_fix"},
+				},
+			},
+		},
+	}
+
+	chatRequest := api.ChatRequest{
+		Model: c.model,
+		Tools: tools,
+	}
+
+	for _, message := range messages {
+		for _, part := range message.Parts {
+			role := string(message.Role)
+			if role == string(llms.ChatMessageTypeHuman) {
+				role = "user"
+			} else if role == string(llms.ChatMessageTypeAI) {
+				role = "assistant"
+			}
+
+			chatRequest.Messages = append(chatRequest.Messages, api.Message{
+				Role:    role,
+				Content: fmt.Sprintf("%s", part),
+			})
+		}
+	}
+
+	var llmResponse strings.Builder
+	if err := c.Client.Chat(ctx, &chatRequest, func(resp api.ChatResponse) error {
+		_, err := llmResponse.WriteString(resp.Message.Content)
+		return err
+	}); err != nil {
+		return nil, err
+	}
+
+	contentResponse := &llms.ContentResponse{
+		Choices: []*llms.ContentChoice{{Content: llmResponse.String()}},
+	}
+	return contentResponse, nil
+}
+
+// Call is a simplified interface for a text-only Model, generating a single
+// string response from a single string prompt.
+//
+// Deprecated: this method is retained for backwards compatibility. Use the
+// more general [GenerateContent] instead. You can also use
+// the [GenerateFromSinglePrompt] function which provides a similar capability
+// to Call and is built on top of the new interface.
+func (c *OllamaClient) Call(ctx context.Context, prompt string, options ...llms.CallOption) (string, error) {
+	return "", nil
+}

playbook/actions/ai.go

@@ -183,15 +183,15 @@ func (t *aiAction) Run(ctx context.Context, spec v1.AIAction) (*AIActionResult,
 	llmConf := llm.Config{AIActionClient: spec.AIActionClient, ResponseFormat: llm.ResponseFormatDiagnosis}
 	response, conversation, genInfo, err := llm.Prompt(ctx, llmConf, spec.SystemPrompt, prompt...)
 	if err != nil {
-		return nil, ctx.Oops().Wrapf(err, "failed to generate response")
+		return &result, ctx.Oops().Wrapf(err, "failed to generate response")
 	}
 	result.Prompt.WriteString(strings.Join(prompt, "\n"))
 	result.JSON = response
 	result.GenerationInfo = append(result.GenerationInfo, genInfo...)
 
 	var diagnosisReport llm.DiagnosisReport
 	if err := json.Unmarshal([]byte(response), &diagnosisReport); err != nil {
-		return nil, ctx.Oops().With("response", response).Wrapf(err, "failed to unmarshal diagnosis report")
+		return &result, ctx.Oops().With("response", response).Wrapf(err, "failed to unmarshal diagnosis report")
 	}
 
 	for _, format := range lo.Uniq(spec.Formats) {
@@ -243,14 +243,14 @@ func (t *aiAction) Run(ctx context.Context, spec v1.AIAction) (*AIActionResult,
 			llmConf.ResponseFormat = llm.ResponseFormatPlaybookRecommendations
 			response, _, genInfo, err := llm.PromptWithHistory(ctx, llmConf, conversation, prompt.String())
 			if err != nil {
-				return nil, fmt.Errorf("failed to generate playbook recommendation: %w", err)
+				return &result, fmt.Errorf("failed to generate playbook recommendation: %w", err)
 			}
 			result.Prompt.WriteString(prompt.String())
 			result.GenerationInfo = append(result.GenerationInfo, genInfo...)
 
 			var recommendations llm.PlaybookRecommendations
 			if err := json.Unmarshal([]byte(response), &recommendations); err != nil {
-				return nil, ctx.Oops().With("response", response).Wrapf(err, "failed to unmarshal playbook recommendations")
+				return &result, ctx.Oops().With("response", response).Wrapf(err, "failed to unmarshal playbook recommendations")
 			}
 
 			groupedResources, err := getGroupedResources(ctx, t.RunID)