LangGraphのテストを書く

## Overview

https://langchain-ai.github.io/langgraph/tutorials/rag/langgraph_adaptive_rag/

上の例のようなLangGraphで実装したコードのTestを書く

LangChainの OpenAIChatがどのように呼ばれるかを知る必要がある


## 基本 

### `chat.with_structured_output`

source

```py
structured_llm_grader = chat.with_structured_output(GradeDocuments)
retrieval_grader = grade_prompt | structured_llm_grader

score = retrieval_grader.invoke({"question": question, "document": d.page_content})
```

test

```py

chat_mock = MagicMock()
chat_mock.with_structured_output.return_value.side_effect = [
    GradeDocuments(binary_score="yes"),  # retrieval_grader for the first document
    GradeDocuments(binary_score="no"),   # retrieval_grader for the second document
]
```

### `chat.invoke`

source

```py
generate_prompt = ChatPromptTemplate.from_template(generate_prompt_template)
rag_chain = generate_prompt | chat | StrOutputParser()
generation = rag_chain.invoke({"context": documents, "question": question})
```

test

```py
    chat_mock.side_effect = [
        AIMessage(
            content="""以下のページを見つけました:
1. <http://example.com/page1|Page 1>
2. <http://example.com/page2|Page 2>"""
        ),  # rag_chain generate
    ]
```

### ToolNode

source:

```py
import operator
from typing import Annotated, Sequence, TypedDict

from langchain_core.messages import BaseMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from langgraph.prebuilt import ToolNode
from src.libs.tools import TOOL_GOOGLE


def create_example(llm=ChatOpenAI(temperature=0, streaming=True)):
    # 1. Set up tools
    tools = [
        TOOL_GOOGLE,
    ]
    tool_node = ToolNode(tools=tools)

    # 3. Bind the tools to the model
    llm = llm.bind_tools(tools)

    # 4. Define AgentState
    class AgentState(TypedDict):
        messages: Annotated[Sequence[BaseMessage], operator.add]

    # 5. Define nodes
    # node: function or runnable (e.g. agent, tool)

    # Define the function that determines whether to continue or not
    def should_continue(state):
        messages = state["messages"]
        last_message = messages[-1]
        if last_message.tool_calls:
            return "continue"
        else:
            return "end"

    # Define the function that calls the model
    def call_model(state):
        messages = state["messages"]
        response = llm.invoke(messages)
        # We return a list, because this will get added to the existing list
        return {"messages": [response]}

    # 6. Define the graph

    # Define a new graph
    workflow = StateGraph(AgentState)

    # Define the two nodes we will cycle between
    workflow.add_node("agent", call_model)
    workflow.add_node("action", tool_node)

    # Set the entrypoint as `agent`
    # This means that this node is the first one called
    workflow.set_entry_point("agent")

    # We now add a conditional edge
    workflow.add_conditional_edges(
        # First, we define the start node. We use `agent`.
        # This means these are the edges taken after the `agent` node is called.
        "agent",
        # Next, we pass in the function that will determine which node is called next.
        should_continue,
        # Finally we pass in a mapping.
        # The keys are strings, and the values are other nodes.
        # END is a special node marking that the graph should finish.
        # What will happen is we will call `should_continue`, and then the output of that
        # will be matched against the keys in this mapping.
        # Based on which one it matches, that node will then be called.
        {
            # If `tools`, then we call the tool node.
            "continue": "action",
            # Otherwise we finish.
            "end": END,
        },
    )

    # We now add a normal edge from `tools` to `agent`.
    # This means that after `tools` is called, `agent` node is called next.
    workflow.add_edge("action", "agent")

    # Finally, we compile it!
    # This compiles it into a LangChain Runnable,
    # meaning you can use it as you would any other runnable
    app = workflow.compile()
    return app


if __name__ == "__main__":
    app = create_example()
    # 7. Use
    inputs = {"messages": [HumanMessage(content="what is the weather in sf")]}
    print(app.invoke(inputs))
```

test:

```py
from unittest.mock import MagicMock, call, patch

from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from src.examples.langgraph_sample import create_example


@patch("src.libs.tools.google")
def test_langgraph(google_mock):
    google_mock.results.return_value = [
        {
            "title": "10-Day Weather Forecast for San Francisco, CA - The Weather ...",
            "link": "https://weather.com/weather/tenday/l/San+Francisco+CA+USCA0987:1:US",
            "snippet": "Be prepared with the most accurate 10-day forecast for San Francisco, CA with highs, lows, chance of precipitation from The Weather Channel and Weather.com.",
        },
    ]

    llm = MagicMock()
    llm.bind_tools.return_value.invoke.side_effect = [
        AIMessage(
            content="",
            additional_kwargs={
                "tool_calls": [{"index": 0, "id": "call_aJsmWxcJT2zEqvmqDeMddm4B", "function": {"arguments": '{"__arg1":"weather in San Francisco"}', "name": "google-search"}, "type": "function"}]
            },
            response_metadata={"finish_reason": "tool_calls", "model_name": "gpt-3.5-turbo-0125"},
            id="run-e5a2304c-36b7-43e7-9983-29a0b7f49857-0",
            tool_calls=[{"name": "google-search", "args": {"__arg1": "weather in San Francisco"}, "id": "call_aJsmWxcJT2zEqvmqDeMddm4B", "type": "tool_call"}],
        ),  # call_model
        AIMessage(
            content="The weather in San Francisco can vary, but you can check the 10-day weather forecast on The Weather Channel's website [here](https://weather.com/weather/tenday/l/San+Francisco+CA+USCA0987:1:US). Additionally, you can find the current weather conditions and forecast on AccuWeather's website [here](https://www.accuweather.com/en/us/san-francisco/94103/weather-forecast/347629).",  # noqa: E501
            tool_calls=[],
        ),  # should_continue
    ]

    app = create_example(llm)
    inputs = {"messages": [HumanMessage(content="what is the weather in sf")]}
    res = app.invoke(inputs)
    print(res)

    assert llm.bind_tools.return_value.invoke.call_args_list == [
        call(
            [
                HumanMessage(content="what is the weather in sf"),
            ]
        ),
        call(
            [
                HumanMessage(content="what is the weather in sf", additional_kwargs={}, response_metadata={}),
                AIMessage(
                    content="",
                    additional_kwargs={
                        "tool_calls": [
                            {"index": 0, "id": "call_aJsmWxcJT2zEqvmqDeMddm4B", "function": {"arguments": '{"__arg1":"weather in San Francisco"}', "name": "google-search"}, "type": "function"}
                        ]
                    },
                    response_metadata={"finish_reason": "tool_calls", "model_name": "gpt-3.5-turbo-0125"},
                    id="run-e5a2304c-36b7-43e7-9983-29a0b7f49857-0",
                    tool_calls=[{"name": "google-search", "args": {"__arg1": "weather in San Francisco"}, "id": "call_aJsmWxcJT2zEqvmqDeMddm4B", "type": "tool_call"}],
                ),
                ToolMessage(
                    content='[{"title": "10-Day Weather Forecast for San Francisco, CA - The Weather ...", "link": "https://weather.com/weather/tenday/l/San+Francisco+CA+USCA0987:1:US", "snippet": "Be prepared with the most accurate 10-day forecast for San Francisco, CA with highs, lows, chance of precipitation from The Weather Channel and Weather.com."}]',
                    name="google-search",
                    tool_call_id="call_aJsmWxcJT2zEqvmqDeMddm4B",
                ),
            ]
        ),
    ]
    assert llm.bind_tools.return_value.invoke.call_count == 2
```


## Example langgraph code

```py
from typing import List

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel, Field
from typing_extensions import TypedDict


retrieve_grader_system = """You are a grader assessing relevance of a retrieved document to a user question. \n
    If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
    It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n
    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""

generate_prompt_template = """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
Question: {question}
Context: {context}
Answer:"""  # noqa E501


hallucination_system = """You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \n
Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts."""  # noqa E501

answer_grader_system = """You are a grader assessing whether an answer addresses / resolves a question \n
Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question."""


question_rewriter_system = """You a question re-writer that converts an input question to a better version that is optimized \n
for vectorstore retrieval. Look at the input and try to reason about the underlying semantic intent / meaning."""  # noqa E501


class GradeDocuments(BaseModel):
    """Binary score for relevance check on retrieved documents."""

    binary_score: str = Field(description="Documents are relevant to the question, 'yes' or 'no'")


class GradeHallucinations(BaseModel):
    """Binary score for hallucination present in generation answer."""

    binary_score: str = Field(description="Answer is grounded in the facts, 'yes' or 'no'")


class GradeAnswer(BaseModel):
    """Binary score to assess answer addresses question."""

    binary_score: str = Field(description="Answer addresses the question, 'yes' or 'no'")


class GraphState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        question: question
        generation: LLM generation
        documents: list of documents
    """

    question: str
    generation: str
    documents: List[str]


def create_graph(chat, retriever):

    # Retrieval Grader
    structured_llm_grader = chat.with_structured_output(GradeDocuments)

    grade_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", retrieve_grader_system),
            ("human", "Retrieved document: \n\n {document} \n\n User question: {question}"),
        ]
    )

    retrieval_grader = grade_prompt | structured_llm_grader

    # Generate
    generate_prompt = ChatPromptTemplate.from_template(generate_prompt_template)
    rag_chain = generate_prompt | chat | StrOutputParser()

    # Hallucination Grader
    structured_llm_grader = chat.with_structured_output(GradeHallucinations)
    hallucination_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", hallucination_system),
            ("human", "Set of facts: \n\n {documents} \n\n LLM generation: {generation}"),
        ]
    )
    hallucination_grader = hallucination_prompt | structured_llm_grader

    # Answer Grader
    structured_llm_grader = chat.with_structured_output(GradeAnswer)
    answer_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", answer_grader_system),
            ("human", "User question: \n\n {question} \n\n LLM generation: {generation}"),
        ]
    )

    answer_grader = answer_prompt | structured_llm_grader

    # Question Re-writer
    re_write_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", question_rewriter_system),
            (
                "human",
                "Here is the initial question: \n\n {question} \n Formulate an improved question.",
            ),
        ]
    )

    question_rewriter = re_write_prompt | chat | StrOutputParser()

    def retrieve(state):
        """
        Retrieve documents

        Args:
            state (dict): The current graph state

        Returns:
            state (dict): New key added to state, documents, that contains retrieved documents
        """
        print("---RETRIEVE---")
        question = state["question"]

        # Retrieval
        documents = retriever.invoke(question)
        print(f"{len(documents)} documents retrieved")
        return {"documents": documents, "question": question}

    def generate(state):
        """
        Generate answer

        Args:
            state (dict): The current graph state

        Returns:
            state (dict): New key added to state, generation, that contains LLM generation
        """
        print("---GENERATE---")
        question = state["question"]
        documents = state["documents"]

        # RAG generation
        generation = rag_chain.invoke({"context": documents, "question": question})
        return {"documents": documents, "question": question, "generation": generation}

    def grade_documents(state):
        """
        Determines whether the retrieved documents are relevant to the question.

        Args:
            state (dict): The current graph state

        Returns:
            state (dict): Updates documents key with only filtered relevant documents
        """

        print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
        question = state["question"]
        documents = state["documents"]

        # Score each doc
        filtered_docs = []
        for d in documents:
            score = retrieval_grader.invoke({"question": question, "document": d.page_content})
            grade = score.binary_score
            if grade == "yes":
                print("---GRADE: DOCUMENT RELEVANT---")
                filtered_docs.append(d)
            else:
                print("---GRADE: DOCUMENT NOT RELEVANT---")
        print(f"filtered documents ({len(documents)} -> {len(filtered_docs)})")
        return {"documents": filtered_docs, "question": question}

    def transform_query(state):
        """
        Transform the query to produce a better question.

        Args:
            state (dict): The current graph state

        Returns:
            state (dict): Updates question key with a re-phrased question
        """

        print("---TRANSFORM QUERY---")
        question = state["question"]
        documents = state["documents"]

        # Re-write question
        better_question = question_rewriter.invoke({"question": question})
        return {"documents": documents, "question": better_question}

    ### Edges ###

    def decide_to_generate(state):
        """
        Determines whether to generate an answer, or re-generate a question.

        Args:
            state (dict): The current graph state

        Returns:
            str: Binary decision for next node to call
        """

        print("---ASSESS GRADED DOCUMENTS---")
        state["question"]
        filtered_documents = state["documents"]

        if not filtered_documents:
            # All documents have been filtered check_relevance
            # We will re-generate a new query
            print("---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---")
            return "transform_query"
        else:
            # We have relevant documents, so generate answer
            print("---DECISION: GENERATE---")
            return "generate"

    def grade_generation_v_documents_and_question(state):
        """
        Determines whether the generation is grounded in the document and answers question.

        Args:
            state (dict): The current graph state

        Returns:
            str: Decision for next node to call
        """

        print("---CHECK HALLUCINATIONS---")
        question = state["question"]
        documents = state["documents"]
        generation = state["generation"]

        score = hallucination_grader.invoke({"documents": documents, "generation": generation})
        grade = score.binary_score

        # Check hallucination
        if grade == "yes":
            print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
            # Check question-answering
            print("---GRADE GENERATION vs QUESTION---")
            score = answer_grader.invoke({"question": question, "generation": generation})
            grade = score.binary_score
            if grade == "yes":
                print("---DECISION: GENERATION ADDRESSES QUESTION---")
                return "useful"
            else:
                print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
                return "not useful"
        else:
            print("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")
            return "not supported"

    workflow = StateGraph(GraphState)

    # Define the nodes
    workflow.add_node("retrieve", retrieve)  # retrieve
    workflow.add_node("grade_documents", grade_documents)  # grade documents
    workflow.add_node("generate", generate)  # generatae
    workflow.add_node("transform_query", transform_query)  # transform_query

    # Build graph
    workflow.add_edge(START, "retrieve")
    workflow.add_edge("retrieve", "grade_documents")
    workflow.add_conditional_edges(
        "grade_documents",
        decide_to_generate,
        {
            "transform_query": "transform_query",
            "generate": "generate",
        },
    )
    workflow.add_edge("transform_query", "retrieve")
    workflow.add_conditional_edges(
        "generate",
        grade_generation_v_documents_and_question,
        {
            "not supported": "generate",
            "useful": END,
            "not useful": "transform_query",
        },
    )

    # Compile
    app = workflow.compile()

    return app

```

## Example test

```py
from unittest.mock import MagicMock
from langchain.schema import Document
from langchain_core.messages import AIMessage

from src.examples.langgraph_self_reflection import create_graph, GradeDocuments, GradeHallucinations, GradeAnswer # import from the code above

def test_graph():
    retriever = MagicMock()
    retriever.invoke.side_effect = [
        [
            Document(metadata={"title": "Page 1", "source": "http://example.com/page1"}, page_content=""),
            Document(metadata={"title": "Page 2", "source": "http://example.com/page2"}, page_content=""),
        ],  # retrieve
    ]

    chat_mock = MagicMock()
    chat_mock.side_effect = [
        AIMessage(
            content="""以下のページを見つけました:
1. <http://example.com/page1|Page 1>"""
        ),  # rag_chain generate
    ]

    chat_mock.with_structured_output.return_value.side_effect = [
        GradeDocuments(binary_score="yes"),  # retrieval_grader for the first document
        GradeDocuments(binary_score="no"),  # retrieval_grader for the second document
        GradeHallucinations(binary_score="yes"),  # hallucination_grader
        GradeAnswer(binary_score="yes"),  # answer_grader
    ]

    graph = create_graph(chat=chat_mock, retriever=retriever)

    res = graph.invoke({"question": "What's the origin of the name 'Sofia'?"})

    # Verify the expected behavior
    expected_output = """以下のページを見つけました:
1. <http://example.com/page1|Page 1>"""

    assert res == {
        "question": "What's the origin of the name 'Sofia'?",
        "generation": expected_output,
        "documents": [
            Document(metadata={"title": "Page 1", "source": "http://example.com/page1"}, page_content=""),
        ],  # the second document is filtered out by the retrieval_grader
    }
```

## Example test execution

```
pytest
```

```
tests/examples/test_langgraph_self_reflection.py::test_graph ---RETRIEVE---
2 documents retrieved
---CHECK DOCUMENT RELEVANCE TO QUESTION---
---GRADE: DOCUMENT RELEVANT---
---GRADE: DOCUMENT NOT RELEVANT---
filtered documents (2 -> 1)
---ASSESS GRADED DOCUMENTS---
---DECISION: GENERATE---
---GENERATE---
---CHECK HALLUCINATIONS---
---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---
---GRADE GENERATION vs QUESTION---
---DECISION: GENERATION ADDRESSES QUESTION---
PASSED
```