From 20256bd6b4f1e21a6ffeb961ef9fc9183f43f48a Mon Sep 17 00:00:00 2001
From: Bel LaPointe <153096461+breel-render@users.noreply.github.com>
Date: Fri, 19 Apr 2024 13:18:12 -0600
Subject: [PATCH] try and raise ollama timeout

---
 ai.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/ai.go b/ai.go
index ba103aa..3c42521 100644
--- a/ai.go
+++ b/ai.go
@@ -2,6 +2,8 @@ package main
 
 import (
 	"context"
+	"net/http"
+	"time"
 
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/llms/ollama"
@@ -32,9 +34,20 @@ func NewAIOllama(url, model string) AIOllama {
 }
 
 func (ai AIOllama) Do(ctx context.Context, prompt string) (string, error) {
+	c := &http.Client{
+		Timeout: time.Hour,
+		Transport: &http.Transport{
+			//DisableKeepAlives:     true,
+			IdleConnTimeout:       time.Hour,
+			ResponseHeaderTimeout: time.Hour,
+			ExpectContinueTimeout: time.Hour,
+		},
+	}
+	defer c.CloseIdleConnections()
 	llm, err := ollama.New(
 		ollama.WithModel(ai.model),
 		ollama.WithServerURL(ai.url),
+		ollama.WithHTTPClient(c),
 	)
 	if err != nil {
 		return "", err