elasticsearch_server.py

"""
MCP server for Elasticsearch-backed event search and semantic event queries.
This server exposes tools for searching, filtering, and semantically matching events stored in Elasticsearch.
"""

from fastmcp import FastMCP
from elasticsearch import Elasticsearch
from config import ELASTIC_ENDPOINT, ELASTIC_API_KEY, EVENT_INDEX

import logging

logging.basicConfig(level=logging.INFO)

es = Elasticsearch(hosts=ELASTIC_ENDPOINT, api_key=ELASTIC_API_KEY)


mcp = FastMCP(
    name="Elasticsearch Server",
)


@mcp.tool(
    name="search_events",
    description="Search events by time range, topic, title, location, or semantic description. Returns a list of matching event docs.",
)
def search_events(
    start_time: str = None,
    end_time: str = None,
    topic: str = None,
    title: str = None,
    location: str = None,
    description_query: str = None,
    presenting: bool = None,
    top_k: int = 10,
) -> dict:
    """
    Search for events in Elasticsearch using structured and/or semantic filters.
    Search by both description_query and topic for maximum relevance.
    """
    must = []
    if start_time and end_time:
        must.append({"range": {"start_time": {"gte": start_time, "lte": end_time}}})
    if topic:
        must.append({"match": {"topic": topic}})
    if title:
        must.append({"match": {"title": title}})
    if location:
        must.append({"match": {"location": location}})
    if presenting is not None:
        must.append({"term": {"presenting": presenting}})

    if description_query:
        # Prepare the retriever block for hybrid RRF search
        retriever = {
            "rrf": {
                "retrievers": [
                    {
                        "standard": {
                            "query": {
                                "multi_match": {
                                    "query": description_query,
                                    "fields": ["title", "description", "topic"],
                                }
                            }
                        }
                    },
                    {
                        "standard": {
                            "query": {
                                "semantic": {
                                    "field": "description_vector",
                                    "query": description_query,
                                }
                            }
                        }
                    },
                ],
                "rank_window_size": top_k,
                "rank_constant": 20,
            }
        }
        resp = es.search(
            index=EVENT_INDEX,
            retriever=retriever,
            size=top_k,
            _source={"excludes": ["description_vector"]},
        )
    else:
        query = {"bool": {"must": must}} if must else {"match_all": {}}
        resp = es.search(
            index=EVENT_INDEX,
            query=query,
            size=top_k,
            _source={"excludes": ["description_vector"]},
        )
    events = [hit["_source"] for hit in resp["hits"]["hits"]]
    return {"events": events}


@mcp.tool(
    name="create_event",
    description="Create a new event in Elasticsearch. Requires title, description, location, topic, start_time, end_time, url, presenting (bool), and talk_title.",
)
def create_event(
    title: str,
    description: str,
    location: str,
    topic: str,
    start_time: str,
    end_time: str,
    url: str = None,
    presenting: bool = False,
    talk_title: str = None,
) -> dict:
    """
    Create a new event in Elasticsearch. The description is copied to description_vector
    which automatically generates semantic embeddings using the semantic_text field type.
    """

    event_doc = {
        "title": title,
        "description": description,
        "location": location,
        "topic": topic,
        "start_time": start_time,
        "end_time": end_time,
        "url": url,
        "presenting": presenting,
        "talk_title": talk_title,
    }
    resp = es.index(index=EVENT_INDEX, document=event_doc)
    return {"event_id": resp["_id"], "event": event_doc}


if __name__ == "__main__":
    mcp.run()