diff --git a/api/src/org/labkey/api/mcp/McpService.java b/api/src/org/labkey/api/mcp/McpService.java index 6ad04872537..85d080e2b27 100644 --- a/api/src/org/labkey/api/mcp/McpService.java +++ b/api/src/org/labkey/api/mcp/McpService.java @@ -183,4 +183,6 @@ default List sendMessageEx(ChatClient chat, String message) * CONSIDER: Is it possible to implement VectorStoreRetriever wrapper for SearchService??? */ VectorStore getVectorStore(); + + void saveVectorStore(); } diff --git a/api/src/org/labkey/api/mcp/NoopMcpService.java b/api/src/org/labkey/api/mcp/NoopMcpService.java index f6f63534ce0..d2d5fa37d3e 100644 --- a/api/src/org/labkey/api/mcp/NoopMcpService.java +++ b/api/src/org/labkey/api/mcp/NoopMcpService.java @@ -84,4 +84,9 @@ public VectorStore getVectorStore() { return null; } + + @Override + public void saveVectorStore() + { + } } diff --git a/core/src/org/labkey/core/CoreMcp.java b/core/src/org/labkey/core/CoreMcp.java index fa18a2a55c1..ceff5a4df00 100644 --- a/core/src/org/labkey/core/CoreMcp.java +++ b/core/src/org/labkey/core/CoreMcp.java @@ -4,11 +4,13 @@ import io.modelcontextprotocol.spec.McpSchema.ReadResourceResult; import org.apache.commons.io.IOUtils; import org.apache.commons.lang3.StringUtils; +import org.json.JSONArray; import org.json.JSONObject; import org.labkey.api.collections.LabKeyCollectors; import org.labkey.api.data.Container; import org.labkey.api.data.ContainerManager; import org.labkey.api.mcp.McpService; +import org.labkey.api.module.ModuleLoader; import org.labkey.api.security.RequiresNoPermission; import org.labkey.api.security.RequiresPermission; import org.labkey.api.security.User; @@ -130,6 +132,27 @@ String setContainer(ToolContext context, @ToolParam(description = "Container pat return message; } + + // TODO replace/augment with available feature list + @Tool(description = "List the modules installed on this server, this may be useful in inferring the available funcitonality. For instance, " + + "the presence of the `premium` module implies the availability of premium featues.") + @RequiresNoPermission + public String listModules(ToolContext context) + { + JSONArray modules = new JSONArray(); + ModuleLoader.getInstance().getModules().stream() + .map(module -> { + JSONObject obj = new JSONObject(); + obj.put("name", module.getName()); + if (StringUtils.isNotEmpty(module.getLabel())) + obj.put("label", module.getLabel()); + return obj; + }) + .forEach(modules::put); + return new JSONObject(Map.of("modules",modules)).toString(); + } + + @McpResource( uri = "resource://org/labkey/core/FileBasedModules.md", mimeType = "application/markdown", diff --git a/devtools/src/org/labkey/devtools/DevtoolsModule.java b/devtools/src/org/labkey/devtools/DevtoolsModule.java index d391f62baa0..178b254ca33 100644 --- a/devtools/src/org/labkey/devtools/DevtoolsModule.java +++ b/devtools/src/org/labkey/devtools/DevtoolsModule.java @@ -18,6 +18,7 @@ import org.jetbrains.annotations.NotNull; import org.labkey.api.exp.property.Domain; +import org.labkey.api.mcp.McpService; import org.labkey.api.module.CodeOnlyModule; import org.labkey.api.module.ModuleContext; import org.labkey.api.security.AuthenticationManager; @@ -71,6 +72,7 @@ protected void init() @Override public void doStartup(ModuleContext moduleContext) { + McpService.get().register(new TestController.DocumentationMCP()); } @Override diff --git a/devtools/src/org/labkey/devtools/TestController.java b/devtools/src/org/labkey/devtools/TestController.java index 5c8cf796240..b5e808a2a38 100644 --- a/devtools/src/org/labkey/devtools/TestController.java +++ b/devtools/src/org/labkey/devtools/TestController.java @@ -17,8 +17,12 @@ package org.labkey.devtools; import jakarta.servlet.http.HttpServletResponse; +import org.apache.commons.io.IOUtils; +import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.jetbrains.annotations.NotNull; +import org.json.JSONArray; +import org.json.JSONObject; import org.labkey.api.action.ApiResponse; import org.labkey.api.action.ApiSimpleResponse; import org.labkey.api.action.ConfirmAction; @@ -28,8 +32,16 @@ import org.labkey.api.action.SimpleResponse; import org.labkey.api.action.SimpleViewAction; import org.labkey.api.action.SpringActionController; +import org.labkey.api.announcements.CommSchema; +import org.labkey.api.collections.CaseInsensitiveHashMap; +import org.labkey.api.collections.LabKeyCollectors; import org.labkey.api.data.Container; import org.labkey.api.data.ContainerManager; +import org.labkey.api.data.SQLFragment; +import org.labkey.api.data.SimpleFilter; +import org.labkey.api.data.SqlSelector; +import org.labkey.api.data.TableInfo; +import org.labkey.api.data.TableSelector; import org.labkey.api.mcp.AbstractAgentAction; import org.labkey.api.mcp.McpService; import org.labkey.api.security.CSRF; @@ -52,6 +64,7 @@ import org.labkey.api.util.HtmlString; import org.labkey.api.util.HtmlStringBuilder; import org.labkey.api.util.PageFlowUtil; +import org.labkey.api.util.Path; import org.labkey.api.util.URLHelper; import org.labkey.api.view.ActionURL; import org.labkey.api.view.HtmlView; @@ -64,7 +77,11 @@ import org.labkey.api.view.template.ClientDependency; import org.labkey.api.view.template.PageConfig; import org.labkey.api.wiki.WikiService; +import org.springframework.ai.chat.model.ToolContext; import org.springframework.ai.document.Document; +import org.springframework.ai.tool.annotation.Tool; +import org.springframework.ai.tool.annotation.ToolParam; +import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.SimpleVectorStore; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.dao.PessimisticLockingFailureException; @@ -76,14 +93,16 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Collection; import java.util.LinkedHashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; -import java.util.stream.Gatherers; +import static org.apache.commons.lang3.StringUtils.isNotBlank; import static org.labkey.api.util.DOM.Attribute.name; import static org.labkey.api.util.DOM.Attribute.src; import static org.labkey.api.util.DOM.Attribute.style; @@ -1368,7 +1387,7 @@ public boolean handlePost(Object o, BindException errors) throw new NotFoundException(); VectorStore vs = McpService.get().getVectorStore(); if (null == vs) - throw new NotFoundException("/Documentation project was not found"); + throw new NotFoundException("VectorStore not enabled."); ActionURL wikiBase = new ActionURL("wiki","page",documentsContainer); @@ -1382,19 +1401,26 @@ public boolean handlePost(Object o, BindException errors) count.incrementAndGet(); var metadata = Map.of( "Content-Type", "text/html", - "filename", wiki.name() + ".html", + "filename", wiki.name() + ".html", // CONSIDER add path information "title", (Object)wiki.title(), "source", wikiBase.clone().addParameter("name",wiki.name()).getURIString() ); return new Document(wiki.entityId(), wiki.html().toString(), metadata); }) - .gather(Gatherers.windowFixed(50)) - .forEach(vs); + .forEach(d -> { + try + { + vs.accept(List.of(d)); + } + catch (IllegalArgumentException x) + { + LogManager.getLogger(TestController.class).info(d.getMetadata().get("filename"),x); + } + }); - var db = FileUtil.getTempDirectoryFileLike().resolveChild("VectorStore.database"); try { - ((SimpleVectorStore)vs).save(db.toNioPathForRead().toFile()); + McpService.get().saveVectorStore(); return true; } catch (Exception x) @@ -1404,4 +1430,149 @@ public boolean handlePost(Object o, BindException errors) } } } + + public static class DocumentationMCP implements McpService.McpImpl + { + static JSONObject full_index = null; + + static + { + try + { + full_index = new JSONObject(IOUtils.resourceToString("org/labkey/devtools/FULL_INDEX.json", null, DevtoolsModule.class.getClassLoader())); + } + catch(Exception x) + { + } + } + + + @Tool(description = "List of available documents from the LabKey user and administration manuals.") + @RequiresNoPermission + String listDocuments(ToolContext toolContext, + @ToolParam(description = "Index to start listing for paginatation (staring at 0)") Integer start, + @ToolParam(description = "Count of listings to return for pagination") Integer count) + { + Container documentsContainer = ContainerManager.getForPath("/Documentation"); + if (null == documentsContainer) + return new JSONObject(Map.of("error","There is no /Documentation project on this server")).toString(); + + if (null == full_index) + { + // CONSIDER include hierarchy or paths + // TODO WikiService doesn't expose this, just do a query for now (even though this info is cached) + TableInfo currentWikiVersions = CommSchema.getInstance().getSchema().getTable("CurrentWikiVersions"); + SimpleFilter filter = SimpleFilter.createContainerFilter(documentsContainer); + Collection> rows = new TableSelector(currentWikiVersions, Set.of("Name","Title","RowId","Parent","EntityId"), filter, null).getMapCollection(); + + JSONArray array = new JSONArray(); + for (var row : rows) + { + CaseInsensitiveHashMap copy = new CaseInsensitiveHashMap<>(row); + copy.put("id", String.valueOf(copy.get("EntityId"))); + copy.remove("EntityId"); + array.put(new JSONObject(copy)); + } + var j = new JSONObject(); + j.put("pages", array); + full_index = j; + } + + int index = start instanceof Integer i && i >= 0 ? i : 0; + int num = count instanceof Integer i && i >= 0 ? i : Integer.MAX_VALUE; + + JSONArray pages = full_index.getJSONArray("pages"); + int total = pages.length(); + int end = (int) Math.min((long) index + num, total); + + JSONArray subset = new JSONArray(); + for (int i = index; i < end; i++) + subset.put(pages.get(i)); + + var ret = new JSONObject(); + ret.put("total", total); + ret.put("start", index); + ret.put("count", subset.length()); + ret.put("pages", subset); + return ret.toString(); + } + + @Tool(description = "Return the entire document from the LabKey documentation using the `id` as returned by `searchDocumentation`.") + @RequiresNoPermission + String retrieveDocument( + ToolContext context, + @ToolParam(description = "Id of the document to return") String id) + { + WikiService service = Objects.requireNonNull(WikiService.get()); + Container documentsContainer = ContainerManager.getForPath("/Documentation"); + if (null == documentsContainer) + return new JSONObject(Map.of("error","There is not /Documentation project on this server")).toString(); + + ActionURL wikiBase = new ActionURL("wiki","page",documentsContainer); + var sql = new SQLFragment("SELECT Name FROM ").append(CommSchema.getInstance().getTableInfoPages(), "p").append(" WHERE EntityId = ").appendValue(id); + var name = new SqlSelector(CommSchema.getInstance().getSchema(), sql).getObject(String.class); + var wiki = service.getRenderedWiki(documentsContainer, name); + if (null == wiki) + throw new NotFoundException(); + + var ret = new JSONObject(); + ret.put("Content-Type", "text/html"); + ret.put("filename", wiki.name() + ".html"); + ret.put("id", wiki.entityId()); + ret.put("title", wiki.title()); + ret.put("source", wikiBase.clone().addParameter("name",wiki.name()).getURIString()); + ret.put("contents", wiki.html().toString()); + return ret.toString(); + } + + @Tool(description = "Search the LabKey documentation for documents semantically similar to a natural language query. " + + "Returns matching documents with their content, metadata (title, source URL, content type), and similarity scores.") + @RequiresNoPermission + String searchDocumentation( + ToolContext context, + @ToolParam(description = "Natural language search query describing what you're looking for") String query, + @ToolParam(required = false, description = "Maximum number of results to return, defaults to 5") String topK) + { + VectorStore vs = McpService.get().getVectorStore(); + if (vs == null) + throw new IllegalStateException("Vector store is not available. An embedding model may not be configured."); + + int k = 5; + if (isNotBlank(topK)) + { + try { k = Math.clamp(Integer.parseInt(topK), 1, 20); } + catch (NumberFormatException ignored) {} + } + + SearchRequest request = SearchRequest.builder() + .query(query) + .topK(k) + .build(); + + List results = vs.similaritySearch(request); + + var docs = results.stream() + .map(doc -> { + var obj = new JSONObject(); + obj.put("id", doc.getId()); + String text = doc.getText(); + if (text != null && text.length() > 2000) + text = text.substring(0, 2000) + "..."; + obj.put("content", text); + obj.put("metadata", new JSONObject(doc.getMetadata())); + if (doc.getScore() != null) + obj.put("score", doc.getScore()); + return obj; + }) + .collect(LabKeyCollectors.toJSONArray()); + + var ret = new JSONObject(Map.of( + "query", query, + "resultCount", results.size(), + "results", docs + )); +// LogManager.getLogger(TestController.class).info("Search: " + query + "\nResult: " +ret); + return ret.toString(); + } + } }