DuckDB で構築したハイブリッド検索システムを MCP 経由で Raycast から呼び出す

from mcp.server.fastmcp import FastMCP

import duckdb
import torch
import json
from lindera_py import Segmenter, Tokenizer, load_dictionary
from sentence_transformers import CrossEncoder
from transformers import AutoModel, AutoTokenizer
import pandas as pd

device = "cuda" if torch.cuda.is_available() else "cpu"

v_tokenizer = AutoTokenizer.from_pretrained(
    "pfnet/plamo-embedding-1b", trust_remote_code=True
)
v_model = AutoModel.from_pretrained("pfnet/plamo-embedding-1b", trust_remote_code=True)
v_model = v_model.to(device)

dictionary = load_dictionary("ipadic")
segmenter = Segmenter("normal", dictionary)
tokenizer = Tokenizer(segmenter)

reranker = CrossEncoder(
    "hotchpotch/japanese-bge-reranker-v2-m3-v1", device=device, max_length=512
)

con = duckdb.connect("article_search.duckdb")
con.install_extension("vss")
con.load_extension("vss")
con.install_extension("fts")
con.load_extension("fts")

mcp = FastMCP("arcticle-search")


def ja_tokens(text: str) -> str:
    return " ".join(t.text for t in tokenizer.tokenize(text))


def fts_search(conn, query, k=5):
    q = ja_tokens(query)
    return conn.sql(f"""
        SELECT id, title,
               fts_main_articles.match_bm25(id, '{q}') AS score
        FROM articles
        WHERE score IS NOT NULL
        ORDER BY score DESC
        LIMIT {k}
    """).fetchdf()


def vss_search(conn, query, k=5):
    with torch.inference_mode():
        q_emb = v_model.encode_query(query, v_tokenizer)
    return conn.sql(
        f"""
        SELECT id, title,
               array_cosine_distance(embedding, ?::FLOAT[2048]) AS dist
        FROM articles
        ORDER BY dist ASC
        LIMIT {k}
        """,
        params=[q_emb.cpu().squeeze().numpy().tolist()],
    ).fetchdf()


def hybrid_search(conn, query, k_fts=10, k_vss=10, top_n=5):
    fts_df = fts_search(conn, query, k_fts)
    vss_df = vss_search(conn, query, k_vss)
    pool = pd.concat([fts_df, vss_df]).drop_duplicates("id")
    pairs = [(query, txt) for txt in pool["title"]]
    scores = reranker.predict(pairs)
    pool["score"] = scores
    result_df = pool.sort_values("score", ascending=False).head(top_n)[
        ["id", "title", "score"]
    ]
    return result_df


def hybrid_search_json(conn, query, k_fts=10, k_vss=10, top_n=5):
    df = hybrid_search(conn, query, k_fts, k_vss, top_n)
    # MEMO: 検証時に url を含めていなかったため id から url を作成する
    df["url"] = df["id"].apply(lambda id: f"https://read.readwise.io/read/{id}")
    results = df.to_dict(orient="records")
    return json.dumps(results, ensure_ascii=False, indent=2)


@mcp.tool()
def query_data(query: str) -> str:
    """Query artickes from the database."""
    try:
        hybrid_search_result = hybrid_search_json(con, query, top_n=20)
        return hybrid_search_result
    except Exception as e:
        return f"Error: {str(e)}"


if __name__ == "__main__":
    print("Starting FastMCP server...")
    mcp.run(transport="stdio")