diff --git a/go/plugins/ollama/ollama.go b/go/plugins/ollama/ollama.go index ea6ef9631c..97f9fe1617 100644 --- a/go/plugins/ollama/ollama.go +++ b/go/plugins/ollama/ollama.go @@ -87,7 +87,7 @@ func (o *Ollama) DefineModel(g *genkit.Genkit, model ModelDefinition, opts *ai.M Supports: modelOpts.Supports, Versions: []string{}, } - gen := &generator{model: model, serverAddress: o.ServerAddress} + gen := &generator{model: model, serverAddress: o.ServerAddress, timeout: o.Timeout} return genkit.DefineModel(g, api.NewName(provider, model.Name), meta, gen.generate) } @@ -111,6 +111,7 @@ type ModelDefinition struct { type generator struct { model ModelDefinition serverAddress string + timeout int } type ollamaMessage struct { @@ -196,6 +197,7 @@ type ollamaModelResponse struct { // Ollama provides configuration options for the Init function. type Ollama struct { ServerAddress string // Server address of oLLama. + Timeout int // Response timeout in seconds (defaulted to 30 seconds) mu sync.Mutex // Mutex to control access. initted bool // Whether the plugin has been initialized. @@ -218,6 +220,9 @@ func (o *Ollama) Init(ctx context.Context) []api.Action { panic("ollama: need ServerAddress") } o.initted = true + if o.Timeout == 0 { + o.Timeout = 30 + } return []api.Action{} } @@ -274,7 +279,7 @@ func (g *generator) generate(ctx context.Context, input *ai.ModelRequest, cb fun payload = chatReq } - client := &http.Client{Timeout: 30 * time.Second} + client := &http.Client{Timeout: time.Duration(g.timeout) * time.Second} payloadBytes, err := json.Marshal(payload) if err != nil { return nil, err diff --git a/go/samples/ollama-tools/main.go b/go/samples/ollama-tools/main.go index 1f4f46d858..2dd4dd2cdf 100644 --- a/go/samples/ollama-tools/main.go +++ b/go/samples/ollama-tools/main.go @@ -43,6 +43,7 @@ func main() { // Initialize Genkit with the Ollama plugin ollamaPlugin := &ollama.Ollama{ ServerAddress: "http://localhost:11434", // Default Ollama server address + Timeout: 60, // Response timeout in seconds } g := genkit.Init(ctx, genkit.WithPlugins(ollamaPlugin)) @@ -80,7 +81,6 @@ func main() { ai.WithTools(weatherTool), ai.WithToolChoice(ai.ToolChoiceAuto), ) - if err != nil { fmt.Printf("Error: %v\n", err) return