Retrieval as a tool
Put the search behind a tool and the agent decides when to look something up, then answers from what it found. When nothing matches, it says so instead of guessing.
This is retrieval-augmented generation: find the passages that answer a question, then answer from them. Here the search is a tool, so the model calls it like any other.
from langchain.tools import tool
from langchain_core.vectorstores import InMemoryVectorStore
from policies import chunks
from word_embeddings import WordEmbeddings
store = InMemoryVectorStore(WordEmbeddings())
store.add_documents(chunks)
@tool
def search_policies(query: str) -> str:
"""Search the shop's policies on refunds, shipping and accounts."""
found = [doc for doc, score in store.similarity_search_with_score(query, k=2) if score >= 0.3]
if not found:
return "No policy covers this."
return "\n".join(f"[{doc.metadata['source']}] {doc.page_content}" for doc in found)The tool searches, keeps the chunks scoring 0.3 or more, and returns them with their sources. Below that, it returns a fixed sentence. A model shown weak matches tends to answer from them anyway; returning nothing gives it nothing to invent from.
Pick one to watch it run, step by step.
A model for questions about policies
from langchain.messages import AIMessage
from shop_model import ShopModel
class HelpModel(ShopModel):
def decide(self, messages):
last = messages[-1]
if last.type == "tool" and last.text == "No policy covers this.":
return AIMessage("Our policies do not cover that. A person from the team will reply.")
if last.type == "tool":
return AIMessage(f"From our policies:\n{last.text}")
search = {"name": "search_policies", "args": {"query": last.text}, "id": "call_search"}
return AIMessage("", tool_calls=[search])HelpModel keeps lesson 6's machinery and changes the decisions: every question goes to search_policies, a result is quoted with its source, and the fixed sentence becomes a reply that admits the gap.
from langchain.agents import create_agent
from help_model import HelpModel
from search import search_policies
agent = create_agent(HelpModel(), tools=[search_policies])for question in ["How long does a refund take?", "Can I pay with bitcoin?"]:
result = agent.invoke({"messages": [{"role": "user", "content": question}]})
print(result["messages"][-1].text, end="\n\n")The refund question found the refund policy and the answer names its file. The bitcoin question scored under the cut, so the reply says the policies do not cover it. That second answer is the one that keeps a support desk trustworthy.
- Lower the cut to 0.2 and ask about bitcoin again.
- Ask "Is express shipping free?" and read which chunks come back.
- Return the score with each chunk and print what the model receives.
Slow is fine. Stopping is the only problem.