# How the Query Rewriting Node Improves RAG Retrieval in LangGraph

> Optimize RAG retrieval with the query rewriting node. This LangGraph component refines vague questions using LLMs, enhancing document relevance and answer accuracy for better RAG pipelines.

- Repository: [jamwithai/production-agentic-rag-course](https://github.com/jamwithai/production-agentic-rag-course)
- Tags: deep-dive
- Published: 2026-03-23

---

**The query rewriting node optimizes vague user questions using a structured LLM call before re-attempting retrieval, creating a feedback loop that dramatically improves document relevance and answer accuracy in RAG pipelines.**

The query rewriting node serves as an intelligent preprocessing layer in production agentic RAG architectures. In the `jamwithai/production-agentic-rag-course` repository, this LangGraph component intercepts under-specified queries and transforms them into high-intent search terms using schema-validated LLM outputs. By implementing robust fallback logic and comprehensive observability, the node ensures that your vector database receives optimized queries that match the user's true information needs.

## How Query Rewriting Works in the RAG Pipeline

The query rewriting node in [`src/services/agents/nodes/rewrite_query_node.py`](https://github.com/jamwithai/production-agentic-rag-course/blob/main/src/services/agents/nodes/rewrite_query_node.py) operates as a dedicated preprocessing step within the LangGraph workflow. When initial retrieval yields insufficient results, the node reads `state["original_query"]` and invokes an LLM with a structured output schema to generate an optimized search query.

### Structured LLM Invocation with Schema Validation

The node enforces type-safe outputs using `QueryRewriteOutput` defined in [`src/services/agents/models.py`](https://github.com/jamwithai/production-agentic-rag-course/blob/main/src/services/agents/models.py). This Pydantic schema requires the LLM to return both a `rewritten_query` string and a `reasoning` explanation, ensuring transparency in the transformation process. The prompt template stored in `REWRITE_PROMPT` within [`src/services/agents/prompts.py`](https://github.com/jamwithai/production-agentic-rag-course/blob/main/src/services/agents/prompts.py) guides the model to expand ambiguous terms like "transformer" into specific, searchable phrases such as "latest transformer architectures for image classification in 2024".

### Resilient Fallback Logic

If the LLM call fails or returns malformed output, the node implements a deterministic fallback strategy to guarantee pipeline continuity. The system automatically appends contextual keywords to the original query, generating `"{original_question} research paper arxiv machine learning"` as a failsafe search term. This ensures that the retrieval node always receives a valid query, even during LLM service degradation or timeout scenarios.

## LangGraph State Management and Graph Routing

The query rewriting node integrates deeply with the LangGraph state machine to create a powerful retrieval feedback loop.

### Updating State with Optimized Queries

After successful transformation, the node returns a `HumanMessage` containing the rewritten query and updates `state["rewritten_query"]` in the shared agent state. This state mutation preserves the original user intent while providing downstream nodes with an enhanced search term stored in [`src/services/agents/state.py`](https://github.com/jamwithai/production-agentic-rag-course/blob/main/src/services/agents/state.py).

### The Retrieval Feedback Loop

In [`src/services/agents/agentic_rag.py`](https://github.com/jamwithai/production-agentic-rag-course/blob/main/src/services/agents/agentic_rag.py), the graph construction creates a critical optimization cycle through `workflow.add_edge("rewrite_query", "retrieve")`. This edge routes the optimized query back to the retrieval node for a second pass against the OpenSearch index. By converting vague requests into precise search terms, the embedding-based similarity search can match a larger and more relevant document set, directly improving the context provided to the generation phase.

## Observability with Langfuse Tracing

When tracing is enabled, the node creates a Langfuse span named `query_rewriting` that captures execution metadata including the original query, rewrite attempt count, model name (e.g., "llama3:8b"), temperature settings, and the final rewritten query. This span is closed with timing metadata, making the rewrite step visible in performance dashboards and allowing developers to track how query transformations impact retrieval quality across different LLM configurations.

## Code Implementation Examples

### Directly Invoking the Rewrite Node

You can invoke the query rewriting logic independently using the `ainvoke_rewrite_query_step` function for testing or custom workflows:

```python
from src.services.agents.nodes.rewrite_query_node import ainvoke_rewrite_query_step
from src.services.agents.state import AgentState
from src.services.agents.context import Context
from langgraph.runtime import Runtime

# Configure runtime with required clients

runtime = Runtime[Context](
    context=Context(
        ollama_client=...,          # Ollama client instance

        opensearch_client=...,      # OpenSearch client instance

        embeddings_client=...,      # Jina embeddings client

        langfuse_tracer=None,       # Optional LangfuseTracer

        trace=None,
        langfuse_enabled=False,
        model_name="llama3:8b",
        temperature=0.3,
        top_k=5,
        max_retrieval_attempts=3,
        guardrail_threshold=70,
    )
)

# Initialize state with vague query

state: AgentState = {
    "messages": [],
    "original_query": "transformer models",
    "retrieval_attempts": 0,
    "guardrail_result": None,
    "routing_decision": None,
    "sources": None,
    "relevant_sources": [],
    "relevant_tool_artefacts": None,
    "grading_results": [],
    "metadata": {},
    "rewritten_query": None,
}

result = await ainvoke_rewrite_query_step(state, runtime)
print(result["rewritten_query"])

```

**Result:** The node transforms "transformer models" into a specific query like `"What are the most recent transformer architectures for image classification (2023-2024) with benchmark results?"`

### Full Agentic RAG Integration

In production, the rewrite node triggers automatically when document grading detects low relevance:

```python
from src.services.agents.agentic_rag import AgenticRAGService
from src.services.opensearch.client import OpenSearchClient
from src.services.ollama.client import OllamaClient
from src.services.embeddings.jina_client import JinaEmbeddingsClient

# Initialize clients

opensearch = OpenSearchClient(...)
ollama = OllamaClient(...)
embeddings = JinaEmbeddingsClient(...)

service = AgenticRAGService(
    opensearch_client=opensearch,
    ollama_client=ollama,
    embeddings_client=embeddings,
)

# Vague query triggers internal rewrite loop

answer = await service.ask("transformer models")
print(answer["rewritten_query"])  # Shows the optimized query

print(answer["answer"])           # Answer from improved retrieval

```

The output demonstrates how the **rewritten query** generated by the node drives a second, richer retrieval pass through OpenSearch, resulting in higher-quality context for the final answer generation.

## Summary

- The query rewriting node in [`src/services/agents/nodes/rewrite_query_node.py`](https://github.com/jamwithai/production-agentic-rag-course/blob/main/src/services/agents/nodes/rewrite_query_node.py) transforms ambiguous inputs using structured LLM outputs defined in `QueryRewriteOutput` before re-attempting retrieval.
- **Robust fallback logic** ensures pipeline continuity by appending keywords like "research paper arxiv machine learning" when LLM calls fail.
- **LangGraph integration** via `workflow.add_edge("rewrite_query", "retrieve")` creates a feedback loop that routes optimized queries back to the retrieval node.
- **Langfuse tracing** provides observability into rewrite operations through the `query_rewriting` span, capturing model parameters and execution timing.
- By updating `state["rewritten_query"]` while preserving `state["original_query"]`, the node maintains audit trails while improving OpenSearch embedding search accuracy.

## Frequently Asked Questions

### What happens if the LLM fails to rewrite the query?

If the structured LLM call fails or returns invalid output, the node automatically falls back to a deterministic keyword expansion strategy. It constructs a new query by appending contextual terms—`"{original_question} research paper arxiv machine learning"`—ensuring the retrieval step always receives a valid search term even during service degradation.

### How does the query rewriting node integrate with LangGraph?

The node integrates as a discrete step in the LangGraph workflow defined in [`src/services/agents/agentic_rag.py`](https://github.com/jamwithai/production-agentic-rag-course/blob/main/src/services/agents/agentic_rag.py). It accepts the current `AgentState`, processes `state["original_query"]`, and returns updates to `state["rewritten_query"]`. The graph configuration uses `workflow.add_edge("rewrite_query", "retrieve")` to route the optimized query back to the retrieval node, creating a retry loop that improves results before answer generation.

### What prompt template does the node use for optimization?

The node uses `REWRITE_PROMPT` defined in [`src/services/agents/prompts.py`](https://github.com/jamwithai/production-agentic-rag-course/blob/main/src/services/agents/prompts.py) to guide the LLM. This template instructs the model to analyze the original query's intent, add missing temporal or domain context, and output both the rewritten query and explanatory reasoning through the `QueryRewriteOutput` schema defined in [`src/services/agents/models.py`](https://github.com/jamwithai/production-agentic-rag-course/blob/main/src/services/agents/models.py).

### When should query rewriting be triggered in a RAG pipeline?

According to the `jamwithai/production-agentic-rag-course` implementation, query rewriting triggers automatically when the document grading node (`grade_documents`) determines that initial retrieval results fall below the relevance threshold. This conditional routing ensures computationally expensive LLM rewrites only occur when the initial vector search fails to return sufficiently relevant documents from OpenSearch.