Testing

docsassistant/
    graph.py          # rewrite -> retrieve -> draft | no_answer
    agent.py          # docs_search tool + make_agent(model, checkpointer)
tests/
    fakes.py          # ScriptedChat
    conftest.py       # ollama_endpoint fixture (Testcontainers)
    test_agent.py
    test_graph.py
    test_snapshots.py
    test_integration.py
pytest.ini

from typing import Callable, Literal

from langchain_core.language_models import BaseChatModel
from langgraph.graph import END, START, StateGraph
from typing_extensions import TypedDict

Retriever = Callable[[str], list[str]]


class State(TypedDict, total=False):
    question: str
    query: str
    context: list[str]
    answer: str


def make_nodes(model: BaseChatModel, retriever: Retriever):
    def rewrite(state: State) -> dict:
        query = model.invoke(f"Rewrite as a search query: {state['question']}").text
        return {"query": query.strip()}

    def retrieve(state: State) -> dict:
        return {"context": retriever(state["query"])}

    def draft(state: State) -> dict:
        prompt = "Answer from this context only:\n" + "\n".join(state["context"]) + f"\nQ: {state['question']}"
        return {"answer": model.invoke(prompt).text}

    def no_answer(state: State) -> dict:
        return {"answer": "I could not find that in the documentation."}

    return rewrite, retrieve, draft, no_answer


def route_after_retrieve(state: State) -> Literal["draft", "no_answer"]:
    return "draft" if state.get("context") else "no_answer"


def build_graph(model: BaseChatModel, retriever: Retriever) -> StateGraph:
    rewrite, retrieve, draft, no_answer = make_nodes(model, retriever)
    graph = StateGraph(State)
    graph.add_node("rewrite", rewrite)
    graph.add_node("retrieve", retrieve)
    graph.add_node("draft", draft)
    graph.add_node("no_answer", no_answer)
    graph.add_edge(START, "rewrite")
    graph.add_edge("rewrite", "retrieve")
    graph.add_conditional_edges("retrieve", route_after_retrieve)
    graph.add_edge("draft", END)
    graph.add_edge("no_answer", END)
    return graph

from langchain.agents import create_agent
from langchain.tools import tool

CORPUS = {
    "vector-rag/chunking.adoc": "Chunk size and overlap control retrieval quality.",
    "vector-rag/pgvector.adoc": "pgvector stores embeddings in PostgreSQL.",
}


@tool
def docs_search(query: str) -> str:
    """Search the documentation and return matching page paths with a one-line snippet."""
    hits = [f"{path}: {text}" for path, text in CORPUS.items() if query.lower() in text.lower()]
    return "\n".join(hits) or "No results."


def make_agent(model, checkpointer=None):
    return create_agent(
        model,
        tools=[docs_search],
        system_prompt="You are DocsAssistant. Use docs_search before answering.",
        checkpointer=checkpointer,
    )

from langchain_core.language_models.fake_chat_models import GenericFakeChatModel


class ScriptedChat(GenericFakeChatModel):
    """GenericFakeChatModel that tolerates bind_tools, which create_agent calls when tools are given."""

    def bind_tools(self, tools, **kwargs):
        return self


def scripted(*messages):
    return ScriptedChat(messages=iter(messages))

from langchain.messages import AIMessage, HumanMessage, ToolCall, ToolMessage

from docsassistant.agent import docs_search, make_agent
from tests.fakes import scripted


def test_docs_search_tool_alone():
    assert "pgvector.adoc" in docs_search.invoke({"query": "pgvector"})
    assert docs_search.invoke({"query": "zzz"}) == "No results."


def test_agent_calls_the_search_tool_then_answers():
    model = scripted(
        AIMessage(content="", tool_calls=[ToolCall(name="docs_search", args={"query": "pgvector"}, id="call_1")]),
        AIMessage(content="pgvector stores embeddings in PostgreSQL."),
    )
    result = make_agent(model).invoke({"messages": [HumanMessage("Where are embeddings stored?")]})

    kinds = [type(m).__name__ for m in result["messages"]]
    assert kinds == ["HumanMessage", "AIMessage", "ToolMessage", "AIMessage"]
    tool_msg = next(m for m in result["messages"] if isinstance(m, ToolMessage))
    assert tool_msg.tool_call_id == "call_1"
    assert "pgvector.adoc" in tool_msg.content
    assert result["messages"][-1].content.startswith("pgvector stores")

from langchain_core.language_models.fake_chat_models import GenericFakeChatModel

from docsassistant.graph import build_graph, make_nodes, route_after_retrieve


def fake(*replies):
    return GenericFakeChatModel(messages=iter(replies))


def test_route_is_a_pure_function():
    assert route_after_retrieve({"context": ["x"]}) == "draft"
    assert route_after_retrieve({"context": []}) == "no_answer"
    assert route_after_retrieve({}) == "no_answer"


def test_a_node_function_without_the_graph():
    rewrite, *_ = make_nodes(fake("  chunk overlap  "), retriever=lambda q: [])
    assert rewrite({"question": "How much should chunks overlap?"}) == {"query": "chunk overlap"}

def test_one_node_of_the_compiled_graph():
    graph = build_graph(fake(), retriever=lambda q: [f"hit for {q}"]).compile()
    assert graph.nodes["retrieve"].invoke({"query": "pgvector"}) == {"context": ["hit for pgvector"]}

def test_end_to_end_happy_path():
    graph = build_graph(fake("pgvector", "It is PostgreSQL."), retriever=lambda q: ["pgvector stores embeddings"]).compile()
    result = graph.invoke({"question": "What is pgvector?"})
    assert result["answer"] == "It is PostgreSQL."
    assert result["query"] == "pgvector"


def test_empty_retrieval_takes_the_fallback_branch():
    # Only one scripted reply: a second model call would exhaust the script and fail the test, proving draft never ran.
    graph = build_graph(fake("nothing"), retriever=lambda q: []).compile()
    assert graph.invoke({"question": "?"})["answer"].startswith("I could not find")

from langchain.messages import AIMessage, HumanMessage, ToolCall

from docsassistant.agent import make_agent
from docsassistant.graph import build_graph
from tests.fakes import scripted


def test_graph_topology(snapshot):
    graph = build_graph(scripted(), retriever=lambda q: []).compile()
    assert graph.get_graph().draw_mermaid() == snapshot


def test_agent_trajectory(snapshot):
    model = scripted(
        AIMessage(content="", tool_calls=[ToolCall(name="docs_search", args={"query": "chunk"}, id="call_1")]),
        AIMessage(content="Tune chunk size and overlap."),
    )
    result = make_agent(model).invoke({"messages": [HumanMessage("How do I tune retrieval?")]})
    trajectory = [
        {"type": m.type, "content": m.content, "tool_calls": [tc["name"] for tc in getattr(m, "tool_calls", [])]}
        for m in result["messages"]
    ]
    assert trajectory == snapshot

graph TD;
	__start__([<p>__start__</p>]):::first
	rewrite(rewrite)
	retrieve(retrieve)
	draft(draft)
	no_answer(no_answer)
	__end__([<p>__end__</p>]):::last
	__start__ --> rewrite;
	retrieve -.-> draft;
	retrieve -.-> no_answer;
	rewrite --> retrieve;
	draft --> __end__;
	no_answer --> __end__;

E       AssertionError: assert [+ received] == [- snapshot]
E           '''
E           ...
E           	no_answer(no_answer)
E         + 	audit(audit)
E           	__end__([<p>__end__</p>]):::last
E           	__start__ --> rewrite;
E         + 	audit --> retrieve;...
FAILED tests/test_snapshots.py::test_graph_topology

[pytest]
markers =
    integration: tests that need Docker and a real Ollama model
addopts = -m "not integration"

import os

import pytest


@pytest.fixture(scope="session")
def ollama_endpoint():
    """Start Ollama in Docker with Testcontainers, or reuse OLLAMA_BASE_URL when CI already runs one."""
    if url := os.environ.get("OLLAMA_BASE_URL"):
        yield url
        return
    try:
        from testcontainers.community.ollama import OllamaContainer
    except ImportError:
        pytest.skip("testcontainers is not installed")
    model = os.environ.get("TEST_OLLAMA_MODEL", "llama3.1:8b")
    try:
        container = OllamaContainer(image="ollama/ollama:0.12.6")
        container.start()
    except Exception as exc:  # Docker not running, image cannot be pulled, ...
        pytest.skip(f"Docker/Ollama unavailable: {exc}")
    try:
        container.pull_model(model)
        yield container.get_endpoint()
    finally:
        container.stop()

import os

import pytest
from langchain.agents import create_agent
from langchain.messages import AIMessage, HumanMessage
from langchain_ollama import ChatOllama

from docsassistant.agent import docs_search

pytestmark = pytest.mark.integration


def test_real_model_asks_for_the_search_tool(ollama_endpoint):
    model = ChatOllama(model=os.environ.get("TEST_OLLAMA_MODEL", "llama3.1:8b"), base_url=ollama_endpoint, temperature=0)
    agent = create_agent(model, tools=[docs_search], system_prompt="Always call docs_search before answering.")
    result = agent.invoke({"messages": [HumanMessage("Where are embeddings stored? Search for 'pgvector'.")]})

    # Assert on structure, not wording: the exact text differs from run to run.
    calls = [tc["name"] for m in result["messages"] if isinstance(m, AIMessage) for tc in m.tool_calls]
    assert "docs_search" in calls
    assert isinstance(result["messages"][-1], AIMessage)
    assert result["messages"][-1].content

pytest -m integration                        # starts Ollama with Testcontainers, pulls llama3.1:8b
OLLAMA_BASE_URL=http://localhost:11434 pytest -m integration   # reuse an Ollama that is already running