feat(clp-mcp-server): Add search_by_kql MCP tool call. (#1436)

Co-authored-by: Lin Zhihao <[email protected]>

Author: 20001020ycx

components/clp-mcp-server/clp_mcp_server/clp_connector.py

@@ -50,10 +50,11 @@ async def submit_query(
         :param begin_ts: The beginning timestamp of the query range.
         :param end_ts: The end timestamp of the query range.
         :return: The ID assigned to the query.
-        :raise ValueError: If ``end_ts`` is smaller than ``begin_ts``.
-        :raise aiomysql.Error: If there is an error connecting to or querying MariaDB.
-        :raise pymongo.errors.PyMongoError: If there is an error interacting with MongoDB.
-        :raise Exception: For any other unexpected errors.
+        :raise: ValueError if `end_ts` is smaller than `begin_ts`.
+        :raise: RuntimeError if it fails to retrieve the ID of the submitted query.
+        :raise: aiomysql.Error if there is an error connecting to or querying MariaDB.
+        :raise: pymongo.errors.PyMongoError if there is an error interacting with MongoDB.
+        :raise: Exception for any other unexpected errors.
         """
         if begin_ts is not None and end_ts is not None and end_ts < begin_ts:
             err_msg = f"end_ts {end_ts} is smaller than begin_ts {begin_ts}."
@@ -123,9 +124,10 @@ async def wait_query_completion(self, query_id: str, timeout: float | None = Non
 
         :param query_id: The ID of the query.
         :param timeout: Maximum time to wait in seconds, or None for no timeout.
-        :raise aiomysql.Error: If there is an error connecting to or querying MariaDB.
-        :raise ValueError: When the query is not found.
-        :raise RuntimeError: When the query fails or is cancelled.
+        :raise: aiomysql.Error if there is an error connecting to or querying MariaDB.
+        :raise: ValueError if the query is not found.
+        :raise: RuntimeError if the query fails or is cancelled.
+        :raise: TimeoutError if the timeout is reached before the query completes.
         """
         waiting_states = {QueryJobStatus.PENDING, QueryJobStatus.RUNNING, QueryJobStatus.CANCELLING}
         error_states = {QueryJobStatus.FAILED, QueryJobStatus.CANCELLED, QueryJobStatus.KILLED}

components/clp-mcp-server/clp_mcp_server/clp_mcp_server.py

@@ -4,8 +4,12 @@
 import logging
 import socket
 import sys
+from pathlib import Path
 
 import click
+from clp_py_utils.clp_config import CLPConfig
+from clp_py_utils.core import read_yaml_config_file
+from pydantic import ValidationError
 
 from .server import create_mcp_server
 
@@ -15,22 +19,32 @@
     "--host", type=str, default="127.0.0.1", help="The server's host address (default: 127.0.0.1)."
 )
 @click.option("--port", type=int, default=8000, help="The server's port number (default: 8000).")
-def main(host: str, port: int) -> None:
+@click.option(
+    "--config-path",
+    type=click.Path(exists=True),
+    default="/etc/clp-config.yml",
+    help="The path to server's configuration file (default: /etc/clp-config.yml).",
+)
+def main(host: str, port: int, config_path: Path) -> int:
     """
     Runs the CLP MCP server with HTTP transport.
 
     :param host: The server's host address (IP address or hostname).
     :param port: The server's port number (1-65535).
+    :param config_path: The path to server's configuration file.
+    :return: Exit code (0 for success, non-zero for failure).
     """
     logging.basicConfig(
         level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
     )
     logger = logging.getLogger(__name__)
 
+    exit_code = 0
+
     # Validate host and port
     if len(host.strip()) == 0:
         logger.error("Host cannot be empty.")
-        sys.exit(1)
+        exit_code = 1
 
     # Validate host format (IP address or resolvable hostname)
     try:
@@ -44,21 +58,32 @@ def main(host: str, port: int) -> None:
                 "Host validation failed: '%s' is not a valid IP address and DNS resolution failed.",
                 host,
             )
-            sys.exit(1)
+            exit_code = 1
 
     max_port = 65535
     if port <= 0 or port > max_port:
         logger.error("Port must be between 1 and %d, got: %d.", max_port, port)
-        sys.exit(1)
+        exit_code = 1
 
     try:
-        mcp = create_mcp_server()
+        clp_config = CLPConfig.model_validate(read_yaml_config_file(config_path))
+    except ValidationError:
+        logger.exception("Configuration validation failed.")
+        exit_code = 1
+    except Exception:
+        logger.exception("Failed to load configuration.")
+        exit_code = 1
+
+    try:
+        mcp = create_mcp_server(clp_config)
         logger.info("Starting CLP MCP Server on %s:%d.", host, port)
         mcp.run(transport="streamable-http", host=host, port=port)
     except Exception:
         logger.exception("Failed to start MCP server.")
-        sys.exit(1)
+        exit_code = 1
+
+    return exit_code
 
 
 if __name__ == "__main__":
-    main()
+    sys.exit(main())

components/clp-mcp-server/clp_mcp_server/server/constants.py

@@ -6,6 +6,8 @@
 # 10 minutes
 SESSION_TTL_SECONDS = 600
 
+TIMESTAMP_NOT_AVAILABLE = "N/A"
+
 SERVER_NAME = "clp-mcp-server"
 
 # fmt: off

components/clp-mcp-server/clp_mcp_server/server/server.py

@@ -2,16 +2,21 @@
 
 from typing import Any
 
+from clp_py_utils.clp_config import CLPConfig
 from fastmcp import Context, FastMCP
 
+from clp_mcp_server.clp_connector import ClpConnector
+
 from . import constants
 from .session_manager import SessionManager
+from .utils import format_query_results, sort_by_timestamp
 
 
-def create_mcp_server() -> FastMCP:
+def create_mcp_server(clp_config: CLPConfig) -> FastMCP:
     """
     Creates and defines API tool calls for the CLP MCP server.
 
+    :param clp_config:
     :return: A configured `FastMCP` instance.
     :raise: Propagates `FastMCP.__init__`'s exceptions.
     :raise: Propagates `FastMCP.tool`'s exceptions.
@@ -20,6 +25,8 @@ def create_mcp_server() -> FastMCP:
 
     session_manager = SessionManager(session_ttl_seconds=constants.SESSION_TTL_SECONDS)
 
+    connector = ClpConnector(clp_config)
+
     @mcp.tool
     async def get_instructions(ctx: Context) -> str:
         """
@@ -67,4 +74,38 @@ def hello_world(name: str = "clp-mcp-server user") -> dict[str, Any]:
             "status": "running",
         }
 
+    @mcp.tool
+    async def search_by_kql(kql_query: str, ctx: Context) -> dict[str, Any]:
+        """
+        Searches log events that match the given Kibana Query Language (KQL) query. The resulting
+        events are ordered by timestamp in descending order (latest to oldest), cached for
+        subsequent pagination, and returned with the first page of results.
+
+        :param kql_query:
+        :param ctx: The `FastMCP` context containing the metadata of the underlying MCP session.
+        :return: A dictionary containing the following key-value pairs on success:
+            - "items": A list of log entries in the requested page.
+            - "num_total_pages": Total number of pages available from the query as an integer.
+            - "num_total_items": Total number of log entries available from the query as an integer.
+            - "num_items_per_page": Number of log entries per page.
+            - "has_next": Whether a page exists after the returned one.
+            - "has_previous": Whether a page exists before the returned one.
+        :return: A dictionary with the following key-value pair on failures:
+            - "Error": An error message describing the failure.
+        """
+        await session_manager.start()
+
+        try:
+            query_id = await connector.submit_query(kql_query)
+            await connector.wait_query_completion(query_id)
+            results = await connector.read_results(query_id)
+        except (ValueError, RuntimeError, TimeoutError) as e:
+            return {"Error": str(e)}
+
+        sorted_results = sort_by_timestamp(results)
+        formatted_results = format_query_results(sorted_results)
+        return session_manager.cache_query_result_and_get_first_page(
+            ctx.session_id, formatted_results
+        )
+
     return mcp

components/clp-mcp-server/clp_mcp_server/server/session_manager.py

@@ -109,7 +109,8 @@ def get_page_data(self, page_index: int) -> dict[str, Any]:
         Retrieves the n-th page of a paginated response with the paging metadata from the previous
         query.
 
-        NOTE: This docstring must be synchronized with `get_nth_page`'s MCP tool call.
+        NOTE: This docstring must be synchronized with MCP tool calls: `get_nth_page` and
+        `search_by_kql`.
 
         :param page_index: Zero-based index, e.g., 0 for the first page.
         :return: A dictionary containing the following key-value pairs on success:

components/clp-mcp-server/clp_mcp_server/server/utils.py

@@ -0,0 +1,80 @@
+"""Utility functions for CLP MCP server."""
+
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+from .constants import TIMESTAMP_NOT_AVAILABLE
+
+logger = logging.getLogger(__name__)
+
+
+def convert_epoch_to_date_string(epoch_ts: int) -> str:
+    """
+    :param epoch_ts: Unix epoch timestamp in milliseconds.
+    :return: ISO 8601 formatted date string with millisecond precision (YYYY-MM-DDTHH:mm:ss.fffZ).
+    :raise: ValueError if `epoch_ts` cannot be converted to a valid date string.
+    """
+    try:
+        epoch_seconds = epoch_ts / 1000.0
+        dt = datetime.fromtimestamp(epoch_seconds, tz=timezone.utc)
+        return dt.isoformat(timespec="milliseconds").replace("+00:00", "Z")
+    except (ValueError, OSError, OverflowError) as e:
+        err_msg = f"Invalid timestamp {epoch_ts}: {e}."
+        raise ValueError(err_msg) from e
+
+
+def format_query_results(query_results: list[dict[str, Any]]) -> list[str]:
+    """
+    Formats the query results. For a log event to be formatted, it must contain the following
+    kv-pairs:
+    - "timestamp": An integer representing the epoch timestamp in milliseconds.
+    - "message": A string representing the log message.
+
+    The message will be formatted as `timestamp: <date string>, message: <message>`:
+
+    :param query_results: A list of dictionaries representing kv-pair log events.
+    :return: A list of strings representing formatted log events.
+    """
+    formatted_log_events = []
+    for obj in query_results:
+        epoch = obj.get("timestamp")
+        timestamp_str = TIMESTAMP_NOT_AVAILABLE
+
+        if isinstance(epoch, int):
+            try:
+                timestamp_str = convert_epoch_to_date_string(epoch)
+            except (TypeError, ValueError) as e:
+                logger.warning("Failed to convert epoch timestamp=%s to date string: %s.", epoch, e)
+
+        message = obj.get("message", "")
+        if not message:
+            logger.warning("Empty message attached to a log event: %s.", obj)
+            continue
+
+        formatted_log_events.append(f"timestamp: {timestamp_str}, message: {message}")
+
+    return formatted_log_events
+
+
+def sort_by_timestamp(query_results: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """
+    Sorts the query results in-place by timestamp in descending order (latest to oldest).
+
+    Note:
+    - Timestamp is expected to be an integer representing the epoch timestamp in milliseconds,
+      stored under the `timestamp` key.
+    - If `timestamp` is missing or not an integer, it is treated as the oldest possible timestamp.
+
+    :param query_results: A list of dictionaries representing kv-pair log events.
+    :return: The input list sorted in-place.
+
+    """
+
+    def _key(log_entry: dict[str, Any]) -> int:
+        ts = log_entry.get("timestamp")
+        return ts if isinstance(ts, int) else -1
+
+    query_results.sort(key=_key, reverse=True)
+
+    return query_results

components/clp-mcp-server/tests/server/__init__.py

@@ -0,0 +1 @@
+"""Tests for server module."""