LangChain源码学习路线：第 16 篇 Agentic RAG 与 Retrieval Runtime源码解剖#

1
class BaseRetriever(RunnableSerializable[str, list[Document]]): ...

1
from langchain_core.documents import Document
2
from langchain_core.retrievers import BaseRetriever
3

4

5
class TravelPolicyRetriever(BaseRetriever):
6
    index: TravelIndex
7
    city: str
8

9
    def _get_relevant_documents(self, query: str, *, run_manager) -> list[Document]:
10
        raw_hits = self.index.search(query=query, filter={"city": self.city}, limit=20)
11

12
        documents: list[Document] = []
13
        for hit in raw_hits:
14
            if hit.score < 0.35:
15
                continue
16
            documents.append(Document(
17
                page_content=hit.text,
18
                metadata={
19
                    "source": hit.url,
20
                    "title": hit.title,
21
                    "score": hit.score,
22
                    "city": self.city,
23
                },
24
            ))
25

26
        return documents
27

28

29
retriever = TravelPolicyRetriever(index=travel_index, city="Osaka")

1
create_retriever_tool(retriever, name, description)

1
from langchain_core.tools import create_retriever_tool
2

3

4
retrieval_tool = create_retriever_tool(
5
    retriever=retriever,
6
    name="travel_knowledge_search",
7
    description=(
8
        "Search verified travel documents. Use this when the answer depends "
9
        "on current opening hours, transportation rules, ticket policy, or local constraints."
10
    ),
11
    document_prompt=document_prompt,
12
    document_separator="\n\n---\n\n",
13
)

1
def simplified_retriever_tool(query: str) -> str:
2
    documents = retriever.invoke(query)
3
    rendered_documents = []
4

5
    for doc in documents:
6
        rendered_documents.append(document_prompt.format(
7
            page_content=doc.page_content,
8
            source=doc.metadata.get("source", "unknown"),
9
            title=doc.metadata.get("title", "untitled"),
10
        ))
11

12
    return document_separator.join(rendered_documents)

1
from typing import TypedDict
2
from langgraph.graph import StateGraph, START, END
3

4

5
class TravelRagState(TypedDict):
6
    question: str
7
    rewritten_query: str | None
8
    documents: list[Document]
9
    evidence: str
10
    answer: str | None
11
    retrieval_attempts: int
12

13

14
builder = StateGraph(TravelRagState)
15
builder.add_node("rewrite_query", rewrite_query_node)
16
builder.add_node("retrieve", retrieve_node)
17
builder.add_node("rerank", rerank_node)
18
builder.add_node("answer", grounded_answer_node)
19

20
builder.add_edge(START, "rewrite_query")
21
builder.add_edge("rewrite_query", "retrieve")
22
builder.add_edge("retrieve", "rerank")
23
builder.add_conditional_edges("rerank", route_after_rerank, {
24
    "answer": "answer",
25
    "rewrite_again": "rewrite_query",
26
    "refuse": END,
27
})
28
builder.add_edge("answer", END)
29

30
rag_graph = builder.compile()

1
def retrieve_node(state: TravelRagState, config: RunnableConfig) -> dict:
2
    query = state["rewritten_query"] or state["question"]
3

4
    documents = retriever.invoke(
5
        query,
6
        config={
7
            **config,
8
            "metadata": {
9
                **config.get("metadata", {}),
10
                "rag_stage": "retrieve",
11
                "query": query,
12
            },
13
        },
14
    )
15

16
    return {
17
        "documents": documents,
18
        "retrieval_attempts": state.get("retrieval_attempts", 0) + 1,
19
    }

1
def rewrite_query_node(state: TravelRagState) -> dict:
2
    rewrite_input = {
3
        "question": state["question"],
4
        "failed_query": state.get("rewritten_query"),
5
        "retrieval_attempts": state.get("retrieval_attempts", 0),
6
        "known_sources": [doc.metadata.get("source") for doc in state.get("documents", [])],
7
    }
8

9
    rewritten = rewrite_chain.invoke(rewrite_input)
10

11
    if not rewritten.query.strip():
12
        return {"rewritten_query": state["question"]}
13

14
    return {
15
        "rewritten_query": rewritten.query,
16
        "query_reason": rewritten.reason,
17
    }

1
def rerank_node(state: TravelRagState) -> dict:
2
    query = state["rewritten_query"] or state["question"]
3
    documents = state.get("documents", [])
4

5
    scored = []
6
    for doc in documents:
7
        score = reranker.score(query=query, text=doc.page_content)
8
        scored.append((score, doc))
9

10
    selected = []
11
    for score, doc in sorted(scored, key=lambda item: item[0], reverse=True):
12
        if score < 0.45:
13
            continue
14
        selected.append(Document(
15
            page_content=compress(doc.page_content, query=query),
16
            metadata={**doc.metadata, "rerank_score": score},
17
        ))
18
        if token_count(selected) >= 1800:
19
            break
20

21
    evidence = format_evidence(selected)
22
    return {"documents": selected, "evidence": evidence}

1
def grounded_answer_node(state: TravelRagState) -> dict:
2
    if not state.get("evidence"):
3
        return {"answer": "我没有检索到足够可靠的资料，不能基于证据回答。"}
4

5
    prompt_value = grounded_prompt.invoke({
6
        "question": state["question"],
7
        "evidence": state["evidence"],
8
        "citation_policy": "Every factual claim must cite a source id from evidence.",
9
    })
10

11
    ai_message = model.invoke(prompt_value)
12

13
    if not citations_are_supported(ai_message.content, state["documents"]):
14
        return {"answer": repair_or_refuse(ai_message.content, state["documents"])}
15

16
    return {"answer": ai_message.content}

1
def route_after_model(ai_message: AIMessage, state: dict) -> str:
2
    if ai_message.tool_calls:
3
        return "run_tools"
4

5
    if requires_external_facts(state["question"]):
6
        return "force_rewrite_then_retrieve"
7

8
    return "final"

1
def route_after_retrieve(state: TravelRagState) -> str:
2
    documents = state.get("documents", [])
3
    attempts = state.get("retrieval_attempts", 0)
4

5
    if not documents and attempts < 2:
6
        return "rewrite_query"
7

8
    if not documents:
9
        return "refuse"
10

11
    best_score = max(doc.metadata.get("score", 0.0) for doc in documents)
12
    if best_score < 0.35 and attempts < 2:
13
        return "rewrite_query"
14

15
    if best_score < 0.35:
16
        return "refuse"
17

18
    return "rerank"

1
def prepare_evidence(documents: list[Document], query: str, budget: int) -> tuple[list[Document], str]:
2
    ranked = rerank(query=query, documents=documents)
3
    selected: list[Document] = []
4

5
    for doc in ranked:
6
        compressed = compress(doc.page_content, query=query)
7
        candidate = Document(page_content=compressed, metadata=doc.metadata)
8

9
        if token_count(format_evidence([*selected, candidate])) > budget:
10
            continue
11

12
        selected.append(candidate)
13

14
    return selected, format_evidence(selected)

1
class TenantAwareRetriever(BaseRetriever):
2
    index: TravelIndex
3

4
    def _get_relevant_documents(self, query: str, *, run_manager) -> list[Document]:
5
        tenant_id = current_config()["metadata"]["tenant_id"]
6
        allowed_sources = load_allowed_sources(tenant_id)
7

8
        hits = self.index.search(query=query, sources=allowed_sources)
9

10
        return [
11
            Document(
12
                page_content=hit.text,
13
                metadata={
14
                    "source": hit.source,
15
                    "tenant_id": tenant_id,
16
                    "score": hit.score,
17
                },
18
            )
19
            for hit in hits
20
            if hit.source in allowed_sources
21
        ]

1
def build_retrieval_query(question: str, memory: dict, policy: dict) -> dict:
2
    query = rewrite_chain.invoke({
3
        "question": question,
4
        "user_preferences": memory.get("travel_preferences", []),
5
        "locale": memory.get("locale", "zh-CN"),
6
    })
7

8
    return {
9
        "query": query.text,
10
        "filters": {
11
            "city": query.city,
12
            "language": policy.get("allowed_languages", ["zh", "en"]),
13
            "source_type": policy.get("allowed_source_types", ["official", "partner"]),
14
        },
15
    }

1
def with_rag_trace(config: RunnableConfig, state: TravelRagState) -> RunnableConfig:
2
    return {
3
        **config,
4
        "tags": [*config.get("tags", []), "agentic-rag"],
5
        "metadata": {
6
            **config.get("metadata", {}),
7
            "question": state["question"],
8
            "query": state.get("rewritten_query"),
9
            "retrieval_attempts": state.get("retrieval_attempts", 0),
10
            "document_sources": [doc.metadata.get("source") for doc in state.get("documents", [])],
11
        },
12
    }