Scriptone

pip install -U langchain-nimble

import pathlib
import logging
import tomllib
from typing_extensions import Any, TypedDict

import fire
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_nimble import NimbleSearchRetriever
from langgraph.graph import START, END, StateGraph
from pydantic import BaseModel, Field


logger = logging.getLogger(__name__)


class SummaryData(BaseModel):
    summary: str = Field(
        ...,
        description="取得したデータを要約したテキストを格納する。",
    )


class State(TypedDict):
    query: str
    k: int
    docs: list[str]
    summary: str
    config: dict[str, Any]


def get_config() -> dict[str, Any]:
    this_dir = pathlib.Path(__file__).parent
    config_file = this_dir / "config.toml"
    with config_file.open("rb") as f:
        return tomllib.load(f)


def retrieve(state: State) -> dict[str, Any]:
    """クエリに基づいてドキュメントを検索し、関連するコンテンツをリスト形式で返す関数。NimbleSearchRetrieverを使用。"""
    retriever = NimbleSearchRetriever(k=state["k"])
    example_docs = retriever.invoke(state["query"])
    doc_list: list[str] = [doc.page_content for doc in example_docs]
    return {"docs": doc_list}


def summarize(state: State) -> dict[str, Any]:
    """ドキュメントのリストを要約する関数。ChatPromptTemplateとChatOpenAIを使用して要約を生成。"""
    prompt = ChatPromptTemplate.from_template(
        template=state["config"]["summarize"]["prompt"]
    )
    llm = ChatOpenAI(model_name=state["config"]["summarize"]["model"])
    chain = prompt | llm.with_structured_output(SummaryData)
    context = "\n\n".join(doc for doc in state["docs"])
    logger.debug(context)
    res: SummaryData = chain.invoke({"context": context})
    return {"summary": res.summary}


def proc(q: str, k: int = 5):
    """指定されたクエリとk値で処理を実行し、要約を取得する関数。kが正の整数であることを検証。"""
    if not isinstance(k, int) or k < 1:
        raise ValueError("kは1以上の整数でなければなりません。")
    # graphを作る
    graph_builder = StateGraph(State)
    # Nodeを追加する
    graph_builder.add_node("retrieve", retrieve)
    graph_builder.add_node("summarize", summarize)
    # Edgeを追加する
    graph_builder.add_edge(START, "retrieve")
    graph_builder.add_edge("retrieve", "summarize")
    graph_builder.add_edge("summarize", END)
    # Compile
    app = graph_builder.compile()
    # appを実行する
    state: State = {
        "query": q,
        "k": k,
        "docs": [],
        "summary": "",
        "config": get_config(),
    }
    res = app.invoke(state)
    # 結果を出力する
    logger.info(res["summary"])


def main():
    fire.Fire(proc)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    main()

[summarize]
model = "gpt-4.1-nano"
prompt = """
コンテキストを20センテンス以内にかつ矛盾生じないようにして要約してください。

Context:
{context}
"""

def retrieve(state: State) -> dict[str, Any]:
    """クエリに基づいてドキュメントを検索し、関連するコンテンツをリスト形式で返す関数。NimbleSearchRetrieverを使用。"""
    retriever = NimbleSearchRetriever(k=state["k"])
    example_docs = retriever.invoke(state["query"])
    doc_list: list[str] = [doc.page_content for doc in example_docs]
    return {"docs": doc_list}