6bb1c84 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

Improve Sourcegraph tool with context window and fix diagnostics

6bb1c84Kujtim Hoxha2 years ago

cmd/lsp/main.go+12

compare with main

unifiedsplit

@@ -1,4 +1,16 @@

1package main

3+import (

4+ "context"

5+ "fmt"

7+ "github.com/kujtimiihoxha/termai/internal/llm/tools"

8+)

3func main() {

10func main() {

11+ t := tools.NewSourcegraphTool()

12+ r, _ := t.Run(context.Background(), tools.ToolCall{

13+ Input: `{"query": "context.WithCancel lang:go"}`,

14+ })

15+ fmt.Println(r.Content)

16}

internal/llm/agent/task.go+1

@@ -34,6 +34,7 @@ func NewTaskAgent(app *app.App) (Agent, error) {

34 tools.NewGlobTool(),

35 tools.NewGrepTool(),

36 tools.NewLsTool(),

37+ tools.NewSourcegraphTool(),

37 tools.NewViewTool(app.LSPClients),

38 tools.NewViewTool(app.LSPClients),

38 },

39 },

39 model: model,

40 model: model,

internal/llm/tools/diagnostics.go+58−4

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

4 "context"

5 "encoding/json"

6 "fmt"

7+ "maps"

7 "sort"

8 "sort"

8 "strings"

9 "strings"

9 "time"

10 "time"

@@ -50,7 +51,7 @@ func (b *diagnosticsTool) Run(ctx context.Context, call ToolCall) (ToolResponse,

50 return NewTextErrorResponse("no LSP clients available"), nil

51 return NewTextErrorResponse("no LSP clients available"), nil

51 }

52 }

53− if params.FilePath == "" {

54+ if params.FilePath != "" {

54 notifyLspOpenFile(ctx, params.FilePath, lsps)

55 notifyLspOpenFile(ctx, params.FilePath, lsps)

55 }

56 }

@@ -60,15 +61,68 @@ func (b *diagnosticsTool) Run(ctx context.Context, call ToolCall) (ToolResponse,

60}

61}

62func notifyLspOpenFile(ctx context.Context, filePath string, lsps map[string]*lsp.Client) {

63func notifyLspOpenFile(ctx context.Context, filePath string, lsps map[string]*lsp.Client) {

64+ // Create a channel to receive diagnostic notifications

65+ diagChan := make(chan struct{}, 1)

66+

67+ // Register a temporary diagnostic handler for each client

63 for _, client := range lsps {

68 for _, client := range lsps {

69+ // Store the original diagnostics map to detect changes

70+ originalDiags := make(map[protocol.DocumentUri][]protocol.Diagnostic)

71+ maps.Copy(originalDiags, client.GetDiagnostics())

72+

73+ // Create a notification handler that will signal when diagnostics are received

74+ handler := func(params json.RawMessage) {

75+ var diagParams protocol.PublishDiagnosticsParams

76+ if err := json.Unmarshal(params, &diagParams); err != nil {

77+ return

78+ }

79+

80+ // If this is for our file or we've received any new diagnostics, signal completion

81+ if diagParams.URI.Path() == filePath || hasDiagnosticsChanged(client.GetDiagnostics(), originalDiags) {

82+ select {

83+ case diagChan <- struct{}{}:

84+ // Signal sent

85+ default:

86+ // Channel already has a value, no need to send again

87+ }

88+ }

89+ }

90+

91+ // Register our temporary handler

92+ client.RegisterNotificationHandler("textDocument/publishDiagnostics", handler)

93+

94+ // Open the file

64 err := client.OpenFile(ctx, filePath)

95 err := client.OpenFile(ctx, filePath)

65 if err != nil {

96 if err != nil {

66− // Wait for the file to be opened and diagnostics to be received

97+ // If there's an error opening the file, continue to the next client

67− // TODO: see if we can do this in a more efficient way

98+ continue

68− time.Sleep(3 * time.Second)

69 }

99 }

100+ }

101+

102+ // Wait for diagnostics with a reasonable timeout

103+ select {

104+ case <-diagChan:

105+ // Diagnostics received

106+ case <-time.After(5 * time.Second):

107+ // Timeout after 2 seconds - this is a fallback in case no diagnostics are published

108+ case <-ctx.Done():

109+ // Context cancelled

110+ }

111

112+ // Note: We're not unregistering our handler because the Client.RegisterNotificationHandler

113+ // replaces any existing handler, and we'll be replaced by the original handler when

114+ // the LSP client is reinitialized or when a new handler is registered.

115+}

116+

117+// hasDiagnosticsChanged checks if there are any new diagnostics compared to the original set

118+func hasDiagnosticsChanged(current, original map[protocol.DocumentUri][]protocol.Diagnostic) bool {

119+ for uri, diags := range current {

120+ origDiags, exists := original[uri]

121+ if !exists || len(diags) != len(origDiags) {

122+ return true

123+ }

71 }

124 }

125+ return false

72}

126}

127

74func appendDiagnostics(filePath string, lsps map[string]*lsp.Client) string {

128func appendDiagnostics(filePath string, lsps map[string]*lsp.Client) string {

internal/llm/tools/sourcegraph.go+16−16

@@ -111,15 +111,10 @@ TIPS:

111)

112

113type SourcegraphParams struct {

114− Query string `json:"query"`

114+ Query string `json:"query"`

115− Count int `json:"count,omitempty"`

115+ Count int `json:"count,omitempty"`

116− Timeout int `json:"timeout,omitempty"`

117−}

118−

119−type SourcegraphPermissionsParams struct {

120− Query string `json:"query"`

121− Count int `json:"count,omitempty"`

116+ ContextWindow int `json:"context_window,omitempty"`

122− Timeout int `json:"timeout,omitempty"`

117+ Timeout int `json:"timeout,omitempty"`

123}

118}

124

119

125type sourcegraphTool struct {

120type sourcegraphTool struct {

@@ -147,6 +142,10 @@ func (t *sourcegraphTool) Info() ToolInfo {

147 "type": "number",

142 "type": "number",

148 "description": "Optional number of results to return (default: 10, max: 20)",

143 "description": "Optional number of results to return (default: 10, max: 20)",

149 },

144 },

145+ "context_window": map[string]any{

146+ "type": "number",

147+ "description": "The context around the match to return (default: 10 lines)",

148+ },

150 "timeout": map[string]any{

149 "timeout": map[string]any{

151 "type": "number",

150 "type": "number",

152 "description": "Optional timeout in seconds (max 120)",

151 "description": "Optional timeout in seconds (max 120)",

@@ -173,6 +172,9 @@ func (t *sourcegraphTool) Run(ctx context.Context, call ToolCall) (ToolResponse,

173 params.Count = 20 // Limit to 20 results

172 params.Count = 20 // Limit to 20 results

174 }

173 }

175

174

175+ if params.ContextWindow <= 0 {

176+ params.ContextWindow = 10 // Default context window

177+ }

176 client := t.client

178 client := t.client

177 if params.Timeout > 0 {

179 if params.Timeout > 0 {

178 maxTimeout := 120 // 2 minutes

180 maxTimeout := 120 // 2 minutes

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

194 }

196 }

195

197

196 request := graphqlRequest{

198 request := graphqlRequest{

197− Query: "query Search($query: String!) { search(query: $query, version: V2, patternType: standard ) { results { matchCount, limitHit, resultCount, approximateResultCount, missing { name }, timedout { name }, indexUnavailable, results { __typename, ... on FileMatch { repository { name }, file { path, url, content }, lineMatches { preview, lineNumber, offsetAndLengths } } } } } }",

199+ Query: "query Search($query: String!) { search(query: $query, version: V2, patternType: keyword ) { results { matchCount, limitHit, resultCount, approximateResultCount, missing { name }, timedout { name }, indexUnavailable, results { __typename, ... on FileMatch { repository { name }, file { path, url, content }, lineMatches { preview, lineNumber, offsetAndLengths } } } } } }",

198 }

200 }

199 request.Variables.Query = params.Query

201 request.Variables.Query = params.Query

200

202

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

246 }

248 }

247

249

248 // Format the results in a readable way

250 // Format the results in a readable way

249− formattedResults, err := formatSourcegraphResults(result)

251+ formattedResults, err := formatSourcegraphResults(result, params.ContextWindow)

250 if err != nil {

252 if err != nil {

251 return NewTextErrorResponse("Failed to format results: " + err.Error()), nil

253 return NewTextErrorResponse("Failed to format results: " + err.Error()), nil

252 }

254 }

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

254 return NewTextResponse(formattedResults), nil

256 return NewTextResponse(formattedResults), nil

255}

257}

256

258

257−func formatSourcegraphResults(result map[string]any) (string, error) {

259+func formatSourcegraphResults(result map[string]any, contextWindow int) (string, error) {

258 var buffer strings.Builder

260 var buffer strings.Builder

259

261

260 // Check for errors in the GraphQL response

262 // Check for errors in the GraphQL response

@@ -364,8 +366,7 @@ func formatSourcegraphResults(result map[string]any) (string, error) {

364 buffer.WriteString("```\n")

366 buffer.WriteString("```\n")

365

367

366 // Display context before the match (up to 10 lines)

368 // Display context before the match (up to 10 lines)

367− contextBefore := 10

368− startLine := max(1, int(lineNumber)-contextBefore)

369+ startLine := max(1, int(lineNumber)-contextWindow)

369

370

370 for j := startLine - 1; j < int(lineNumber)-1 && j < len(lines); j++ {

371 for j := startLine - 1; j < int(lineNumber)-1 && j < len(lines); j++ {

371 if j >= 0 {

372 if j >= 0 {

@@ -377,8 +378,7 @@ func formatSourcegraphResults(result map[string]any) (string, error) {

377 buffer.WriteString(fmt.Sprintf("%d| %s\n", int(lineNumber), preview))

378 buffer.WriteString(fmt.Sprintf("%d| %s\n", int(lineNumber), preview))

378

379

379 // Display context after the match (up to 10 lines)

380 // Display context after the match (up to 10 lines)

380− contextAfter := 10

381− endLine := int(lineNumber) + contextAfter

381+ endLine := int(lineNumber) + contextWindow

382

383 for j := int(lineNumber); j < endLine && j < len(lines); j++ {

384 if j < len(lines) {