LightRAGで取得したコンテキストをＪＳＯＮで取得する方法

LightRAGで生成をせずに取得したコンテキストをJSONで取得する方法をメモしておく。

まとめ

• only_need_context=Trueをパラメータに追加し、queryする
• queryのレスポンスから正規表現でエンティティ、リレーション、データソースを取得する

ソースコード付きのざっくり説明

まずは、LightRAGのセットアップ

import os
import asyncio
import logging
import logging.config
from lightrag import LightRAG, QueryParam
from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
from lightrag.kg.shared_storage import initialize_pipeline_status
from lightrag.utils import logger, set_verbose_debug

WORKING_DIR = "./dickens"

# Colabで動かしているため、環境変数へAPIキーをセット
from google.colab import userdata
os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')

async def initialize_rag():
    rag = LightRAG(
        working_dir=WORKING_DIR,
        embedding_func=openai_embed,
        llm_model_func=gpt_4o_mini_complete,
    )

    await rag.initialize_storages()
    await initialize_pipeline_status()

    return rag

rag = await initialize_rag()

with open("./output.txt", "r", encoding="utf-8") as f:
    await rag.ainsert(f.read())

以下のようにクエリパラメータとレスポンスを処理することで、 取得したエンティティをJSONとして扱える

# クエリ時にonly_need_contextをパラメータとして追加
jsonString=await rag.aquery(
    "クエリの内容?", param=QueryParam(mode="local",only_need_context=True)
)

import re
import json
pattern = r"```json\s*(\[\s*[\s\S]*?\])\s*```"
json_blocks = re.findall(pattern, jsonString)

# 抽出したブロックごとに json.loads でパース

entities = json.loads(json_blocks[0])
relationships = json.loads(json_blocks[1])
sources = json.loads(json_blocks[2])