Skip to content

Commit 774dba6

Browse files
4ndrelim4ndrelim
committed
feat: Add DynamicTool to represent generic XtraMCP tool
1 parent 622ed60 commit 774dba6

2 files changed

Lines changed: 183 additions & 10 deletions

File tree

‎internal/services/toolkit/client/client.go‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import (
99
"paperdebugger/internal/services"
1010
"paperdebugger/internal/services/toolkit/handler"
1111
"paperdebugger/internal/services/toolkit/registry"
12-
"paperdebugger/internal/services/toolkit/tools/xtragpt"
12+
"paperdebugger/internal/services/toolkit/tools/xtramcp"
1313

1414
"github.com/openai/openai-go/v2"
1515
"github.com/openai/openai-go/v2/option"
@@ -42,7 +42,6 @@ func NewAIClient(
4242
option.WithAPIKey(cfg.OpenAIAPIKey),
4343
)
4444
CheckOpenAIWorks(oaiClient, logger)
45-
toolSearchPapers := xtragpt.NewSearchPapersTool(db, projectService)
4645
// toolPaperScore := tools.NewPaperScoreTool(db, projectService)
4746
// toolPaperScoreComment := tools.NewPaperScoreCommentTool(db, projectService, reverseCommentService)
4847

@@ -53,16 +52,17 @@ func NewAIClient(
5352
// toolRegistry.Register("paper_score", toolPaperScore.Description, toolPaperScore.Call)
5453
// toolRegistry.Register("paper_score_comment", toolPaperScoreComment.Description, toolPaperScoreComment.Call)
5554

56-
// toolRegistry.Register("export_papers")
57-
// toolRegistry.Register("get_conference_papers")
58-
// toolRegistry.Register("get_user_papers")
59-
toolRegistry.Register("search_relevant_papers", toolSearchPapers.Description, toolSearchPapers.Call)
60-
// toolRegistry.Register("search_user")
61-
// toolRegistry.Register("identify_improvements")
62-
// toolRegistry.Register("suggest_improvement")
55+
// Load tools dynamically from backend (TODO: Make URL configurable / Xtramcp url)
56+
xtraMCPLoader := xtramcp.NewXtraMCPLoader(db, projectService, "http://localhost:8080/mcp")
57+
err := xtraMCPLoader.LoadToolsFromBackend(toolRegistry)
58+
if err != nil {
59+
logger.Errorf("[AI Client] Failed to load XtraMCP tools: %v", err)
60+
// Fallback to static tools or return error based on your preference
61+
} else {
62+
logger.Info("[AI Client] Successfully loaded XtraMCP tools")
63+
}
6364

6465
toolCallHandler := handler.NewToolCallHandler(toolRegistry)
65-
6666
client := &AIClient{
6767
openaiClient: &oaiClient,
6868
toolCallHandler: toolCallHandler,

‎internal/services/toolkit/tools/xtramcp/tool.go‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
package xtramcp
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"encoding/json"
7+
"fmt"
8+
"io"
9+
"net/http"
10+
"paperdebugger/internal/libs/db"
11+
"paperdebugger/internal/services"
12+
toolCallRecordDB "paperdebugger/internal/services/toolkit/db"
13+
"strings"
14+
"time"
15+
16+
"github.com/openai/openai-go/v2"
17+
"github.com/openai/openai-go/v2/packages/param"
18+
"github.com/openai/openai-go/v2/responses"
19+
"github.com/samber/lo"
20+
)
21+
22+
// ToolSchema represents the schema from your backend
23+
type ToolSchema struct {
24+
Name string `json:"name"`
25+
Description string `json:"description"`
26+
InputSchema map[string]interface{} `json:"inputSchema"`
27+
OutputSchema map[string]interface{} `json:"outputSchema"`
28+
}
29+
30+
// MCPRequest represents the JSON-RPC request structure
31+
type MCPRequest struct {
32+
JSONRPC string `json:"jsonrpc"`
33+
Method string `json:"method"`
34+
ID int `json:"id"`
35+
Params MCPParams `json:"params"`
36+
}
37+
38+
// MCPParams represents the parameters for the MCP request
39+
type MCPParams struct {
40+
Name string `json:"name"`
41+
Arguments map[string]interface{} `json:"arguments"`
42+
}
43+
44+
// DynamicTool represents a generic tool that can handle any schema
45+
type DynamicTool struct {
46+
Name string
47+
Description responses.ToolUnionParam
48+
toolCallRecordDB *toolCallRecordDB.ToolCallRecordDB
49+
projectService *services.ProjectService
50+
coolDownTime time.Duration
51+
baseURL string
52+
client *http.Client
53+
schema map[string]interface{}
54+
sessionID string // Reuse the session ID from initialization
55+
}
56+
57+
// NewDynamicTool creates a new dynamic tool from a schema
58+
func NewDynamicTool(db *db.DB, projectService *services.ProjectService, toolSchema ToolSchema, baseURL string, sessionID string) *DynamicTool {
59+
// Create tool description with the schema
60+
description := responses.ToolUnionParam{
61+
OfFunction: &responses.FunctionToolParam{
62+
Name: toolSchema.Name,
63+
Description: param.NewOpt(toolSchema.Description),
64+
Parameters: openai.FunctionParameters(toolSchema.InputSchema),
65+
},
66+
}
67+
68+
toolCallRecordDB := toolCallRecordDB.NewToolCallRecordDB(db)
69+
return &DynamicTool{
70+
Name: toolSchema.Name,
71+
Description: description,
72+
toolCallRecordDB: toolCallRecordDB,
73+
projectService: projectService,
74+
coolDownTime: 5 * time.Minute,
75+
baseURL: baseURL,
76+
client: &http.Client{},
77+
schema: toolSchema.InputSchema,
78+
sessionID: sessionID, // Store the session ID for reuse
79+
}
80+
}
81+
82+
// Call handles the tool execution (generic for any tool)
83+
func (t *DynamicTool) Call(ctx context.Context, toolCallId string, args json.RawMessage) (string, string, error) {
84+
// Parse arguments as generic map since we don't know the structure
85+
var argsMap map[string]interface{}
86+
err := json.Unmarshal(args, &argsMap)
87+
if err != nil {
88+
return "", "", err
89+
}
90+
91+
// Create function call record
92+
record, err := t.toolCallRecordDB.Create(ctx, toolCallId, t.Name, argsMap)
93+
if err != nil {
94+
return "", "", err
95+
}
96+
97+
// Execute the tool via MCP
98+
respStr, err := t.executeTool(argsMap)
99+
if err != nil {
100+
err = fmt.Errorf("failed to execute tool %s: %v", t.Name, err)
101+
t.toolCallRecordDB.OnError(ctx, record, err)
102+
return "", "", err
103+
}
104+
105+
rawJson, err := json.Marshal(respStr)
106+
if err != nil {
107+
err = fmt.Errorf("failed to marshal tool result: %v", err)
108+
t.toolCallRecordDB.OnError(ctx, record, err)
109+
return "", "", err
110+
}
111+
t.toolCallRecordDB.OnSuccess(ctx, record, string(rawJson))
112+
113+
return respStr, "", nil
114+
}
115+
116+
// executeTool makes the MCP request (generic for any tool)
117+
func (t *DynamicTool) executeTool(args map[string]interface{}) (string, error) {
118+
// Use the stored session ID - no need to re-initialize!
119+
fmt.Printf("Using existing sessionId for %s: %s\n", t.Name, t.sessionID)
120+
121+
request := MCPRequest{
122+
JSONRPC: "2.0",
123+
Method: "tools/call",
124+
ID: int(time.Now().Unix()), // to ensure unique ID; TODO: consider better ID generation
125+
Params: MCPParams{
126+
Name: t.Name,
127+
Arguments: args,
128+
},
129+
}
130+
131+
// Marshal request to JSON
132+
jsonData, err := json.Marshal(request)
133+
if err != nil {
134+
return "", fmt.Errorf("failed to marshal MCP request: %w", err)
135+
}
136+
137+
// Create HTTP request
138+
req, err := http.NewRequest("POST", t.baseURL, bytes.NewBuffer(jsonData))
139+
if err != nil {
140+
return "", fmt.Errorf("failed to create HTTP request: %w", err)
141+
}
142+
143+
// Set headers
144+
req.Header.Set("Content-Type", "application/json")
145+
req.Header.Set("Accept", "application/json, text/event-stream")
146+
req.Header.Set("mcp-session-id", t.sessionID) // Use the stored session ID
147+
148+
// Make the request
149+
resp, err := t.client.Do(req)
150+
if err != nil {
151+
return "", fmt.Errorf("failed to make request: %w", err)
152+
}
153+
defer resp.Body.Close()
154+
155+
// Read response
156+
body, err := io.ReadAll(resp.Body)
157+
if err != nil {
158+
return "", fmt.Errorf("failed to read response: %w", err)
159+
}
160+
fmt.Printf("Response body for %s: %s\n", t.Name, string(body))
161+
162+
// Parse response (assuming stream format)
163+
lines := strings.Split(string(body), "\n")
164+
lines = lo.Filter(lines, func(line string, _ int) bool {
165+
return strings.HasPrefix(line, "data:")
166+
})
167+
if len(lines) == 0 {
168+
return "", fmt.Errorf("no data line found")
169+
}
170+
line := lines[0]
171+
line = strings.TrimPrefix(line, "data: ")
172+
return line, nil
173+
}

0 commit comments

Comments
 (0)