토스페이먼츠 결제 시스템 연동을 돕는 MCP 서버 구현기

"BM25는 정보 검색에 사용되는 랭킹 함수입니다."

["bm25는", "정보", "검색에", "사용되는", "랭킹", "함수입니다"]

const tree = unified().use(remarkParse).parse(markdown);

visit(tree, (node) => {
  if (node.type === "heading" && node.depth <= 2) {
    chunks.push(...)
  }
  ...
});

export function joinShortChunks(chunks: string[], minWords = 30): string[] {
  const result: string[] = [];

  let buffer = "";
  let bufferCount = 0;

  for (const chunk of chunks) {
    const wc = chunk.split(/\s+/).length;
    if (wc < minWords) {
      buffer += (buffer ? "\n\n" : "") + chunk;
      bufferCount += wc;
      continue;
    }

    if (buffer) {
      result.push(buffer.trim());
      buffer = "";
      bufferCount = 0;
    }

    result.push(chunk.trim());
  }

  if (buffer) {
    result.push(buffer.trim());
  }

  return result;
}

export class TossPaymentsDocument {
  private readonly chunks: DocumentChunk[] = [];

  constructor(
    private readonly keywordSet: Set<string>,
    private readonly remoteMarkdownDocument: RemoteMarkdownDocument,
    private readonly _version: string | undefined,
    public readonly id: number
  ) {
    remoteMarkdownDocument.chunks.forEach((chunk, index) => {
      this.chunks.push({
        id: this.id,
        chunkId: this.id * 1000 + index,
        originTitle: remoteMarkdownDocument.metadata.title,
        text: chunk,
        wordCount: chunk.split(/\s+/).length,
      });
    });
  }

  getChunkWithWindow(chunkId: number, windowSize: number): DocumentChunk[] {
    const chunkIndex = this.chunks.findIndex(
      (chunk) => chunk.chunkId === chunkId
    );
    if (chunkIndex === -1) {
      return [];
    }

    const start = Math.max(0, chunkIndex - windowSize);
    const end = Math.min(this.chunks.length, chunkIndex + windowSize + 1);

    return this.chunks.slice(start, end);
  }

  getChunks(): DocumentChunk[] {
    return this.chunks;
  }

  get content(): string {
    return this.remoteMarkdownDocument.markdown;
  }

  get title() {
    return this.remoteMarkdownDocument.metadata.title;
  }

  get version(): string | undefined {
    return this._version;
  }

  get description() {
    return this.remoteMarkdownDocument.metadata.description;
  }

  toString() {
    return this.remoteMarkdownDocument.markdown;
  }

  toJSON() {
    return {
      version: this.version,
      id: this.id,
      title: this.title,
      link: this.remoteMarkdownDocument.link,
      description: this.description,
      keywords: Array.from(this.keywordSet),
    };
  }
}

export class TossPaymentsBM25Calculator {
  private readonly allChunks: DocumentChunk[];
  private readonly totalCount: number;
  private readonly averageDocLength: number;
  private readonly N: number;

  constructor(
    documents: TossPaymentsDocument[],
    private readonly k1: number = 1.2,
    private readonly b: number = 0.75
  ) {
    this.allChunks = documents.flatMap((doc) => doc.getChunks());
    this.totalCount = this.allChunks.reduce(
      (count, doc) => count + doc.wordCount,
      0
    );
    this.averageDocLength = this.totalCount / this.allChunks.length;
    this.N = this.allChunks.length;
  }

  calculate(keywords: string): Result[] {
    const { termFrequencies, docFrequencies } =
      this.calculateFrequencies(keywords);

    const scores = this.calculateScore(termFrequencies, docFrequencies);

    scores.sort((a, b) =>
      b.score !== a.score ? b.score - a.score : b.totalTF - a.totalTF
    );

    return scores.map(({ id, score, chunkId }) => ({ id, chunkId, score }));
  }

  private calculateFrequencies(query: string) {
    const pattern = new RegExp(query, "gi");
    const termFrequencies: Record<number, Record<string, number>> = {};
    const docFrequencies: Record<string, number> = {};

    for (const doc of this.allChunks) {
      const text = doc.text;

      const matches = Array.from(text.matchAll(pattern));
      const termCounts: Record<string, number> = {};

      for (const match of matches) {
        const term = match[0].toLowerCase();
        termCounts[term] = (termCounts[term] || 0) + 1;
      }

      if (Object.keys(termCounts).length > 0) {
        termFrequencies[doc.chunkId] = termCounts;
        for (const term of Object.keys(termCounts)) {
          docFrequencies[term] = (docFrequencies[term] || 0) + 1;
        }
      }
    }

    return { termFrequencies, docFrequencies };
  }

  private calculateScore(
    termFrequencies: Record<number, Record<string, number>>,
    docFrequencies: Record<string, number>
  ) {
    return this.allChunks
      .filter((chunk) => termFrequencies[chunk.chunkId])
      .map((chunk) => {
        const tf = termFrequencies[chunk.chunkId];
        const len = chunk.wordCount;

        const score = Object.keys(tf)
          .map((term) => {
            const df = docFrequencies[term];
            const idf = Math.log((this.N - df + 0.5) / (df + 0.5));
            const numerator = tf[term] * (this.k1 + 1);
            const denominator =
              tf[term] +
              this.k1 * (1 - this.b + this.b * (len / this.averageDocLength));
            return idf * (numerator / denominator);
          })
          .reduce((sum, v) => sum + v, 0);

        const totalTF = Object.values(tf).reduce((sum, v) => sum + v, 0);

        return {
          id: chunk.id,
          chunkId: chunk.chunkId,
          score,
          totalTF,
        };
      });
  }
}

export interface Result {
  id: number;
  chunkId: number;
  score: number;
}

export class TossPaymentDocsRepository {
  private readonly documentV1BM25Calculator: TossPaymentsBM25Calculator;
  private readonly documentV2BM25Calculator: TossPaymentsBM25Calculator;

  constructor(private readonly documents: TossPaymentsDocument[]) {
    const v1Documents = documents.filter(
      (document) => document.version === "v1"
    );

    const v2Documents = documents.filter(
      (document) => document.version === "v2"
    );

    const generalDocuments = documents.filter(
      (document) => document.version == null
    );

    this.documentV1BM25Calculator = new TossPaymentsBM25Calculator(
      v1Documents.concat(generalDocuments)
    );
    this.documentV2BM25Calculator = new TossPaymentsBM25Calculator(
      v2Documents.concat(generalDocuments)
    );
  }

  async findV1DocumentsByKeyword(
    keywords: string[],
    topN: number = 10
  ): Promise<string> {
    const results = this.documentV1BM25Calculator.calculate(
      keywords.map((keyword) => escapeRegExp(keyword.trim())).join("|")
    );

    const docs = results.slice(0, topN);

    return this.toString(docs);
  }

  async findV2DocumentsByKeyword(
    keywords: string[],
    topN: number = 10
  ): Promise<string> {
    const results = this.documentV2BM25Calculator.calculate(
      keywords.map((keyword) => escapeRegExp(keyword.trim())).join("|")
    );

    const docs = results.slice(0, topN);

    return this.toString(docs);
  }

  findOneById(id: number) {
    return this.documents[id];
  }

  private toString(results: Result[]): string {
    const docs = results
      .map((item) => {
        const document = this.findOneById(item.id);
        return document.getChunkWithWindow(item.chunkId, 1);
      })
      .filter((item) => item !== undefined)
      .map((items) => this.normalizeChunks(items));

    return docs.join("\n\n");
  }

  private normalizeChunks(chunks: DocumentChunk[]): string {
    return `## 원본문서 제목 : ${chunks[0].originTitle}\n* 원본문서 ID : ${
      chunks[0].id
    }\n\n${chunks.map((chunk) => chunk.text).join("\n\n")}`;
  }
}

function escapeRegExp(string: string) {
  return string.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); // $&는 일치한 전체 문자열을 의미합니다.
}