6bb1c84 · crush · git.iko.soy

git.iko.soy/crush

Improve Sourcegraph tool with context window and fix diagnostics

- Add context_window parameter to control code context display - Fix LSP diagnostics notification handling with proper async waiting - Switch to keyword search pattern for better results - Add Sourcegraph tool to task agent 🤖 Generated with termai Co-Authored-By: termai <noreply@termai.io>

commit6bb1c84f7f7f0430f2808d50c533e923aae4c787 authorKujtim Hoxha <kujtimii.h@gmail.com> dateFri 4 Apr 2025, 14:23 (2 years ago) parenteb9877e

4 files changed+87−20

compare with main

unifiedsplit

cmd/lsp/main.go+12 internal/llm/agent/task.go+1 internal/llm/tools/diagnostics.go+58−4 internal/llm/tools/sourcegraph.go+16−16

cmd/lsp/main.go+12

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

11package main

3+import (

4+ "context"

5+ "fmt"

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

8+)

310func 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)

416}

internal/llm/agent/task.go+1

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

3434 tools.NewGlobTool(),

3535 tools.NewGrepTool(),

3636 tools.NewLsTool(),

37+ tools.NewSourcegraphTool(),

3738 tools.NewViewTool(app.LSPClients),

3839 },

3940 model: model,

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

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

44 "context"

55 "encoding/json"

66 "fmt"

7+ "maps"

78 "sort"

89 "strings"

910 "time"

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

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

5152 }

5253

53− if params.FilePath == "" {

54+ if params.FilePath != "" {

5455 notifyLspOpenFile(ctx, params.FilePath, lsps)

5556 }

5657

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

6061}

6162

6263func 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

6368 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

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

6596 if err != nil {

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

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

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

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

98+ continue

6999 }

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+ }

70111

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+ }

71124 }

125+ return false

72126}

73127

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

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

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

111111)

112112

113113type SourcegraphParams struct {

114− Query string `json:"query"`

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"`

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

114+ Query string `json:"query"`

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

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

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

123118}

124119

125120type sourcegraphTool struct {

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

147142 "type": "number",

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

149144 },

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

146+ "type": "number",

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

148+ },

150149 "timeout": map[string]any{

151150 "type": "number",

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

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

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

174173 }

175174

175+ if params.ContextWindow <= 0 {

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

177+ }

176178 client := t.client

177179 if params.Timeout > 0 {

178180 maxTimeout := 120 // 2 minutes

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

194196 }

195197

196198 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 } } } } } }",

198200 }

199201 request.Variables.Query = params.Query

200202

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

246248 }

247249

248250 // Format the results in a readable way

249− formattedResults, err := formatSourcegraphResults(result)

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

250252 if err != nil {

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

252254 }

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

254256 return NewTextResponse(formattedResults), nil

255257}

256258

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

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

258260 var buffer strings.Builder

259261

260262 // Check for errors in the GraphQL response

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

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

365367

366368 // 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)

369370

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

371372 if j >= 0 {

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

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

378379

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

380− contextAfter := 10

381− endLine := int(lineNumber) + contextAfter

381+ endLine := int(lineNumber) + contextWindow

382382

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

384384 if j < len(lines) {