@mseep/atlas-mcp-server/src/mcp/tools/atlas_unified_search/index.ts

A Model Context Protocol (MCP) server for ATLAS, a Neo4j-powered task management system for LLM Agents - implementing a three-tier architecture (Projects, Tasks, Knowledge) to manage complex workflows.

import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { z } from "zod";
import { ResponseFormat, createResponseFormatEnum, createToolResponse } from "../../../types/mcp.js";
import { createToolExample, createToolMetadata, registerTool } from "../../../types/tool.js";
import { formatUnifiedSearchResponse } from "./responseFormat.js";
import { UnifiedSearchRequestInput } from "./types.js"; // Corrected type import
import { atlasUnifiedSearch } from "./unifiedSearch.js";

export const registerAtlasUnifiedSearchTool = (server: McpServer) => {
  registerTool(
    server,
    "atlas_unified_search",
    "Performs a unified search across specified entity types (node labels) with relevance scoring and flexible filtering options", // Updated description
    {
      property: z.string().optional().describe(
        "Specific property to search within (e.g., name, description, text)"
      ),
      value: z.string().describe(
        "Search term or phrase to find across the knowledge base (required)"
      ),
      entityTypes: z.array(
        z.string() // Allow any string label
      ).optional().describe(
        "Array of entity types (node labels) to include in search (Default: project, task, knowledge if omitted)" // Updated description
      ),
      caseInsensitive: z.boolean().optional().default(true).describe(
        "Boolean flag to ignore case sensitivity when searching for better recall (Default: true)"
      ),
      fuzzy: z.boolean().optional().default(false).describe(
        "Boolean flag to enable approximate matching for typos, spelling variations, and similar terms (Default: false)"
      ),
      taskType: z.string().optional().describe(
        "Optional filter by project/task classification type for more targeted results (applies only if searching Project or Task types)" // Clarified description
      ),
      page: z.number().optional().describe(
        "Page number for paginated results (Default: 1)"
      ),
      limit: z.number().optional().describe(
        "Number of results per page, maximum 100 (Default: 20)"
      ),
      responseFormat: createResponseFormatEnum().optional().default(ResponseFormat.FORMATTED).describe(
        "Desired response format: 'formatted' (default string) or 'json' (raw object)"
      ),
    },
    async (input, context) => {
      // Process unified search request
      const validatedInput = input as unknown as UnifiedSearchRequestInput & { responseFormat?: ResponseFormat };
      
      // Provide default entityTypes if not specified
      const searchInputWithDefaults = {
        ...validatedInput,
        entityTypes: validatedInput.entityTypes && validatedInput.entityTypes.length > 0 
          ? validatedInput.entityTypes 
          : ['project', 'task', 'knowledge'] // Default if empty or undefined
      };

      const result = await atlasUnifiedSearch(searchInputWithDefaults, context);

      // Conditionally format response
      if (validatedInput.responseFormat === ResponseFormat.JSON) {
        return createToolResponse(JSON.stringify(result, null, 2));
      } else {
        // Return the result using the formatter for rich display
        return formatUnifiedSearchResponse(result, false);
      }
    },
    createToolMetadata({
      examples: [
        createToolExample(
          {
            value: "authentication",
            entityTypes: ["project", "task"], // Example still uses specific types
            fuzzy: true
          },
          `{
            "results": [
              {
                "id": "task_auth123",
                "type": "task", // Example still uses specific types
                "entityType": "implementation",
                "title": "Implement OAuth Authentication",
                "description": "Create secure authentication system using OAuth 2.0 protocol",
                "matchedProperty": "title",
                "matchedValue": "Implement OAuth Authentication",
                "createdAt": "2025-03-15T10:22:44.123Z",
                "updatedAt": "2025-03-15T10:22:44.123Z",
                "projectId": "proj_backend42",
                "projectName": "API Platform Modernization",
                "score": 9.5
              },
              {
                "id": "proj_auth456",
                "type": "project", // Example still uses specific types
                "entityType": "security",
                "title": "Authentication Microservice",
                "description": "Build standalone authentication microservice with JWT, refresh tokens and multi-factor support",
                "matchedProperty": "name",
                "matchedValue": "Authentication Microservice",
                "createdAt": "2025-03-10T08:30:12.456Z",
                "updatedAt": "2025-03-10T08:30:12.456Z",
                "score": 10
              }
              // ... potentially other types if they match
            ],
            "total": 2, // Example total might change if other types match
            "page": 1,
            "limit": 20,
            "totalPages": 1
          }`,
          "Search for authentication-related projects and tasks with fuzzy matching"
        ),
        createToolExample(
          {
            value: "performance",
            property: "description",
            entityTypes: ["knowledge"] // Example still uses specific types
          },
          `{
            "results": [
              {
                "id": "know_perf123",
                "type": "knowledge", // Example still uses specific types
                "entityType": "technical",
                "title": "React Performance Optimiz...",
                "description": "Techniques for optimizing React component performance including memoization, virtualization, and code splitting",
                "matchedProperty": "text",
                "matchedValue": "Techniques for optimizing React component performance including memoization, virtualization, and code splitting",
                "createdAt": "2025-03-18T14:05:33.789Z",
                "updatedAt": "2025-03-18T14:05:33.789Z",
                "projectId": "proj_frontend42",
                "projectName": "Frontend Modernization",
                "score": 8.2
              }
              // ... potentially other types if they match
            ],
            "total": 1, // Example total might change if other types match
            "page": 1,
            "limit": 20,
            "totalPages": 1
          }`,
          "Search knowledge items containing 'performance' in the description"
        ),
        createToolExample(
          {
            value: "api"
            // No entityTypes specified, defaults to project, task, knowledge
          },
          `{
            "results": [
              {
                "id": "proj_api789",
                "type": "project", // Example still uses specific types
                "entityType": "integration",
                "title": "API Gateway Implementation",
                "description": "Create centralized API gateway for service integration with rate limiting, monitoring and authentication",
                "matchedProperty": "name",
                "matchedValue": "API Gateway Implementation",
                "createdAt": "2025-03-01T09:45:22.321Z",
                "updatedAt": "2025-03-05T15:12:44.456Z",
                "score": 10
              },
              {
                "id": "task_api456",
                "type": "task", // Example still uses specific types
                "entityType": "development",
                "title": "Document REST API Endpoints",
                "description": "Create comprehensive documentation for all REST API endpoints using OpenAPI specification",
                "matchedProperty": "title",
                "matchedValue": "Document REST API Endpoints",
                "createdAt": "2025-03-08T11:20:15.654Z",
                "updatedAt": "2025-03-08T11:20:15.654Z",
                "projectId": "proj_api789",
                "projectName": "API Gateway Implementation",
                "score": 9.8
              },
              {
                "id": "know_api321",
                "type": "knowledge", // Example still uses specific types
                "entityType": "technical",
                "title": "API Design Best Practices...",
                "description": "Best practices for RESTful API design including versioning, error handling, and resource naming conventions",
                "matchedProperty": "text",
                "matchedValue": "Best practices for RESTful API design including versioning, error handling, and resource naming conventions",
                "createdAt": "2025-03-12T16:30:45.987Z",
                "updatedAt": "2025-03-12T16:30:45.987Z",
                "projectId": "proj_api789",
                "projectName": "API Gateway Implementation",
                "score": 8.5
              }
              // ... potentially other types if they match
            ],
            "total": 8, // Example total might change if other types match
            "page": 1,
            "limit": 10,
            "totalPages": 1
          }`,
          "Search for 'api' across default entity types (project, task, knowledge) with pagination"
        )
      ],
      requiredPermission: "search:read",
      returnSchema: z.object({
        results: z.array(z.object({
          id: z.string().describe("Unique identifier"),
          type: z.string().describe("Entity type (node label)"), // Allow any string
          entityType: z.string().optional().describe("Specific classification of the entity (e.g., taskType, domain)"), // Made optional as it might not apply to all types
          title: z.string().describe("Entity title or name (might be generated for some types)"),
          description: z.string().optional().describe("Entity description text (might be primary text for some types)"), // Made optional
          matchedProperty: z.string().describe("Property where the match was found"),
          matchedValue: z.string().describe("Value containing the match (potentially truncated)"),
          createdAt: z.string().optional().describe("Creation timestamp (if available)"), // Made optional
          updatedAt: z.string().optional().describe("Last update timestamp (if available)"), // Made optional
          projectId: z.string().optional().describe("Associated Project ID (if applicable)"),
          projectName: z.string().optional().describe("Associated Project name (if applicable)"),
          score: z.number().describe("Relevance score")
        })),
        total: z.number().int().describe("Total number of matching results"),
        page: z.number().int().describe("Current page number"),
        limit: z.number().int().describe("Results per page"),
        totalPages: z.number().int().describe("Total number of pages")
      }),
      rateLimit: {
        windowMs: 60 * 1000, // 1 minute
        maxRequests: 20 // 20 requests per minute
      }
    })
  );
};