Skip to content

Commit e776108

Browse files
authored
opt: better json schema (#51)
1 parent c8b30ac commit e776108

File tree

2 files changed

+13
-4
lines changed

2 files changed

+13
-4
lines changed

llm/mcp/handler.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ func NewTool[R any, T any](name string, desc string, schema json.RawMessage, han
3939
if resp, err := handler(ctx, req); err != nil {
4040
isError = true
4141
final = err.Error()
42-
} else if js, err := utils.MarshalJSONIndent(resp); err != nil {
42+
} else if js, err := utils.MarshalJSONBytes(resp); err != nil {
4343
isError = true
4444
final = err.Error()
4545
} else {

llm/tool/ast_read.go

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,14 @@ func NewNodeID(id uniast.Identity) NodeID {
234234
}
235235
}
236236

237+
func (n NodeID) Identity() uniast.Identity {
238+
return uniast.Identity{
239+
ModPath: n.ModPath,
240+
PkgPath: n.PkgPath,
241+
Name: n.Name,
242+
}
243+
}
244+
237245
func (t *ASTReadTools) getRepoAST(repoName string) (*uniast.Repository, error) {
238246
repo, ok := t.repos.Load(repoName)
239247
if !ok {
@@ -419,8 +427,8 @@ func (t *ASTReadTools) getFileStructure(_ context.Context, req GetFileStructReq,
419427
}
420428

421429
type GetASTNodeReq struct {
422-
RepoName string `json:"repo_name" jsonschema:"description=the name of the repository"`
423-
NodeIDs []uniast.Identity `json:"node_ids" jsonschema:"description=the identities of the ast node"`
430+
RepoName string `json:"repo_name" jsonschema:"description=the name of the repository"`
431+
NodeIDs []NodeID `json:"node_ids" jsonschema:"description=the identities of the ast node"`
424432
}
425433

426434
type GetASTNodeResp struct {
@@ -439,7 +447,8 @@ func (t *ASTReadTools) GetASTNode(_ context.Context, params GetASTNodeReq) (*Get
439447
}
440448

441449
resp := new(GetASTNodeResp)
442-
for _, id := range params.NodeIDs {
450+
for _, nid := range params.NodeIDs {
451+
id := nid.Identity()
443452
log.Debug("query ast node %v", id.Full())
444453
node := repo.GetNode(id)
445454
if node == nil {

0 commit comments

Comments
 (0)