機密情報漏洩の心配なし！ローカルLLMでAI秘書をつくろう！

import os
import io
import base64
from pathlib import Path


import fitz  # PyMuPDF
import requests
from PIL import Image


OLLAMA_BASE_URL = os.getenv("OLLAMA_HOST", "http://localhost:11434")
VISION_MODEL = "qwen2.5vl:7b"


OCR_PROMPT = (
   "ページに含まれる全ての文字を、表や箇条書きの構造を可能な限り保持して書き起こしてください。\n"
   "- 表は Markdown 表形式で出力してください\n"
   "- 装飾・説明・注釈の追加は不要です\n"
   "- 画像に存在しない情報は出力しないでください\n"
   "- テキストのみを出力してください"
)




def _ollama_vision_extract(image: Image.Image) -> str:
   buf = io.BytesIO()
   image.save(buf, format="PNG")
   b64 = base64.b64encode(buf.getvalue()).decode()
   res = requests.post(
       f"{OLLAMA_BASE_URL}/api/generate",
       json={
           "model": VISION_MODEL,
           "prompt": OCR_PROMPT,
           "images": [b64],
           "stream": False,
           # コンテキストを広げて生成途中の KV キャッシュ・シフトを回避する。
           # 256k は別の 500 エラーを誘発する報告があるため、ほどほどの値に留める。
           "options": {"num_ctx": 16384},
       },
       timeout=600,
   )
   res.raise_for_status()
   return res.json()["response"]




def ocr_pdf_pages(pdf: Path, dpi: int = 200) -> list[str]:
   pages: list[str] = []
   with fitz.open(pdf) as doc:
       total = doc.page_count
       for i, page in enumerate(doc, start=1):
           pix = page.get_pixmap(dpi=dpi)
           img = Image.open(io.BytesIO(pix.tobytes("png")))
           print(f"  OCR page {i}/{total} ...")
           pages.append(_ollama_vision_extract(img))
   return pages

[
 "
   1. 概要 / Overview\n
   \n
   本ドキュメントはOCR(光学文字認識)の動作検証を目的として生成された\n
   サンプルPDFです。日本語と英語が混在しており、見出し・本文・箇条書き・\n
   表・数値などの代表的なレイアウト要素を含んでいます。\n
   \n
   This document is a sample PDF generated to verify OCR behavior.\n
   It contains mixed Japanese and English text, headings, body text,\n
   bullet lists, a table, and numeric values.\n
   \n
   2. 検証ポイント / Verification Points\n
   \n
   - 日本語テキストの認識精度 (ひらがな、カタカナ、漢字)\n
   - 英数字と記号の認識: ABC abc 0123456789 !@#$%^&*()\n
   - 全角・半角の区別: A B C a b c 0 1 2 3 / ABCabc0123\n
   - 句読点と記号: 、 。「」『』（）～・…\n
   - 改行・段落の構造保持
 "
]

import json
import hashlib
from pathlib import Path


from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_ollama import OllamaEmbeddings


from ocr import VISION_MODEL, ocr_pdf_pages


EMBEDDING_MODEL = "bge-m3"


OCR_CACHE_DIR = Path("./ocr_cache")
CHROMA_BASE_DIR = Path("./chroma_db")




def _pdf_hash(pdf: Path) -> str:
   return hashlib.md5(pdf.read_bytes()).hexdigest()




def _cache_status(label: str, location: Path, *, populated: bool) -> bool:
   """キャッシュの有無をログ出力し、ヒットしたかどうかを返す。"""
   if populated:
       print(f"{label} cache hit: {location}")
       return True
   print(f"{label} cache miss: {location}")
   return False


def _get_chunked_documents(pdf: Path) -> list[Document]:
   cache_file = OCR_CACHE_DIR / f"{pdf.stem}.{_pdf_hash(pdf)}.json"


   # OCRキャッシュの有無を判定、あればロードしてテキストを返す
   if _cache_status("OCR", cache_file, populated=cache_file.exists()):
       pages = json.loads(cache_file.read_text(encoding="utf-8"))
   # OCRキャッシュがない場合はOCR処理して保存
   else:
       print(f"  Running Vision OCR ({VISION_MODEL}) on {pdf} ...")
       pages = ocr_pdf_pages(pdf)
       cache_file.parent.mkdir(parents=True, exist_ok=True)
       cache_file.write_text(json.dumps(pages, ensure_ascii=False), encoding="utf-8")
       print(f"  OCR cache saved: {cache_file}")


   docs = [
       Document(page_content=t, metadata={"source": str(pdf), "page": i + 1})
       for i, t in enumerate(pages)
   ]
   splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
   return splitter.split_documents(docs)




def find_or_create_vector_store(pdf_path: str | Path) -> Chroma:
   pdf = Path(pdf_path)
   embedding = OllamaEmbeddings(model=EMBEDDING_MODEL)
   # PDF の中身が変わったら自動で別ディレクトリに作り直されるよう、ハッシュをパスに含める
   persist_dir = CHROMA_BASE_DIR / f"{pdf.stem}_{_pdf_hash(pdf)[:12]}"


   # Chromaキャッシュの有無を判定、あればロードして返す
   populated = persist_dir.exists() and any(persist_dir.iterdir())
   if _cache_status("Chroma", persist_dir, populated=populated):
       return Chroma(persist_directory=str(persist_dir), embedding_function=embedding)


   # Chromaキャッシュがない場合は新規作成して保存
   print(f"  Building vector store at {persist_dir} ...")
   documents = _get_chunked_documents(pdf)
   return Chroma.from_documents(
       documents=documents,
       embedding=embedding,
       persist_directory=str(persist_dir),
   )

import argparse


from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import OllamaLLM


from vector_store import find_or_create_vector_store


LLM_MODEL = "qwen3.5:4b"




def get_rag_chain(retriever, llm, system_prompt):
   question_answer_chain = create_stuff_documents_chain(
       llm,
       ChatPromptTemplate.from_messages([
           ("system", system_prompt),
           ("human", "{input}"),
       ])
   )
   return create_retrieval_chain(retriever, question_answer_chain)




if __name__ == "__main__":
   parser = argparse.ArgumentParser(description="PDF に対する RAG 質問応答")
   parser.add_argument("pdf_path", help="質問対象の PDF ファイルのパス")
   parser.add_argument("input", help="LLM に投げる質問文")
   args = parser.parse_args()


   vector_store = find_or_create_vector_store(args.pdf_path)
   retriever = vector_store.as_retriever(search_kwargs={"k": 3})


   system_prompt = (
       "以下の【参考コンテキスト】の情報に基づいて、"
       "ユーザーの質問に正確に答えてください。\n\n"
       "【参考コンテキスト】:\n{context}"
   )


   rag_chain = get_rag_chain(
       retriever,
       OllamaLLM(model=LLM_MODEL, reasoning=False),
       system_prompt,
   )


   response = rag_chain.invoke({"input": args.input})


   print("回答:", response["answer"])