OpenMetadata/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/DefaultToolContext.java at f2e4bf42302bb0eb6f0251a0aa5e5d37cdc5fe9a · open-metadata/OpenMetadata · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
package org.openmetadata.mcp.tools;

import static org.openmetadata.mcp.McpUtils.getToolProperties;

import io.modelcontextprotocol.spec.McpSchema;
import java.util.List;
import java.util.Map;
import lombok.extern.slf4j.Slf4j;
import org.openmetadata.schema.utils.JsonUtils;
import org.openmetadata.service.limits.Limits;
import org.openmetadata.service.security.AuthorizationException;
import org.openmetadata.service.security.Authorizer;
import org.openmetadata.service.security.auth.CatalogSecurityContext;

@Slf4j
public class DefaultToolContext {
  public DefaultToolContext() {}

  /**
   * Loads tool definitions from a JSON file located at the specified path. The JSON file should
   * contain an array of tool definitions under the "tools" key.
   *
   * @return List of McpSchema.Tool objects loaded from the JSON file.
   */
  public List<McpSchema.Tool> loadToolsDefinitionsFromJson(String toolFilePath) {
    return getToolProperties(toolFilePath);
  }

  public McpSchema.CallToolResult callTool(
      Authorizer authorizer,
      Limits limits,
      String toolName,
      CatalogSecurityContext securityContext,
      McpSchema.CallToolRequest request) {
    LOG.info(
        "Catalog Principal: {} is trying to call the tool: {}",
        securityContext.getUserPrincipal().getName(),
        toolName);
    Map<String, Object> params = request.arguments();
    Object result;
    try {
      McpTool tool;
      switch (toolName) {
        case "search_metadata":
          tool = new SearchMetadataTool();
          result = tool.execute(authorizer, securityContext, params);
          break;
        case "semantic_search":
          result = new SemanticSearchTool().execute(authorizer, securityContext, params);
          break;
        case "get_entity_details":
          tool = new GetEntityTool();
          result = tool.execute(authorizer, securityContext, params);
          break;
        case "create_glossary":
          tool = new GlossaryTool();
          result = tool.execute(authorizer, limits, securityContext, params);
          break;
        case "create_glossary_term":
          tool = new GlossaryTermTool();
          result = tool.execute(authorizer, limits, securityContext, params);
          break;
        case "patch_entity":
          tool = new PatchEntityTool();
          result = tool.execute(authorizer, securityContext, params);
          break;
        case "get_entity_lineage":
          tool = new GetLineageTool();
          result = tool.execute(authorizer, securityContext, params);
          break;
        case "create_lineage":
          result = new LineageTool().execute(authorizer, securityContext, params);
          break;
        case "get_test_definitions":
          result = new TestDefinitionsTool().execute(authorizer, securityContext, params);
          break;
        case "create_test_case":
          result = new CreateTestCaseTool().execute(authorizer, limits, securityContext, params);
          break;
        case "root_cause_analysis":
          result = new RootCauseAnalysisTool().execute(authorizer, securityContext, params);
          break;
        case "create_metric":
          result = new CreateMetricTool().execute(authorizer, limits, securityContext, params);
          break;
        case "list_ingestion_pipelines":
          result = new ListIngestionPipelinesTool().execute(authorizer, securityContext, params);
          break;
        case "get_pipeline_status":
          result = new GetPipelineStatusTool().execute(authorizer, securityContext, params);
          break;
        case "trigger_ingestion_pipeline":
          result = new TriggerIngestionPipelineTool().execute(authorizer, securityContext, params);
          break;
        default:
          return McpSchema.CallToolResult.builder()
              .content(
                  List.of(
                      new McpSchema.TextContent(
                          JsonUtils.pojoToJson(Map.of("error", "Unknown function: " + toolName)))))
              .isError(true)
              .build();
      }

      return McpSchema.CallToolResult.builder()
          .content(List.of(new McpSchema.TextContent(JsonUtils.pojoToJson(result))))
          .isError(false)
          .build();
    } catch (AuthorizationException ex) {
      LOG.warn("Authorization error: {}", ex.getMessage());
      return McpSchema.CallToolResult.builder()
          .content(
              List.of(
                  new McpSchema.TextContent(
                      JsonUtils.pojoToJson(
                          Map.of(
                              "error",
                              String.format("Authorization error: %s", ex.getMessage()),
                              "statusCode",
                              403)))))
          .isError(true)
          .build();
    } catch (Exception ex) {
      LOG.error("Error executing tool '{}': {}", toolName, ex.getMessage(), ex);
      return McpSchema.CallToolResult.builder()
          .content(
              List.of(
                  new McpSchema.TextContent(
                      JsonUtils.pojoToJson(
                          Map.of(
                              "error",
                              String.format("Error executing tool: %s", ex.getMessage()),
                              "statusCode",
                              500)))))
          .isError(true)
          .build();
    }
  }
}