Praison AI

QueryRewriterAgent Module

Transform user queries to improve retrieval quality in RAG applications

QueryRewriterAgent Module

The QueryRewriterAgent class transforms user queries to improve retrieval quality in RAG applications.

Supported Strategies

  • BASIC: Simple rephrasing for clarity and keyword optimization
  • HYDE: Hypothetical Document Embeddings – generates a hypothetical answer
  • STEP_BACK: Generates higher-level concept questions for complex queries
  • SUB_QUERIES: Decomposes multi-part questions into focused sub-queries
  • MULTI_QUERY: Generates multiple paraphrased versions for ensemble retrieval
  • CONTEXTUAL: Uses conversation history to resolve references and context
  • AUTO: Automatically selects the best strategy

Import

python
from praisonaiagents import QueryRewriterAgent, RewriteStrategy

Quick Example

python
from praisonaiagents import QueryRewriterAgent, RewriteStrategy

agent = QueryRewriterAgent()

# Basic rewriting
result = agent.rewrite("AI trends")
print(result.rewritten_queries)

# HyDE for better semantic matching
result = agent.rewrite(
    "What is quantum computing?",
    strategy=RewriteStrategy.HYDE
)
print(result.hypothetical_document)

Constructor

QueryRewriterAgent()

Creates a new QueryRewriterAgent instance.

Parameters:

ParameterTypeDefaultDescription
namestr"Query Rewriter"Agent name
modelstr"gpt-4o-mini"LLM model to use
verboseboolFalseEnable verbose logging
toolslist[]Optional tools for context-aware rewriting

Methods

rewrite(query, strategy, chat_history)

Rewrites a query using the specified strategy.

Parameters:

  • query (str): The original user query
  • strategy (RewriteStrategy): Rewriting strategy to use (default: AUTO)
  • chat_history (list): Optional conversation history for contextual rewriting

Returns: RewriteResult – Contains rewritten queries and metadata

RewriteStrategy Enum

python
class RewriteStrategy(Enum):
    BASIC = "basic"
    HYDE = "hyde"
    STEP_BACK = "step_back"
    SUB_QUERIES = "sub_queries"
    MULTI_QUERY = "multi_query"
    CONTEXTUAL = "contextual"
    AUTO = "auto"

RewriteResult

python
@dataclass
class RewriteResult:
    original_query: str
    rewritten_queries: List[str]
    strategy_used: RewriteStrategy
    hypothetical_document: Optional[str] = None
    step_back_question: Optional[str] = None
    sub_queries: Optional[List[str]] = None
    metadata: Dict[str, Any]

    @property
    def primary_query(self) -> str:
        """Returns the primary rewritten query."""

    @property
    def all_queries(self) -> List[str]:
        """Returns all queries including original and rewritten."""

Examples

Contextual Rewriting with Chat History

python
from praisonaiagents import QueryRewriterAgent

agent = QueryRewriterAgent()

chat_history = [
    {"role": "user", "content": "Tell me about Python"},
    {"role": "assistant", "content": "Python is a programming language..."}
]

result = agent.rewrite(
    "What about its performance?",
    chat_history=chat_history
)
# Resolves "its" to "Python's"

With Search Tools

python
from praisonaiagents import QueryRewriterAgent
from praisonaiagents import internet_search

agent = QueryRewriterAgent(tools=[internet_search], )
result = agent.rewrite("latest AI developments")
# Searches first, then rewrites with context

Related