Skip to content

Commit 2dedda6

Browse files
SamMorrowDrumsSamMorrowDrums
committed
add agent search
1 parent 495c0cb commit 2dedda6

File tree

2 files changed

+239
-0
lines changed

2 files changed

+239
-0
lines changed

‎pkg/github/blackbird-agent.go

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
package github
2+
3+
import (
4+
"bufio"
5+
"bytes"
6+
"context"
7+
"encoding/json"
8+
"fmt"
9+
"io"
10+
"net/http"
11+
"os"
12+
"strings"
13+
14+
"github.com/github/github-mcp-server/pkg/translations"
15+
"github.com/mark3labs/mcp-go/mcp"
16+
"github.com/mark3labs/mcp-go/server"
17+
)
18+
19+
// POST https://api.githubcopilot.com/agents/github-search-agent
20+
// Authorization: Bearer <API_TOKEN>
21+
// Copilot-Integration-Id: copilot-pr-reviews
22+
// Content-Type: application/json
23+
24+
// {
25+
// "messages": [
26+
// {
27+
// "role": "user",
28+
// "copilot_references": [
29+
// {
30+
// "type": "github.searchInputs",
31+
// "data": {
32+
// "type": "search-inputs",
33+
// "id": 0,
34+
// "depth": 1,
35+
// "task": "InheritDocstrings metaclass doesn't work for properties\nInside the InheritDocstrings metaclass it uses `inspect.isfunction` which returns `False` for properties.\n",
36+
// "owner": "astropy",
37+
// "repositoryName": "astropy",
38+
// "numSnippets": 20
39+
// }
40+
// }
41+
// ],
42+
// "name": "question"
43+
// }
44+
// ]
45+
// }
46+
47+
type SearchAgentResult struct {
48+
Contents string `json:"contents"`
49+
LanguageName string `json:"languageName"`
50+
Path string `json:"path"`
51+
Range *Range `json:"range"`
52+
Ref string `json:"ref"`
53+
Type string `json:"type"`
54+
URL string `json:"url"`
55+
}
56+
type Range struct {
57+
Start int `json:"start"`
58+
End int `json:"end"`
59+
}
60+
61+
// searchAgent creates a tool wrap the GitHub Blackbird Search Agent.
62+
func SearchAgent(getClient GetClientFn, t translations.TranslationHelperFunc) (tool mcp.Tool, handler server.ToolHandlerFunc) {
63+
return mcp.NewTool("search_repository_with_agent",
64+
mcp.WithDescription(t("TOOL_SEARCH_AGENT_DESCRIPTION", `Search for relevent code within a GitHub repository, when you need to find code relevent to
65+
performing a specific task.`)),
66+
mcp.WithString("owner",
67+
mcp.Required(),
68+
mcp.Description("Owner of the repository"),
69+
),
70+
mcp.WithString("repo",
71+
mcp.Required(),
72+
mcp.Description("Name of the repository"),
73+
),
74+
mcp.WithString("task",
75+
mcp.Required(),
76+
mcp.Description("Task to perform"),
77+
),
78+
mcp.WithNumber("numSnippets",
79+
mcp.Description("Number of snippets to return"),
80+
mcp.Min(1),
81+
mcp.Max(100),
82+
),
83+
),
84+
func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
85+
// api.githubcopilot.com
86+
// currently behind this feature flag:
87+
// https://devportal.githubapp.com/feature-flags/capi-search-agent-enabled/targeting-rules/actors-view?stamp=dotcom
88+
89+
owner, err := requiredParam[string](request, "owner")
90+
if err != nil {
91+
return mcp.NewToolResultError(err.Error()), nil
92+
}
93+
repo, err := requiredParam[string](request, "repo")
94+
if err != nil {
95+
return mcp.NewToolResultError(err.Error()), nil
96+
}
97+
task, err := requiredParam[string](request, "task")
98+
if err != nil {
99+
return mcp.NewToolResultError(err.Error()), nil
100+
}
101+
numSnippets, err := OptionalIntParamWithDefault(request, "numSnippets", 1)
102+
if err != nil {
103+
return mcp.NewToolResultError(err.Error()), nil
104+
}
105+
// Construct the request payload
106+
payload := map[string]interface{}{
107+
"messages": []map[string]interface{}{
108+
{
109+
"role": "user",
110+
"copilot_references": []map[string]interface{}{
111+
{
112+
"type": "github.searchInputs",
113+
"data": map[string]interface{}{
114+
"type": "search-inputs",
115+
"id": 0,
116+
"depth": 2,
117+
"task": task,
118+
"owner": owner,
119+
"repositoryName": repo,
120+
"numSnippets": numSnippets,
121+
},
122+
},
123+
},
124+
"name": "question",
125+
},
126+
},
127+
}
128+
129+
// Encode the payload as JSON
130+
payloadBytes, err := json.Marshal(payload)
131+
if err != nil {
132+
return nil, fmt.Errorf("failed to marshal payload: %w", err)
133+
}
134+
135+
capiHost := "api.githubcopilot.com"
136+
capiProtocol := "https"
137+
capiPort := "443"
138+
// Check if the CAPI_HOST environment variable is set
139+
if host := os.Getenv("CAPI_HOST"); host != "" {
140+
capiHost = host
141+
}
142+
// Check if the CAPI_PROTOCOL environment variable is set
143+
if protocol := os.Getenv("CAPI_PROTOCOL"); protocol != "" {
144+
capiProtocol = protocol
145+
}
146+
if capiProtocol == "http" {
147+
capiPort = "80"
148+
}
149+
150+
// Check if the CAPI_PORT environment variable is set
151+
if port := os.Getenv("CAPI_PORT"); port != "" {
152+
capiPort = port
153+
}
154+
155+
capiURL := fmt.Sprintf("%s://%s:%s/agents/github-search-agent", capiProtocol, capiHost, capiPort)
156+
157+
req, err := http.NewRequest("POST", capiURL, bytes.NewBuffer(payloadBytes))
158+
159+
req.Header.Set("Copilot-Integration-Id", "copilot-pr-reviews")
160+
if err != nil {
161+
return nil, fmt.Errorf("failed to create request: %w", err)
162+
}
163+
client, err := getClient(ctx)
164+
if err != nil {
165+
return nil, fmt.Errorf("failed to get client: %w", err)
166+
}
167+
168+
resp, err := client.Client().Do(req)
169+
if err != nil {
170+
return nil, fmt.Errorf("failed to send request: %w", err)
171+
}
172+
defer resp.Body.Close()
173+
174+
if resp.StatusCode != http.StatusOK {
175+
return nil, fmt.Errorf("failed to get response: %s", resp.Status)
176+
}
177+
178+
// Create a buffered reader to read the response body line by line
179+
reader := bufio.NewReader(resp.Body)
180+
181+
var results []SearchAgentResult
182+
183+
for {
184+
line, err := reader.ReadString('\n') // Read until the next newline
185+
if err != nil {
186+
if err == io.EOF {
187+
break // End of the response
188+
}
189+
return nil, fmt.Errorf("failed to read response: %w", err)
190+
}
191+
192+
// Check if the line starts with "data: "
193+
if strings.HasPrefix(line, "data: ") {
194+
data := strings.TrimPrefix(line, "data: ")
195+
data = strings.TrimSpace(data)
196+
// Check if the line is empty
197+
if data == "" {
198+
continue // Skip empty lines
199+
}
200+
201+
// Stop processing if "data: [DONE]" is encountered
202+
if data == "[DONE]" {
203+
break
204+
}
205+
206+
// Unmarshal the JSON data
207+
var result map[string]interface{}
208+
if err := json.Unmarshal([]byte(data), &result); err != nil {
209+
return nil, fmt.Errorf("failed to read response: %w", err)
210+
}
211+
// Check if the result contains "copilot_references" key this in an array, and for each entry extract into []SearchAgentResult
212+
if copilotReferences, ok := result["copilot_references"].([]interface{}); ok {
213+
for _, ref := range copilotReferences {
214+
if refMap, ok := ref.(map[string]interface{}); ok {
215+
if data, ok := refMap["data"].(map[string]interface{}); ok {
216+
if resultData, err := json.Marshal(data); err == nil {
217+
var searchResult SearchAgentResult
218+
if err := json.Unmarshal(resultData, &searchResult); err == nil {
219+
results = append(results, searchResult)
220+
}
221+
}
222+
}
223+
}
224+
}
225+
}
226+
}
227+
}
228+
r, err := json.Marshal(results)
229+
if err != nil {
230+
return nil, fmt.Errorf("failed to marshal response: %w", err)
231+
}
232+
return mcp.NewToolResultText(string(r)), nil
233+
}
234+
}

‎pkg/github/tools.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,10 @@ func InitToolsets(passedToolsets []string, readOnly bool, getClient GetClientFn,
7878
toolsets.NewServerTool(GetSecretScanningAlert(getClient, t)),
7979
toolsets.NewServerTool(ListSecretScanningAlerts(getClient, t)),
8080
)
81+
codeSearch := toolsets.NewToolset("code_search", "Advanced code search tools").
82+
AddReadTools(
83+
toolsets.NewServerTool(SearchAgent(getClient, t)),
84+
)
8185
// Keep experiments alive so the system doesn't error out when it's always enabled
8286
experiments := toolsets.NewToolset("experiments", "Experimental features that are not considered stable yet")
8387

@@ -88,6 +92,7 @@ func InitToolsets(passedToolsets []string, readOnly bool, getClient GetClientFn,
8892
tsg.AddToolset(pullRequests)
8993
tsg.AddToolset(codeSecurity)
9094
tsg.AddToolset(secretProtection)
95+
tsg.AddToolset(codeSearch)
9196
tsg.AddToolset(experiments)
9297
// Enable the requested features
9398

0 commit comments

Comments
 (0)