Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
9 changes: 7 additions & 2 deletions go/plugins/ollama/ollama.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand All @@ -111,6 +111,7 @@ type ModelDefinition struct {
type generator struct {
model ModelDefinition
serverAddress string
timeout int
}

type ollamaMessage struct {
Expand Down Expand Up @@ -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.
Expand All @@ -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{}
}

Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion go/samples/ollama-tools/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -80,7 +81,6 @@ func main() {
ai.WithTools(weatherTool),
ai.WithToolChoice(ai.ToolChoiceAuto),
)

if err != nil {
fmt.Printf("Error: %v\n", err)
return
Expand Down
Loading