【GraphRAG】langchain x Ollama x Neo4jを使って、グラフデータベースを活用したRAGを実装してみる

import base64
from collections import defaultdict
import io
import os
import pickle
import re

from PIL import Image
from tqdm.notebook import tqdm
from unstructured.partition.pdf import partition_pdf


# 定数定義
DATA_PAR_PATH = os.path.join('..','..','data')
INPUT_DATA_PATH = os.path.join(DATA_PAR_PATH,'安全なウェブサイトの作り方.pdf')
OUTPUT_DATA_PATH = os.path.join(DATA_PAR_PATH,'output.pkl')

%%time

raw_pdf_elements = partition_pdf(
    languages=['jpn'],
    filename=INPUT_DATA_PATH,
    infer_table_structure=True,
    strategy='hi_res',
    extract_images_in_pdf=True,
    extract_image_block_types=['Image', 'Table'],
    extract_image_block_to_payload=True
)

res = defaultdict(list)

for elem in tqdm(raw_pdf_elements):
    page_no = elem.metadata.page_number
    cat = elem.category
    print(f"{page_no = }")
    
    if cat in ['Image', 'Table']:
        image_base64 = elem.metadata.image_base64
        decoded_image = base64.b64decode(image_base64)
        binary_image = io.BytesIO(decoded_image)
        image = Image.open(binary_image)
        display(image)

        res[page_no].append(
            {
                'category': 'image',
                'data': decoded_image  # バイナリ形式で出力させる
            }
        )
    else:
        text = elem.text
        cleaned_text = text.replace(' ', '')

        # URLは除去
        url_pattern = 'https?://[\w/:%#\$&\?\(\)~\.=\+\-]+'
        text_without_url = re.sub(url_pattern, '', cleaned_text)
        
        print(text_without_url)

        res[page_no].append(
            {
                'category': 'text',
                'detail': cat,
                'data': text_without_url
            }
        )

defaultdict(int,
            {'Title': 512,
             'UncategorizedText': 116,
             'FigureCaption': 15,
             'NarrativeText': 629,
             'Header': 100,
             'Footer': 53,
             'ListItem': 118})

with open(OUTPUT_DATA_PATH, 'wb') as wf:
    pickle.dump(res, wf)

FROM [ダウンロードしたElyza-8Bのパス名]
TEMPLATE "{{ if .System }}<|start_header_id|>system<|end_header_id|>

{{ .System }}<|eot_id|>{{ end }}{{ if .Prompt }}<|start_header_id|>user<|end_header_id|>

{{ .Prompt }}<|eot_id|>{{ end }}<|start_header_id|>assistant<|end_header_id|>

{{ .Response }}<|eot_id|>"
PARAMETER stop <|start_header_id|>
PARAMETER stop <|end_header_id|>
PARAMETER stop <|eot_id|>
PARAMETER num_keep 24

from collections import defaultdict
import io
import os
import pickle

from langchain_community.vectorstores.neo4j_vector import Neo4jVector
from langchain_ollama import OllamaEmbeddings, ChatOllama
from PIL import Image
from tqdm.notebook import tqdm

import settings as sets


# 定数定義
DATA_PAR_PATH = os.path.join('..','..','data')
INPUT_DATA_PATH = os.path.join(DATA_PAR_PATH,'output.pkl')

NEO4J_URI: str = sets.NEO4J_URI
NEO4J_USER: str = sets.NEO4J_USER
NEO4J_PASSWORD: str = sets.NEO4J_PASSWORD

MODEL_NAME = 'elyza:8b'
CHUNK_SIZE = 500

import os

from dotenv import load_dotenv


env_file_path: str = os.path.join('.','.env')
load_dotenv(env_file_path)


NEO4J_URI: str = os.environ['NEO4J_URI']
NEO4J_USER: str = os.environ['NEO4J_USER']
NEO4J_PASSWORD: str = os.environ['NEO4J_PASSWORD']

with open(INPUT_DATA_PATH, 'rb') as rf:
    elements = pickle.load(rf)

# データを確認してみる
page_no = 1
for elem in elements[page_no]:
    cat = elem['category']
    data = elem['data']
    print(cat)
    
    if cat == 'image':
        binary_image = io.BytesIO(data)
        image = Image.open(binary_image)
        display(image)
    else:
        print(data)

embeddings = OllamaEmbeddings(model=MODEL_NAME)

total = len(elements.keys())
embed_info = []
text = ''
reference_text = ''

def add_embed_info(text):
    text = text.rstrip('\n')
    embed_text = embeddings.embed_query(text)  # str -> vector(dim=4096)
    embed_info.append(
        {
            'text': text,
            'reference_text': reference_text,
            'embedding': embed_text
        }
    )

for page_no in tqdm(elements.keys(), total=total):    
    for elem in elements[page_no]:
        cat = elem['category']
        data = elem['data']
        
        if cat == 'text':
            detail_cat = elem['detail']
            
            if detail_cat == 'Header':
                if len(text) == 0:
                    continue
                
                add_embed_info(text)
                
                text = data
                reference_text = data
            elif len(text) < CHUNK_SIZE:
                text += data + '\n'
            else:
                add_embed_info(text)
                
                text = data

if len(text) > 0:
    add_embed_info(text)

# データの確認
embed_info[1]

%%time

text_embeddings = [(e['text'], e['embedding']) for e in embed_info]
metadatas = [{'reference_text': e['reference_text']} for e in embed_info]

graph_db = Neo4jVector.from_embeddings(
    text_embeddings=text_embeddings,
    embedding=embeddings,
    metadatas=metadatas,
    url=NEO4J_URI,
    username=NEO4J_USER,
    password=NEO4J_PASSWORD
)

MATCH (n) RETURN n

MATCH (n)-[r]-(m) RETURN n,r,m

headers = set([e['reference_text'] for e in embed_info])
headers

{'',
 '1.1SQLインジェクション',
 '1.2OSコマンド・インジェクション',
 '1.3パス名パラメータの未チェック/ディレクトリラバーサル',
 '1.4セッション管理の不備',
 '1.5クロスサイト・スクリプティング',
 '1.6CSRF',
 '1.7HTTPヘッダ・インジェクション',
 '1.8メールヘッダ・インジェクション',
 '1.9アクセス制御や認可制御の欠落',
 '2.1ウェブサーバのセキュリティ対策',
 '2.2DNS情報の設定不備',
 '2.3ネットワーク盗聴への対策',
 '2.5フィッシング詐欺を助長しないための対策',
 '2.6WAFによるウェブアプリケーションの保護',
 '2.7携帯ウェブ向けのサイトにおける注意点',
 '3.1失敗例(SQLインジェクション)',
 '3.2失敗例(OSコマンド・インジェクション)',
 '3.3失敗例(パス名パラメータの未チェック)',
 '3.4失敗例(不適切なセッション管理)',
 '3.5失敗例(クロスサイト・スクリプティング)',
 '3.6失敗例(CSRF)',
 '3.7失敗例(HTTPヘッダ・インジェクション',
 '3.7失敗例(HTTPヘッダ・インジェクション)',
 '3.8失敗例(メールヘッダ・インジェクション)',
 'CWE対応表',
 'おわりに',
 'はじめに',
 'チェックリスト',
 '参考資料',
 '用語集'}

create_node_query = 'CREATE' + '\n'

for header in headers:
   query_elem = f"(:Topic {{topic: '{header}'}}),"
   
   create_node_query += query_elem + '\n'

create_node_query = create_node_query.rstrip(',\n')
print(create_node_query)

CREATE
(:Topic {topic: ''}),
(:Topic {topic: '1.6CSRF'}),
(:Topic {topic: '3.4失敗例(不適切なセッション管理)'}),
(:Topic {topic: '1.1SQLインジェクション'}),
(:Topic {topic: '1.3パス名パラメータの未チェック/ディレクトリラバーサル'}),
(:Topic {topic: 'はじめに'}),
(:Topic {topic: '2.5フィッシング詐欺を助長しないための対策'}),
(:Topic {topic: '用語集'}),
(:Topic {topic: '3.8失敗例(メールヘッダ・インジェクション)'}),
(:Topic {topic: '1.2OSコマンド・インジェクション'}),
(:Topic {topic: '2.1ウェブサーバのセキュリティ対策'}),
(:Topic {topic: '1.7HTTPヘッダ・インジェクション'}),
(:Topic {topic: '3.5失敗例(クロスサイト・スクリプティング)'}),
(:Topic {topic: '3.3失敗例(パス名パラメータの未チェック)'}),
(:Topic {topic: '2.3ネットワーク盗聴への対策'}),
(:Topic {topic: '2.6WAFによるウェブアプリケーションの保護'}),
(:Topic {topic: '3.2失敗例(OSコマンド・インジェクション)'}),
(:Topic {topic: '2.7携帯ウェブ向けのサイトにおける注意点'}),
(:Topic {topic: '3.1失敗例(SQLインジェクション)'}),
(:Topic {topic: '3.6失敗例(CSRF)'}),
(:Topic {topic: '1.4セッション管理の不備'}),
(:Topic {topic: '1.8メールヘッダ・インジェクション'}),
(:Topic {topic: '1.9アクセス制御や認可制御の欠落'}),
(:Topic {topic: 'チェックリスト'}),
(:Topic {topic: '3.7失敗例(HTTPヘッダ・インジェクション'}),
(:Topic {topic: 'CWE対応表'}),
(:Topic {topic: '3.7失敗例(HTTPヘッダ・インジェクション)'}),
(:Topic {topic: '参考資料'}),
(:Topic {topic: '1.5クロスサイト・スクリプティング'}),
(:Topic {topic: 'おわりに'}),
(:Topic {topic: '2.2DNS情報の設定不備'})

graph_db.query(create_node_query)

MATCH (n) RETURN n

create_relation_query = '''
MATCH (c:Chunk), (t:Topic)
WHERE c.reference_text = t.topic
CREATE (t)-[:reference]->(c)'''

graph_db.query(create_relation_query)

MATCH (n)-[r]->(m) RETURN n,r,m

from langchain.chains import GraphCypherQAChain
from langchain_community.graphs import Neo4jGraph
from langchain_community.vectorstores.neo4j_vector import Neo4jVector
from langchain_ollama import ChatOllama, OllamaEmbeddings

import settings as sets


# 定数定義
NEO4J_URI: str = sets.NEO4J_URI
NEO4J_USER: str = sets.NEO4J_USER
NEO4J_PASSWORD: str = sets.NEO4J_PASSWORD

MODEL_NAME = 'elyza:8b'

user_query = 'SQLインジェクションについて教えて'

%%time

# グラフDB（リレーション含まない、すなわちただのベクトルDBと同等）
exist_db = Neo4jVector.from_existing_graph(
    embedding=OllamaEmbeddings(model=MODEL_NAME),
    node_label='Chunk',
    url=NEO4J_URI,
    username=NEO4J_USER,
    password=NEO4J_PASSWORD,
    text_node_properties=['text'],
    embedding_node_property='embedding'
)

exist_db.similarity_search_with_relevance_scores(user_query)

[(Document(metadata={'reference_text': '3.4失敗例(不適切なセッション管理)'}, page_content='\ntext: 「093a2031a79cb4904b1466ee7ad5faaa3afe7b787db66712f407326b213cc2a4」等です。このプログラムではセッションIDにハッシュ値を使っているため、一見、安全そうに見えるかもしれません。しかし、セッションIDの生成方法が第三者に知られた場合光には、第三者がセッションIDを推測できる余地があります。\n罠のウェブサイトへの誘導等によって、攻撃者は利用者の接続元IPアドレスを入手できますぎ。一方、ポート番号は攻撃者が知り得えない情報です。しかし、接続元ポート番号の範囲は1024から65535であり、利用者のネットワーク環境によってはこの範囲が2万通り程度に限定される場合があります。\nこのプログラムがオープンソースで開発されている場合、またはソースコードが漏えいしてしまった場合等を想定できます。5ウェブサイトに到達するまでのネットワーク経路によっては、特定のipアドレスとならない場合があります\n。\nな'),
  0.7720184326171875),
 (Document(metadata={'reference_text': '2.7携帯ウェブ向けのサイトにおける注意点'}, page_content='\ntext: 4ここで使用するセッションIDは、第三者に推測されない必要があります。詳しくは18ページの根本的解決4-()を参照してください。60\nな'),
  0.7718048095703125),
 (Document(metadata={'reference_text': '1.5クロスサイト・スクリプティング'}, page_content='\ntext: 1.5クロスサイト・スクリプティング2)に該当するウェブアプリケーションの例には、自由度の高い掲示板やブログ等が挙げられます。たとえば、利用者が入力文字の色やサイズを指定できる機能等を実装するために、HTMLテキストの入力を許可する場合があるかもしれません。\n3)は、1、2)の両者のウェブアプリケーションに共通して必要な対策です。\n1.5.1HTMLテキストの入力を許可しない場合の対策\n園根本的解決\n5-(|\nウェブページに出力する全ての要素に対して、エスケープ処理を施す。\nウェブページを構成する要素として、ウェブページの本文やHTMLタグの属性値等に相当する全ての出力要素にエスケープ処理を行います。エスケープ処理には、ウェブページの表示に影響する特別な記号文字(「く<」、「>」、「&」等)を、HTMLエンティティ(「&It:」、「&gt:」、「&amp:」等)に置換する方法があります。また、HTMLタグを出力する場合は、その属性値を必ず「"」(ダブルクォート)で括るようにします。そして、「"」で括られた属性値に含まれる「"」を、HTMLエンティティ「&quot:\'」にエスケープします'),
  0.7661819458007812),
 (Document(metadata={'reference_text': '3.6失敗例(CSRF)'}, page_content='\ntext: 3.6失敗例(CSRF)3.6CSRF(クロスサイト・リクエスト・フォージェリ)うの例\nCSRF(クロスサイト・リクエスト・フォージェリ)の脆弱性を考慮できていない例として、登録情報編集画面のプログラムを紹介します。\n園PHPによる登録情報編集機能\n【脆弱な実装】\n下図は、会員制ウェブサイトにおける、ユーザの会員登録情報を変更する機能の、典型的な画面皿移の例です。ここでは、ユーザが住所を東京都から大阪府に変更するときの操作を例にしています。\n画面1\n画面2\nこのサイトの構成では、まず画面1(view.php)で登録情報を確認し、編集の必要があれば編集ボタンを押して画面2(edit.php)へ進み、画面2で必要な情報を入力して、さらにパスワードを入力した上で、画面3(confirm.php)へ進むようになっています。画面3で入力した情報を確認し、実行ボタンを押して画面4(commit.php)に進むと、ここで登録情報の更新処理が実行されて、その旨が表示されます。\n画面2で、本人確認のためにパスワードを入力するようになっており、パスワードが正しい場合にしか画面3に進めないようになっています。'),
  0.7476806640625)]

# グラフDB（リレーション含む）
graph_db = Neo4jGraph(
    url=NEO4J_URI,
    username=NEO4J_USER,
    password=NEO4J_PASSWORD
)

graph_chain = GraphCypherQAChain.from_llm(
    ChatOllama(model=MODEL_NAME, temperature=0.0),
    graph=graph_db,
    verbose=True,
    allow_dangerous_requests=True
)

graph_chain.invoke(user_query)

> Entering new GraphCypherQAChain chain...
Generated Cypher:
MATCH (t:Topic)-[:reference]->(c:Chunk) RETURN c.text
Full Context:
[{'c.text': 'ウェブアプリケーションのセキュリティ実装とウェブサイトの安全性向上のための取り組み\nIDA独立行政法人情報処理推進機構セキュリティセンター\n2012年12月\n本書は、以下のURLからダウンロードできます。「安全なウェブサイトの作り方」\nおわりに..…参考資料..eoメールヘッダ・インジェクションの例'},
 {'c.text': "プレースホルダに実際の値を割り当てる処理をバインドと呼びます。バインドの方式には、プレースホルダのままSQL文をコンパイルしておき、データベースエンジン側で値を割り当てる方式(静的プレースホルダ)と、アプリケーション側のデータベース接続ライブラリ内で値をエスケーブプ処理してプレースホルダにはめ込む方式(動的プレースホルダ)があります。静的プレースホルダは、SQLのISO/JIS規格では、準備された文(PreparedStatement)と呼ばれます。どちらを用いてもSQLインジェクション脆弱性を解消できますが、原理的にSQLインジェクション脆弱性の可能性がなくなるという点で、静的プレースホルダの方が優ります。詳しくは本書別冊の「安全なSQLの呼び出し方」のプレースホルダの項(3.2節)を参照してください。\nごも\nSQL文の組み立てを文字列連結により行う場合は、エスケーブプ処理等を行うデータベースエンジンのAPIlを用いて、SQL文のリテラルを正しく構成する。\nSQL文の組み立てを文字列連結により行う場合は、SQL文中で可変となる値をリテラル(定数)の形で埋め込みます。値を文字列型として埋め込む場合は、値をシングルクォートで囲んで記述しますが、その際に文字列リテラル内で特別な意味を持つ記号文字をエスケープ処理します(たとえば、「'」つつ「'「」、「\\」つ「」等)。値を数値型として埋め込む場合は、数値リテラルであることを確実にする処"},
 {'c.text': '$最新情報は、下記URLを参照してください。\n脆弱性関連情報に関する届出状況:\n7\n|'}, 
 {'c.text': '1.1SQLインジェクション理(数値型へのキャスト等)を行います。\nこうした処理で具体的に何をすべきかは、データベースエンジンの種類や設定によって異なるため、それにあわせた実装が必要です。データベースエンジンによっては、リテラルを文字列として生成する専用のAPI\'を提供しているものがありますので、それを利用することをお勧めします。詳しくは、「安全なSQLの呼び出し方」の4.1節を参照してください。\nなお、この処理は、外部からの入力の影響を受ける値のみに限定して行うのではなく、SQL文を構成する全てのリテラル生成に対して行うべきです。\n|\n1-(ii)\nRn"ウェブアプリケーションに渡されるパラメータにSQL文を直接指定しない\n。\nこれは、いわば「論外」の実装ですが、hiddenパラメータ等にSQL文をそのまま指定するという事例の届出がありましたので、避けるべき実装として紹介します。\nウェブアプリケーションに渡されるパラメータにSQL文を直接指定する実装は、そのパラメータ値の改変により、データベースの不正利用につながる可能性があります。\n園保険的対策\n1-(iii)\nエラーメッツセージをそのままブラウザに表示しない。'},
 {'c.text': 'エラーメッセージの内容に、データベースの種類やエラーの原因、実行エラーを起こしたSQL文等の情報が含まれる場合、これらはSQLインジェクション攻撃につながる有用な情報となりえます。また、エラーメッセージは、攻撃の手がかりを与えるだけでなく、実際に攻撃された結果を表示する情報源として悪用される場合があります。データベースに関連するエラーメッセージは、利用者のブラウザ上に表示させないことをお勧めします。1-(iv)\n|\n1-(iv)hn"データベースアカウントに適切な権限を与える。\nウェブアプリケーションがデータベースに接続する際に使用するアカウントの権限が必要以上に高い場合、攻撃による被害が深刻化する恐れがあります。ウェブアプリケーションからデータベースに渡す命令文を洗い出し、その命令文の実行に必要な最小限の権限をデータベースアカウントに与えてください。\n以上の対策により、SQLインジェクション攻撃に対する安全性の向上が期待できます。データベースと連携したウェブアプリケーションの構築や、SQLインジェクションの脆弱性に関する情報については、次の資料も参考にしてください。\n7′実行環境によっては、エスケープ処理を適切に行わない脆弱性が指摘されているAPIもあります。その場合は修正パッチを適用するか、別の方法を検討して下さい。'},
 {'c.text': '8|\n|\n画参考URL\nIPA:安全なSQLの呼び出し方\nIPA:知っていますか?脆弱性(ぜいじゃくせい)「1.SQLインジェクション」\nIPA:セキュア・プログラミング講座「SQL注入:#実装における対策」\nA:セキュア・プログラミング講座「SQL注入:#2設定における対策」D://www.ipa.go.jp/security/awareness/vendor/programmingv2/contents/503.htmlき\nIPA:情報セキュリティ白書2009第2部「10大脅威攻撃手法の『多様化』が進む」\nIPA:情報セキュリティ白書2008第2部「10大脅威ますます進む『見えない化』」'},
 {'c.text': '1.1SQLインジェクション'},
 {'c.text': '1.1SQLインジェクション1.ウェブアプリケーションのセキュリティ実装\n本章では、ウェブアプリケーションのセキュリティ実装として、下記の脆弱性を取り上げ*、発生しうる脅威、注意が必要なサイト、根本的解決および保険的対策を示します。\n1)SQLインジェクション\n2)0Sコマンド・インジェクション\n3)パス名パラメータの未チェックンディレクトリ・トラバーサル\n4)セッション管理の不備\n5)クロスサイト・スクリプティング\n6)CSRF(クロスサイト・リクエスト・フォージェリ)\n7)HTTPヘッダ・インジェクション\n8)メールヘッダ・インジェクション\n9)アクセス制御や認可制御の欠落\n*資料の構成上、脆弱性の深刻度や攻撃による影響を考慮して項番を割り当てていますが、これは対策の優先順位を示すものではありません。優先順位は運営するウェブサイトの状況に合わせてご検討ください。\n5'},
 {'c.text': '1.1SQLインジェクション1.1SQLインジェクション\nデータベースと連携したウェブアプリケーションの多くは、利用者からの入力情報を基にSQL文(データベースへの命令文)を組み立てています。ここで、SQL文の組み立て方法に問題がある場合、攻撃によってデータベースの不正利用をまねく可能性があります。このような問題を「SQLインジェクションの脆弱性」と呼び、問題を悪用した攻撃を、「SQLインジェクション攻撃」と呼びます。\n画発生しうる脅威\nSQLインジェクション攻撃により、発生しうる脅威は次のとおりです。\nデータベースに蓄積された非公開情報の閲覧\n個人情報の漏えい等\nデータベースに蓄積された情報の改ざん、消去\nウェブページの改ざん、パスワード変更、システム停止等\n認証回避による不正ログイン*\nログインした利用者に許可されている全ての操作を不正に行われる\nストアドプロシージャ等を利用した0Sコマンドの実行\nシステムの乗っ取り、他への攻撃の踏み台としての悪用等\n回注意が必要なウェブサイトの特徴\n運営主体やウェブサイトの性質を問わず、データベース?を利用するウェブアプリケーションを設置しているウェブサイトに存在しうる問題です。個人情報等の重要情報をデータベースに格納しているウェブサイトは、特に注意が必要です。'},
 {'c.text': '*後述「1.3セッション管理の不備」で解説する「発生しうる脅威」と同じ内容です。5代表的なデータベースエンジンには、MySQL,PostgreSQL,Oracle,MicrosoftSQLServer,DB2等が挙げられます。\n6'}]

> Finished chain.
{'query': 'SQLインジェクションについて教えて',
 'result': "SQLインジェクションとは、Webアプリケーションの入力フォームなどを通じて、悪意のあるユーザーがデータベースに不正な命令を送信し、データを改ざんや漏えいさせる攻撃手法です。\n\n具体的には、Webアプリケーションがデータベースとの通信で使用するSQL文に、特別な文字列を埋め込むことで、意図的にデータベースの動作を変更したり、情報を取得したりすることができます。\n\n例えば、入力フォームに「' OR 1=1 --」などの特殊な文字列を入力し、Webアプリケーション安全等/tech/sqlinjection.html"}

user_query2 = 'WAFの効果について例を教えて'
graph_chain.invoke(user_query2)

> Entering new GraphCypherQAChain chain...
Generated Cypher:
MATCH (t:Topic {topic:"WAF"})-[:reference]->(c:Chunk) RETURN c.text
Full Context:
[]

> Finished chain.
{'query': 'WAFの効果について例を教えて',
 'result': 'WAFは、Web Application Firewallの略で、Webアプリケーション層での攻撃を防ぐためのセキュリティ対策です。'}