RetrievalQA 与 LLMChain 迁移到 LCEL：对等性与调用次数实测

"""Offline test doubles: a keyword retriever and a chat model that records every prompt."""
from typing import Any, List, Optional
 
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.retrievers import BaseRetriever
 
DOCS = [
    Document(page_content="个人套餐购买后七天内可以申请全额退款。", metadata={"source": "refund"}),
    Document(page_content="发票在订单详情页申请，一至三个工作日内开具。", metadata={"source": "invoice"}),
    Document(page_content="分享链接默认有效期为七天，创建者可以随时取消分享。", metadata={"source": "share"}),
    Document(page_content="删除的文件会在回收站保留三十天。", metadata={"source": "trash"}),
]
 
 
class KeywordRetriever(BaseRetriever):
    """Rank documents by how many distinct query characters they contain; ties keep corpus order."""
 
    k: int = 2
 
    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
        chars = {c for c in query if not c.isspace()}
        scored = [(len(chars & set(d.page_content)), -i, d) for i, d in enumerate(DOCS)]
        scored.sort(key=lambda item: (item[0], item[1]), reverse=True)
        return [d for score, _, d in scored[: self.k] if score > 0]
 
 
class RecordingModel(BaseChatModel):
    """Answers 'answer N' and keeps every message list it received."""
 
    calls: int = 0
    seen: List[List[BaseMessage]] = []
 
    @property
    def _llm_type(self) -> str:
        return "recording"
 
    def get_token_ids(self, text: str) -> List[int]:
        """Count characters as tokens so token-budget logic works without a tokenizer download."""
        return [ord(c) for c in text]
 
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        self.calls += 1
        self.seen.append(list(messages))
        return ChatResult(generations=[ChatGeneration(message=AIMessage(content=f"answer {self.calls}"))])

"""Run the old chains and their LCEL replacements side by side and compare what really differs."""
import warnings
 
from fakes import KeywordRetriever, RecordingModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnableParallel, RunnablePassthrough
 
try:
    from langchain_classic.chains import LLMChain, RetrievalQA
 
    IMPORT_PATH = "langchain_classic.chains"
except ImportError:
    from langchain.chains import LLMChain, RetrievalQA
 
    IMPORT_PATH = "langchain.chains"
 
warnings.filterwarnings("ignore")  # the chains emit deprecation warnings on every use
 
print(f"old chains imported from: {IMPORT_PATH}")
QUESTION = "个人套餐七天内能退款吗"
 
 
def texts(model):
    return [(m.type, m.content) for m in model.seen[0]]
 
 
print()
print("1. LLMChain versus prompt | model | parser")
prompt = ChatPromptTemplate.from_messages([("human", "请用一句话回答：{question}")])
old_model, new_model = RecordingModel(), RecordingModel()
old = LLMChain(llm=old_model, prompt=prompt).invoke({"question": QUESTION})
new = (prompt | new_model | StrOutputParser()).invoke({"question": QUESTION})
print("  old result:", old)
print("  new result:", repr(new))
print("  model saw identical messages:", texts(old_model) == texts(new_model))
 
print()
print("2. RetrievalQA (stuff) versus an LCEL retrieval chain")
retriever = KeywordRetriever(k=2)
old_model, new_model = RecordingModel(), RecordingModel()
old_qa = RetrievalQA.from_chain_type(llm=old_model, chain_type="stuff", retriever=retriever, return_source_documents=True)
old = old_qa.invoke({"query": QUESTION})
 
rag_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "Use the following pieces of context to answer the user's question.\n"
            "If you don't know the answer, just say that you don't know, don't try to make up an answer.\n"
            "----------------\n{context}",
        ),
        ("human", "{question}"),
    ]
)
 
 
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)
 
 
answer = (
    RunnableLambda(lambda x: {"context": format_docs(x["docs"]), "question": x["question"]})
    | rag_prompt
    | new_model
    | StrOutputParser()
)
rag = RunnableParallel(docs=retriever, question=RunnablePassthrough()).assign(answer=answer)
new = rag.invoke(QUESTION)
 
print("  old keys:", list(old))
print("  new keys:", list(new))
print("  model saw identical messages:", texts(old_model) == texts(new_model))
print("  same answer text:", old["result"] == new["answer"])
print("  same source documents:", [d.metadata for d in old["source_documents"]] == [d.metadata for d in new["docs"]])
 
print()
print("3. A thin adapter keeps the old output shape for existing callers")
as_old_shape = rag | RunnableLambda(
    lambda x: {"query": x["question"], "result": x["answer"], "source_documents": x["docs"]}
)
adapted = as_old_shape.invoke(QUESTION)
print("  adapted keys:", list(adapted))
same_values = all(adapted[key] == old[key] for key in ("query", "source_documents"))
print("  query and source_documents equal the old output:", same_values)
 
print()
print("4. Input shape")
for label, call in (
    ("old chain, {'question': ...}", lambda: old_qa.invoke({"question": QUESTION})),
    ("old chain, plain string", lambda: old_qa.invoke(QUESTION)),
    ("new chain, plain string", lambda: rag.invoke(QUESTION)),
):
    try:
        call()
        print(f"  {label:<30} accepted")
    except Exception as exc:
        print(f"  {label:<30} {type(exc).__name__}: {exc}")

old chains imported from: langchain.chains

1. LLMChain versus prompt | model | parser
  old result: {'question': '个人套餐七天内能退款吗', 'text': 'answer 1'}
  new result: 'answer 1'
  model saw identical messages: True

2. RetrievalQA (stuff) versus an LCEL retrieval chain
  old keys: ['query', 'result', 'source_documents']
  new keys: ['docs', 'question', 'answer']
  model saw identical messages: True
  same answer text: True
  same source documents: True

3. A thin adapter keeps the old output shape for existing callers
  adapted keys: ['query', 'result', 'source_documents']
  query and source_documents equal the old output: True

4. Input shape
  old chain, {'question': ...}   ValueError: Missing some input keys: {'query'}
  old chain, plain string        accepted
  new chain, plain string        accepted

"""How many model calls does one RetrievalQA question cost, per chain_type? And with no documents?"""
import warnings
 
from fakes import KeywordRetriever, RecordingModel
 
try:
    from langchain_classic.chains import RetrievalQA
except ImportError:
    from langchain.chains import RetrievalQA
 
warnings.filterwarnings("ignore")
 
QUESTION = "个人套餐七天内能退款吗"
 
print("three retrieved documents")
for chain_type in ("stuff", "map_reduce", "refine"):
    model = RecordingModel()
    qa = RetrievalQA.from_chain_type(llm=model, chain_type=chain_type, retriever=KeywordRetriever(k=3))
    result = qa.invoke({"query": QUESTION})
    print(f"  chain_type={chain_type:<11} model_calls={model.calls}  result={result['result']!r}")
 
print()
print("a question that matches no document")
model = RecordingModel()
qa = RetrievalQA.from_chain_type(llm=model, retriever=KeywordRetriever(k=2))
result = qa.invoke({"query": "zzz"})
system_prompt = model.seen[0][0].content
print(f"  model_calls={model.calls}  result={result['result']!r}")
print(f"  system prompt ends with: {system_prompt[-20:]!r}")

three retrieved documents
  chain_type=stuff       model_calls=1  result='answer 1'
  chain_type=map_reduce  model_calls=4  result='answer 4'
  chain_type=refine      model_calls=3  result='answer 3'

a question that matches no document
  model_calls=1  result='answer 1'
  system prompt ends with: 'r.\n----------------\n'

"""The LCEL chain needs only langchain-core: it runs where langchain-classic is not installed."""
import importlib.util
 
from fakes import KeywordRetriever, RecordingModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
 
print("langchain_classic installed:", importlib.util.find_spec("langchain_classic") is not None)
print("langchain.chains importable:", importlib.util.find_spec("langchain.chains") is not None)
 
retriever = KeywordRetriever(k=2)
prompt = ChatPromptTemplate.from_messages(
    [("system", "Answer using only this context:\n{context}"), ("human", "{question}")]
)
 
 
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)
 
 
def build(model):
    return (
        {"context": retriever | RunnableLambda(format_docs), "question": RunnablePassthrough()}
        | prompt
        | model
        | StrOutputParser()
    )
 
 
def build_guarded(model):
    """Skip the model call when nothing was retrieved, instead of asking with an empty context."""
    chain = build(model)
 
    def route(question):
        if not retriever.invoke(question):
            return "no relevant document found"
        return chain.invoke(question)
 
    return RunnableLambda(route)
 
 
print()
print("a question with matching documents")
model = RecordingModel()
print("answer:", repr(build(model).invoke("个人套餐七天内能退款吗")))
print("system prompt the model saw:", repr(model.seen[0][0].content))
 
print()
print("a question that matches no document")
model = RecordingModel()
print("plain chain  -> model_calls:", (build(model).invoke("zzz"), model.calls)[1])
model = RecordingModel()
result = build_guarded(model).invoke("zzz")
print(f"guarded chain -> result: {result!r}, model_calls: {model.calls}")

langchain_classic installed: False
langchain.chains importable: False

a question with matching documents
answer: 'answer 1'
system prompt the model saw: 'Answer using only this context:\n个人套餐购买后七天内可以申请全额退款。\n\n发票在订单详情页申请，一至三个工作日内开具。'

a question that matches no document
plain chain  -> model_calls: 1
guarded chain -> result: 'no relevant document found', model_calls: 0