Dify + Mathpix + LLM で論文 QA システム

Dify + Mathpix + LLM で論文 QA システムを作った話

背景

論文を読むのは時間がかかる作業です。特に、数式が多い論文では記号の定義を追ったり、細かい数式の導出を確認したりするのに多くの時間を費やしてしまいます。自分が知りたい情報を取り出すために、ちまちまスクロールして頑張って目力で頑張って探すのは、時間の浪費以外の何物でもありません。RAGも試しましたが、全然精度の良い回答が返ってこないので、使い物になりません。そこで、PDF を投入するだけで、その内容に関する質問に精密に答えてくれるチャットボットを開発しました。

システム概要

Dify を活用した LLM ベースの QA システムを作りました。Mathpix を使って PDF を OCR 解析し、検索性の高い形式に変換することで、論文に対する高度な質問応答を実現しています。

Dify のフロー図は以下の通りです。シンプルで、特筆すべき点は以下の点くらいです。

• URL or pdf file の分岐
• http リクエストで論文 pdf を markdown 形式に変換
• 論文 markdown を丸ごと LLM のシステムプロンプトに投げる。「丸ごと」がみそで RAG では得られない圧倒的な回答精度を実現できます。

技術構成

1. Dify を用いた LLM ベースの QA チャットボット

Dify は、ユーザーが簡単にカスタマイズ可能な LLM アプリケーションを構築できるプラットフォームです。このシステムでは、Dify のフローを活用し、ユーザーからの質問に対して適切に応答できるように設計しています。

今回は Dify をローカルマシン上で docker compose によってデプロイしました。

2. PDF の処理とキャッシュ

本システムでは、Dify のほかに 同じ物理マシン上に HTTP サーバーを立て、PDF を受け取って Mathpix API に送信する仕組みを導入しました。Dify はキャッシュする仕組みがなかったためです。

このサーバーは

• Mathpix API を活用して OCR 変換を行い、論文を Markdown 形式で保存
• 毎回 Mathpix に API コールをすると時間とコストがかかるため、キャッシュを実装し、既に処理済みの論文は再処理しないように

するものです。

論文を丸ごと markdown 化するので、実は Dify のデフォルトの文字数制限に引っかかることがあります。この場合 dify/docker の .env ファイルの CODE_MAX_STRING_LENGTH を適切に設定しなおしてあげると事なきを得ます。僕は特に理由なく 2000000 にしています。

3. ネットワーク構成とセキュリティ対策

個人のプライベートネットワーク上に構築するのであれば何も考えず、関連するすべてのポートを any で公開するだけでいいので楽ですが、組織内のネットワークの場合適切なファイアウォール設定が必要になります。２回、別の環境で別々の要請に合わせて構築したため知見がたまりました。正直開発の９割の時間はここに費やしたと思います。

IP アドレスでアクセス制限をかける場合

特にグローバルIPが付与されているようなサーバー上で構築する場合、悪用されないように細心の注意が必要です。今回自分が使ったのは主に ufw でファイアーウォールを構成している Ubuntu マシンでしたが、docker compose をすると勝手に穴があけられてしまいます。これはかなり面倒で、苦戦しましたが色々調べた結果 https://github.com/chaifeng/ufw-docker の解決策を使って解決できました。

具体的には次の手順です。

Modify the UFW configuration file /etc/ufw/after.rules and add the following rules at the end of the file:

# BEGIN UFW AND DOCKER
*filter
:ufw-user-forward - [0:0]
:ufw-docker-logging-deny - [0:0]
:DOCKER-USER - [0:0]
-A DOCKER-USER -j ufw-user-forward

-A DOCKER-USER -j RETURN -s 10.0.0.0/8
-A DOCKER-USER -j RETURN -s 172.16.0.0/12
-A DOCKER-USER -j RETURN -s 192.168.0.0/16

-A DOCKER-USER -p udp -m udp --sport 53 --dport 1024:65535 -j RETURN

-A DOCKER-USER -j ufw-docker-logging-deny -p tcp -m tcp --tcp-flags FIN,SYN,RST,ACK SYN -d 192.168.0.0/16
-A DOCKER-USER -j ufw-docker-logging-deny -p tcp -m tcp --tcp-flags FIN,SYN,RST,ACK SYN -d 10.0.0.0/8
-A DOCKER-USER -j ufw-docker-logging-deny -p tcp -m tcp --tcp-flags FIN,SYN,RST,ACK SYN -d 172.16.0.0/12
-A DOCKER-USER -j ufw-docker-logging-deny -p udp -m udp --dport 0:32767 -d 192.168.0.0/16
-A DOCKER-USER -j ufw-docker-logging-deny -p udp -m udp --dport 0:32767 -d 10.0.0.0/8
-A DOCKER-USER -j ufw-docker-logging-deny -p udp -m udp --dport 0:32767 -d 172.16.0.0/12

-A DOCKER-USER -j RETURN

-A ufw-docker-logging-deny -m limit --limit 3/min --limit-burst 10 -j LOG --log-prefix "[UFW DOCKER BLOCK] "
-A ufw-docker-logging-deny -j DROP

COMMIT
# END UFW AND DOCKER

Using command sudo systemctl restart ufw or sudo ufw reload to restart UFW after changing the file.

このあとは通常と同じように ufw route allow proto tcp from [IPアドレス範囲] to any port 80 などとすることで適切なアクセス制限をかけることができます。

SSH ログインできるユーザーのみに公開したい

Docker 初心者の自分にはかなり難しかったです。問題は次のような感じでした。

• Dify の立ち上げを docker compose 経由以外ですることは自分の能力的に厳しかったので、Dify の立ち上げ方法は docker compose に限定して考えていました。

• Dify で立ち上がったコンテナたちは彼ら同士だけで通信できるように設定されており（？）、上記キャッシュ用の http サーバーにアクセスできなかった。

  • おそらく ufw で適切にアクセス許可すればできたと思われますが、このときは root を持っていなかったため、それはできなかった。
  • （私の linux のネットワーク設定に関する知識がないのが原因で、ほかにも解決策は色々あるのかもしれない）

そこで解決策として、Dify で立ち上がった docker network につなげる形で http サーバーをつなげました。具体的には Dify が立ち上がっている状態で、以下の設定で docker compose を行ってキャッシュサーバーを立ち上げます。

version: '3.8'

services:
  flask-mathpix:
    build:
      context: path/to/the/directory/where/the/httpserver/python/file/is
      dockerfile: Dockerfile
    container_name: flask-mathpix
    ports:
      - "5000:5000"
    volumes:
      - path/to/where/you/wish/to/save/logs/on/your/hostmachine:/app/logs
      - path/to/where/you/wish/to/save/cache/on/your/hostmachine:/app/cache
    networks:
      - ssrf_proxy_network
    environment:
      HTTP_PROXY: http://ssrf_proxy:3128
      HTTPS_PROXY: http://ssrf_proxy:3128
      NO_PROXY: localhost,127.0.0.1
    command: python app.py

networks:
  ssrf_proxy_network:
    external: true

Dify は現時点では ssrf_proxy という名前のネットワーク上で外への通信を行うように設定されているので、それを使ってやる設定をしています。

運用結果と評価

実際に運用を開始したところ、個人的には誇張抜きで論文を読む時間が 1/10 に短縮されるレベルの圧倒的性能を感じています。近くの研究者仲間からも「めっちゃいい」とフィードバックもらってます。

あとがき

正直こういうちまちましたことは基盤モデルの進化で競争優位性がなくなりそうなので、公開してもいいかなと思い自分の備忘録を兼ねて書きました。日本語論文の OCR パッケージがあったらより捗りそうです。mathpix は英語しかできないので。。。

参考情報

HTTP サーバーの実装

以下が本システムの HTTP サーバーの実装コードです。（chatgptくんに生成してもらったものをほぼそのまま使っています。）

from flask import Flask, request, jsonify, send_from_directory
import requests
import os
import json
import hashlib
import logging

app = Flask(__name__)

MATHPIX_APP_ID = os.getenv('MATHPIX_APP_ID')
MATHPIX_APP_KEY = os.getenv('MATHPIX_APP_KEY')
CACHE_DIR = "cache"
os.makedirs(CACHE_DIR, exist_ok=True)

logging.basicConfig(filename='app.log', level=logging.INFO)

def get_cache_key(arxiv_url):
    return hashlib.md5(arxiv_url.encode('utf-8')).hexdigest()

# キャッシュから結果を読み込む関数
def load_from_cache(cache_key):
    cache_path = os.path.join(CACHE_DIR, f"{cache_key}.json")
    if os.path.exists(cache_path):
        with open(cache_path, 'r') as f:
            return json.load(f)
    return None

# キャッシュに結果を保存する関数
def save_to_cache(cache_key, data):
    cache_path = os.path.join(CACHE_DIR, f"{cache_key}.json")
    with open(cache_path, 'w') as f:
        json.dump(data, f)

# PDFのダウンロード関数
def download_pdf(arxiv_url, save_path):
    response = requests.get(arxiv_url)
    with open(save_path, 'wb') as f:
        f.write(response.content)

# Mathpix APIを使ってPDFをアップロードし、pdf_idを取得する関数
def upload_pdf_to_mathpix(pdf_path):
    options = {
        "conversion_formats": {"md": True},
        "math_inline_delimiters": ["$", "$"],
        "rm_spaces": True
    }
    
    headers = {
        'app_id': MATHPIX_APP_ID,
        'app_key': MATHPIX_APP_KEY
    }
    
    with open(pdf_path, 'rb') as f:
        files = {'file': f}
        data = {'options_json': json.dumps(options)}
        response = requests.post('https://api.mathpix.com/v3/pdf', headers=headers, data=data, files=files)
    
    response_data = response.json()
    return response_data.get('pdf_id', ''), response_data.get('error', '')

# Mathpix APIを使ってPDFの処理ステータスを取得する関数
def get_pdf_status(pdf_id):
    headers = {
        'app_id': MATHPIX_APP_ID,
        'app_key': MATHPIX_APP_KEY
    }
    
    response = requests.get(f'https://api.mathpix.com/v3/pdf/{pdf_id}', headers=headers)
    return response.json()

# Mathpix APIを使ってPDFの結果を取得する関数
def get_pdf_results(pdf_id):
    headers = {
        'app_id': MATHPIX_APP_ID,
        'app_key': MATHPIX_APP_KEY
    }
    
    url = f'https://api.mathpix.com/v3/pdf/{pdf_id}.md'
    response = requests.get(url, headers=headers)
    
    results = {
        "md": response.text
    }
    
    return results

@app.route('/convert', methods=['POST'])
def convert():
    try:
        pdf_path = None
        cache_key = None
        logging.info(f"Request files: {request.files}")
        logging.info(f"File keys: {list(request.files.keys())}")
        # ファイルがアップロードされた場合
        if 'file' in request.files:
            uploaded_file = request.files['file']
            if uploaded_file.filename == '':
                return jsonify({'error': 'No file uploaded'}), 400
            
            # 一時ファイルとして保存
            temp_path = os.path.join(CACHE_DIR, uploaded_file.filename)
            uploaded_file.save(temp_path)
            
            # ファイルのハッシュを計算
            with open(temp_path, 'rb') as f:
                file_hash = hashlib.md5(f.read()).hexdigest()
            
            cache_key = file_hash
            cached_result = load_from_cache(cache_key)
            
            # キャッシュに結果が存在する場合はそれを返す
            if cached_result:
                logging.info(f"Found cached result for {uploaded_file.filename} (cache: {cache_key})")
                os.remove(temp_path)  # 一時ファイルを削除
                return jsonify(cached_result)
            
            # ファイルを処理用ディレクトリに移動
            pdf_path = os.path.join(CACHE_DIR, f"{cache_key}.pdf")
            os.rename(temp_path, pdf_path)
        
        # URLが提供された場合
        elif request.json:
            data = request.json
            if not data or 'arxiv_url' not in data:
                return jsonify({'error': 'No arXiv URL or file provided'}), 400
            
            arxiv_url = data.get('arxiv_url')
            if not arxiv_url:
                return jsonify({'error': 'Invalid URL provided'}), 400
            
            # キャッシュキーを生成
            cache_key = get_cache_key(arxiv_url)
            cached_result = load_from_cache(cache_key)
            if cached_result:
                return jsonify(cached_result)
            
            # URLからPDFをダウンロード
            pdf_path = os.path.join(CACHE_DIR, f"{cache_key}.pdf")
            try:
                download_pdf(arxiv_url, pdf_path)
            except Exception as e:
                logging.error(f"Failed to download PDF: {e}")
                return jsonify({'error': f"Failed to download PDF from URL: {e}"}), 500
        
        else:
            return jsonify({'error': 'No input provided (file or URL required)'}), 400
        
        # PDFをMathpixにアップロード
        pdf_id, error = upload_pdf_to_mathpix(pdf_path)
        if error:
            return jsonify({'error': error}), 500
        
        # ステータスをポーリングして進行状況を確認
        while True:
            status_data = get_pdf_status(pdf_id)
            if status_data.get('status') == 'completed':
                break
            time.sleep(5)  # 5秒待機して再度確認
        
        # PDFの結果を取得
        results = get_pdf_results(pdf_id)
        
        # キャッシュに結果を保存
        if cache_key:
            save_to_cache(cache_key, results)
        
        return jsonify(results)
    
    except Exception as e:
        logging.error(f"Error in /convert: {e}")
        return jsonify({'error': str(e)}), 500



if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)