5056d44 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix: remove pool logic

5056d44Raphael Amorim1 year ago

internal/llm/tools/fetch.go+15−49

compare with main

unifiedsplit

@@ -7,7 +7,6 @@ import (

7 "io"

8 "net/http"

9 "strings"

10− "sync"

11 "time"

10 "time"

13 md "github.com/JohannesKaufmann/html-to-markdown"

12 md "github.com/JohannesKaufmann/html-to-markdown"

@@ -29,10 +28,8 @@ type FetchPermissionsParams struct {

29}

28}

31type fetchTool struct {

30type fetchTool struct {

32− client *http.Client

31+ client *http.Client

33− clientPool map[int]*http.Client

34− clientPoolMu sync.RWMutex

35− permissions permission.Service

32+ permissions permission.Service

36}

33}

38const (

35const (

@@ -78,51 +75,10 @@ func NewFetchTool(permissions permission.Service) BaseTool {

78 IdleConnTimeout: 90 * time.Second,

75 IdleConnTimeout: 90 * time.Second,

79 },

76 },

80 },

77 },

81− clientPool: make(map[int]*http.Client),

82 permissions: permissions,

78 permissions: permissions,

83 }

79 }

84}

80}

86−// getClientForTimeout returns a cached client for the given timeout or the default client

87−func (t *fetchTool) getClientForTimeout(timeout int) *http.Client {

88− if timeout <= 0 {

89− return t.client

90− }

91−

92− maxTimeout := 120 // 2 minutes

93− if timeout > maxTimeout {

94− timeout = maxTimeout

95− }

96−

97− // Check if we have a cached client for this timeout

98− t.clientPoolMu.RLock()

99− if client, exists := t.clientPool[timeout]; exists {

100− t.clientPoolMu.RUnlock()

101− return client

102− }

103− t.clientPoolMu.RUnlock()

104−

105− // Create and cache a new client

106− t.clientPoolMu.Lock()

107− defer t.clientPoolMu.Unlock()

108−

109− // Double-check in case another goroutine created it

110− if client, exists := t.clientPool[timeout]; exists {

111− return client

112− }

113−

114− client := &http.Client{

115− Timeout: time.Duration(timeout) * time.Second,

116− Transport: &http.Transport{

117− MaxIdleConns: 100,

118− MaxIdleConnsPerHost: 10,

119− IdleConnTimeout: 90 * time.Second,

120− },

121− }

122− t.clientPool[timeout] = client

123− return client

124−}

125−

126func (t *fetchTool) Info() ToolInfo {

82func (t *fetchTool) Info() ToolInfo {

127 return ToolInfo{

83 return ToolInfo{

128 Name: FetchToolName,

84 Name: FetchToolName,

@@ -185,16 +141,26 @@ func (t *fetchTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error

185 return ToolResponse{}, permission.ErrorPermissionDenied

141 return ToolResponse{}, permission.ErrorPermissionDenied

186 }

142 }

187

143

188− client := t.getClientForTimeout(params.Timeout)

144+ // Handle timeout with context

145+ requestCtx := ctx

146+ if params.Timeout > 0 {

147+ maxTimeout := 120 // 2 minutes

148+ if params.Timeout > maxTimeout {

149+ params.Timeout = maxTimeout

150+ }

151+ var cancel context.CancelFunc

152+ requestCtx, cancel = context.WithTimeout(ctx, time.Duration(params.Timeout)*time.Second)

153+ defer cancel()

154+ }

189

155

190− req, err := http.NewRequestWithContext(ctx, "GET", params.URL, nil)

156+ req, err := http.NewRequestWithContext(requestCtx, "GET", params.URL, nil)

191 if err != nil {

157 if err != nil {

192 return ToolResponse{}, fmt.Errorf("failed to create request: %w", err)

158 return ToolResponse{}, fmt.Errorf("failed to create request: %w", err)

193 }

159 }

194

160

195 req.Header.Set("User-Agent", "crush/1.0")

161 req.Header.Set("User-Agent", "crush/1.0")

196

162

197− resp, err := client.Do(req)

163+ resp, err := t.client.Do(req)

198 if err != nil {

164 if err != nil {

199 return ToolResponse{}, fmt.Errorf("failed to fetch URL: %w", err)

165 return ToolResponse{}, fmt.Errorf("failed to fetch URL: %w", err)

200 }

166 }

internal/llm/tools/grep.go+10−9

@@ -402,19 +402,20 @@ func fileContainsPattern(filePath string, pattern *regexp.Regexp) (bool, int, st

402 return false, 0, "", scanner.Err()

403}

404

405+var binaryExts = map[string]struct{}{

406+ ".exe": {}, ".dll": {}, ".so": {}, ".dylib": {},

407+ ".bin": {}, ".obj": {}, ".o": {}, ".a": {},

408+ ".zip": {}, ".tar": {}, ".gz": {}, ".bz2": {},

409+ ".jpg": {}, ".jpeg": {}, ".png": {}, ".gif": {},

410+ ".pdf": {}, ".doc": {}, ".docx": {}, ".xls": {},

411+ ".mp3": {}, ".mp4": {}, ".avi": {}, ".mov": {},

412+}

413+

405// isBinaryFile performs a quick check to determine if a file is binary

414// isBinaryFile performs a quick check to determine if a file is binary

406func isBinaryFile(filePath string) bool {

415func isBinaryFile(filePath string) bool {

407 // Check file extension first (fastest)

416 // Check file extension first (fastest)

408 ext := strings.ToLower(filepath.Ext(filePath))

417 ext := strings.ToLower(filepath.Ext(filePath))

409− binaryExts := map[string]bool{

410− ".exe": true, ".dll": true, ".so": true, ".dylib": true,

411− ".bin": true, ".obj": true, ".o": true, ".a": true,

412− ".zip": true, ".tar": true, ".gz": true, ".bz2": true,

413− ".jpg": true, ".jpeg": true, ".png": true, ".gif": true,

414− ".pdf": true, ".doc": true, ".docx": true, ".xls": true,

415− ".mp3": true, ".mp4": true, ".avi": true, ".mov": true,

416− }

417− if binaryExts[ext] {

418+ if _, isBinary := binaryExts[ext]; isBinary {

418 return true

419 return true

419 }

420 }

420

421

internal/llm/tools/sourcegraph.go+15−48

@@ -8,7 +8,6 @@ import (

8 "io"

9 "net/http"

10 "strings"

11− "sync"

12 "time"

11 "time"

13)

12)

@@ -25,9 +24,7 @@ type SourcegraphResponseMetadata struct {

25}

24}

27type sourcegraphTool struct {

26type sourcegraphTool struct {

28− client *http.Client

27+ client *http.Client

29− clientPool map[int]*http.Client

30− clientPoolMu sync.RWMutex

31}

28}

33const (

30const (

@@ -138,50 +135,9 @@ func NewSourcegraphTool() BaseTool {

138 IdleConnTimeout: 90 * time.Second,

135 IdleConnTimeout: 90 * time.Second,

139 },

136 },

140 },

137 },

141− clientPool: make(map[int]*http.Client),

142 }

138 }

143}

139}

144

140

145−// getClientForTimeout returns a cached client for the given timeout or the default client

146−func (t *sourcegraphTool) getClientForTimeout(timeout int) *http.Client {

147− if timeout <= 0 {

148− return t.client

149− }

150−

151− maxTimeout := 120 // 2 minutes

152− if timeout > maxTimeout {

153− timeout = maxTimeout

154− }

155−

156− // Check if we have a cached client for this timeout

157− t.clientPoolMu.RLock()

158− if client, exists := t.clientPool[timeout]; exists {

159− t.clientPoolMu.RUnlock()

160− return client

161− }

162− t.clientPoolMu.RUnlock()

163−

164− // Create and cache a new client

165− t.clientPoolMu.Lock()

166− defer t.clientPoolMu.Unlock()

167−

168− // Double-check in case another goroutine created it

169− if client, exists := t.clientPool[timeout]; exists {

170− return client

171− }

172−

173− client := &http.Client{

174− Timeout: time.Duration(timeout) * time.Second,

175− Transport: &http.Transport{

176− MaxIdleConns: 100,

177− MaxIdleConnsPerHost: 10,

178− IdleConnTimeout: 90 * time.Second,

179− },

180− }

181− t.clientPool[timeout] = client

182− return client

183−}

184−

185func (t *sourcegraphTool) Info() ToolInfo {

141func (t *sourcegraphTool) Info() ToolInfo {

186 return ToolInfo{

142 return ToolInfo{

187 Name: SourcegraphToolName,

143 Name: SourcegraphToolName,

@@ -227,7 +183,18 @@ func (t *sourcegraphTool) Run(ctx context.Context, call ToolCall) (ToolResponse,

227 if params.ContextWindow <= 0 {

183 if params.ContextWindow <= 0 {

228 params.ContextWindow = 10 // Default context window

184 params.ContextWindow = 10 // Default context window

229 }

185 }

230− client := t.getClientForTimeout(params.Timeout)

186+

187+ // Handle timeout with context

188+ requestCtx := ctx

189+ if params.Timeout > 0 {

190+ maxTimeout := 120 // 2 minutes

191+ if params.Timeout > maxTimeout {

192+ params.Timeout = maxTimeout

193+ }

194+ var cancel context.CancelFunc

195+ requestCtx, cancel = context.WithTimeout(ctx, time.Duration(params.Timeout)*time.Second)

196+ defer cancel()

197+ }

231

198

232 type graphqlRequest struct {

199 type graphqlRequest struct {

233 Query string `json:"query"`

200 Query string `json:"query"`

@@ -248,7 +215,7 @@ func (t *sourcegraphTool) Run(ctx context.Context, call ToolCall) (ToolResponse,

248 graphqlQuery := string(graphqlQueryBytes)

215 graphqlQuery := string(graphqlQueryBytes)

249

216

250 req, err := http.NewRequestWithContext(

217 req, err := http.NewRequestWithContext(

251− ctx,

218+ requestCtx,

252 "POST",

219 "POST",

253 "https://sourcegraph.com/.api/graphql",

220 "https://sourcegraph.com/.api/graphql",

254 bytes.NewBuffer([]byte(graphqlQuery)),

221 bytes.NewBuffer([]byte(graphqlQuery)),

@@ -260,7 +227,7 @@ func (t *sourcegraphTool) Run(ctx context.Context, call ToolCall) (ToolResponse,

260 req.Header.Set("Content-Type", "application/json")

227 req.Header.Set("Content-Type", "application/json")

261 req.Header.Set("User-Agent", "crush/1.0")

228 req.Header.Set("User-Agent", "crush/1.0")

262

229

263− resp, err := client.Do(req)

230+ resp, err := t.client.Do(req)

264 if err != nil {

231 if err != nil {

265 return ToolResponse{}, fmt.Errorf("failed to fetch URL: %w", err)

232 return ToolResponse{}, fmt.Errorf("failed to fetch URL: %w", err)

266 }

233 }