Controllable RAG Agent

The agent uses a deterministic graph to act as a reasoning engine for complex RAG tasks. The workflow follows these stages:

1. Data Processing: PDF loading, chapter splitting, text preprocessing, and LLM-based summarization of chapters.
2. Indexing: Creation of a Book Quotes Database and encoding content/summaries into vector stores (FAISS).
3. Question Processing:
   • Anonymization: Replacing named entities with variables to prevent LLM bias from pre-trained knowledge.
   • Planning: Generating a high-level plan for the anonymized question, then de-anonymizing it into specific tasks.
4. Task Execution: For each task, the agent decides whether to retrieve information from vector stores or answer directly using Chain-of-Thought reasoning.
5. Verification & Re-planning: Verifying that generated content is grounded in the context and updating the plan based on new information.
6. Final Answer: Generating the final response using accumulated context and Chain-of-Thought reasoning.

The sophisticated RAG pipeline uses a PlanExecute state (a TypedDict) to pass information between nodes in the LangGraph workflow. This state tracks the lifecycle of a query from anonymization to final answer.

Key fields include:

• question: The original user query.
• anonymized_question: The query with named entities replaced by variables.
• mapping: A dictionary mapping variables back to original entities.
• plan: A list of steps to follow.
• past_steps: A history of completed tasks.
• curr_context: The context for the current task.
• aggregated_context: The cumulative information retrieved from all steps.
• tool: The identifier for the tool currently being used or selected.
• response: The final generated answer.

class PlanExecute(TypedDict):
    curr_state: str
    question: str
    anonymized_question: str
    query_to_retrieve_or_answer: str
    plan: List[str]
    past_steps: List[str]
    mapping: dict 
    curr_context: str
    aggregated_context: str
    tool: str
    response: str

The core of the advanced RAG system is a langgraph workflow designed to ensure high-quality, hallucination-free answers. The graph follows this logic:

1. Retrieve Context: Fetches chunks, chapter summaries, and quotes.
2. Keep Only Relevant Content: Distills the retrieved context using an LLM.
3. Relevance Check: If content is relevant, proceed to answer; otherwise, Rewrite Question and retry retrieval.
4. Answer Question: Generates an answer using Chain-of-Thought reasoning.
5. Grade Generation: Checks if the answer is a hallucination or if it's useful. If it's a hallucination, it retries answering; if not useful, it rewrites the question; if useful, it ends.

# The graph structure defined in the notebook
qualitative_retrieval_answer_workflow.add_edge("retrieve_context_per_question", "keep_only_relevant_content")
qualitative_retrieval_answer_workflow.add_conditional_edges(
    "keep_only_relevant_content",
    is_relevant_content,
    {"relevant":"answer_question_from_context", "not relevant":"rewrite_question"}
)
# ... and so on

The task_handler_chain acts as a router that decides which tool to use for a specific task in the plan. It selects from four primary capabilities:

• retrieve_chunks: Searches for information in a vector store of raw book chunks.
• retrieve_summaries: Searches for information in a vector store of chapter summaries.
• retrieve_quotes: Searches for information in a vector store of book quotes.
• answer_from_context: Answers a question using the aggregated_context already collected.

The handler receives the curr_task, aggregated_context, last_tool, past_steps, and the original question to make an informed decision.

class TaskHandlerOutput(BaseModel):
    """Output schema for the task handler."""
    query: str = Field(description="The query to be either retrieved from the vector store, or the question that should be answered from context.")
    curr_context: str = Field(description="The context to be based on in order to answer the query.")
    tool: str = Field(description="The tool to be used should be either retrieve_chunks, retrieve_summaries, retrieve_quotes, or answer_from_context.")

Conditional edges allow the agent to branch its logic based on runtime evaluations (e.g., checking for hallucinations or deciding which tool to use).

When using add_conditional_edges(source_node, routing_function, mapping_dict), the routing_function must return a key that exists in the mapping_dict. The value associated with that key in the dictionary determines the next node to visit.

Common decision patterns in this agent:

• Tool Selection: Routing to specific retrieval tools (chunks, summaries, or quotes) or an answering node.
• Groundedness Checks: Routing back to a previous node (e.g., replan or keep_only_relevant_content) if a hallucination is detected, or proceeding if the content is grounded.
• Task Completion: Checking if a question can_be_answered to decide between getting a final answer or breaking down the plan further.

# Example: Routing based on tool selection
agent_workflow.add_conditional_edges(
    "task_handler", 
    retrieve_or_answer, 
    {
        "chosen_tool_is_retrieve_chunks": "retrieve_book_chunks", 
        "chosen_tool_is_retrieve_summaries": "retrieve_summaries", 
        "chosen_tool_is_retrieve_quotes": "retrieve_book_quotes", 
        "chosen_tool_is_answer": "answer"
    }
)

# Example: Routing based on grounding/hallucination check
agent_workflow.add_conditional_edges(
    "answer",
    is_answer_grounded_on_context,
    {
        "hallucination": "answer", 
        "grounded on context": "replan"
    }
)

To prevent the LLM from using prior knowledge (biases) when generating a plan, the pipeline first anonymizes the question. It replaces named entities (e.g., "Harry Potter") with variables (e.g., "X") and stores the mapping.

Example workflow:

1. Anonymize: anonymize_question_chain converts "Who is Harry Potter?" to "Who is X?" with mapping {"X": "harry potter"}.
2. Plan: The planner generates steps based on the anonymized query.
3. De-anonymize: de_anonymize_plan_chain restores the original entities into the plan steps using the mapping.

# Anonymize the question
anonymized_question_output = anonymize_question_chain.invoke(state['question'])
anonymized_question = anonymized_question_output["anonymized_question"]
mapping = anonymized_question_output["mapping"]

# Generate plan using anonymized question
plan = planner.invoke({"question": anonymized_question})

# Restore original entities to the plan
deanonimzed_plan = de_anonymize_plan_chain.invoke({"plan": plan.steps, "mapping": mapping})

The controllable RAG agent is implemented using langgraph.graph.StateGraph. You define the agent's logic by adding nodes (representing individual processing steps or functions) and edges (representing the flow of control between steps).

Key components of the workflow include:

• Nodes: Functions that perform tasks like anonymize_question, planner, retrieve_book_chunks, or answer.
• Edges: Direct transitions from one node to another using add_edge.
• Conditional Edges: Decision points that route the workflow to different nodes based on the output of a function using add_conditional_edges.
• Entry Point: The starting node of the graph, defined via set_entry_point.
• Compilation: The graph must be compiled using .compile() to create an executable application.

To visualize the resulting graph, you can use the .get_graph(xray=True).draw_mermaid_png() method.

from langgraph.graph import StateGraph, END

# 1. Initialize the graph with a state schema
agent_workflow = StateGraph(PlanExecute)

# 2. Add nodes
agent_workflow.add_node("anonymize_question", anonymize_queries)
agent_workflow.add_node("planner", plan_step)

# 3. Set entry point
agent_workflow.set_entry_point("anonymize_question")

# 4. Add edges
agent_workflow.add_edge("anonymize_question", "planner")

# 5. Add conditional edges
agent_workflow.add_conditional_edges(
    "task_handler", 
    retrieve_or_answer, 
    {
        "chosen_tool_is_retrieve_chunks": "retrieve_book_chunks", 
        "chosen_tool_is_retrieve_summaries": "retrieve_summaries", 
        "chosen_tool_is_retrieve_quotes": "retrieve_book_quotes", 
        "chosen_tool_is_answer": "answer"
    }
)

# 6. Compile the graph
plan_and_execute_app = agent_workflow.compile()