add mcp
This commit is contained in:
@@ -0,0 +1,262 @@
|
||||
#!/usr/bin/env python3
|
||||
"""MCP server (stdio) exposing mm_search.py search capabilities.
|
||||
|
||||
Cursor / Claude Desktop config example (~/.cursor/mcp.json):
|
||||
|
||||
{
|
||||
"mcpServers": {
|
||||
"mattermost-search": {
|
||||
"command": "/path/to/mattermost-dl/.venv/bin/python",
|
||||
"args": ["/path/to/mattermost-dl/mm_search_mcp.py"],
|
||||
"cwd": "/path/to/mattermost-dl",
|
||||
"env": {
|
||||
"MM_SEARCH_ROOT": "results",
|
||||
"MM_SEARCH_DB": "mattermost_index.sqlite"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Requires: pip install fastmcp
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
import os
|
||||
from typing import Any, Literal, Optional
|
||||
|
||||
from fastmcp import FastMCP
|
||||
|
||||
import mm_search as mm_search_lib
|
||||
|
||||
mcp = FastMCP("mattermost-search")
|
||||
|
||||
SearchMode = Literal["message", "file_name", "file_content"]
|
||||
|
||||
|
||||
def _env_root() -> str:
|
||||
return os.environ.get("MM_SEARCH_ROOT", "results")
|
||||
|
||||
|
||||
def _env_db() -> str:
|
||||
return os.environ.get("MM_SEARCH_DB", "mattermost_index.sqlite")
|
||||
|
||||
|
||||
def _parse_date(value: str) -> dt.datetime:
|
||||
try:
|
||||
return mm_search_lib.parse_iso_date(value)
|
||||
except Exception as exc:
|
||||
raise ValueError(f"invalid date '{value}': {exc}") from exc
|
||||
|
||||
|
||||
def _collect_matches(
|
||||
*,
|
||||
query: str,
|
||||
root: str,
|
||||
user: Optional[str],
|
||||
channel: Optional[str],
|
||||
team: Optional[str],
|
||||
after: Optional[dt.datetime],
|
||||
before: Optional[dt.datetime],
|
||||
regex: bool,
|
||||
search_mode: SearchMode,
|
||||
ignore_case: Optional[bool],
|
||||
limit: Optional[int],
|
||||
db: str,
|
||||
no_index: bool,
|
||||
) -> list[dict[str, Any]]:
|
||||
if search_mode != "message" and regex:
|
||||
raise ValueError("search_mode file_name/file_content do not support regex=True")
|
||||
|
||||
effective_ignore_case = ignore_case if ignore_case is not None else (not regex)
|
||||
matcher = mm_search_lib.make_matcher(query, regex, effective_ignore_case)
|
||||
use_index = (not no_index) and os.path.exists(db)
|
||||
|
||||
if search_mode in ("file_name", "file_content"):
|
||||
if not use_index:
|
||||
raise ValueError(
|
||||
"Attachment search (file_name/file_content) requires sqlite index. "
|
||||
"Create it with mm_index.py and do not set no_index=True."
|
||||
)
|
||||
mode = "name" if search_mode == "file_name" else "content"
|
||||
matches_iter = mm_search_lib.search_attachments_index(
|
||||
db_path=db,
|
||||
query=query,
|
||||
user=user,
|
||||
channel=channel,
|
||||
team=team,
|
||||
after=after,
|
||||
before=before,
|
||||
limit=limit,
|
||||
mode=mode,
|
||||
)
|
||||
output_limit = None
|
||||
elif use_index and not regex:
|
||||
matches_iter = mm_search_lib.search_index(
|
||||
db_path=db,
|
||||
query=query,
|
||||
user=user,
|
||||
channel=channel,
|
||||
team=team,
|
||||
after=after,
|
||||
before=before,
|
||||
limit=limit,
|
||||
ignore_case=effective_ignore_case,
|
||||
)
|
||||
output_limit = None
|
||||
else:
|
||||
files = list(mm_search_lib.iter_result_files(root))
|
||||
matches_iter = (
|
||||
m
|
||||
for path in files
|
||||
for m in mm_search_lib.search_file(
|
||||
path=path,
|
||||
matcher=matcher,
|
||||
user=user,
|
||||
channel=channel,
|
||||
team=team,
|
||||
after=after,
|
||||
before=before,
|
||||
)
|
||||
)
|
||||
output_limit = limit
|
||||
|
||||
out: list[dict[str, Any]] = []
|
||||
count = 0
|
||||
for match in matches_iter:
|
||||
if output_limit is not None and count >= output_limit:
|
||||
break
|
||||
out.append(match)
|
||||
count += 1
|
||||
return out
|
||||
|
||||
|
||||
def _format_human(matches: list[dict[str, Any]], show_file: bool) -> str:
|
||||
lines: list[str] = []
|
||||
for m in matches:
|
||||
p = m["post"]
|
||||
created = p.get("created", "")
|
||||
username = p.get("username", "")
|
||||
msg = (p.get("message") or "").replace("\n", "\\n")
|
||||
attachment = m.get("attachment")
|
||||
attachment_suffix = ""
|
||||
if isinstance(attachment, dict) and attachment.get("filename"):
|
||||
attachment_suffix = f" [file: {attachment['filename']}]"
|
||||
if show_file:
|
||||
file_rel = os.path.relpath(m["file"])
|
||||
idx = p.get("idx", "?")
|
||||
prefix = f"{file_rel}:{idx}: "
|
||||
else:
|
||||
prefix = ""
|
||||
channel = m.get("channel") or "?"
|
||||
team = m.get("team") or ""
|
||||
team_prefix = f"[{team}]" if team else ""
|
||||
lines.append(f"{prefix}{team_prefix}[{channel}] {created} {username}: {msg}{attachment_suffix}")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
def mm_search(
|
||||
query: str,
|
||||
root: Optional[str] = None,
|
||||
user: Optional[str] = None,
|
||||
channel: Optional[str] = None,
|
||||
team: Optional[str] = None,
|
||||
after: Optional[str] = None,
|
||||
before: Optional[str] = None,
|
||||
regex: bool = False,
|
||||
search_mode: SearchMode = "message",
|
||||
ignore_case: Optional[bool] = None,
|
||||
limit: Optional[int] = None,
|
||||
db: Optional[str] = None,
|
||||
no_index: bool = False,
|
||||
output: Literal["json", "human"] = "json",
|
||||
show_file: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Search Mattermost export JSON files produced by mattermost-dl.py.
|
||||
|
||||
Mirrors mm_search.py CLI. Returns structured matches or human-readable lines.
|
||||
|
||||
Args:
|
||||
query: Search string or regex (when regex=True).
|
||||
root: Root directory with export results (default: MM_SEARCH_ROOT or "results").
|
||||
user: Filter by username (exact match).
|
||||
channel: Filter by channel display_name or name (exact match).
|
||||
team: Filter by team name (exact match).
|
||||
after: Only messages on/after this date (YYYY-MM-DD or ISO).
|
||||
before: Only messages on/before this date (YYYY-MM-DD or ISO).
|
||||
regex: Treat query as regular expression (message scan only).
|
||||
search_mode: "message" (default), "file_name", or "file_content".
|
||||
ignore_case: Case-insensitive search. Default: True for substring, False for regex.
|
||||
limit: Maximum number of matches to return.
|
||||
db: SQLite index path (default: MM_SEARCH_DB or "mattermost_index.sqlite").
|
||||
no_index: Do not use sqlite index even if DB exists.
|
||||
output: "json" returns match objects; "human" returns CLI-style text lines.
|
||||
show_file: Include source filename and per-file index in human output.
|
||||
"""
|
||||
resolved_root = root or _env_root()
|
||||
resolved_db = db or _env_db()
|
||||
parsed_after = _parse_date(after) if after else None
|
||||
parsed_before = _parse_date(before) if before else None
|
||||
|
||||
matches = _collect_matches(
|
||||
query=query,
|
||||
root=resolved_root,
|
||||
user=user,
|
||||
channel=channel,
|
||||
team=team,
|
||||
after=parsed_after,
|
||||
before=parsed_before,
|
||||
regex=regex,
|
||||
search_mode=search_mode,
|
||||
ignore_case=ignore_case,
|
||||
limit=limit,
|
||||
db=resolved_db,
|
||||
no_index=no_index,
|
||||
)
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"count": len(matches),
|
||||
"used_index": (not no_index)
|
||||
and os.path.exists(resolved_db)
|
||||
and (search_mode in ("file_name", "file_content") or not regex),
|
||||
"search_mode": search_mode,
|
||||
"root": resolved_root,
|
||||
"db": resolved_db,
|
||||
}
|
||||
|
||||
if output == "human":
|
||||
result["text"] = _format_human(matches, show_file=show_file)
|
||||
else:
|
||||
result["matches"] = matches
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@mcp.tool
|
||||
def mm_search_index_status(
|
||||
db: Optional[str] = None,
|
||||
root: Optional[str] = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Report whether the sqlite index and export root are available."""
|
||||
resolved_db = db or _env_db()
|
||||
resolved_root = root or _env_root()
|
||||
db_exists = os.path.exists(resolved_db)
|
||||
root_exists = os.path.isdir(resolved_root)
|
||||
|
||||
json_count = 0
|
||||
if root_exists:
|
||||
json_count = sum(1 for _ in mm_search_lib.iter_result_files(resolved_root))
|
||||
|
||||
return {
|
||||
"db": resolved_db,
|
||||
"db_exists": db_exists,
|
||||
"root": resolved_root,
|
||||
"root_exists": root_exists,
|
||||
"json_files": json_count,
|
||||
"index_required_for": ["file_name", "file_content"],
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
mcp.run(transport="stdio")
|
||||
Reference in New Issue
Block a user