Skip to content

Commit dff04ac

Browse files
labkey-adamlabkey-adam
committed
Simple validation/guidance for missing parameters. Fix getSourceForSavedQuery().
1 parent f21d9fd commit dff04ac

4 files changed

Lines changed: 141 additions & 10 deletions

File tree

‎api/src/org/labkey/api/mcp/McpService.java‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,14 @@
44
import io.modelcontextprotocol.server.McpServerFeatures;
55
import jakarta.servlet.http.HttpSession;
66
import org.jetbrains.annotations.NotNull;
7+
import org.jetbrains.annotations.Nullable;
78
import org.jspecify.annotations.NonNull;
89
import org.labkey.api.data.Container;
910
import org.labkey.api.security.User;
1011
import org.labkey.api.services.ServiceRegistry;
1112
import org.labkey.api.util.HtmlString;
13+
import org.labkey.api.util.StringUtilsLabKey;
14+
import org.labkey.api.view.NotFoundException;
1215
import org.labkey.api.writer.ContainerUser;
1316
import org.springframework.ai.chat.client.ChatClient;
1417
import org.springframework.ai.chat.model.ToolContext;
@@ -19,7 +22,9 @@
1922
import org.springframework.ai.vectorstore.VectorStore;
2023

2124
import java.util.Arrays;
25+
import java.util.HashMap;
2226
import java.util.List;
27+
import java.util.Map;
2328
import java.util.function.Supplier;
2429

2530
/**
@@ -55,11 +60,47 @@ default void incrementResourceRequestCount(String resource)
5560
{
5661
get().incrementResourceRequestCount(resource);
5762
}
63+
64+
// These methods throw a NotFoundException listing all missing parameters. Apparently, even though parameters
65+
// are marked as required, the LLM may not send them or send them with a different name. Best to check them all.
66+
default void validateRequiredParameters(String k1, @Nullable Object v1)
67+
{
68+
validateRequiredParameters(new HashMap<>(){{put(k1, v1);}});
69+
}
70+
71+
default void validateRequiredParameters(String k1, @Nullable Object v1, String k2, @Nullable Object v2)
72+
{
73+
validateRequiredParameters(new HashMap<>(){{put(k1, v1);put(k2, v2);}});
74+
}
75+
76+
default void validateRequiredParameters(String k1, @Nullable Object v1, String k2, @Nullable Object v2, String k3, @Nullable Object v3)
77+
{
78+
validateRequiredParameters(new HashMap<>(){{put(k1, v1);put(k2, v2);put(k3, v3);}});
79+
}
80+
81+
default void validateRequiredParameters(Map<String, Object> parameters)
82+
{
83+
List<String> missing = parameters.entrySet().stream()
84+
.filter(entry -> entry.getValue() == null || entry.getValue().equals(""))
85+
.map(Map.Entry::getKey)
86+
.toList();
87+
88+
if (!missing.isEmpty())
89+
{
90+
if (missing.size() == 1)
91+
throw new NotFoundException(missing.getFirst() + " parameter is required");
92+
else
93+
throw new NotFoundException("The following parameters are required: " + StringUtilsLabKey.joinWithConjunction(missing, "and"));
94+
}
95+
}
5896
}
5997

6098
static @NotNull McpService get()
6199
{
62-
return ServiceRegistry.get().getService(McpService.class);
100+
McpService svc = ServiceRegistry.get().getService(McpService.class);
101+
if (svc == null)
102+
svc = NoopMcpService.get();
103+
return svc;
63104
}
64105

65106
static void setInstance(McpService service)

‎api/src/org/labkey/api/mcp/NoopMcpService.java‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
package org.labkey.api.mcp;
2+
3+
import io.modelcontextprotocol.server.McpServerFeatures;
4+
import jakarta.servlet.http.HttpSession;
5+
import org.jetbrains.annotations.NotNull;
6+
import org.jspecify.annotations.NonNull;
7+
import org.labkey.api.data.Container;
8+
import org.springframework.ai.chat.client.ChatClient;
9+
import org.springframework.ai.chat.model.ToolContext;
10+
import org.springframework.ai.tool.ToolCallback;
11+
import org.springframework.ai.vectorstore.VectorStore;
12+
13+
import java.util.List;
14+
import java.util.function.Supplier;
15+
16+
class NoopMcpService implements McpService
17+
{
18+
private static final McpService INSTANCE = new NoopMcpService();
19+
20+
static McpService get()
21+
{
22+
return INSTANCE;
23+
}
24+
25+
@Override
26+
public boolean isReady()
27+
{
28+
return false;
29+
}
30+
31+
@Override
32+
public void registerTools(@NotNull List<ToolCallback> tools, McpImpl mcp)
33+
{
34+
35+
}
36+
37+
@Override
38+
public void registerPrompts(@NotNull List<McpServerFeatures.SyncPromptSpecification> prompts)
39+
{
40+
41+
}
42+
43+
@Override
44+
public void registerResources(@NotNull List<McpServerFeatures.SyncResourceSpecification> resources)
45+
{
46+
47+
}
48+
49+
@Override
50+
public ToolCallback @NonNull [] getToolCallbacks()
51+
{
52+
return new ToolCallback[0];
53+
}
54+
55+
@Override
56+
public void saveSessionContainer(ToolContext context, Container container)
57+
{
58+
}
59+
60+
@Override
61+
public void incrementResourceRequestCount(String resource)
62+
{
63+
}
64+
65+
@Override
66+
public ChatClient getChat(HttpSession session, String agentName, Supplier<String> systemPromptSupplier, boolean createIfNotExists)
67+
{
68+
return null;
69+
}
70+
71+
@Override
72+
public void close(HttpSession session, ChatClient chat)
73+
{
74+
}
75+
76+
@Override
77+
public MessageResponse sendMessage(ChatClient chat, String message)
78+
{
79+
return null;
80+
}
81+
82+
@Override
83+
public VectorStore getVectorStore()
84+
{
85+
return null;
86+
}
87+
}

‎core/src/org/labkey/core/CoreMcp.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ String setContainer(ToolContext context, @ToolParam(description = "Container pat
9595

9696
if (containerPath == null)
9797
{
98-
message = "Container path was null. Please enter a valid container path. Try using listContainers to see them.";
98+
message = "Container path was null. Please provide a valid containerPath parameter. Try using the listContainers tool to see them.";
9999
}
100100
else
101101
{

‎query/src/org/labkey/query/controllers/QueryMcp.java‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -59,17 +59,19 @@ public McpSchema.ReadResourceResult getLabKeySQLDocumentation() throws IOExcepti
5959

6060
@Tool(description = "Provide column metadata for a sql table. This tool will also return SQL source for saved queries.")
6161
@RequiresPermission(ReadPermission.class)
62-
String listColumns(ToolContext toolContext, @ToolParam(description = "Fully qualified table name as it would appear in SQL e.g. \"schema\".\"table\"") String fullQuotedTableName)
62+
String listColumns(ToolContext toolContext, @ToolParam(description = "Fully qualified table name as it would appear in SQL e.g. \"schema\".\"table\"") String tableName)
6363
{
64-
var json = _listColumns(fullQuotedTableName, toolContext);
64+
validateRequiredParameters("tableName", tableName);
65+
var json = _listColumns(tableName, toolContext);
6566
return json.toString();
6667
}
6768

6869
@Tool(description = "Provide list of tables within the provided schema.")
6970
@RequiresPermission(ReadPermission.class)
70-
String listTables(ToolContext toolContext, @ToolParam(description = "Fully qualified schema name as it would appear in SQL e.g. \"schema\"") String quotedSchemaName)
71+
String listTables(ToolContext toolContext, @ToolParam(description = "Fully qualified schema name as it would appear in SQL e.g. \"schema\"") String schemaName)
7172
{
72-
var json = _listTables(quotedSchemaName, getContext(toolContext));
73+
validateRequiredParameters("schemaName", schemaName);
74+
var json = _listTables(schemaName, getContext(toolContext));
7375
return json.toString();
7476
}
7577

@@ -94,16 +96,17 @@ String listSchemas(ToolContext toolContext)
9496

9597
@Tool(description = "Provide the SQL source for a saved query.")
9698
@RequiresPermission(ReadPermission.class)
97-
String getSourceForSavedQuery(ToolContext toolContext, @ToolParam(description = "Fully qualified query name as it would appear in SQL e.g. \"schema\".\"table or query\"") String fullQuotedTableName)
99+
String getSourceForSavedQuery(ToolContext toolContext, @ToolParam(description = "Fully qualified query name as it would appear in SQL e.g. \"schema\".\"table or query\"") String tableName)
98100
{
99-
var json = _listTables(fullQuotedTableName, getContext(toolContext));
101+
validateRequiredParameters("tableName", tableName);
102+
var json = _listColumns(tableName, toolContext);
100103
if (json.has("sql"))
101104
return "```sql\n" + json.getString("sql") + "\n```\n";
102105
else
103-
return "I could not find the source for " + fullQuotedTableName;
106+
return "I could not find the source for " + tableName;
104107
}
105108

106-
/* For now, list all schemas. CONSIDER support incremental querying. */
109+
/* For now, list all schemas. CONSIDER support incremental querying. */
107110
public static Map<SchemaKey, UserSchema> _listAllSchemas(DefaultSchema root)
108111
{
109112
SimpleSchemaTreeVisitor<Map<SchemaKey,UserSchema>, Void> visitor = new SimpleSchemaTreeVisitor<>(false)

0 commit comments

Comments
 (0)